Compare commits
8
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3894e8e5fd | ||
|
|
751533f853 | ||
|
|
79a503f824 | ||
|
|
2f06c0e513 | ||
|
|
6b7750a0dd | ||
|
|
a133270725 | ||
|
|
8932502768 | ||
|
|
92733317a9 |
No files matched your search
+36
@@ -1065,6 +1065,42 @@ training yet; quality gating of a published adapter is by hand (ctl eval).
|
||||
|
||||
---
|
||||
|
||||
### F5. Large imports (M) — measured 2026-10-09, open
|
||||
The metallurgy corpus (2,314 files, 108,779 chunks, bge-m3 on the one
|
||||
GPU) went in at about 25 chunks/s, the same over a phone hotspot and the
|
||||
LAN. The engine's slots were the limit: llama.cpp `--parallel 4` embeds
|
||||
four texts per pass, while llmgw-kb sends 32 a request from 6 workers. The
|
||||
usage rows show 4.6 s a request on average, 1.8 s of it in the gateway's
|
||||
queue and 0.04 s in the guard. The GPU sat at 18–50% and the store machine
|
||||
at a third of its CPUs.
|
||||
- An embedding model's profile with many slots by default (`--parallel`
|
||||
32–64, `--ctx-size` to match, `--batch-size`/`--ubatch-size` at least the
|
||||
longest chunk), or `ctl model serve` choosing that for a model with
|
||||
`capabilities.embeddings`. With the default 512, one 516-token chunk failed.
|
||||
- Embedding on the gateway's side of a document PUT, so the client sends
|
||||
text only: today an embeddings answer is ~21 KB of JSON a chunk, and
|
||||
the text and vector go up again (~11 KB).
|
||||
- The store next: HNSW inserts took half of one of pg01's two cores; at a
|
||||
million chunks the vectors alone outgrow its 4 GB. halfvec (pgvector ≥
|
||||
0.7) and a larger store machine.
|
||||
- Done: a batch of embeddings checked text by text against the context
|
||||
(branch `embeddings-batch`); it was refused as one over-long prompt.
|
||||
- Done: `query_prefix` on an embedding model, before questions only.
|
||||
- Qwen3-Embedding-4B on the RTX 6000 (2026-10-09): llama.cpp Q8 1,575
|
||||
tok/s as first set up (32k cache, 16 slots), 2,700 with a 16k cache,
|
||||
FlashAttention and ubatch 1–2k; f16 no better. vLLM fp16 (`--runner
|
||||
pooling`) 6,200 direct, 4,700–4,900 through the gateway in the import,
|
||||
and its vectors match llama.cpp's (cosine 0.9995). The 103,579 chunks
|
||||
took 3 h. Wanted: an embedding model's GPU profile on vLLM and CPU
|
||||
profile on llama.cpp, which needs a source per profile (the catalog has
|
||||
one per model; the import swapped the entry by hand), and an admission
|
||||
queue limit that suits bulk imports (30 s turned some away).
|
||||
- Retrieval on 40 questions written from the corpus, document in the top
|
||||
5: bge-m3 31/40 (Turkish questions 13/21), Qwen3-Embedding-4B 37/40
|
||||
(18/21), with its query prefix 36/40 and first place 31/40 instead of
|
||||
28 (MRR 0.738 → 0.789 → 0.838). Next: translating the question with the
|
||||
chat model, and a reranker.
|
||||
|
||||
## D. Facts worth knowing before starting
|
||||
|
||||
- Admin API: `127.0.0.1:9090`, no auth, `ctlClient` in `internal/cli/ctl.go`.
|
||||
|
||||
@@ -8565,6 +8565,10 @@
|
||||
},
|
||||
"embeddings": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"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."
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -8605,6 +8609,14 @@
|
||||
"peer": {
|
||||
"type": "string",
|
||||
"description": "Peer names the federated llmgw cluster this model is offered by, for a model that federation added (`llmgw ctl federation add`) rather than one written by hand. Such a model is named PEER/MODEL, comes and goes with the peer's own list, and is sent only while the peer is up."
|
||||
},
|
||||
"api": {
|
||||
"type": "string",
|
||||
"description": "API is what the provider speaks: OpenAI's, unless this says comfyui — a ComfyUI server, asked for pictures at /v1/images/generations and sent the workflow below with the request's prompt in it. Model is then the checkpoint file the workflow loads, as ComfyUI lists it."
|
||||
},
|
||||
"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."
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -8710,6 +8722,10 @@
|
||||
"$ref": "#/components/schemas/catalog.Capabilities",
|
||||
"description": "Capabilities is what this model can do beyond plain completion."
|
||||
},
|
||||
"query_prefix": {
|
||||
"type": "string",
|
||||
"description": "QueryPrefix is put before a question an embedding model embeds for a knowledge-base search; documents are embedded without it. Models trained with an instruction on the query side want it, Qwen3-Embedding \"Instruct: Given a web search query, retrieve relevant passages that answer the query\\nQuery: \" for one."
|
||||
},
|
||||
"role": {
|
||||
"type": "string",
|
||||
"description": "Role \"guard\" is a model that judges the gateway's traffic for the guardrails rather than serving clients: hidden from /v1/models and the chat page, refused to any key that names it, critical by default, and always running. Guard says how its answers are read."
|
||||
|
||||
@@ -1270,6 +1270,51 @@ prices and budgets, rate limits, usage records, and key allowlists.
|
||||
whose allowlist leaves the external model out cannot reach it, as a fallback or
|
||||
by name. The cluster page and `ctl models` mark it `external`.
|
||||
|
||||
### Pictures from ComfyUI
|
||||
|
||||
A ComfyUI server is an external model with `api: comfyui`. Clients ask it for
|
||||
pictures at OpenAI's `/v1/images/generations`; llmgw fills a ComfyUI workflow
|
||||
with the request, queues it, waits, and sends the pictures back as `b64_json`:
|
||||
|
||||
```yaml
|
||||
- name: flux-schnell
|
||||
source: external
|
||||
external:
|
||||
url: http://10.0.0.147:8188 # ComfyUI, started with --listen
|
||||
model: flux1-schnell-fp8.safetensors # the checkpoint, as ComfyUI lists it
|
||||
api: comfyui
|
||||
workflow: flux-schnell # or flux-dev, or a JSON file
|
||||
```
|
||||
|
||||
`flux-schnell` (the default) and `flux-dev` are built in. They suit the
|
||||
single-file FLUX checkpoints that ComfyUI loads with `CheckpointLoaderSimple`.
|
||||
For any other model, export the workflow from ComfyUI with **Save (API)** and
|
||||
name the file. In the file, put placeholders where the request's values go:
|
||||
|
||||
- `{{prompt}}` and `{{negative_prompt}}` are the prompts.
|
||||
- `{{width}}` and `{{height}}` are the size.
|
||||
- `{{n}}` is the batch size.
|
||||
- `{{seed}}` is the seed.
|
||||
- `{{steps|4}}` and `{{guidance|3.5}}` take the value after the `|` when the
|
||||
request leaves them out.
|
||||
- `{{checkpoint}}` is the `model` above.
|
||||
|
||||
A request is OpenAI's: `prompt`, `n` (at most 4), `size` (each side a multiple
|
||||
of 16, from 256 to 2048) and `response_format: b64_json`. Beyond those, it may
|
||||
send `seed`, `steps`, `negative_prompt` and `guidance`. The seed used comes
|
||||
back in `x-llmgw-seed`. `url` is not offered, because pictures are kept
|
||||
nowhere. ComfyUI also writes the workflow into each PNG; llmgw strips it
|
||||
before the picture is sent.
|
||||
|
||||
Requests for pictures pass the same checks as any other request: keys,
|
||||
allowlists, limits, admission, and the guardrails on the prompt. A client that
|
||||
leaves takes its workflow off the ComfyUI queue. Usage is recorded, but for
|
||||
now it is charged as the prompt's tokens; a per-picture price is not yet
|
||||
available.
|
||||
|
||||
Start ComfyUI with `--disable-smart-memory` on a card that llmgw also places
|
||||
models on. Without it, ComfyUI keeps the model on the card after each picture.
|
||||
|
||||
### Federating clusters
|
||||
|
||||
Another llmgw cluster — another site, another team's GPUs — can be a provider
|
||||
|
||||
@@ -45,6 +45,15 @@ is available to everyone who may use them.
|
||||
A collection stays with the model it was made with: vectors from two
|
||||
models cannot be compared.
|
||||
|
||||
Some embedding models are trained with an instruction before the
|
||||
question and none before the documents. Give it as the entry's
|
||||
`query_prefix`; llmgw puts it before questions only. For
|
||||
Qwen3-Embedding:
|
||||
|
||||
```yaml
|
||||
query_prefix: "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: "
|
||||
```
|
||||
|
||||
With pgvector 0.7 or later a collection's index is kept at half
|
||||
precision (`halfvec`), which halves its memory and takes models of up to
|
||||
4000 dimensions; with an older pgvector the limit is 2000. A collection
|
||||
|
||||
@@ -60,6 +60,7 @@ rather than repeated.
|
||||
| `tools` | `bool` | [Capabilities](#capabilities) |
|
||||
| `vision` | `bool` | [Capabilities](#capabilities) |
|
||||
| `embeddings` | `bool` | [Capabilities](#capabilities) |
|
||||
| `images` | `bool` | [Capabilities](#capabilities) |
|
||||
| `interval` | `time.Duration` | [Control](#control) |
|
||||
| `base_backoff` | `time.Duration` | [Control](#control) |
|
||||
| `max_backoff` | `time.Duration` | [Control](#control) |
|
||||
@@ -76,6 +77,7 @@ rather than repeated.
|
||||
| `api_key` | `string` | [Endpoint](#endpoint) |
|
||||
| `models` | `[]string` | [Endpoint](#endpoint) |
|
||||
| `engine` | `string` | [Endpoint](#endpoint) |
|
||||
| `workflow` | `string` | [Endpoint](#endpoint) |
|
||||
| `pool` | `string` | [Endpoint](#endpoint) |
|
||||
| `weight` | `float64` | [Endpoint](#endpoint) |
|
||||
| `priority` | `int` | [Endpoint](#endpoint) |
|
||||
@@ -90,6 +92,8 @@ rather than repeated.
|
||||
| `model` | `string` | [External](#external) |
|
||||
| `api_key_file` | `string` | [External](#external) |
|
||||
| `peer` | `string` | [External](#external) |
|
||||
| `api` | `string` | [External](#external) |
|
||||
| `workflow` | `string` | [External](#external) |
|
||||
| `budgets` | `[]FacilityBudget` | [Facility](#facility) |
|
||||
| `margin` | `float64` | [Facility](#facility) |
|
||||
| `gpu_watts` | `float64` | [Facility](#facility) |
|
||||
@@ -145,6 +149,7 @@ rather than repeated.
|
||||
| `context_length` | `int` | [Model](#model) |
|
||||
| `max_output` | `int` | [Model](#model) |
|
||||
| `capabilities` | `Capabilities` | [Model](#model) |
|
||||
| `query_prefix` | `string` | [Model](#model) |
|
||||
| `role` | `string` | [Model](#model) |
|
||||
| `guard` | `*GuardModel` | [Model](#model) |
|
||||
| `origin` | `*Origin` | [Model](#model) |
|
||||
@@ -715,6 +720,14 @@ token is MaxTTFTRatio times this model's. Zero is 50, 0.05 and 1.5.
|
||||
|
||||
`bool`
|
||||
|
||||
### `images`
|
||||
|
||||
`bool`
|
||||
|
||||
Images is a model that makes pictures, asked for at
|
||||
/v1/images/generations. Implied by an external model whose API is
|
||||
ComfyUI's.
|
||||
|
||||
## Control
|
||||
|
||||
### `interval`
|
||||
@@ -828,6 +841,13 @@ with one.
|
||||
Engine names the inference engine, which decides what its metrics are
|
||||
called. Empty means vLLM's names.
|
||||
|
||||
### `workflow`
|
||||
|
||||
`string`
|
||||
|
||||
Workflow is the ComfyUI workflow for engine: comfyui, a built-in's
|
||||
name or its JSON; flux-schnell unless set.
|
||||
|
||||
### `pool`
|
||||
|
||||
`string`
|
||||
@@ -919,6 +939,23 @@ a model that federation added (`llmgw ctl federation add`) rather
|
||||
than one written by hand. Such a model is named PEER/MODEL, comes and
|
||||
goes with the peer's own list, and is sent only while the peer is up.
|
||||
|
||||
### `api`
|
||||
|
||||
`string`
|
||||
|
||||
API is what the provider speaks: OpenAI's, unless this says comfyui —
|
||||
a ComfyUI server, asked for pictures at /v1/images/generations and
|
||||
sent the workflow below with the request's prompt in it. Model is then
|
||||
the checkpoint file the workflow loads, as ComfyUI lists it.
|
||||
|
||||
### `workflow`
|
||||
|
||||
`string`
|
||||
|
||||
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.
|
||||
|
||||
## Facility
|
||||
|
||||
### `budgets`
|
||||
@@ -1389,6 +1426,16 @@ Capabilities is what this model can do beyond plain completion.
|
||||
|
||||
Takes a [Capabilities](#capabilities) block.
|
||||
|
||||
### `query_prefix`
|
||||
|
||||
`string`
|
||||
|
||||
QueryPrefix is put before a question an embedding model embeds for a
|
||||
knowledge-base search; documents are embedded without it. Models
|
||||
trained with an instruction on the query side want it, Qwen3-Embedding
|
||||
"Instruct: Given a web search query, retrieve relevant passages that
|
||||
answer the query\nQuery: " for one.
|
||||
|
||||
### `role`
|
||||
|
||||
`string`
|
||||
|
||||
@@ -0,0 +1,529 @@
|
||||
// Package comfy has images made by ComfyUI: a workflow filled in with what
|
||||
// was asked for, queued, waited on, and its pictures fetched.
|
||||
//
|
||||
// ComfyUI speaks no OpenAI. A workflow is a graph of nodes that is queued with
|
||||
// POST /prompt, is done when GET /history/{id} lists it, and leaves its
|
||||
// pictures to be read one by one from GET /view. What the graph is depends on
|
||||
// the model — FLUX schnell is not FLUX dev, and neither is SDXL — so the graph
|
||||
// is a template the catalog names, with placeholders where a request's
|
||||
// prompt, size and seed go.
|
||||
package comfy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"embed"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
//go:embed workflows/*.json
|
||||
var builtin embed.FS
|
||||
|
||||
// DefaultWorkflow is the template used when the catalog names none.
|
||||
const DefaultWorkflow = "flux-schnell"
|
||||
|
||||
// Builtins are the workflows that come with llmgw, by name.
|
||||
func Builtins() []string {
|
||||
entries, _ := builtin.ReadDir("workflows")
|
||||
out := make([]string, 0, len(entries))
|
||||
for _, e := range entries {
|
||||
out = append(out, strings.TrimSuffix(e.Name(), ".json"))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Template is a workflow in ComfyUI's API format, with placeholders.
|
||||
//
|
||||
// A placeholder is a string value "{{name}}", or "{{name|default}}" for one a
|
||||
// request may leave out. One that is the whole value becomes the value
|
||||
// itself, a number where the node wants a number; one inside a longer string
|
||||
// is written into it. The names are those of Params.
|
||||
type Template struct {
|
||||
graph map[string]any
|
||||
}
|
||||
|
||||
// Load reads a template: a built-in's name, or the workflow's JSON itself.
|
||||
func Load(nameOrJSON string) (*Template, error) {
|
||||
s := strings.TrimSpace(nameOrJSON)
|
||||
if s == "" {
|
||||
s = DefaultWorkflow
|
||||
}
|
||||
raw := []byte(s)
|
||||
if !strings.HasPrefix(s, "{") {
|
||||
b, err := builtin.ReadFile("workflows/" + s + ".json")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("no built-in workflow %q (there are %s)", s, strings.Join(Builtins(), ", "))
|
||||
}
|
||||
raw = b
|
||||
}
|
||||
var g map[string]any
|
||||
if err := json.Unmarshal(raw, &g); err != nil {
|
||||
return nil, fmt.Errorf("the workflow is not JSON: %w", err)
|
||||
}
|
||||
if len(g) == 0 {
|
||||
return nil, errors.New("the workflow has no nodes")
|
||||
}
|
||||
for id, n := range g {
|
||||
node, ok := n.(map[string]any)
|
||||
if !ok || node["class_type"] == nil {
|
||||
// The UI's own save format has a "nodes" list instead, and is
|
||||
// not something /prompt accepts.
|
||||
return nil, fmt.Errorf("node %q has no class_type: the workflow must be saved in API format", id)
|
||||
}
|
||||
}
|
||||
return &Template{graph: g}, nil
|
||||
}
|
||||
|
||||
// Params are what a request decides. A zero number is one not given, which
|
||||
// leaves it to the template's default.
|
||||
type Params struct {
|
||||
Checkpoint string
|
||||
Prompt string
|
||||
NegativePrompt string
|
||||
Width, Height int
|
||||
N int
|
||||
Seed int64
|
||||
Steps int
|
||||
Guidance float64
|
||||
}
|
||||
|
||||
func (p Params) value(name string) (any, bool) {
|
||||
switch name {
|
||||
case "checkpoint":
|
||||
return p.Checkpoint, p.Checkpoint != ""
|
||||
case "prompt":
|
||||
return p.Prompt, true
|
||||
case "negative_prompt":
|
||||
return p.NegativePrompt, true
|
||||
case "width":
|
||||
return p.Width, p.Width > 0
|
||||
case "height":
|
||||
return p.Height, p.Height > 0
|
||||
case "n":
|
||||
return max(p.N, 1), true
|
||||
case "seed":
|
||||
return p.Seed, true
|
||||
case "steps":
|
||||
return p.Steps, p.Steps > 0
|
||||
case "guidance":
|
||||
return p.Guidance, p.Guidance > 0
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// Fill is the workflow for one request.
|
||||
func (t *Template) Fill(p Params) (map[string]any, error) {
|
||||
var missing []string
|
||||
var fill func(v any) any
|
||||
fill = func(v any) any {
|
||||
switch v := v.(type) {
|
||||
case map[string]any:
|
||||
out := make(map[string]any, len(v))
|
||||
for k, x := range v {
|
||||
out[k] = fill(x)
|
||||
}
|
||||
return out
|
||||
case []any:
|
||||
out := make([]any, len(v))
|
||||
for i, x := range v {
|
||||
out[i] = fill(x)
|
||||
}
|
||||
return out
|
||||
case string:
|
||||
return fillString(v, p, &missing)
|
||||
}
|
||||
return v
|
||||
}
|
||||
g := fill(t.graph).(map[string]any)
|
||||
if len(missing) > 0 {
|
||||
slices.Sort(missing)
|
||||
return nil, fmt.Errorf("the workflow wants %s, which the request did not give and the workflow has no default for",
|
||||
strings.Join(slices.Compact(missing), ", "))
|
||||
}
|
||||
return g, nil
|
||||
}
|
||||
|
||||
func fillString(s string, p Params, missing *[]string) any {
|
||||
if !strings.Contains(s, "{{") {
|
||||
return s
|
||||
}
|
||||
resolve := func(ph string) (any, bool) {
|
||||
name, def, hasDef := strings.Cut(ph, "|")
|
||||
name = strings.TrimSpace(name)
|
||||
if v, ok := p.value(name); ok {
|
||||
return v, true
|
||||
}
|
||||
if hasDef {
|
||||
def = strings.TrimSpace(def)
|
||||
if f, err := strconv.ParseFloat(def, 64); err == nil {
|
||||
return f, true
|
||||
}
|
||||
return def, true
|
||||
}
|
||||
*missing = append(*missing, name)
|
||||
return nil, false
|
||||
}
|
||||
// The whole value: what the node gets is the value, not its text.
|
||||
if strings.HasPrefix(s, "{{") && strings.HasSuffix(s, "}}") && strings.Count(s, "{{") == 1 {
|
||||
v, _ := resolve(s[2 : len(s)-2])
|
||||
return v
|
||||
}
|
||||
var b strings.Builder
|
||||
for {
|
||||
i := strings.Index(s, "{{")
|
||||
j := strings.Index(s, "}}")
|
||||
if i < 0 || j < i {
|
||||
b.WriteString(s)
|
||||
return b.String()
|
||||
}
|
||||
b.WriteString(s[:i])
|
||||
if v, ok := resolve(s[i+2 : j]); ok {
|
||||
fmt.Fprint(&b, v)
|
||||
}
|
||||
s = s[j+2:]
|
||||
}
|
||||
}
|
||||
|
||||
// RandomSeed is a seed for a request that named none. ComfyUI's are 64-bit,
|
||||
// but JSON numbers are read as doubles on the way: kept to 2^53.
|
||||
func RandomSeed() int64 {
|
||||
n, err := rand.Int(rand.Reader, big.NewInt(1<<53))
|
||||
if err != nil {
|
||||
return time.Now().UnixNano() & (1<<53 - 1)
|
||||
}
|
||||
return n.Int64()
|
||||
}
|
||||
|
||||
// Image is one picture made.
|
||||
type Image struct {
|
||||
Data []byte
|
||||
ContentType string
|
||||
}
|
||||
|
||||
// Error is a failure, said as what it means for the request.
|
||||
type Error struct {
|
||||
Msg string
|
||||
// Retryable is a failure of this server, which another one might not
|
||||
// have: it could not be reached, or it answered 5xx.
|
||||
Retryable bool
|
||||
}
|
||||
|
||||
func (e *Error) Error() string { return e.Msg }
|
||||
|
||||
// Client queues workflows and waits for them.
|
||||
type Client struct {
|
||||
HTTP *http.Client
|
||||
// Poll is how often a queued workflow is asked after.
|
||||
Poll time.Duration
|
||||
}
|
||||
|
||||
// Generate runs a workflow on the ComfyUI at base and returns its pictures,
|
||||
// in the order of the nodes that made them. A context that ends takes the
|
||||
// workflow out of the queue, or interrupts it if it is running, so a client
|
||||
// that went away does not leave the card working for nobody.
|
||||
func (c *Client) Generate(ctx context.Context, base, apiKey string, workflow map[string]any) ([]Image, error) {
|
||||
base = strings.TrimRight(base, "/")
|
||||
client := make([]byte, 8)
|
||||
_, _ = rand.Read(client)
|
||||
body, err := json.Marshal(map[string]any{"prompt": workflow, "client_id": "llmgw-" + hex.EncodeToString(client)})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var queued struct {
|
||||
PromptID string `json:"prompt_id"`
|
||||
Error json.RawMessage `json:"error"`
|
||||
NodeErrors json.RawMessage `json:"node_errors"`
|
||||
}
|
||||
status, err := c.do(ctx, http.MethodPost, base+"/prompt", apiKey, body, &queued)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if status != http.StatusOK || queued.PromptID == "" {
|
||||
// A workflow ComfyUI will not run: a checkpoint it does not have, a
|
||||
// node it does not know. The catalog's to fix, not the client's, and
|
||||
// no other replica set up the same way would take it either.
|
||||
return nil, &Error{Msg: "ComfyUI refused the workflow: " + refusal(queued.Error, queued.NodeErrors)}
|
||||
}
|
||||
id := queued.PromptID
|
||||
|
||||
poll := c.Poll
|
||||
if poll <= 0 {
|
||||
poll = 250 * time.Millisecond
|
||||
}
|
||||
t := time.NewTicker(poll)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
c.cancel(context.WithoutCancel(ctx), base, apiKey, id)
|
||||
return nil, ctx.Err()
|
||||
case <-t.C:
|
||||
}
|
||||
var hist map[string]historyEntry
|
||||
if _, err := c.do(ctx, http.MethodGet, base+"/history/"+url.PathEscape(id), apiKey, nil, &hist); err != nil {
|
||||
if ctx.Err() != nil {
|
||||
continue // said by the select
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
h, ok := hist[id]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if h.Status.StatusStr == "error" {
|
||||
return nil, &Error{Msg: "the workflow failed in ComfyUI: " + h.failure()}
|
||||
}
|
||||
if !h.Status.Completed {
|
||||
continue
|
||||
}
|
||||
return c.fetch(ctx, base, apiKey, h)
|
||||
}
|
||||
}
|
||||
|
||||
type historyEntry struct {
|
||||
Status struct {
|
||||
StatusStr string `json:"status_str"`
|
||||
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"`
|
||||
}
|
||||
|
||||
// failure is what ComfyUI said went wrong, from its execution_error.
|
||||
func (h historyEntry) failure() string {
|
||||
for _, m := range h.Status.Messages {
|
||||
if len(m) < 2 || string(m[0]) != `"execution_error"` {
|
||||
continue
|
||||
}
|
||||
var e struct {
|
||||
NodeType string `json:"node_type"`
|
||||
Message string `json:"exception_message"`
|
||||
}
|
||||
if json.Unmarshal(m[1], &e) == nil && e.Message != "" {
|
||||
return strings.TrimSpace(e.NodeType + ": " + strings.TrimSpace(e.Message))
|
||||
}
|
||||
}
|
||||
return "no reason given"
|
||||
}
|
||||
|
||||
func (c *Client) fetch(ctx context.Context, base, apiKey string, h historyEntry) ([]Image, error) {
|
||||
nodes := make([]string, 0, len(h.Outputs))
|
||||
for id := range h.Outputs {
|
||||
nodes = append(nodes, id)
|
||||
}
|
||||
// Numerically where they are numbers, which is how ComfyUI names them.
|
||||
slices.SortFunc(nodes, func(a, b string) int {
|
||||
x, ea := strconv.Atoi(a)
|
||||
y, eb := strconv.Atoi(b)
|
||||
if ea == nil && eb == nil {
|
||||
return x - y
|
||||
}
|
||||
return strings.Compare(a, b)
|
||||
})
|
||||
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
|
||||
}
|
||||
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 out, nil
|
||||
}
|
||||
|
||||
// cancel takes a workflow out of the queue, or stops it if it is the one
|
||||
// running. Best effort: the server may be gone, which also stops it.
|
||||
func (c *Client) cancel(ctx context.Context, base, apiKey, id string) {
|
||||
ctx, done := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer done()
|
||||
body, _ := json.Marshal(map[string]any{"delete": []string{id}})
|
||||
_, _ = c.do(ctx, http.MethodPost, base+"/queue", apiKey, body, nil)
|
||||
var q struct {
|
||||
Running [][]json.RawMessage `json:"queue_running"`
|
||||
}
|
||||
if _, err := c.do(ctx, http.MethodGet, base+"/queue", apiKey, nil, &q); err != nil {
|
||||
return
|
||||
}
|
||||
for _, r := range q.Running {
|
||||
if len(r) > 1 && string(r[1]) == strconv.Quote(id) {
|
||||
_, _ = c.do(ctx, http.MethodPost, base+"/interrupt", apiKey, []byte("{}"), nil)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) client() *http.Client {
|
||||
if c.HTTP != nil {
|
||||
return c.HTTP
|
||||
}
|
||||
return http.DefaultClient
|
||||
}
|
||||
|
||||
func (c *Client) request(ctx context.Context, method, u, apiKey string, body []byte) (*http.Response, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, method, u, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
if apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
resp, err := c.client().Do(req)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
return nil, &Error{Msg: "ComfyUI could not be reached: " + err.Error(), Retryable: true}
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// do sends a request and reads a JSON answer into out. A 5xx is an error;
|
||||
// any other status is returned with whatever body could be read.
|
||||
func (c *Client) do(ctx context.Context, method, u, apiKey string, body []byte, out any) (int, error) {
|
||||
resp, err := c.request(ctx, method, u, apiKey, body)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(io.LimitReader(resp.Body, 16<<20))
|
||||
if err != nil {
|
||||
return 0, &Error{Msg: "reading ComfyUI's answer: " + err.Error(), Retryable: true}
|
||||
}
|
||||
if resp.StatusCode >= 500 {
|
||||
return resp.StatusCode, &Error{Msg: fmt.Sprintf("ComfyUI answered %d: %s", resp.StatusCode, snippet(b)), Retryable: true}
|
||||
}
|
||||
if out != nil && len(b) > 0 {
|
||||
if err := json.Unmarshal(b, out); err != nil && resp.StatusCode == http.StatusOK {
|
||||
return resp.StatusCode, &Error{Msg: "ComfyUI's answer is not JSON: " + snippet(b)}
|
||||
}
|
||||
}
|
||||
return resp.StatusCode, nil
|
||||
}
|
||||
|
||||
func (c *Client) get(ctx context.Context, u, apiKey string) (Image, error) {
|
||||
resp, err := c.request(ctx, http.MethodGet, u, apiKey, nil)
|
||||
if err != nil {
|
||||
return Image{}, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(io.LimitReader(resp.Body, 64<<20))
|
||||
if err != nil {
|
||||
return Image{}, &Error{Msg: "reading a picture 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)}
|
||||
}
|
||||
ct := resp.Header.Get("Content-Type")
|
||||
if ct == "" || ct == "application/octet-stream" {
|
||||
ct = http.DetectContentType(b)
|
||||
}
|
||||
if ct == "image/png" {
|
||||
b = withoutPNGText(b)
|
||||
}
|
||||
return Image{Data: b, ContentType: ct}, nil
|
||||
}
|
||||
|
||||
// withoutPNGText is a PNG without its text chunks. ComfyUI writes the whole
|
||||
// workflow into every picture it saves — the prompt, and the checkpoint and
|
||||
// nodes behind it — and what a picture says about how it was made is the
|
||||
// operator's to publish, not something to hand every client. A file that is
|
||||
// not a well-formed PNG is returned as it came.
|
||||
func withoutPNGText(b []byte) []byte {
|
||||
const sig = "\x89PNG\r\n\x1a\n"
|
||||
if len(b) < len(sig) || string(b[:len(sig)]) != sig {
|
||||
return b
|
||||
}
|
||||
out := make([]byte, 0, len(b))
|
||||
out = append(out, b[:len(sig)]...)
|
||||
for i := len(sig); i < len(b); {
|
||||
if i+12 > len(b) {
|
||||
return b
|
||||
}
|
||||
n := int(binary.BigEndian.Uint32(b[i:]))
|
||||
end := i + 12 + n
|
||||
if n < 0 || end > len(b) {
|
||||
return b
|
||||
}
|
||||
switch string(b[i+4 : i+8]) {
|
||||
case "tEXt", "iTXt", "zTXt":
|
||||
default:
|
||||
out = append(out, b[i:end]...)
|
||||
}
|
||||
i = end
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// refusal is what /prompt said when it would not queue a workflow.
|
||||
func refusal(e, nodes json.RawMessage) string {
|
||||
var top struct {
|
||||
Message string `json:"message"`
|
||||
Details string `json:"details"`
|
||||
}
|
||||
msg := snippet(e)
|
||||
if json.Unmarshal(e, &top) == nil && top.Message != "" {
|
||||
msg = top.Message
|
||||
if top.Details != "" {
|
||||
msg += ": " + top.Details
|
||||
}
|
||||
}
|
||||
var perNode map[string]struct {
|
||||
ClassType string `json:"class_type"`
|
||||
Errors []struct {
|
||||
Message string `json:"message"`
|
||||
Details string `json:"details"`
|
||||
} `json:"errors"`
|
||||
}
|
||||
if json.Unmarshal(nodes, &perNode) == nil {
|
||||
var parts []string
|
||||
for _, n := range perNode {
|
||||
for _, er := range n.Errors {
|
||||
parts = append(parts, strings.TrimSpace(n.ClassType+": "+er.Message+" "+er.Details))
|
||||
}
|
||||
}
|
||||
slices.Sort(parts)
|
||||
if len(parts) > 0 {
|
||||
msg += " (" + strings.Join(parts, "; ") + ")"
|
||||
}
|
||||
}
|
||||
if msg == "" {
|
||||
msg = "no reason given"
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
func snippet(b []byte) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if len(s) > 300 {
|
||||
s = s[:300] + "…"
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
package comfy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestTheBuiltInWorkflowsFillIn(t *testing.T) {
|
||||
for _, name := range Builtins() {
|
||||
tpl, err := Load(name)
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
g, err := tpl.Fill(Params{Checkpoint: "flux.safetensors", Prompt: `a "red" fox`, Width: 768, Height: 512, N: 2, Seed: 42})
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
b, _ := json.Marshal(g)
|
||||
s := string(b)
|
||||
for _, want := range []string{`"ckpt_name":"flux.safetensors"`, `"text":"a \"red\" fox"`, `"width":768`,
|
||||
`"height":512`, `"batch_size":2`, `"seed":42`} {
|
||||
if !strings.Contains(s, want) {
|
||||
t.Errorf("%s: %s not in %s", name, want, s)
|
||||
}
|
||||
}
|
||||
if strings.Contains(s, "{{") {
|
||||
t.Errorf("%s: a placeholder was left: %s", name, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAStepCountGivenReplacesTheDefault(t *testing.T) {
|
||||
tpl, _ := Load("flux-schnell")
|
||||
g, _ := tpl.Fill(Params{Checkpoint: "c", Prompt: "p", Width: 512, Height: 512})
|
||||
if steps := g["5"].(map[string]any)["inputs"].(map[string]any)["steps"]; steps != 4.0 {
|
||||
t.Errorf("default steps = %v, want 4", steps)
|
||||
}
|
||||
g, _ = tpl.Fill(Params{Checkpoint: "c", Prompt: "p", Width: 512, Height: 512, Steps: 8})
|
||||
if steps := g["5"].(map[string]any)["inputs"].(map[string]any)["steps"]; steps != 8 {
|
||||
t.Errorf("steps = %v, want 8", steps)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPlaceholderWithoutAValueOrADefaultIsSaid(t *testing.T) {
|
||||
tpl, err := Load(`{"1":{"class_type":"X","inputs":{"a":"{{checkpoint}}","b":"size {{width}}x{{height}}"}}}`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := tpl.Fill(Params{Width: 5}); err == nil || !strings.Contains(err.Error(), "checkpoint, height") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
g, err := tpl.Fill(Params{Checkpoint: "c", Width: 5, Height: 6})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if b := g["1"].(map[string]any)["inputs"].(map[string]any)["b"]; b != "size 5x6" {
|
||||
t.Errorf("b = %v", b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWorkflowInTheUIsFormatIsRefused(t *testing.T) {
|
||||
if _, err := Load(`{"nodes":[{"id":1}],"links":[]}`); err == nil || !strings.Contains(err.Error(), "API format") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
if _, err := Load("sdxl-turbo"); err == nil {
|
||||
t.Fatal("an unknown built-in was loaded")
|
||||
}
|
||||
}
|
||||
|
||||
// fakeComfy is just enough of ComfyUI: /prompt queues, /history answers
|
||||
// after a few asks, /view serves a picture.
|
||||
type fakeComfy struct {
|
||||
mu sync.Mutex
|
||||
asked int
|
||||
fail string
|
||||
refuse bool
|
||||
deleted []string
|
||||
queued map[string]any
|
||||
readyAt int
|
||||
srv *httptest.Server
|
||||
interrupted bool
|
||||
}
|
||||
|
||||
func newFakeComfy(t *testing.T) *fakeComfy {
|
||||
f := &fakeComfy{readyAt: 2}
|
||||
f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
switch {
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/prompt":
|
||||
if f.refuse {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
fmt.Fprint(w, `{"error":{"message":"Prompt outputs failed validation","details":""},"node_errors":{"1":{"class_type":"CheckpointLoaderSimple","errors":[{"message":"Value not in list","details":"ckpt_name: 'x' not in []"}]}}}`)
|
||||
return
|
||||
}
|
||||
var b struct {
|
||||
Prompt map[string]any `json:"prompt"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&b)
|
||||
f.queued = b.Prompt
|
||||
fmt.Fprint(w, `{"prompt_id":"p1","number":1,"node_errors":{}}`)
|
||||
case r.URL.Path == "/history/p1":
|
||||
f.asked++
|
||||
if f.asked < f.readyAt {
|
||||
fmt.Fprint(w, `{}`)
|
||||
return
|
||||
}
|
||||
if f.fail != "" {
|
||||
fmt.Fprintf(w, `{"p1":{"status":{"status_str":"error","completed":false,"messages":[["execution_start",{}],["execution_error",{"node_type":"KSampler","exception_message":%q}]]},"outputs":{}}}`, f.fail)
|
||||
return
|
||||
}
|
||||
fmt.Fprint(w, `{"p1":{"status":{"status_str":"success","completed":true,"messages":[]},"outputs":{"7":{"images":[{"filename":"a.png","subfolder":"","type":"temp"},{"filename":"b.png","subfolder":"","type":"temp"}]}}}}`)
|
||||
case r.URL.Path == "/view":
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
fmt.Fprint(w, "PNG:"+r.URL.Query().Get("filename"))
|
||||
case r.URL.Path == "/queue" && r.Method == http.MethodPost:
|
||||
var b struct {
|
||||
Delete []string `json:"delete"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&b)
|
||||
f.deleted = append(f.deleted, b.Delete...)
|
||||
case r.URL.Path == "/queue":
|
||||
fmt.Fprint(w, `{"queue_running":[[1,"p1",{},{},[]]],"queue_pending":[]}`)
|
||||
case r.URL.Path == "/interrupt":
|
||||
f.interrupted = true
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
t.Cleanup(f.srv.Close)
|
||||
return f
|
||||
}
|
||||
|
||||
func TestGenerateWaitsForTheWorkflowAndFetchesItsPictures(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
c := &Client{Poll: time.Millisecond}
|
||||
images, err := c.Generate(t.Context(), f.srv.URL, "", map[string]any{"1": map[string]any{"class_type": "X"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(images) != 2 || string(images[0].Data) != "PNG:a.png" || images[1].ContentType != "image/png" {
|
||||
t.Fatalf("images = %+v", images)
|
||||
}
|
||||
if f.queued["1"] == nil {
|
||||
t.Error("the workflow was not what was queued")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWorkflowThatFailsSaysWhy(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
f.fail = "CUDA out of memory"
|
||||
_, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{})
|
||||
if err == nil || !strings.Contains(err.Error(), "KSampler: CUDA out of memory") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWorkflowComfyUIRefusesIsNotRetried(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
f.refuse = true
|
||||
_, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{})
|
||||
ce, ok := err.(*Error)
|
||||
if !ok || ce.Retryable || !strings.Contains(ce.Msg, "Value not in list") {
|
||||
t.Fatalf("err = %#v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAComfyUIThatCannotBeReachedIsRetryable(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
f.srv.Close()
|
||||
_, err := (&Client{Poll: time.Millisecond}).Generate(t.Context(), f.srv.URL, "", map[string]any{})
|
||||
if ce, ok := err.(*Error); !ok || !ce.Retryable {
|
||||
t.Fatalf("err = %#v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAClientThatLeavesTakesItsWorkflowOffTheCard(t *testing.T) {
|
||||
f := newFakeComfy(t)
|
||||
f.readyAt = 1 << 30
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 30*time.Millisecond)
|
||||
defer cancel()
|
||||
_, err := (&Client{Poll: time.Millisecond}).Generate(ctx, f.srv.URL, "", map[string]any{})
|
||||
if err == nil {
|
||||
t.Fatal("no error")
|
||||
}
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
if len(f.deleted) != 1 || f.deleted[0] != "p1" || !f.interrupted {
|
||||
t.Errorf("deleted %v, interrupted %v", f.deleted, f.interrupted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPictureLosesTheWorkflowComfyUIWroteIntoIt(t *testing.T) {
|
||||
chunk := func(kind, data string) []byte {
|
||||
b := make([]byte, 0, 12+len(data))
|
||||
b = binary.BigEndian.AppendUint32(b, uint32(len(data)))
|
||||
b = append(b, kind...)
|
||||
b = append(b, data...)
|
||||
return append(b, 0, 0, 0, 0) // a CRC nobody here checks
|
||||
}
|
||||
png := []byte("\x89PNG\r\n\x1a\n")
|
||||
png = append(png, chunk("IHDR", "0123456789abc")...)
|
||||
png = append(png, chunk("tEXt", "prompt\x00{\"1\":...}")...)
|
||||
png = append(png, chunk("IDAT", "pixels")...)
|
||||
png = append(png, chunk("iTXt", "workflow\x00...")...)
|
||||
png = append(png, chunk("IEND", "")...)
|
||||
got := string(withoutPNGText(png))
|
||||
if strings.Contains(got, "prompt") || strings.Contains(got, "workflow") {
|
||||
t.Errorf("text left in: %q", got)
|
||||
}
|
||||
for _, kind := range []string{"IHDR", "IDAT", "pixels", "IEND"} {
|
||||
if !strings.Contains(got, kind) {
|
||||
t.Errorf("%s gone: %q", kind, got)
|
||||
}
|
||||
}
|
||||
if broken := png[:30]; string(withoutPNGText(broken)) != string(broken) {
|
||||
t.Error("a truncated PNG was changed")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "{{checkpoint}}"}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{prompt}}", "clip": ["1", 1]}},
|
||||
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{negative_prompt}}", "clip": ["1", 1]}},
|
||||
"8": {"class_type": "FluxGuidance", "inputs": {"guidance": "{{guidance|3.5}}", "conditioning": ["2", 0]}},
|
||||
"4": {"class_type": "EmptySD3LatentImage", "inputs": {"width": "{{width}}", "height": "{{height}}", "batch_size": "{{n}}"}},
|
||||
"5": {"class_type": "KSampler", "inputs": {
|
||||
"seed": "{{seed}}", "steps": "{{steps|20}}", "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1,
|
||||
"model": ["1", 0], "positive": ["8", 0], "negative": ["3", 0], "latent_image": ["4", 0]}},
|
||||
"6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
|
||||
"7": {"class_type": "PreviewImage", "inputs": {"images": ["6", 0]}}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "{{checkpoint}}"}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{prompt}}", "clip": ["1", 1]}},
|
||||
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": "{{negative_prompt}}", "clip": ["1", 1]}},
|
||||
"4": {"class_type": "EmptySD3LatentImage", "inputs": {"width": "{{width}}", "height": "{{height}}", "batch_size": "{{n}}"}},
|
||||
"5": {"class_type": "KSampler", "inputs": {
|
||||
"seed": "{{seed}}", "steps": "{{steps|4}}", "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1,
|
||||
"model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["4", 0]}},
|
||||
"6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
|
||||
"7": {"class_type": "PreviewImage", "inputs": {"images": ["6", 0]}}
|
||||
}
|
||||
@@ -723,7 +723,10 @@ type Endpoint struct {
|
||||
Models []string `koanf:"models"`
|
||||
// Engine names the inference engine, which decides what its metrics are
|
||||
// called. Empty means vLLM's names.
|
||||
Engine string `koanf:"engine"`
|
||||
Engine string `koanf:"engine"`
|
||||
// Workflow is the ComfyUI workflow for engine: comfyui, a built-in's
|
||||
// name or its JSON; flux-schnell unless set.
|
||||
Workflow string `koanf:"workflow"`
|
||||
Pool string `koanf:"pool"`
|
||||
Weight float64 `koanf:"weight"`
|
||||
Priority int `koanf:"priority"`
|
||||
|
||||
@@ -499,7 +499,11 @@ func (a *API) kbSearch(w http.ResponseWriter, r *http.Request) {
|
||||
// Said, not swallowed: a search that quietly lost its meaning half
|
||||
// looks like a collection that does not know the answer.
|
||||
if a.kb.Embed != nil {
|
||||
v, err := a.kb.Embed(r.Context(), c.EmbedModel, in.Query)
|
||||
question := in.Query
|
||||
if m, ok := a.ctl.Catalog().Resolve(c.EmbedModel); ok {
|
||||
question = m.QueryPrefix + question
|
||||
}
|
||||
v, err := a.kb.Embed(r.Context(), c.EmbedModel, question)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, c.EmbedModel+" did not embed the question: "+err.Error())
|
||||
return
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/comfy"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/facility"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/placement"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
||||
@@ -363,6 +364,13 @@ type Model struct {
|
||||
// Capabilities is what this model can do beyond plain completion.
|
||||
Capabilities Capabilities `json:"capabilities,omitempty" koanf:"capabilities"`
|
||||
|
||||
// QueryPrefix is put before a question an embedding model embeds for a
|
||||
// knowledge-base search; documents are embedded without it. Models
|
||||
// trained with an instruction on the query side want it, Qwen3-Embedding
|
||||
// "Instruct: Given a web search query, retrieve relevant passages that
|
||||
// answer the query\nQuery: " for one.
|
||||
QueryPrefix string `json:"query_prefix,omitempty" koanf:"query_prefix"`
|
||||
|
||||
// Role "guard" is a model that judges the gateway's traffic for the
|
||||
// guardrails rather than serving clients: hidden from /v1/models and the
|
||||
// chat page, refused to any key that names it, critical by default, and
|
||||
@@ -450,6 +458,15 @@ type Capabilities struct {
|
||||
Tools bool `json:"tools,omitempty" koanf:"tools"`
|
||||
Vision bool `json:"vision,omitempty" koanf:"vision"`
|
||||
Embeddings bool `json:"embeddings,omitempty" koanf:"embeddings"`
|
||||
// Images is a model that makes pictures, asked for at
|
||||
// /v1/images/generations. Implied by an external model whose API is
|
||||
// ComfyUI's.
|
||||
Images bool `json:"images,omitempty" koanf:"images"`
|
||||
}
|
||||
|
||||
// MakesImages reports whether a model answers /v1/images/generations.
|
||||
func (m Model) MakesImages() bool {
|
||||
return m.Capabilities.Images || m.External != nil && m.External.API == APIComfyUI
|
||||
}
|
||||
|
||||
// Catalog is the set of models the cluster should serve.
|
||||
@@ -1094,8 +1111,20 @@ type External struct {
|
||||
// than one written by hand. Such a model is named PEER/MODEL, comes and
|
||||
// goes with the peer's own list, and is sent only while the peer is up.
|
||||
Peer string `json:"peer,omitempty" koanf:"peer"`
|
||||
// API is what the provider speaks: OpenAI's, unless this says comfyui —
|
||||
// a ComfyUI server, asked for pictures at /v1/images/generations and
|
||||
// sent the workflow below with the request's prompt in it. Model is then
|
||||
// the checkpoint file the workflow loads, as ComfyUI lists it.
|
||||
API string `json:"api,omitempty" koanf:"api"`
|
||||
// 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.
|
||||
Workflow string `json:"workflow,omitempty" koanf:"workflow"`
|
||||
}
|
||||
|
||||
// APIComfyUI is External.API for a ComfyUI server.
|
||||
const APIComfyUI = "comfyui"
|
||||
|
||||
// IsPeer reports whether a model is offered by a federated cluster.
|
||||
func (m Model) IsPeer() bool { return m.External != nil && m.External.Peer != "" }
|
||||
|
||||
@@ -1110,6 +1139,23 @@ func (e *External) validate(model string, profiles int) error {
|
||||
if strings.TrimSpace(e.Model) == "" {
|
||||
return fmt.Errorf("model %q: external model needs the name the provider knows it by", model)
|
||||
}
|
||||
switch e.API {
|
||||
case "", "openai":
|
||||
if e.Workflow != "" {
|
||||
return fmt.Errorf("model %q: a workflow is for api: comfyui", model)
|
||||
}
|
||||
case APIComfyUI:
|
||||
if e.Peer != "" {
|
||||
return fmt.Errorf("model %q: a federated model speaks OpenAI's API, not ComfyUI's", model)
|
||||
}
|
||||
// A file is read where the router runs; a name is checked here.
|
||||
if w := e.Workflow; w != "" && !strings.ContainsAny(w, "/.") && !slices.Contains(comfy.Builtins(), w) {
|
||||
return fmt.Errorf("model %q: no built-in workflow %q (there are %s; or name a JSON file)",
|
||||
model, w, strings.Join(comfy.Builtins(), ", "))
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("model %q: external api %q is not openai or comfyui", model, e.API)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -412,12 +413,27 @@ func (c *Controller) externalEndpoints() []endpoint.Info {
|
||||
if m.IsPeer() {
|
||||
id = "peer:" + m.Name
|
||||
}
|
||||
engine, workflow := endpoint.EngineExternal, ""
|
||||
if m.External.API == catalog.APIComfyUI {
|
||||
engine, workflow = endpoint.EngineComfyUI, m.External.Workflow
|
||||
// A file is read here, on the control plane the catalog says it
|
||||
// is on, and sent as the workflow itself.
|
||||
if strings.ContainsAny(workflow, "/.") {
|
||||
b, err := os.ReadFile(workflow)
|
||||
if err != nil {
|
||||
c.log.Warn("a ComfyUI model's workflow could not be read; not offered", "model", m.Name, "err", err)
|
||||
continue
|
||||
}
|
||||
workflow = string(b)
|
||||
}
|
||||
}
|
||||
out = append(out, endpoint.Info{
|
||||
ID: id, URL: strings.TrimRight(m.External.URL, "/"),
|
||||
Models: append([]string{m.Name}, m.Aliases...), APIKey: key,
|
||||
Engine: endpoint.EngineExternal, Pool: "external", Weight: 1,
|
||||
Engine: engine, Pool: "external", Weight: 1,
|
||||
UpstreamModel: m.External.Model,
|
||||
Peer: m.External.Peer, PeerCluster: cluster, PeerTags: tags,
|
||||
Workflow: workflow,
|
||||
})
|
||||
}
|
||||
return out
|
||||
@@ -427,8 +443,9 @@ func (c *Controller) externalEndpoints() []endpoint.Info {
|
||||
func fingerprint(infos []endpoint.Info) string {
|
||||
var b strings.Builder
|
||||
for _, i := range infos {
|
||||
fmt.Fprintf(&b, "%s|%s|%s|%s|%v|%g|%d|%s|%s|%s|%t|%x;", i.ID, i.URL, i.Engine, i.Pool, i.Models, i.Weight,
|
||||
i.Priority, i.UpstreamModel, i.Peer, i.PeerCluster, i.PeerTags, sha256.Sum256([]byte(i.APIKey)))
|
||||
fmt.Fprintf(&b, "%s|%s|%s|%s|%v|%g|%d|%s|%s|%s|%t|%x|%x;", i.ID, i.URL, i.Engine, i.Pool, i.Models, i.Weight,
|
||||
i.Priority, i.UpstreamModel, i.Peer, i.PeerCluster, i.PeerTags, sha256.Sum256([]byte(i.APIKey)),
|
||||
sha256.Sum256([]byte(i.Workflow)))
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
@@ -64,6 +64,7 @@ var Reference = []Block{
|
||||
{Key: "tools", Type: "bool", Doc: "", Nested: ""},
|
||||
{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: ""},
|
||||
}},
|
||||
{Package: "config", Name: "Control", Title: "Control", Doc: "", Fields: []Field{
|
||||
{Key: "interval", Type: "time.Duration", Doc: "Interval is how often desired and actual state are compared.", Nested: ""},
|
||||
@@ -84,6 +85,7 @@ var Reference = []Block{
|
||||
{Key: "api_key", Type: "string", Doc: "APIKey is sent to the engine as a bearer token, for an engine started\nwith one.", Nested: ""},
|
||||
{Key: "models", Type: "[]string", Doc: "", Nested: ""},
|
||||
{Key: "engine", Type: "string", Doc: "Engine names the inference engine, which decides what its metrics are\ncalled. Empty means vLLM's names.", Nested: ""},
|
||||
{Key: "workflow", Type: "string", Doc: "Workflow is the ComfyUI workflow for engine: comfyui, a built-in's\nname or its JSON; flux-schnell unless set.", Nested: ""},
|
||||
{Key: "pool", Type: "string", Doc: "", Nested: ""},
|
||||
{Key: "weight", Type: "float64", Doc: "", Nested: ""},
|
||||
{Key: "priority", Type: "int", Doc: "", Nested: ""},
|
||||
@@ -104,6 +106,8 @@ var Reference = []Block{
|
||||
{Key: "model", Type: "string", Doc: "Model is what the provider calls it; requests are rewritten to it.", Nested: ""},
|
||||
{Key: "api_key_file", Type: "string", Doc: "APIKeyFile holds the provider's key, on the control plane. A file\nrather than the key itself, so the key is in neither the catalog, the\nstore nor the audit log; empty for a provider that asks for none.", Nested: ""},
|
||||
{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: ""},
|
||||
}},
|
||||
{Package: "config", Name: "Facility", Title: "Facility", Doc: "", Fields: []Field{
|
||||
{Key: "budgets", Type: "[]FacilityBudget", Doc: "Budgets are the domains with a power cap.", Nested: ""},
|
||||
@@ -181,6 +185,7 @@ var Reference = []Block{
|
||||
{Key: "context_length", Type: "int", Doc: "ContextLength is how many tokens one request may use, prompt and\ncompletion together, and MaxOutput bounds the completion alone. Zero\nmeans unknown, which is what every model was until now.\n\nThe gateway knew neither, so an over-long prompt was not refused: it was\nadmitted, given a replica, and failed at the engine, and a client sizing\nits prompts from /v1/models had nothing to read. Worse, the number is not\nthe one an operator writes in the engine's arguments — llama.cpp divides\n--ctx-size across its parallel slots, so a pool of 8192 across four slots\nis 2048 a request, and nothing said so.", Nested: ""},
|
||||
{Key: "max_output", Type: "int", Doc: "", Nested: ""},
|
||||
{Key: "capabilities", Type: "Capabilities", Doc: "Capabilities is what this model can do beyond plain completion.", Nested: "catalog.Capabilities"},
|
||||
{Key: "query_prefix", Type: "string", Doc: "QueryPrefix is put before a question an embedding model embeds for a\nknowledge-base search; documents are embedded without it. Models\ntrained with an instruction on the query side want it, Qwen3-Embedding\n\"Instruct: Given a web search query, retrieve relevant passages that\nanswer the query\\nQuery: \" for one.", Nested: ""},
|
||||
{Key: "role", Type: "string", Doc: "Role \"guard\" is a model that judges the gateway's traffic for the\nguardrails rather than serving clients: hidden from /v1/models and the\nchat page, refused to any key that names it, critical by default, and\nalways running. Guard says how its answers are read.", Nested: ""},
|
||||
{Key: "guard", Type: "*GuardModel", Doc: "", Nested: "catalog.GuardModel"},
|
||||
{Key: "origin", Type: "*Origin", Doc: "Origin is the registry repository the weights came from, at which\ncommit, and each file's checksum there. Written when a model is served\nfrom the registry; for one that was not, worked out from the mirror's\nrecord of where it fetched the file. It is what lets the control plane\nsay that the repository has published different weights since, and\nmove the model to them.", Nested: "catalog.Origin"},
|
||||
|
||||
@@ -42,7 +42,7 @@ func (r *Router) budgetFallback(p auth.Principal, req *parsedRequest) (accountin
|
||||
// answer than saying the prompt is too large, and it would be done
|
||||
// silently, to a request that was already being quietly moved.
|
||||
canonical, cand := r.resolve(name)
|
||||
if why := checkContextFits(cand, req.PromptTokens, req.MaxTokens); why != "" {
|
||||
if why := checkContextFits(cand, req.contextTokens(), req.MaxTokens); why != "" {
|
||||
continue
|
||||
}
|
||||
res, err := r.acct.Check(p, canonical, req.estimatedTokens())
|
||||
|
||||
@@ -48,12 +48,23 @@ type Info struct {
|
||||
Peer string
|
||||
PeerCluster string
|
||||
PeerTags bool
|
||||
// Workflow is the ComfyUI workflow sent for an endpoint whose engine is
|
||||
// EngineComfyUI: a built-in's name or the workflow's JSON.
|
||||
Workflow string
|
||||
}
|
||||
|
||||
// EngineExternal marks an endpoint that is a provider outside the cluster.
|
||||
// It publishes no metrics, so it is never scraped and never reads as stale.
|
||||
const EngineExternal = "external"
|
||||
|
||||
// EngineComfyUI marks a ComfyUI server. The router speaks ComfyUI's API to
|
||||
// it, not OpenAI's, and it publishes no metrics either.
|
||||
const EngineComfyUI = "comfyui"
|
||||
|
||||
// Scraped reports whether an engine publishes metrics the router reads. One
|
||||
// that does not is never read as stale for having none.
|
||||
func Scraped(engine string) bool { return engine != EngineExternal && engine != EngineComfyUI }
|
||||
|
||||
// Serves reports whether the replica serves the given model.
|
||||
func (i Info) Serves(model string) bool {
|
||||
return slices.Contains(i.Models, model)
|
||||
@@ -257,7 +268,7 @@ func (r *Registry) Snapshots(model string) []Snapshot {
|
||||
continue
|
||||
}
|
||||
var age time.Duration
|
||||
stale := l.info.Engine != EngineExternal
|
||||
stale := Scraped(l.info.Engine)
|
||||
if !l.scrapedAt.IsZero() {
|
||||
age = now.Sub(l.scrapedAt)
|
||||
stale = age > r.cfg.StaleAfter
|
||||
|
||||
@@ -102,6 +102,8 @@ func (r *Router) Handler() http.Handler {
|
||||
mux.HandleFunc("POST /v1/score", r.handleInference("/v1/score"))
|
||||
mux.HandleFunc("POST /v1/audio/transcriptions", r.handleInference("/v1/audio/transcriptions"))
|
||||
mux.HandleFunc("POST /v1/audio/translations", r.handleInference("/v1/audio/translations"))
|
||||
// Pictures, made by ComfyUI from a workflow the catalog names.
|
||||
mux.HandleFunc("POST "+imagesPath, r.handleInference(imagesPath))
|
||||
// Answered by the guardrails' judge model.
|
||||
mux.HandleFunc("POST /v1/moderations", r.handleModerations)
|
||||
// Codex, and whatever else speaks only the Responses API, translated to
|
||||
@@ -391,7 +393,11 @@ type parsedRequest struct {
|
||||
ProfileOverride string
|
||||
PromptPrefix string
|
||||
PromptTokens int
|
||||
MaxTokens int
|
||||
// ContextTokens is what has to fit in one context: the prompt, except
|
||||
// for embeddings, where each input is embedded on its own and the
|
||||
// longest one is what a context must hold.
|
||||
ContextTokens int
|
||||
MaxTokens int
|
||||
|
||||
// ThreadID is the conversation this request continues, empty for the
|
||||
// stateless path every request took before threads existed.
|
||||
@@ -473,11 +479,19 @@ func (r *Router) handleInference(path string) http.HandlerFunc {
|
||||
return
|
||||
}
|
||||
|
||||
// A picture asked of a model that writes text, or the other way
|
||||
// round, is refused before it waits for either.
|
||||
if why := r.imageRouteRefusal(path, parsed.Model, parsed.UpstreamBody); why != "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", why)
|
||||
r.m.observeRequest(r.metricModel(parsed.Model), "", http.StatusBadRequest, time.Since(start), nil, 0)
|
||||
return
|
||||
}
|
||||
|
||||
// Before authentication work and before admission: a request that
|
||||
// cannot fit should not consume a quota reservation or a queue slot on
|
||||
// its way to being refused.
|
||||
if _, entry := r.resolve(parsed.Model); entry.ContextLength > 0 {
|
||||
if why := checkContextFits(entry, parsed.PromptTokens, parsed.MaxTokens); why != "" {
|
||||
if why := checkContextFits(entry, parsed.contextTokens(), parsed.MaxTokens); why != "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", why)
|
||||
r.m.observeRequest(r.metricModel(parsed.Model), "", http.StatusBadRequest, time.Since(start), nil, 0)
|
||||
return
|
||||
@@ -825,6 +839,12 @@ func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) {
|
||||
// plan §14). Measured over the whole prompt, not the routing prefix —
|
||||
// see promptBytes.
|
||||
p.PromptTokens = promptBytes(&b) / 4
|
||||
p.ContextTokens = p.PromptTokens
|
||||
if path == "/v1/embeddings" {
|
||||
if longest, ok := longestInput(b.Input); ok {
|
||||
p.ContextTokens = longest / 4
|
||||
}
|
||||
}
|
||||
|
||||
// Tags are the gateway's: an engine has no use for them, and its logs
|
||||
// are no place for what a client wrote there.
|
||||
@@ -897,6 +917,29 @@ func promptPrefix(b *inferenceBody) string {
|
||||
// bounded head of the conversation; a token estimate wants all of it, and using
|
||||
// one for the other meant TPM limits undercounted exactly the requests that
|
||||
// consume the most, and nothing could tell whether a prompt fit in a context.
|
||||
// longestInput is the length of an embeddings request's longest input, when
|
||||
// input is a list of strings: a batch of texts, each its own context.
|
||||
func longestInput(raw json.RawMessage) (int, bool) {
|
||||
var list []string
|
||||
if len(raw) == 0 || raw[0] != '[' || json.Unmarshal(raw, &list) != nil || len(list) == 0 {
|
||||
return 0, false
|
||||
}
|
||||
n := 0
|
||||
for _, s := range list {
|
||||
n = max(n, len(s))
|
||||
}
|
||||
return n, true
|
||||
}
|
||||
|
||||
// contextTokens is what has to fit in one context; requests built inside
|
||||
// the router set only PromptTokens.
|
||||
func (p *parsedRequest) contextTokens() int {
|
||||
if p.ContextTokens > 0 {
|
||||
return p.ContextTokens
|
||||
}
|
||||
return p.PromptTokens
|
||||
}
|
||||
|
||||
func promptBytes(b *inferenceBody) int {
|
||||
n := 0
|
||||
add := func(raw json.RawMessage) {
|
||||
@@ -1088,7 +1131,7 @@ func (r *Router) handleModels(w http.ResponseWriter, req *http.Request) {
|
||||
}
|
||||
entry["capabilities"] = map[string]bool{
|
||||
"tools": m.Capabilities.Tools, "vision": m.Capabilities.Vision,
|
||||
"embeddings": m.Capabilities.Embeddings,
|
||||
"embeddings": m.Capabilities.Embeddings, "images": m.MakesImages(),
|
||||
}
|
||||
data = append(data, entry)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,200 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"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"
|
||||
)
|
||||
|
||||
// /v1/images/generations, OpenAI's shape, made by ComfyUI: the request goes
|
||||
// through the same checks as any other — key, allowlist, limits, admission,
|
||||
// guardrails on its prompt — and only at the replica is it turned into a
|
||||
// workflow, queued and waited on, and its pictures sent back as base64.
|
||||
|
||||
const imagesPath = "/v1/images/generations"
|
||||
|
||||
// imageRequest is what a request for pictures asks. Beyond OpenAI's fields,
|
||||
// those a diffusion model has and DALL·E does not: seed, steps,
|
||||
// negative_prompt and guidance.
|
||||
type imageRequest struct {
|
||||
Prompt string `json:"prompt"`
|
||||
N int `json:"n"`
|
||||
Size string `json:"size"`
|
||||
ResponseFormat string `json:"response_format"`
|
||||
Seed *int64 `json:"seed"`
|
||||
Steps int `json:"steps"`
|
||||
NegativePrompt string `json:"negative_prompt"`
|
||||
Guidance *float64 `json:"guidance"`
|
||||
|
||||
width, height int
|
||||
}
|
||||
|
||||
// The bounds of a request. A FLUX latent is 1/16 of the picture, so a side
|
||||
// that is not a multiple of 16 is not one it can make; the rest keep one
|
||||
// request from holding a card for the length of a film.
|
||||
const (
|
||||
maxImages = 4
|
||||
minSide = 256
|
||||
maxSide = 2048
|
||||
maxSteps = 100
|
||||
defaultSide = 1024
|
||||
)
|
||||
|
||||
func parseImageRequest(body []byte) (imageRequest, error) {
|
||||
var in imageRequest
|
||||
if err := json.Unmarshal(body, &in); err != nil {
|
||||
return in, err
|
||||
}
|
||||
if strings.TrimSpace(in.Prompt) == "" {
|
||||
return in, errors.New("prompt is required")
|
||||
}
|
||||
switch {
|
||||
case in.N == 0:
|
||||
in.N = 1
|
||||
case in.N < 0 || in.N > maxImages:
|
||||
return in, fmt.Errorf("n is between 1 and %d", maxImages)
|
||||
}
|
||||
switch in.ResponseFormat {
|
||||
case "", "b64_json":
|
||||
case "url":
|
||||
// A URL would have to be served by somebody and kept for some time,
|
||||
// and a picture made from a prompt is the prompt's owner's: it is
|
||||
// not kept here.
|
||||
return in, errors.New(`response_format "url" is not offered: the picture comes back as b64_json and is kept nowhere`)
|
||||
default:
|
||||
return in, fmt.Errorf("response_format %q is not b64_json", in.ResponseFormat)
|
||||
}
|
||||
in.width, in.height = defaultSide, defaultSide
|
||||
if s := strings.TrimSpace(in.Size); s != "" && s != "auto" {
|
||||
w, h, ok := strings.Cut(strings.ToLower(s), "x")
|
||||
wi, err1 := strconv.Atoi(w)
|
||||
hi, err2 := strconv.Atoi(h)
|
||||
if !ok || err1 != nil || err2 != nil {
|
||||
return in, fmt.Errorf("size %q is not WIDTHxHEIGHT", s)
|
||||
}
|
||||
in.width, in.height = wi, hi
|
||||
}
|
||||
for _, side := range []int{in.width, in.height} {
|
||||
if side < minSide || side > maxSide || side%16 != 0 {
|
||||
return in, fmt.Errorf("size %dx%d: each side is a multiple of 16 from %d to %d", in.width, in.height, minSide, maxSide)
|
||||
}
|
||||
}
|
||||
if in.Steps < 0 || in.Steps > maxSteps {
|
||||
return in, fmt.Errorf("steps is between 1 and %d", maxSteps)
|
||||
}
|
||||
if in.Guidance != nil && (*in.Guidance <= 0 || *in.Guidance > 30) {
|
||||
return in, errors.New("guidance is above 0 and at most 30")
|
||||
}
|
||||
return in, nil
|
||||
}
|
||||
|
||||
// imageRouteRefusal says why a request cannot be answered at the path it was
|
||||
// sent to, before it takes a place in a queue: a picture asked of a model
|
||||
// that writes text, or text of one that makes pictures. Empty when it can,
|
||||
// and for a model the catalog does not know, which a hand-configured
|
||||
// endpoint may serve as it likes.
|
||||
func (r *Router) imageRouteRefusal(path, model string, body []byte) string {
|
||||
m, known := r.modelMap()[model]
|
||||
if path == imagesPath {
|
||||
if known && !m.MakesImages() {
|
||||
return fmt.Sprintf("%s does not make pictures", m.Name)
|
||||
}
|
||||
if _, err := parseImageRequest(body); err != nil {
|
||||
return err.Error()
|
||||
}
|
||||
return ""
|
||||
}
|
||||
if known && m.External != nil && m.External.API == catalog.APIComfyUI {
|
||||
return fmt.Sprintf("%s makes pictures: ask for them at %s", m.Name, imagesPath)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// generateImages answers a request at a ComfyUI replica. It returns what the
|
||||
// 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 != imagesPath {
|
||||
return proxy.Outcome{}, fmt.Errorf("%s is a ComfyUI server: it answers %s only", ep.ID, imagesPath)
|
||||
}
|
||||
in, err := parseImageRequest(body)
|
||||
if err != nil {
|
||||
return proxy.Outcome{}, err
|
||||
}
|
||||
var upstream struct {
|
||||
Model string `json:"model"`
|
||||
}
|
||||
_ = json.Unmarshal(body, &upstream)
|
||||
tpl, err := comfy.Load(ep.Workflow)
|
||||
if err != nil {
|
||||
return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err)
|
||||
}
|
||||
seed := comfy.RandomSeed()
|
||||
if in.Seed != nil {
|
||||
seed = *in.Seed
|
||||
}
|
||||
p := comfy.Params{
|
||||
Checkpoint: upstream.Model, Prompt: in.Prompt, NegativePrompt: in.NegativePrompt,
|
||||
Width: in.width, Height: in.height, N: in.N, Seed: seed, Steps: in.Steps,
|
||||
}
|
||||
if in.Guidance != nil {
|
||||
p.Guidance = *in.Guidance
|
||||
}
|
||||
workflow, err := tpl.Fill(p)
|
||||
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()
|
||||
images, err := r.comfy.Generate(ctx, ep.URL, ep.APIKey, workflow)
|
||||
if err != nil {
|
||||
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 picture within %s", ep.ID, r.cfg.Timeouts.Total)
|
||||
}
|
||||
return proxy.Outcome{}, err
|
||||
}
|
||||
|
||||
data := make([]map[string]any, 0, len(images))
|
||||
for _, im := range images {
|
||||
data = append(data, map[string]any{"b64_json": base64.StdEncoding.EncodeToString(im.Data)})
|
||||
}
|
||||
out, err := json.Marshal(map[string]any{"created": time.Now().Unix(), "data": data})
|
||||
if err != nil {
|
||||
return proxy.Outcome{}, err
|
||||
}
|
||||
// The seed that made them, so a picture liked can be made again.
|
||||
w.Header().Set("x-llmgw-seed", strconv.FormatInt(seed, 10))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(out)
|
||||
took := time.Since(start)
|
||||
return proxy.Outcome{Wrote: true, Status: http.StatusOK, Bytes: int64(len(out)), TTFT: took, Decode: took}, nil
|
||||
}
|
||||
|
||||
// ctxCauseIsClient reports whether a context ended because the client went
|
||||
// away, rather than because a deadline set here passed.
|
||||
func ctxCauseIsClient(ctx context.Context) bool {
|
||||
return !errors.Is(context.Cause(ctx), context.DeadlineExceeded)
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/comfy"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
)
|
||||
|
||||
// imageRouter has a ComfyUI model, flux, beside a text model, and returns
|
||||
// what the ComfyUI server was sent.
|
||||
func imageRouter(t *testing.T) (*Router, *map[string]any, *sync.Mutex) {
|
||||
t.Helper()
|
||||
var mu sync.Mutex
|
||||
var queued map[string]any
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/metrics":
|
||||
t.Error("ComfyUI was scraped for metrics")
|
||||
case "/prompt":
|
||||
var b struct {
|
||||
Prompt map[string]any `json:"prompt"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&b)
|
||||
mu.Lock()
|
||||
queued = b.Prompt
|
||||
mu.Unlock()
|
||||
fmt.Fprint(w, `{"prompt_id":"p1"}`)
|
||||
case "/history/p1":
|
||||
fmt.Fprint(w, `{"p1":{"status":{"status_str":"success","completed":true},"outputs":{"7":{"images":[{"filename":"a.png","type":"temp"}]}}}}`)
|
||||
case "/view":
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
fmt.Fprint(w, "png-bytes")
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
cfg := baseConfig()
|
||||
cfg.Models = []catalog.Model{catModel("chat", nil), {Name: "flux", Source: "external",
|
||||
External: &catalog.External{URL: srv.URL, Model: "flux1-schnell-fp8.safetensors", API: catalog.APIComfyUI}}}
|
||||
rt := newTestRouter(t, cfg)
|
||||
rt.comfy = &comfy.Client{Poll: time.Millisecond}
|
||||
rt.SetEndpoints(t.Context(), []endpoint.Info{
|
||||
{ID: "external:flux", URL: srv.URL, Models: []string{"flux"}, Engine: endpoint.EngineComfyUI,
|
||||
Pool: "external", Weight: 1, UpstreamModel: "flux1-schnell-fp8.safetensors"},
|
||||
})
|
||||
return rt, &queued, &mu
|
||||
}
|
||||
|
||||
func postImages(t *testing.T, h http.Handler, path string, body any) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
b, _ := json.Marshal(body)
|
||||
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(b))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestAPictureIsMadeByComfyUIAndComesBackAsBase64(t *testing.T) {
|
||||
rt, queued, mu := imageRouter(t)
|
||||
rec := postImages(t, rt.Handler(), imagesPath, map[string]any{
|
||||
"model": "flux", "prompt": "a lighthouse at dusk", "size": "768x512", "seed": 7})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var out struct {
|
||||
Data []struct {
|
||||
B64 string `json:"b64_json"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || len(out.Data) != 1 {
|
||||
t.Fatalf("body %s: %v", rec.Body.String(), err)
|
||||
}
|
||||
if b, _ := base64.StdEncoding.DecodeString(out.Data[0].B64); string(b) != "png-bytes" {
|
||||
t.Errorf("picture = %q", b)
|
||||
}
|
||||
if rec.Header().Get("x-llmgw-seed") != "7" {
|
||||
t.Errorf("seed header = %q", rec.Header().Get("x-llmgw-seed"))
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
wf, _ := json.Marshal(*queued)
|
||||
for _, want := range []string{`"ckpt_name":"flux1-schnell-fp8.safetensors"`, `"text":"a lighthouse at dusk"`,
|
||||
`"width":768`, `"height":512`, `"seed":7`} {
|
||||
if !strings.Contains(string(wf), want) {
|
||||
t.Errorf("%s not in the workflow sent: %s", want, wf)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPicturesAndTextAreAskedForAtTheirOwnPaths(t *testing.T) {
|
||||
rt, _, _ := imageRouter(t)
|
||||
h := rt.Handler()
|
||||
for name, c := range map[string]struct {
|
||||
path string
|
||||
body map[string]any
|
||||
want string
|
||||
}{
|
||||
"text of flux": {"/v1/chat/completions", map[string]any{"model": "flux", "messages": []any{}}, "ask for them at " + imagesPath},
|
||||
"picture of chat": {imagesPath, map[string]any{"model": "chat", "prompt": "x"}, "does not make pictures"},
|
||||
"no prompt": {imagesPath, map[string]any{"model": "flux"}, "prompt is required"},
|
||||
"odd size": {imagesPath, map[string]any{"model": "flux", "prompt": "x", "size": "1000x1000"}, "multiple of 16"},
|
||||
"url wanted": {imagesPath, map[string]any{"model": "flux", "prompt": "x", "response_format": "url"}, "b64_json"},
|
||||
"too many at once": {imagesPath, map[string]any{"model": "flux", "prompt": "x", "n": 9}, "n is between"},
|
||||
} {
|
||||
rec := postImages(t, h, c.path, c.body)
|
||||
if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), c.want) {
|
||||
t.Errorf("%s: %d %s", name, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAComfyUIModelSaysItMakesPictures(t *testing.T) {
|
||||
rt, _, _ := imageRouter(t)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
rt.Handler().ServeHTTP(rec, req)
|
||||
if !strings.Contains(rec.Body.String(), `"images":true`) {
|
||||
t.Errorf("/v1/models: %s", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAComfyUIModelIsCheckedInTheCatalog(t *testing.T) {
|
||||
ok := catalog.Model{Name: "f", Source: "external", External: &catalog.External{
|
||||
URL: "http://10.0.0.147:8188", Model: "flux1-dev.safetensors", API: catalog.APIComfyUI, Workflow: "flux-dev"}}
|
||||
if _, err := catalog.New([]catalog.Model{ok}); err != nil {
|
||||
t.Fatalf("refused: %v", err)
|
||||
}
|
||||
for name, e := range map[string]catalog.External{
|
||||
"unknown api": {URL: "http://a:1", Model: "m", API: "a1111"},
|
||||
"unknown workflow": {URL: "http://a:1", Model: "m", API: catalog.APIComfyUI, Workflow: "sdxl"},
|
||||
"workflow, openai": {URL: "http://a:1", Model: "m", Workflow: "flux-dev"},
|
||||
"comfyui from peer": {URL: "http://a:1", Model: "m", API: catalog.APIComfyUI, Peer: "other"},
|
||||
} {
|
||||
e := e
|
||||
if _, err := catalog.New([]catalog.Model{{Name: "f", Source: "external", External: &e}}); err == nil {
|
||||
t.Errorf("%s: accepted", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -276,6 +276,10 @@ func splitList(v string) []string {
|
||||
// ordinary path, so the caller's limits, budget and guardrails apply to it
|
||||
// and it is counted to them.
|
||||
func (r *Router) embedFor(ctx context.Context, p auth.Principal, model, text string) ([]float32, error) {
|
||||
// A question, not a document: what the model wants before a query.
|
||||
if _, m := r.resolve(model); m.QueryPrefix != "" {
|
||||
text = m.QueryPrefix + text
|
||||
}
|
||||
body, _ := json.Marshal(map[string]any{"model": model, "input": []string{text}})
|
||||
hreq, err := http.NewRequestWithContext(auth.WithPrincipal(ctx, p), http.MethodPost, "/v1/embeddings", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
|
||||
@@ -9,7 +9,9 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
@@ -253,3 +255,25 @@ func TestAVLLMReplicaGetsASaltOnlyWithPassages(t *testing.T) {
|
||||
t.Errorf("a salt without passages:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAQuestionIsEmbeddedWithTheModelsQueryPrefix(t *testing.T) {
|
||||
// Qwen3-Embedding and its like are trained with an instruction before
|
||||
// the query, never before the documents: the documents come through
|
||||
// /v1/embeddings as they are; the question gets the prefix here.
|
||||
e := newFakeEngine(t, "emb")
|
||||
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"emb"}})
|
||||
m := catModel("emb", nil)
|
||||
m.Capabilities.Embeddings = true
|
||||
m.QueryPrefix = "Instruct: retrieve passages\nQuery: "
|
||||
cfg.Models = []catalog.Model{m}
|
||||
rt := newTestRouter(t, cfg)
|
||||
// The test engine answers with the input's length in bytes and words.
|
||||
v, err := rt.embedFor(t.Context(), auth.Principal{}, "emb", "kaç gün izin var")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := "Instruct: retrieve passages\nQuery: kaç gün izin var"
|
||||
if int(v[0]) != len(want) || int(v[1]) != len(strings.Fields(want)) {
|
||||
t.Errorf("embedded %v, not %q", v, want)
|
||||
}
|
||||
}
|
||||
@@ -29,6 +29,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/accounting"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"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/control/reconciler"
|
||||
@@ -53,7 +54,9 @@ type Router struct {
|
||||
reg *endpoint.Registry
|
||||
poller *scrape.Poller
|
||||
px *proxy.Proxy
|
||||
m *metrics
|
||||
// comfy asks ComfyUI replicas for pictures.
|
||||
comfy *comfy.Client
|
||||
m *metrics
|
||||
// adm is nil when admission control is disabled.
|
||||
adm *admission.Controller
|
||||
|
||||
@@ -311,6 +314,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
|
||||
kv: kvindex.New(),
|
||||
kvClient: newKVClient(),
|
||||
m: metrics,
|
||||
comfy: &comfy.Client{HTTP: &http.Client{Transport: &http.Transport{DialContext: (&net.Dialer{Timeout: cfg.Router.Timeouts.Dial}).DialContext}}},
|
||||
px: proxy.New(proxy.Config{
|
||||
DialTimeout: cfg.Router.Timeouts.Dial,
|
||||
ResponseHeaderTimeout: cfg.Router.Timeouts.ResponseHeader,
|
||||
@@ -487,6 +491,7 @@ func toInfos(eps []config.Endpoint) []endpoint.Info {
|
||||
Pool: e.Pool,
|
||||
Weight: e.Weight,
|
||||
Priority: e.Priority,
|
||||
Workflow: e.Workflow,
|
||||
})
|
||||
}
|
||||
return out
|
||||
@@ -1162,7 +1167,13 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa
|
||||
}
|
||||
// A filter per attempt, for the model this attempt is served by.
|
||||
filter := r.outputFilter(uctx, req, model, &preq)
|
||||
out, perr := r.px.Do(uctx, w, preq)
|
||||
var out proxy.Outcome
|
||||
var perr error
|
||||
if dec.Endpoint.Engine == endpoint.EngineComfyUI {
|
||||
out, perr = r.generateImages(uctx, w, req, dec.Endpoint.Info, body)
|
||||
} else {
|
||||
out, perr = r.px.Do(uctx, w, preq)
|
||||
}
|
||||
releaseAdapter()
|
||||
uspan.SetAttributes(attribute.Int("http.response.status_code", out.Status))
|
||||
endSpan(uspan, perr)
|
||||
|
||||
@@ -783,3 +783,28 @@ func TestAStoppedServiceAnswers503AndStillListsModels(t *testing.T) {
|
||||
t.Errorf("started again, still refused: %s", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestABatchOfEmbeddingsIsCheckedTextByText(t *testing.T) {
|
||||
// llmgw-kb sends many passages in one request; each is embedded on its
|
||||
// own, so their sum over the context is no reason to refuse them. One
|
||||
// passage over the context still is.
|
||||
e := newFakeEngine(t, "emb")
|
||||
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"emb"}})
|
||||
m := catModel("emb", nil)
|
||||
m.ContextLength = 2048
|
||||
cfg.Models = []catalog.Model{m}
|
||||
rt := newTestRouter(t, cfg)
|
||||
send := func(input []string) *httptest.ResponseRecorder {
|
||||
b, _ := json.Marshal(map[string]any{"model": "emb", "input": input})
|
||||
rec := httptest.NewRecorder()
|
||||
rt.Handler().ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/embeddings", bytes.NewReader(b)))
|
||||
return rec
|
||||
}
|
||||
passage := strings.Repeat("x", 6000) // about 1500 tokens
|
||||
if rec := send([]string{passage, passage, passage, passage}); rec.Code == http.StatusBadRequest {
|
||||
t.Errorf("four passages that each fit were refused: %s", rec.Body.String())
|
||||
}
|
||||
if rec := send([]string{"short", strings.Repeat("x", 40_000)}); rec.Code != http.StatusBadRequest {
|
||||
t.Errorf("a passage over the context was let through: %d", rec.Code)
|
||||
}
|
||||
}
|
||||
@@ -96,8 +96,8 @@ func (p *Poller) Sync(ctx context.Context) {
|
||||
|
||||
want := make(map[string]endpoint.Info, len(infos))
|
||||
for _, in := range infos {
|
||||
if in.Engine == endpoint.EngineExternal {
|
||||
continue // a provider elsewhere publishes no /metrics
|
||||
if !endpoint.Scraped(in.Engine) {
|
||||
continue // a provider elsewhere, or ComfyUI, publishes no /metrics
|
||||
}
|
||||
want[in.ID] = in
|
||||
}
|
||||
|
||||
@@ -327,6 +327,10 @@ func (p *AccountPages) logout(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
clearSessionCookie(w, r)
|
||||
// What the pages kept in this browser goes with the session: the chat's
|
||||
// conversation, htmx's copies of the pages visited. The chat keys its
|
||||
// own by account as well, for a browser that ignores this header.
|
||||
w.Header().Set("Clear-Site-Data", `"storage"`)
|
||||
http.Redirect(w, r, "/login", http.StatusSeeOther)
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package web
|
||||
|
||||
import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -113,6 +115,43 @@ func TestTheChatPageOffersWhatItWasGiven(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheKeptConversationBelongsToWhoeverKeptIt(t *testing.T) {
|
||||
// One browser, two people: an administrator signs out, a new account signs
|
||||
// in, and the chat must not open on the administrator's conversation.
|
||||
h := Chat(store.User{ID: "u-42", Email: "a@example.com"}, []ChatModel{{Name: "m"}}, nil, "", "v1", ChatKnowledge{})
|
||||
w := httptest.NewRecorder()
|
||||
h(w, httptest.NewRequest(http.MethodGet, "http://x/chat", nil))
|
||||
page := w.Body.String()
|
||||
if !strings.Contains(page, `<meta name="chat-owner" content="u-42">`) {
|
||||
t.Fatal("the page does not say whose conversation it keeps")
|
||||
}
|
||||
if !strings.Contains(page, `const KEY = PREFIX + "." + document.querySelector('meta[name="chat-owner"]').content;`) {
|
||||
t.Error("the conversation is not kept under the account")
|
||||
}
|
||||
// The unowned key from before, and anyone else's, are dropped before the
|
||||
// kept conversation is read.
|
||||
drop, read := strings.Index(page, "localStorage.removeItem(k)"), strings.Index(page, "const kept = load_()")
|
||||
if drop < 0 || read < 0 || drop > read {
|
||||
t.Error("someone else's conversation is not dropped before this one is put back")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSigningOutClearsWhatTheBrowserKept(t *testing.T) {
|
||||
p := NewAccountPages(fakeAccounts{user: store.User{ID: "u1", Email: "a@example.com"}}, nil, nil, "v0",
|
||||
slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
r := httptest.NewRequest(http.MethodPost, "http://x/logout", strings.NewReader("csrf=tok"))
|
||||
r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
r.AddCookie(&http.Cookie{Name: SessionCookie, Value: "good"})
|
||||
w := httptest.NewRecorder()
|
||||
p.logout(w, r)
|
||||
if w.Code != http.StatusSeeOther {
|
||||
t.Fatalf("status %d", w.Code)
|
||||
}
|
||||
if got := w.Header().Get("Clear-Site-Data"); got != `"storage"` {
|
||||
t.Errorf("Clear-Site-Data = %q, want \"storage\"", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheSimulatedModelSaysSo(t *testing.T) {
|
||||
// big-model returns Greek letters at a fixed rate. Offered next to a real
|
||||
// model with nothing to tell them apart, it reads as a broken model rather
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
<script defer src="/_static/katex/auto-render.min.js"></script>
|
||||
<script defer src="/_static/pages.js"></script>
|
||||
<meta name="csrf" content="{{.CSRF}}">
|
||||
<meta name="chat-owner" content="{{.User.ID}}">
|
||||
</head><body><div class="wrap">
|
||||
<header>
|
||||
<h1>{{t "chat.title"}}</h1>
|
||||
@@ -131,10 +132,22 @@
|
||||
// In this browser's storage, and only there: the gateway keeps no
|
||||
// conversation. Every read and write is guarded — a private window, or a
|
||||
// browser that refuses storage, simply starts empty each time.
|
||||
//
|
||||
// Kept under the account's id: a browser is often shared, and the next
|
||||
// person to sign in on it must not be handed the last one's conversation.
|
||||
// Anyone else's left here is dropped on the way in, as is the one from
|
||||
// before the key carried an owner.
|
||||
const sys = document.getElementById("system");
|
||||
const temp = document.getElementById("temperature");
|
||||
const tempOut = document.getElementById("temp-out");
|
||||
const KEY = "llmgw.chat.v1";
|
||||
const PREFIX = "llmgw.chat.v1";
|
||||
const KEY = PREFIX + "." + document.querySelector('meta[name="chat-owner"]').content;
|
||||
try {
|
||||
for (let i = localStorage.length - 1; i >= 0; i--) {
|
||||
const k = localStorage.key(i);
|
||||
if (k && (k === PREFIX || k.startsWith(PREFIX + ".")) && k !== KEY) localStorage.removeItem(k);
|
||||
}
|
||||
} catch (_) { /* no storage: nothing to drop */ }
|
||||
const load_ = () => { try { return JSON.parse(localStorage.getItem(KEY) || "null"); } catch (_) { return null; } };
|
||||
const save = () => {
|
||||
try {
|
||||
|
||||
Reference in new issue
Block a user