diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index aec49aa..2f9a16f 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -7599,6 +7599,10 @@ "type": "integer", "format": "int64" }, + "meshes": { + "type": "integer", + "format": "int64" + }, "cost_micros": { "type": "integer", "format": "int64" @@ -8573,6 +8577,10 @@ "images": { "type": "boolean", "description": "Images is a model that makes pictures, asked for at /v1/images/generations. Implied by an external model whose API is ComfyUI's." + }, + "meshes": { + "type": "boolean", + "description": "Meshes is a model that makes 3D models from a picture, asked for as a job at /v1/3d/generations. Implied by a ComfyUI model whose output is meshes." } } }, @@ -8587,6 +8595,10 @@ "checkpoint": { "type": "string", "description": "Checkpoint is the file the workflow's {{checkpoint}} names: the source's own name unless set, which a source that is a set of files must." + }, + "output": { + "type": "string", + "description": "Output is what the workflow makes: pictures (the default), asked for at /v1/images/generations, or meshes, asked for as a job at /v1/3d/generations from a picture." } } }, @@ -8635,6 +8647,10 @@ "workflow": { "type": "string", "description": "Workflow is the ComfyUI workflow, in its API format: a built-in's name (flux-schnell, the default, or flux-dev) or a JSON file on the control plane." + }, + "output": { + "type": "string", + "description": "Output is what it makes: pictures, or meshes (see ComfyUI.Output)." } } }, @@ -8655,6 +8671,23 @@ } } }, + "catalog.License": { + "type": "object", + "description": "License is the terms what a model makes comes under: said beside every result, and written into a mesh, which travels without the API.", + "properties": { + "name": { + "type": "string", + "description": "Name is the licence's own: MIT, Apache-2.0." + }, + "url": { + "type": "string" + }, + "notice": { + "type": "string", + "description": "Notice is what must accompany what was made: attribution, terms on the output, what the weights were trained on." + } + } + }, "catalog.Model": { "type": "object", "description": "Model is a catalog entry as clients see it.", @@ -8701,6 +8734,10 @@ "$ref": "#/components/schemas/catalog.ComfyUI", "description": "ComfyUI is how a model whose profiles run engine comfyui makes its pictures: the workflow, and the file it loads." }, + "license": { + "$ref": "#/components/schemas/catalog.License", + "description": "License is the terms of what the model makes, for a model whose output somebody ships: shown with every result and written into a mesh." + }, "adapters": { "type": "array", "items": { @@ -10296,6 +10333,10 @@ "Images": { "type": "integer", "format": "int64" + }, + "Meshes": { + "type": "integer", + "format": "int64" } } }, diff --git a/docs/reference.md b/docs/reference.md index 62a960d..01a44ba 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -61,8 +61,10 @@ rather than repeated. | `vision` | `bool` | [Capabilities](#capabilities) | | `embeddings` | `bool` | [Capabilities](#capabilities) | | `images` | `bool` | [Capabilities](#capabilities) | +| `meshes` | `bool` | [Capabilities](#capabilities) | | `workflow` | `string` | [ComfyUI](#comfyui) | | `checkpoint` | `string` | [ComfyUI](#comfyui) | +| `output` | `string` | [ComfyUI](#comfyui) | | `dir` | `string` | [ComfyUIConfig](#comfyuiconfig) | | `python` | `string` | [ComfyUIConfig](#comfyuiconfig) | | `interval` | `time.Duration` | [Control](#control) | @@ -98,6 +100,7 @@ rather than repeated. | `peer` | `string` | [External](#external) | | `api` | `string` | [External](#external) | | `workflow` | `string` | [External](#external) | +| `output` | `string` | [External](#external) | | `budgets` | `[]FacilityBudget` | [Facility](#facility) | | `margin` | `float64` | [Facility](#facility) | | `gpu_watts` | `float64` | [Facility](#facility) | @@ -126,6 +129,9 @@ rather than repeated. | `notice` | `string` | [Landing](#landing) | | `language` | `string` | [Landing](#landing) | | `docs_sign_in` | `bool` | [Landing](#landing) | +| `name` | `string` | [License](#license) | +| `url` | `string` | [License](#license) | +| `notice` | `string` | [License](#license) | | `bin` | `string` | [LlamaCppConfig](#llamacppconfig) | | `parallel` | `int` | [LlamaCppConfig](#llamacppconfig) | | `threads` | `int` | [LlamaCppConfig](#llamacppconfig) | @@ -144,6 +150,7 @@ rather than repeated. | `owner` | `string` | [Model](#model) | | `external` | `*External` | [Model](#model) | | `comfyui` | `*ComfyUI` | [Model](#model) | +| `license` | `*License` | [Model](#model) | | `adapters` | `[]Adapter` | [Model](#model) | | `max_loaded_adapters` | `int` | [Model](#model) | | `fallbacks` | `[]string` | [Model](#model) | @@ -354,6 +361,7 @@ rather than repeated. | `prompt_per_1k` | `int64` | [accounting.Price](#accountingprice) | | `completion_per_1k` | `int64` | [accounting.Price](#accountingprice) | | `per_image` | `int64` | [accounting.Price](#accountingprice) | +| `per_mesh` | `int64` | [accounting.Price](#accountingprice) | | `name` | `string` | [admission.Band](#admissionband) | | `priority` | `int` | [admission.Band](#admissionband) | | `saturation_ceiling` | `float64` | [admission.Band](#admissionband) | @@ -391,6 +399,7 @@ rather than repeated. | `prompt_per_1k` | `int64` | [config.Price](#configprice) | | `completion_per_1k` | `int64` | [config.Price](#configprice) | | `per_image` | `int64` | [config.Price](#configprice) | +| `per_mesh` | `int64` | [config.Price](#configprice) | | `disabled` | `[]string` | [engines.Config](#enginesconfig) | | `vllm` | `VLLMConfig` | [engines.Config](#enginesconfig) | | `sglang` | `SGLangConfig` | [engines.Config](#enginesconfig) | @@ -739,6 +748,14 @@ Images is a model that makes pictures, asked for at /v1/images/generations. Implied by an external model whose API is ComfyUI's. +### `meshes` + +`bool` + +Meshes is a model that makes 3D models from a picture, asked for as a +job at /v1/3d/generations. Implied by a ComfyUI model whose output is +meshes. + ## ComfyUI ### `workflow` @@ -757,6 +774,14 @@ Checkpoint is the file the workflow's {{checkpoint}} names: the source's own name unless set, which a source that is a set of files must. +### `output` + +`string` + +Output is what the workflow makes: pictures (the default), asked for +at /v1/images/generations, or meshes, asked for as a job at +/v1/3d/generations from a picture. + ## ComfyUIConfig ### `dir` @@ -1000,6 +1025,12 @@ Workflow is the ComfyUI workflow, in its API format: a built-in's name (flux-schnell, the default, or flux-dev) or a JSON file on the control plane. +### `output` + +`string` + +Output is what it makes: pictures, or meshes (see ComfyUI.Output). + ## Facility ### `budgets` @@ -1230,6 +1261,25 @@ llmgw in general — no address, model or key of this installation — which the project's repository publishes anyway. On, a visitor is sent to sign in, and without accounts the portal needs an API key. +## License + +### `name` + +`string` + +Name is the licence's own: MIT, Apache-2.0. + +### `url` + +`string` + +### `notice` + +`string` + +Notice is what must accompany what was made: attribution, terms on +the output, what the weights were trained on. + ## LlamaCppConfig ### `bin` @@ -1395,6 +1445,16 @@ pictures: the workflow, and the file it loads. Takes a [ComfyUI](#comfyui) block. +### `license` + +`*License` + +License is the terms of what the model makes, for a model whose +output somebody ships: shown with every result and written into a +mesh. + +Takes a [License](#license) block. + ### `adapters` `[]Adapter` @@ -3007,6 +3067,13 @@ PerImage is what one picture costs, for a model that makes them: a picture's cost is the card's time, which its prompt's few tokens do not measure. +### `per_mesh` + +`int64` + +PerMesh is what one 3D model costs: minutes of a card, which no +token count measures. + ## admission.Band ### `name` @@ -3273,6 +3340,12 @@ had rather wait behind people than be turned away. PerImage is the cost of one picture, for a model that makes them (/v1/images/generations), on top of its prompt's tokens. +### `per_mesh` + +`int64` + +PerMesh is the cost of one 3D model (/v1/3d/generations). + ## engines.Config ### `disabled` diff --git a/internal/accounting/accounting.go b/internal/accounting/accounting.go index 0359268..c510852 100644 --- a/internal/accounting/accounting.go +++ b/internal/accounting/accounting.go @@ -323,7 +323,7 @@ func (a *Accountant) Record(res Reservation, u store.Usage) { if u.ServedModel != "" && u.ServedModel != u.Model && a.pricing != nil { price = a.pricing.For(u.ServedModel) } - u.CostMicros = price.Cost(u.PromptTokens, u.CompletionTokens, u.Images) + u.CostMicros = price.Cost(u.PromptTokens, u.CompletionTokens, u.Images, u.Meshes) for _, b := range res.Budgets { a.budgets.Add(b.Counter, b.Period, u.CostMicros) } diff --git a/internal/accounting/accounting_test.go b/internal/accounting/accounting_test.go index 0b27457..e00b909 100644 --- a/internal/accounting/accounting_test.go +++ b/internal/accounting/accounting_test.go @@ -39,20 +39,20 @@ func TestCostRoundsUpSoSmallRequestsAreNotFree(t *testing.T) { // Rounding down would make a stream of tiny requests cost nothing, which is // exactly the traffic shape a quota exists to bound. p := Price{PromptPer1k: 1000, CompletionPer1k: 3000} - if got := p.Cost(1, 0, 0); got != 1 { + if got := p.Cost(1, 0, 0, 0); got != 1 { t.Errorf("one prompt token cost %d", got) } - if got := p.Cost(1000, 1000, 0); got != 4000 { + if got := p.Cost(1000, 1000, 0, 0); got != 4000 { t.Errorf("a thousand of each cost %d, want 4000", got) } - if got := p.Cost(0, 0, 0); got != 0 { + if got := p.Cost(0, 0, 0, 0); got != 0 { t.Errorf("an empty request cost %d", got) } } func TestAPictureCostsItsPriceOnTopOfItsPrompt(t *testing.T) { p := Price{PromptPer1k: 1000, PerImage: 50000} - if got := p.Cost(20, 0, 3); got != 20+150000 { + if got := p.Cost(20, 0, 3, 0); got != 20+150000 { t.Errorf("three pictures from a 20-token prompt cost %d", got) } if (Price{PerImage: 1}).Zero() { diff --git a/internal/accounting/cost.go b/internal/accounting/cost.go index a2423d5..f7797fd 100644 --- a/internal/accounting/cost.go +++ b/internal/accounting/cost.go @@ -22,6 +22,9 @@ type Price struct { // picture's cost is the card's time, which its prompt's few tokens do // not measure. PerImage int64 `koanf:"per_image"` + // PerMesh is what one 3D model costs: minutes of a card, which no + // token count measures. + PerMesh int64 `koanf:"per_mesh"` } // Zero reports whether this model spends nothing. @@ -29,15 +32,17 @@ type Price struct { // It matters beyond arithmetic: a zero-cost model stays available to a key // whose budget is spent, which is what makes the plan's budget fallback work // (§7.7). -func (p Price) Zero() bool { return p.PromptPer1k == 0 && p.CompletionPer1k == 0 && p.PerImage == 0 } +func (p Price) Zero() bool { + return p.PromptPer1k == 0 && p.CompletionPer1k == 0 && p.PerImage == 0 && p.PerMesh == 0 +} // Cost computes what a request spent, rounding up. // // Up rather than to nearest: rounding down would make a stream of tiny requests // free, which is exactly the shape of traffic a quota exists to bound. -func (p Price) Cost(promptTokens, completionTokens, images int64) int64 { +func (p Price) Cost(promptTokens, completionTokens, images, meshes int64) int64 { return divUp(promptTokens*p.PromptPer1k, 1000) + divUp(completionTokens*p.CompletionPer1k, 1000) + - max(images, 0)*p.PerImage + max(images, 0)*p.PerImage + max(meshes, 0)*p.PerMesh } func divUp(n, d int64) int64 { diff --git a/internal/cli/ctl_request.go b/internal/cli/ctl_request.go index 0b6c814..727fb5a 100644 --- a/internal/cli/ctl_request.go +++ b/internal/cli/ctl_request.go @@ -59,6 +59,9 @@ func newRequestCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { if v.Images > 0 { fmt.Fprintf(out, "images %d\n", v.Images) } + if v.Meshes > 0 { + fmt.Fprintf(out, "meshes %d\n", v.Meshes) + } if v.CacheHit { fmt.Fprintf(out, "cache answered from the project's answer cache: nothing generated, nothing charged\n") } diff --git a/internal/comfy/comfy.go b/internal/comfy/comfy.go index fd45bc4..2e30905 100644 --- a/internal/comfy/comfy.go +++ b/internal/comfy/comfy.go @@ -21,8 +21,10 @@ import ( "fmt" "io" "math/big" + "mime/multipart" "net/http" "net/url" + "path" "slices" "strconv" "strings" @@ -98,6 +100,11 @@ type Params struct { Seed int64 Steps int Guidance float64 + // Image is the input picture, by the name Upload gave it in ComfyUI's + // input directory: {{image}}. + Image string + // Extra are a workflow's own placeholders, by name: {{resolution}}. + Extra map[string]any } func (p Params) value(name string) (any, bool) { @@ -120,8 +127,11 @@ func (p Params) value(name string) (any, bool) { return p.Steps, p.Steps > 0 case "guidance": return p.Guidance, p.Guidance > 0 + case "image": + return p.Image, p.Image != "" } - return nil, false + v, ok := p.Extra[name] + return v, ok } // Fill is the workflow for one request. @@ -207,10 +217,12 @@ func RandomSeed() int64 { return n.Int64() } -// Image is one picture made. +// Image is one file a workflow made: a picture, a mesh. type Image struct { Data []byte ContentType string + // Name is the file's own, as ComfyUI saved it. + Name string } // Error is a failure, said as what it means for the request. @@ -299,13 +311,16 @@ type historyEntry struct { Completed bool `json:"completed"` Messages [][]json.RawMessage `json:"messages"` } `json:"status"` - Outputs map[string]struct { - Images []struct { - Filename string `json:"filename"` - Subfolder string `json:"subfolder"` - Type string `json:"type"` - } `json:"images"` - } `json:"outputs"` + // Outputs are each node's, by kind: "images" for a picture, "3d" for + // the mesh a 3D node saves, and whatever else a custom node names. + Outputs map[string]map[string]json.RawMessage `json:"outputs"` +} + +// fileRef is a file an output names, to be read from /view. +type fileRef struct { + Filename string `json:"filename"` + Subfolder string `json:"subfolder"` + Type string `json:"type"` } // failure is what ComfyUI said went wrong, from its execution_error. @@ -341,17 +356,38 @@ func (c *Client) fetch(ctx context.Context, base, apiKey string, h historyEntry) }) var out []Image for _, n := range nodes { - for _, im := range h.Outputs[n].Images { - q := url.Values{"filename": {im.Filename}, "subfolder": {im.Subfolder}, "type": {im.Type}} - img, err := c.get(ctx, base+"/view?"+q.Encode(), apiKey) - if err != nil { - return nil, err + kinds := make([]string, 0, len(h.Outputs[n])) + for k := range h.Outputs[n] { + kinds = append(kinds, k) + } + slices.Sort(kinds) + for _, k := range kinds { + // Any list of files, whatever the node calls it; text and + // numbers a node reports beside them are not files. + var refs []fileRef + if json.Unmarshal(h.Outputs[n][k], &refs) != nil { + continue + } + for _, f := range refs { + if f.Filename == "" { + continue + } + q := url.Values{"filename": {f.Filename}, "subfolder": {f.Subfolder}, "type": {f.Type}} + img, err := c.get(ctx, base+"/view?"+q.Encode(), apiKey) + if err != nil { + return nil, err + } + img.Name = f.Filename + if t := typeByName(f.Filename); t != "" && (img.ContentType == "" || + strings.HasPrefix(img.ContentType, "application/octet-stream") || strings.HasPrefix(img.ContentType, "text/plain")) { + img.ContentType = t + } + out = append(out, img) } - out = append(out, img) } } if len(out) == 0 { - return nil, &Error{Msg: "the workflow finished without a picture: it has no SaveImage or PreviewImage node"} + return nil, &Error{Msg: "the workflow finished without a file: it has no node that saves one"} } return out, nil } @@ -434,12 +470,12 @@ func (c *Client) get(ctx context.Context, u, apiKey string) (Image, error) { return Image{}, err } defer resp.Body.Close() - b, err := io.ReadAll(io.LimitReader(resp.Body, 64<<20)) + b, err := io.ReadAll(io.LimitReader(resp.Body, maxFile)) if err != nil { - return Image{}, &Error{Msg: "reading a picture from ComfyUI: " + err.Error(), Retryable: true} + return Image{}, &Error{Msg: "reading a file from ComfyUI: " + err.Error(), Retryable: true} } if resp.StatusCode != http.StatusOK { - return Image{}, &Error{Msg: fmt.Sprintf("ComfyUI answered %d for a picture it made", resp.StatusCode)} + return Image{}, &Error{Msg: fmt.Sprintf("ComfyUI answered %d for a file it made", resp.StatusCode)} } ct := resp.Header.Get("Content-Type") if ct == "" || ct == "application/octet-stream" { @@ -527,3 +563,78 @@ func snippet(b []byte) string { } return s } + +// maxFile bounds one file read from ComfyUI: a mesh with its textures runs +// to tens of megabytes, a picture to a few. +const maxFile = 256 << 20 + +// typeByName is a file's type by its extension, for those a server sends +// as octet-stream: the meshes a 3D node saves. +func typeByName(name string) string { + switch strings.ToLower(path.Ext(name)) { + case ".glb": + return "model/gltf-binary" + case ".gltf": + return "model/gltf+json" + case ".obj": + return "model/obj" + case ".stl": + return "model/stl" + case ".ply": + return "application/ply" + case ".png": + return "image/png" + } + return "" +} + +// Upload puts a picture in ComfyUI's input directory, where a LoadImage node +// reads it, and returns the name it was given there. +func (c *Client) Upload(ctx context.Context, base, apiKey, name string, data []byte) (string, error) { + base = strings.TrimRight(base, "/") + var body bytes.Buffer + mw := multipart.NewWriter(&body) + fw, err := mw.CreateFormFile("image", name) + if err != nil { + return "", err + } + if _, err := fw.Write(data); err != nil { + return "", err + } + _ = mw.WriteField("type", "input") + _ = mw.WriteField("overwrite", "true") + if err := mw.Close(); err != nil { + return "", err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/upload/image", &body) + if err != nil { + return "", err + } + req.Header.Set("Content-Type", mw.FormDataContentType()) + if apiKey != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) + } + resp, err := c.client().Do(req) + if err != nil { + if ctx.Err() != nil { + return "", ctx.Err() + } + return "", &Error{Msg: "ComfyUI could not be reached: " + err.Error(), Retryable: true} + } + defer resp.Body.Close() + b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if resp.StatusCode >= 500 { + return "", &Error{Msg: fmt.Sprintf("ComfyUI answered %d to the picture: %s", resp.StatusCode, snippet(b)), Retryable: true} + } + var out struct { + Name string `json:"name"` + Subfolder string `json:"subfolder"` + } + if resp.StatusCode != http.StatusOK || json.Unmarshal(b, &out) != nil || out.Name == "" { + return "", &Error{Msg: fmt.Sprintf("ComfyUI did not take the picture (%d): %s", resp.StatusCode, snippet(b))} + } + if out.Subfolder != "" { + return out.Subfolder + "/" + out.Name, nil + } + return out.Name, nil +} diff --git a/internal/comfy/gate.go b/internal/comfy/gate.go index 8d22fa8..f965e72 100644 --- a/internal/comfy/gate.go +++ b/internal/comfy/gate.go @@ -63,7 +63,10 @@ type GateConfig struct { // The requests the router makes of ComfyUI: queueing, waiting, fetching and // cancelling. Nothing that uploads, installs or changes ComfyUI. -var gatePaths = []string{"/prompt", "/history/", "/view", "/queue", "/interrupt", "/system_stats", "/object_info"} +var gatePaths = []string{"/prompt", "/history/", "/view", "/queue", "/interrupt", "/system_stats", "/object_info", + // The input picture of a 3D model: into ComfyUI's input directory, + // which is the replica's own. + "/upload/image"} // modelFolders are the kinds of file ComfyUI's loaders look for. Each is // shown the same directory: a workflow names its files, so a file being diff --git a/internal/comfy/gate_test.go b/internal/comfy/gate_test.go index 9bd1383..38aebd3 100644 --- a/internal/comfy/gate_test.go +++ b/internal/comfy/gate_test.go @@ -45,12 +45,12 @@ func TestTheGateWantsTheKeyAndLetsThroughOnlyTheRoutersRequests(t *testing.T) { if code, _ := do("POST", "/prompt", "wrong"); code != http.StatusUnauthorized { t.Errorf("wrong key: %d", code) } - for _, p := range []string{"/upload/image", "/api/manager/queue/install", "/userdata/x", "/internal/logs"} { + for _, p := range []string{"/upload/mask", "/api/manager/queue/install", "/userdata/x", "/internal/logs"} { if code, _ := do("POST", p, "k"); code != http.StatusNotFound { t.Errorf("%s: %d, want it refused", p, code) } } - for _, p := range []string{"/prompt", "/history/p1", "/view?filename=a.png", "/queue", "/interrupt"} { + for _, p := range []string{"/prompt", "/history/p1", "/view?filename=a.png", "/queue", "/interrupt", "/upload/image"} { if code, _ := do("POST", p, "k"); code != http.StatusOK { t.Errorf("%s: %d", p, code) } diff --git a/internal/config/config.go b/internal/config/config.go index 0624d05..de9b0fc 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -496,6 +496,8 @@ type Price struct { // PerImage is the cost of one picture, for a model that makes them // (/v1/images/generations), on top of its prompt's tokens. PerImage int64 `koanf:"per_image"` + // PerMesh is the cost of one 3D model (/v1/3d/generations). + PerMesh int64 `koanf:"per_mesh"` } // Control configures the control plane. diff --git a/internal/control/adminapi/api.go b/internal/control/adminapi/api.go index 19e0de8..55d8d53 100644 --- a/internal/control/adminapi/api.go +++ b/internal/control/adminapi/api.go @@ -1612,6 +1612,7 @@ type RequestView struct { PromptTokens int64 `json:"prompt_tokens"` CompletionTokens int64 `json:"completion_tokens"` Images int64 `json:"images,omitempty"` + Meshes int64 `json:"meshes,omitempty"` CostMicros int64 `json:"cost_micros"` // What the guardrails did, by name. GuardAction string `json:"guard_action,omitempty"` @@ -1645,7 +1646,7 @@ func (a *API) request(w http.ResponseWriter, r *http.Request) { RequestID: u.RequestID, At: u.At, KeyID: u.KeyID, UserID: u.UserID, ProjectID: u.ProjectID, Model: u.Model, ServedModel: u.ServedModel, Replica: u.Replica, Status: u.Status, LatencyMs: u.LatencyMs, QueueMs: u.QueueMs, TTFTMs: u.TTFTMs, - PromptTokens: u.PromptTokens, CompletionTokens: u.CompletionTokens, Images: u.Images, CostMicros: u.CostMicros, + PromptTokens: u.PromptTokens, CompletionTokens: u.CompletionTokens, Images: u.Images, Meshes: u.Meshes, CostMicros: u.CostMicros, GuardAction: u.GuardAction, GuardPolicy: u.GuardPolicy, GuardCategory: u.GuardCategory, GuardRule: u.GuardRule, GuardMs: u.GuardMs, EndUser: u.EndUser, Tags: u.Tags, CacheHit: u.CacheHit, }) diff --git a/internal/control/adminapi/export.go b/internal/control/adminapi/export.go index 44cbe54..7e01f51 100644 --- a/internal/control/adminapi/export.go +++ b/internal/control/adminapi/export.go @@ -75,11 +75,11 @@ func (a *API) usageExport(w http.ResponseWriter, r *http.Request) { func WriteExportCSV(w io.Writer, rows []store.ExportRow) { c := csv.NewWriter(w) _ = c.Write([]string{"period_start", "org", "team", "project", "user", "key", "model", - "requests", "errors", "limited", "prompt_tokens", "completion_tokens", "cost_micros", "cost", "images"}) + "requests", "errors", "limited", "prompt_tokens", "completion_tokens", "cost_micros", "cost", "images", "meshes"}) for _, r := range rows { _ = c.Write([]string{r.Period.Format(time.RFC3339), safe(r.Org), safe(r.Team), safe(r.Project), safe(r.User), r.Key, safe(r.Model), i(r.Requests), i(r.Errors), i(r.Limited), i(r.Prompt), i(r.Completion), - i(r.CostMicros), strconv.FormatFloat(float64(r.CostMicros)/1e6, 'f', 6, 64), i(r.Images)}) + i(r.CostMicros), strconv.FormatFloat(float64(r.CostMicros)/1e6, 'f', 6, 64), i(r.Images), i(r.Meshes)}) } c.Flush() } diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 79cea98..e5201c2 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -332,6 +332,10 @@ type Model struct { // ComfyUI is how a model whose profiles run engine comfyui makes its // pictures: the workflow, and the file it loads. ComfyUI *ComfyUI `json:"comfyui,omitempty" koanf:"comfyui"` + // License is the terms of what the model makes, for a model whose + // output somebody ships: shown with every result and written into a + // mesh. + License *License `json:"license,omitempty" koanf:"license"` // Adapters are LoRA adapters served on this model's replicas, each asked // for by its own name as though it were a model. Every replica loads // every adapter, applying one only to a request that names it — unless @@ -472,11 +476,52 @@ type Capabilities struct { // /v1/images/generations. Implied by an external model whose API is // ComfyUI's. Images bool `json:"images,omitempty" koanf:"images"` + // Meshes is a model that makes 3D models from a picture, asked for as a + // job at /v1/3d/generations. Implied by a ComfyUI model whose output is + // meshes. + Meshes bool `json:"meshes,omitempty" koanf:"meshes"` +} + +// The kinds of thing a ComfyUI model makes. +const ( + OutputPictures = "pictures" + OutputMeshes = "meshes" +) + +// comfyOutput is what a ComfyUI model makes: pictures unless it says. +func (m Model) comfyOutput() string { + out := "" + switch { + case m.External != nil: + out = m.External.Output + case m.ComfyUI != nil: + out = m.ComfyUI.Output + } + if out == "" { + return OutputPictures + } + return out } // MakesImages reports whether a model answers /v1/images/generations. func (m Model) MakesImages() bool { - return m.Capabilities.Images || m.ServesComfyUI() + return m.Capabilities.Images || m.ServesComfyUI() && m.comfyOutput() == OutputPictures +} + +// MakesMeshes reports whether a model answers /v1/3d/generations. +func (m Model) MakesMeshes() bool { + return m.Capabilities.Meshes || m.ServesComfyUI() && m.comfyOutput() == OutputMeshes +} + +// License is the terms what a model makes comes under: said beside every +// result, and written into a mesh, which travels without the API. +type License struct { + // Name is the licence's own: MIT, Apache-2.0. + Name string `json:"name" koanf:"name"` + URL string `json:"url,omitempty" koanf:"url"` + // Notice is what must accompany what was made: attribution, terms on + // the output, what the weights were trained on. + Notice string `json:"notice,omitempty" koanf:"notice"` } // EngineComfyUI is the engine that runs ComfyUI on a node. @@ -501,6 +546,10 @@ type ComfyUI struct { // source's own name unless set, which a source that is a set of files // must. Checkpoint string `json:"checkpoint,omitempty" koanf:"checkpoint"` + // Output is what the workflow makes: pictures (the default), asked for + // at /v1/images/generations, or meshes, asked for as a job at + // /v1/3d/generations from a picture. + Output string `json:"output,omitempty" koanf:"output"` } // ComfyWorkflow is what a ComfyUI model's replicas are sent: the workflow @@ -555,6 +604,12 @@ func (m Model) validateComfyUI() error { if err := checkWorkflow(m.Name, w); err != nil { return err } + if out := m.comfyOutput(); out != OutputPictures && out != OutputMeshes { + return fmt.Errorf("model %q: comfyui.output %q is not pictures or meshes", m.Name, out) + } + if m.comfyOutput() == OutputMeshes && w == "" { + return fmt.Errorf("model %q: a model that makes meshes names its workflow: no built-in makes them", m.Name) + } if _, ckpt := m.ComfyWorkflow(); ckpt == "" { return fmt.Errorf("model %q: its source is a set of files, so comfyui.checkpoint must name the one the workflow loads", m.Name) @@ -1257,6 +1312,8 @@ type External struct { // (flux-schnell, the default, or flux-dev) or a JSON file on the control // plane. Workflow string `json:"workflow,omitempty" koanf:"workflow"` + // Output is what it makes: pictures, or meshes (see ComfyUI.Output). + Output string `json:"output,omitempty" koanf:"output"` } // APIComfyUI is External.API for a ComfyUI server. @@ -1278,8 +1335,8 @@ func (e *External) validate(model string, profiles int) error { } switch e.API { case "", "openai": - if e.Workflow != "" { - return fmt.Errorf("model %q: a workflow is for api: comfyui", model) + if e.Workflow != "" || e.Output != "" { + return fmt.Errorf("model %q: a workflow and an output are for api: comfyui", model) } case APIComfyUI: if e.Peer != "" { @@ -1288,6 +1345,12 @@ func (e *External) validate(model string, profiles int) error { if err := checkWorkflow(model, e.Workflow); err != nil { return err } + if e.Output != "" && e.Output != OutputPictures && e.Output != OutputMeshes { + return fmt.Errorf("model %q: external output %q is not pictures or meshes", model, e.Output) + } + if e.Output == OutputMeshes && e.Workflow == "" { + return fmt.Errorf("model %q: a model that makes meshes names its workflow: no built-in makes them", model) + } default: return fmt.Errorf("model %q: external api %q is not openai or comfyui", model, e.API) } diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index 805a4d3..4a0e6fc 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -65,10 +65,12 @@ var Reference = []Block{ {Key: "vision", Type: "bool", Doc: "", Nested: ""}, {Key: "embeddings", Type: "bool", Doc: "", Nested: ""}, {Key: "images", Type: "bool", Doc: "Images is a model that makes pictures, asked for at\n/v1/images/generations. Implied by an external model whose API is\nComfyUI's.", Nested: ""}, + {Key: "meshes", Type: "bool", Doc: "Meshes is a model that makes 3D models from a picture, asked for as a\njob at /v1/3d/generations. Implied by a ComfyUI model whose output is\nmeshes.", Nested: ""}, }}, {Package: "catalog", Name: "ComfyUI", Title: "ComfyUI", Doc: "", Fields: []Field{ {Key: "workflow", Type: "string", Doc: "Workflow is a built-in's name (flux-schnell, the default, or\nflux-dev) or a JSON file on the control plane, in ComfyUI's API\nformat.", Nested: ""}, {Key: "checkpoint", Type: "string", Doc: "Checkpoint is the file the workflow's {{checkpoint}} names: the\nsource's own name unless set, which a source that is a set of files\nmust.", Nested: ""}, + {Key: "output", Type: "string", Doc: "Output is what the workflow makes: pictures (the default), asked for\nat /v1/images/generations, or meshes, asked for as a job at\n/v1/3d/generations from a picture.", Nested: ""}, }}, {Package: "engines", Name: "ComfyUIConfig", Title: "ComfyUIConfig", Doc: "", Fields: []Field{ {Key: "dir", Type: "string", Doc: "Dir is ComfyUI's directory, where main.py is. Readable by the user\nengines run as.", Nested: ""}, @@ -116,6 +118,7 @@ var Reference = []Block{ {Key: "peer", Type: "string", Doc: "Peer names the federated llmgw cluster this model is offered by, for\na model that federation added (`llmgw ctl federation add`) rather\nthan one written by hand. Such a model is named PEER/MODEL, comes and\ngoes with the peer's own list, and is sent only while the peer is up.", Nested: ""}, {Key: "api", Type: "string", Doc: "API is what the provider speaks: OpenAI's, unless this says comfyui —\na ComfyUI server, asked for pictures at /v1/images/generations and\nsent the workflow below with the request's prompt in it. Model is then\nthe checkpoint file the workflow loads, as ComfyUI lists it.", Nested: ""}, {Key: "workflow", Type: "string", Doc: "Workflow is the ComfyUI workflow, in its API format: a built-in's name\n(flux-schnell, the default, or flux-dev) or a JSON file on the control\nplane.", Nested: ""}, + {Key: "output", Type: "string", Doc: "Output is what it makes: pictures, or meshes (see ComfyUI.Output).", Nested: ""}, }}, {Package: "config", Name: "Facility", Title: "Facility", Doc: "", Fields: []Field{ {Key: "budgets", Type: "[]FacilityBudget", Doc: "Budgets are the domains with a power cap.", Nested: ""}, @@ -161,6 +164,11 @@ var Reference = []Block{ {Key: "language", Type: "string", Doc: "Language is what the web portal speaks to someone who has not chosen\na language with the switcher at the foot of every page and whose\nbrowser asks for none it speaks: \"en\" or \"tr\". Empty is English.", Nested: ""}, {Key: "docs_sign_in", Type: "bool", Doc: "DocsSignIn serves the docs portal at /docs only to someone signed in.\nOff, the portal is public like this page: its client guide is for the\nperson who has no account yet, and the operator documents describe\nllmgw in general — no address, model or key of this installation —\nwhich the project's repository publishes anyway. On, a visitor is sent\nto sign in, and without accounts the portal needs an API key.", Nested: ""}, }}, + {Package: "catalog", Name: "License", Title: "License", Doc: "", Fields: []Field{ + {Key: "name", Type: "string", Doc: "Name is the licence's own: MIT, Apache-2.0.", Nested: ""}, + {Key: "url", Type: "string", Doc: "", Nested: ""}, + {Key: "notice", Type: "string", Doc: "Notice is what must accompany what was made: attribution, terms on\nthe output, what the weights were trained on.", Nested: ""}, + }}, {Package: "engines", Name: "LlamaCppConfig", Title: "LlamaCppConfig", Doc: "", Fields: []Field{ {Key: "bin", Type: "string", Doc: "Bin is the llama-server binary. It is the server, not the CLI: the\ngateway talks OpenAI over HTTP and needs the one that serves it.\nEmpty means the agent installs llama.cpp itself, at the version the\ncontrol plane asks for. Setting it opts out of that.", Nested: ""}, {Key: "parallel", Type: "int", Doc: "Parallel is how many requests a server handles at once, when a profile\ndoes not say. llama.cpp defaults to one, which would make every second\nrequest wait behind the first no matter what the router decided.", Nested: ""}, @@ -184,6 +192,7 @@ var Reference = []Block{ {Key: "owner", Type: "string", Doc: "Owner is the part of the tenancy tree whose capacity quota this model\ncounts against, written org:ID, team:ID or project:ID. Empty is shared:\nheld by the cluster, counted against nobody.", Nested: ""}, {Key: "external", Type: "*External", Doc: "External is a model served by someone else — an OpenAI-compatible API\noutside the cluster — rather than by replicas placed here. Such a model\nhas no profiles and nothing is placed for it; the router sends its\nrequests to the provider. Most useful named in another model's\nfallbacks: used only when the cluster has no capacity left. Its\nprompts leave the cluster, so which keys may name it is worth an\nallowlist.", Nested: "catalog.External"}, {Key: "comfyui", Type: "*ComfyUI", Doc: "ComfyUI is how a model whose profiles run engine comfyui makes its\npictures: the workflow, and the file it loads.", Nested: "catalog.ComfyUI"}, + {Key: "license", Type: "*License", Doc: "License is the terms of what the model makes, for a model whose\noutput somebody ships: shown with every result and written into a\nmesh.", Nested: "catalog.License"}, {Key: "adapters", Type: "[]Adapter", Doc: "Adapters are LoRA adapters served on this model's replicas, each asked\nfor by its own name as though it were a model. Every replica loads\nevery adapter, applying one only to a request that names it — unless\nMaxLoadedAdapters is set.", Nested: ""}, {Key: "max_loaded_adapters", Type: "int", Doc: "MaxLoadedAdapters, above zero, has each replica hold at most this many\nadapters at once, loaded when a request first asks for one and unloaded\nleast recently used first to make room. The router sends a request for\nan adapter to a replica that already holds it while that one has room\nto take it. vLLM profiles only: llama-server loads adapters at start\nand cannot load one afterwards. Zero loads every adapter on every\nreplica at start.", Nested: ""}, {Key: "fallbacks", Type: "[]string", Doc: "", Nested: ""}, @@ -460,6 +469,7 @@ var Reference = []Block{ {Key: "prompt_per_1k", Type: "int64", Doc: "", Nested: ""}, {Key: "completion_per_1k", Type: "int64", Doc: "", Nested: ""}, {Key: "per_image", Type: "int64", Doc: "PerImage is what one picture costs, for a model that makes them: a\npicture's cost is the card's time, which its prompt's few tokens do\nnot measure.", Nested: ""}, + {Key: "per_mesh", Type: "int64", Doc: "PerMesh is what one 3D model costs: minutes of a card, which no\ntoken count measures.", Nested: ""}, }}, {Package: "admission", Name: "Band", Title: "admission.Band", Doc: "", Fields: []Field{ {Key: "name", Type: "string", Doc: "", Nested: ""}, @@ -507,6 +517,7 @@ var Reference = []Block{ {Key: "prompt_per_1k", Type: "int64", Doc: "", Nested: ""}, {Key: "completion_per_1k", Type: "int64", Doc: "", Nested: ""}, {Key: "per_image", Type: "int64", Doc: "PerImage is the cost of one picture, for a model that makes them\n(/v1/images/generations), on top of its prompt's tokens.", Nested: ""}, + {Key: "per_mesh", Type: "int64", Doc: "PerMesh is the cost of one 3D model (/v1/3d/generations).", Nested: ""}, }}, {Package: "engines", Name: "Config", Title: "engines.Config", Doc: "", Fields: []Field{ {Key: "disabled", Type: "[]string", Doc: "Disabled are engines this node never installs or advertises, whatever\nthe control plane asks: vLLM on a node of cards no build of it runs on,\nthat is meant for llama.cpp alone.", Nested: ""}, diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 21e1fe9..f10c3bf 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -763,6 +763,7 @@ func (r *Router) account(req *parsedRequest, res dispatchResult, status int, too u.PromptTokens = int64(req.PromptTokens) } u.Images = int64(res.outcome.Images) + u.Meshes = int64(res.outcome.Meshes) r.m.tenant.observe(req, status, u.PromptTokens, u.CompletionTokens, res.outcome.TTFT) r.acct.Record(req.reservation, u) } @@ -1142,9 +1143,12 @@ func (r *Router) handleModels(w http.ResponseWriter, req *http.Request) { if m.Description != "" { entry["description"] = m.Description } + if m.License != nil { + entry["license"] = m.License + } entry["capabilities"] = map[string]bool{ "tools": m.Capabilities.Tools, "vision": m.Capabilities.Vision, - "embeddings": m.Capabilities.Embeddings, "images": m.MakesImages(), + "embeddings": m.Capabilities.Embeddings, "images": m.MakesImages(), "meshes": m.MakesMeshes(), } data = append(data, entry) } diff --git a/internal/router/images.go b/internal/router/images.go index a649870..eaafa74 100644 --- a/internal/router/images.go +++ b/internal/router/images.go @@ -112,15 +112,21 @@ func parseImageRequest(body []byte) (imageRequest, error) { func (r *Router) imageRouteRefusal(path, model string, body []byte) string { m, known := r.modelMap()[model] if path == meshesPath { - // Until a model makes them (BACKLOG G3), every model the catalog - // knows is one that does not; one it does not know is a - // hand-configured endpoint's to answer. - if known { + if known && !m.MakesMeshes() { + if m.MakesImages() { + return fmt.Sprintf("%s makes pictures, not 3D models: ask for them at %s", m.Name, imagesPath) + } return fmt.Sprintf("%s does not make 3D models", m.Name) } + if _, err := parseMeshRequest(body); err != nil { + return err.Error() + } return "" } if path == imagesPath { + if known && m.MakesMeshes() && !m.MakesImages() { + return fmt.Sprintf("%s makes 3D models, not pictures: ask for one at %s, as a job", m.Name, meshesPath) + } if known && !m.MakesImages() { return fmt.Sprintf("%s does not make pictures", m.Name) } @@ -130,6 +136,9 @@ func (r *Router) imageRouteRefusal(path, model string, body []byte) string { return "" } if known && m.ServesComfyUI() { + if m.MakesMeshes() { + return fmt.Sprintf("%s makes 3D models: ask for one at %s, as a job", m.Name, meshesPath) + } return fmt.Sprintf("%s makes pictures: ask for them at %s", m.Name, imagesPath) } return "" @@ -139,8 +148,11 @@ func (r *Router) imageRouteRefusal(path, model string, body []byte) string { // proxy would: a failure of this replica that another may not have is // retryable, and nothing is written until the pictures are all in hand. func (r *Router) generateImages(ctx context.Context, w http.ResponseWriter, req *parsedRequest, ep endpoint.Info, body []byte) (proxy.Outcome, error) { + if req.Path == meshesPath { + return r.generateMeshes(ctx, w, req, ep, body) + } if req.Path != imagesPath { - return proxy.Outcome{}, fmt.Errorf("%s is a ComfyUI server: it answers %s only", ep.ID, imagesPath) + return proxy.Outcome{}, fmt.Errorf("%s is a ComfyUI server: it answers %s and %s only", ep.ID, imagesPath, meshesPath) } in, err := parseImageRequest(body) if err != nil { diff --git a/internal/router/jobs.go b/internal/router/jobs.go index a568b61..80cbbb9 100644 --- a/internal/router/jobs.go +++ b/internal/router/jobs.go @@ -50,12 +50,17 @@ func (r *Router) jobJSON(ctx context.Context, j store.Job) map[string]any { if j.Error != "" { out["error"] = map[string]string{"message": j.Error} } + // The terms what it makes comes under, beside the result: whoever ships + // the mesh answers for them. + if m, ok := r.modelMap()[j.Model]; ok && m.License != nil { + out["license"] = m.License + } // What the work cost, from the usage record the replayed request left // under the job's id. if store.JobDone(j.Status) && r.store != nil { if u, ok, err := r.store.UsageByRequest(ctx, j.ID); err == nil && ok { out["usage"] = map[string]int64{"prompt_tokens": u.PromptTokens, "completion_tokens": u.CompletionTokens, - "images": u.Images} + "images": u.Images, "meshes": u.Meshes} } } return out diff --git a/internal/router/meshes.go b/internal/router/meshes.go new file mode 100644 index 0000000..c835b4b --- /dev/null +++ b/internal/router/meshes.go @@ -0,0 +1,240 @@ +package router + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/base64" + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/comfy" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" +) + +// 3D models from a picture, made by ComfyUI. Clients ask for them as a job +// (jobs.go); the job worker replays the request through the ordinary path, +// and at a ComfyUI replica it arrives here: the picture is put in ComfyUI's +// input directory, the model's workflow is filled in with its name and run, +// and the mesh it saves comes back as the answer — the file the job keeps. + +// meshRequest is what a request for a 3D model asks. The picture is the +// whole of the prompt: base64, or a data URL. +type meshRequest struct { + Model string `json:"model"` + Image string `json:"image"` + Seed *int64 `json:"seed"` + Steps int `json:"steps"` + // Resolution is the model's own grid, where it has a choice: TRELLIS.2 + // makes 512, 1024 or 1536. Zero leaves it to the workflow. + Resolution int `json:"resolution"` + + picture []byte + ext string +} + +// maxPicture bounds the picture a mesh is made from. +const maxPicture = 16 << 20 + +func parseMeshRequest(body []byte) (meshRequest, error) { + var in meshRequest + if err := json.Unmarshal(body, &in); err != nil { + return in, err + } + s := strings.TrimSpace(in.Image) + if s == "" { + return in, errors.New("image is required: the picture to make a 3D model of, base64 or a data URL") + } + if strings.HasPrefix(s, "data:") { + _, after, ok := strings.Cut(s, ";base64,") + if !ok { + return in, errors.New("image is a data URL without base64 in it") + } + s = after + } + if base64.StdEncoding.DecodedLen(len(s)) > maxPicture+3 { + return in, fmt.Errorf("image is past %d MB", maxPicture>>20) + } + b, err := base64.StdEncoding.DecodeString(s) + if err != nil { + return in, errors.New("image is not base64: " + err.Error()) + } + switch http.DetectContentType(b) { + case "image/png": + in.ext = ".png" + case "image/jpeg": + in.ext = ".jpg" + case "image/webp": + in.ext = ".webp" + default: + return in, errors.New("image is not a PNG, JPEG or WebP picture") + } + in.picture = b + if in.Steps < 0 || in.Steps > maxSteps { + return in, fmt.Errorf("steps is between 1 and %d", maxSteps) + } + switch in.Resolution { + case 0, 512, 1024, 1536: + default: + return in, fmt.Errorf("resolution %d is not 512, 1024 or 1536", in.Resolution) + } + return in, nil +} + +// generateMeshes answers a request for a 3D model at a ComfyUI replica. It +// writes the mesh itself, with its type, which is what the job keeps. +func (r *Router) generateMeshes(ctx context.Context, w http.ResponseWriter, req *parsedRequest, ep endpoint.Info, body []byte) (proxy.Outcome, error) { + in, err := parseMeshRequest(body) + if err != nil { + return proxy.Outcome{}, err + } + tpl, err := comfy.Load(ep.Workflow) + if err != nil { + return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err) + } + if t := r.cfg.Timeouts.Total; t > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, t) + defer cancel() + } + start := time.Now() + fail := func(err error) (proxy.Outcome, error) { + var ce *comfy.Error + switch { + case errors.As(err, &ce): + return proxy.Outcome{Retryable: ce.Retryable}, err + case errors.Is(err, context.Canceled) && ctxCauseIsClient(ctx): + return proxy.Outcome{}, proxy.ErrClientGone + case errors.Is(err, context.DeadlineExceeded): + return proxy.Outcome{}, fmt.Errorf("%s made no 3D model within %s", ep.ID, r.cfg.Timeouts.Total) + } + return proxy.Outcome{}, err + } + // A name of its own: two requests' pictures must not overwrite each + // other in the replica's input directory. + rnd := make([]byte, 8) + _, _ = rand.Read(rnd) + name, err := r.comfy.Upload(ctx, ep.URL, ep.APIKey, "llmgw-"+hex.EncodeToString(rnd)+in.ext, in.picture) + if err != nil { + return fail(err) + } + seed := comfy.RandomSeed() + if in.Seed != nil { + seed = *in.Seed + } + var upstream struct { + Model string `json:"model"` + } + _ = json.Unmarshal(body, &upstream) + p := comfy.Params{Checkpoint: upstream.Model, Image: name, Seed: seed, Steps: in.Steps, Extra: map[string]any{}} + if in.Resolution > 0 { + p.Extra["resolution"] = in.Resolution + } + workflow, err := tpl.Fill(p) + if err != nil { + return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err) + } + files, err := r.comfy.Generate(ctx, ep.URL, ep.APIKey, workflow) + if err != nil { + return fail(err) + } + var mesh *comfy.Image + for i := range files { + if isMeshType(files[i].ContentType) { + mesh = &files[i] + break + } + } + if mesh == nil { + return proxy.Outcome{}, fmt.Errorf("%s: the workflow saved no mesh (it saved %d other files): "+ + "its last node should save a .glb", ep.ID, len(files)) + } + data := mesh.Data + m := r.modelMap()[req.Model] + if m.License != nil && mesh.ContentType == "model/gltf-binary" { + if with, err := glbWithLicense(data, *m.License); err == nil { + data = with + } else { + // Refused rather than sent without it: a mesh travels without + // the API, and the notice is what the terms ask to travel too. + return proxy.Outcome{}, fmt.Errorf("%s: the licence could not be written into the mesh: %w", ep.ID, err) + } + } + w.Header().Set("x-llmgw-seed", strconv.FormatInt(seed, 10)) + if m.License != nil { + w.Header().Set("x-llmgw-license", m.License.Name) + } + w.Header().Set("Content-Type", mesh.ContentType) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(data) + took := time.Since(start) + return proxy.Outcome{Wrote: true, Status: http.StatusOK, Bytes: int64(len(data)), TTFT: took, Decode: took, + Meshes: 1}, nil +} + +func isMeshType(ct string) bool { + return strings.HasPrefix(ct, "model/") || ct == "application/ply" +} + +// glbWithLicense is a GLB with the licence in asset.extras.license, where a +// file's terms travel with it: glTF's own place for what an application +// adds, kept by every tool that loads and saves the file. +func glbWithLicense(glb []byte, lic catalog.License) ([]byte, error) { + if len(glb) < 20 || string(glb[:4]) != "glTF" || binary.LittleEndian.Uint32(glb[4:]) != 2 { + return nil, errors.New("not a glTF 2 binary") + } + n := int(binary.LittleEndian.Uint32(glb[12:])) + if string(glb[16:20]) != "JSON" || 20+n > len(glb) { + return nil, errors.New("its first chunk is not its JSON") + } + var doc map[string]any + d := json.NewDecoder(bytes.NewReader(glb[20 : 20+n])) + d.UseNumber() // every number as it was written + if err := d.Decode(&doc); err != nil { + return nil, err + } + asset, _ := doc["asset"].(map[string]any) + if asset == nil { + asset = map[string]any{"version": "2.0"} + } + extras, _ := asset["extras"].(map[string]any) + if extras == nil { + extras = map[string]any{} + } + l := map[string]any{"name": lic.Name} + if lic.URL != "" { + l["url"] = lic.URL + } + if lic.Notice != "" { + l["notice"] = lic.Notice + } + extras["license"] = l + asset["extras"] = extras + doc["asset"] = asset + js, err := json.Marshal(doc) + if err != nil { + return nil, err + } + // A chunk is padded to four bytes, the JSON one with spaces. + for len(js)%4 != 0 { + js = append(js, ' ') + } + rest := glb[20+n:] + out := make([]byte, 0, 20+len(js)+len(rest)) + out = append(out, "glTF"...) + out = binary.LittleEndian.AppendUint32(out, 2) + out = binary.LittleEndian.AppendUint32(out, uint32(20+len(js)+len(rest))) //nolint:gosec // bounded by the file read + out = binary.LittleEndian.AppendUint32(out, uint32(len(js))) //nolint:gosec // bounded by the file read + out = append(out, "JSON"...) + out = append(out, js...) + return append(out, rest...), nil +} diff --git a/internal/router/meshes_test.go b/internal/router/meshes_test.go new file mode 100644 index 0000000..3fad7fe --- /dev/null +++ b/internal/router/meshes_test.go @@ -0,0 +1,247 @@ +package router + +import ( + "bytes" + "crypto/sha256" + "encoding/base64" + "encoding/binary" + "encoding/json" + "fmt" + "image" + "image/png" + "io" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/comfy" + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// testGLB is the smallest glTF binary: an asset, and a binary chunk after +// its JSON that must come through untouched. +func testGLB(t *testing.T) []byte { + t.Helper() + js := []byte(`{"asset":{"version":"2.0","generator":"test"},"buffers":[{"byteLength":4}]}`) + for len(js)%4 != 0 { + js = append(js, ' ') + } + bin := []byte{1, 2, 3, 4} + var b []byte + b = append(b, "glTF"...) + b = binary.LittleEndian.AppendUint32(b, 2) + b = binary.LittleEndian.AppendUint32(b, uint32(12+8+len(js)+8+len(bin))) + b = binary.LittleEndian.AppendUint32(b, uint32(len(js))) + b = append(b, "JSON"...) + b = append(b, js...) + b = binary.LittleEndian.AppendUint32(b, uint32(len(bin))) + b = append(b, "BIN\x00"...) + return append(b, bin...) +} + +// meshComfy is a ComfyUI that takes a picture and saves a mesh, as a 3D +// node's SaveGLB does: under "3d", not "images". +type meshComfy struct { + srv *httptest.Server + mu sync.Mutex + uploaded []string + workflow map[string]any +} + +func newMeshComfy(t *testing.T) *meshComfy { + glb := testGLB(t) + c := &meshComfy{} + c.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/upload/image": + f, h, err := r.FormFile("image") + if err != nil { + http.Error(w, err.Error(), 400) + return + } + b, _ := io.ReadAll(f) + c.mu.Lock() + c.uploaded = append(c.uploaded, h.Filename+":"+http.DetectContentType(b)) + c.mu.Unlock() + fmt.Fprintf(w, `{"name":%q,"subfolder":"","type":"input"}`, h.Filename) + case "/prompt": + var b struct { + Prompt map[string]any `json:"prompt"` + } + _ = json.NewDecoder(r.Body).Decode(&b) + c.mu.Lock() + c.workflow = b.Prompt + c.mu.Unlock() + fmt.Fprint(w, `{"prompt_id":"m1"}`) + case "/history/m1": + fmt.Fprint(w, `{"m1":{"status":{"status_str":"success","completed":true},"outputs":{ + "3":{"text":["a preview"]},"4":{"3d":[{"filename":"mesh_00001_.glb","subfolder":"3d","type":"output"}]}}}}`) + case "/view": + // Served as ComfyUI serves a file it has no type for. + w.Header().Set("Content-Type", "application/octet-stream") + _, _ = w.Write(glb) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(c.srv.Close) + return c +} + +const meshWorkflow = `{ + "1": {"class_type": "LoadImage", "inputs": {"image": "{{image}}"}}, + "2": {"class_type": "Trellis2ImageToMesh", "inputs": {"image": ["1", 0], "seed": "{{seed}}", "resolution": "{{resolution|512}}"}}, + "4": {"class_type": "SaveGLB", "inputs": {"mesh": ["2", 0], "filename_prefix": "3d/mesh"}} +}` + +func testPicture(t *testing.T) string { + var b bytes.Buffer + _ = png.Encode(&b, image.NewRGBA(image.Rect(0, 0, 8, 8))) + return base64.StdEncoding.EncodeToString(b.Bytes()) +} + +func meshRouter(t *testing.T, c *meshComfy) (*Router, string) { + t.Helper() + cfg := baseConfig() + mit := &catalog.License{Name: "MIT", URL: "https://github.com/microsoft/TRELLIS.2/blob/main/LICENSE", + Notice: "Made with TRELLIS.2 (MIT)."} + cfg.Models = []catalog.Model{ + {Name: "trellis2", Source: "external", License: mit, External: &catalog.External{URL: c.srv.URL, + Model: "trellis2", API: catalog.APIComfyUI, Output: catalog.OutputMeshes, Workflow: "w.json"}}, + {Name: "flux", Source: "external", External: &catalog.External{URL: c.srv.URL, + Model: "flux1-schnell-fp8.safetensors", API: catalog.APIComfyUI}}, + } + cfg.Auth.Enabled = true + cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db") + cfg.Auth.Pricing = map[string]config.Price{"trellis2": {PerMesh: 250000}} + cfg.Auth.FlushInterval = 20 * time.Millisecond + cfg.Auth.BatchSize = 1 + rt := newTestRouter(t, cfg) + rt.comfy = &comfy.Client{Poll: time.Millisecond} + rt.SetEndpoints(t.Context(), []endpoint.Info{ + {ID: "external:trellis2", URL: c.srv.URL, Models: []string{"trellis2"}, Engine: endpoint.EngineComfyUI, + Pool: "external", Weight: 1, UpstreamModel: "trellis2", Workflow: meshWorkflow}, + {ID: "external:flux", URL: c.srv.URL, Models: []string{"flux"}, Engine: endpoint.EngineComfyUI, + Pool: "external", Weight: 1, UpstreamModel: "flux1-schnell-fp8.safetensors"}, + }) + org, _ := rt.store.CreateOrg(t.Context(), "studio") + team, _ := rt.store.CreateTeam(t.Context(), org, "art") + project, err := rt.store.CreateProject(t.Context(), team, "props") + if err != nil { + t.Fatal(err) + } + secret := store.KeyPrefix + strings.Repeat("m", 43) + sum := sha256.Sum256([]byte(secret)) + if err := rt.store.CreateKey(t.Context(), store.Key{ID: "key-m", Hash: sum[:], Prefix: "mmmm", Name: "art", + ProjectID: project, CreatedAt: time.Now()}); err != nil { + t.Fatal(err) + } + return rt, secret +} + +func TestA3DModelIsMadeFromAPictureWithItsLicenceInside(t *testing.T) { + c := newMeshComfy(t) + rt, secret := meshRouter(t, c) + h := rt.Handler() + + code, j, raw := call(t, h, secret, http.MethodPost, meshesPath, + map[string]any{"model": "trellis2", "image": "data:image/png;base64," + testPicture(t), "seed": 11, "resolution": 1024}) + if code != http.StatusAccepted { + t.Fatalf("%d %s", code, raw) + } + got := waitJob(t, h, secret, j["id"].(string), jobDone) + if got["status"] != "completed" || got["content_type"] != "model/gltf-binary" { + t.Fatalf("job %v", got) + } + if lic, _ := got["license"].(map[string]any); lic["name"] != "MIT" { + t.Errorf("the job does not say the licence: %v", got["license"]) + } + code, _, file := call(t, h, secret, http.MethodGet, "/v1/files/"+got["output_file_id"].(string)+"/content", nil) + if code != http.StatusOK || !strings.HasPrefix(file, "glTF") { + t.Fatalf("the result is not the mesh: %d %q", code, file[:min(len(file), 40)]) + } + b := []byte(file) + n := binary.LittleEndian.Uint32(b[12:]) + var doc struct { + Asset struct { + Generator string `json:"generator"` + Extras struct { + License catalog.License `json:"license"` + } `json:"extras"` + } `json:"asset"` + } + if err := json.Unmarshal(b[20:20+n], &doc); err != nil { + t.Fatal(err) + } + if doc.Asset.Extras.License.Notice != "Made with TRELLIS.2 (MIT)." || doc.Asset.Generator != "test" { + t.Errorf("asset %+v: the notice is not in the mesh, or what was there is lost", doc.Asset) + } + if int(binary.LittleEndian.Uint32(b[8:])) != len(b) || !bytes.HasSuffix(b, []byte{1, 2, 3, 4}) { + t.Error("the mesh's length or its binary chunk did not survive") + } + + c.mu.Lock() + wf, _ := json.Marshal(c.workflow) + uploaded := strings.Join(c.uploaded, ",") + c.mu.Unlock() + if !strings.HasSuffix(uploaded, ".png:image/png") || !strings.HasPrefix(uploaded, "llmgw-") { + t.Errorf("uploaded %q", uploaded) + } + name := strings.TrimSuffix(uploaded, ":image/png") + for _, want := range []string{`"image":"` + name + `"`, `"seed":11`, `"resolution":1024`} { + if !strings.Contains(string(wf), want) { + t.Errorf("%s not in the workflow sent: %s", want, wf) + } + } + deadline := time.Now().Add(5 * time.Second) + for { + u, ok, _ := rt.store.UsageByRequest(t.Context(), j["id"].(string)) + if ok { + if u.Meshes != 1 || u.CostMicros < 250000 { + t.Errorf("recorded %d meshes, %d micros; want one at its price", u.Meshes, u.CostMicros) + } + break + } + if time.Now().After(deadline) { + t.Fatal("no usage recorded") + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestA3DModelIsAskedOfAModelThatMakesThem(t *testing.T) { + c := newMeshComfy(t) + rt, secret := meshRouter(t, c) + h := rt.Handler() + for name, c := range map[string]struct { + path string + body map[string]any + want string + }{ + "a mesh of flux": {meshesPath, map[string]any{"model": "flux", "image": testPicture(t)}, "makes pictures, not 3D models"}, + "no picture": {meshesPath, map[string]any{"model": "trellis2"}, "image is required"}, + "not a picture": {meshesPath, map[string]any{"model": "trellis2", "image": base64.StdEncoding.EncodeToString([]byte("hello"))}, "not a PNG"}, + "an odd resolution": {meshesPath, map[string]any{"model": "trellis2", "image": testPicture(t), "resolution": 700}, "not 512, 1024 or 1536"}, + "a picture of trellis2": {imagesPath, map[string]any{"model": "trellis2", "prompt": "a cube"}, "makes 3D models, not pictures"}, + "a chat with trellis2": {"/v1/chat/completions", map[string]any{"model": "trellis2", "messages": []any{}}, "makes 3D models"}, + } { + code, _, raw := call(t, h, secret, http.MethodPost, c.path, c.body) + if code != http.StatusBadRequest || !strings.Contains(raw, c.want) { + t.Errorf("%s: %d %s", name, code, raw) + } + } + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + req.Header.Set("Authorization", "Bearer "+secret) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if !strings.Contains(rec.Body.String(), `"meshes":true`) || !strings.Contains(rec.Body.String(), `"license":{"name":"MIT"`) { + t.Errorf("/v1/models: %s", rec.Body.String()) + } +} diff --git a/internal/router/proxy/proxy.go b/internal/router/proxy/proxy.go index 61872d8..7cf681b 100644 --- a/internal/router/proxy/proxy.go +++ b/internal/router/proxy/proxy.go @@ -226,6 +226,8 @@ type Outcome struct { Usage *Usage // Images is how many pictures were made, by an endpoint that makes them. Images int + // Meshes is how many 3D models were made. + Meshes int // Bytes is how many bytes were relayed to the client. Bytes int64 // Reply is the assistant's text, when CaptureReply asked for it. Truncated diff --git a/internal/router/router.go b/internal/router/router.go index 403e949..34e1095 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -377,7 +377,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { prices := make(map[string]accounting.Price, len(cfg.Auth.Pricing)) for name, p := range cfg.Auth.Pricing { prices[name] = accounting.Price{ - PromptPer1k: p.PromptPer1k, CompletionPer1k: p.CompletionPer1k, PerImage: p.PerImage, + PromptPer1k: p.PromptPer1k, CompletionPer1k: p.CompletionPer1k, PerImage: p.PerImage, PerMesh: p.PerMesh, } } r.recorder = accounting.NewRecorder(st, r.auth, accounting.RecorderConfig{ @@ -394,6 +394,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { PromptPer1k: cfg.Auth.DefaultPrice.PromptPer1k, CompletionPer1k: cfg.Auth.DefaultPrice.CompletionPer1k, PerImage: cfg.Auth.DefaultPrice.PerImage, + PerMesh: cfg.Auth.DefaultPrice.PerMesh, }), r.recorder, log) r.threads = r.store diff --git a/internal/store/export.go b/internal/store/export.go index 7048a70..89f4ef2 100644 --- a/internal/store/export.go +++ b/internal/store/export.go @@ -26,6 +26,7 @@ type ExportRow struct { Completion int64 CostMicros int64 Images int64 + Meshes int64 } // ExportDims are what an export can be broken down by. @@ -52,18 +53,18 @@ func (s *DB) UsageExport(ctx context.Context, from, to time.Time, period string, bucket int64 key, user, project, model string requests, errors, limited, prompt, completion, cost int64 - images int64 + images, meshes int64 } //nolint:gosec // table is one of two literals above rows, err := queryAll(ctx, s.db, `SELECT bucket, key_id, user_id, project_id, served_model, SUM(requests), SUM(errors), SUM(limited), SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_micros), - SUM(images) + SUM(images), SUM(meshes) FROM `+table+` WHERE bucket >= ? AND bucket < ? GROUP BY bucket, key_id, user_id, project_id, served_model`, []any{millis(from), millis(to)}, func(r *sql.Rows) (raw, error) { var x raw err := r.Scan(&x.bucket, &x.key, &x.user, &x.project, &x.model, - &x.requests, &x.errors, &x.limited, &x.prompt, &x.completion, &x.cost, &x.images) + &x.requests, &x.errors, &x.limited, &x.prompt, &x.completion, &x.cost, &x.images, &x.meshes) return x, err }) if err != nil { @@ -141,6 +142,7 @@ func (s *DB) UsageExport(ctx context.Context, from, to time.Time, period string, t.Completion += x.completion t.CostMicros += x.cost t.Images += x.images + t.Meshes += x.meshes } out := make([]ExportRow, 0, len(order)) for _, k := range order { diff --git a/internal/store/migrations/postgres/0045_usage_meshes.sql b/internal/store/migrations/postgres/0045_usage_meshes.sql new file mode 100644 index 0000000..6a77ae1 --- /dev/null +++ b/internal/store/migrations/postgres/0045_usage_meshes.sql @@ -0,0 +1,7 @@ +-- Ported from the SQLite migration of the same number. +-- How many 3D models a request made, as images says how many pictures: what +-- per_mesh is charged by, in the request rows and in the sums. +ALTER TABLE usage ADD COLUMN meshes BIGINT NOT NULL DEFAULT 0; +ALTER TABLE usage_5m ADD COLUMN meshes BIGINT NOT NULL DEFAULT 0; +ALTER TABLE usage_1h ADD COLUMN meshes BIGINT NOT NULL DEFAULT 0; +ALTER TABLE usage_1d ADD COLUMN meshes BIGINT NOT NULL DEFAULT 0; diff --git a/internal/store/migrations/sqlite/0045_usage_meshes.sql b/internal/store/migrations/sqlite/0045_usage_meshes.sql new file mode 100644 index 0000000..942c47b --- /dev/null +++ b/internal/store/migrations/sqlite/0045_usage_meshes.sql @@ -0,0 +1,6 @@ +-- How many 3D models a request made, as images says how many pictures: what +-- per_mesh is charged by, in the request rows and in the sums. +ALTER TABLE usage ADD COLUMN meshes INTEGER NOT NULL DEFAULT 0; +ALTER TABLE usage_5m ADD COLUMN meshes INTEGER NOT NULL DEFAULT 0; +ALTER TABLE usage_1h ADD COLUMN meshes INTEGER NOT NULL DEFAULT 0; +ALTER TABLE usage_1d ADD COLUMN meshes INTEGER NOT NULL DEFAULT 0; diff --git a/internal/store/rollup.go b/internal/store/rollup.go index fe0749d..20cae08 100644 --- a/internal/store/rollup.go +++ b/internal/store/rollup.go @@ -69,7 +69,7 @@ type rollupKey struct { // rollupRow is what a batch adds to one row. type rollupRow struct { requests, errors, limited, prompt int64 - completion, cost, images int64 + completion, cost, images, meshes int64 latency, queue, ttft, ttftN, last int64 lat, ttftH []int64 } @@ -99,6 +99,7 @@ func rollupBatch(t tier, records []Usage) (map[rollupKey]*rollupRow, []rollupKey row.completion += r.CompletionTokens row.cost += r.CostMicros row.images += r.Images + row.meshes += r.Meshes row.latency += r.LatencyMs row.queue += r.QueueMs row.lat[latencyBucket(r.LatencyMs)]++ @@ -117,7 +118,7 @@ var rollupColumns = func() []string { n := len(LatencyBounds) + 1 cols := []string{"bucket", "key_id", "user_id", "project_id", "model", "served_model", "requests", "errors", "limited", "prompt_tokens", "completion_tokens", "cost_micros", - "latency_ms_sum", "queue_ms_sum", "ttft_ms_sum", "ttft_n", "last_ts", "images"} + "latency_ms_sum", "queue_ms_sum", "ttft_ms_sum", "ttft_n", "last_ts", "images", "meshes"} for i := range n { cols = append(cols, fmt.Sprintf("lat_%d", i)) } @@ -157,7 +158,7 @@ func writeRollup(ctx context.Context, tx conn, records []Usage) error { r := rows[k] args := []any{k.bucket, k.keyID, k.userID, k.projectID, k.model, k.servedModel, r.requests, r.errors, r.limited, r.prompt, r.completion, r.cost, - r.latency, r.queue, r.ttft, r.ttftN, r.last, r.images} + r.latency, r.queue, r.ttft, r.ttftN, r.last, r.images, r.meshes} for _, v := range r.lat { args = append(args, v) } @@ -188,11 +189,11 @@ func usageSince(since time.Time) (string, []any) { from := millis(since) to5m, to1h, to1d := tier5m.ceil(since), tier1h.ceil(since), tier1d.ceil(since) const cols = `key_id, user_id, project_id, model, served_model, requests, errors, limited, - prompt_tokens, completion_tokens, cost_micros, latency_ms_sum, last_ts, images` + prompt_tokens, completion_tokens, cost_micros, latency_ms_sum, last_ts, images, meshes` return `(SELECT key_id, user_id, project_id, model, served_model, 1 AS requests, CASE WHEN status >= 400 THEN 1 ELSE 0 END AS errors, CASE WHEN status = 429 THEN 1 ELSE 0 END AS limited, - prompt_tokens, completion_tokens, cost_micros, latency_ms AS latency_ms_sum, ts AS last_ts, images + prompt_tokens, completion_tokens, cost_micros, latency_ms AS latency_ms_sum, ts AS last_ts, images, meshes FROM usage WHERE ts >= ? AND ts < ? UNION ALL SELECT ` + cols + ` FROM usage_5m WHERE bucket >= ? AND bucket < ? UNION ALL SELECT ` + cols + ` FROM usage_1h WHERE bucket >= ? AND bucket < ? diff --git a/internal/store/sqlite.go b/internal/store/sqlite.go index 62ad830..ecdda69 100644 --- a/internal/store/sqlite.go +++ b/internal/store/sqlite.go @@ -442,8 +442,8 @@ func (s *DB) WriteUsage(ctx context.Context, records []Usage) error { `INSERT INTO usage (ts, key_id, project_id, model, served_model, prompt_tokens, completion_tokens, cost_micros, status, latency_ms, user_id, request_id, replica, queue_ms, ttft_ms, - guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit, images) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`) + guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit, images, meshes) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`) if err != nil { return err } @@ -457,7 +457,7 @@ func (s *DB) WriteUsage(ctx context.Context, records []Usage) error { r.Model, r.ServedModel, r.PromptTokens, r.CompletionTokens, r.CostMicros, r.Status, r.LatencyMs, r.UserID, r.RequestID, r.Replica, r.QueueMs, r.TTFTMs, - r.GuardAction, r.GuardPolicy, r.GuardCategory, r.GuardRule, r.GuardMs, r.EndUser, tagsJSON(r.Tags), hit, r.Images); err != nil { + r.GuardAction, r.GuardPolicy, r.GuardCategory, r.GuardRule, r.GuardMs, r.EndUser, tagsJSON(r.Tags), hit, r.Images, r.Meshes); err != nil { return err } } @@ -871,12 +871,12 @@ func (s *DB) UsageByRequest(ctx context.Context, requestID string) (Usage, bool, err := s.db.QueryRowContext(ctx, `SELECT ts, key_id, project_id, COALESCE(user_id, ''), model, served_model, prompt_tokens, completion_tokens, cost_micros, status, latency_ms, request_id, replica, queue_ms, ttft_ms, - guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit, images + guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit, images, meshes FROM usage WHERE request_id = ? ORDER BY ts DESC LIMIT 1`, requestID). Scan(&ts, &u.KeyID, &u.ProjectID, &u.UserID, &u.Model, &u.ServedModel, &u.PromptTokens, &u.CompletionTokens, &u.CostMicros, &u.Status, &u.LatencyMs, &u.RequestID, &u.Replica, &u.QueueMs, &u.TTFTMs, &u.GuardAction, &u.GuardPolicy, &u.GuardCategory, &u.GuardRule, &u.GuardMs, - &u.EndUser, &tags, &hit, &u.Images) + &u.EndUser, &tags, &hit, &u.Images, &u.Meshes) if errors.Is(err, sql.ErrNoRows) { return Usage{}, false, nil } diff --git a/internal/store/store.go b/internal/store/store.go index 73d0c5f..18382b1 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -243,6 +243,8 @@ type Usage struct { CompletionTokens int64 // Images is how many pictures were made, for a request for them. Images int64 + // Meshes is how many 3D models were made. + Meshes int64 // CostMicros is what this request spent against the budget. CostMicros int64 Status int