From 495165266f2d2e523cf8b144f61bf953ce06776b Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Sun, 20 Sep 2026 19:39:06 +0100 Subject: [PATCH] feat(kimodocpp): track usage through generic backend metadata (#12162) Report input tokens and frame-step output units in response metadata and record them through the existing usage accounting pipeline. Preserve the accounting rule and model-specific dimensions as JSON without extending the gRPC schema for each modality. Expose animation usage only under metadata.usage, validate counts before recording, and document the response contract and loaded-model location. Add coverage for transport, defaults, failures, persistence, and recording requests once with statistics enabled or disabled. Assisted-by: Codex:GPT-6 Signed-off-by: Richard Palethorpe --- backend/backend.proto | 3 + backend/go/kimodocpp/README.md | 13 ++ backend/go/kimodocpp/integration_test.go | 15 ++- backend/go/kimodocpp/kimodo.go | 82 +++++++++---- backend/go/kimodocpp/kimodo_test.go | 116 ++++++++++++++++++ backend/go/kimodocpp/main.go | 3 + backend/go/kimodocpp/native/bridge.cpp | 38 ++++++ core/backend/animation.go | 12 +- core/http/auth/usage.go | 3 + core/http/auth/usage_test.go | 12 ++ .../endpoints/localai/model3d_animation.go | 14 ++- .../model3d_animation_internal_test.go | 55 ++++++++- core/http/middleware/usage.go | 4 + core/http/middleware/usage_stamp.go | 23 +++- core/http/middleware/usage_test.go | 25 ++++ core/http/routes/localai.go | 1 + core/schema/openai.go | 3 +- docs/content/features/3d-animation.md | 87 +++++++++++++ docs/content/operations/activity.md | 5 + pkg/grpc/animation_usage_test.go | 51 ++++++++ pkg/grpc/interface.go | 6 + pkg/grpc/metadata/metadata.go | 68 ++++++++++ pkg/grpc/metadata/metadata_test.go | 55 +++++++++ pkg/grpc/server.go | 7 ++ 24 files changed, 657 insertions(+), 44 deletions(-) create mode 100644 pkg/grpc/animation_usage_test.go create mode 100644 pkg/grpc/metadata/metadata.go create mode 100644 pkg/grpc/metadata/metadata_test.go diff --git a/backend/backend.proto b/backend/backend.proto index 7fcfe183a..a27b153ca 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -539,6 +539,9 @@ message ProxyOptions { message Result { string message = 1; bool success = 2; + // JSON object containing arbitrary backend response metadata. Application + // conventions (such as metadata.usage) are independent of this schema. + bytes metadata = 4; } // EmbeddingLayout describes whether embeddings contains one final vector or diff --git a/backend/go/kimodocpp/README.md b/backend/go/kimodocpp/README.md index 79ce88e8c..e5e442651 100644 --- a/backend/go/kimodocpp/README.md +++ b/backend/go/kimodocpp/README.md @@ -33,3 +33,16 @@ the pinned upstream commit, verify the C ABI layout/version and all three skeleton families. Pin changes update a clean cached checkout automatically; local source modifications stop the update rather than being discarded. Preserve any such changes before using `make clean` to replace that generated checkout. + +## Response metadata + +`Animate3DWithMetadata` returns UTF-8 JSON bytes in the generic gRPC +`Result.metadata` field. Kimodo populates `usage.input_units` with text token +count (including BOS), `usage.output_units` with frames × sampling steps, and +`usage.accounting_rule` with `frame_steps_v1`. `usage.details` retains +`output_frames` and `sampling_steps`. There is no separate protobuf usage type. + +The HTTP handler returns this object under `metadata`, with usage only at +`metadata.usage`. Internal accounting reads those counts and records the request +once; no top-level HTTP usage summary is emitted. See the [usage accounting documentation](../../../docs/content/features/3d-animation.md#usage-accounting) +for the distinct backend and HTTP formats, validation, and persistence behavior. diff --git a/backend/go/kimodocpp/integration_test.go b/backend/go/kimodocpp/integration_test.go index b2ccd8dd3..ff54d7e99 100644 --- a/backend/go/kimodocpp/integration_test.go +++ b/backend/go/kimodocpp/integration_test.go @@ -5,6 +5,7 @@ import ( "os" "path/filepath" + "github.com/mudler/LocalAI/pkg/grpc/metadata" pb "github.com/mudler/LocalAI/pkg/grpc/proto" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -32,11 +33,19 @@ var _ = Describe("real-model generation", Label("real-models"), func() { for index := range 2 { path := filepath.Join(GinkgoT().TempDir(), "animation.glb") By("generating a clip with the existing session") - Expect(backend.Animate3D(&pb.Animate3DRequest{Dst: path, + data, err := backend.Animate3DWithMetadata(&pb.Animate3DRequest{Dst: path, Inputs: map[string]*pb.AnimationInput{"prompt": {Type: "text", Data: "A person walks forward."}}, Params: map[string]string{"frames": "60", "steps": "1", "seed": "42"}, - })).To(Succeed(), "clip %d", index) - data, err := os.ReadFile(path) + }) + Expect(err).NotTo(HaveOccurred(), "clip %d", index) + usage, err := metadata.ParseUsage(data) + Expect(err).NotTo(HaveOccurred()) + Expect(usage).NotTo(BeNil()) + Expect(usage.InputUnits).To(BeNumerically(">", 1)) + Expect(usage.OutputUnits).To(Equal(60)) + Expect(usage.Details).To(MatchJSON(`{"output_frames":60,"sampling_steps":1}`)) + Expect(usage.AccountingRule).To(Equal("frame_steps_v1")) + data, err = os.ReadFile(path) Expect(err).NotTo(HaveOccurred()) Expect(len(data)).To(BeNumerically(">", 1000)) Expect(string(data[:4])).To(Equal("glTF")) diff --git a/backend/go/kimodocpp/kimodo.go b/backend/go/kimodocpp/kimodo.go index c2f2a41bd..5d05071b8 100644 --- a/backend/go/kimodocpp/kimodo.go +++ b/backend/go/kimodocpp/kimodo.go @@ -3,6 +3,7 @@ package main import ( "bytes" + "encoding/json" "fmt" "maps" "math" @@ -16,6 +17,7 @@ import ( "unsafe" "github.com/mudler/LocalAI/pkg/grpc/base" + "github.com/mudler/LocalAI/pkg/grpc/metadata" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/xlog" ) @@ -31,26 +33,30 @@ type generationOptions struct { } var ( - nativeABI func() int32 - nativeLoad func(string, string, string, uintptr, *byte, int32) uintptr - nativeFree func(uintptr) - nativeGenerate func(uintptr, string, *generationOptions, *byte, int32) uintptr - nativeMotionFree func(uintptr) - nativeFrames func(uintptr) int32 - nativeJoints func(uintptr) int32 - nativeRotations func(uintptr) *float32 - nativeRoots func(uintptr) *float32 - nativeConfigure func(string, int32, int32) int32 - nativeJointName func(int32, int32) string - nativeJointParent func(int32, int32) int32 - nativeJointOffset func(int32, int32) *float32 + nativeABI func() int32 + nativeLoad func(string, string, string, uintptr, *byte, int32) uintptr + nativeFree func(uintptr) + nativeGenerate func(uintptr, string, *generationOptions, *byte, int32) uintptr + nativeMotionFree func(uintptr) + nativeFrames func(uintptr) int32 + nativeJoints func(uintptr) int32 + nativeRotations func(uintptr) *float32 + nativeRoots func(uintptr) *float32 + nativeConfigure func(string, int32, int32) int32 + nativeJointName func(int32, int32) string + nativeJointParent func(int32, int32) int32 + nativeJointOffset func(int32, int32) *float32 + nativeTokenizerLoad func(string, *byte, int32) uintptr + nativeTokenizerFree func(uintptr) + nativePromptTokens func(uintptr, string) int32 ) type Kimodo struct { base.Base - mu sync.Mutex - model uintptr - defaults map[string]string + mu sync.Mutex + model uintptr + tokenizer uintptr + defaults map[string]string } func parseGeneration(params map[string]string) (generationOptions, error) { @@ -152,9 +158,18 @@ func (k *Kimodo) Load(options *pb.ModelOptions) error { if loaded == 0 { return fmt.Errorf("loading kimodo: %s", nativeError(errorBuffer)) } + tokenizer := nativeTokenizerLoad(textBundle, &errorBuffer[0], int32(len(errorBuffer))) + if tokenizer == 0 { + nativeFree(loaded) + return fmt.Errorf("loading kimodo tokenizer: %s", nativeError(errorBuffer)) + } if k.model != 0 { nativeFree(k.model) } + if k.tokenizer != 0 { + nativeTokenizerFree(k.tokenizer) + } + k.tokenizer = tokenizer k.model, k.defaults = loaded, defaults xlog.Info("Kimodo loaded", "device", device, "threads", threads, "text_layer_chunk", chunk) return nil @@ -174,23 +189,32 @@ func (k *Kimodo) Free() error { nativeFree(k.model) k.model = 0 } + if k.tokenizer != 0 { + nativeTokenizerFree(k.tokenizer) + k.tokenizer = 0 + } return nil } func (k *Kimodo) Animate3D(request *pb.Animate3DRequest) error { + _, err := k.Animate3DWithMetadata(request) + return err +} + +func (k *Kimodo) Animate3DWithMetadata(request *pb.Animate3DRequest) ([]byte, error) { k.mu.Lock() defer k.mu.Unlock() if k.model == 0 { - return fmt.Errorf("kimodo model is not loaded") + return nil, fmt.Errorf("kimodo model is not loaded") } prompt := request.Inputs["prompt"] if len(request.Inputs) != 1 || prompt == nil || prompt.Type != "text" || strings.TrimSpace(prompt.Data) == "" || len(prompt.Data) > 4096 || !utf8.ValidString(prompt.Data) || strings.ContainsRune(prompt.Data, 0) { - return fmt.Errorf("kimodo requires one UTF-8 text prompt of 1..4096 bytes without NUL characters") + return nil, fmt.Errorf("kimodo requires one UTF-8 text prompt of 1..4096 bytes without NUL characters") } if request.Dst == "" { - return fmt.Errorf("animation destination is required") + return nil, fmt.Errorf("animation destination is required") } params := maps.Clone(k.defaults) if params == nil { @@ -199,29 +223,37 @@ func (k *Kimodo) Animate3D(request *pb.Animate3DRequest) error { maps.Copy(params, request.Params) options, err := parseGeneration(params) if err != nil { - return err + return nil, err + } + promptTokens := nativePromptTokens(k.tokenizer, prompt.Data) + if promptTokens < 2 || promptTokens > 512 { + return nil, fmt.Errorf("cannot count kimodo prompt tokens (expected 1..511 tokens excluding BOS)") } errorBuffer := make([]byte, 1024) motion := nativeGenerate(k.model, prompt.Data, &options, &errorBuffer[0], int32(len(errorBuffer))) if motion == 0 { - return fmt.Errorf("generating kimodo motion: %s", nativeError(errorBuffer)) + return nil, fmt.Errorf("generating kimodo motion: %s", nativeError(errorBuffer)) } defer nativeMotionFree(motion) frames, joints := nativeFrames(motion), nativeJoints(motion) if frames != int32(options.Frames) || (joints != 22 && joints != 30 && joints != 34) { - return fmt.Errorf("unexpected kimodo motion dimensions: %d frames, %d joints", frames, joints) + return nil, fmt.Errorf("unexpected kimodo motion dimensions: %d frames, %d joints", frames, joints) } roots, rotations := nativeRoots(motion), nativeRotations(motion) if roots == nil || rotations == nil { - return fmt.Errorf("kimodo returned empty motion buffers") + return nil, fmt.Errorf("kimodo returned empty motion buffers") } skeleton := make([]animationJoint, joints) for joint := range joints { offset := nativeJointOffset(joints, joint) if offset == nil { - return fmt.Errorf("missing skeleton joint %d", joint) + return nil, fmt.Errorf("missing skeleton joint %d", joint) } skeleton[joint] = animationJoint{Name: nativeJointName(joints, joint), Parent: int(nativeJointParent(joints, joint)), Offset: [3]float32(unsafe.Slice(offset, 3))} } - return writeAnimationGLB(request.Dst, unsafe.Slice(roots, int(frames)*3), unsafe.Slice(rotations, int(frames*joints)*4), skeleton) + if err := writeAnimationGLB(request.Dst, unsafe.Slice(roots, int(frames)*3), unsafe.Slice(rotations, int(frames*joints)*4), skeleton); err != nil { + return nil, err + } + details, _ := json.Marshal(map[string]int32{"output_frames": frames, "sampling_steps": int32(options.Steps)}) + return metadata.EncodeUsage(metadata.Usage{InputUnits: int(promptTokens), OutputUnits: int(frames) * int(options.Steps), AccountingRule: "frame_steps_v1", Details: details}) } diff --git a/backend/go/kimodocpp/kimodo_test.go b/backend/go/kimodocpp/kimodo_test.go index abd38dc52..282c78a2e 100644 --- a/backend/go/kimodocpp/kimodo_test.go +++ b/backend/go/kimodocpp/kimodo_test.go @@ -4,12 +4,14 @@ package main import ( "encoding/binary" "encoding/json" + "fmt" "math" "os" "path/filepath" "testing" "unsafe" + "github.com/mudler/LocalAI/pkg/grpc/metadata" pb "github.com/mudler/LocalAI/pkg/grpc/proto" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -53,6 +55,10 @@ var _ = Describe("text encoder loading", func() { var textPath string BeforeEach(func() { configure, load, free := nativeConfigure, nativeLoad, nativeFree + loadTokenizer, freeTokenizer := nativeTokenizerLoad, nativeTokenizerFree + DeferCleanup(func() { nativeTokenizerLoad, nativeTokenizerFree = loadTokenizer, freeTokenizer }) + nativeTokenizerLoad = func(string, *byte, int32) uintptr { return 3 } + nativeTokenizerFree = func(uintptr) {} DeferCleanup(func() { nativeConfigure, nativeLoad, nativeFree = configure, load, free }) chunk, textPath = 0, "" nativeConfigure = func(_ string, threads, layers int32) int32 { @@ -74,6 +80,31 @@ var _ = Describe("text encoder loading", func() { Expect(chunk).To(Equal(int32(32))) Expect(textPath).To(Equal(filepath.Join("/models", text))) }, Entry("Q8 monolith", "kimodo/text/Llama-3-Kimodo-Q8_0.gguf"), Entry("low-bit monolith", "kimodo/text/Llama-3-Kimodo-Q4_K_M.gguf"), Entry("legacy directory", "kimodo/text")) + It("preserves the previous session and releases the new model if tokenizer loading fails", func() { + backend := &Kimodo{model: 10, tokenizer: 30} + var released []uintptr + nativeFree = func(handle uintptr) { released = append(released, handle) } + nativeTokenizerLoad = func(string, *byte, int32) uintptr { return 0 } + Expect(backend.Load(&pb.ModelOptions{ModelFile: "motion.gguf", Threads: 8, + Options: []string{"text_bundle:encoder.gguf"}})).To(MatchError(ContainSubstring("loading kimodo tokenizer"))) + Expect(released).To(Equal([]uintptr{1})) + Expect(backend.model).To(Equal(uintptr(10))) + Expect(backend.tokenizer).To(Equal(uintptr(30))) + }) + It("releases the old model and tokenizer on reload and frees the new pair only once", func() { + backend := &Kimodo{model: 10, tokenizer: 30} + var models, tokenizers []uintptr + nativeFree = func(handle uintptr) { models = append(models, handle) } + nativeTokenizerFree = func(handle uintptr) { tokenizers = append(tokenizers, handle) } + Expect(backend.Load(&pb.ModelOptions{ModelFile: "motion.gguf", Threads: 8, + Options: []string{"text_bundle:encoder.gguf"}})).To(Succeed()) + Expect(models).To(Equal([]uintptr{10})) + Expect(tokenizers).To(Equal([]uintptr{30})) + Expect(backend.Free()).To(Succeed()) + Expect(backend.Free()).To(Succeed()) + Expect(models).To(Equal([]uintptr{10, 1})) + Expect(tokenizers).To(Equal([]uintptr{30, 3})) + }) DescribeTable("honors the layer residency option", func(value string, expected int32) { backend := &Kimodo{} DeferCleanup(backend.Free) @@ -136,6 +167,9 @@ var _ = Describe("skeleton GLB export", func() { var _ = Describe("native resource lifetime", func() { It("releases the native motion when exporting fails", func() { + oldCount := nativePromptTokens + DeferCleanup(func() { nativePromptTokens = oldCount }) + nativePromptTokens = func(uintptr, string) int32 { return 3 } oldGenerate, oldFree := nativeGenerate, nativeMotionFree oldFrames, oldJoints := nativeFrames, nativeJoints DeferCleanup(func() { @@ -153,3 +187,85 @@ var _ = Describe("native resource lifetime", func() { Expect(freed).To(BeTrue()) }) }) + +var _ = Describe("animation usage", func() { + var backend *Kimodo + var request *pb.Animate3DRequest + var freed bool + BeforeEach(func() { + count, generate, free := nativePromptTokens, nativeGenerate, nativeMotionFree + frames, joints, roots, rotations := nativeFrames, nativeJoints, nativeRoots, nativeRotations + name, parent, offset := nativeJointName, nativeJointParent, nativeJointOffset + DeferCleanup(func() { + nativePromptTokens, nativeGenerate, nativeMotionFree = count, generate, free + nativeFrames, nativeJoints, nativeRoots, nativeRotations = frames, joints, roots, rotations + nativeJointName, nativeJointParent, nativeJointOffset = name, parent, offset + }) + backend = &Kimodo{model: 1, tokenizer: 3} + request = &pb.Animate3DRequest{Dst: filepath.Join(GinkgoT().TempDir(), "clip.glb"), + Inputs: map[string]*pb.AnimationInput{"prompt": {Type: "text", Data: "A person walks forward."}}} + freed = false + var generatedFrames int32 + nativePromptTokens = func(handle uintptr, text string) int32 { + Expect(handle).To(Equal(uintptr(3))) + Expect(text).To(Equal(request.Inputs["prompt"].Data)) + return 7 + } + nativeGenerate = func(_ uintptr, _ string, options *generationOptions, _ *byte, _ int32) uintptr { + generatedFrames = int32(options.Frames) + return 2 + } + nativeMotionFree = func(handle uintptr) { Expect(handle).To(Equal(uintptr(2))); freed = true } + nativeFrames = func(uintptr) int32 { return generatedFrames } + nativeJoints = func(uintptr) int32 { return 22 } + rootBuffer := make([]float32, 150*3) + rotationBuffer := make([]float32, 150*22*4) + for i := 3; i < len(rotationBuffer); i += 4 { + rotationBuffer[i] = 1 + } + nativeRoots = func(uintptr) *float32 { return &rootBuffer[0] } + nativeRotations = func(uintptr) *float32 { return &rotationBuffer[0] } + nativeJointName = func(int32, int32) string { return "joint" } + nativeJointParent = func(_ int32, joint int32) int32 { return joint - 1 } + nativeJointOffset = func(int32, int32) *float32 { return &rootBuffer[0] } + }) + DescribeTable("reports effective frame-steps and actual prompt tokens", func(defaults, params map[string]string, frames, steps int32) { + backend.defaults, request.Params = defaults, params + data, err := backend.Animate3DWithMetadata(request) + Expect(err).NotTo(HaveOccurred()) + usage, err := metadata.ParseUsage(data) + Expect(err).NotTo(HaveOccurred()) + Expect(usage).NotTo(BeNil()) + Expect(usage.InputUnits).To(Equal(7)) + Expect(usage.OutputUnits).To(Equal(int(frames * steps))) + Expect(usage.Details).To(MatchJSON(fmt.Sprintf(`{"output_frames":%d,"sampling_steps":%d}`, frames, steps))) + Expect(usage.AccountingRule).To(Equal("frame_steps_v1")) + Expect(freed).To(BeTrue()) + }, Entry("reference defaults", nil, nil, int32(150), int32(100)), + Entry("model defaults", map[string]string{"frames": "90", "steps": "20"}, nil, int32(90), int32(20)), + Entry("request overrides", map[string]string{"frames": "90", "steps": "20"}, map[string]string{"frames": "60", "steps": "1"}, int32(60), int32(1)), + Entry("maximum units", nil, map[string]string{"steps": "1000"}, int32(150), int32(1000))) + It("does not report usage when tokenization fails", func() { + nativePromptTokens = func(uintptr, string) int32 { return -1 } + nativeGenerate = func(uintptr, string, *generationOptions, *byte, int32) uintptr { + Fail("inference must not run") + return 0 + } + data, err := backend.Animate3DWithMetadata(request) + Expect(err).To(HaveOccurred()) + Expect(data).To(BeNil()) + }) + It("does not report usage when native generation fails", func() { + nativeGenerate = func(uintptr, string, *generationOptions, *byte, int32) uintptr { return 0 } + data, err := backend.Animate3DWithMetadata(request) + Expect(err).To(HaveOccurred()) + Expect(data).To(BeNil()) + }) + It("does not report usage when GLB export fails", func() { + request.Dst = filepath.Join(GinkgoT().TempDir(), "missing", "clip.glb") + data, err := backend.Animate3DWithMetadata(request) + Expect(err).To(HaveOccurred()) + Expect(data).To(BeNil()) + Expect(freed).To(BeTrue()) + }) +}) diff --git a/backend/go/kimodocpp/main.go b/backend/go/kimodocpp/main.go index 1b05d5d4b..04c0667cc 100644 --- a/backend/go/kimodocpp/main.go +++ b/backend/go/kimodocpp/main.go @@ -43,6 +43,9 @@ func loadNativeLibrary(path string) error { {&nativeJointName, "localai_kimodo_joint_name"}, {&nativeJointParent, "localai_kimodo_joint_parent"}, {&nativeJointOffset, "localai_kimodo_joint_offset"}, + {&nativeTokenizerLoad, "localai_kimodo_tokenizer_load"}, + {&nativeTokenizerFree, "localai_kimodo_tokenizer_free"}, + {&nativePromptTokens, "localai_kimodo_prompt_tokens"}, } { purego.RegisterLibFunc(binding.function, lib, binding.name) } diff --git a/backend/go/kimodocpp/native/bridge.cpp b/backend/go/kimodocpp/native/bridge.cpp index 22ca7478c..8cc1ee630 100644 --- a/backend/go/kimodocpp/native/bridge.cpp +++ b/backend/go/kimodocpp/native/bridge.cpp @@ -1,10 +1,12 @@ // SPDX-License-Identifier: MIT #include #include "../sources/kimodo.cpp/src/skeleton.hpp" +#include "../sources/kimodo.cpp/src/llm_tokenizer.hpp" #include #include #include +#include #ifdef KIMODO_HAVE_GGML_VULKAN #include @@ -22,6 +24,42 @@ const kimodo::detail::skeleton_spec *skeleton(int joints) { } extern "C" { +// Use the encoder's tokenizer implementation and vocabulary, including BOS. +// Keep it loaded so accounting does not reload the vocabulary on every call. +KIMODO_API void *localai_kimodo_tokenizer_load(const char *source, char *error, int capacity) { + try { + auto path = std::filesystem::path(source); + if (!std::filesystem::is_directory(path)) path = path.parent_path(); + auto tokenizer = kimodo::detail::llm_tokenizer::load((path / "tokenizer.gguf").string()); + if (!tokenizer) { + if (error && capacity > 0) std::snprintf(error, capacity, "%s", tokenizer.error().c_str()); + return nullptr; + } + return tokenizer->release(); + } catch (const std::exception &e) { + if (error && capacity > 0) std::snprintf(error, capacity, "%s", e.what()); + return nullptr; + } catch (...) { + if (error && capacity > 0) std::snprintf(error, capacity, "unknown tokenizer error"); + return nullptr; + } +} + +KIMODO_API void localai_kimodo_tokenizer_free(void *tokenizer) { + delete static_cast(tokenizer); +} + +KIMODO_API int localai_kimodo_prompt_tokens(void *tokenizer, const char *prompt) { + if (!tokenizer || !prompt) return -1; + try { + auto ids = static_cast(tokenizer)->encode(prompt); + if (!ids || ids->size() < 2 || ids->size() > 512) return -1; + return static_cast(ids->size()); + } catch (...) { + return -1; + } +} + // The upstream runtime currently reads environment variables instead of its // C API runtime_options. Set the native environment before creating a session. KIMODO_API int localai_kimodo_configure(const char *device, int threads, int chunk) { diff --git a/core/backend/animation.go b/core/backend/animation.go index b29986fa1..530a86a85 100644 --- a/core/backend/animation.go +++ b/core/backend/animation.go @@ -12,15 +12,15 @@ import ( "github.com/mudler/LocalAI/pkg/model" ) -func Model3DAnimation(ctx context.Context, request *proto.Animate3DRequest, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (err error) { +func Model3DAnimation(ctx context.Context, request *proto.Animate3DRequest, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (responseMetadata []byte, err error) { inferenceModel, err := loader.Load(ModelOptions(modelConfig, appConfig)...) if err != nil { recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil) - return err + return nil, err } release, err := AcquireGlobalBackendSlot() if err != nil { - return err + return nil, err } defer release() if appConfig.EnableTracing { @@ -40,10 +40,10 @@ func Model3DAnimation(ctx context.Context, request *proto.Animate3DRequest, load request.ModelIdentity = modelConfig.Model result, err := inferenceModel.Animate3D(ctx, request) if err != nil { - return err + return nil, err } if result == nil || !result.Success { - return fmt.Errorf("animation backend failed: %s", result.GetMessage()) + return nil, fmt.Errorf("animation backend failed: %s", result.GetMessage()) } - return nil + return result.Metadata, nil } diff --git a/core/http/auth/usage.go b/core/http/auth/usage.go index 1ab308d52..0cceed697 100644 --- a/core/http/auth/usage.go +++ b/core/http/auth/usage.go @@ -50,6 +50,9 @@ type UsageRecord struct { Duration int64 // milliseconds CreatedAt time.Time `gorm:"index:idx_usage_user_time"` + // Metadata preserves the backend accounting convention and its original dimensions. + Metadata string `gorm:"type:text"` + // Routing extension fields. Nullable / zero-valued for legacy rows. RequestedModel string `gorm:"size:255;index"` ServedModel string `gorm:"size:255;index"` diff --git a/core/http/auth/usage_test.go b/core/http/auth/usage_test.go index bb3dc945d..faa07254c 100644 --- a/core/http/auth/usage_test.go +++ b/core/http/auth/usage_test.go @@ -14,6 +14,18 @@ import ( var _ = Describe("Usage", func() { Describe("RecordUsage", func() { + It("persists animation accounting dimensions with token-compatible units", func() { + db := testDB() + record := &auth.UsageRecord{UserID: "user-1", Model: "kimodo", Endpoint: "/3d/animate", + PromptTokens: 32, CompletionTokens: 15000, TotalTokens: 15032, + Metadata: `{"usage":{"input_units":32,"output_units":15000,"details":{"output_frames":150,"sampling_steps":100}}}`, CreatedAt: time.Now()} + Expect(auth.RecordUsage(db, record)).To(Succeed()) + var stored auth.UsageRecord + Expect(db.First(&stored, record.ID).Error).NotTo(HaveOccurred()) + Expect(stored.PromptTokens).To(Equal(int64(32))) + Expect(stored.CompletionTokens).To(Equal(int64(15000))) + Expect(stored.Metadata).To(Equal(record.Metadata)) + }) It("inserts a usage record", func() { db := testDB() record := &auth.UsageRecord{ diff --git a/core/http/endpoints/localai/model3d_animation.go b/core/http/endpoints/localai/model3d_animation.go index 8eb46b327..be4cf7678 100644 --- a/core/http/endpoints/localai/model3d_animation.go +++ b/core/http/endpoints/localai/model3d_animation.go @@ -3,6 +3,7 @@ package localai import ( "encoding/base64" + "encoding/json" "fmt" "net/http" "net/url" @@ -20,6 +21,7 @@ import ( "github.com/mudler/LocalAI/core/schema" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/xlog" ) func validateAnimationRequest(input *schema.Model3DAnimationRequest, cfg *config.ModelConfig) error { @@ -128,7 +130,8 @@ func Model3DAnimationEndpoint(ml *model.ModelLoader, appConfig *config.Applicati return err } request.Dst = file.Name() - if err := backend.Model3DAnimation(c.Request().Context(), request, ml, *cfg, appConfig); err != nil { + responseMetadata, err := backend.Model3DAnimation(c.Request().Context(), request, ml, *cfg, appConfig) + if err != nil { return mapBackendError(err) } item := schema.Item{} @@ -145,6 +148,13 @@ func Model3DAnimationEndpoint(ml *model.ModelLoader, appConfig *config.Applicati } preserve = true } - return c.JSON(http.StatusOK, schema.OpenAIResponse{ID: uuid.NewString(), Created: int(time.Now().Unix()), Data: []schema.Item{item}}) + response := schema.OpenAIResponse{ID: uuid.NewString(), Model: input.Model, Created: int(time.Now().Unix()), Data: []schema.Item{item}} + metadataErr := middleware.StampResponseMetadata(c, input.Model, responseMetadata) + if metadataErr != nil { + xlog.Warn("ignoring invalid animation response metadata", "model", input.Model, "error", metadataErr) + } else { + response.Metadata = json.RawMessage(responseMetadata) + } + return c.JSON(http.StatusOK, response) } } diff --git a/core/http/endpoints/localai/model3d_animation_internal_test.go b/core/http/endpoints/localai/model3d_animation_internal_test.go index 94332eeab..52b04e709 100644 --- a/core/http/endpoints/localai/model3d_animation_internal_test.go +++ b/core/http/endpoints/localai/model3d_animation_internal_test.go @@ -13,8 +13,10 @@ import ( "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/auth" "github.com/mudler/LocalAI/core/http/middleware" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/routing/billing" grpcPkg "github.com/mudler/LocalAI/pkg/grpc" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/model" @@ -26,8 +28,9 @@ import ( type animationEndpointBackend struct { grpcPkg.Backend - success bool - seen *pb.Animate3DRequest + metadata []byte + success bool + seen *pb.Animate3DRequest } func (*animationEndpointBackend) HealthCheck(context.Context) (bool, error) { return true, nil } @@ -37,14 +40,18 @@ func (b *animationEndpointBackend) Animate3D(_ context.Context, request *pb.Anim if err := os.WriteFile(request.Dst, []byte("glTF fixture"), 0o600); err != nil { return nil, err } - return &pb.Result{Success: b.success, Message: "fixture failure"}, nil + return &pb.Result{Success: b.success, Metadata: b.metadata, Message: "fixture failure"}, nil } +const animationTestMetadata = `{"usage":{"input_units":32,"output_units":15000,"accounting_rule":"frame_steps_v1","details":{"output_frames":150,"sampling_steps":100}},"custom":true}` + var _ = Describe("3D animation HTTP output", func() { - DescribeTable("returns an asset and cleans temporary output on base64 or backend failure", func(format string, success bool) { + DescribeTable("returns an asset and cleans temporary output on base64 or backend failure", func(format string, success bool, responseMetadata string, statsEnabled bool) { + reportUsage := responseMetadata == animationTestMetadata state := &system.SystemState{} loader := model.NewModelLoader(state) fixture := &animationEndpointBackend{success: success} + fixture.metadata = []byte(responseMetadata) loader.SetModelRouter(func(_ context.Context, id string, _, _, _, _ string, _ *pb.ModelOptions, _ bool) (*model.Model, error) { return model.NewModelWithClient(id, "test://animation", fixture), nil }) @@ -61,7 +68,25 @@ var _ = Describe("3D animation HTTP output", func() { c := e.NewContext(httptest.NewRequest(http.MethodPost, "/3d/animate", nil), recorder) c.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, request) c.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, cfg) - err := Model3DAnimationEndpoint(loader, appConfig)(c) + stats := billing.NewMemoryBackend(10) + DeferCleanup(func() { Expect(stats.Close()).To(Succeed()) }) + var statsRecorder *billing.Recorder + if statsEnabled { + statsRecorder = billing.NewRecorder(stats) + } + handler := middleware.UsageMiddleware(statsRecorder, &auth.User{ID: "local"})(Model3DAnimationEndpoint(loader, appConfig)) + err := handler(c) + records, statsErr := stats.Aggregate(context.Background(), billing.AggregateQuery{UserID: "local", Period: "all"}) + Expect(statsErr).NotTo(HaveOccurred()) + if success && reportUsage && statsEnabled { + Expect(records).To(HaveLen(1)) + Expect(records[0].RequestCount).To(Equal(int64(1))) + Expect(records[0].PromptTokens).To(Equal(int64(32))) + Expect(records[0].CompletionTokens).To(Equal(int64(15000))) + Expect(records[0].TotalTokens).To(Equal(int64(15032))) + } else { + Expect(records).To(BeEmpty()) + } Expect(fixture.seen).NotTo(BeNil()) Expect(fixture.seen.ModelIdentity).To(Equal("motion.gguf")) files, readErr := filepath.Glob(filepath.Join(appConfig.GeneratedContentDir, "3d", "*")) @@ -69,6 +94,7 @@ var _ = Describe("3D animation HTTP output", func() { if !success { Expect(err).To(HaveOccurred()) Expect(files).To(BeEmpty()) + Expect(c.Get(middleware.ContextKeyPromptTokens)).To(BeNil()) return } Expect(err).NotTo(HaveOccurred()) @@ -76,6 +102,21 @@ var _ = Describe("3D animation HTTP output", func() { var response schema.OpenAIResponse Expect(json.Unmarshal(recorder.Body.Bytes(), &response)).To(Succeed()) Expect(response.Data).To(HaveLen(1)) + Expect(response.Model).To(Equal("motion")) + var body map[string]json.RawMessage + Expect(json.Unmarshal(recorder.Body.Bytes(), &body)).To(Succeed()) + Expect(body).NotTo(HaveKey("usage")) + if reportUsage { + Expect(response.Metadata).To(MatchJSON(fixture.metadata)) + Expect(c.Get(middleware.ContextKeyCompletionTokens)).To(Equal(int64(15000))) + } else { + if responseMetadata == `{"custom":true}` { + Expect(response.Metadata).To(MatchJSON(responseMetadata)) + } else { + Expect(response.Metadata).To(BeEmpty()) + } + Expect(c.Get(middleware.ContextKeyPromptTokens)).To(BeNil()) + } if format == "b64_json" { Expect(response.Data[0].B64JSON).To(Equal(base64.StdEncoding.EncodeToString([]byte("glTF fixture")))) Expect(files).To(BeEmpty()) @@ -83,7 +124,9 @@ var _ = Describe("3D animation HTTP output", func() { Expect(response.Data[0].URL).To(ContainSubstring("/generated-3d/animation-")) Expect(files).To(HaveLen(1)) } - }, Entry("URL", "url", true), Entry("base64", "b64_json", true), Entry("backend failure", "url", false)) + }, Entry("URL", "url", true, animationTestMetadata, true), Entry("base64", "b64_json", true, animationTestMetadata, true), Entry("backend failure", "url", false, animationTestMetadata, true), Entry("no metadata", "url", true, "", true), + Entry("unrelated metadata", "url", true, `{"custom":true}`, true), Entry("invalid JSON", "url", true, `{`, true), Entry("invalid units", "url", true, `{"usage":{"input_units":1,"output_units":-1}}`, true), + Entry("statistics disabled", "url", true, animationTestMetadata, false)) }) var _ = Describe("3D animation validation", func() { diff --git a/core/http/middleware/usage.go b/core/http/middleware/usage.go index ff0820db6..e383a28df 100644 --- a/core/http/middleware/usage.go +++ b/core/http/middleware/usage.go @@ -129,6 +129,10 @@ func UsageMiddleware(recorder *billing.Recorder, fallbackUser *auth.User) echo.M CorrelationID: correlationIDFromContext(c), } + if data, ok := c.Get(responseMetadataKey).(string); ok { + record.Metadata = data + } + if key := auth.GetAPIKey(c); key != nil { id := key.ID record.APIKeyID = &id diff --git a/core/http/middleware/usage_stamp.go b/core/http/middleware/usage_stamp.go index 7e82ab744..642619902 100644 --- a/core/http/middleware/usage_stamp.go +++ b/core/http/middleware/usage_stamp.go @@ -1,6 +1,27 @@ package middleware -import "github.com/labstack/echo/v4" +import ( + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/pkg/grpc/metadata" +) + +const responseMetadataKey = "localai_response_metadata" + +// StampResponseMetadata extracts the shared usage convention while preserving +// arbitrary metadata. Invalid counts never enter the accounting pipeline. +func StampResponseMetadata(c echo.Context, model string, data []byte) error { + usage, err := metadata.ParseUsage(data) + if err != nil { + return err + } + if c != nil { + c.Set(responseMetadataKey, string(data)) + if usage != nil { + StampUsage(c, model, usage.InputUnits, usage.OutputUnits) + } + } + return nil +} // StampUsage records the canonical token counts on the echo context so // UsageMiddleware can attribute the request without parsing the response diff --git a/core/http/middleware/usage_test.go b/core/http/middleware/usage_test.go index e95110988..c314bf9ed 100644 --- a/core/http/middleware/usage_test.go +++ b/core/http/middleware/usage_test.go @@ -33,6 +33,31 @@ func (c *captureBackend) Aggregate(_ context.Context, _ billing.AggregateQuery) func (c *captureBackend) Close() error { return nil } var _ = Describe("UsageMiddleware", func() { + DescribeTable("records animation units and their accounting basis only on success", func(status int) { + cap := &captureBackend{} + e := echo.New() + e.POST("/3d/animate", func(c echo.Context) error { + err := httpMiddleware.StampResponseMetadata(c, "kimodo", []byte(`{"usage":{"input_units":32,"output_units":15000,"accounting_rule":"frame_steps_v1","details":{"output_frames":150,"sampling_steps":100}},"custom":true}`)) + Expect(err).NotTo(HaveOccurred()) + return c.JSON(status, map[string]string{"model": "kimodo"}) + }, httpMiddleware.UsageMiddleware(billing.NewRecorder(cap), &auth.User{ID: "local"})) + w := httptest.NewRecorder() + e.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/3d/animate", nil)) + Expect(w.Code).To(Equal(status)) + if status != http.StatusOK { + Expect(cap.records).To(BeEmpty()) + return + } + Expect(cap.records).To(HaveLen(1)) + r := cap.records[0] + Expect(r.Model).To(Equal("kimodo")) + Expect(r.Endpoint).To(Equal("/3d/animate")) + Expect(r.PromptTokens).To(Equal(int64(32))) + Expect(r.CompletionTokens).To(Equal(int64(15000))) + Expect(r.TotalTokens).To(Equal(int64(15032))) + Expect(r.Metadata).To(MatchJSON(`{"usage":{"input_units":32,"output_units":15000,"accounting_rule":"frame_steps_v1","details":{"output_frames":150,"sampling_steps":100}},"custom":true}`)) + }, Entry("success", http.StatusOK), Entry("failure", http.StatusInternalServerError)) + mockChat := func(usage string) echo.HandlerFunc { return func(c echo.Context) error { c.Response().Header().Set("Content-Type", "application/json") diff --git a/core/http/routes/localai.go b/core/http/routes/localai.go index 0b1325455..601d71c6c 100644 --- a/core/http/routes/localai.go +++ b/core/http/routes/localai.go @@ -231,6 +231,7 @@ func RegisterLocalAIRoutes(router *echo.Echo, model3dHandler := localai.Model3DEndpoint(cl, ml, appConfig) router.POST("/3d/animate", localai.Model3DAnimationEndpoint(ml, appConfig), + middleware.UsageMiddleware(app.StatsRecorder(), app.FallbackUser()), echomiddleware.BodyLimit("45M"), requestExtractor.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_3D_ANIMATION)), requestExtractor.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.Model3DAnimationRequest) })) diff --git a/core/schema/openai.go b/core/schema/openai.go index e6394a602..2aa69969b 100644 --- a/core/schema/openai.go +++ b/core/schema/openai.go @@ -98,7 +98,8 @@ type OpenAIResponse struct { // `"usage":{"prompt_tokens":0,...}` on every chunk and break // OpenAI-SDK consumers that filter on a truthy `result.usage` // (continuedev/continue, Kilo Code, Roo Code, etc.). - Usage *OpenAIUsage `json:"usage,omitempty"` + Usage *OpenAIUsage `json:"usage,omitempty"` + Metadata json.RawMessage `json:"metadata,omitempty"` } // StreamOptions mirrors OpenAI's `stream_options` request field. The only diff --git a/docs/content/features/3d-animation.md b/docs/content/features/3d-animation.md index 5e2823eac..bf14df4e6 100644 --- a/docs/content/features/3d-animation.md +++ b/docs/content/features/3d-animation.md @@ -124,3 +124,90 @@ Kimodo accepts one prompt with 60–150 frames at 30 FPS. The reference defaults are 150 frames, 100 sampling steps, and text guidance 2. Parameters are strings; unsupported inputs and parameters are rejected. Multi-prompt transitions and mesh retargeting are not currently exposed by this adapter. + +## Usage accounting + +Kimodo reports one set of usage measurements through **generic backend metadata**. +The gRPC `Result` has a `metadata` bytes field containing a JSON object; there is +no separate gRPC usage field or model-specific usage message. + +### Backend metadata + +The JSON stored in `Result.metadata` looks like this: + +```json +{ + "usage": { + "input_units": 32, + "output_units": 15000, + "accounting_rule": "frame_steps_v1", + "details": { + "output_frames": 150, + "sampling_steps": 100 + } + } +} +``` + +`usage` is an application convention **inside** the generic metadata object. +Other metadata keys and model-specific `details` can be added without changing +the gRPC definitions or database schema. + +For kimodo, input units are the tokens produced by its text encoder's tokenizer, +including the beginning-of-sequence (BOS) token. Output units are actual generated +frames multiplied by effective sampling steps, after applying model defaults and +request overrides. In the example, 150 frames × 100 steps = 15,000 output units. + +### HTTP response and recorded usage + +For `POST /3d/animate`, LocalAI returns the backend object under `metadata`. +Usage appears only at `metadata.usage`; there is no top-level `usage` summary. +The following is an excerpt from either a `url` or `b64_json` response: + +```json +{ + "metadata": { + "usage": { + "input_units": 32, + "output_units": 15000, + "accounting_rule": "frame_steps_v1", + "details": { + "output_frames": 150, + "sampling_steps": 100 + } + } + } +} +``` + +Internal accounting reads `metadata.usage` directly and records each successful +request once. The existing token-named database columns store the unit counts: + +| Response metadata | Usage record column | +|---|---| +| `metadata.usage.input_units` | `PromptTokens` | +| `metadata.usage.output_units` | `CompletionTokens` | +| Sum of input and output units | `TotalTokens` | + +The record's single `Metadata` column retains the complete backend JSON, +including the rule and frame/step breakdown. With statistics disabled, response +metadata is still returned but no usage record is created. + +Both unit counts must be present, nonnegative integers whose sum fits a Go `int`. +Explicit zero counts are valid. If metadata is absent, the response omits +`metadata`. Metadata without a `usage` member is returned, but creates no usage +record. Malformed metadata or invalid usage counts cause metadata to be omitted +and a warning to be logged; the generated asset is still returned. Failed +generations produce no usage record. Counts are never estimated for backends +that do not report them. + +### Pricing units + +Kimodo's output units are **frame-steps**, not text tokens. They approximate +computational work rather than runtime or hardware cost. A price per million +output tokens therefore means dollars per million frame-steps for this model: + +`cost = (input_units × input_price + output_units × output_price) / 1,000,000` + +This describes how to interpret the existing token-based pricing dimensions; +it does not set a price or charge money automatically. diff --git a/docs/content/operations/activity.md b/docs/content/operations/activity.md index f6172c329..cc4977df3 100644 --- a/docs/content/operations/activity.md +++ b/docs/content/operations/activity.md @@ -12,6 +12,11 @@ of the app, and the **Activity** page, which holds the full picture. Both are admin-only. +To see models currently loaded for inference, open **Home** (`/app`) and scroll +to **Loaded models**. That list includes each model’s backend and a stop control. +A loaded model can be idle; Activity reports management jobs rather than whether +a model is resident or generating output. + ## The operations strip While there is anything to report, a single line appears at the top of the app. diff --git a/pkg/grpc/animation_usage_test.go b/pkg/grpc/animation_usage_test.go new file mode 100644 index 000000000..b1aeacab3 --- /dev/null +++ b/pkg/grpc/animation_usage_test.go @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: MIT +package grpc + +import ( + "context" + "errors" + + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "google.golang.org/protobuf/proto" +) + +type usageAnimationBackend struct { + modalityBackend + metadata []byte + err error +} + +func (b *usageAnimationBackend) Animate3DWithMetadata(*pb.Animate3DRequest) ([]byte, error) { + return b.metadata, b.err +} + +var _ = Describe("animation usage transport", func() { + It("preserves usage across protobuf serialization without calling legacy inference", func() { + backend := &usageAnimationBackend{metadata: []byte(`{"usage":{"input_units":32,"output_units":15000},"custom":{"value":true}}`)} + result, err := (&server{llm: backend}).Animate3D(context.Background(), &pb.Animate3DRequest{}) + Expect(err).NotTo(HaveOccurred()) + Expect(backend.served).To(BeZero()) + data, err := proto.Marshal(result) + Expect(err).NotTo(HaveOccurred()) + decoded := &pb.Result{} + Expect(proto.Unmarshal(data, decoded)).To(Succeed()) + Expect(decoded.Success).To(BeTrue()) + Expect(decoded.Metadata).To(Equal(backend.metadata)) + }) + It("propagates failures without usage", func() { + backend := &usageAnimationBackend{err: errors.New("generation failed")} + result, err := (&server{llm: backend}).Animate3D(context.Background(), &pb.Animate3DRequest{}) + Expect(err).To(MatchError("generation failed")) + Expect(result).To(BeNil()) + }) + It("supports backends without usage reporting", func() { + backend := &modalityBackend{} + result, err := (&server{llm: backend}).Animate3D(context.Background(), &pb.Animate3DRequest{}) + Expect(err).NotTo(HaveOccurred()) + Expect(result.Success).To(BeTrue()) + Expect(result.Metadata).To(BeEmpty()) + Expect(backend.served).To(Equal(1)) + }) +}) diff --git a/pkg/grpc/interface.go b/pkg/grpc/interface.go index 8370602a3..715a0acf5 100644 --- a/pkg/grpc/interface.go +++ b/pkg/grpc/interface.go @@ -6,6 +6,12 @@ import ( pb "github.com/mudler/LocalAI/pkg/grpc/proto" ) +// AnimationMetadataModel optionally reports JSON metadata without changing the legacy +// animation interface implemented by other backends. +type AnimationMetadataModel interface { + Animate3DWithMetadata(*pb.Animate3DRequest) ([]byte, error) +} + type AIModel interface { Busy() bool Lock() diff --git a/pkg/grpc/metadata/metadata.go b/pkg/grpc/metadata/metadata.go new file mode 100644 index 000000000..067974028 --- /dev/null +++ b/pkg/grpc/metadata/metadata.go @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: MIT +// Package metadata defines application conventions for opaque backend metadata. +// New conventions do not require changes to the gRPC transport schema. +package metadata + +import ( + "encoding/json" + "fmt" + "math" +) + +type Usage struct { + InputUnits int `json:"input_units"` + OutputUnits int `json:"output_units"` + AccountingRule string `json:"accounting_rule,omitempty"` + Details json.RawMessage `json:"details,omitempty"` +} + +func (u Usage) validate() error { + if u.InputUnits < 0 || u.OutputUnits < 0 || u.InputUnits > math.MaxInt-u.OutputUnits { + return fmt.Errorf("usage units must be nonnegative integers with a representable total") + } + return nil +} + +func EncodeUsage(u Usage) ([]byte, error) { + if err := u.validate(); err != nil { + return nil, err + } + return json.Marshal(map[string]any{"usage": u}) +} + +// ParseUsage ignores unrelated metadata, but refuses incomplete or invalid +// counters so malformed metadata cannot silently produce a bill. +func ParseUsage(data []byte) (*Usage, error) { + if len(data) == 0 { + return nil, nil + } + var envelope map[string]json.RawMessage + if err := json.Unmarshal(data, &envelope); err != nil { + return nil, err + } + if envelope == nil { + return nil, fmt.Errorf("metadata must be a JSON object") + } + raw, found := envelope["usage"] + if !found { + return nil, nil + } + var fields map[string]json.RawMessage + if err := json.Unmarshal(raw, &fields); err != nil { + return nil, err + } + for _, key := range []string{"input_units", "output_units"} { + value, found := fields[key] + if !found || string(value) == "null" { + return nil, fmt.Errorf("usage requires %s", key) + } + } + var u Usage + if err := json.Unmarshal(raw, &u); err != nil { + return nil, err + } + if err := u.validate(); err != nil { + return nil, err + } + return &u, nil +} diff --git a/pkg/grpc/metadata/metadata_test.go b/pkg/grpc/metadata/metadata_test.go new file mode 100644 index 000000000..f5579dc92 --- /dev/null +++ b/pkg/grpc/metadata/metadata_test.go @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: MIT +package metadata + +import ( + "encoding/json" + "fmt" + "math" + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestMetadata(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Backend response metadata") +} + +var _ = Describe("usage convention", func() { + It("round trips usage with arbitrary nested details", func() { + original := Usage{InputUnits: 6, OutputUnits: 60, AccountingRule: "frame_steps_v1", Details: json.RawMessage(`{"frames":60,"future":{"mode":"fast","weights":[1,2]}}`)} + data, err := EncodeUsage(original) + Expect(err).NotTo(HaveOccurred()) + decoded, err := ParseUsage(data) + Expect(err).NotTo(HaveOccurred()) + Expect(decoded).To(Equal(&original)) + }) + DescribeTable("ignores metadata without usage", func(data string) { + usage, err := ParseUsage([]byte(data)) + Expect(err).NotTo(HaveOccurred()) + Expect(usage).To(BeNil()) + }, Entry("absent", ""), Entry("empty object", "{}"), Entry("unrelated", `{"timings":{"load":1},"custom":[true,"value"]}`)) + DescribeTable("rejects invalid counters", func(data string) { + usage, err := ParseUsage([]byte(data)) + Expect(err).To(HaveOccurred()) + Expect(usage).To(BeNil()) + }, Entry("invalid JSON", `{`), Entry("array", `[]`), Entry("null metadata", `null`), + Entry("null usage", `{"usage":null}`), Entry("missing input", `{"usage":{"output_units":1}}`), + Entry("missing output", `{"usage":{"input_units":1}}`), Entry("null count", `{"usage":{"input_units":null,"output_units":1}}`), + Entry("negative", `{"usage":{"input_units":1,"output_units":-1}}`), + Entry("fraction", `{"usage":{"input_units":1.5,"output_units":1}}`), + Entry("string", `{"usage":{"input_units":"1","output_units":1}}`), + Entry("overflow", fmt.Sprintf(`{"usage":{"input_units":%d,"output_units":1}}`, math.MaxInt))) + It("accepts explicit zero counts", func() { + usage, err := ParseUsage([]byte(`{"usage":{"input_units":0,"output_units":0}}`)) + Expect(err).NotTo(HaveOccurred()) + Expect(usage).To(Equal(&Usage{})) + }) + It("rejects invalid data when producing metadata", func() { + _, err := EncodeUsage(Usage{InputUnits: -1}) + Expect(err).To(HaveOccurred()) + _, err = EncodeUsage(Usage{Details: json.RawMessage(`{`)}) + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/pkg/grpc/server.go b/pkg/grpc/server.go index 092e68b65..67af9d60a 100644 --- a/pkg/grpc/server.go +++ b/pkg/grpc/server.go @@ -188,6 +188,13 @@ func (s *server) Animate3D(ctx context.Context, in *pb.Animate3DRequest) (*pb.Re s.llm.Lock() defer s.llm.Unlock() } + if model, ok := s.llm.(AnimationMetadataModel); ok { + metadata, err := model.Animate3DWithMetadata(in) + if err != nil { + return nil, err + } + return &pb.Result{Message: "3D animation generated", Success: true, Metadata: metadata}, nil + } if err := s.llm.Animate3D(in); err != nil { return nil, err }