From 6c0c5bd9078c6674d77e423715da650b1e0c116f Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Thu, 2 Jul 2026 15:53:21 +0100 Subject: [PATCH 01/10] feat(router): make KNN a first-class classifier with a persisted, curated corpus MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `classifier: knn` — similarity-weighted voting over labelled example prompts. Unlike score/colbert it needs no classifier model: label knowledge lives in a corpus seeded and curated through the admin API, so routing decisions are deterministic, auditable, and grounded in graded experience rather than a model's opinion. Epistemic gate: corpus entries below knn.similarity_threshold cannot vote; when none clears it the classifier activates no labels and the router uses the fallback — a prompt unlike all labelled experience is treated as undecidable, not guessed. Decisions record nearest_similarity (also on fallback rows) so admins can see how far the nearest labelled experience was; the Routing tab explains out-of-corpus fallbacks and shows per-label corpus counts. Persistence: one JSONL file per router under /router-corpus (text, labels, vector, embedder fingerprint). The file is the source of truth; the local-store index is rebuilt from it at classifier build time and stays a pure in-memory index. Entries recorded under a different embedding model re-embed on load. Also corrects the docs' false claim that local-store collections persist — the embedding cache never survived restarts (and still doesn't); the corpus does. Corpus input is API-only by design (entries may contain example user content): POST /api/router/{name}/corpus seeds (labels validated against declared policies, embedded server-side, indexed immediately), GET .../corpus/stats inspects — label counts only, entry texts are never returned by any surface — DELETE .../corpus wipes. Admin-gated like the sibling router endpoints, and exposed as MCP tools (seed_router_corpus / get_router_corpus_stats / clear_router_corpus) in both the httpapi and inproc clients with coverage-test route mappings. Plumbing: VectorStore gains SearchK (top-K was hardcoded to 1); local-store gets InsertBatch/Delete as optional fast paths; RouterConfig gains a knn block (embedding_model, k, similarity_threshold, vote_threshold, store_name) with meta-registry fields; the classifier dropdown now offers knn and the previously-missing colbert; embedding_cache is ignored (with a warning) for knn — it IS an embedding-KNN lookup; the stale /api/instructions intelligent-routing entry is rewritten (it described a classifier that no longer exists); swagger regenerated. Tests: KNN vote/gate specs with hand-computed vote shares, corpus manager suite (restart reload without re-embedding, fingerprint re-embed, dedupe, hostile store names), middleware specs (corpus routing, gate fallback, config validation, cache-wrap refusal), corpus endpoint specs pinning the texts-never-returned contract, MCP catalog + route-mapping gates, and a Playwright spec for corpus stats and the out-of-corpus decision detail. Assisted-by: Claude:claude-fable-5 [Claude Code] --- core/application/application.go | 9 + core/application/router_factories.go | 14 + core/backend/stores.go | 93 ++++- core/config/meta/registry.go | 49 ++- core/config/model_config.go | 48 +++ .../endpoints/localai/api_instructions.go | 2 +- core/http/endpoints/localai/router_corpus.go | 189 +++++++++ .../endpoints/localai/router_corpus_test.go | 210 ++++++++++ core/http/endpoints/localai/router_decide.go | 19 +- .../endpoints/mcp/localai_assistant_test.go | 12 + core/http/endpoints/openai/realtime_model.go | 1 + core/http/middleware/route_model.go | 70 +++- core/http/middleware/route_model_test.go | 183 +++++++++ .../http/react-ui/e2e/middleware-page.spec.js | 51 ++- core/http/react-ui/src/pages/Middleware.jsx | 49 ++- core/http/routes/anthropic.go | 1 + core/http/routes/middleware.go | 54 ++- core/http/routes/openai.go | 1 + core/schema/localai.go | 126 ++++-- .../routing/corpus/corpus_suite_test.go | 13 + core/services/routing/corpus/manager.go | 375 ++++++++++++++++++ core/services/routing/corpus/manager_test.go | 240 +++++++++++ core/services/routing/router/decisions.go | 29 +- .../routing/router/embedding_cache_test.go | 26 ++ core/services/routing/router/knn.go | 222 +++++++++++ core/services/routing/router/knn_test.go | 174 ++++++++ core/services/routing/router/resolve.go | 1 + core/services/routing/router/types.go | 14 + docs/content/features/middleware.md | 141 ++++++- pkg/mcp/localaitools/client.go | 12 + pkg/mcp/localaitools/coverage_test.go | 37 +- pkg/mcp/localaitools/dto.go | 44 ++ pkg/mcp/localaitools/fakes_test.go | 15 + pkg/mcp/localaitools/httpapi/client.go | 31 ++ pkg/mcp/localaitools/inproc/client.go | 119 ++++++ pkg/mcp/localaitools/server_test.go | 4 + pkg/mcp/localaitools/tools.go | 35 +- pkg/mcp/localaitools/tools_middleware.go | 57 ++- swagger/docs.go | 272 +++++++++++++ swagger/swagger.json | 272 +++++++++++++ swagger/swagger.yaml | 189 +++++++++ 41 files changed, 3388 insertions(+), 115 deletions(-) create mode 100644 core/http/endpoints/localai/router_corpus.go create mode 100644 core/http/endpoints/localai/router_corpus_test.go create mode 100644 core/services/routing/corpus/corpus_suite_test.go create mode 100644 core/services/routing/corpus/manager.go create mode 100644 core/services/routing/corpus/manager_test.go create mode 100644 core/services/routing/router/knn.go create mode 100644 core/services/routing/router/knn_test.go diff --git a/core/application/application.go b/core/application/application.go index ab7f34a41f54..d7b51c1fe687 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -19,6 +19,7 @@ import ( "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/routing/admission" "github.com/mudler/LocalAI/core/services/routing/billing" + "github.com/mudler/LocalAI/core/services/routing/corpus" "github.com/mudler/LocalAI/core/services/routing/pii" "github.com/mudler/LocalAI/core/services/routing/piidetector" "github.com/mudler/LocalAI/core/services/routing/router" @@ -76,6 +77,8 @@ type Application struct { mitmHostConflicts atomic.Pointer[map[string][]string] routerDecisions router.DecisionStore routerRegistry *router.Registry + routerCorpus *corpus.Manager + routerCorpusOnce sync.Once admissionLimiter *admission.Limiter watchdogMutex sync.Mutex watchdogStop chan bool @@ -549,6 +552,12 @@ func (a *Application) start() error { assistantClient.PIIRedactor = a.piiRedactor assistantClient.PIIEvents = a.piiEvents assistantClient.RouterDecisions = a.routerDecisions + // Router corpus tools — same factories the RouteModel middleware + // uses, so the assistant and the request path agree on store + // namespaces and model resolution. + assistantClient.RouterCorpus = a.RouterCorpus() + assistantClient.RouterEmbedder = a.Embedder + assistantClient.RouterVectorStore = a.VectorStore if err := holder.Initialize(a.applicationConfig.Context, assistantClient, localaitools.Options{}); err != nil { // Why log+continue instead of fail: the assistant is an optional // feature; a failure here must not take down the whole server. diff --git a/core/application/router_factories.go b/core/application/router_factories.go index 879c43a835ee..0c4bb936950a 100644 --- a/core/application/router_factories.go +++ b/core/application/router_factories.go @@ -1,11 +1,14 @@ package application import ( + "cmp" "context" "fmt" + "path/filepath" "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/routing/corpus" ) // adapterConfig resolves a model name to its runtime ModelConfig, or nil when @@ -118,3 +121,14 @@ func (l *lazyEmbedder) Embed(ctx context.Context, text string) ([]float32, error func (a *Application) VectorStore(storeName string) backend.VectorStore { return backend.NewVectorStore(a.modelLoader, a.applicationConfig, storeName) } + +// RouterCorpus returns the process-wide KNN corpus manager. Corpus +// files live under /router-corpus (same DataPath → +// DynamicConfigsDir precedence the agent pool uses for its state). +func (a *Application) RouterCorpus() *corpus.Manager { + a.routerCorpusOnce.Do(func() { + root := cmp.Or(a.applicationConfig.DataPath, a.applicationConfig.DynamicConfigsDir, ".") + a.routerCorpus = corpus.NewManager(filepath.Join(root, "router-corpus")) + }) + return a.routerCorpus +} diff --git a/core/backend/stores.go b/core/backend/stores.go index 480400f428f2..9daa5c4a9b46 100644 --- a/core/backend/stores.go +++ b/core/backend/stores.go @@ -14,13 +14,23 @@ import ( ) // VectorStore is the narrowed KNN store used by the router's embedding -// cache. Search returns the top-1 match (cosine similarity in [-1, 1]) -// and the serialised payload, or ok=false on a clean miss. +// cache and the KNN classifier. Search returns the top-1 match (cosine +// similarity in [-1, 1]) and the serialised payload, or ok=false on a +// clean miss. SearchK returns up to k nearest neighbours ordered by +// descending similarity; an empty slice is a clean miss. type VectorStore interface { Search(ctx context.Context, vec []float32) (similarity float64, payload []byte, ok bool, err error) + SearchK(ctx context.Context, vec []float32, k int) ([]Neighbor, error) Insert(ctx context.Context, vec []float32, payload []byte) error } +// Neighbor is one SearchK result — the stored payload and its cosine +// similarity to the query vector. +type Neighbor struct { + Similarity float64 + Payload []byte +} + // NewVectorStore returns a VectorStore backed by the local-store // gRPC backend, namespaced by storeName so two routers don't collide. func NewVectorStore(loader *model.ModelLoader, appConfig *config.ApplicationConfig, storeName string) VectorStore { @@ -63,6 +73,35 @@ func (s *localVectorStore) Search(ctx context.Context, vec []float32) (sim float return float64(similarities[0]), values[0], true, nil } +func (s *localVectorStore) SearchK(ctx context.Context, vec []float32, k int) (neighbors []Neighbor, err error) { + start := time.Now() + outcome := "hit" + sim := 0.0 + defer func() { + s.recordTrace(start, "search", len(vec), sim, outcome, err) + }() + be, berr := s.backend(ctx) + if berr != nil { + outcome = "backend_load_error" + return nil, fmt.Errorf("vector store load: %w", berr) + } + _, values, similarities, ferr := store.Find(ctx, be, vec, k) + if ferr != nil { + outcome = "find_error" + return nil, fmt.Errorf("vector store find: %w", ferr) + } + if len(values) == 0 { + outcome = "miss" + return nil, nil + } + neighbors = make([]Neighbor, 0, len(values)) + for i, v := range values { + neighbors = append(neighbors, Neighbor{Similarity: float64(similarities[i]), Payload: v}) + } + sim = neighbors[0].Similarity + return neighbors, nil +} + func (s *localVectorStore) Insert(ctx context.Context, vec []float32, payload []byte) (err error) { start := time.Now() outcome := "ok" @@ -81,6 +120,56 @@ func (s *localVectorStore) Insert(ctx context.Context, vec []float32, payload [] return nil } +// InsertBatch upserts many vectors in one gRPC round-trip. Not part of +// the VectorStore interface — the corpus manager type-asserts for it +// and falls back to per-entry Insert on stores that lack it. +func (s *localVectorStore) InsertBatch(ctx context.Context, vecs [][]float32, payloads [][]byte) (err error) { + start := time.Now() + outcome := "ok" + dim := 0 + if len(vecs) > 0 { + dim = len(vecs[0]) + } + defer func() { + s.recordTrace(start, "insert_batch", dim, 0, outcome, err) + }() + be, berr := s.backend(ctx) + if berr != nil { + outcome = "backend_load_error" + return fmt.Errorf("vector store load: %w", berr) + } + if serr := store.SetCols(ctx, be, vecs, payloads); serr != nil { + outcome = "insert_error" + return serr + } + return nil +} + +// Delete removes vectors by key. Optional capability like InsertBatch; +// used by the corpus manager's Clear so a wiped corpus also leaves the +// live index. +func (s *localVectorStore) Delete(ctx context.Context, vecs [][]float32) (err error) { + start := time.Now() + outcome := "ok" + dim := 0 + if len(vecs) > 0 { + dim = len(vecs[0]) + } + defer func() { + s.recordTrace(start, "delete", dim, 0, outcome, err) + }() + be, berr := s.backend(ctx) + if berr != nil { + outcome = "backend_load_error" + return fmt.Errorf("vector store load: %w", berr) + } + if serr := store.DeleteCols(ctx, be, vecs); serr != nil { + outcome = "delete_error" + return serr + } + return nil +} + // recordTrace surfaces vector-store calls in /api/backend-traces, including // the backend-load-failure path that otherwise vanishes into an xlog.Warn. // modelName uses the store namespace (e.g. "router-cache-smart-router") so diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index 2ad64f4d3f0f..a362d9b676aa 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -847,17 +847,19 @@ func DefaultRegistry() map[string]FieldMetaOverride { "router.classifier": { Section: "router", Label: "Classifier", - Description: "Picks a candidate by scoring every policy label against the prompt. Only \"score\" is shipped today; it asks the classifier_model to rank each label and reads off the softmax. Empty defaults to \"score\".", + Description: "How the router picks labels for a prompt. \"score\" asks the classifier_model to rank each policy label and reads off the softmax; \"colbert\" reranks policy descriptions against the prompt via a reranker model; \"knn\" votes over a curated corpus of labelled example prompts (seeded via the corpus API) and routes to the fallback when the prompt is unlike all corpus entries. Empty defaults to \"score\".", Component: "select", Options: []FieldOption{ {Value: "score", Label: "Score (Arch-Router-style)"}, + {Value: "colbert", Label: "Colbert (reranker)"}, + {Value: "knn", Label: "KNN (labelled corpus)"}, }, Order: 230, }, "router.classifier_model": { Section: "router", Label: "Classifier Model", - Description: "Loaded LocalAI model the score classifier asks to rank each policy label as a continuation. Must support the Score gRPC primitive (today: llama-cpp, vLLM) and use the ChatML template. Arch-Router-1.5B Q4_K_M is the canonical choice; any small ChatML instruct model also works at a higher activation_threshold.", + Description: "Loaded LocalAI model the score classifier asks to rank each policy label as a continuation (for colbert: the reranker model). Must support the Score gRPC primitive (today: llama-cpp, vLLM) and use the ChatML template. Arch-Router-1.5B Q4_K_M is the canonical choice; any small ChatML instruct model also works at a higher activation_threshold. Not used by the knn classifier.", Component: "model-select", AutocompleteProvider: ProviderModelsScore, Order: 231, @@ -949,5 +951,48 @@ func DefaultRegistry() map[string]FieldMetaOverride { Component: "input", Order: 240, }, + "router.knn.embedding_model": { + Section: "router", + Label: "KNN: Embedding Model", + Description: "Embedding model the knn classifier uses for corpus entries and incoming prompts. Required when classifier is \"knn\". Changing it invalidates stored vectors — entries recorded under a different embedder are re-embedded on load. nomic-embed-text-v1.5 is the recommended default.", + Component: "model-select", + AutocompleteProvider: ProviderModels, + Order: 241, + }, + "router.knn.k": { + Section: "router", + Label: "KNN: Neighbours (K)", + Description: "How many nearest corpus entries vote on a prompt. 0 picks the default (3). K=1 routes on the single nearest example; larger K tolerates a mislabelled exemplar but needs denser corpus coverage per label.", + Component: "number", + Min: f64(0), + Order: 242, + }, + "router.knn.similarity_threshold": { + Section: "router", + Label: "KNN: Similarity Threshold", + Description: "Cosine-similarity floor a corpus entry must clear to vote. When no entry clears it the router uses the fallback model — a prompt unlike all labelled examples is treated as undecidable rather than guessed. 0 picks the default (0.80).", + Component: "slider", + Min: f64(0), + Max: f64(1), + Step: f64(0.01), + Order: 243, + }, + "router.knn.vote_threshold": { + Section: "router", + Label: "KNN: Vote Threshold", + Description: "Similarity-weighted vote share a label needs to activate. 0 picks the default (0.5, a weighted majority). Lower values allow multi-label activations from minority neighbours; higher values demand near-unanimous neighbourhoods.", + Component: "slider", + Min: f64(0), + Max: f64(1), + Step: f64(0.05), + Order: 244, + }, + "router.knn.store_name": { + Section: "router", + Label: "KNN: Store Name", + Description: "Optional override for the local-store collection holding the corpus vectors. Empty defaults to \"router-corpus-\".", + Component: "input", + Order: 245, + }, } } diff --git a/core/config/model_config.go b/core/config/model_config.go index a71a2c683372..2c76446e9b67 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -348,7 +348,17 @@ type RouterConfig struct { // embeddings to past decisions, so semantically-similar prompts // reuse a classification instead of re-running the classifier // model. Omit the block to disable. See router/embedding_cache.go. + // Ignored (with a warning) for the knn classifier — that IS a + // KNN lookup already; wrapping it in another would embed twice + // for no additional information. EmbeddingCache *EmbeddingCacheConfig `yaml:"embedding_cache,omitempty" json:"embedding_cache,omitempty"` + + // KNN configures the "knn" classifier: nearest-neighbour voting + // over a curated corpus of labelled example prompts. Required when + // classifier is "knn", ignored otherwise. The corpus is seeded and + // curated through the router corpus API (never through the UI); + // see router/knn.go for the decision semantics. + KNN *RouterKNNConfig `yaml:"knn,omitempty" json:"knn,omitempty"` } // EmbeddingCacheConfig configures the L2 embedding-similarity decision @@ -382,6 +392,44 @@ type EmbeddingCacheConfig struct { StoreName string `yaml:"store_name,omitempty" json:"store_name,omitempty"` } +// RouterKNNConfig configures the knn classifier. It shares the +// embedding + local-store plumbing with EmbeddingCacheConfig but the +// two are deliberately separate blocks: the cache stores another +// classifier's decisions opportunistically, while the KNN corpus is +// explicit labelled ground truth — different lifecycle, different +// store namespace, different failure story. +type RouterKNNConfig struct { + // EmbeddingModel names the loaded LocalAI model used to embed + // both corpus entries and incoming probes. Required. Changing it + // invalidates the stored vectors — the corpus loader re-embeds + // entries recorded under a different embedder fingerprint. + EmbeddingModel string `yaml:"embedding_model" json:"embedding_model"` + + // K is how many nearest corpus entries vote on a probe. 0 picks + // the package default (3). K=1 reproduces exact nearest-entry + // routing; larger K tolerates mislabelled exemplars at the cost + // of needing denser corpus coverage per label region. + K int `yaml:"k,omitempty" json:"k,omitempty"` + + // SimilarityThreshold is the epistemic gate: corpus entries less + // similar than this to the probe cannot vote, and when none clear + // it the router uses the fallback model — a probe unlike all + // labelled experience is undecidable, not a guess. 0 picks the + // package default (0.80). + SimilarityThreshold float64 `yaml:"similarity_threshold,omitempty" json:"similarity_threshold,omitempty"` + + // VoteThreshold is the similarity-weighted vote share a label + // needs to activate. 0 picks the package default (0.5, a weighted + // majority). Lower values let minority-label neighbours activate + // additional labels (multi-label routing); higher values demand + // near-unanimous neighbourhoods. + VoteThreshold float64 `yaml:"vote_threshold,omitempty" json:"vote_threshold,omitempty"` + + // StoreName overrides the local-store collection holding the + // corpus vectors. Empty defaults to "router-corpus-". + StoreName string `yaml:"store_name,omitempty" json:"store_name,omitempty"` +} + // RouterPolicy is one entry in the label vocabulary. The label string // is what the classifier model emits and what candidates reference in // their Labels field; the description is the natural-language hint diff --git a/core/http/endpoints/localai/api_instructions.go b/core/http/endpoints/localai/api_instructions.go index 02ff85d1d088..ca5acf2d68d4 100644 --- a/core/http/endpoints/localai/api_instructions.go +++ b/core/http/endpoints/localai/api_instructions.go @@ -121,7 +121,7 @@ var instructionDefs = []instructionDef{ Name: "intelligent-routing", Description: "Per-model `router:` configuration that classifies requests and rewrites the served model", Tags: []string{"router"}, - Intro: "Add a `router:` block to a ModelConfig to turn it into a routing model. The block declares a classifier (today: `feature` — handcrafted rules over prompt length and code-fence presence), a list of candidates (label + downstream model + optional rule), and a fallback. When a client addresses the routing model, the RouteModel middleware invokes the classifier, picks a candidate, and rewrites input.Model — the standard model-resolution path then runs ACL, disabled-state, and per-model PII against the chosen target. Depth-1 invariant: candidates must NOT themselves carry a `router:` block; runtime check returns 500 on violation. Decisions are logged to GET /api/router/decisions and surfaced in the /app/middleware Routing tab. POST /api/router/decide is the programmatic decision-oracle: external routers (e.g. an organisation-wide router service) send `{router, input}` and receive the classifier's label set + candidate model WITHOUT LocalAI rewriting, forwarding, or recording the call. Shares the classifier cache with the in-band path so warm-up costs are paid once.", + Intro: "Add a `router:` block to a ModelConfig to turn it into a routing model. The block declares a classifier (`score` — a small model ranks each policy label, Arch-Router-style; `colbert` — a reranker scores policy descriptions against the prompt; `knn` — similarity-weighted vote over a curated corpus of labelled example prompts), `policies` (the label vocabulary), `candidates` (downstream model + labels it serves; first candidate whose labels cover the active set wins, so order small → large), and a `fallback`. The knn classifier needs a `knn: { embedding_model }` block instead of a classifier_model, and reads a persisted corpus seeded via POST /api/router/{name}/corpus with `{entries: [{text, labels}]}` (admin-only; texts are embedded server-side, persisted under the state dir, and NEVER returned by any endpoint — GET /api/router/{name}/corpus/stats reports label counts only, DELETE /api/router/{name}/corpus wipes it). knn routes to the fallback whenever the prompt is less similar than knn.similarity_threshold to every corpus entry — out-of-corpus prompts are treated as undecidable rather than guessed. When a client addresses the routing model, the RouteModel middleware invokes the classifier, picks a candidate, and rewrites input.Model — the standard model-resolution path then runs ACL, disabled-state, and per-model PII against the chosen target. Depth-1 invariant: candidates must NOT themselves carry a `router:` block; runtime check returns 500 on violation. Decisions are logged to GET /api/router/decisions and surfaced in the /app/middleware Routing tab. POST /api/router/decide is the programmatic decision-oracle: external routers (e.g. an organisation-wide router service) send `{router, input}` and receive the classifier's label set + candidate model WITHOUT LocalAI rewriting, forwarding, or recording the call. Shares the classifier cache with the in-band path so warm-up costs are paid once.", }, } diff --git a/core/http/endpoints/localai/router_corpus.go b/core/http/endpoints/localai/router_corpus.go new file mode 100644 index 000000000000..bf2a8223055c --- /dev/null +++ b/core/http/endpoints/localai/router_corpus.go @@ -0,0 +1,189 @@ +package localai + +import ( + "fmt" + "net/http" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/routing/corpus" +) + +// The router corpus endpoints manage the labelled exemplar corpus +// behind the knn classifier. Corpus input is API-only by design: +// entries can contain example user content, so they are seeded and +// curated programmatically and never entered through or displayed in +// the UI — the inspection surface returns label counts, never texts. +// +// All three handlers resolve the router model by path param and +// require it to declare a `router.knn` block; the store name and +// embedding model come from that config so the API can't desync from +// what the classifier actually queries. + +// resolveKNNRouter loads the named model config and returns its router +// KNN settings (store name defaulted the same way buildClassifier +// defaults it). Echo-shaped errors for the three failure modes. +func resolveKNNRouter(c echo.Context, loader *config.ModelConfigLoader, appConfig *config.ApplicationConfig) (*config.ModelConfig, string, error) { + name := c.Param("name") + if name == "" { + return nil, "", echo.NewHTTPError(http.StatusBadRequest, "router name is required") + } + cfg, err := loader.LoadModelConfigFileByNameDefaultOptions(name, appConfig) + if err != nil { + return nil, "", echo.NewHTTPError(http.StatusInternalServerError, "failed to load model config: "+err.Error()) + } + // A synthetic stub (no Name) means the model is unknown — see + // RouterDecideEndpoint for the discrimination rationale. + if cfg == nil || cfg.Name == "" { + return nil, "", echo.NewHTTPError(http.StatusNotFound, fmt.Sprintf("model %q not found", name)) + } + if cfg.Router.KNN == nil || cfg.Router.KNN.EmbeddingModel == "" { + return nil, "", echo.NewHTTPError(http.StatusBadRequest, + fmt.Sprintf("model %q has no router.knn block (set classifier: knn and knn.embedding_model first)", name)) + } + storeName := cfg.Router.KNN.StoreName + if storeName == "" { + storeName = "router-corpus-" + cfg.Name + } + return cfg, storeName, nil +} + +// RouterCorpusAddEndpoint bulk-seeds a router's KNN corpus. Entries +// are validated against the router's policy labels, embedded with the +// router's knn.embedding_model, persisted to the corpus file, and +// upserted into the live vector index — routing sees them immediately, +// no reload required. +// +// @Summary Seed the KNN routing corpus with labelled example prompts +// @Tags router +// @Accept json +// @Produce json +// @Param name path string true "router model name" +// @Param request body schema.RouterCorpusAddRequest true "labelled exemplars" +// @Success 200 {object} schema.RouterCorpusAddResponse +// @Failure 400 {object} map[string]string +// @Failure 404 {object} map[string]string +// @Failure 500 {object} map[string]string +// @Router /api/router/{name}/corpus [post] +func RouterCorpusAddEndpoint(loader *config.ModelConfigLoader, appConfig *config.ApplicationConfig, mgr *corpus.Manager, deps middleware.ClassifierDeps) echo.HandlerFunc { + return func(c echo.Context) error { + cfg, storeName, err := resolveKNNRouter(c, loader, appConfig) + if err != nil { + return err + } + var req schema.RouterCorpusAddRequest + if err := c.Bind(&req); err != nil { + return echo.NewHTTPError(http.StatusBadRequest, "invalid request body: "+err.Error()) + } + if len(req.Entries) == 0 { + return echo.NewHTTPError(http.StatusBadRequest, "entries is required") + } + + // Labels must be declared policies — the same invariant + // candidate tables are validated against. Catching it here + // keeps a typo from silently creating an unroutable label. + declared := map[string]struct{}{} + for _, p := range cfg.Router.Policies { + declared[p.Label] = struct{}{} + } + entries := make([]corpus.Entry, 0, len(req.Entries)) + for i, e := range req.Entries { + for _, l := range e.Labels { + if _, ok := declared[l]; !ok { + return echo.NewHTTPError(http.StatusBadRequest, + fmt.Sprintf("entry %d: label %q is not declared in router policies", i, l)) + } + } + entries = append(entries, corpus.Entry{Text: e.Text, Labels: e.Labels}) + } + + embedder := deps.Embedder(cfg.Router.KNN.EmbeddingModel) + if embedder == nil { + return echo.NewHTTPError(http.StatusBadRequest, + fmt.Sprintf("embedding_model %q not loadable", cfg.Router.KNN.EmbeddingModel)) + } + store := deps.VectorStore(storeName) + + added, skipped, err := mgr.Add(c.Request().Context(), storeName, cfg.Router.KNN.EmbeddingModel, embedder, store, entries) + if err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, err.Error()) + } + stats, err := mgr.Stats(storeName) + if err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, err.Error()) + } + return c.JSON(http.StatusOK, schema.RouterCorpusAddResponse{ + Router: cfg.Name, + Added: added, + Skipped: skipped, + Total: stats.Total, + LabelCounts: stats.LabelCounts, + }) + } +} + +// RouterCorpusStatsEndpoint reports corpus size and per-label counts. +// Deliberately count-only: corpus texts never leave the server. +// +// @Summary Inspect a router's KNN corpus (label counts only, never texts) +// @Tags router +// @Produce json +// @Param name path string true "router model name" +// @Success 200 {object} schema.RouterCorpusStatsResponse +// @Failure 400 {object} map[string]string +// @Failure 404 {object} map[string]string +// @Failure 500 {object} map[string]string +// @Router /api/router/{name}/corpus/stats [get] +func RouterCorpusStatsEndpoint(loader *config.ModelConfigLoader, appConfig *config.ApplicationConfig, mgr *corpus.Manager) echo.HandlerFunc { + return func(c echo.Context) error { + cfg, storeName, err := resolveKNNRouter(c, loader, appConfig) + if err != nil { + return err + } + stats, err := mgr.Stats(storeName) + if err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, err.Error()) + } + return c.JSON(http.StatusOK, schema.RouterCorpusStatsResponse{ + Router: cfg.Name, + StoreName: storeName, + EmbeddingModel: cfg.Router.KNN.EmbeddingModel, + Total: stats.Total, + LabelCounts: stats.LabelCounts, + EmbeddingModels: stats.EmbeddingModels, + }) + } +} + +// RouterCorpusClearEndpoint wipes a router's corpus — file and live +// index. Reseed with POST afterwards; there is intentionally no +// per-entry delete (exemplars are curated as a set; partial edits by +// vector identity are how mislabelled neighbourhoods linger). +// +// @Summary Clear a router's KNN corpus +// @Tags router +// @Produce json +// @Param name path string true "router model name" +// @Success 200 {object} schema.RouterCorpusClearResponse +// @Failure 400 {object} map[string]string +// @Failure 404 {object} map[string]string +// @Failure 500 {object} map[string]string +// @Router /api/router/{name}/corpus [delete] +func RouterCorpusClearEndpoint(loader *config.ModelConfigLoader, appConfig *config.ApplicationConfig, mgr *corpus.Manager, deps middleware.ClassifierDeps) echo.HandlerFunc { + return func(c echo.Context) error { + cfg, storeName, err := resolveKNNRouter(c, loader, appConfig) + if err != nil { + return err + } + cleared, err := mgr.Clear(c.Request().Context(), storeName, deps.VectorStore(storeName)) + if err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, err.Error()) + } + return c.JSON(http.StatusOK, schema.RouterCorpusClearResponse{ + Router: cfg.Name, + Cleared: cleared, + }) + } +} diff --git a/core/http/endpoints/localai/router_corpus_test.go b/core/http/endpoints/localai/router_corpus_test.go new file mode 100644 index 000000000000..bfb042730b79 --- /dev/null +++ b/core/http/endpoints/localai/router_corpus_test.go @@ -0,0 +1,210 @@ +package localai_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/endpoints/localai" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/services/routing/corpus" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" +) + +// The corpus endpoints manage the knn classifier's labelled exemplar +// store. These specs pin the validation surface (knn-only, declared +// labels), the seed → stats → clear lifecycle, and the privacy +// contract: corpus texts never appear in any response body. + +type corpusTestEmbedder struct{} + +func (corpusTestEmbedder) Embed(_ context.Context, text string) ([]float32, error) { + return []float32{float32(len(text)), 1}, nil +} + +// corpusTestStore implements only the narrow VectorStore interface — +// no batch/delete fast paths — so the specs also cover the manager's +// per-entry fallbacks. +type corpusTestStore struct { + inserted int +} + +func (s *corpusTestStore) Search(_ context.Context, _ []float32) (float64, []byte, bool, error) { + return 0, nil, false, nil +} + +func (s *corpusTestStore) SearchK(_ context.Context, _ []float32, _ int) ([]backend.Neighbor, error) { + return nil, nil +} + +func (s *corpusTestStore) Insert(_ context.Context, _ []float32, _ []byte) error { + s.inserted++ + return nil +} + +func writeKNNRouter(modelDir, name string) { + cfg := &config.ModelConfig{ + Name: name, + Router: config.RouterConfig{ + Classifier: "knn", + Fallback: "small-model", + KNN: &config.RouterKNNConfig{EmbeddingModel: "embed-model"}, + Policies: []config.RouterPolicy{ + {Label: "code-generation", Description: "writing or debugging code"}, + {Label: "casual-chat", Description: "small talk"}, + }, + Candidates: []config.RouterCandidate{ + {Model: "small-model", Labels: []string{"casual-chat"}}, + {Model: "big-model", Labels: []string{"code-generation", "casual-chat"}}, + }, + }, + } + b, err := yaml.Marshal(cfg) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(filepath.Join(modelDir, name+".yaml"), []byte(b), 0o644)).To(Succeed()) +} + +var _ = Describe("Router corpus endpoints", func() { + var ( + modelDir string + corpusDir string + appConfig *config.ApplicationConfig + loader *config.ModelConfigLoader + mgr *corpus.Manager + store *corpusTestStore + e *echo.Echo + ) + + const seededText = "please debug this stack trace for me" + + BeforeEach(func() { + d, err := os.MkdirTemp("", "router-corpus-ep-*") + Expect(err).NotTo(HaveOccurred()) + modelDir = d + corpusDir = filepath.Join(d, "corpus") + appConfig = &config.ApplicationConfig{ + Context: context.Background(), + SystemState: &system.SystemState{Model: system.Model{ModelsPath: modelDir}}, + } + loader = config.NewModelConfigLoader(modelDir) + mgr = corpus.NewManager(corpusDir) + store = &corpusTestStore{} + + deps := middleware.ClassifierDeps{ + Embedder: func(string) backend.Embedder { return corpusTestEmbedder{} }, + VectorStore: func(string) backend.VectorStore { return store }, + } + e = echo.New() + e.POST("/api/router/:name/corpus", localai.RouterCorpusAddEndpoint(loader, appConfig, mgr, deps)) + e.GET("/api/router/:name/corpus/stats", localai.RouterCorpusStatsEndpoint(loader, appConfig, mgr)) + e.DELETE("/api/router/:name/corpus", localai.RouterCorpusClearEndpoint(loader, appConfig, mgr, deps)) + }) + + AfterEach(func() { + _ = os.RemoveAll(modelDir) + }) + + do := func(method, path, body string) *httptest.ResponseRecorder { + var rd *strings.Reader + if body == "" { + rd = strings.NewReader("") + } else { + rd = strings.NewReader(body) + } + req := httptest.NewRequest(method, path, rd) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + e.ServeHTTP(rec, req) + return rec + } + + seedBody := `{"entries":[ + {"text":"` + seededText + `","labels":["code-generation"]}, + {"text":"how is your day going","labels":["casual-chat"]} + ]}` + + It("returns 404 for an unknown router", func() { + rec := do(http.MethodPost, "/api/router/nope/corpus", seedBody) + Expect(rec.Code).To(Equal(http.StatusNotFound)) + }) + + It("returns 400 for a router without a knn block", func() { + writeScoreRouter(modelDir, "score-router") + rec := do(http.MethodPost, "/api/router/score-router/corpus", seedBody) + Expect(rec.Code).To(Equal(http.StatusBadRequest)) + Expect(rec.Body.String()).To(ContainSubstring("router.knn")) + }) + + It("rejects entries with undeclared labels", func() { + writeKNNRouter(modelDir, "knn-router") + rec := do(http.MethodPost, "/api/router/knn-router/corpus", + `{"entries":[{"text":"hello","labels":["not-a-policy"]}]}`) + Expect(rec.Code).To(Equal(http.StatusBadRequest)) + Expect(rec.Body.String()).To(ContainSubstring("not declared")) + }) + + It("rejects an empty entries list", func() { + writeKNNRouter(modelDir, "knn-router") + rec := do(http.MethodPost, "/api/router/knn-router/corpus", `{"entries":[]}`) + Expect(rec.Code).To(Equal(http.StatusBadRequest)) + }) + + It("seeds, reports counts only, and clears", func() { + writeKNNRouter(modelDir, "knn-router") + + rec := do(http.MethodPost, "/api/router/knn-router/corpus", seedBody) + Expect(rec.Code).To(Equal(http.StatusOK), rec.Body.String()) + var addResp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &addResp)).To(Succeed()) + Expect(addResp["added"]).To(BeNumerically("==", 2)) + Expect(addResp["total"]).To(BeNumerically("==", 2)) + Expect(store.inserted).To(Equal(2)) + + // The seed response must not echo the entry texts back. + Expect(rec.Body.String()).NotTo(ContainSubstring(seededText)) + + rec = do(http.MethodGet, "/api/router/knn-router/corpus/stats", "") + Expect(rec.Code).To(Equal(http.StatusOK)) + var stats map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &stats)).To(Succeed()) + Expect(stats["total"]).To(BeNumerically("==", 2)) + Expect(stats["label_counts"]).To(HaveKeyWithValue("code-generation", BeNumerically("==", 1))) + Expect(stats["embedding_model"]).To(Equal("embed-model")) + // Privacy contract: the inspection surface never returns texts. + Expect(rec.Body.String()).NotTo(ContainSubstring(seededText)) + + rec = do(http.MethodDelete, "/api/router/knn-router/corpus", "") + Expect(rec.Code).To(Equal(http.StatusOK)) + var clr map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &clr)).To(Succeed()) + Expect(clr["cleared"]).To(BeNumerically("==", 2)) + + rec = do(http.MethodGet, "/api/router/knn-router/corpus/stats", "") + Expect(rec.Code).To(Equal(http.StatusOK)) + var after map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &after)).To(Succeed()) + Expect(after["total"]).To(BeNumerically("==", 0)) + }) + + It("skips duplicate texts on reseed", func() { + writeKNNRouter(modelDir, "knn-router") + Expect(do(http.MethodPost, "/api/router/knn-router/corpus", seedBody).Code).To(Equal(http.StatusOK)) + rec := do(http.MethodPost, "/api/router/knn-router/corpus", seedBody) + Expect(rec.Code).To(Equal(http.StatusOK)) + var resp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + Expect(resp["added"]).To(BeNumerically("==", 0)) + Expect(resp["skipped"]).To(BeNumerically("==", 2)) + Expect(resp["total"]).To(BeNumerically("==", 2)) + }) +}) diff --git a/core/http/endpoints/localai/router_decide.go b/core/http/endpoints/localai/router_decide.go index 11fddcf574f7..b13b3a7acbc0 100644 --- a/core/http/endpoints/localai/router_decide.go +++ b/core/http/endpoints/localai/router_decide.go @@ -95,15 +95,16 @@ func RouterDecideEndpoint(loader *config.ModelConfigLoader, appConfig *config.Ap } return c.JSON(http.StatusOK, schema.RouterDecideResponse{ - Router: req.Router, - Classifier: classifierName, - Labels: decision.Labels, - Candidate: candidate, - Fallback: fallback, - Score: decision.Score, - LatencyMs: decision.Latency.Milliseconds(), - Cached: decision.Cached, - CacheSimilarity: decision.CacheSimilarity, + Router: req.Router, + Classifier: classifierName, + Labels: decision.Labels, + Candidate: candidate, + Fallback: fallback, + Score: decision.Score, + LatencyMs: decision.Latency.Milliseconds(), + Cached: decision.Cached, + CacheSimilarity: decision.CacheSimilarity, + NearestSimilarity: decision.NearestSimilarity, }) } } diff --git a/core/http/endpoints/mcp/localai_assistant_test.go b/core/http/endpoints/mcp/localai_assistant_test.go index 2231350d075a..4483406d13d5 100644 --- a/core/http/endpoints/mcp/localai_assistant_test.go +++ b/core/http/endpoints/mcp/localai_assistant_test.go @@ -186,3 +186,15 @@ var _ = Describe("LocalAIAssistantHolder", func() { Expect(exec.HasTools()).To(BeFalse()) }) }) + +func (stubClient) GetRouterCorpusStats(_ context.Context, routerModel string) (*localaitools.RouterCorpusStats, error) { + return &localaitools.RouterCorpusStats{Router: routerModel, LabelCounts: map[string]int{}}, nil +} + +func (stubClient) SeedRouterCorpus(_ context.Context, req localaitools.RouterCorpusSeedRequest) (*localaitools.RouterCorpusSeedResult, error) { + return &localaitools.RouterCorpusSeedResult{Router: req.Router, LabelCounts: map[string]int{}}, nil +} + +func (stubClient) ClearRouterCorpus(_ context.Context, routerModel string) (*localaitools.RouterCorpusClearResult, error) { + return &localaitools.RouterCorpusClearResult{Router: routerModel}, nil +} diff --git a/core/http/endpoints/openai/realtime_model.go b/core/http/endpoints/openai/realtime_model.go index 0449daee3740..d980e8c34889 100644 --- a/core/http/endpoints/openai/realtime_model.go +++ b/core/http/endpoints/openai/realtime_model.go @@ -550,6 +550,7 @@ func buildRealtimeRoutingContext(a *application.Application, sessionID string) * } deps := &middleware.ClassifierDeps{ Scorer: a.Scorer, + Corpus: a.RouterCorpus(), TokenCounter: a.TokenCounter, Embedder: a.Embedder, VectorStore: a.VectorStore, diff --git a/core/http/middleware/route_model.go b/core/http/middleware/route_model.go index 470bd05f5aa2..a7e3255fb39d 100644 --- a/core/http/middleware/route_model.go +++ b/core/http/middleware/route_model.go @@ -49,6 +49,16 @@ type RerankerFactory func(modelName string) backend.Reranker // lives in ModelConfig.Validate() and runs at config load/save time. type ModelConfigLookup func(modelName string) *config.ModelConfig +// CorpusLoader syncs a router's persisted KNN corpus into the live +// vector index when the knn classifier is built. Implemented by +// corpus.Manager; declared here (consumer side) so the middleware +// doesn't depend on the corpus package. Optional — when nil, the knn +// classifier serves whatever the index already holds (tests, embedded +// callers). +type CorpusLoader interface { + EnsureLoaded(ctx context.Context, storeName, embeddingModel string, embedder backend.Embedder, store backend.VectorStore) (int, error) +} + // ClassifierDeps bundles the backend factories the router middleware // needs to build a classifier and its optional L2 cache. Bundled into // one struct because RouteModel already takes many positional @@ -66,6 +76,10 @@ type ClassifierDeps struct { VectorStore VectorStoreFactory Reranker RerankerFactory + // Corpus loads the persisted KNN corpus into the vector index when + // a knn classifier is built. Optional; nil skips the load. + Corpus CorpusLoader + // ModelLookup resolves the classifier_model name to its config so // buildClassifier can reject misconfigurations that would // otherwise crash the llama-cpp backend at request time. Optional @@ -375,13 +389,61 @@ func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Class rerankClassifier = rerankClassifier.WithTokenTrim(count, ctxTokens) } inner = rerankClassifier + case router.ClassifierKNN: + if rc.KNN == nil || rc.KNN.EmbeddingModel == "" { + return nil, fmt.Errorf("router classifier knn requires a knn block with embedding_model") + } + if deps.Embedder == nil || deps.VectorStore == nil { + return nil, fmt.Errorf("router classifier knn unavailable: embedder/vector-store factories not wired") + } + embedder := deps.Embedder(rc.KNN.EmbeddingModel) + if embedder == nil { + return nil, fmt.Errorf("router classifier knn: embedding_model %q not loadable", rc.KNN.EmbeddingModel) + } + storeName := rc.KNN.StoreName + if storeName == "" { + storeName = "router-corpus-" + cfg.Name + } + vstore := deps.VectorStore(storeName) + if vstore == nil { + return nil, fmt.Errorf("router classifier knn: vector store %q not loadable", storeName) + } + if deps.Corpus != nil { + // A failed load must not break routing — the classifier + // still works against whatever the index holds, and the + // epistemic gate falls back safely on an empty index. + if n, err := deps.Corpus.EnsureLoaded(context.Background(), storeName, rc.KNN.EmbeddingModel, embedder, vstore); err != nil { + xlog.Warn("router: knn corpus load failed; routing continues on the live index", + "router_model", cfg.Name, "store", storeName, "error", err) + } else if n > 0 { + xlog.Info("router: knn corpus loaded", + "router_model", cfg.Name, "store", storeName, "entries", n) + } + } + knnClassifier := router.NewKNNClassifier(embedder, vstore, router.KNNClassifierOptions{ + K: rc.KNN.K, + SimilarityThreshold: rc.KNN.SimilarityThreshold, + VoteThreshold: rc.KNN.VoteThreshold, + }) + if count, ctxTokens := modelTokenTrim(rc.KNN.EmbeddingModel, deps); count != nil { + knnClassifier = knnClassifier.WithTokenTrim(count, ctxTokens) + } + inner = knnClassifier default: - return nil, fmt.Errorf("router: unknown classifier %q (supported: %s)", name, strings.Join([]string{router.ClassifierScore, router.ClassifierColbert}, ", ")) + return nil, fmt.Errorf("router: unknown classifier %q (supported: %s)", name, strings.Join([]string{router.ClassifierScore, router.ClassifierColbert, router.ClassifierKNN}, ", ")) } if rc.EmbeddingCache == nil { return inner, nil } + if name == router.ClassifierKNN { + // The knn classifier IS an embedding-KNN lookup — wrapping it in + // the embedding cache would embed every probe twice to answer the + // same question. Ignore the block rather than fail routing. + xlog.Warn("router: embedding_cache ignored for knn classifier", + "router_model", cfg.Name) + return inner, nil + } wrapped, err := wrapWithEmbeddingCache(cfg, inner, deps) if err != nil { // Caching plumbing problems must not break routing — log, @@ -428,7 +490,11 @@ func assertClassifierDeclaresScore(classifierModel string, lookup ModelConfigLoo // returns the parsed []ScorePolicy. Both Score and Rerank classifiers // take the same policy shape. func validateRouterPolicies(classifierName string, rc config.RouterConfig) ([]router.ScorePolicy, error) { - if rc.ClassifierModel == "" { + // The knn classifier has no scoring model — its label knowledge + // lives in the corpus. It still declares policies: they are the + // label vocabulary corpus entries are validated against and the + // "categories" surface the Routing tab shows. + if rc.ClassifierModel == "" && classifierName != router.ClassifierKNN { return nil, fmt.Errorf("router classifier %s requires classifier_model", classifierName) } if len(rc.Policies) == 0 { diff --git a/core/http/middleware/route_model_test.go b/core/http/middleware/route_model_test.go index 4a9be2b12fb3..c8f85b8f3a05 100644 --- a/core/http/middleware/route_model_test.go +++ b/core/http/middleware/route_model_test.go @@ -2,6 +2,7 @@ package middleware_test import ( "context" + "errors" "net/http" "net/http/httptest" "os" @@ -540,3 +541,185 @@ template: ` Expect(os.WriteFile(filepath.Join(modelDir, name+".yaml"), []byte(body), 0o644)).To(Succeed()) } + +// --- knn classifier middleware specs --- + +// fixedEmbedder returns the same vector for every probe — the KNN +// middleware specs script the neighbourhood in the store fake, so the +// query vector itself is irrelevant. +type fixedEmbedder struct{} + +func (fixedEmbedder) Embed(_ context.Context, _ string) ([]float32, error) { + return []float32{1, 0, 0}, nil +} + +// scriptedVectorStore returns a fixed neighbour list from SearchK. The +// classifier must never call Search (top-1 is the cache's shape) nor +// Insert (the KNN corpus grows only through explicit curation) — both +// error loudly so a regression shows up as a failed spec. +type scriptedVectorStore struct { + neighbors []backend.Neighbor +} + +func (s *scriptedVectorStore) SearchK(_ context.Context, _ []float32, k int) ([]backend.Neighbor, error) { + if len(s.neighbors) > k { + return s.neighbors[:k], nil + } + return s.neighbors, nil +} + +func (s *scriptedVectorStore) Search(_ context.Context, _ []float32) (float64, []byte, bool, error) { + return 0, nil, false, errTestKNNSearch +} + +func (s *scriptedVectorStore) Insert(_ context.Context, _ []float32, _ []byte) error { + return errTestKNNInsert +} + +var ( + errTestKNNSearch = errors.New("knn classifier must use SearchK, not Search") + errTestKNNInsert = errors.New("knn classifier must never insert into the corpus") +) + +func corpusPayload(labels ...string) []byte { + b, err := router.EncodeCorpusEntry(labels) + Expect(err).NotTo(HaveOccurred()) + return b +} + +// newKNNRouterModel mirrors newScoreRouterModel but with the knn +// classifier: same policy vocabulary and candidate table, no +// classifier_model, corpus semantics supplied by the store fake. +func newKNNRouterModel(modelDir, name string) *config.ModelConfig { + cfg := &config.ModelConfig{ + Name: name, + Router: config.RouterConfig{ + Classifier: "knn", + Fallback: "qwen3-0.6b", + KNN: &config.RouterKNNConfig{EmbeddingModel: "embed-model"}, + Policies: []config.RouterPolicy{ + {Label: "code-generation", Description: "writing or debugging code"}, + {Label: "casual-chat", Description: "small talk"}, + {Label: "math-reasoning", Description: "arithmetic and word problems"}, + }, + Candidates: []config.RouterCandidate{ + {Model: "small-model", Labels: []string{"casual-chat"}}, + {Model: "big-model", Labels: []string{"code-generation", "casual-chat", "math-reasoning"}}, + }, + }, + } + Expect(os.WriteFile(filepath.Join(modelDir, name+".yaml"), []byte(toYAML(cfg)), 0o644)).To(Succeed()) + return cfg +} + +var _ = Describe("RouteModel middleware (knn classifier)", func() { + var ( + modelDir string + appConfig *config.ApplicationConfig + loader *config.ModelConfigLoader + store *fakeDecisionStore + vstore *scriptedVectorStore + storeName string + ) + + BeforeEach(func() { + d, err := os.MkdirTemp("", "router-knn-test-*") + Expect(err).NotTo(HaveOccurred()) + modelDir = d + appConfig = &config.ApplicationConfig{ + Context: context.Background(), + SystemState: &system.SystemState{Model: system.Model{ModelsPath: modelDir}}, + } + loader = config.NewModelConfigLoader(modelDir) + store = &fakeDecisionStore{} + vstore = &scriptedVectorStore{} + storeName = "" + }) + + AfterEach(func() { + _ = os.RemoveAll(modelDir) + }) + + knnDeps := func() ClassifierDeps { + return ClassifierDeps{ + Embedder: func(string) backend.Embedder { return fixedEmbedder{} }, + VectorStore: func(name string) backend.VectorStore { + storeName = name + return vstore + }, + } + } + + It("routes to the candidate covering the corpus vote", func() { + routerCfg := newKNNRouterModel(modelDir, "smart-router") + writeCandidate(modelDir, "small-model") + writeCandidate(modelDir, "big-model") + vstore.neighbors = []backend.Neighbor{ + {Similarity: 0.92, Payload: corpusPayload("code-generation")}, + {Similarity: 0.88, Payload: corpusPayload("code-generation")}, + } + + rec, err := runRouterWithDeps(loader, appConfig, store, routerCfg, + openAIChat("debug my Go null pointer"), knnDeps()) + Expect(err).NotTo(HaveOccurred()) + Expect(rec.Body.String()).To(Equal("served:big-model")) + Expect(storeName).To(Equal("router-corpus-smart-router"), + "corpus store must be namespaced per router, separate from the decision cache") + Expect(store.records).To(HaveLen(1)) + Expect(store.records[0].Classifier).To(Equal("knn")) + Expect(store.records[0].Label).To(ContainSubstring("code-generation")) + Expect(store.records[0].NearestSimilarity).To(BeNumerically("~", 0.92, 1e-9)) + }) + + It("falls back when the probe is out of corpus range (epistemic gate)", func() { + routerCfg := newKNNRouterModel(modelDir, "smart-router") + writeCandidate(modelDir, "small-model") + writeCandidate(modelDir, "big-model") + writeCandidate(modelDir, "qwen3-0.6b") + // Nearest labelled experience is far below the 0.80 default gate. + vstore.neighbors = []backend.Neighbor{ + {Similarity: 0.42, Payload: corpusPayload("code-generation")}, + } + + rec, err := runRouterWithDeps(loader, appConfig, store, routerCfg, + openAIChat("translate this sanskrit poem"), knnDeps()) + Expect(err).NotTo(HaveOccurred()) + Expect(rec.Body.String()).To(Equal("served:qwen3-0.6b")) + Expect(store.records).To(HaveLen(1)) + Expect(store.records[0].Label).To(Equal(router.LabelFallback)) + // The decision log must still carry the epistemic signal — how + // close the nearest corpus entry was to routing this probe. + Expect(store.records[0].NearestSimilarity).To(BeNumerically("~", 0.42, 1e-9)) + }) + + It("rejects a knn router without a knn block at build time", func() { + routerCfg := newKNNRouterModel(modelDir, "smart-router") + routerCfg.Router.KNN = nil + writeCandidate(modelDir, "small-model") + writeCandidate(modelDir, "big-model") + + _, err := runRouterWithDeps(loader, appConfig, store, routerCfg, + openAIChat("hello"), knnDeps()) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("knn")) + }) + + It("ignores an embedding_cache block instead of double-embedding", func() { + routerCfg := newKNNRouterModel(modelDir, "smart-router") + routerCfg.Router.EmbeddingCache = &config.EmbeddingCacheConfig{EmbeddingModel: "embed-model"} + writeCandidate(modelDir, "small-model") + writeCandidate(modelDir, "big-model") + vstore.neighbors = []backend.Neighbor{ + {Similarity: 0.95, Payload: corpusPayload("casual-chat")}, + } + + rec, err := runRouterWithDeps(loader, appConfig, store, routerCfg, + openAIChat("hi"), knnDeps()) + Expect(err).NotTo(HaveOccurred()) + // Routed by the corpus (small-model covers casual-chat) and NOT + // wrapped: a cache wrap would have called the fake's Search, + // which errors loudly. + Expect(rec.Body.String()).To(Equal("served:small-model")) + Expect(store.records[0].Cached).To(BeFalse()) + }) +}) diff --git a/core/http/react-ui/e2e/middleware-page.spec.js b/core/http/react-ui/e2e/middleware-page.spec.js index 98e011c1e1d4..a12733a99659 100644 --- a/core/http/react-ui/e2e/middleware-page.spec.js +++ b/core/http/react-ui/e2e/middleware-page.spec.js @@ -57,9 +57,31 @@ const MOCK_STATUS = { }, }, }, + { + name: 'knn-router', + classifier: 'knn', + fallback: 'qwen-7b', + policies: [ + { label: 'casual-chat', description: 'small talk' }, + { label: 'code-generation', description: 'writing or debugging code' }, + ], + candidates: [ + { model: 'qwen-3b', labels: ['casual-chat'] }, + { model: 'qwen-coder', labels: ['code-generation', 'casual-chat'] }, + ], + knn: { + embedding_model: 'nomic-embed-text-v1.5', + k: 3, + similarity_threshold: 0.80, + vote_threshold: 0.5, + store_name: 'router-corpus-knn-router', + // Counts only — the status endpoint never sends corpus texts. + corpus: { total: 12, label_counts: { 'code-generation': 7, 'casual-chat': 5 } }, + }, + }, ], recent_decision_count: 1, - available_classifiers: ['score'], + available_classifiers: ['score', 'colbert', 'knn'], }, } @@ -72,6 +94,13 @@ const MOCK_DECISIONS = { cached: true, cache_similarity: 0.92, created_at: '2026-05-06T11:00:00Z', }, + { + id: 'rd_a2', correlation_id: 'corr-2', user_id: 'local', + router_model: 'knn-router', requested_model: 'knn-router', served_model: 'qwen-7b', + classifier: 'knn', label: 'fallback', score: 0, latency_ms: 9, + cached: false, nearest_similarity: 0.42, + created_at: '2026-05-06T11:01:00Z', + }, ], } @@ -229,6 +258,26 @@ test.describe('Middleware page — admin in no-auth mode', () => { await expect(page.getByText('casual-chat').first()).toBeVisible() }) + test('Routing tab renders knn corpus stats and out-of-corpus fallback detail', async ({ page }) => { + await page.goto('/app/middleware') + await page.getByRole('button', { name: /Routing/i }).click() + + // KNN router row: corpus size, K, and gate threshold in the + // Cache / corpus column. + await expect(page.getByText('knn-router').first()).toBeVisible() + await expect(page.getByText(/12 exemplars · k=3 · sim ≥ 0\.8/).first()).toBeVisible() + + // Per-label exemplar counts — counts only, never corpus texts. + await expect(page.getByText(/code-generation: 7/).first()).toBeVisible() + await expect(page.getByText(/casual-chat: 5/).first()).toBeVisible() + + // Expanding the knn fallback decision explains the epistemic gate + // and surfaces how far away the nearest labelled experience was. + await page.getByText('fallback', { exact: true }).first().click() + await expect(page.getByText(/Out-of-corpus fallback/i).first()).toBeVisible() + await expect(page.getByText(/similarity 0\.42/).first()).toBeVisible() + }) + test('Routing tab renders embedding-cache stats and similarity histogram', async ({ page }) => { await page.goto('/app/middleware') await page.getByRole('button', { name: /Routing/i }).click() diff --git a/core/http/react-ui/src/pages/Middleware.jsx b/core/http/react-ui/src/pages/Middleware.jsx index 027967f694d8..0fb10ae61903 100644 --- a/core/http/react-ui/src/pages/Middleware.jsx +++ b/core/http/react-ui/src/pages/Middleware.jsx @@ -512,7 +512,9 @@ function DecisionDetail({ d }) {
{d.cached ? 'Cached decision — per-label scores not recorded (the cache stores only the resulting label set).' - : 'No per-label scores recorded for this decision (likely a fallback row).'} + : d.nearest_similarity + ? `Out-of-corpus fallback — the nearest labelled corpus entry was at similarity ${d.nearest_similarity.toFixed(2)}, below the router's gate. Seed exemplars near this kind of prompt to route it.` + : 'No per-label scores recorded for this decision (likely a fallback row).'}
) } @@ -611,7 +613,7 @@ function RoutingTab({ status, decisions }) { Model Classifier Candidates - Embedding cache + Cache / corpus Fallback @@ -632,7 +634,7 @@ function RoutingTab({ status, decisions }) { ))} - + {m.knn ? : } {m.fallback || '—'} @@ -1104,6 +1106,47 @@ function EventsTab({ events }) { // for configured caches, shows hit/miss/near-miss counters plus a // similarity histogram with a marker at the configured threshold so // admins can tell at a glance whether the threshold is well-placed. +// RouterKNNCell summarises a knn router's corpus for the Active +// routers table: embedding model, corpus size, per-label exemplar +// counts, and the epistemic-gate threshold. Counts only — corpus +// texts never reach the UI (the status endpoint doesn't send them, +// by design; seeding/curation is API-only). +function RouterKNNCell({ knn }) { + if (!knn) { + return — + } + const corpus = knn.corpus || {} + const total = corpus.total || 0 + const counts = corpus.label_counts || {} + const k = knn.k || 3 + const sim = knn.similarity_threshold || 0.80 + return ( +
+
{knn.embedding_model}
+
+ {total === 0 ? ( + + empty corpus — seed via API + + ) : ( + + {total} exemplars · k={k} · sim ≥ {sim} + + )} +
+ {total > 0 && ( +
+ {Object.entries(counts).sort((a, b) => b[1] - a[1] || a[0].localeCompare(b[0])).map(([label, n]) => ( + + {label}: {n} + + ))} +
+ )} +
+ ) +} + function RouterCacheCell({ cache }) { if (!cache) { return — diff --git a/core/http/routes/anthropic.go b/core/http/routes/anthropic.go index 288aa6f57b1a..4565a4e7815a 100644 --- a/core/http/routes/anthropic.go +++ b/core/http/routes/anthropic.go @@ -57,6 +57,7 @@ func RegisterAnthropicRoutes(app *echo.Echo, router.SourceAnthropic, middleware.ClassifierDeps{ Scorer: application.Scorer, + Corpus: application.RouterCorpus(), TokenCounter: application.TokenCounter, Embedder: application.Embedder, VectorStore: application.VectorStore, diff --git a/core/http/routes/middleware.go b/core/http/routes/middleware.go index 6d130863ad5e..b22c7874c4db 100644 --- a/core/http/routes/middleware.go +++ b/core/http/routes/middleware.go @@ -136,6 +136,7 @@ func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { app.ApplicationConfig(), middleware.ClassifierDeps{ Scorer: app.Scorer, + Corpus: app.RouterCorpus(), TokenCounter: app.TokenCounter, Embedder: app.Embedder, VectorStore: app.VectorStore, @@ -155,6 +156,35 @@ func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { } return decideHandler(c) }) + + // Router KNN corpus management. Corpus input/curation is API-only + // by design — entries can contain example user content, so the UI + // never sends or renders them; the stats endpoint returns label + // counts only. Admin-gated like /api/router/decide: seeding the + // corpus changes routing behaviour and Add runs embedding + // inference on arbitrary input. + corpusDeps := middleware.ClassifierDeps{ + Embedder: app.Embedder, + VectorStore: app.VectorStore, + } + corpusAdd := localai.RouterCorpusAddEndpoint(app.ModelConfigLoader(), app.ApplicationConfig(), app.RouterCorpus(), corpusDeps) + corpusStats := localai.RouterCorpusStatsEndpoint(app.ModelConfigLoader(), app.ApplicationConfig(), app.RouterCorpus()) + corpusClear := localai.RouterCorpusClearEndpoint(app.ModelConfigLoader(), app.ApplicationConfig(), app.RouterCorpus(), corpusDeps) + adminGate := func(h echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + viewer := resolveUsageUser(c, app) + if viewer == nil { + return c.JSON(http.StatusUnauthorized, map[string]string{"error": "not authenticated"}) + } + if viewer.Role != auth.RoleAdmin { + return c.JSON(http.StatusForbidden, map[string]string{"error": "admin access required"}) + } + return h(c) + } + } + e.POST("/api/router/:name/corpus", adminGate(corpusAdd)) + e.GET("/api/router/:name/corpus/stats", adminGate(corpusStats)) + e.DELETE("/api/router/:name/corpus", adminGate(corpusClear)) } // buildRouterStatus inventories every model that declares a Router @@ -210,6 +240,28 @@ func buildRouterStatus(app *application.Application) map[string]any { } entry["embedding_cache"] = cacheEntry } + if kc := cfg.Router.KNN; kc != nil { + storeName := kc.StoreName + if storeName == "" { + storeName = "router-corpus-" + cfg.Name + } + knnEntry := map[string]any{ + "embedding_model": kc.EmbeddingModel, + "k": kc.K, + "similarity_threshold": kc.SimilarityThreshold, + "vote_threshold": kc.VoteThreshold, + "store_name": storeName, + } + // Corpus stats are counts only — entry texts never reach + // the UI (or any API surface). + if s, err := app.RouterCorpus().Stats(storeName); err == nil { + knnEntry["corpus"] = map[string]any{ + "total": s.Total, + "label_counts": s.LabelCounts, + } + } + entry["knn"] = knnEntry + } models = append(models, entry) } @@ -224,7 +276,7 @@ func buildRouterStatus(app *application.Application) map[string]any { "configured": hasAny, "models": models, "recent_decision_count": recentCount, - "available_classifiers": []string{router.ClassifierScore}, + "available_classifiers": []string{router.ClassifierScore, router.ClassifierColbert, router.ClassifierKNN}, } if !hasAny { out["note"] = "No router models configured. Add a `router:` block to a model YAML to enable intelligent routing." diff --git a/core/http/routes/openai.go b/core/http/routes/openai.go index 9c118bdbc269..191a54ca1f1f 100644 --- a/core/http/routes/openai.go +++ b/core/http/routes/openai.go @@ -73,6 +73,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, router.SourceChat, middleware.ClassifierDeps{ Scorer: application.Scorer, + Corpus: application.RouterCorpus(), TokenCounter: application.TokenCounter, Embedder: application.Embedder, VectorStore: application.VectorStore, diff --git a/core/schema/localai.go b/core/schema/localai.go index 47d71ec00dac..6c51273fd68b 100644 --- a/core/schema/localai.go +++ b/core/schema/localai.go @@ -51,33 +51,33 @@ type GalleryResponse struct { type VideoRequest struct { BasicModelRequest - Prompt string `json:"prompt" yaml:"prompt"` // text description of the video to generate - NegativePrompt string `json:"negative_prompt" yaml:"negative_prompt"` // things to avoid in the output - StartImage string `json:"start_image" yaml:"start_image"` // URL or base64 of the first frame - EndImage string `json:"end_image" yaml:"end_image"` // URL or base64 of the last frame - Audio string `json:"audio,omitempty" yaml:"audio,omitempty"` // URL or base64 audio for audio-conditioned generation - Width int32 `json:"width" yaml:"width"` // output width in pixels - Height int32 `json:"height" yaml:"height"` // output height in pixels - NumFrames int32 `json:"num_frames" yaml:"num_frames"` // total number of frames to generate - FPS int32 `json:"fps" yaml:"fps"` // frames per second - Seconds string `json:"seconds,omitempty" yaml:"seconds,omitempty"` // duration in seconds (alternative to num_frames) - Size string `json:"size,omitempty" yaml:"size,omitempty"` // WxH shorthand (e.g. "512x512") - InputReference string `json:"input_reference,omitempty" yaml:"input_reference,omitempty"` // reference image or video URL - Seed int32 `json:"seed" yaml:"seed"` // random seed for reproducibility - CFGScale float32 `json:"cfg_scale" yaml:"cfg_scale"` // classifier-free guidance scale - Step int32 `json:"step" yaml:"step"` // number of diffusion steps - ResponseFormat string `json:"response_format" yaml:"response_format"` // output format (url or b64_json) - Params map[string]string `json:"params,omitempty" yaml:"params,omitempty"` // backend-specific generation parameters + Prompt string `json:"prompt" yaml:"prompt"` // text description of the video to generate + NegativePrompt string `json:"negative_prompt" yaml:"negative_prompt"` // things to avoid in the output + StartImage string `json:"start_image" yaml:"start_image"` // URL or base64 of the first frame + EndImage string `json:"end_image" yaml:"end_image"` // URL or base64 of the last frame + Audio string `json:"audio,omitempty" yaml:"audio,omitempty"` // URL or base64 audio for audio-conditioned generation + Width int32 `json:"width" yaml:"width"` // output width in pixels + Height int32 `json:"height" yaml:"height"` // output height in pixels + NumFrames int32 `json:"num_frames" yaml:"num_frames"` // total number of frames to generate + FPS int32 `json:"fps" yaml:"fps"` // frames per second + Seconds string `json:"seconds,omitempty" yaml:"seconds,omitempty"` // duration in seconds (alternative to num_frames) + Size string `json:"size,omitempty" yaml:"size,omitempty"` // WxH shorthand (e.g. "512x512") + InputReference string `json:"input_reference,omitempty" yaml:"input_reference,omitempty"` // reference image or video URL + Seed int32 `json:"seed" yaml:"seed"` // random seed for reproducibility + CFGScale float32 `json:"cfg_scale" yaml:"cfg_scale"` // classifier-free guidance scale + Step int32 `json:"step" yaml:"step"` // number of diffusion steps + ResponseFormat string `json:"response_format" yaml:"response_format"` // output format (url or b64_json) + Params map[string]string `json:"params,omitempty" yaml:"params,omitempty"` // backend-specific generation parameters } // @Description TTS request body type TTSRequest struct { BasicModelRequest - Input string `json:"input" yaml:"input"` // text input - Voice string `json:"voice" yaml:"voice"` // voice audio file or speaker id - Backend string `json:"backend" yaml:"backend"` // backend engine override - Language string `json:"language,omitempty" yaml:"language,omitempty"` // (optional) language to use with TTS model - Format string `json:"response_format,omitempty" yaml:"response_format,omitempty"` // (optional) output format + Input string `json:"input" yaml:"input"` // text input + Voice string `json:"voice" yaml:"voice"` // voice audio file or speaker id + Backend string `json:"backend" yaml:"backend"` // backend engine override + Language string `json:"language,omitempty" yaml:"language,omitempty"` // (optional) language to use with TTS model + Format string `json:"response_format,omitempty" yaml:"response_format,omitempty"` // (optional) output format Stream bool `json:"stream,omitempty" yaml:"stream,omitempty"` // (optional) enable streaming TTS SampleRate int `json:"sample_rate,omitempty" yaml:"sample_rate,omitempty"` // (optional) desired output sample rate // Instructions is a free-form, per-request style/voice description. It maps to @@ -181,10 +181,10 @@ type SystemInformationResponse struct { type DetectionRequest struct { BasicModelRequest Image string `json:"image"` // URL or base64-encoded image to analyze - Prompt string `json:"prompt,omitempty"` // Text prompt (for SAM 3 PCS mode) - Points []float32 `json:"points,omitempty"` // Point coordinates as [x,y,label,...] triples (label: 1=pos, 0=neg) - Boxes []float32 `json:"boxes,omitempty"` // Box coordinates as [x1,y1,x2,y2,...] quads - Threshold float32 `json:"threshold,omitempty"` // Detection confidence threshold + Prompt string `json:"prompt,omitempty"` // Text prompt (for SAM 3 PCS mode) + Points []float32 `json:"points,omitempty"` // Point coordinates as [x,y,label,...] triples (label: 1=pos, 0=neg) + Boxes []float32 `json:"boxes,omitempty"` // Box coordinates as [x1,y1,x2,y2,...] quads + Threshold float32 `json:"threshold,omitempty"` // Detection confidence threshold } type DetectionResponse struct { @@ -256,14 +256,14 @@ type FaceVerifyRequest struct { } type FaceVerifyResponse struct { - Verified bool `json:"verified"` - Distance float32 `json:"distance"` - Threshold float32 `json:"threshold"` - Confidence float32 `json:"confidence"` - Model string `json:"model"` - Img1Area FacialArea `json:"img1_area"` - Img2Area FacialArea `json:"img2_area"` - ProcessingTimeMs float32 `json:"processing_time_ms,omitempty"` + Verified bool `json:"verified"` + Distance float32 `json:"distance"` + Threshold float32 `json:"threshold"` + Confidence float32 `json:"confidence"` + Model string `json:"model"` + Img1Area FacialArea `json:"img1_area"` + Img2Area FacialArea `json:"img2_area"` + ProcessingTimeMs float32 `json:"processing_time_ms,omitempty"` // Liveness fields are only populated when the request set // anti_spoofing=true. Pointers keep them fully absent from the // JSON response otherwise, so callers can tell "not checked" @@ -544,6 +544,64 @@ type RouterDecideResponse struct { // CacheSimilarity carries the cosine similarity of the cache hit // (0 when not cached). CacheSimilarity float64 `json:"cache_similarity,omitempty"` + // NearestSimilarity is the cosine similarity of the closest KNN + // corpus entry — populated by the knn classifier even when the + // decision fell back because the probe was out of corpus range. + // 0 for other classifiers. + NearestSimilarity float64 `json:"nearest_similarity,omitempty"` +} + +// RouterCorpusEntry is one labelled exemplar submitted to +// POST /api/router/{name}/corpus. The text is embedded server-side +// with the router's knn.embedding_model; labels must be declared in +// the router's policies. +type RouterCorpusEntry struct { + Text string `json:"text"` + Labels []string `json:"labels"` +} + +// RouterCorpusAddRequest is the input for POST /api/router/{name}/corpus — +// bulk-seeds the KNN routing corpus. Corpus input is API-only by +// design: entries may contain example user content, so they are never +// entered through (or displayed in) the UI. +type RouterCorpusAddRequest struct { + Entries []RouterCorpusEntry `json:"entries"` +} + +// RouterCorpusAddResponse reports the outcome of a corpus seed call. +type RouterCorpusAddResponse struct { + Router string `json:"router"` + // Added is how many entries were embedded, persisted, and indexed. + Added int `json:"added"` + // Skipped counts entries whose text was already in the corpus — + // duplicates are rejected rather than double-weighted. + Skipped int `json:"skipped"` + // Total is the corpus size after the call. + Total int `json:"total"` + // LabelCounts is the per-label exemplar count after the call. + LabelCounts map[string]int `json:"label_counts"` +} + +// RouterCorpusStatsResponse is the inspection surface for a router's +// KNN corpus: counts and configuration only — entry texts are never +// returned by any endpoint. +type RouterCorpusStatsResponse struct { + Router string `json:"router"` + StoreName string `json:"store_name"` + EmbeddingModel string `json:"embedding_model"` + Total int `json:"total"` + LabelCounts map[string]int `json:"label_counts"` + // EmbeddingModels lists the embedder fingerprints present in the + // persisted corpus; more than one means part of the corpus is + // pending re-embedding on the next load. + EmbeddingModels []string `json:"embedding_models,omitempty"` +} + +// RouterCorpusClearResponse reports how many entries a +// DELETE /api/router/{name}/corpus removed. +type RouterCorpusClearResponse struct { + Router string `json:"router"` + Cleared int `json:"cleared"` } // PIIDecideRequest is the input for POST /api/pii/decide — the diff --git a/core/services/routing/corpus/corpus_suite_test.go b/core/services/routing/corpus/corpus_suite_test.go new file mode 100644 index 000000000000..36e2abea42ba --- /dev/null +++ b/core/services/routing/corpus/corpus_suite_test.go @@ -0,0 +1,13 @@ +package corpus_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestCorpus(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Router corpus manager suite") +} diff --git a/core/services/routing/corpus/manager.go b/core/services/routing/corpus/manager.go new file mode 100644 index 000000000000..e61d8fd84bd8 --- /dev/null +++ b/core/services/routing/corpus/manager.go @@ -0,0 +1,375 @@ +// Package corpus persists and serves the labelled exemplar corpora +// behind KNN routing (router classifier "knn"). +// +// The corpus FILE is the source of truth: one JSONL file per store +// name under , each line an Entry {text, labels, vector, +// embedding_model}. The local-store vector backend is a pure in-memory +// index rebuilt from the file — it has no persistence of its own (its +// Load is an explicit no-op), and keeping it that way preserves the +// documented swap point for external vector backends. Vectors are +// cached in the file alongside the text so a restart re-indexes +// without re-embedding; entries recorded under a different embedding +// model are re-embedded on load. +// +// Texts in the corpus file never leave the server: Stats exposes label +// counts only, and there is deliberately no API that returns entries. +package corpus + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "sync" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/services/routing/router" + "github.com/mudler/xlog" +) + +// Entry is one labelled exemplar. Vector and EmbeddingModel are the +// embedding cache — absent (or stale) entries are re-embedded when the +// corpus is loaded or added to. +type Entry struct { + Text string `json:"text"` + Labels []string `json:"labels"` + Vector []float32 `json:"vector,omitempty"` + EmbeddingModel string `json:"embedding_model,omitempty"` +} + +// Stats is the inspection surface — counts only, never texts. The +// router corpus API and the Routing tab render this. +type Stats struct { + StoreName string `json:"store_name"` + Total int `json:"total"` + LabelCounts map[string]int `json:"label_counts"` + // EmbeddingModels lists the distinct embedder fingerprints present + // in the file. More than one entry means a re-embed is pending for + // part of the corpus (it happens lazily on the next load). + EmbeddingModels []string `json:"embedding_models,omitempty"` +} + +// batchIndex and deleter are optional fast paths the local-store +// implementation provides; the Manager degrades to per-entry Insert +// (and to "entries persist in the index until restart" on Clear) when +// a store doesn't. +type batchIndex interface { + InsertBatch(ctx context.Context, vecs [][]float32, payloads [][]byte) error +} + +type deleter interface { + Delete(ctx context.Context, vecs [][]float32) error +} + +// Manager owns the corpus files and their sync state into the +// in-memory vector index. One per process, shared by the classifier +// build path (EnsureLoaded) and the corpus API (Add/Stats/Clear). +type Manager struct { + dir string + + mu sync.Mutex + // loadedModel records which embedding model a store name was synced + // into the live index under. Guards double-loading and detects the + // embedding-model-changed-on-a-live-index case, which local-store + // cannot serve (it enforces one key length per store). + loadedModel map[string]string +} + +// NewManager roots corpus files at dir (created lazily on first +// write). +func NewManager(dir string) *Manager { + return &Manager{dir: dir, loadedModel: map[string]string{}} +} + +// path maps a store name to its corpus file. Store names come from +// YAML (store_name) or model names; sanitise so they can't escape the +// corpus dir or collide with path syntax. +func (m *Manager) path(storeName string) string { + safe := strings.Map(func(r rune) rune { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '-', r == '_', r == '.': + return r + } + return '_' + }, storeName) + return filepath.Join(m.dir, safe+".jsonl") +} + +// EnsureLoaded syncs the persisted corpus for storeName into the live +// vector index, once per (store, embedding model) per process. Entries +// recorded under a different embedding model are re-embedded and the +// file rewritten. Returns the number of entries now indexed. A missing +// file is an empty corpus, not an error. +func (m *Manager) EnsureLoaded(ctx context.Context, storeName, embeddingModel string, embedder backend.Embedder, store backend.VectorStore) (int, error) { + m.mu.Lock() + defer m.mu.Unlock() + + if prev, ok := m.loadedModel[storeName]; ok { + if prev == embeddingModel { + return 0, nil + } + // The in-memory index already holds vectors from another + // embedder; local-store enforces a single key length, so mixing + // would fail on insert (or silently corrupt neighbourhoods when + // dimensions happen to match). A restart re-indexes cleanly. + return 0, fmt.Errorf("corpus %q was indexed with embedding model %q this process; restart LocalAI to re-index it with %q", storeName, prev, embeddingModel) + } + + entries, err := m.read(storeName) + if err != nil { + return 0, err + } + if len(entries) == 0 { + m.loadedModel[storeName] = embeddingModel + return 0, nil + } + + dirty := false + for i := range entries { + if len(entries[i].Vector) > 0 && entries[i].EmbeddingModel == embeddingModel { + continue + } + if embedder == nil { + return 0, fmt.Errorf("corpus %q has entries needing (re-)embedding but no embedder is available", storeName) + } + vec, err := embedder.Embed(ctx, entries[i].Text) + if err != nil { + return 0, fmt.Errorf("corpus %q: re-embedding entry %d: %w", storeName, i, err) + } + entries[i].Vector = vec + entries[i].EmbeddingModel = embeddingModel + dirty = true + } + if dirty { + if err := m.write(storeName, entries); err != nil { + return 0, err + } + } + + if err := insertAll(ctx, store, entries); err != nil { + return 0, err + } + m.loadedModel[storeName] = embeddingModel + return len(entries), nil +} + +// Add validates, embeds, persists, and indexes new exemplars. Entries +// whose text is already in the corpus are skipped (an exemplar's +// labels are corrected via Clear + reseed, not silent overwrite). +// Returns (added, skipped). The file write happens before the index +// insert: if indexing fails the entries are still durable and the next +// EnsureLoaded (or restart) syncs them. +func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, embedder backend.Embedder, store backend.VectorStore, entries []Entry) (int, int, error) { + if embedder == nil { + return 0, 0, fmt.Errorf("corpus %q: no embedder available", storeName) + } + for i, e := range entries { + if strings.TrimSpace(e.Text) == "" { + return 0, 0, fmt.Errorf("corpus entry %d: empty text", i) + } + if len(e.Labels) == 0 { + return 0, 0, fmt.Errorf("corpus entry %d: at least one label required", i) + } + } + + m.mu.Lock() + defer m.mu.Unlock() + + existing, err := m.read(storeName) + if err != nil { + return 0, 0, err + } + seen := make(map[string]struct{}, len(existing)) + for _, e := range existing { + seen[e.Text] = struct{}{} + } + + added := make([]Entry, 0, len(entries)) + skipped := 0 + for _, e := range entries { + if _, dup := seen[e.Text]; dup { + skipped++ + continue + } + seen[e.Text] = struct{}{} + vec, err := embedder.Embed(ctx, e.Text) + if err != nil { + return 0, skipped, fmt.Errorf("corpus %q: embedding %q-labelled entry: %w", storeName, e.Labels[0], err) + } + e.Vector = vec + e.EmbeddingModel = embeddingModel + added = append(added, e) + } + if len(added) == 0 { + return 0, skipped, nil + } + + if err := m.write(storeName, append(existing, added...)); err != nil { + return 0, skipped, err + } + if store != nil { + if err := insertAll(ctx, store, added); err != nil { + // Durable but not indexed — routing won't see the new + // entries until the next successful load. Surface loudly. + return len(added), skipped, fmt.Errorf("corpus %q: entries persisted but indexing failed (they will index on next load/restart): %w", storeName, err) + } + } + return len(added), skipped, nil +} + +// Stats reports label counts for the persisted corpus. Never returns +// entry texts. +func (m *Manager) Stats(storeName string) (Stats, error) { + m.mu.Lock() + defer m.mu.Unlock() + + entries, err := m.read(storeName) + if err != nil { + return Stats{}, err + } + s := Stats{StoreName: storeName, Total: len(entries), LabelCounts: map[string]int{}} + models := map[string]struct{}{} + for _, e := range entries { + for _, l := range e.Labels { + s.LabelCounts[l]++ + } + if e.EmbeddingModel != "" { + models[e.EmbeddingModel] = struct{}{} + } + } + for mn := range models { + s.EmbeddingModels = append(s.EmbeddingModels, mn) + } + sort.Strings(s.EmbeddingModels) + return s, nil +} + +// Clear deletes the corpus file and removes its vectors from the live +// index when the store supports deletion. When it doesn't, the index +// keeps serving stale entries until restart — the returned count and +// the warning log make that visible rather than silent. +func (m *Manager) Clear(ctx context.Context, storeName string, store backend.VectorStore) (int, error) { + m.mu.Lock() + defer m.mu.Unlock() + + entries, err := m.read(storeName) + if err != nil { + return 0, err + } + if err := os.Remove(m.path(storeName)); err != nil && !os.IsNotExist(err) { + return 0, err + } + delete(m.loadedModel, storeName) + if len(entries) == 0 || store == nil { + return len(entries), nil + } + if d, ok := store.(deleter); ok { + vecs := make([][]float32, 0, len(entries)) + for _, e := range entries { + if len(e.Vector) > 0 { + vecs = append(vecs, e.Vector) + } + } + if len(vecs) > 0 { + if err := d.Delete(ctx, vecs); err != nil { + return len(entries), fmt.Errorf("corpus %q: file cleared but live index deletion failed (stale entries served until restart): %w", storeName, err) + } + } + } else { + xlog.Warn("corpus: vector store does not support deletion; cleared corpus stays in the live index until restart", + "store", storeName, "entries", len(entries)) + } + return len(entries), nil +} + +func insertAll(ctx context.Context, store backend.VectorStore, entries []Entry) error { + vecs := make([][]float32, 0, len(entries)) + payloads := make([][]byte, 0, len(entries)) + for _, e := range entries { + payload, err := router.EncodeCorpusEntry(e.Labels) + if err != nil { + return err + } + vecs = append(vecs, e.Vector) + payloads = append(payloads, payload) + } + if b, ok := store.(batchIndex); ok { + return b.InsertBatch(ctx, vecs, payloads) + } + for i := range vecs { + if err := store.Insert(ctx, vecs[i], payloads[i]); err != nil { + return err + } + } + return nil +} + +// read returns the persisted entries for storeName; a missing file is +// an empty corpus. Callers hold m.mu. +func (m *Manager) read(storeName string) ([]Entry, error) { + f, err := os.Open(m.path(storeName)) + if os.IsNotExist(err) { + return nil, nil + } + if err != nil { + return nil, err + } + defer func() { _ = f.Close() }() + var entries []Entry + sc := bufio.NewScanner(f) + // Vectors inline in JSON push line length well past the default + // 64KiB scanner cap (a 4096-dim float32 vector is ~50KiB of JSON + // alone). 16MiB bounds any realistic embedding width. + sc.Buffer(make([]byte, 0, 1<<20), 16<<20) + line := 0 + for sc.Scan() { + line++ + raw := strings.TrimSpace(sc.Text()) + if raw == "" { + continue + } + var e Entry + if err := json.Unmarshal([]byte(raw), &e); err != nil { + return nil, fmt.Errorf("corpus file %s line %d: %w", m.path(storeName), line, err) + } + entries = append(entries, e) + } + if err := sc.Err(); err != nil { + return nil, err + } + return entries, nil +} + +// write atomically replaces the corpus file (tmp + rename) so a crash +// mid-write can't truncate the corpus. Callers hold m.mu. +func (m *Manager) write(storeName string, entries []Entry) error { + if err := os.MkdirAll(m.dir, 0o750); err != nil { + return err + } + target := m.path(storeName) + tmp, err := os.CreateTemp(m.dir, filepath.Base(target)+".tmp-*") + if err != nil { + return err + } + defer func() { _ = os.Remove(tmp.Name()) }() + w := bufio.NewWriter(tmp) + enc := json.NewEncoder(w) + for _, e := range entries { + if err := enc.Encode(e); err != nil { + _ = tmp.Close() + return err + } + } + if err := w.Flush(); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + return os.Rename(tmp.Name(), target) +} diff --git a/core/services/routing/corpus/manager_test.go b/core/services/routing/corpus/manager_test.go new file mode 100644 index 000000000000..9c69581bd65b --- /dev/null +++ b/core/services/routing/corpus/manager_test.go @@ -0,0 +1,240 @@ +package corpus_test + +import ( + "context" + "os" + "path/filepath" + "strings" + "sync" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/services/routing/corpus" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// countingEmbedder returns a deterministic vector per (model, text) +// and counts calls, so specs can assert when re-embedding happened vs +// when the cached vectors were reused. +type countingEmbedder struct { + mu sync.Mutex + model float32 // baked into the vector so specs can tell models apart + calls int +} + +func (e *countingEmbedder) Embed(_ context.Context, text string) ([]float32, error) { + e.mu.Lock() + defer e.mu.Unlock() + e.calls += 1 + return []float32{float32(len(text)), e.model}, nil +} + +// capturingStore records index mutations. Search/SearchK are +// irrelevant to the manager and return clean misses. +type capturingStore struct { + mu sync.Mutex + payloads [][]byte + batches int + deleted [][]float32 +} + +func (s *capturingStore) Search(_ context.Context, _ []float32) (float64, []byte, bool, error) { + return 0, nil, false, nil +} + +func (s *capturingStore) SearchK(_ context.Context, _ []float32, _ int) ([]backend.Neighbor, error) { + return nil, nil +} + +func (s *capturingStore) Insert(_ context.Context, _ []float32, payload []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + s.payloads = append(s.payloads, payload) + return nil +} + +func (s *capturingStore) InsertBatch(_ context.Context, vecs [][]float32, payloads [][]byte) error { + s.mu.Lock() + defer s.mu.Unlock() + s.batches++ + s.payloads = append(s.payloads, payloads...) + _ = vecs + return nil +} + +func (s *capturingStore) Delete(_ context.Context, vecs [][]float32) error { + s.mu.Lock() + defer s.mu.Unlock() + s.deleted = append(s.deleted, vecs...) + return nil +} + +var _ = Describe("corpus.Manager", func() { + var ( + dir string + mgr *corpus.Manager + embedder *countingEmbedder + store *capturingStore + ctx context.Context + ) + + const storeName = "router-corpus-smart-router" + + seed := []corpus.Entry{ + {Text: "debug this Go null pointer", Labels: []string{"code-generation"}}, + {Text: "what is 12 * 42?", Labels: []string{"math-reasoning"}}, + {Text: "refactor this and explain the math", Labels: []string{"code-generation", "math-reasoning"}}, + } + + BeforeEach(func() { + d, err := os.MkdirTemp("", "corpus-test-*") + Expect(err).NotTo(HaveOccurred()) + dir = d + mgr = corpus.NewManager(dir) + embedder = &countingEmbedder{model: 1} + store = &capturingStore{} + ctx = context.Background() + }) + + AfterEach(func() { + _ = os.RemoveAll(dir) + }) + + It("adds entries: embeds, persists, and indexes them", func() { + added, skipped, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + Expect(added).To(Equal(3)) + Expect(skipped).To(BeZero()) + Expect(embedder.calls).To(Equal(3)) + Expect(store.payloads).To(HaveLen(3)) + Expect(store.batches).To(Equal(1), "should use the batch fast path") + + // Payloads are the label sets the classifier votes over. + Expect(string(store.payloads[0])).To(ContainSubstring("code-generation")) + + // Persisted on disk under a sanitised name. + _, err = os.Stat(filepath.Join(dir, storeName+".jsonl")) + Expect(err).NotTo(HaveOccurred()) + }) + + It("skips duplicate texts instead of double-weighting them", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + added, skipped, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed[:2]) + Expect(err).NotTo(HaveOccurred()) + Expect(added).To(BeZero()) + Expect(skipped).To(Equal(2)) + }) + + It("rejects empty text and label-less entries", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, []corpus.Entry{{Text: " ", Labels: []string{"x"}}}) + Expect(err).To(HaveOccurred()) + _, _, err = mgr.Add(ctx, storeName, "embed-1", embedder, store, []corpus.Entry{{Text: "hello", Labels: nil}}) + Expect(err).To(HaveOccurred()) + }) + + It("reloads a persisted corpus into a fresh index without re-embedding", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + + // Simulate restart: fresh manager over the same dir, empty index. + mgr2 := corpus.NewManager(dir) + store2 := &capturingStore{} + embedder2 := &countingEmbedder{model: 1} + n, err := mgr2.EnsureLoaded(ctx, storeName, "embed-1", embedder2, store2) + Expect(err).NotTo(HaveOccurred()) + Expect(n).To(Equal(3)) + Expect(store2.payloads).To(HaveLen(3)) + Expect(embedder2.calls).To(BeZero(), "cached vectors must be reused") + + // Second call is a no-op — already synced this process. + n, err = mgr2.EnsureLoaded(ctx, storeName, "embed-1", embedder2, store2) + Expect(err).NotTo(HaveOccurred()) + Expect(n).To(BeZero()) + }) + + It("re-embeds entries recorded under a different embedding model", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + + mgr2 := corpus.NewManager(dir) + newEmbedder := &countingEmbedder{model: 2} + store2 := &capturingStore{} + n, err := mgr2.EnsureLoaded(ctx, storeName, "embed-2", newEmbedder, store2) + Expect(err).NotTo(HaveOccurred()) + Expect(n).To(Equal(3)) + Expect(newEmbedder.calls).To(Equal(3), "fingerprint mismatch must re-embed") + + // The rewrite is durable: a third manager loads under embed-2 + // without touching the embedder again. + mgr3 := corpus.NewManager(dir) + embedder3 := &countingEmbedder{model: 2} + n, err = mgr3.EnsureLoaded(ctx, storeName, "embed-2", embedder3, &capturingStore{}) + Expect(err).NotTo(HaveOccurred()) + Expect(n).To(Equal(3)) + Expect(embedder3.calls).To(BeZero()) + }) + + It("refuses to mix embedding models in a live index", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + _, err = mgr.EnsureLoaded(ctx, storeName, "embed-1", embedder, store) + Expect(err).NotTo(HaveOccurred()) + + _, err = mgr.EnsureLoaded(ctx, storeName, "embed-2", &countingEmbedder{model: 2}, store) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("restart")) + }) + + It("reports label counts and never texts", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + st, err := mgr.Stats(storeName) + Expect(err).NotTo(HaveOccurred()) + Expect(st.Total).To(Equal(3)) + Expect(st.LabelCounts).To(Equal(map[string]int{ + "code-generation": 2, + "math-reasoning": 2, + })) + Expect(st.EmbeddingModels).To(Equal([]string{"embed-1"})) + }) + + It("clears the file and the live index", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + + n, err := mgr.Clear(ctx, storeName, store) + Expect(err).NotTo(HaveOccurred()) + Expect(n).To(Equal(3)) + Expect(store.deleted).To(HaveLen(3)) + + st, err := mgr.Stats(storeName) + Expect(err).NotTo(HaveOccurred()) + Expect(st.Total).To(BeZero()) + + // And a load after clear indexes nothing. + loaded, err := corpus.NewManager(dir).EnsureLoaded(ctx, storeName, "embed-1", embedder, &capturingStore{}) + Expect(err).NotTo(HaveOccurred()) + Expect(loaded).To(BeZero()) + }) + + It("treats a missing file as an empty corpus", func() { + n, err := mgr.EnsureLoaded(ctx, "never-seeded", "embed-1", embedder, store) + Expect(err).NotTo(HaveOccurred()) + Expect(n).To(BeZero()) + st, err := mgr.Stats("never-seeded") + Expect(err).NotTo(HaveOccurred()) + Expect(st.Total).To(BeZero()) + }) + + It("sanitises hostile store names into the corpus dir", func() { + hostile := "../../etc/passwd" + _, _, err := mgr.Add(ctx, hostile, "embed-1", embedder, store, seed[:1]) + Expect(err).NotTo(HaveOccurred()) + entries, err := os.ReadDir(dir) + Expect(err).NotTo(HaveOccurred()) + Expect(entries).To(HaveLen(1)) + Expect(strings.Contains(entries[0].Name(), "/")).To(BeFalse()) + }) +}) diff --git a/core/services/routing/router/decisions.go b/core/services/routing/router/decisions.go index d446ac29a63e..23f512813720 100644 --- a/core/services/routing/router/decisions.go +++ b/core/services/routing/router/decisions.go @@ -11,18 +11,19 @@ import ( // Prompt is NEVER stored — admins audit by Hash if they need to // dedupe recurring routing patterns. type DecisionRecord struct { - ID string `json:"id"` - CorrelationID string `json:"correlation_id"` - UserID string `json:"user_id"` - RouterModel string `json:"router_model"` // The smart-router model name the client asked for. - RequestedModel string `json:"requested_model"`// Same as RouterModel for now; reserved for chained routers. - ServedModel string `json:"served_model"` // The candidate the classifier picked. - Classifier string `json:"classifier"` // Classifier.Name(), e.g. "score". - Label string `json:"label"` - Score float64 `json:"score"` - LatencyMs int64 `json:"latency_ms"` - Cached bool `json:"cached"` // True when the decision came from the L2 embedding cache. - CacheSimilarity float64 `json:"cache_similarity,omitempty"` // Cosine similarity of the cache hit, 0 when not cached. + ID string `json:"id"` + CorrelationID string `json:"correlation_id"` + UserID string `json:"user_id"` + RouterModel string `json:"router_model"` // The smart-router model name the client asked for. + RequestedModel string `json:"requested_model"` // Same as RouterModel for now; reserved for chained routers. + ServedModel string `json:"served_model"` // The candidate the classifier picked. + Classifier string `json:"classifier"` // Classifier.Name(), e.g. "score". + Label string `json:"label"` + Score float64 `json:"score"` + LatencyMs int64 `json:"latency_ms"` + Cached bool `json:"cached"` // True when the decision came from the L2 embedding cache. + CacheSimilarity float64 `json:"cache_similarity,omitempty"` // Cosine similarity of the cache hit, 0 when not cached. + NearestSimilarity float64 `json:"nearest_similarity,omitempty"` // KNN classifier: similarity of the closest corpus entry, set even on fallback decisions. 0 for other classifiers. // LabelScores carries the full per-label score distribution so the // admin UI can show how close inactive labels got to the activation // threshold. Empty on cache hits (only the final label set is cached). @@ -32,8 +33,8 @@ type DecisionRecord struct { // the admin page can split realtime / chat / anthropic streams. Empty // string is treated as "chat" for backward compatibility with rows // written before the field existed. - Source string `json:"source,omitempty"` - CreatedAt time.Time `json:"created_at"` + Source string `json:"source,omitempty"` + CreatedAt time.Time `json:"created_at"` } // Source values for DecisionRecord.Source. Kept as constants so callers diff --git a/core/services/routing/router/embedding_cache_test.go b/core/services/routing/router/embedding_cache_test.go index e36b049c3d87..41be408a00f7 100644 --- a/core/services/routing/router/embedding_cache_test.go +++ b/core/services/routing/router/embedding_cache_test.go @@ -9,6 +9,7 @@ import ( "sync" "time" + "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/services/routing/router" . "github.com/onsi/ginkgo/v2" @@ -88,6 +89,31 @@ func (s *memVectorStore) Search(_ context.Context, vec []float32) (float64, []by return 0, nil, false, nil } +// SearchK reuses the same synthetic metric as Search: 1.0 exact, 0.8 +// leading-element, 0.0 otherwise — zero-similarity entries are omitted. +func (s *memVectorStore) SearchK(_ context.Context, vec []float32, k int) ([]backend.Neighbor, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.failOps > 0 { + s.failOps-- + return nil, errors.New("store offline") + } + var exact, close []backend.Neighbor + for _, e := range s.entries { + switch { + case vecEqual(e.vec, vec): + exact = append(exact, backend.Neighbor{Similarity: 1.0, Payload: e.payload}) + case len(vec) > 0 && len(e.vec) > 0 && vec[0] == e.vec[0]: + close = append(close, backend.Neighbor{Similarity: 0.80, Payload: e.payload}) + } + } + neighbors := append(exact, close...) + if len(neighbors) > k { + neighbors = neighbors[:k] + } + return neighbors, nil +} + func (s *memVectorStore) Insert(_ context.Context, vec []float32, payload []byte) error { s.mu.Lock() defer s.mu.Unlock() diff --git a/core/services/routing/router/knn.go b/core/services/routing/router/knn.go new file mode 100644 index 000000000000..ddde30f43679 --- /dev/null +++ b/core/services/routing/router/knn.go @@ -0,0 +1,222 @@ +package router + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "time" + + "github.com/mudler/LocalAI/core/backend" +) + +// KNNClassifier routes by nearest-neighbour vote over a curated, +// labelled corpus of example prompts. It is the first-class form of +// what EmbeddingCacheClassifier does opportunistically: instead of +// caching another classifier's decisions, the corpus is seeded and +// curated explicitly (via the router corpus API), each entry carrying +// the policy labels a matching prompt should activate. +// +// Classify embeds the probe, fetches the K nearest corpus entries, and +// activates every label whose similarity-weighted vote share clears +// VoteThreshold. Neighbours below SimilarityThreshold are discarded +// first — that cutoff is the epistemic gate: a probe dissimilar from +// *all* labelled experience is undecidable by construction, so the +// classifier returns an empty label set and the middleware falls back +// to cfg.Router.Fallback (the assumed-best model) rather than guessing. +// +// The classifier never inserts into the corpus on its own. Routing +// outcomes only become corpus entries through explicit curation — a +// mislabelled exemplar poisons every future neighbourhood around it, +// so growth is an admin decision, not a side effect. +type KNNClassifier struct { + embedder backend.Embedder + store backend.VectorStore + k int + similarityThreshold float64 + voteThreshold float64 + + // budget trims the conversation to the embedder model's own context + // before embedding; nil embeds Probe.Prompt as built by the caller. + budget *lazyBudget +} + +// Defaults. K=3 keeps a weighted majority meaningful on small corpora +// while tolerating one mislabelled neighbour; the similarity default +// matches the embedding cache so one threshold intuition serves both. +// VoteThreshold 0.5 means a label activates on a weighted majority — +// with K=1 this degenerates to "the nearest entry's labels", the same +// contract the embedding cache implements. +const ( + defaultKNNK = 3 + defaultKNNSimilarity = 0.80 + defaultKNNVote = 0.5 +) + +// KNNClassifierOptions carries the tunables; zero values pick the +// package defaults above. +type KNNClassifierOptions struct { + K int + SimilarityThreshold float64 + VoteThreshold float64 +} + +// NewKNNClassifier builds a KNN classifier over the given embedder and +// vector store. Panics on nil embedder/store — same fail-fast posture +// as the other classifiers; buildClassifier validates config before +// construction. +func NewKNNClassifier(embedder backend.Embedder, store backend.VectorStore, opts KNNClassifierOptions) *KNNClassifier { + if embedder == nil { + panic("router/knn: embedder is required") + } + if store == nil { + panic("router/knn: vector store is required") + } + if opts.K <= 0 { + opts.K = defaultKNNK + } + if opts.SimilarityThreshold <= 0 { + opts.SimilarityThreshold = defaultKNNSimilarity + } + if opts.VoteThreshold <= 0 { + opts.VoteThreshold = defaultKNNVote + } + return &KNNClassifier{ + embedder: embedder, + store: store, + k: opts.K, + similarityThreshold: opts.SimilarityThreshold, + voteThreshold: opts.VoteThreshold, + } +} + +// WithTokenTrim wires the embedder model's own tokenizer and context so +// the probe embeds the most recent turns that fit instead of a +// caller-chosen size. nil tokenizer / non-positive context leaves +// trimming off. Returns the receiver for chaining at construction. +func (c *KNNClassifier) WithTokenTrim(tokenize func(string) (int, error), maxContextTokens int) *KNNClassifier { + c.budget = &lazyBudget{tokenize: tokenize, maxContext: maxContextTokens} + return c +} + +func (c *KNNClassifier) Name() string { return ClassifierKNN } + +func (c *KNNClassifier) Classify(ctx context.Context, p Probe) (Decision, error) { + start := time.Now() + + vec, err := c.embedder.Embed(ctx, trimmedProbeText(p, c.budget, identityRender)) + if err != nil { + return errDecision(start, fmt.Errorf("knn classifier embed: %w", err)) + } + neighbors, err := c.store.SearchK(ctx, vec, c.k) + if err != nil { + return errDecision(start, fmt.Errorf("knn classifier search: %w", err)) + } + + // Epistemic gate: only neighbours the probe is genuinely close to + // may vote. Keeping sub-threshold neighbours out of the vote (rather + // than merely gating on the best one) stops far-away corpus regions + // from diluting a clear local majority. + best := 0.0 + usable := neighbors[:0] + for _, n := range neighbors { + if n.Similarity > best { + best = n.Similarity + } + if n.Similarity >= c.similarityThreshold { + usable = append(usable, n) + } + } + if len(usable) == 0 { + // Out of corpus range — empty label set routes to the fallback + // via MatchCandidate's empty-active-set contract. Surfacing the + // best similarity in the decision log tells the admin whether + // the corpus needs entries near this probe or the threshold is + // simply too tight. + return Decision{ + NearestSimilarity: best, + ActivationThreshold: c.voteThreshold, + Latency: time.Since(start), + }, nil + } + + votes := map[string]float64{} + total := 0.0 + for _, n := range usable { + entry, ok := decodeCorpusEntry(n.Payload) + if !ok { + // A corrupt payload can't vote; it still counted toward K. + continue + } + total += n.Similarity + for _, l := range entry.Labels { + votes[l] += n.Similarity + } + } + if total == 0 { + return Decision{ + NearestSimilarity: best, + ActivationThreshold: c.voteThreshold, + Latency: time.Since(start), + }, nil + } + + // Vote shares in descending order; ties broken lexicographically so + // the decision log is deterministic. + labels := make([]string, 0, len(votes)) + for l := range votes { + labels = append(labels, l) + } + sort.Slice(labels, func(i, j int) bool { + if votes[labels[i]] != votes[labels[j]] { + return votes[labels[i]] > votes[labels[j]] + } + return labels[i] < labels[j] + }) + scores := make([]float64, len(labels)) + active := []string{} + for i, l := range labels { + scores[i] = votes[l] / total + if scores[i] >= c.voteThreshold { + active = append(active, l) + } + } + + d := Decision{ + Labels: active, + ActivationThreshold: c.voteThreshold, + LabelScores: NewLabelScores(labels, scores), + NearestSimilarity: best, + Latency: time.Since(start), + } + if len(active) > 0 { + d.Score = votes[active[0]] / total + } + return d, nil +} + +// corpusEntry is the stored shape of one labelled exemplar. Kept +// deliberately minimal: the vector key lives in the store, the text +// lives only in the corpus file (never returned by inspection APIs), +// so the store payload is just the label set. +type corpusEntry struct { + Labels []string `json:"labels"` +} + +// EncodeCorpusEntry serialises the labels of one corpus exemplar into +// the vector-store payload shape Classify votes over. Exported for the +// corpus loader/API in core, which owns insertion. +func EncodeCorpusEntry(labels []string) ([]byte, error) { + if len(labels) == 0 { + return nil, fmt.Errorf("corpus entry needs at least one label") + } + return json.Marshal(corpusEntry{Labels: labels}) +} + +func decodeCorpusEntry(b []byte) (corpusEntry, bool) { + var e corpusEntry + if err := json.Unmarshal(b, &e); err != nil || len(e.Labels) == 0 { + return corpusEntry{}, false + } + return e, true +} diff --git a/core/services/routing/router/knn_test.go b/core/services/routing/router/knn_test.go new file mode 100644 index 000000000000..49bfd9a4c1e6 --- /dev/null +++ b/core/services/routing/router/knn_test.go @@ -0,0 +1,174 @@ +package router_test + +import ( + "context" + "errors" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/services/routing/router" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// scriptedKNNStore returns a fixed neighbour list from SearchK, +// letting tests exercise the vote/gate math without a real store. +type scriptedKNNStore struct { + neighbors []backend.Neighbor + err error + lastK int +} + +func (s *scriptedKNNStore) SearchK(_ context.Context, _ []float32, k int) ([]backend.Neighbor, error) { + s.lastK = k + if s.err != nil { + return nil, s.err + } + if len(s.neighbors) > k { + return s.neighbors[:k], s.err + } + return s.neighbors, s.err +} + +func (s *scriptedKNNStore) Search(_ context.Context, _ []float32) (float64, []byte, bool, error) { + return 0, nil, false, errors.New("knn classifier must use SearchK") +} + +func (s *scriptedKNNStore) Insert(_ context.Context, _ []float32, _ []byte) error { + return errors.New("knn classifier must never insert") +} + +func mustEntry(labels ...string) []byte { + b, err := router.EncodeCorpusEntry(labels) + Expect(err).ToNot(HaveOccurred()) + return b +} + +var _ = Describe("KNNClassifier", func() { + var ( + embedder *fakeEmbedder + probe router.Probe + ctx context.Context + ) + + BeforeEach(func() { + ctx = context.Background() + embedder = &fakeEmbedder{table: map[string][]float32{"prompt": {1, 0, 0}}} + probe = router.Probe{Prompt: "prompt"} + }) + + classify := func(store *scriptedKNNStore, opts router.KNNClassifierOptions) router.Decision { + c := router.NewKNNClassifier(embedder, store, opts) + d, err := c.Classify(ctx, probe) + Expect(err).ToNot(HaveOccurred()) + return d + } + + It("computes similarity-weighted vote shares exactly", func() { + // Hand-computed: usable sims 0.90 {code}, 0.85 {code,math}, + // 0.82 {math}; total = 2.57. + // share(code) = (0.90+0.85)/2.57 = 0.68093... + // share(math) = (0.85+0.82)/2.57 = 0.64980... + // Both clear the 0.5 majority → both active; candidate matching + // then requires a model labelled for both. + store := &scriptedKNNStore{neighbors: []backend.Neighbor{ + {Similarity: 0.90, Payload: mustEntry("code")}, + {Similarity: 0.85, Payload: mustEntry("code", "math")}, + {Similarity: 0.82, Payload: mustEntry("math")}, + }} + d := classify(store, router.KNNClassifierOptions{K: 3}) + Expect(d.Labels).To(Equal([]string{"code", "math"})) + Expect(d.LabelScores).To(HaveLen(2)) + Expect(d.LabelScores[0].Label).To(Equal("code")) + Expect(d.LabelScores[0].Score).To(BeNumerically("~", 1.75/2.57, 1e-9)) + Expect(d.LabelScores[1].Label).To(Equal("math")) + Expect(d.LabelScores[1].Score).To(BeNumerically("~", 1.67/2.57, 1e-9)) + Expect(d.Score).To(BeNumerically("~", 1.75/2.57, 1e-9)) + Expect(d.NearestSimilarity).To(BeNumerically("~", 0.90, 1e-9)) + Expect(store.lastK).To(Equal(3)) + }) + + It("does not activate a minority label", func() { + // share(chat) = 0.81/2.57 < 0.5 → inactive, but still reported + // in LabelScores so the decision log shows how close it came. + store := &scriptedKNNStore{neighbors: []backend.Neighbor{ + {Similarity: 0.90, Payload: mustEntry("code")}, + {Similarity: 0.86, Payload: mustEntry("code")}, + {Similarity: 0.81, Payload: mustEntry("chat")}, + }} + d := classify(store, router.KNNClassifierOptions{K: 3}) + Expect(d.Labels).To(Equal([]string{"code"})) + Expect(d.LabelScores).To(HaveLen(2)) + Expect(d.LabelScores[1].Label).To(Equal("chat")) + Expect(d.LabelScores[1].Score).To(BeNumerically("<", 0.5)) + }) + + It("gates out-of-corpus probes to the fallback (empty labels)", func() { + store := &scriptedKNNStore{neighbors: []backend.Neighbor{ + {Similarity: 0.55, Payload: mustEntry("code")}, + {Similarity: 0.40, Payload: mustEntry("math")}, + }} + d := classify(store, router.KNNClassifierOptions{SimilarityThreshold: 0.80}) + Expect(d.Labels).To(BeEmpty()) + // The admin-facing epistemic signal: how far away the nearest + // labelled experience was. + Expect(d.NearestSimilarity).To(BeNumerically("~", 0.55, 1e-9)) + }) + + It("excludes sub-threshold neighbours from the vote", func() { + // The 0.3-sim {chat} neighbour must not dilute the local + // majority: with it, share(code) would be 0.9/1.2 = 0.75; the + // vote must instead be over the single usable neighbour. + store := &scriptedKNNStore{neighbors: []backend.Neighbor{ + {Similarity: 0.90, Payload: mustEntry("code")}, + {Similarity: 0.30, Payload: mustEntry("chat")}, + }} + d := classify(store, router.KNNClassifierOptions{K: 2, SimilarityThreshold: 0.80}) + Expect(d.Labels).To(Equal([]string{"code"})) + Expect(d.LabelScores).To(HaveLen(1)) + Expect(d.LabelScores[0].Score).To(BeNumerically("~", 1.0, 1e-9)) + }) + + It("degenerates to nearest-entry labels at K=1", func() { + store := &scriptedKNNStore{neighbors: []backend.Neighbor{ + {Similarity: 0.95, Payload: mustEntry("reasoning", "math")}, + }} + d := classify(store, router.KNNClassifierOptions{K: 1}) + Expect(d.Labels).To(Equal([]string{"math", "reasoning"})) + Expect(d.Score).To(BeNumerically("~", 1.0, 1e-9)) + }) + + It("skips corrupt payloads and falls back when nothing can vote", func() { + store := &scriptedKNNStore{neighbors: []backend.Neighbor{ + {Similarity: 0.90, Payload: []byte("not json")}, + }} + d := classify(store, router.KNNClassifierOptions{}) + Expect(d.Labels).To(BeEmpty()) + Expect(d.NearestSimilarity).To(BeNumerically("~", 0.90, 1e-9)) + }) + + It("returns the embed error so the middleware can fall back", func() { + embedder.table = nil // unknown prompt → fakeEmbedder errors + store := &scriptedKNNStore{} + c := router.NewKNNClassifier(embedder, store, router.KNNClassifierOptions{}) + _, err := c.Classify(ctx, probe) + Expect(err).To(HaveOccurred()) + }) + + It("returns the store error so the middleware can fall back", func() { + store := &scriptedKNNStore{err: errors.New("store offline")} + c := router.NewKNNClassifier(embedder, store, router.KNNClassifierOptions{}) + _, err := c.Classify(ctx, probe) + Expect(err).To(HaveOccurred()) + }) + + It("rejects corpus entries without labels at encode time", func() { + _, err := router.EncodeCorpusEntry(nil) + Expect(err).To(HaveOccurred()) + }) + + It("panics on missing embedder or store", func() { + Expect(func() { router.NewKNNClassifier(nil, &scriptedKNNStore{}, router.KNNClassifierOptions{}) }).To(Panic()) + Expect(func() { router.NewKNNClassifier(embedder, nil, router.KNNClassifierOptions{}) }).To(Panic()) + }) +}) diff --git a/core/services/routing/router/resolve.go b/core/services/routing/router/resolve.go index a474d6d4854f..486eedfdea5d 100644 --- a/core/services/routing/router/resolve.go +++ b/core/services/routing/router/resolve.go @@ -165,6 +165,7 @@ func (r *ResolveResult) ToDecisionRecord(id, correlationID, userID, source strin LatencyMs: r.Decision.Latency.Milliseconds(), Cached: r.Decision.Cached, CacheSimilarity: r.Decision.CacheSimilarity, + NearestSimilarity: r.Decision.NearestSimilarity, LabelScores: r.Decision.LabelScores, ActivationThreshold: r.Decision.ActivationThreshold, Source: source, diff --git a/core/services/routing/router/types.go b/core/services/routing/router/types.go index 05efdc3497fa..8e4d8f1da086 100644 --- a/core/services/routing/router/types.go +++ b/core/services/routing/router/types.go @@ -71,6 +71,13 @@ type Decision struct { // the cosine similarity of the cache hit (0 when not cached). Cached bool `json:"cached,omitempty"` CacheSimilarity float64 `json:"cache_similarity,omitempty"` + + // NearestSimilarity is the cosine similarity of the closest corpus + // entry the KNN classifier saw — set even when the decision fell + // through to the fallback because the probe was out of corpus range, + // which is exactly when an admin wants to know how far off the + // nearest labelled experience was. 0 for other classifiers. + NearestSimilarity float64 `json:"nearest_similarity,omitempty"` } // LabelScore is one entry in Decision.LabelScores — a policy label and @@ -127,6 +134,13 @@ const ( // `type:` field on that model's YAML controls which Reranker // library mode loads. See router/rerank.go. ClassifierColbert = "colbert" + + // ClassifierKNN picks labels by similarity-weighted vote over a + // curated corpus of labelled example prompts, with an epistemic + // gate: probes dissimilar from all corpus entries activate no + // labels and route to the fallback. Needs an embedding model and + // a seeded corpus, not a classifier_model. See router/knn.go. + ClassifierKNN = "knn" ) // LabelFallback is the synthetic label written to the decision diff --git a/docs/content/features/middleware.md b/docs/content/features/middleware.md index 397af3c9207d..c3733891c078 100644 --- a/docs/content/features/middleware.md +++ b/docs/content/features/middleware.md @@ -334,17 +334,20 @@ silent-bypass. ### Available classifiers -LocalAI ships two classifier implementations. Pick one with `classifier:` +LocalAI ships three classifier implementations. Pick one with `classifier:` in the router YAML: | Classifier | When to use | Underlying primitive | |---|---|---| | `score` (default) | Small classifier-tuned LM (Arch-Router-style). Best when label vocabulary is well-covered by next-token continuation. | `Score` gRPC primitive (llama-cpp, vLLM). | | `colbert` | When label descriptions are abstract or short and a next-token classifier produces flat distributions. Robust on long-form policy descriptions. | rerankers backend in ColBERT mode (e.g. `bge-m3-colbert` from the gallery). | +| `knn` | When you have (or can generate) labelled example prompts — including outcome-labelled production traffic. Deterministic, auditable, cheapest per request, and the only classifier with an explicit out-of-distribution fallback. | embeddings backend + local-store KNN over a persisted, curated corpus. | -Both classifiers share the same YAML shape: `classifier_model`, -`policies`, `candidates`, `fallback`, `activation_threshold`, -`classifier_cache_size`, and the optional `embedding_cache` block. +All three share `policies`, `candidates`, `fallback`, and +`classifier_cache_size`. `score` and `colbert` take a +`classifier_model` (+ `activation_threshold`, optional +`embedding_cache`); `knn` instead takes a `knn:` block and a corpus +seeded through the API. ### The Score classifier @@ -425,6 +428,109 @@ underlying scoring head loads — `colbert` for late-interaction MaxSim, `cross-encoder` for cross-attention scoring. The classifier itself is indifferent; pick the head that fits your latency / quality budget. +### The KNN classifier + +The `knn` classifier routes by **similarity-weighted vote over a +curated corpus of labelled example prompts**. Where `score` and +`colbert` ask a model's opinion per request, `knn` consults recorded +experience: each corpus entry is an example prompt plus the policy +labels it should activate. It needs no classifier model — just an +embedding model and a seeded corpus. + +```yaml +router: + classifier: knn + fallback: gpt-4o-proxy # used whenever the prompt is unlike all corpus entries + knn: + embedding_model: nomic-embed-text-v1.5 + k: 3 # neighbours that vote (default 3) + similarity_threshold: 0.80 # the epistemic gate (default 0.80) + vote_threshold: 0.5 # weighted vote share a label needs (default 0.5) + # store_name: router-corpus-smart-router # default "router-corpus-" + policies: + - label: code-generation + description: writing or debugging code + - label: casual-chat + description: small talk + candidates: + - model: qwen3-0.6b + labels: [casual-chat] + - model: qwen-coder + labels: [code-generation, casual-chat] +``` + +For each request: + +1. Embed the prompt with `knn.embedding_model`. +2. Fetch the `k` nearest corpus entries (cosine similarity). +3. **Epistemic gate**: entries below `similarity_threshold` cannot + vote. If none clears it, the classifier activates **no** labels and + the router uses `fallback` — a prompt unlike all labelled + experience is treated as *undecidable*, not guessed. The decision + log records `nearest_similarity` so you can see how far away the + closest labelled example was. +4. Each surviving neighbour votes for its labels, weighted by its + similarity; every label whose vote share clears `vote_threshold` + joins the active set. Candidate matching then proceeds exactly as + for the other classifiers. With `k: 1` this degenerates to + "nearest example's labels". + +#### Seeding and curating the corpus (API-only) + +Corpus entries may contain example user content, so they are managed +exclusively through the admin API — the UI never sends or displays +them, and no endpoint returns entry texts (inspection is label counts +only): + +```bash +# Seed labelled exemplars (embedded server-side; indexed immediately) +curl -X POST http://localhost:8080/api/router/smart-router/corpus \ + -H "Content-Type: application/json" \ + -d '{"entries": [ + {"text": "why does this segfault when I free the buffer twice", "labels": ["code-generation"]}, + {"text": "hey hows it going", "labels": ["casual-chat"]} + ]}' + +# Inspect — counts only, never texts +curl http://localhost:8080/api/router/smart-router/corpus/stats + +# Wipe (file + live index); reseed afterwards +curl -X DELETE http://localhost:8080/api/router/smart-router/corpus +``` + +Entry labels must be declared in `policies` (same invariant as +candidate labels), and duplicate texts are skipped rather than +double-weighted. Label your exemplars with *outcomes*, not topics, +when routing for difficulty: an entry recording "the small model +handled prompts like this" is exactly as useful as one recording that +it failed — grade a sample of production traffic against your +candidates and seed both. + +#### Persistence + +The corpus is persisted as one JSONL file per router under +`/router-corpus/` (text, labels, vector, embedding-model +fingerprint) — **the file is the source of truth** and survives +restarts; the local-store index is rebuilt from it at classifier build +time without re-embedding. Changing `knn.embedding_model` re-embeds +the corpus on the next load (restart LocalAI if the old index was +already live — mixed embedding spaces cannot be served). + +#### Tuning notes + +- **`similarity_threshold` is the safety knob.** Too low and the + router confidently extrapolates from unrelated exemplars; too high + and everything falls back. Watch `nearest_similarity` in the + decision log: fallback rows clustering just under the threshold mean + the corpus needs entries near that traffic (or the gate is too + tight). +- **`k` trades robustness for corpus density**: `k: 3` tolerates one + mislabelled neighbour; raise it only when every label region has + several exemplars. +- **`embedding_cache` is ignored** for `knn` (with a warning) — the + classifier is already an embedding KNN lookup; wrapping it in + another would embed twice for no additional information. + ### YAML reference ```yaml @@ -432,8 +538,9 @@ name: smart-router known_usecases: - chat router: - # `score` (Arch-Router-style next-token scoring) or `colbert` - # (rerank policy descriptions). See "Available classifiers" above. + # `score` (Arch-Router-style next-token scoring), `colbert` (rerank + # policy descriptions), or `knn` (vote over a labelled corpus). + # See "Available classifiers" above. classifier: score # A model loaded by LocalAI that supports the Score gRPC primitive @@ -553,8 +660,11 @@ For each request: The local-store collection is named `router-cache-` by default — each router gets its own collection so two routers can't -cross-contaminate. Collections persist on disk (local-store is the -canonical persistent vector backend), so the cache survives restarts. +cross-contaminate. The collection is **in-memory only**: local-store +keeps no on-disk artefact, so the embedding cache starts empty on every +restart and re-learns from live traffic. (The KNN classifier's corpus +does NOT have this limitation — its corpus file is the source of truth +and re-indexes on startup; see "The KNN classifier" above.) #### Tuning notes @@ -595,6 +705,9 @@ canonical usage log lives in `/api/usage` and correlates by request ID. |---|---|---|---| | GET | `/api/router/status` | any | Router configuration: each router model's classifier, policies, candidates. | | GET | `/api/router/decisions` | admin | Decision log with optional filters (`correlation_id`, `user_id`, `router_model`, `limit`). | +| POST | `/api/router/{name}/corpus` | admin | Seed the KNN corpus with labelled exemplars: `{"entries": [{"text": "...", "labels": ["..."]}]}`. Embedded server-side, persisted, indexed immediately. | +| GET | `/api/router/{name}/corpus/stats` | admin | KNN corpus size and per-label counts. Counts only — entry texts are never returned. | +| DELETE | `/api/router/{name}/corpus` | admin | Wipe the KNN corpus (file + live index). | | POST | `/api/score` | admin | Direct access to the `Score` gRPC primitive — useful for offline threshold tuning. Body: `{"model": "", "prompt": "", "candidates": ["label-a", ...], "length_normalize": true}`. The llama-cpp and vLLM backends implement Score; other backends return `UNIMPLEMENTED`. | ### MCP tools @@ -603,10 +716,14 @@ canonical usage log lives in `/api/usage` and correlates by request ID. |---|---|---| | `get_router_decisions` | read | Recent decision log with optional filters. | | `get_middleware_status` | read | Includes the router section listing configured router models. | - -Mutating routing config — adding a candidate, changing the classifier -model — is YAML-only today; reload with `POST /models/reload` to pick -up edits without restarting. +| `get_router_corpus_stats` | read | KNN corpus size and per-label counts (never texts). | +| `seed_router_corpus` | write | Add labelled exemplars to a KNN router's corpus. | +| `clear_router_corpus` | write | Wipe a KNN router's corpus. | + +Mutating the rest of the routing config — adding a candidate, changing +the classifier model — goes through the model-config surface +(`edit_model_config` / `PATCH /api/models/config-json/:name`); reload +with `POST /models/reload` to pick up YAML edits without restarting. ### Operational notes diff --git a/pkg/mcp/localaitools/client.go b/pkg/mcp/localaitools/client.go index 9d76bee5885c..c53892d77f0b 100644 --- a/pkg/mcp/localaitools/client.go +++ b/pkg/mcp/localaitools/client.go @@ -114,4 +114,16 @@ type LocalAIClient interface { // /app/middleware Routing tab and for agent-driven introspection. // Admin-required when auth is on. GetRouterDecisions(ctx context.Context, q RouterDecisionsQuery) ([]RouterDecision, error) + + // GetRouterCorpusStats reports a knn router's corpus size and + // per-label counts — counts only, texts are never exposed. + GetRouterCorpusStats(ctx context.Context, routerModel string) (*RouterCorpusStats, error) + + // SeedRouterCorpus adds labelled exemplars to a knn router's + // corpus (embedded server-side, persisted, indexed immediately). + SeedRouterCorpus(ctx context.Context, req RouterCorpusSeedRequest) (*RouterCorpusSeedResult, error) + + // ClearRouterCorpus wipes a knn router's corpus — file and live + // index. + ClearRouterCorpus(ctx context.Context, routerModel string) (*RouterCorpusClearResult, error) } diff --git a/pkg/mcp/localaitools/coverage_test.go b/pkg/mcp/localaitools/coverage_test.go index d068180aacb7..7ab18dee306b 100644 --- a/pkg/mcp/localaitools/coverage_test.go +++ b/pkg/mcp/localaitools/coverage_test.go @@ -26,23 +26,24 @@ import ( // the contributor explicitly acknowledges the asymmetry. var toolToHTTPRoute = map[string]string{ // Read-only tools. - ToolGallerySearch: "GET /models/available", - ToolListInstalledModels: "GET / (welcome JSON, ModelsConfig field)", - ToolListGalleries: "GET /models/galleries", - ToolGetJobStatus: "GET /models/jobs/:uuid", - ToolGetModelConfig: "(none) — no JSON-only REST yet; httpapi.Client returns a documented stub", - ToolListBackends: "GET /backends", - ToolListKnownBackends: "GET /backends/known", - ToolSystemInfo: "GET / (welcome JSON)", - ToolListNodes: "GET /api/nodes", - ToolVRAMEstimate: "POST /api/models/vram-estimate", - ToolGetBranding: "GET /api/branding", - ToolGetUsageStats: "GET /api/usage (or /api/usage/all when all=true)", - ToolGetPIIEvents: "GET /api/pii/events", - ToolGetMiddlewareStatus: "GET /api/middleware/status", - ToolGetRouterDecisions: "GET /api/router/decisions", - ToolListAliases: "GET /api/aliases", - ToolListVoiceProfiles: "GET /api/voice-profiles", + ToolGallerySearch: "GET /models/available", + ToolListInstalledModels: "GET / (welcome JSON, ModelsConfig field)", + ToolListGalleries: "GET /models/galleries", + ToolGetJobStatus: "GET /models/jobs/:uuid", + ToolGetModelConfig: "(none) — no JSON-only REST yet; httpapi.Client returns a documented stub", + ToolListBackends: "GET /backends", + ToolListKnownBackends: "GET /backends/known", + ToolSystemInfo: "GET / (welcome JSON)", + ToolListNodes: "GET /api/nodes", + ToolVRAMEstimate: "POST /api/models/vram-estimate", + ToolGetBranding: "GET /api/branding", + ToolGetUsageStats: "GET /api/usage (or /api/usage/all when all=true)", + ToolGetPIIEvents: "GET /api/pii/events", + ToolGetMiddlewareStatus: "GET /api/middleware/status", + ToolGetRouterDecisions: "GET /api/router/decisions", + ToolGetRouterCorpusStats: "GET /api/router/:name/corpus/stats", + ToolListAliases: "GET /api/aliases", + ToolListVoiceProfiles: "GET /api/voice-profiles", // Mutating tools. ToolInstallModel: "POST /models/apply", @@ -57,6 +58,8 @@ var toolToHTTPRoute = map[string]string{ ToolToggleModelPinned: "PUT /models/toggle-pinned/:name/:action", ToolSetBranding: "POST /api/settings (instance_name, instance_tagline)", ToolSetAlias: "PATCH /api/models/config-json/:name (swap) or POST /models/import (create)", + ToolSeedRouterCorpus: "POST /api/router/:name/corpus", + ToolClearRouterCorpus: "DELETE /api/router/:name/corpus", ToolCreateVoiceProfile: "POST /api/voice-profiles", ToolDeleteVoiceProfile: "DELETE /api/voice-profiles/:id", ToolSetNodeVRAMBudget: "PUT /api/nodes/:id/vram-budget", diff --git a/pkg/mcp/localaitools/dto.go b/pkg/mcp/localaitools/dto.go index ecc23204a734..9de392b519f2 100644 --- a/pkg/mcp/localaitools/dto.go +++ b/pkg/mcp/localaitools/dto.go @@ -311,6 +311,50 @@ type RouterDecision struct { CreatedAt string `json:"created_at"` } +// RouterCorpusEntry is one labelled exemplar for seed_router_corpus. +type RouterCorpusEntry struct { + Text string `json:"text" jsonschema:"Example prompt text. Embedded server-side and persisted; NEVER returned by any tool or endpoint."` + Labels []string `json:"labels" jsonschema:"Policy labels this exemplar activates. Every label must be declared in the router's policies."` +} + +// RouterCorpusSeedRequest is the input for seed_router_corpus. +type RouterCorpusSeedRequest struct { + Router string `json:"router" jsonschema:"Router model name — the ModelConfig with classifier: knn and a router.knn block."` + Entries []RouterCorpusEntry `json:"entries" jsonschema:"Labelled exemplars to add. Duplicate texts are skipped, not double-weighted."` +} + +// RouterCorpusSeedResult reports the outcome of seed_router_corpus. +type RouterCorpusSeedResult struct { + Router string `json:"router"` + Added int `json:"added"` + Skipped int `json:"skipped"` + Total int `json:"total"` + LabelCounts map[string]int `json:"label_counts"` +} + +// RouterCorpusQuery names the router whose corpus to inspect or clear. +type RouterCorpusQuery struct { + Router string `json:"router" jsonschema:"Router model name."` +} + +// RouterCorpusStats is the count-only inspection surface for a +// router's KNN corpus. Entry texts are never exposed. +type RouterCorpusStats struct { + Router string `json:"router"` + StoreName string `json:"store_name"` + EmbeddingModel string `json:"embedding_model"` + Total int `json:"total"` + LabelCounts map[string]int `json:"label_counts"` + EmbeddingModels []string `json:"embedding_models,omitempty"` +} + +// RouterCorpusClearResult reports how many entries clear_router_corpus +// removed. +type RouterCorpusClearResult struct { + Router string `json:"router"` + Cleared int `json:"cleared"` +} + // VRAMEstimateRequest is the input for vram_estimate. The output type is // pkg/vram.EstimateResult — used directly via the LocalAIClient interface // so the LLM sees the same shape (size_bytes/size_display/vram_bytes/ diff --git a/pkg/mcp/localaitools/fakes_test.go b/pkg/mcp/localaitools/fakes_test.go index ed328352c73b..8a9cb7220c1d 100644 --- a/pkg/mcp/localaitools/fakes_test.go +++ b/pkg/mcp/localaitools/fakes_test.go @@ -342,3 +342,18 @@ func (f *fakeClient) GetMiddlewareStatus(_ context.Context) (*MiddlewareStatus, Router: MiddlewareRouterStatus{Configured: false, Models: []string{}}, }, nil } + +func (f *fakeClient) GetRouterCorpusStats(_ context.Context, routerModel string) (*RouterCorpusStats, error) { + f.record("GetRouterCorpusStats", routerModel) + return &RouterCorpusStats{Router: routerModel, LabelCounts: map[string]int{}}, nil +} + +func (f *fakeClient) SeedRouterCorpus(_ context.Context, req RouterCorpusSeedRequest) (*RouterCorpusSeedResult, error) { + f.record("SeedRouterCorpus", req) + return &RouterCorpusSeedResult{Router: req.Router, Added: len(req.Entries), LabelCounts: map[string]int{}}, nil +} + +func (f *fakeClient) ClearRouterCorpus(_ context.Context, routerModel string) (*RouterCorpusClearResult, error) { + f.record("ClearRouterCorpus", routerModel) + return &RouterCorpusClearResult{Router: routerModel}, nil +} diff --git a/pkg/mcp/localaitools/httpapi/client.go b/pkg/mcp/localaitools/httpapi/client.go index 771d2908a870..f1ee78b1c865 100644 --- a/pkg/mcp/localaitools/httpapi/client.go +++ b/pkg/mcp/localaitools/httpapi/client.go @@ -752,3 +752,34 @@ func containsTagExact(tags []string, lowerNeedle string) bool { } return false } + +func (c *Client) GetRouterCorpusStats(ctx context.Context, routerModel string) (*localaitools.RouterCorpusStats, error) { + var out localaitools.RouterCorpusStats + path := fmt.Sprintf("/api/router/%s/corpus/stats", url.PathEscape(routerModel)) + if err := c.do(ctx, http.MethodGet, path, nil, &out); err != nil { + return nil, err + } + return &out, nil +} + +func (c *Client) SeedRouterCorpus(ctx context.Context, req localaitools.RouterCorpusSeedRequest) (*localaitools.RouterCorpusSeedResult, error) { + // The REST body carries entries only; the router rides in the path. + body := struct { + Entries []localaitools.RouterCorpusEntry `json:"entries"` + }{Entries: req.Entries} + var out localaitools.RouterCorpusSeedResult + path := fmt.Sprintf("/api/router/%s/corpus", url.PathEscape(req.Router)) + if err := c.do(ctx, http.MethodPost, path, body, &out); err != nil { + return nil, err + } + return &out, nil +} + +func (c *Client) ClearRouterCorpus(ctx context.Context, routerModel string) (*localaitools.RouterCorpusClearResult, error) { + var out localaitools.RouterCorpusClearResult + path := fmt.Sprintf("/api/router/%s/corpus", url.PathEscape(routerModel)) + if err := c.do(ctx, http.MethodDelete, path, nil, &out); err != nil { + return nil, err + } + return &out, nil +} diff --git a/pkg/mcp/localaitools/inproc/client.go b/pkg/mcp/localaitools/inproc/client.go index 9f6eb8fe20f6..a74968928b5b 100644 --- a/pkg/mcp/localaitools/inproc/client.go +++ b/pkg/mcp/localaitools/inproc/client.go @@ -24,6 +24,7 @@ import ( "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/modeladmin" "github.com/mudler/LocalAI/core/services/routing/billing" + "github.com/mudler/LocalAI/core/services/routing/corpus" "github.com/mudler/LocalAI/core/services/routing/pii" "github.com/mudler/LocalAI/core/services/routing/router" "github.com/mudler/LocalAI/core/services/voiceprofile" @@ -68,6 +69,16 @@ type Client struct { // returns when stats are disabled. RouterDecisions router.DecisionStore + // RouterCorpus + the two factories back the corpus tools + // (seed_router_corpus / get_router_corpus_stats / + // clear_router_corpus). nil RouterCorpus makes them return an + // "unavailable" error. The factories mirror the middleware's + // ClassifierDeps so the tools and the request path resolve models + // and store namespaces identically. + RouterCorpus *corpus.Manager + RouterEmbedder func(modelName string) backend.Embedder + RouterVectorStore func(storeName string) backend.VectorStore + modelAdmin *modeladmin.ConfigService } @@ -884,3 +895,111 @@ func capabilityFlagsOf(m *config.ModelConfig) []string { } return out } + +// resolveKNNRouter mirrors the REST endpoint's resolution: the model +// must exist and declare a router.knn block; the store name defaults +// the same way buildClassifier defaults it. +func (c *Client) resolveKNNRouter(routerModel string) (*config.ModelConfig, string, error) { + cfg, err := c.ConfigLoader.LoadModelConfigFileByNameDefaultOptions(routerModel, c.AppConfig) + if err != nil { + return nil, "", fmt.Errorf("load model config: %w", err) + } + if cfg == nil || cfg.Name == "" { + return nil, "", fmt.Errorf("model %q not found", routerModel) + } + if cfg.Router.KNN == nil || cfg.Router.KNN.EmbeddingModel == "" { + return nil, "", fmt.Errorf("model %q has no router.knn block (set classifier: knn and knn.embedding_model first)", routerModel) + } + storeName := cfg.Router.KNN.StoreName + if storeName == "" { + storeName = "router-corpus-" + cfg.Name + } + return cfg, storeName, nil +} + +func (c *Client) GetRouterCorpusStats(_ context.Context, routerModel string) (*localaitools.RouterCorpusStats, error) { + if c.RouterCorpus == nil { + return nil, errors.New("router corpus manager unavailable") + } + cfg, storeName, err := c.resolveKNNRouter(routerModel) + if err != nil { + return nil, err + } + stats, err := c.RouterCorpus.Stats(storeName) + if err != nil { + return nil, err + } + return &localaitools.RouterCorpusStats{ + Router: cfg.Name, + StoreName: storeName, + EmbeddingModel: cfg.Router.KNN.EmbeddingModel, + Total: stats.Total, + LabelCounts: stats.LabelCounts, + EmbeddingModels: stats.EmbeddingModels, + }, nil +} + +func (c *Client) SeedRouterCorpus(ctx context.Context, req localaitools.RouterCorpusSeedRequest) (*localaitools.RouterCorpusSeedResult, error) { + if c.RouterCorpus == nil || c.RouterEmbedder == nil || c.RouterVectorStore == nil { + return nil, errors.New("router corpus manager unavailable") + } + cfg, storeName, err := c.resolveKNNRouter(req.Router) + if err != nil { + return nil, err + } + + // Same invariant the REST endpoint enforces: labels must be + // declared policies, so a typo can't create an unroutable label. + declared := map[string]struct{}{} + for _, p := range cfg.Router.Policies { + declared[p.Label] = struct{}{} + } + entries := make([]corpus.Entry, 0, len(req.Entries)) + for i, e := range req.Entries { + for _, l := range e.Labels { + if _, ok := declared[l]; !ok { + return nil, fmt.Errorf("entry %d: label %q is not declared in router policies", i, l) + } + } + entries = append(entries, corpus.Entry{Text: e.Text, Labels: e.Labels}) + } + + embedder := c.RouterEmbedder(cfg.Router.KNN.EmbeddingModel) + if embedder == nil { + return nil, fmt.Errorf("embedding_model %q not loadable", cfg.Router.KNN.EmbeddingModel) + } + added, skipped, err := c.RouterCorpus.Add(ctx, storeName, cfg.Router.KNN.EmbeddingModel, embedder, c.RouterVectorStore(storeName), entries) + if err != nil { + return nil, err + } + stats, err := c.RouterCorpus.Stats(storeName) + if err != nil { + return nil, err + } + return &localaitools.RouterCorpusSeedResult{ + Router: cfg.Name, + Added: added, + Skipped: skipped, + Total: stats.Total, + LabelCounts: stats.LabelCounts, + }, nil +} + +func (c *Client) ClearRouterCorpus(ctx context.Context, routerModel string) (*localaitools.RouterCorpusClearResult, error) { + if c.RouterCorpus == nil { + return nil, errors.New("router corpus manager unavailable") + } + cfg, storeName, err := c.resolveKNNRouter(routerModel) + if err != nil { + return nil, err + } + var store backend.VectorStore + if c.RouterVectorStore != nil { + store = c.RouterVectorStore(storeName) + } + cleared, err := c.RouterCorpus.Clear(ctx, storeName, store) + if err != nil { + return nil, err + } + return &localaitools.RouterCorpusClearResult{Router: cfg.Name, Cleared: cleared}, nil +} diff --git a/pkg/mcp/localaitools/server_test.go b/pkg/mcp/localaitools/server_test.go index a47b84ba3bc3..99d51495e479 100644 --- a/pkg/mcp/localaitools/server_test.go +++ b/pkg/mcp/localaitools/server_test.go @@ -81,6 +81,7 @@ var expectedFullCatalog = sortedStrings( ToolGetMiddlewareStatus, ToolGetModelConfig, ToolGetPIIEvents, + ToolGetRouterCorpusStats, ToolGetRouterDecisions, ToolGetUsageStats, ToolImportModelURI, @@ -95,6 +96,8 @@ var expectedFullCatalog = sortedStrings( ToolListVoiceProfiles, ToolLoadModel, ToolReloadModels, + ToolSeedRouterCorpus, + ToolClearRouterCorpus, ToolSetAlias, ToolSetBranding, ToolSystemInfo, @@ -115,6 +118,7 @@ var expectedReadOnlyCatalog = sortedStrings( ToolGetMiddlewareStatus, ToolGetModelConfig, ToolGetPIIEvents, + ToolGetRouterCorpusStats, ToolGetRouterDecisions, ToolGetUsageStats, ToolListAliases, diff --git a/pkg/mcp/localaitools/tools.go b/pkg/mcp/localaitools/tools.go index 0d6890e34680..7245eb0a2e61 100644 --- a/pkg/mcp/localaitools/tools.go +++ b/pkg/mcp/localaitools/tools.go @@ -8,22 +8,23 @@ package localaitools // SafetyAnchors guards that those strings stay aligned. const ( // Read-only tools. - ToolGallerySearch = "gallery_search" - ToolListInstalledModels = "list_installed_models" - ToolListGalleries = "list_galleries" - ToolGetJobStatus = "get_job_status" - ToolGetModelConfig = "get_model_config" - ToolListBackends = "list_backends" - ToolListKnownBackends = "list_known_backends" - ToolSystemInfo = "system_info" - ToolListNodes = "list_nodes" - ToolVRAMEstimate = "vram_estimate" - ToolGetBranding = "get_branding" - ToolGetUsageStats = "get_usage_stats" - ToolGetPIIEvents = "get_pii_events" - ToolGetMiddlewareStatus = "get_middleware_status" - ToolGetRouterDecisions = "get_router_decisions" - ToolListVoiceProfiles = "list_voice_profiles" + ToolGallerySearch = "gallery_search" + ToolListInstalledModels = "list_installed_models" + ToolListGalleries = "list_galleries" + ToolGetJobStatus = "get_job_status" + ToolGetModelConfig = "get_model_config" + ToolListBackends = "list_backends" + ToolListKnownBackends = "list_known_backends" + ToolSystemInfo = "system_info" + ToolListNodes = "list_nodes" + ToolVRAMEstimate = "vram_estimate" + ToolGetBranding = "get_branding" + ToolGetUsageStats = "get_usage_stats" + ToolGetPIIEvents = "get_pii_events" + ToolGetMiddlewareStatus = "get_middleware_status" + ToolGetRouterDecisions = "get_router_decisions" + ToolGetRouterCorpusStats = "get_router_corpus_stats" + ToolListVoiceProfiles = "list_voice_profiles" // Mutating tools — guarded by Options.DisableMutating and the // LLM-side safety prompt (see prompts/10_safety.md). @@ -39,6 +40,8 @@ const ( ToolToggleModelPinned = "toggle_model_pinned" ToolSetBranding = "set_branding" ToolSetAlias = "set_alias" + ToolSeedRouterCorpus = "seed_router_corpus" + ToolClearRouterCorpus = "clear_router_corpus" ToolCreateVoiceProfile = "create_voice_profile" ToolDeleteVoiceProfile = "delete_voice_profile" ToolSetNodeVRAMBudget = "set_node_vram_budget" diff --git a/pkg/mcp/localaitools/tools_middleware.go b/pkg/mcp/localaitools/tools_middleware.go index 5dd8066fd4fb..d761768a80d6 100644 --- a/pkg/mcp/localaitools/tools_middleware.go +++ b/pkg/mcp/localaitools/tools_middleware.go @@ -18,7 +18,13 @@ import ( // PII detection policy lives on each detector model's pii_detection // block, edited via the model-config tools — there is no global pattern // set to mutate here anymore. -func registerMiddlewareTools(s *mcp.Server, client LocalAIClient, _ Options) { +// +// The router corpus tools manage the knn classifier's labelled +// exemplar store: get_router_corpus_stats is read-only (counts, never +// texts); seed_router_corpus / clear_router_corpus mutate routing +// behaviour and sit behind the DisableMutating gate like the +// model-config tools. +func registerMiddlewareTools(s *mcp.Server, client LocalAIClient, opts Options) { mcp.AddTool(s, &mcp.Tool{ Name: ToolGetMiddlewareStatus, Description: "Aggregated routing-module status: per-model resolved PII state and the NER detector models each one references, recent event count, plus the active router models and their classifier configs. Read-only.", @@ -40,4 +46,53 @@ func registerMiddlewareTools(s *mcp.Server, client LocalAIClient, _ Options) { } return jsonResult(decisions), nil, nil }) + + mcp.AddTool(s, &mcp.Tool{ + Name: ToolGetRouterCorpusStats, + Description: "Size and per-label exemplar counts of a knn router's corpus. Counts only — corpus texts are never exposed by any tool. Read-only.", + }, func(ctx context.Context, _ *mcp.CallToolRequest, args RouterCorpusQuery) (*mcp.CallToolResult, any, error) { + if args.Router == "" { + return errorResultf("router is required"), nil, nil + } + stats, err := client.GetRouterCorpusStats(ctx, args.Router) + if err != nil { + return errorResult(err), nil, nil + } + return jsonResult(stats), nil, nil + }) + + if opts.DisableMutating { + return + } + + mcp.AddTool(s, &mcp.Tool{ + Name: ToolSeedRouterCorpus, + Description: "Add labelled example prompts to a knn router's corpus. Entries are embedded server-side with the router's knn.embedding_model, persisted, and indexed immediately — routing behaviour changes right away, so confirm with the user first (safety rule 1). Labels must be declared in the router's policies; duplicate texts are skipped.", + }, func(ctx context.Context, _ *mcp.CallToolRequest, args RouterCorpusSeedRequest) (*mcp.CallToolResult, any, error) { + if args.Router == "" { + return errorResultf("router is required"), nil, nil + } + if len(args.Entries) == 0 { + return errorResultf("entries is required"), nil, nil + } + res, err := client.SeedRouterCorpus(ctx, args) + if err != nil { + return errorResult(err), nil, nil + } + return jsonResult(res), nil, nil + }) + + mcp.AddTool(s, &mcp.Tool{ + Name: ToolClearRouterCorpus, + Description: "Wipe a knn router's corpus — the persisted file and the live index. The router falls back for every prompt until reseeded. Destructive; requires explicit user confirmation per safety rule 1.", + }, func(ctx context.Context, _ *mcp.CallToolRequest, args RouterCorpusQuery) (*mcp.CallToolResult, any, error) { + if args.Router == "" { + return errorResultf("router is required"), nil, nil + } + res, err := client.ClearRouterCorpus(ctx, args.Router) + if err != nil { + return errorResult(err), nil, nil + } + return jsonResult(res), nil, nil + }) } diff --git a/swagger/docs.go b/swagger/docs.go index e6792fb09b2b..e11e2f5fb1cc 100644 --- a/swagger/docs.go +++ b/swagger/docs.go @@ -1368,6 +1368,181 @@ const docTemplate = `{ } } }, + "/api/router/{name}/corpus": { + "post": { + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "router" + ], + "summary": "Seed the KNN routing corpus with labelled example prompts", + "parameters": [ + { + "type": "string", + "description": "router model name", + "name": "name", + "in": "path", + "required": true + }, + { + "description": "labelled exemplars", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/schema.RouterCorpusAddRequest" + } + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/schema.RouterCorpusAddResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "404": { + "description": "Not Found", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + }, + "delete": { + "produces": [ + "application/json" + ], + "tags": [ + "router" + ], + "summary": "Clear a router's KNN corpus", + "parameters": [ + { + "type": "string", + "description": "router model name", + "name": "name", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/schema.RouterCorpusClearResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "404": { + "description": "Not Found", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, + "/api/router/{name}/corpus/stats": { + "get": { + "produces": [ + "application/json" + ], + "tags": [ + "router" + ], + "summary": "Inspect a router's KNN corpus (label counts only, never texts)", + "parameters": [ + { + "type": "string", + "description": "router model name", + "name": "name", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/schema.RouterCorpusStatsResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "404": { + "description": "Not Found", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, "/api/traces": { "get": { "description": "Returns captured API exchange traces (request/response pairs) in reverse chronological order", @@ -6607,6 +6782,99 @@ const docTemplate = `{ } } }, + "schema.RouterCorpusAddRequest": { + "type": "object", + "properties": { + "entries": { + "type": "array", + "items": { + "$ref": "#/definitions/schema.RouterCorpusEntry" + } + } + } + }, + "schema.RouterCorpusAddResponse": { + "type": "object", + "properties": { + "added": { + "description": "Added is how many entries were embedded, persisted, and indexed.", + "type": "integer" + }, + "label_counts": { + "description": "LabelCounts is the per-label exemplar count after the call.", + "type": "object", + "additionalProperties": { + "type": "integer" + } + }, + "router": { + "type": "string" + }, + "skipped": { + "description": "Skipped counts entries whose text was already in the corpus —\nduplicates are rejected rather than double-weighted.", + "type": "integer" + }, + "total": { + "description": "Total is the corpus size after the call.", + "type": "integer" + } + } + }, + "schema.RouterCorpusClearResponse": { + "type": "object", + "properties": { + "cleared": { + "type": "integer" + }, + "router": { + "type": "string" + } + } + }, + "schema.RouterCorpusEntry": { + "type": "object", + "properties": { + "labels": { + "type": "array", + "items": { + "type": "string" + } + }, + "text": { + "type": "string" + } + } + }, + "schema.RouterCorpusStatsResponse": { + "type": "object", + "properties": { + "embedding_model": { + "type": "string" + }, + "embedding_models": { + "description": "EmbeddingModels lists the embedder fingerprints present in the\npersisted corpus; more than one means part of the corpus is\npending re-embedding on the next load.", + "type": "array", + "items": { + "type": "string" + } + }, + "label_counts": { + "type": "object", + "additionalProperties": { + "type": "integer" + } + }, + "router": { + "type": "string" + }, + "store_name": { + "type": "string" + }, + "total": { + "type": "integer" + } + } + }, "schema.RouterDecideRequest": { "type": "object", "properties": { @@ -6654,6 +6922,10 @@ const docTemplate = `{ "description": "LatencyMs is the classifier's wall-clock cost.", "type": "integer" }, + "nearest_similarity": { + "description": "NearestSimilarity is the cosine similarity of the closest KNN\ncorpus entry — populated by the knn classifier even when the\ndecision fell back because the probe was out of corpus range.\n0 for other classifiers.", + "type": "number" + }, "router": { "description": "Router echoes the requested router model.", "type": "string" diff --git a/swagger/swagger.json b/swagger/swagger.json index 17eba941d909..94c791ca2165 100644 --- a/swagger/swagger.json +++ b/swagger/swagger.json @@ -1365,6 +1365,181 @@ } } }, + "/api/router/{name}/corpus": { + "post": { + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "router" + ], + "summary": "Seed the KNN routing corpus with labelled example prompts", + "parameters": [ + { + "type": "string", + "description": "router model name", + "name": "name", + "in": "path", + "required": true + }, + { + "description": "labelled exemplars", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/schema.RouterCorpusAddRequest" + } + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/schema.RouterCorpusAddResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "404": { + "description": "Not Found", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + }, + "delete": { + "produces": [ + "application/json" + ], + "tags": [ + "router" + ], + "summary": "Clear a router's KNN corpus", + "parameters": [ + { + "type": "string", + "description": "router model name", + "name": "name", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/schema.RouterCorpusClearResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "404": { + "description": "Not Found", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, + "/api/router/{name}/corpus/stats": { + "get": { + "produces": [ + "application/json" + ], + "tags": [ + "router" + ], + "summary": "Inspect a router's KNN corpus (label counts only, never texts)", + "parameters": [ + { + "type": "string", + "description": "router model name", + "name": "name", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/schema.RouterCorpusStatsResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "404": { + "description": "Not Found", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, "/api/traces": { "get": { "description": "Returns captured API exchange traces (request/response pairs) in reverse chronological order", @@ -6604,6 +6779,99 @@ } } }, + "schema.RouterCorpusAddRequest": { + "type": "object", + "properties": { + "entries": { + "type": "array", + "items": { + "$ref": "#/definitions/schema.RouterCorpusEntry" + } + } + } + }, + "schema.RouterCorpusAddResponse": { + "type": "object", + "properties": { + "added": { + "description": "Added is how many entries were embedded, persisted, and indexed.", + "type": "integer" + }, + "label_counts": { + "description": "LabelCounts is the per-label exemplar count after the call.", + "type": "object", + "additionalProperties": { + "type": "integer" + } + }, + "router": { + "type": "string" + }, + "skipped": { + "description": "Skipped counts entries whose text was already in the corpus —\nduplicates are rejected rather than double-weighted.", + "type": "integer" + }, + "total": { + "description": "Total is the corpus size after the call.", + "type": "integer" + } + } + }, + "schema.RouterCorpusClearResponse": { + "type": "object", + "properties": { + "cleared": { + "type": "integer" + }, + "router": { + "type": "string" + } + } + }, + "schema.RouterCorpusEntry": { + "type": "object", + "properties": { + "labels": { + "type": "array", + "items": { + "type": "string" + } + }, + "text": { + "type": "string" + } + } + }, + "schema.RouterCorpusStatsResponse": { + "type": "object", + "properties": { + "embedding_model": { + "type": "string" + }, + "embedding_models": { + "description": "EmbeddingModels lists the embedder fingerprints present in the\npersisted corpus; more than one means part of the corpus is\npending re-embedding on the next load.", + "type": "array", + "items": { + "type": "string" + } + }, + "label_counts": { + "type": "object", + "additionalProperties": { + "type": "integer" + } + }, + "router": { + "type": "string" + }, + "store_name": { + "type": "string" + }, + "total": { + "type": "integer" + } + } + }, "schema.RouterDecideRequest": { "type": "object", "properties": { @@ -6651,6 +6919,10 @@ "description": "LatencyMs is the classifier's wall-clock cost.", "type": "integer" }, + "nearest_similarity": { + "description": "NearestSimilarity is the cosine similarity of the closest KNN\ncorpus entry — populated by the knn classifier even when the\ndecision fell back because the probe was out of corpus range.\n0 for other classifiers.", + "type": "number" + }, "router": { "description": "Router echoes the requested router model.", "type": "string" diff --git a/swagger/swagger.yaml b/swagger/swagger.yaml index 62c655ac8140..b1795aa170ba 100644 --- a/swagger/swagger.yaml +++ b/swagger/swagger.yaml @@ -2173,6 +2173,73 @@ definitions: redacted_text: type: string type: object + schema.RouterCorpusAddRequest: + properties: + entries: + items: + $ref: '#/definitions/schema.RouterCorpusEntry' + type: array + type: object + schema.RouterCorpusAddResponse: + properties: + added: + description: Added is how many entries were embedded, persisted, and indexed. + type: integer + label_counts: + additionalProperties: + type: integer + description: LabelCounts is the per-label exemplar count after the call. + type: object + router: + type: string + skipped: + description: |- + Skipped counts entries whose text was already in the corpus — + duplicates are rejected rather than double-weighted. + type: integer + total: + description: Total is the corpus size after the call. + type: integer + type: object + schema.RouterCorpusClearResponse: + properties: + cleared: + type: integer + router: + type: string + type: object + schema.RouterCorpusEntry: + properties: + labels: + items: + type: string + type: array + text: + type: string + type: object + schema.RouterCorpusStatsResponse: + properties: + embedding_model: + type: string + embedding_models: + description: |- + EmbeddingModels lists the embedder fingerprints present in the + persisted corpus; more than one means part of the corpus is + pending re-embedding on the next load. + items: + type: string + type: array + label_counts: + additionalProperties: + type: integer + type: object + router: + type: string + store_name: + type: string + total: + type: integer + type: object schema.RouterDecideRequest: properties: input: @@ -2224,6 +2291,13 @@ definitions: latency_ms: description: LatencyMs is the classifier's wall-clock cost. type: integer + nearest_similarity: + description: |- + NearestSimilarity is the cosine similarity of the closest KNN + corpus entry — populated by the knn classifier even when the + decision fell back because the probe was out of corpus range. + 0 for other classifiers. + type: number router: description: Router echoes the requested router model. type: string @@ -3553,6 +3627,121 @@ paths: summary: Redact PII in a string by applying the configured policy. tags: - pii + /api/router/{name}/corpus: + delete: + parameters: + - description: router model name + in: path + name: name + required: true + type: string + produces: + - application/json + responses: + "200": + description: OK + schema: + $ref: '#/definitions/schema.RouterCorpusClearResponse' + "400": + description: Bad Request + schema: + additionalProperties: + type: string + type: object + "404": + description: Not Found + schema: + additionalProperties: + type: string + type: object + "500": + description: Internal Server Error + schema: + additionalProperties: + type: string + type: object + summary: Clear a router's KNN corpus + tags: + - router + post: + consumes: + - application/json + parameters: + - description: router model name + in: path + name: name + required: true + type: string + - description: labelled exemplars + in: body + name: request + required: true + schema: + $ref: '#/definitions/schema.RouterCorpusAddRequest' + produces: + - application/json + responses: + "200": + description: OK + schema: + $ref: '#/definitions/schema.RouterCorpusAddResponse' + "400": + description: Bad Request + schema: + additionalProperties: + type: string + type: object + "404": + description: Not Found + schema: + additionalProperties: + type: string + type: object + "500": + description: Internal Server Error + schema: + additionalProperties: + type: string + type: object + summary: Seed the KNN routing corpus with labelled example prompts + tags: + - router + /api/router/{name}/corpus/stats: + get: + parameters: + - description: router model name + in: path + name: name + required: true + type: string + produces: + - application/json + responses: + "200": + description: OK + schema: + $ref: '#/definitions/schema.RouterCorpusStatsResponse' + "400": + description: Bad Request + schema: + additionalProperties: + type: string + type: object + "404": + description: Not Found + schema: + additionalProperties: + type: string + type: object + "500": + description: Internal Server Error + schema: + additionalProperties: + type: string + type: object + summary: Inspect a router's KNN corpus (label counts only, never texts) + tags: + - router /api/router/decide: post: consumes: From d745bd81418cbd57086a5f3b316fba51c63102e4 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Mon, 6 Jul 2026 21:04:40 +0100 Subject: [PATCH 02/10] feat(router): name consulted corpus neighbours in knn decisions Every knn decision (decision log rows and the /api/router/decide response) now carries neighbors: the K retrieved corpus entries by descending similarity - including ones below the epistemic gate, which is what makes fallback decisions diagnosable - each as {id, similarity, labels}. The id is the entry's content hash (first 8 bytes of the SHA-256 of its text, hex): stable across reseeds and re-embeds, and text-free, so an external platform that seeded the corpus can recompute text->id on its own copy and bucket decisions by corpus region (per- region reliability accounting) without corpus text ever leaving the server. A corrupt index payload surfaces as an id-less neighbour at a real similarity instead of disappearing. Assisted-by: Claude:claude-fable-5 [Claude Code] --- core/http/endpoints/localai/router_decide.go | 10 ++++ core/http/middleware/route_model_test.go | 2 +- core/schema/localai.go | 16 ++++++ core/services/routing/corpus/manager.go | 2 +- core/services/routing/router/decisions.go | 5 ++ core/services/routing/router/knn.go | 49 +++++++++++++++-- core/services/routing/router/knn_test.go | 58 +++++++++++++++----- core/services/routing/router/resolve.go | 1 + core/services/routing/router/types.go | 21 +++++++ docs/content/features/middleware.md | 10 ++++ swagger/docs.go | 24 ++++++++ swagger/swagger.json | 24 ++++++++ swagger/swagger.yaml | 19 +++++++ 13 files changed, 219 insertions(+), 22 deletions(-) diff --git a/core/http/endpoints/localai/router_decide.go b/core/http/endpoints/localai/router_decide.go index b13b3a7acbc0..ca24e69e643c 100644 --- a/core/http/endpoints/localai/router_decide.go +++ b/core/http/endpoints/localai/router_decide.go @@ -94,6 +94,15 @@ func RouterDecideEndpoint(loader *config.ModelConfigLoader, appConfig *config.Ap classifierName = router.ClassifierScore } + neighbors := make([]schema.RouterDecideNeighbor, 0, len(decision.Neighbors)) + for _, n := range decision.Neighbors { + neighbors = append(neighbors, schema.RouterDecideNeighbor{ + ID: n.ID, + Similarity: n.Similarity, + Labels: n.Labels, + }) + } + return c.JSON(http.StatusOK, schema.RouterDecideResponse{ Router: req.Router, Classifier: classifierName, @@ -105,6 +114,7 @@ func RouterDecideEndpoint(loader *config.ModelConfigLoader, appConfig *config.Ap Cached: decision.Cached, CacheSimilarity: decision.CacheSimilarity, NearestSimilarity: decision.NearestSimilarity, + Neighbors: neighbors, }) } } diff --git a/core/http/middleware/route_model_test.go b/core/http/middleware/route_model_test.go index c8f85b8f3a05..ea8dbe46ad01 100644 --- a/core/http/middleware/route_model_test.go +++ b/core/http/middleware/route_model_test.go @@ -582,7 +582,7 @@ var ( ) func corpusPayload(labels ...string) []byte { - b, err := router.EncodeCorpusEntry(labels) + b, err := router.EncodeCorpusEntry(router.EntryID(strings.Join(labels, "+")), labels) Expect(err).NotTo(HaveOccurred()) return b } diff --git a/core/schema/localai.go b/core/schema/localai.go index 6c51273fd68b..9cdb7a020ba6 100644 --- a/core/schema/localai.go +++ b/core/schema/localai.go @@ -549,6 +549,22 @@ type RouterDecideResponse struct { // decision fell back because the probe was out of corpus range. // 0 for other classifiers. NearestSimilarity float64 `json:"nearest_similarity,omitempty"` + // Neighbors lists the corpus entries the knn classifier retrieved, + // by descending similarity, including ones below the similarity + // gate. Empty for other classifiers. + Neighbors []RouterDecideNeighbor `json:"neighbors,omitempty"` +} + +// RouterDecideNeighbor is one KNN corpus neighbour in a decide +// response (mirrors router.NeighborRef; schema stays import-free of +// the routing packages). ID is the corpus entry's content hash — the +// first 8 bytes of the SHA-256 of its text, hex-encoded — so callers +// that seeded the corpus can recompute text→ID locally and bucket +// decisions by corpus region; the text itself never leaves the server. +type RouterDecideNeighbor struct { + ID string `json:"id,omitempty"` + Similarity float64 `json:"similarity"` + Labels []string `json:"labels,omitempty"` } // RouterCorpusEntry is one labelled exemplar submitted to diff --git a/core/services/routing/corpus/manager.go b/core/services/routing/corpus/manager.go index e61d8fd84bd8..8f0bc35cd799 100644 --- a/core/services/routing/corpus/manager.go +++ b/core/services/routing/corpus/manager.go @@ -290,7 +290,7 @@ func insertAll(ctx context.Context, store backend.VectorStore, entries []Entry) vecs := make([][]float32, 0, len(entries)) payloads := make([][]byte, 0, len(entries)) for _, e := range entries { - payload, err := router.EncodeCorpusEntry(e.Labels) + payload, err := router.EncodeCorpusEntry(router.EntryID(e.Text), e.Labels) if err != nil { return err } diff --git a/core/services/routing/router/decisions.go b/core/services/routing/router/decisions.go index 23f512813720..492da3af892d 100644 --- a/core/services/routing/router/decisions.go +++ b/core/services/routing/router/decisions.go @@ -29,6 +29,11 @@ type DecisionRecord struct { // threshold. Empty on cache hits (only the final label set is cached). LabelScores []LabelScore `json:"label_scores,omitempty"` ActivationThreshold float64 `json:"activation_threshold,omitempty"` + // Neighbors names the corpus entries a KNN decision consulted (content- + // hash IDs, similarities, labels — never text). This is the join key + // for platform-side per-region reliability accounting: decisions that + // retrieve the same corpus entries belong to the same region. + Neighbors []NeighborRef `json:"neighbors,omitempty"` // Source groups decisions by the entry point that produced them so // the admin page can split realtime / chat / anthropic streams. Empty // string is treated as "chat" for backward compatibility with rows diff --git a/core/services/routing/router/knn.go b/core/services/routing/router/knn.go index ddde30f43679..f496551a9971 100644 --- a/core/services/routing/router/knn.go +++ b/core/services/routing/router/knn.go @@ -2,6 +2,8 @@ package router import ( "context" + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" "sort" @@ -113,6 +115,21 @@ func (c *KNNClassifier) Classify(ctx context.Context, p Probe) (Decision, error) return errDecision(start, fmt.Errorf("knn classifier search: %w", err)) } + // Record every retrieved neighbour — including sub-gate ones — before + // the filtering below reuses the slice's backing array. The decision + // log needs the full retrieval to make fallbacks diagnosable: "what + // WAS nearby, and how was it labelled". A corrupt payload still names + // its similarity so index corruption is visible rather than silent. + refs := make([]NeighborRef, 0, len(neighbors)) + for _, n := range neighbors { + ref := NeighborRef{Similarity: n.Similarity} + if entry, ok := decodeCorpusEntry(n.Payload); ok { + ref.ID = entry.ID + ref.Labels = entry.Labels + } + refs = append(refs, ref) + } + // Epistemic gate: only neighbours the probe is genuinely close to // may vote. Keeping sub-threshold neighbours out of the vote (rather // than merely gating on the best one) stops far-away corpus regions @@ -135,6 +152,7 @@ func (c *KNNClassifier) Classify(ctx context.Context, p Probe) (Decision, error) // simply too tight. return Decision{ NearestSimilarity: best, + Neighbors: refs, ActivationThreshold: c.voteThreshold, Latency: time.Since(start), }, nil @@ -156,6 +174,7 @@ func (c *KNNClassifier) Classify(ctx context.Context, p Probe) (Decision, error) if total == 0 { return Decision{ NearestSimilarity: best, + Neighbors: refs, ActivationThreshold: c.voteThreshold, Latency: time.Since(start), }, nil @@ -187,6 +206,7 @@ func (c *KNNClassifier) Classify(ctx context.Context, p Probe) (Decision, error) ActivationThreshold: c.voteThreshold, LabelScores: NewLabelScores(labels, scores), NearestSimilarity: best, + Neighbors: refs, Latency: time.Since(start), } if len(active) > 0 { @@ -198,19 +218,36 @@ func (c *KNNClassifier) Classify(ctx context.Context, p Probe) (Decision, error) // corpusEntry is the stored shape of one labelled exemplar. Kept // deliberately minimal: the vector key lives in the store, the text // lives only in the corpus file (never returned by inspection APIs), -// so the store payload is just the label set. +// so the store payload is the label set plus the content-hash ID that +// lets decisions name the exemplar without carrying its text. type corpusEntry struct { + ID string `json:"id,omitempty"` Labels []string `json:"labels"` } -// EncodeCorpusEntry serialises the labels of one corpus exemplar into -// the vector-store payload shape Classify votes over. Exported for the -// corpus loader/API in core, which owns insertion. -func EncodeCorpusEntry(labels []string) ([]byte, error) { +// EntryID returns the stable identifier of a corpus entry: the first 8 +// bytes of the SHA-256 of its text, hex-encoded. Content-derived, so it +// survives reseeds, re-embeds, and restarts; one-way, so the decision +// log can reference an exemplar without leaking it. A platform that +// seeded the corpus can recompute text→ID on its own copy to join +// decisions back to its labelled data. +func EntryID(text string) string { + sum := sha256.Sum256([]byte(text)) + return hex.EncodeToString(sum[:8]) +} + +// EncodeCorpusEntry serialises one corpus exemplar into the +// vector-store payload shape Classify votes over. Exported for the +// corpus loader/API in core, which owns insertion; id comes from +// EntryID so the payload never carries text. +func EncodeCorpusEntry(id string, labels []string) ([]byte, error) { + if id == "" { + return nil, fmt.Errorf("corpus entry needs an id (EntryID of its text)") + } if len(labels) == 0 { return nil, fmt.Errorf("corpus entry needs at least one label") } - return json.Marshal(corpusEntry{Labels: labels}) + return json.Marshal(corpusEntry{ID: id, Labels: labels}) } func decodeCorpusEntry(b []byte) (corpusEntry, bool) { diff --git a/core/services/routing/router/knn_test.go b/core/services/routing/router/knn_test.go index 49bfd9a4c1e6..7b732c813f4f 100644 --- a/core/services/routing/router/knn_test.go +++ b/core/services/routing/router/knn_test.go @@ -38,8 +38,8 @@ func (s *scriptedKNNStore) Insert(_ context.Context, _ []float32, _ []byte) erro return errors.New("knn classifier must never insert") } -func mustEntry(labels ...string) []byte { - b, err := router.EncodeCorpusEntry(labels) +func mustEntry(id string, labels ...string) []byte { + b, err := router.EncodeCorpusEntry(id, labels) Expect(err).ToNot(HaveOccurred()) return b } @@ -72,9 +72,9 @@ var _ = Describe("KNNClassifier", func() { // Both clear the 0.5 majority → both active; candidate matching // then requires a model labelled for both. store := &scriptedKNNStore{neighbors: []backend.Neighbor{ - {Similarity: 0.90, Payload: mustEntry("code")}, - {Similarity: 0.85, Payload: mustEntry("code", "math")}, - {Similarity: 0.82, Payload: mustEntry("math")}, + {Similarity: 0.90, Payload: mustEntry("e1", "code")}, + {Similarity: 0.85, Payload: mustEntry("e2", "code", "math")}, + {Similarity: 0.82, Payload: mustEntry("e3", "math")}, }} d := classify(store, router.KNNClassifierOptions{K: 3}) Expect(d.Labels).To(Equal([]string{"code", "math"})) @@ -86,15 +86,23 @@ var _ = Describe("KNNClassifier", func() { Expect(d.Score).To(BeNumerically("~", 1.75/2.57, 1e-9)) Expect(d.NearestSimilarity).To(BeNumerically("~", 0.90, 1e-9)) Expect(store.lastK).To(Equal(3)) + // The decision names every consulted neighbour so the log can be + // joined back to corpus entries. + Expect(d.Neighbors).To(HaveLen(3)) + Expect(d.Neighbors[0].ID).To(Equal("e1")) + Expect(d.Neighbors[0].Similarity).To(BeNumerically("~", 0.90, 1e-9)) + Expect(d.Neighbors[1].ID).To(Equal("e2")) + Expect(d.Neighbors[1].Labels).To(Equal([]string{"code", "math"})) + Expect(d.Neighbors[2].ID).To(Equal("e3")) }) It("does not activate a minority label", func() { // share(chat) = 0.81/2.57 < 0.5 → inactive, but still reported // in LabelScores so the decision log shows how close it came. store := &scriptedKNNStore{neighbors: []backend.Neighbor{ - {Similarity: 0.90, Payload: mustEntry("code")}, - {Similarity: 0.86, Payload: mustEntry("code")}, - {Similarity: 0.81, Payload: mustEntry("chat")}, + {Similarity: 0.90, Payload: mustEntry("e1", "code")}, + {Similarity: 0.86, Payload: mustEntry("e2", "code")}, + {Similarity: 0.81, Payload: mustEntry("e3", "chat")}, }} d := classify(store, router.KNNClassifierOptions{K: 3}) Expect(d.Labels).To(Equal([]string{"code"})) @@ -105,14 +113,19 @@ var _ = Describe("KNNClassifier", func() { It("gates out-of-corpus probes to the fallback (empty labels)", func() { store := &scriptedKNNStore{neighbors: []backend.Neighbor{ - {Similarity: 0.55, Payload: mustEntry("code")}, - {Similarity: 0.40, Payload: mustEntry("math")}, + {Similarity: 0.55, Payload: mustEntry("e1", "code")}, + {Similarity: 0.40, Payload: mustEntry("e2", "math")}, }} d := classify(store, router.KNNClassifierOptions{SimilarityThreshold: 0.80}) Expect(d.Labels).To(BeEmpty()) // The admin-facing epistemic signal: how far away the nearest // labelled experience was. Expect(d.NearestSimilarity).To(BeNumerically("~", 0.55, 1e-9)) + // Fallback decisions still name the sub-gate neighbours — that is + // exactly what makes them diagnosable. + Expect(d.Neighbors).To(HaveLen(2)) + Expect(d.Neighbors[0].ID).To(Equal("e1")) + Expect(d.Neighbors[1].ID).To(Equal("e2")) }) It("excludes sub-threshold neighbours from the vote", func() { @@ -120,8 +133,8 @@ var _ = Describe("KNNClassifier", func() { // majority: with it, share(code) would be 0.9/1.2 = 0.75; the // vote must instead be over the single usable neighbour. store := &scriptedKNNStore{neighbors: []backend.Neighbor{ - {Similarity: 0.90, Payload: mustEntry("code")}, - {Similarity: 0.30, Payload: mustEntry("chat")}, + {Similarity: 0.90, Payload: mustEntry("e1", "code")}, + {Similarity: 0.30, Payload: mustEntry("e2", "chat")}, }} d := classify(store, router.KNNClassifierOptions{K: 2, SimilarityThreshold: 0.80}) Expect(d.Labels).To(Equal([]string{"code"})) @@ -131,7 +144,7 @@ var _ = Describe("KNNClassifier", func() { It("degenerates to nearest-entry labels at K=1", func() { store := &scriptedKNNStore{neighbors: []backend.Neighbor{ - {Similarity: 0.95, Payload: mustEntry("reasoning", "math")}, + {Similarity: 0.95, Payload: mustEntry("e1", "reasoning", "math")}, }} d := classify(store, router.KNNClassifierOptions{K: 1}) Expect(d.Labels).To(Equal([]string{"math", "reasoning"})) @@ -145,6 +158,11 @@ var _ = Describe("KNNClassifier", func() { d := classify(store, router.KNNClassifierOptions{}) Expect(d.Labels).To(BeEmpty()) Expect(d.NearestSimilarity).To(BeNumerically("~", 0.90, 1e-9)) + // A corrupt payload is still visible in the neighbour list — an + // empty ID at a real similarity flags index corruption. + Expect(d.Neighbors).To(HaveLen(1)) + Expect(d.Neighbors[0].ID).To(BeEmpty()) + Expect(d.Neighbors[0].Similarity).To(BeNumerically("~", 0.90, 1e-9)) }) It("returns the embed error so the middleware can fall back", func() { @@ -163,10 +181,22 @@ var _ = Describe("KNNClassifier", func() { }) It("rejects corpus entries without labels at encode time", func() { - _, err := router.EncodeCorpusEntry(nil) + _, err := router.EncodeCorpusEntry("e1", nil) Expect(err).To(HaveOccurred()) }) + It("rejects corpus entries without an id at encode time", func() { + _, err := router.EncodeCorpusEntry("", []string{"code"}) + Expect(err).To(HaveOccurred()) + }) + + It("derives stable, text-free entry ids", func() { + id := router.EntryID("Is 18857 a prime number? Answer yes or no.") + Expect(id).To(HaveLen(16)) + Expect(id).To(Equal(router.EntryID("Is 18857 a prime number? Answer yes or no."))) + Expect(id).ToNot(Equal(router.EntryID("Is 18858 a prime number? Answer yes or no."))) + }) + It("panics on missing embedder or store", func() { Expect(func() { router.NewKNNClassifier(nil, &scriptedKNNStore{}, router.KNNClassifierOptions{}) }).To(Panic()) Expect(func() { router.NewKNNClassifier(embedder, nil, router.KNNClassifierOptions{}) }).To(Panic()) diff --git a/core/services/routing/router/resolve.go b/core/services/routing/router/resolve.go index 486eedfdea5d..6950bf8caa33 100644 --- a/core/services/routing/router/resolve.go +++ b/core/services/routing/router/resolve.go @@ -168,6 +168,7 @@ func (r *ResolveResult) ToDecisionRecord(id, correlationID, userID, source strin NearestSimilarity: r.Decision.NearestSimilarity, LabelScores: r.Decision.LabelScores, ActivationThreshold: r.Decision.ActivationThreshold, + Neighbors: r.Decision.Neighbors, Source: source, CreatedAt: time.Now().UTC(), } diff --git a/core/services/routing/router/types.go b/core/services/routing/router/types.go index 8e4d8f1da086..bba1b9d399c4 100644 --- a/core/services/routing/router/types.go +++ b/core/services/routing/router/types.go @@ -78,6 +78,27 @@ type Decision struct { // which is exactly when an admin wants to know how far off the // nearest labelled experience was. 0 for other classifiers. NearestSimilarity float64 `json:"nearest_similarity,omitempty"` + + // Neighbors lists the K corpus entries the KNN classifier retrieved, + // ordered by descending similarity — including neighbours below the + // similarity gate, which is what makes fallback decisions diagnosable + // (what WAS nearby, and how was it labelled). Empty for other + // classifiers. + Neighbors []NeighborRef `json:"neighbors,omitempty"` +} + +// NeighborRef identifies one corpus neighbour consulted for a KNN +// decision. ID is the entry's content hash (EntryID) — stable across +// reseeds and re-embeds, and text-free: an external platform that +// seeded the corpus can recompute text→ID on its own copy to bucket +// decisions by corpus region without corpus text ever leaving the +// server. Labels repeats the entry's label set (already public via +// corpus stats). An empty ID with a non-zero similarity marks a +// corrupt index payload. +type NeighborRef struct { + ID string `json:"id,omitempty"` + Similarity float64 `json:"similarity"` + Labels []string `json:"labels,omitempty"` } // LabelScore is one entry in Decision.LabelScores — a policy label and diff --git a/docs/content/features/middleware.md b/docs/content/features/middleware.md index c3733891c078..f4c735b10734 100644 --- a/docs/content/features/middleware.md +++ b/docs/content/features/middleware.md @@ -475,6 +475,16 @@ For each request: for the other classifiers. With `k: 1` this degenerates to "nearest example's labels". +Every knn decision (in the decision log and the `/api/router/decide` +response) also carries `neighbors` — the `k` retrieved corpus entries +by descending similarity, **including** ones below the gate, each as +`{id, similarity, labels}`. The `id` is the entry's content hash (the +first 8 bytes of the SHA-256 of its text, hex-encoded): stable across +reseeds and re-embeds, and text-free — whoever seeded the corpus can +recompute text→id on their own copy to group decisions by corpus +region (e.g. for external per-region reliability accounting) without +corpus text ever leaving the server. + #### Seeding and curating the corpus (API-only) Corpus entries may contain example user content, so they are managed diff --git a/swagger/docs.go b/swagger/docs.go index e11e2f5fb1cc..f3d55d61bf98 100644 --- a/swagger/docs.go +++ b/swagger/docs.go @@ -6875,6 +6875,23 @@ const docTemplate = `{ } } }, + "schema.RouterDecideNeighbor": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "labels": { + "type": "array", + "items": { + "type": "string" + } + }, + "similarity": { + "type": "number" + } + } + }, "schema.RouterDecideRequest": { "type": "object", "properties": { @@ -6926,6 +6943,13 @@ const docTemplate = `{ "description": "NearestSimilarity is the cosine similarity of the closest KNN\ncorpus entry — populated by the knn classifier even when the\ndecision fell back because the probe was out of corpus range.\n0 for other classifiers.", "type": "number" }, + "neighbors": { + "description": "Neighbors lists the corpus entries the knn classifier retrieved,\nby descending similarity, including ones below the similarity\ngate. Empty for other classifiers.", + "type": "array", + "items": { + "$ref": "#/definitions/schema.RouterDecideNeighbor" + } + }, "router": { "description": "Router echoes the requested router model.", "type": "string" diff --git a/swagger/swagger.json b/swagger/swagger.json index 94c791ca2165..fd14899d87f1 100644 --- a/swagger/swagger.json +++ b/swagger/swagger.json @@ -6872,6 +6872,23 @@ } } }, + "schema.RouterDecideNeighbor": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "labels": { + "type": "array", + "items": { + "type": "string" + } + }, + "similarity": { + "type": "number" + } + } + }, "schema.RouterDecideRequest": { "type": "object", "properties": { @@ -6923,6 +6940,13 @@ "description": "NearestSimilarity is the cosine similarity of the closest KNN\ncorpus entry — populated by the knn classifier even when the\ndecision fell back because the probe was out of corpus range.\n0 for other classifiers.", "type": "number" }, + "neighbors": { + "description": "Neighbors lists the corpus entries the knn classifier retrieved,\nby descending similarity, including ones below the similarity\ngate. Empty for other classifiers.", + "type": "array", + "items": { + "$ref": "#/definitions/schema.RouterDecideNeighbor" + } + }, "router": { "description": "Router echoes the requested router model.", "type": "string" diff --git a/swagger/swagger.yaml b/swagger/swagger.yaml index b1795aa170ba..47b75b942913 100644 --- a/swagger/swagger.yaml +++ b/swagger/swagger.yaml @@ -2240,6 +2240,17 @@ definitions: total: type: integer type: object + schema.RouterDecideNeighbor: + properties: + id: + type: string + labels: + items: + type: string + type: array + similarity: + type: number + type: object schema.RouterDecideRequest: properties: input: @@ -2298,6 +2309,14 @@ definitions: decision fell back because the probe was out of corpus range. 0 for other classifiers. type: number + neighbors: + description: |- + Neighbors lists the corpus entries the knn classifier retrieved, + by descending similarity, including ones below the similarity + gate. Empty for other classifiers. + items: + $ref: '#/definitions/schema.RouterDecideNeighbor' + type: array router: description: Router echoes the requested router model. type: string From a8c8031642f67c8f16a08e67f0d7d78722816a22 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Tue, 7 Jul 2026 10:05:35 +0100 Subject: [PATCH 03/10] refactor(router): deduplicate knn plumbing and cut corpus hot-path waste Post-review cleanup of the knn-first-class-router branch; no behaviour changes on the API surface. Reuse/altitude: - RouterKNNConfig.ResolvedStoreName is now the single source of the router-corpus- default (was hand-derived in four files). - corpus.ResolveKNNRouter + corpus.Seed carry the shared model resolution and seed validation; the REST endpoints and the assistant MCP client are thin transport adapters over them, with sentinel errors mapped to HTTP statuses at the echo boundary. - middleware.NewClassifierDeps assembles the classifier dependency set once for all five entry points (OpenAI, Anthropic, realtime, decide, corpus) instead of five hand-copied literals. - router.AllClassifiers feeds both the status endpoint and the unknown-classifier error, ending the classifier-list drift. - Per-classifier requirements moved out of validateRouterPolicies into their buildClassifier arms; the knn arm owns its embedding_cache opt-out instead of a name-check in the shared wrap tail. - adminOnly replaces four inline copies of the admin gate in the middleware routes. - localVectorStore.Search delegates to SearchK (identical traces). Efficiency: - Manager.Add embeds outside the manager mutex and appends to the JSONL file (O(new) instead of O(corpus) rewrite); a torn tail from a crash mid-append is tolerated on read and repaired on next write. - Stats memoises per store keyed on the file's stat fingerprint and no longer takes the manager mutex, so the 5s status poll stops parsing vector-laden JSONL and stops blocking behind seeds. - KNN Classify decodes each neighbour payload once (was twice) and builds refs and votes in a single pass with one fallback return. - Corpus file writes fsync before rename/close. - The corpus manager is built eagerly in newApplication (sync.Once dropped); test helper dead branch removed. Assisted-by: Claude:claude-fable-5 [Claude Code] --- core/application/application.go | 7 +- core/application/router_factories.go | 11 +- core/backend/stores.go | 30 +-- core/config/model_config.go | 12 + core/http/endpoints/localai/router_corpus.go | 78 +++---- .../endpoints/localai/router_corpus_test.go | 8 +- core/http/endpoints/openai/realtime_model.go | 14 +- core/http/middleware/route_model.go | 71 +++--- core/http/routes/anthropic.go | 12 +- core/http/routes/middleware.go | 96 +++----- core/http/routes/openai.go | 12 +- core/services/routing/corpus/manager.go | 215 +++++++++++++++--- core/services/routing/corpus/manager_test.go | 46 ++++ core/services/routing/corpus/resolve.go | 93 ++++++++ core/services/routing/router/knn.go | 60 ++--- core/services/routing/router/types.go | 7 + pkg/mcp/localaitools/inproc/client.go | 53 +---- 17 files changed, 497 insertions(+), 328 deletions(-) create mode 100644 core/services/routing/corpus/resolve.go diff --git a/core/application/application.go b/core/application/application.go index d7b51c1fe687..f2bb5f8bdc69 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -1,8 +1,10 @@ package application import ( + "cmp" "context" "math/rand/v2" + "path/filepath" "sync" "sync/atomic" "time" @@ -78,7 +80,6 @@ type Application struct { routerDecisions router.DecisionStore routerRegistry *router.Registry routerCorpus *corpus.Manager - routerCorpusOnce sync.Once admissionLimiter *admission.Limiter watchdogMutex sync.Mutex watchdogStop chan bool @@ -128,6 +129,10 @@ func newApplication(appConfig *config.ApplicationConfig) *Application { applicationConfig: appConfig, templatesEvaluator: templates.NewEvaluator(appConfig.SystemState.Model.ModelsPath), voiceProfileStore: voiceprofile.NewStore(appConfig.DataPath), + // KNN corpus files live under /router-corpus (same + // DataPath → DynamicConfigsDir precedence the agent pool uses). + routerCorpus: corpus.NewManager(filepath.Join( + cmp.Or(appConfig.DataPath, appConfig.DynamicConfigsDir, "."), "router-corpus")), } // Face-recognition registry backed by LocalAI's built-in vector store. diff --git a/core/application/router_factories.go b/core/application/router_factories.go index 0c4bb936950a..2b86fd4c94f4 100644 --- a/core/application/router_factories.go +++ b/core/application/router_factories.go @@ -1,10 +1,8 @@ package application import ( - "cmp" "context" "fmt" - "path/filepath" "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" @@ -122,13 +120,8 @@ func (a *Application) VectorStore(storeName string) backend.VectorStore { return backend.NewVectorStore(a.modelLoader, a.applicationConfig, storeName) } -// RouterCorpus returns the process-wide KNN corpus manager. Corpus -// files live under /router-corpus (same DataPath → -// DynamicConfigsDir precedence the agent pool uses for its state). +// RouterCorpus returns the process-wide KNN corpus manager, built in +// newApplication. func (a *Application) RouterCorpus() *corpus.Manager { - a.routerCorpusOnce.Do(func() { - root := cmp.Or(a.applicationConfig.DataPath, a.applicationConfig.DynamicConfigsDir, ".") - a.routerCorpus = corpus.NewManager(filepath.Join(root, "router-corpus")) - }) return a.routerCorpus } diff --git a/core/backend/stores.go b/core/backend/stores.go index 9daa5c4a9b46..8321656b0bb9 100644 --- a/core/backend/stores.go +++ b/core/backend/stores.go @@ -50,27 +50,15 @@ func (s *localVectorStore) backend(_ context.Context) (grpc.Backend, error) { return StoreBackend(s.loader, s.appConfig, s.storeName, "") } -func (s *localVectorStore) Search(ctx context.Context, vec []float32) (sim float64, payload []byte, ok bool, err error) { - start := time.Now() - outcome := "hit" - defer func() { - s.recordTrace(start, "search", len(vec), sim, outcome, err) - }() - be, berr := s.backend(ctx) - if berr != nil { - outcome = "backend_load_error" - return 0, nil, false, fmt.Errorf("vector store load: %w", berr) - } - _, values, similarities, ferr := store.Find(ctx, be, vec, 1) - if ferr != nil { - outcome = "find_error" - return 0, nil, false, fmt.Errorf("vector store find: %w", ferr) - } - if len(values) == 0 || len(similarities) == 0 { - outcome = "miss" - return 0, nil, false, nil - } - return float64(similarities[0]), values[0], true, nil +// Search is the top-1 special case of SearchK; delegating keeps the +// backend-load/Find/trace plumbing in one place (SearchK records the +// identically-shaped trace, so /api/backend-traces sees no difference). +func (s *localVectorStore) Search(ctx context.Context, vec []float32) (float64, []byte, bool, error) { + neighbors, err := s.SearchK(ctx, vec, 1) + if err != nil || len(neighbors) == 0 { + return 0, nil, false, err + } + return neighbors[0].Similarity, neighbors[0].Payload, true, nil } func (s *localVectorStore) SearchK(ctx context.Context, vec []float32, k int) (neighbors []Neighbor, err error) { diff --git a/core/config/model_config.go b/core/config/model_config.go index 2c76446e9b67..a59077f22473 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -430,6 +430,18 @@ type RouterKNNConfig struct { StoreName string `yaml:"store_name,omitempty" json:"store_name,omitempty"` } +// ResolvedStoreName is the local-store collection this router's corpus +// actually lives in. The single source of the "router-corpus-" +// default — the classifier build, the corpus API, the status page, and +// the assistant MCP tools all resolve through here so they cannot +// desync on which store they touch. +func (k *RouterKNNConfig) ResolvedStoreName(routerName string) string { + if k.StoreName != "" { + return k.StoreName + } + return "router-corpus-" + routerName +} + // RouterPolicy is one entry in the label vocabulary. The label string // is what the classifier model emits and what candidates reference in // their Labels field; the description is the natural-language hint diff --git a/core/http/endpoints/localai/router_corpus.go b/core/http/endpoints/localai/router_corpus.go index bf2a8223055c..48c05f55f118 100644 --- a/core/http/endpoints/localai/router_corpus.go +++ b/core/http/endpoints/localai/router_corpus.go @@ -1,7 +1,7 @@ package localai import ( - "fmt" + "errors" "net/http" "github.com/labstack/echo/v4" @@ -17,37 +17,30 @@ import ( // curated programmatically and never entered through or displayed in // the UI — the inspection surface returns label counts, never texts. // -// All three handlers resolve the router model by path param and -// require it to declare a `router.knn` block; the store name and -// embedding model come from that config so the API can't desync from -// what the classifier actually queries. +// The handlers are transport adapters over corpus.ResolveKNNRouter and +// corpus.Seed — the same helpers the assistant MCP tools use — so the +// two surfaces cannot drift on model resolution, store naming, or +// seed validation. This layer only binds requests and maps the corpus +// package's sentinel errors to HTTP statuses. -// resolveKNNRouter loads the named model config and returns its router -// KNN settings (store name defaulted the same way buildClassifier -// defaults it). Echo-shaped errors for the three failure modes. +// resolveKNNRouter adapts corpus.ResolveKNNRouter to echo: name from +// the path param, sentinel errors to 404/400, the rest to 500. func resolveKNNRouter(c echo.Context, loader *config.ModelConfigLoader, appConfig *config.ApplicationConfig) (*config.ModelConfig, string, error) { name := c.Param("name") if name == "" { return nil, "", echo.NewHTTPError(http.StatusBadRequest, "router name is required") } - cfg, err := loader.LoadModelConfigFileByNameDefaultOptions(name, appConfig) - if err != nil { - return nil, "", echo.NewHTTPError(http.StatusInternalServerError, "failed to load model config: "+err.Error()) + cfg, storeName, err := corpus.ResolveKNNRouter(loader, appConfig, name) + switch { + case err == nil: + return cfg, storeName, nil + case errors.Is(err, corpus.ErrRouterNotFound): + return nil, "", echo.NewHTTPError(http.StatusNotFound, err.Error()) + case errors.Is(err, corpus.ErrNotKNNRouter): + return nil, "", echo.NewHTTPError(http.StatusBadRequest, err.Error()) + default: + return nil, "", echo.NewHTTPError(http.StatusInternalServerError, err.Error()) } - // A synthetic stub (no Name) means the model is unknown — see - // RouterDecideEndpoint for the discrimination rationale. - if cfg == nil || cfg.Name == "" { - return nil, "", echo.NewHTTPError(http.StatusNotFound, fmt.Sprintf("model %q not found", name)) - } - if cfg.Router.KNN == nil || cfg.Router.KNN.EmbeddingModel == "" { - return nil, "", echo.NewHTTPError(http.StatusBadRequest, - fmt.Sprintf("model %q has no router.knn block (set classifier: knn and knn.embedding_model first)", name)) - } - storeName := cfg.Router.KNN.StoreName - if storeName == "" { - storeName = "router-corpus-" + cfg.Name - } - return cfg, storeName, nil } // RouterCorpusAddEndpoint bulk-seeds a router's KNN corpus. Entries @@ -80,38 +73,19 @@ func RouterCorpusAddEndpoint(loader *config.ModelConfigLoader, appConfig *config if len(req.Entries) == 0 { return echo.NewHTTPError(http.StatusBadRequest, "entries is required") } - - // Labels must be declared policies — the same invariant - // candidate tables are validated against. Catching it here - // keeps a typo from silently creating an unroutable label. - declared := map[string]struct{}{} - for _, p := range cfg.Router.Policies { - declared[p.Label] = struct{}{} - } entries := make([]corpus.Entry, 0, len(req.Entries)) - for i, e := range req.Entries { - for _, l := range e.Labels { - if _, ok := declared[l]; !ok { - return echo.NewHTTPError(http.StatusBadRequest, - fmt.Sprintf("entry %d: label %q is not declared in router policies", i, l)) - } - } + for _, e := range req.Entries { entries = append(entries, corpus.Entry{Text: e.Text, Labels: e.Labels}) } - embedder := deps.Embedder(cfg.Router.KNN.EmbeddingModel) - if embedder == nil { - return echo.NewHTTPError(http.StatusBadRequest, - fmt.Sprintf("embedding_model %q not loadable", cfg.Router.KNN.EmbeddingModel)) - } - store := deps.VectorStore(storeName) - - added, skipped, err := mgr.Add(c.Request().Context(), storeName, cfg.Router.KNN.EmbeddingModel, embedder, store, entries) - if err != nil { - return echo.NewHTTPError(http.StatusInternalServerError, err.Error()) - } - stats, err := mgr.Stats(storeName) + added, skipped, stats, err := corpus.Seed(c.Request().Context(), mgr, cfg, storeName, + deps.Embedder, deps.VectorStore, entries) if err != nil { + // Undeclared labels and an unloadable embedding model are + // request/config mistakes, not server faults. + if errors.Is(err, corpus.ErrUndeclaredLabel) || errors.Is(err, corpus.ErrEmbedderUnavailable) { + return echo.NewHTTPError(http.StatusBadRequest, err.Error()) + } return echo.NewHTTPError(http.StatusInternalServerError, err.Error()) } return c.JSON(http.StatusOK, schema.RouterCorpusAddResponse{ diff --git a/core/http/endpoints/localai/router_corpus_test.go b/core/http/endpoints/localai/router_corpus_test.go index bfb042730b79..35625eca965a 100644 --- a/core/http/endpoints/localai/router_corpus_test.go +++ b/core/http/endpoints/localai/router_corpus_test.go @@ -115,13 +115,7 @@ var _ = Describe("Router corpus endpoints", func() { }) do := func(method, path, body string) *httptest.ResponseRecorder { - var rd *strings.Reader - if body == "" { - rd = strings.NewReader("") - } else { - rd = strings.NewReader(body) - } - req := httptest.NewRequest(method, path, rd) + req := httptest.NewRequest(method, path, strings.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() e.ServeHTTP(rec, req) diff --git a/core/http/endpoints/openai/realtime_model.go b/core/http/endpoints/openai/realtime_model.go index d980e8c34889..e150911edcda 100644 --- a/core/http/endpoints/openai/realtime_model.go +++ b/core/http/endpoints/openai/realtime_model.go @@ -548,23 +548,13 @@ func buildRealtimeRoutingContext(a *application.Application, sessionID string) * if a == nil { return nil } - deps := &middleware.ClassifierDeps{ - Scorer: a.Scorer, - Corpus: a.RouterCorpus(), - TokenCounter: a.TokenCounter, - Embedder: a.Embedder, - VectorStore: a.VectorStore, - Reranker: a.Reranker, - ModelLookup: a.ModelConfigLookup(), - Registry: a.RouterClassifierRegistry(), - Evaluator: a.TemplatesEvaluator(), - } + deps := middleware.NewClassifierDeps(a) userID := "" if u := a.FallbackUser(); u != nil { userID = u.ID } return &RealtimeRoutingContext{ - Deps: deps, + Deps: &deps, Store: a.RouterDecisions(), SessionID: sessionID, UserID: userID, diff --git a/core/http/middleware/route_model.go b/core/http/middleware/route_model.go index a7e3255fb39d..8527f8608f93 100644 --- a/core/http/middleware/route_model.go +++ b/core/http/middleware/route_model.go @@ -11,6 +11,7 @@ import ( "time" "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/auth" @@ -109,6 +110,26 @@ type ClassifierDeps struct { TokenCounter func(modelName string) func(text string) (int, error) } +// NewClassifierDeps assembles the full classifier dependency set from +// the application container. Every entry point that runs the router — +// OpenAI chat, Anthropic messages, realtime, the decision oracle, the +// corpus API — builds its deps here, so a new ClassifierDeps field is +// wired once instead of hand-copied per call site (where a missed site +// silently degrades that path rather than failing). +func NewClassifierDeps(app *application.Application) ClassifierDeps { + return ClassifierDeps{ + Scorer: app.Scorer, + Corpus: app.RouterCorpus(), + TokenCounter: app.TokenCounter, + Embedder: app.Embedder, + VectorStore: app.VectorStore, + Reranker: app.Reranker, + ModelLookup: app.ModelConfigLookup(), + Registry: app.RouterClassifierRegistry(), + Evaluator: app.TemplatesEvaluator(), + } +} + // ProbeExtractor pulls the prompt content out of a parsed request so // the classifier can inspect it without taking a dependency on the // schema package. One extractor per request shape — wired by the @@ -326,6 +347,9 @@ func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Class var inner router.Classifier switch name { case router.ClassifierScore: + if rc.ClassifierModel == "" { + return nil, fmt.Errorf("router classifier score requires classifier_model") + } if deps.Scorer == nil { return nil, fmt.Errorf("router classifier score unavailable: no scorer factory wired") } @@ -377,6 +401,9 @@ func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Class } inner = router.NewScoreClassifier(policies, scorer, opts) case router.ClassifierColbert: + if rc.ClassifierModel == "" { + return nil, fmt.Errorf("router classifier colbert requires classifier_model") + } if deps.Reranker == nil { return nil, fmt.Errorf("router classifier colbert unavailable: no reranker factory wired") } @@ -400,10 +427,7 @@ func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Class if embedder == nil { return nil, fmt.Errorf("router classifier knn: embedding_model %q not loadable", rc.KNN.EmbeddingModel) } - storeName := rc.KNN.StoreName - if storeName == "" { - storeName = "router-corpus-" + cfg.Name - } + storeName := rc.KNN.ResolvedStoreName(cfg.Name) vstore := deps.VectorStore(storeName) if vstore == nil { return nil, fmt.Errorf("router classifier knn: vector store %q not loadable", storeName) @@ -428,22 +452,24 @@ func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Class if count, ctxTokens := modelTokenTrim(rc.KNN.EmbeddingModel, deps); count != nil { knnClassifier = knnClassifier.WithTokenTrim(count, ctxTokens) } - inner = knnClassifier + if rc.EmbeddingCache != nil { + // The knn classifier IS an embedding-KNN lookup — wrapping it + // in the embedding cache would embed every probe twice to + // answer the same question. Ignore the block rather than fail + // routing. Returning from the arm (instead of name-checking in + // the shared wrap below) keeps the opt-out local: the next + // embedding-based classifier makes its own wrapping decision. + xlog.Warn("router: embedding_cache ignored for knn classifier", + "router_model", cfg.Name) + } + return knnClassifier, nil default: - return nil, fmt.Errorf("router: unknown classifier %q (supported: %s)", name, strings.Join([]string{router.ClassifierScore, router.ClassifierColbert, router.ClassifierKNN}, ", ")) + return nil, fmt.Errorf("router: unknown classifier %q (supported: %s)", name, strings.Join(router.AllClassifiers, ", ")) } if rc.EmbeddingCache == nil { return inner, nil } - if name == router.ClassifierKNN { - // The knn classifier IS an embedding-KNN lookup — wrapping it in - // the embedding cache would embed every probe twice to answer the - // same question. Ignore the block rather than fail routing. - xlog.Warn("router: embedding_cache ignored for knn classifier", - "router_model", cfg.Name) - return inner, nil - } wrapped, err := wrapWithEmbeddingCache(cfg, inner, deps) if err != nil { // Caching plumbing problems must not break routing — log, @@ -484,19 +510,12 @@ func assertClassifierDeclaresScore(classifierModel string, lookup ModelConfigLoo return nil } -// validateRouterPolicies checks the shared invariants both classifiers -// rely on (non-empty policies, every candidate label declared as a -// policy, every candidate has a model + at least one label) and -// returns the parsed []ScorePolicy. Both Score and Rerank classifiers -// take the same policy shape. +// validateRouterPolicies checks the invariants every classifier relies +// on (non-empty policies, every candidate label declared as a policy, +// every candidate has a model + at least one label) and returns the +// parsed []ScorePolicy. Per-classifier requirements (classifier_model, +// knn block, ...) live in the buildClassifier arms, not here. func validateRouterPolicies(classifierName string, rc config.RouterConfig) ([]router.ScorePolicy, error) { - // The knn classifier has no scoring model — its label knowledge - // lives in the corpus. It still declares policies: they are the - // label vocabulary corpus entries are validated against and the - // "categories" surface the Routing tab shows. - if rc.ClassifierModel == "" && classifierName != router.ClassifierKNN { - return nil, fmt.Errorf("router classifier %s requires classifier_model", classifierName) - } if len(rc.Policies) == 0 { return nil, fmt.Errorf("router classifier %s requires at least one policy", classifierName) } diff --git a/core/http/routes/anthropic.go b/core/http/routes/anthropic.go index 4565a4e7815a..124557655eb2 100644 --- a/core/http/routes/anthropic.go +++ b/core/http/routes/anthropic.go @@ -55,17 +55,7 @@ func RegisterAnthropicRoutes(app *echo.Echo, application.FallbackUser(), middleware.AnthropicProbe, router.SourceAnthropic, - middleware.ClassifierDeps{ - Scorer: application.Scorer, - Corpus: application.RouterCorpus(), - TokenCounter: application.TokenCounter, - Embedder: application.Embedder, - VectorStore: application.VectorStore, - Reranker: application.Reranker, - ModelLookup: application.ModelConfigLookup(), - Registry: application.RouterClassifierRegistry(), - Evaluator: application.TemplatesEvaluator(), - }, + middleware.NewClassifierDeps(application), ), middleware.AdmissionControl(application.AdmissionLimiter(), application.PIIEvents()), pii.RequestMiddleware(application.PIIRedactor(), application.PIIEvents(), piiadapter.Anthropic(), application.FallbackUser(), pii.WithNERResolver(application.PIINERResolver()), pii.WithPolicyResolver(application.PIIPolicyResolver())), diff --git a/core/http/routes/middleware.go b/core/http/routes/middleware.go index b22c7874c4db..0195f445e255 100644 --- a/core/http/routes/middleware.go +++ b/core/http/routes/middleware.go @@ -29,15 +29,7 @@ import ( // the synthetic local user has Role: admin so the page works without // extra config — same gating shape as the existing /api/usage/all. func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { - e.GET("/api/middleware/status", func(c echo.Context) error { - viewer := resolveUsageUser(c, app) - if viewer == nil { - return c.JSON(http.StatusUnauthorized, map[string]string{"error": "not authenticated"}) - } - if viewer.Role != auth.RoleAdmin { - return c.JSON(http.StatusForbidden, map[string]string{"error": "admin access required"}) - } - + e.GET("/api/middleware/status", adminOnly(app, func(c echo.Context) error { piiSection := buildPIIStatus(app) routerSection := buildRouterStatus(app) mitmSection := buildMITMStatus(app) @@ -49,7 +41,7 @@ func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { "mitm": mitmSection, "admission": admissionSection, }) - }) + })) e.GET("/api/router/status", func(c echo.Context) error { // Read-only — admins want to see classifier configurations @@ -74,18 +66,10 @@ func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { return c.Blob(http.StatusOK, "application/x-pem-file", ca.PublicCertPEM()) }) - e.GET("/api/router/decisions", func(c echo.Context) error { - viewer := resolveUsageUser(c, app) - if viewer == nil { - return c.JSON(http.StatusUnauthorized, map[string]string{"error": "not authenticated"}) - } - // Decision logs may include user ids — admin-only when auth is - // on; the synthetic local user has admin so single-user mode - // works. - if viewer.Role != auth.RoleAdmin { - return c.JSON(http.StatusForbidden, map[string]string{"error": "admin access required"}) - } - + // Decision logs may include user ids — admin-only when auth is + // on; the synthetic local user has admin so single-user mode + // works. + e.GET("/api/router/decisions", adminOnly(app, func(c echo.Context) error { store := app.RouterDecisions() if store == nil { return c.JSON(http.StatusOK, map[string]any{"decisions": []any{}}) @@ -107,7 +91,7 @@ func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { return c.JSON(http.StatusInternalServerError, map[string]string{"error": "failed to list decisions"}) } return c.JSON(http.StatusOK, map[string]any{"decisions": decisions}) - }) + })) // GET /api/router/cache/stats — embedding-cache counters per // router model. Read-only; same auth gating as /api/router/status @@ -134,28 +118,9 @@ func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { decideHandler := localai.RouterDecideEndpoint( app.ModelConfigLoader(), app.ApplicationConfig(), - middleware.ClassifierDeps{ - Scorer: app.Scorer, - Corpus: app.RouterCorpus(), - TokenCounter: app.TokenCounter, - Embedder: app.Embedder, - VectorStore: app.VectorStore, - Reranker: app.Reranker, - ModelLookup: app.ModelConfigLookup(), - Registry: app.RouterClassifierRegistry(), - Evaluator: app.TemplatesEvaluator(), - }, + middleware.NewClassifierDeps(app), ) - e.POST("/api/router/decide", func(c echo.Context) error { - viewer := resolveUsageUser(c, app) - if viewer == nil { - return c.JSON(http.StatusUnauthorized, map[string]string{"error": "not authenticated"}) - } - if viewer.Role != auth.RoleAdmin { - return c.JSON(http.StatusForbidden, map[string]string{"error": "admin access required"}) - } - return decideHandler(c) - }) + e.POST("/api/router/decide", adminOnly(app, decideHandler)) // Router KNN corpus management. Corpus input/curation is API-only // by design — entries can contain example user content, so the UI @@ -163,28 +128,30 @@ func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { // counts only. Admin-gated like /api/router/decide: seeding the // corpus changes routing behaviour and Add runs embedding // inference on arbitrary input. - corpusDeps := middleware.ClassifierDeps{ - Embedder: app.Embedder, - VectorStore: app.VectorStore, - } + corpusDeps := middleware.NewClassifierDeps(app) corpusAdd := localai.RouterCorpusAddEndpoint(app.ModelConfigLoader(), app.ApplicationConfig(), app.RouterCorpus(), corpusDeps) corpusStats := localai.RouterCorpusStatsEndpoint(app.ModelConfigLoader(), app.ApplicationConfig(), app.RouterCorpus()) corpusClear := localai.RouterCorpusClearEndpoint(app.ModelConfigLoader(), app.ApplicationConfig(), app.RouterCorpus(), corpusDeps) - adminGate := func(h echo.HandlerFunc) echo.HandlerFunc { - return func(c echo.Context) error { - viewer := resolveUsageUser(c, app) - if viewer == nil { - return c.JSON(http.StatusUnauthorized, map[string]string{"error": "not authenticated"}) - } - if viewer.Role != auth.RoleAdmin { - return c.JSON(http.StatusForbidden, map[string]string{"error": "admin access required"}) - } - return h(c) + e.POST("/api/router/:name/corpus", adminOnly(app, corpusAdd)) + e.GET("/api/router/:name/corpus/stats", adminOnly(app, corpusStats)) + e.DELETE("/api/router/:name/corpus", adminOnly(app, corpusClear)) +} + +// adminOnly wraps a handler with the admin gate every routing-module +// endpoint shares: 401 when unauthenticated, 403 for non-admins. The +// synthetic local user has Role: admin, so single-user (no-auth) mode +// passes without extra config — same shape as /api/usage/all. +func adminOnly(app *application.Application, h echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + viewer := resolveUsageUser(c, app) + if viewer == nil { + return c.JSON(http.StatusUnauthorized, map[string]string{"error": "not authenticated"}) } + if viewer.Role != auth.RoleAdmin { + return c.JSON(http.StatusForbidden, map[string]string{"error": "admin access required"}) + } + return h(c) } - e.POST("/api/router/:name/corpus", adminGate(corpusAdd)) - e.GET("/api/router/:name/corpus/stats", adminGate(corpusStats)) - e.DELETE("/api/router/:name/corpus", adminGate(corpusClear)) } // buildRouterStatus inventories every model that declares a Router @@ -241,10 +208,7 @@ func buildRouterStatus(app *application.Application) map[string]any { entry["embedding_cache"] = cacheEntry } if kc := cfg.Router.KNN; kc != nil { - storeName := kc.StoreName - if storeName == "" { - storeName = "router-corpus-" + cfg.Name - } + storeName := kc.ResolvedStoreName(cfg.Name) knnEntry := map[string]any{ "embedding_model": kc.EmbeddingModel, "k": kc.K, @@ -276,7 +240,7 @@ func buildRouterStatus(app *application.Application) map[string]any { "configured": hasAny, "models": models, "recent_decision_count": recentCount, - "available_classifiers": []string{router.ClassifierScore, router.ClassifierColbert, router.ClassifierKNN}, + "available_classifiers": router.AllClassifiers, } if !hasAny { out["note"] = "No router models configured. Add a `router:` block to a model YAML to enable intelligent routing." diff --git a/core/http/routes/openai.go b/core/http/routes/openai.go index 191a54ca1f1f..0820f4e31a49 100644 --- a/core/http/routes/openai.go +++ b/core/http/routes/openai.go @@ -71,17 +71,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, application.FallbackUser(), middleware.OpenAIProbe, router.SourceChat, - middleware.ClassifierDeps{ - Scorer: application.Scorer, - Corpus: application.RouterCorpus(), - TokenCounter: application.TokenCounter, - Embedder: application.Embedder, - VectorStore: application.VectorStore, - Reranker: application.Reranker, - ModelLookup: application.ModelConfigLookup(), - Registry: application.RouterClassifierRegistry(), - Evaluator: application.TemplatesEvaluator(), - }, + middleware.NewClassifierDeps(application), ), // Admission control runs after RouteModel so the SERVED // model's limits apply — a router fanout that lands on a diff --git a/core/services/routing/corpus/manager.go b/core/services/routing/corpus/manager.go index 8f0bc35cd799..3ace526e5466 100644 --- a/core/services/routing/corpus/manager.go +++ b/core/services/routing/corpus/manager.go @@ -25,6 +25,7 @@ import ( "sort" "strings" "sync" + "time" "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/services/routing/router" @@ -77,12 +78,39 @@ type Manager struct { // embedding-model-changed-on-a-live-index case, which local-store // cannot serve (it enforces one key length per store). loadedModel map[string]string + + // statsCache memoises Stats per store, keyed by the corpus file's + // stat fingerprint — the file is the source of truth, so a matching + // fingerprint means the counts are current without re-parsing the + // (vector-laden) JSONL. Its own mutex, NOT m.mu: the status page + // polls Stats every few seconds and must not block behind a seed + // that is busy embedding under m.mu. + statsMu sync.Mutex + statsCache map[string]cachedStats +} + +type cachedStats struct { + key fileFingerprint + stats Stats +} + +// fileFingerprint identifies a corpus file version cheaply (one +// os.Stat). Size changes on every append and rename replaces the +// inode's mtime, so any successful write bumps the fingerprint. +type fileFingerprint struct { + exists bool + size int64 + modTime time.Time } // NewManager roots corpus files at dir (created lazily on first // write). func NewManager(dir string) *Manager { - return &Manager{dir: dir, loadedModel: map[string]string{}} + return &Manager{ + dir: dir, + loadedModel: map[string]string{}, + statsCache: map[string]cachedStats{}, + } } // path maps a store name to its corpus file. Store names come from @@ -119,7 +147,7 @@ func (m *Manager) EnsureLoaded(ctx context.Context, storeName, embeddingModel st return 0, fmt.Errorf("corpus %q was indexed with embedding model %q this process; restart LocalAI to re-index it with %q", storeName, prev, embeddingModel) } - entries, err := m.read(storeName) + entries, _, err := m.readAll(storeName) if err != nil { return 0, err } @@ -163,6 +191,12 @@ func (m *Manager) EnsureLoaded(ctx context.Context, storeName, embeddingModel st // Returns (added, skipped). The file write happens before the index // insert: if indexing fails the entries are still durable and the next // EnsureLoaded (or restart) syncs them. +// +// The embedding calls — the slow part of a big seed — run OUTSIDE +// m.mu, so a long seed doesn't freeze EnsureLoaded and other stores' +// corpus operations for its whole duration. The commit re-checks for +// duplicates under the lock, so a racing Add can at worst turn an +// entry into a skip, never a double insert. func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, embedder backend.Embedder, store backend.VectorStore, entries []Entry) (int, int, error) { if embedder == nil { return 0, 0, fmt.Errorf("corpus %q: no embedder available", storeName) @@ -176,19 +210,21 @@ func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, emb } } + // First pass: dedupe against the persisted corpus (and within the + // request) so known texts aren't re-embedded. m.mu.Lock() - defer m.mu.Unlock() - - existing, err := m.read(storeName) + existing, _, err := m.readAll(storeName) if err != nil { + m.mu.Unlock() return 0, 0, err } seen := make(map[string]struct{}, len(existing)) for _, e := range existing { seen[e.Text] = struct{}{} } + m.mu.Unlock() - added := make([]Entry, 0, len(entries)) + candidates := make([]Entry, 0, len(entries)) skipped := 0 for _, e := range entries { if _, dup := seen[e.Text]; dup { @@ -196,19 +232,56 @@ func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, emb continue } seen[e.Text] = struct{}{} - vec, err := embedder.Embed(ctx, e.Text) + candidates = append(candidates, e) + } + if len(candidates) == 0 { + return 0, skipped, nil + } + + for i := range candidates { + vec, err := embedder.Embed(ctx, candidates[i].Text) if err != nil { - return 0, skipped, fmt.Errorf("corpus %q: embedding %q-labelled entry: %w", storeName, e.Labels[0], err) + return 0, skipped, fmt.Errorf("corpus %q: embedding %q-labelled entry: %w", storeName, candidates[i].Labels[0], err) + } + candidates[i].Vector = vec + candidates[i].EmbeddingModel = embeddingModel + } + + m.mu.Lock() + defer m.mu.Unlock() + + // Commit: re-read to catch texts a concurrent Add landed while we + // were embedding, then append only the survivors — JSONL appends + // keep the write O(new entries), not O(corpus). + current, torn, err := m.readAll(storeName) + if err != nil { + return 0, skipped, err + } + seen = make(map[string]struct{}, len(current)) + for _, e := range current { + seen[e.Text] = struct{}{} + } + added := candidates[:0] + for _, e := range candidates { + if _, dup := seen[e.Text]; dup { + skipped++ + continue } - e.Vector = vec - e.EmbeddingModel = embeddingModel added = append(added, e) } if len(added) == 0 { return 0, skipped, nil } - if err := m.write(storeName, append(existing, added...)); err != nil { + if torn { + // A crash mid-append left a torn final line; appending after it + // would bury the damage mid-file. Repair with a full atomic + // rewrite of the readable entries plus the new ones. + err = m.write(storeName, append(current, added...)) + } else { + err = m.appendEntries(storeName, added) + } + if err != nil { return 0, skipped, err } if store != nil { @@ -223,11 +296,24 @@ func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, emb // Stats reports label counts for the persisted corpus. Never returns // entry texts. +// +// Deliberately does NOT take m.mu — the status page polls this every +// few seconds and must not block behind a seed or load. Reading the +// file without the write lock is safe: write() replaces it atomically +// (rename) and appendEntries only extends the tail, whose worst-case +// torn line readAll tolerates. The fingerprint is taken BEFORE the +// read so a write racing the read at worst caches fresher stats under +// an older key, which the next call corrects. func (m *Manager) Stats(storeName string) (Stats, error) { - m.mu.Lock() - defer m.mu.Unlock() + key := m.fingerprint(storeName) + m.statsMu.Lock() + if c, ok := m.statsCache[storeName]; ok && c.key.equal(key) { + m.statsMu.Unlock() + return c.stats.copy(), nil + } + m.statsMu.Unlock() - entries, err := m.read(storeName) + entries, _, err := m.readAll(storeName) if err != nil { return Stats{}, err } @@ -245,9 +331,37 @@ func (m *Manager) Stats(storeName string) (Stats, error) { s.EmbeddingModels = append(s.EmbeddingModels, mn) } sort.Strings(s.EmbeddingModels) + + m.statsMu.Lock() + m.statsCache[storeName] = cachedStats{key: key, stats: s.copy()} + m.statsMu.Unlock() return s, nil } +func (m *Manager) fingerprint(storeName string) fileFingerprint { + fi, err := os.Stat(m.path(storeName)) + if err != nil { + return fileFingerprint{} + } + return fileFingerprint{exists: true, size: fi.Size(), modTime: fi.ModTime()} +} + +func (f fileFingerprint) equal(o fileFingerprint) bool { + return f.exists == o.exists && f.size == o.size && f.modTime.Equal(o.modTime) +} + +// copy protects the cached maps/slices from mutation by callers (and +// callers from later cache updates). +func (s Stats) copy() Stats { + out := s + out.LabelCounts = make(map[string]int, len(s.LabelCounts)) + for k, v := range s.LabelCounts { + out.LabelCounts[k] = v + } + out.EmbeddingModels = append([]string(nil), s.EmbeddingModels...) + return out +} + // Clear deletes the corpus file and removes its vectors from the live // index when the store supports deletion. When it doesn't, the index // keeps serving stale entries until restart — the returned count and @@ -256,7 +370,7 @@ func (m *Manager) Clear(ctx context.Context, storeName string, store backend.Vec m.mu.Lock() defer m.mu.Unlock() - entries, err := m.read(storeName) + entries, _, err := m.readAll(storeName) if err != nil { return 0, err } @@ -308,44 +422,58 @@ func insertAll(ctx context.Context, store backend.VectorStore, entries []Entry) return nil } -// read returns the persisted entries for storeName; a missing file is -// an empty corpus. Callers hold m.mu. -func (m *Manager) read(storeName string) ([]Entry, error) { +// readAll returns the persisted entries for storeName; a missing file +// is an empty corpus. The torn flag reports an unparseable FINAL line +// — the signature a crash mid-append (or a read racing an in-flight +// append) leaves — which is dropped rather than treated as corruption; +// Add repairs it on the next write. An unparseable line anywhere else +// is real corruption and errors. +func (m *Manager) readAll(storeName string) (entries []Entry, torn bool, err error) { f, err := os.Open(m.path(storeName)) if os.IsNotExist(err) { - return nil, nil + return nil, false, nil } if err != nil { - return nil, err + return nil, false, err } defer func() { _ = f.Close() }() - var entries []Entry sc := bufio.NewScanner(f) // Vectors inline in JSON push line length well past the default // 64KiB scanner cap (a 4096-dim float32 vector is ~50KiB of JSON // alone). 16MiB bounds any realistic embedding width. sc.Buffer(make([]byte, 0, 1<<20), 16<<20) line := 0 + var pendingErr error for sc.Scan() { line++ raw := strings.TrimSpace(sc.Text()) if raw == "" { continue } + if pendingErr != nil { + // The failed line has a successor, so it wasn't a torn tail. + return nil, false, pendingErr + } var e Entry if err := json.Unmarshal([]byte(raw), &e); err != nil { - return nil, fmt.Errorf("corpus file %s line %d: %w", m.path(storeName), line, err) + pendingErr = fmt.Errorf("corpus file %s line %d: %w", m.path(storeName), line, err) + continue } entries = append(entries, e) } if err := sc.Err(); err != nil { - return nil, err + return nil, false, err } - return entries, nil + if pendingErr != nil { + xlog.Warn("corpus: dropping torn final line (crash or concurrent append); repaired on next write", + "file", m.path(storeName), "line", line) + return entries, true, nil + } + return entries, false, nil } -// write atomically replaces the corpus file (tmp + rename) so a crash -// mid-write can't truncate the corpus. Callers hold m.mu. +// write atomically replaces the corpus file (tmp + fsync + rename) so +// a crash mid-write can't truncate the corpus. Callers hold m.mu. func (m *Manager) write(storeName string, entries []Entry) error { if err := os.MkdirAll(m.dir, 0o750); err != nil { return err @@ -368,8 +496,43 @@ func (m *Manager) write(storeName string, entries []Entry) error { _ = tmp.Close() return err } + if err := tmp.Sync(); err != nil { + _ = tmp.Close() + return err + } if err := tmp.Close(); err != nil { return err } return os.Rename(tmp.Name(), target) } + +// appendEntries extends the corpus file in place — JSONL's whole point +// — so Add costs O(new entries), not O(corpus). A crash mid-append +// tears at most the final line, which readAll tolerates and the next +// Add rewrite repairs. Callers hold m.mu. +func (m *Manager) appendEntries(storeName string, entries []Entry) error { + if err := os.MkdirAll(m.dir, 0o750); err != nil { + return err + } + f, err := os.OpenFile(m.path(storeName), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o640) + if err != nil { + return err + } + w := bufio.NewWriter(f) + enc := json.NewEncoder(w) + for _, e := range entries { + if err := enc.Encode(e); err != nil { + _ = f.Close() + return err + } + } + if err := w.Flush(); err != nil { + _ = f.Close() + return err + } + if err := f.Sync(); err != nil { + _ = f.Close() + return err + } + return f.Close() +} diff --git a/core/services/routing/corpus/manager_test.go b/core/services/routing/corpus/manager_test.go index 9c69581bd65b..0735d7e11ba3 100644 --- a/core/services/routing/corpus/manager_test.go +++ b/core/services/routing/corpus/manager_test.go @@ -228,6 +228,52 @@ var _ = Describe("corpus.Manager", func() { Expect(st.Total).To(BeZero()) }) + It("tolerates a torn final line and repairs it on the next add", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed[:2]) + Expect(err).NotTo(HaveOccurred()) + + // Simulate a crash mid-append: an unparseable tail. + path := filepath.Join(dir, storeName+".jsonl") + f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o640) + Expect(err).NotTo(HaveOccurred()) + _, err = f.WriteString(`{"text":"torn`) + Expect(err).NotTo(HaveOccurred()) + Expect(f.Close()).To(Succeed()) + + // Reads drop the torn line instead of failing the corpus. + st, err := mgr.Stats(storeName) + Expect(err).NotTo(HaveOccurred()) + Expect(st.Total).To(Equal(2)) + + // The next add rewrites the file whole, clearing the damage. + added, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed[2:]) + Expect(err).NotTo(HaveOccurred()) + Expect(added).To(Equal(1)) + raw, err := os.ReadFile(path) + Expect(err).NotTo(HaveOccurred()) + Expect(string(raw)).NotTo(ContainSubstring("torn")) + + st, err = mgr.Stats(storeName) + Expect(err).NotTo(HaveOccurred()) + Expect(st.Total).To(Equal(3)) + }) + + It("fails on corruption that is not a torn tail", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed[:1]) + Expect(err).NotTo(HaveOccurred()) + + // Garbage with a valid successor line is real corruption, not + // an interrupted append — reads must refuse to guess. + path := filepath.Join(dir, storeName+".jsonl") + raw, err := os.ReadFile(path) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(path, append([]byte("garbage\n"), raw...), 0o640)).To(Succeed()) + + _, err = mgr.Stats(storeName) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("line 1")) + }) + It("sanitises hostile store names into the corpus dir", func() { hostile := "../../etc/passwd" _, _, err := mgr.Add(ctx, hostile, "embed-1", embedder, store, seed[:1]) diff --git a/core/services/routing/corpus/resolve.go b/core/services/routing/corpus/resolve.go new file mode 100644 index 000000000000..13bf976dadf0 --- /dev/null +++ b/core/services/routing/corpus/resolve.go @@ -0,0 +1,93 @@ +package corpus + +import ( + "context" + "errors" + "fmt" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" +) + +// The REST corpus endpoints and the assistant MCP tools are thin +// transport adapters over the two helpers in this file, so the "what +// counts as a KNN router" rule and the seed-time validation cannot +// drift between the two surfaces. Sentinel errors let the HTTP layer +// map failure modes to status codes without string matching. +// The sentinel texts are sentence fragments so the wrapped messages +// read as the API's established error strings. +var ( + // ErrRouterNotFound: the named model doesn't exist (HTTP 404). + ErrRouterNotFound = errors.New("not found") + // ErrNotKNNRouter: the model exists but declares no router.knn + // block (HTTP 400). + ErrNotKNNRouter = errors.New("has no router.knn block (set classifier: knn and knn.embedding_model first)") + // ErrUndeclaredLabel: a seed entry uses a label that is not a + // declared router policy (HTTP 400). + ErrUndeclaredLabel = errors.New("is not declared in router policies") + // ErrEmbedderUnavailable: the router's knn.embedding_model can't + // be loaded (HTTP 400 — a config problem, not a server fault). + ErrEmbedderUnavailable = errors.New("not loadable") +) + +// ResolveKNNRouter loads the named model config, requires it to declare +// a router.knn block, and resolves the corpus store name the classifier +// build uses — the single resolution path for every corpus surface. +// +// LoadModelConfigFileByName returns a synthetic stub (empty Name) when +// the model is unknown; that maps to ErrRouterNotFound so callers can +// distinguish "unknown model" from "known but not a KNN router". +func ResolveKNNRouter(loader *config.ModelConfigLoader, appConfig *config.ApplicationConfig, name string) (*config.ModelConfig, string, error) { + cfg, err := loader.LoadModelConfigFileByNameDefaultOptions(name, appConfig) + if err != nil { + return nil, "", fmt.Errorf("failed to load model config: %w", err) + } + if cfg == nil || cfg.Name == "" { + return nil, "", fmt.Errorf("model %q %w", name, ErrRouterNotFound) + } + if cfg.Router.KNN == nil || cfg.Router.KNN.EmbeddingModel == "" { + return nil, "", fmt.Errorf("model %q %w", name, ErrNotKNNRouter) + } + return cfg, cfg.Router.KNN.ResolvedStoreName(cfg.Name), nil +} + +// Seed validates entries against the router's declared policy labels +// (the same invariant candidate tables are validated against — a typo +// must not silently create an unroutable label), embeds and persists +// them via the Manager, and returns the post-seed stats. +func Seed(ctx context.Context, mgr *Manager, cfg *config.ModelConfig, storeName string, + embedderFor func(string) backend.Embedder, storeFor func(string) backend.VectorStore, + entries []Entry) (added, skipped int, stats Stats, err error) { + + declared := map[string]struct{}{} + for _, p := range cfg.Router.Policies { + declared[p.Label] = struct{}{} + } + for i, e := range entries { + for _, l := range e.Labels { + if _, ok := declared[l]; !ok { + return 0, 0, Stats{}, fmt.Errorf("entry %d: label %q %w", i, l, ErrUndeclaredLabel) + } + } + } + + embeddingModel := cfg.Router.KNN.EmbeddingModel + var embedder backend.Embedder + if embedderFor != nil { + embedder = embedderFor(embeddingModel) + } + if embedder == nil { + return 0, 0, Stats{}, fmt.Errorf("embedding_model %q %w", embeddingModel, ErrEmbedderUnavailable) + } + var store backend.VectorStore + if storeFor != nil { + store = storeFor(storeName) + } + + added, skipped, err = mgr.Add(ctx, storeName, embeddingModel, embedder, store, entries) + if err != nil { + return added, skipped, Stats{}, err + } + stats, err = mgr.Stats(storeName) + return added, skipped, stats, err +} diff --git a/core/services/routing/router/knn.go b/core/services/routing/router/knn.go index f496551a9971..2203278cbd12 100644 --- a/core/services/routing/router/knn.go +++ b/core/services/routing/router/knn.go @@ -115,55 +115,32 @@ func (c *KNNClassifier) Classify(ctx context.Context, p Probe) (Decision, error) return errDecision(start, fmt.Errorf("knn classifier search: %w", err)) } - // Record every retrieved neighbour — including sub-gate ones — before - // the filtering below reuses the slice's backing array. The decision - // log needs the full retrieval to make fallbacks diagnosable: "what - // WAS nearby, and how was it labelled". A corrupt payload still names - // its similarity so index corruption is visible rather than silent. + // One pass does double duty: refs records every retrieved neighbour + // — including sub-gate and corrupt-payload ones — so the decision + // log makes fallbacks diagnosable ("what WAS nearby, and how was it + // labelled"; a corrupt payload still names its similarity so index + // corruption is visible rather than silent). The vote accumulates + // only neighbours that clear the epistemic gate: keeping + // sub-threshold neighbours out of the vote (rather than merely + // gating on the best one) stops far-away corpus regions from + // diluting a clear local majority. refs := make([]NeighborRef, 0, len(neighbors)) + votes := map[string]float64{} + best, total := 0.0, 0.0 for _, n := range neighbors { ref := NeighborRef{Similarity: n.Similarity} - if entry, ok := decodeCorpusEntry(n.Payload); ok { + entry, ok := decodeCorpusEntry(n.Payload) + if ok { ref.ID = entry.ID ref.Labels = entry.Labels } refs = append(refs, ref) - } - - // Epistemic gate: only neighbours the probe is genuinely close to - // may vote. Keeping sub-threshold neighbours out of the vote (rather - // than merely gating on the best one) stops far-away corpus regions - // from diluting a clear local majority. - best := 0.0 - usable := neighbors[:0] - for _, n := range neighbors { if n.Similarity > best { best = n.Similarity } - if n.Similarity >= c.similarityThreshold { - usable = append(usable, n) - } - } - if len(usable) == 0 { - // Out of corpus range — empty label set routes to the fallback - // via MatchCandidate's empty-active-set contract. Surfacing the - // best similarity in the decision log tells the admin whether - // the corpus needs entries near this probe or the threshold is - // simply too tight. - return Decision{ - NearestSimilarity: best, - Neighbors: refs, - ActivationThreshold: c.voteThreshold, - Latency: time.Since(start), - }, nil - } - - votes := map[string]float64{} - total := 0.0 - for _, n := range usable { - entry, ok := decodeCorpusEntry(n.Payload) - if !ok { - // A corrupt payload can't vote; it still counted toward K. + // Sub-gate neighbours and corrupt payloads can't vote (the + // latter still counted toward K). + if !ok || n.Similarity < c.similarityThreshold { continue } total += n.Similarity @@ -172,6 +149,11 @@ func (c *KNNClassifier) Classify(ctx context.Context, p Probe) (Decision, error) } } if total == 0 { + // Out of corpus range — empty label set routes to the fallback + // via MatchCandidate's empty-active-set contract. Surfacing the + // best similarity in the decision log tells the admin whether + // the corpus needs entries near this probe or the threshold is + // simply too tight. return Decision{ NearestSimilarity: best, Neighbors: refs, diff --git a/core/services/routing/router/types.go b/core/services/routing/router/types.go index bba1b9d399c4..02cc67256db0 100644 --- a/core/services/routing/router/types.go +++ b/core/services/routing/router/types.go @@ -164,6 +164,13 @@ const ( ClassifierKNN = "knn" ) +// AllClassifiers is the canonical classifier list. The middleware's +// unknown-classifier error and the /api/router/status +// available_classifiers field both derive from it, so a new classifier +// added to the buildClassifier switch shows up on every surface by +// extending this one slice (colbert once drifted out of both). +var AllClassifiers = []string{ClassifierScore, ClassifierColbert, ClassifierKNN} + // LabelFallback is the synthetic label written to the decision // store when the middleware uses cfg.Router.Fallback rather than a // classifier-picked candidate. diff --git a/pkg/mcp/localaitools/inproc/client.go b/pkg/mcp/localaitools/inproc/client.go index a74968928b5b..3147ad9fa7ad 100644 --- a/pkg/mcp/localaitools/inproc/client.go +++ b/pkg/mcp/localaitools/inproc/client.go @@ -896,32 +896,11 @@ func capabilityFlagsOf(m *config.ModelConfig) []string { return out } -// resolveKNNRouter mirrors the REST endpoint's resolution: the model -// must exist and declare a router.knn block; the store name defaults -// the same way buildClassifier defaults it. -func (c *Client) resolveKNNRouter(routerModel string) (*config.ModelConfig, string, error) { - cfg, err := c.ConfigLoader.LoadModelConfigFileByNameDefaultOptions(routerModel, c.AppConfig) - if err != nil { - return nil, "", fmt.Errorf("load model config: %w", err) - } - if cfg == nil || cfg.Name == "" { - return nil, "", fmt.Errorf("model %q not found", routerModel) - } - if cfg.Router.KNN == nil || cfg.Router.KNN.EmbeddingModel == "" { - return nil, "", fmt.Errorf("model %q has no router.knn block (set classifier: knn and knn.embedding_model first)", routerModel) - } - storeName := cfg.Router.KNN.StoreName - if storeName == "" { - storeName = "router-corpus-" + cfg.Name - } - return cfg, storeName, nil -} - func (c *Client) GetRouterCorpusStats(_ context.Context, routerModel string) (*localaitools.RouterCorpusStats, error) { if c.RouterCorpus == nil { return nil, errors.New("router corpus manager unavailable") } - cfg, storeName, err := c.resolveKNNRouter(routerModel) + cfg, storeName, err := corpus.ResolveKNNRouter(c.ConfigLoader, c.AppConfig, routerModel) if err != nil { return nil, err } @@ -943,36 +922,16 @@ func (c *Client) SeedRouterCorpus(ctx context.Context, req localaitools.RouterCo if c.RouterCorpus == nil || c.RouterEmbedder == nil || c.RouterVectorStore == nil { return nil, errors.New("router corpus manager unavailable") } - cfg, storeName, err := c.resolveKNNRouter(req.Router) + cfg, storeName, err := corpus.ResolveKNNRouter(c.ConfigLoader, c.AppConfig, req.Router) if err != nil { return nil, err } - - // Same invariant the REST endpoint enforces: labels must be - // declared policies, so a typo can't create an unroutable label. - declared := map[string]struct{}{} - for _, p := range cfg.Router.Policies { - declared[p.Label] = struct{}{} - } entries := make([]corpus.Entry, 0, len(req.Entries)) - for i, e := range req.Entries { - for _, l := range e.Labels { - if _, ok := declared[l]; !ok { - return nil, fmt.Errorf("entry %d: label %q is not declared in router policies", i, l) - } - } + for _, e := range req.Entries { entries = append(entries, corpus.Entry{Text: e.Text, Labels: e.Labels}) } - - embedder := c.RouterEmbedder(cfg.Router.KNN.EmbeddingModel) - if embedder == nil { - return nil, fmt.Errorf("embedding_model %q not loadable", cfg.Router.KNN.EmbeddingModel) - } - added, skipped, err := c.RouterCorpus.Add(ctx, storeName, cfg.Router.KNN.EmbeddingModel, embedder, c.RouterVectorStore(storeName), entries) - if err != nil { - return nil, err - } - stats, err := c.RouterCorpus.Stats(storeName) + added, skipped, stats, err := corpus.Seed(ctx, c.RouterCorpus, cfg, storeName, + c.RouterEmbedder, c.RouterVectorStore, entries) if err != nil { return nil, err } @@ -989,7 +948,7 @@ func (c *Client) ClearRouterCorpus(ctx context.Context, routerModel string) (*lo if c.RouterCorpus == nil { return nil, errors.New("router corpus manager unavailable") } - cfg, storeName, err := c.resolveKNNRouter(routerModel) + cfg, storeName, err := corpus.ResolveKNNRouter(c.ConfigLoader, c.AppConfig, routerModel) if err != nil { return nil, err } From 0db780a55022edf4c716d87d8961c18f9f382b98 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Wed, 15 Jul 2026 10:46:38 +0100 Subject: [PATCH 04/10] feat(router): bind knn corpus vectors to an embedder fingerprint and fail closed on mismatch Co-Authored-By: Claude Fable 5 --- core/application/application.go | 1 + core/application/router_factories.go | 60 ++++++ core/application/router_factories_test.go | 31 +++ core/config/meta/registry.go | 15 +- core/config/model_config.go | 17 ++ core/http/endpoints/localai/router_corpus.go | 11 +- .../endpoints/localai/router_corpus_test.go | 30 ++- core/http/middleware/route_model.go | 85 ++++++-- core/http/middleware/route_model_test.go | 39 ++++ core/services/routing/corpus/manager.go | 198 ++++++++++++++---- core/services/routing/corpus/manager_test.go | 121 ++++++++--- core/services/routing/corpus/resolve.go | 51 ++++- core/services/routing/router/knn.go | 24 ++- core/services/routing/router/knn_test.go | 15 ++ docs/content/features/middleware.md | 25 ++- pkg/mcp/localaitools/inproc/client.go | 11 +- 16 files changed, 606 insertions(+), 128 deletions(-) diff --git a/core/application/application.go b/core/application/application.go index f2bb5f8bdc69..e1b43dfa24f8 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -562,6 +562,7 @@ func (a *Application) start() error { // namespaces and model resolution. assistantClient.RouterCorpus = a.RouterCorpus() assistantClient.RouterEmbedder = a.Embedder + assistantClient.RouterEmbedderFingerprint = a.EmbedderFingerprint assistantClient.RouterVectorStore = a.VectorStore if err := holder.Initialize(a.applicationConfig.Context, assistantClient, localaitools.Options{}); err != nil { // Why log+continue instead of fail: the assistant is an optional diff --git a/core/application/router_factories.go b/core/application/router_factories.go index 2b86fd4c94f4..44db891eb8ee 100644 --- a/core/application/router_factories.go +++ b/core/application/router_factories.go @@ -2,11 +2,18 @@ package application import ( "context" + "crypto/sha256" + "encoding/hex" "fmt" + "os" + "path/filepath" + "sort" + "strings" "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/routing/corpus" + "gopkg.in/yaml.v3" ) // adapterConfig resolves a model name to its runtime ModelConfig, or nil when @@ -101,6 +108,59 @@ func (a *Application) Embedder(modelName string) backend.Embedder { return &lazyEmbedder{app: a, modelName: modelName} } +// EmbedderFingerprint returns a stable identity for the embedding space a +// named model currently produces. The effective config covers backend and +// embedding-affecting options; declared download checksums are part of that +// config, while local artifact stat data detects the common in-place file +// replacement case. Remote services without a stable artifact identity can +// use router.knn.embedding_revision to force invalidation explicitly. +func (a *Application) EmbedderFingerprint(modelName string) (string, error) { + cfg := a.adapterConfig(modelName) + if cfg == nil { + return "", fmt.Errorf("embedding model %q not available", modelName) + } + raw, err := yaml.Marshal(cfg) + if err != nil { + return "", fmt.Errorf("fingerprint embedding model %q config: %w", modelName, err) + } + h := sha256.New() + _, _ = h.Write(raw) + + paths := map[string]struct{}{} + for _, name := range append([]string{cfg.Model, cfg.MMProj}, downloadFileNames(cfg)...) { + if name == "" || strings.Contains(name, "://") { + continue + } + if !filepath.IsAbs(name) { + name = filepath.Join(a.applicationConfig.SystemState.Model.ModelsPath, name) + } + paths[filepath.Clean(name)] = struct{}{} + } + ordered := make([]string, 0, len(paths)) + for path := range paths { + ordered = append(ordered, path) + } + sort.Strings(ordered) + for _, path := range ordered { + _, _ = h.Write([]byte("\x00" + path)) + fi, statErr := os.Stat(path) + if statErr != nil { + _, _ = h.Write([]byte("\x00missing")) + continue + } + _, _ = fmt.Fprintf(h, "\x00%d\x00%d\x00%s", fi.Size(), fi.ModTime().UnixNano(), fi.Mode().String()) + } + return hex.EncodeToString(h.Sum(nil)), nil +} + +func downloadFileNames(cfg *config.ModelConfig) []string { + out := make([]string, 0, len(cfg.DownloadFiles)) + for _, f := range cfg.DownloadFiles { + out = append(out, f.Filename) + } + return out +} + type lazyEmbedder struct { app *Application modelName string diff --git a/core/application/router_factories_test.go b/core/application/router_factories_test.go index 5a6988a88fba..ae5495aa0bd2 100644 --- a/core/application/router_factories_test.go +++ b/core/application/router_factories_test.go @@ -94,6 +94,37 @@ var _ = Describe("router_factories lazy config resolution", func() { }) }) + Context("EmbedderFingerprint", func() { + It("changes when the effective config or local artifact changes", func() { + writeCfg("emb-test", "llama-cpp") + artifact := filepath.Join(tmpDir, "emb-test.bin") + Expect(os.WriteFile(artifact, []byte("first"), 0o644)).To(Succeed()) + + first, err := app.EmbedderFingerprint("emb-test") + Expect(err).NotTo(HaveOccurred()) + again, err := app.EmbedderFingerprint("emb-test") + Expect(err).NotTo(HaveOccurred()) + Expect(again).To(Equal(first)) + + Expect(os.WriteFile(artifact, []byte("replacement-with-different-size"), 0o644)).To(Succeed()) + replaced, err := app.EmbedderFingerprint("emb-test") + Expect(err).NotTo(HaveOccurred()) + Expect(replaced).NotTo(Equal(first)) + + app.backendLoader.UpdateModelConfig("emb-test", func(c *config.ModelConfig) { + c.Backend = "rerankers" + }) + updated, err := app.EmbedderFingerprint("emb-test") + Expect(err).NotTo(HaveOccurred()) + Expect(updated).NotTo(Equal(replaced)) + }) + + It("rejects an unknown model", func() { + _, err := app.EmbedderFingerprint("missing") + Expect(err).To(HaveOccurred()) + }) + }) + Context("Scorer", func() { It("returns nil at construction for an unknown model", func() { Expect(app.Scorer("missing")).To(BeNil()) diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index a362d9b676aa..027d5bc1ad3e 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -959,13 +959,20 @@ func DefaultRegistry() map[string]FieldMetaOverride { AutocompleteProvider: ProviderModels, Order: 241, }, + "router.knn.embedding_revision": { + Section: "router", + Label: "KNN: Embedding Revision", + Description: "Optional identity suffix for embedding services or remote models whose weights may change without their LocalAI model name/config changing. Bump this value after replacing such a model so the persisted corpus is re-embedded.", + Component: "input", + Order: 242, + }, "router.knn.k": { Section: "router", Label: "KNN: Neighbours (K)", Description: "How many nearest corpus entries vote on a prompt. 0 picks the default (3). K=1 routes on the single nearest example; larger K tolerates a mislabelled exemplar but needs denser corpus coverage per label.", Component: "number", Min: f64(0), - Order: 242, + Order: 243, }, "router.knn.similarity_threshold": { Section: "router", @@ -975,7 +982,7 @@ func DefaultRegistry() map[string]FieldMetaOverride { Min: f64(0), Max: f64(1), Step: f64(0.01), - Order: 243, + Order: 244, }, "router.knn.vote_threshold": { Section: "router", @@ -985,14 +992,14 @@ func DefaultRegistry() map[string]FieldMetaOverride { Min: f64(0), Max: f64(1), Step: f64(0.05), - Order: 244, + Order: 245, }, "router.knn.store_name": { Section: "router", Label: "KNN: Store Name", Description: "Optional override for the local-store collection holding the corpus vectors. Empty defaults to \"router-corpus-\".", Component: "input", - Order: 245, + Order: 246, }, } } diff --git a/core/config/model_config.go b/core/config/model_config.go index a59077f22473..4e714af47178 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -1,6 +1,8 @@ package config import ( + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" "os" @@ -405,6 +407,12 @@ type RouterKNNConfig struct { // entries recorded under a different embedder fingerprint. EmbeddingModel string `yaml:"embedding_model" json:"embedding_model"` + // EmbeddingRevision is an operator-controlled identity suffix for + // embedding backends whose underlying weights can change without a + // stable local checksum or config change. Bump it when replacing such + // a model in place so persisted corpus vectors are re-embedded. + EmbeddingRevision string `yaml:"embedding_revision,omitempty" json:"embedding_revision,omitempty"` + // K is how many nearest corpus entries vote on a probe. 0 picks // the package default (3). K=1 reproduces exact nearest-entry // routing; larger K tolerates mislabelled exemplars at the cost @@ -442,6 +450,15 @@ func (k *RouterKNNConfig) ResolvedStoreName(routerName string) string { return "router-corpus-" + routerName } +// ResolvedEmbeddingFingerprint binds the resolved embedding-model +// fingerprint to the optional operator revision. Keeping this combination +// in config makes the classifier build path and corpus management surfaces +// derive exactly the same persisted identity. +func (k *RouterKNNConfig) ResolvedEmbeddingFingerprint(modelFingerprint string) string { + sum := sha256.Sum256([]byte(modelFingerprint + "\x00" + k.EmbeddingRevision)) + return hex.EncodeToString(sum[:]) +} + // RouterPolicy is one entry in the label vocabulary. The label string // is what the classifier model emits and what candidates reference in // their Labels field; the description is the natural-language hint diff --git a/core/http/endpoints/localai/router_corpus.go b/core/http/endpoints/localai/router_corpus.go index 48c05f55f118..a5bde7899295 100644 --- a/core/http/endpoints/localai/router_corpus.go +++ b/core/http/endpoints/localai/router_corpus.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/middleware" "github.com/mudler/LocalAI/core/schema" @@ -79,11 +80,11 @@ func RouterCorpusAddEndpoint(loader *config.ModelConfigLoader, appConfig *config } added, skipped, stats, err := corpus.Seed(c.Request().Context(), mgr, cfg, storeName, - deps.Embedder, deps.VectorStore, entries) + deps.Embedder, deps.EmbedderFingerprint, deps.VectorStore, entries) if err != nil { // Undeclared labels and an unloadable embedding model are // request/config mistakes, not server faults. - if errors.Is(err, corpus.ErrUndeclaredLabel) || errors.Is(err, corpus.ErrEmbedderUnavailable) { + if errors.Is(err, corpus.ErrUndeclaredLabel) || errors.Is(err, corpus.ErrInvalidEntry) || errors.Is(err, corpus.ErrEmbedderUnavailable) { return echo.NewHTTPError(http.StatusBadRequest, err.Error()) } return echo.NewHTTPError(http.StatusInternalServerError, err.Error()) @@ -151,7 +152,11 @@ func RouterCorpusClearEndpoint(loader *config.ModelConfigLoader, appConfig *conf if err != nil { return err } - cleared, err := mgr.Clear(c.Request().Context(), storeName, deps.VectorStore(storeName)) + var store backend.VectorStore + if deps.VectorStore != nil { + store = deps.VectorStore(storeName) + } + cleared, err := mgr.Clear(c.Request().Context(), storeName, store) if err != nil { return echo.NewHTTPError(http.StatusInternalServerError, err.Error()) } diff --git a/core/http/endpoints/localai/router_corpus_test.go b/core/http/endpoints/localai/router_corpus_test.go index 35625eca965a..010b1f73df6f 100644 --- a/core/http/endpoints/localai/router_corpus_test.go +++ b/core/http/endpoints/localai/router_corpus_test.go @@ -101,8 +101,9 @@ var _ = Describe("Router corpus endpoints", func() { store = &corpusTestStore{} deps := middleware.ClassifierDeps{ - Embedder: func(string) backend.Embedder { return corpusTestEmbedder{} }, - VectorStore: func(string) backend.VectorStore { return store }, + Embedder: func(string) backend.Embedder { return corpusTestEmbedder{} }, + EmbedderFingerprint: func(string) (string, error) { return "test-embedding-fingerprint", nil }, + VectorStore: func(string) backend.VectorStore { return store }, } e = echo.New() e.POST("/api/router/:name/corpus", localai.RouterCorpusAddEndpoint(loader, appConfig, mgr, deps)) @@ -147,6 +148,31 @@ var _ = Describe("Router corpus endpoints", func() { Expect(rec.Body.String()).To(ContainSubstring("not declared")) }) + It("rejects duplicate labels before they can inflate KNN votes", func() { + writeKNNRouter(modelDir, "knn-router") + rec := do(http.MethodPost, "/api/router/knn-router/corpus", + `{"entries":[{"text":"hello","labels":["casual-chat","casual-chat"]}]}`) + Expect(rec.Code).To(Equal(http.StatusBadRequest)) + Expect(rec.Body.String()).To(ContainSubstring("duplicate label")) + }) + + It("rejects a score router that merely has a stray knn block", func() { + writeScoreRouter(modelDir, "score-router") + path := filepath.Join(modelDir, "score-router.yaml") + raw, err := os.ReadFile(path) + Expect(err).NotTo(HaveOccurred()) + var cfg config.ModelConfig + Expect(yaml.Unmarshal(raw, &cfg)).To(Succeed()) + cfg.Router.KNN = &config.RouterKNNConfig{EmbeddingModel: "embed-model"} + raw, err = yaml.Marshal(cfg) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(path, raw, 0o644)).To(Succeed()) + + rec := do(http.MethodPost, "/api/router/score-router/corpus", seedBody) + Expect(rec.Code).To(Equal(http.StatusBadRequest)) + Expect(rec.Body.String()).To(ContainSubstring("classifier: knn")) + }) + It("rejects an empty entries list", func() { writeKNNRouter(modelDir, "knn-router") rec := do(http.MethodPost, "/api/router/knn-router/corpus", `{"entries":[]}`) diff --git a/core/http/middleware/route_model.go b/core/http/middleware/route_model.go index 8527f8608f93..4c77aa5e3b4f 100644 --- a/core/http/middleware/route_model.go +++ b/core/http/middleware/route_model.go @@ -33,6 +33,11 @@ type ScorerFactory func(modelName string) backend.Scorer // classifier so routing still happens. type EmbedderFactory func(modelName string) backend.Embedder +// EmbedderFingerprintFactory returns the identity of the embedding space a +// named model currently produces. Persisted KNN vectors are only reusable +// while this identity is unchanged. +type EmbedderFingerprintFactory func(modelName string) (string, error) + // VectorStoreFactory returns a backend.VectorStore bound to a named // collection. Each router model's cache lives in its own collection // so two routers can't poison each other's hits. @@ -57,7 +62,7 @@ type ModelConfigLookup func(modelName string) *config.ModelConfig // classifier serves whatever the index already holds (tests, embedded // callers). type CorpusLoader interface { - EnsureLoaded(ctx context.Context, storeName, embeddingModel string, embedder backend.Embedder, store backend.VectorStore) (int, error) + EnsureLoaded(ctx context.Context, storeName, embeddingModel, embeddingFingerprint string, embedder backend.Embedder, store backend.VectorStore) (int, error) } // ClassifierDeps bundles the backend factories the router middleware @@ -72,10 +77,13 @@ type CorpusLoader interface { // score classifier runs unwrapped and the embedding-cache YAML is // ignored with a warning. type ClassifierDeps struct { - Scorer ScorerFactory - Embedder EmbedderFactory - VectorStore VectorStoreFactory - Reranker RerankerFactory + Scorer ScorerFactory + Embedder EmbedderFactory + // EmbedderFingerprint identifies the weights/config behind Embedder so + // KNN corpus vectors cannot be queried across embedding spaces. + EmbedderFingerprint EmbedderFingerprintFactory + VectorStore VectorStoreFactory + Reranker RerankerFactory // Corpus loads the persisted KNN corpus into the vector index when // a knn classifier is built. Optional; nil skips the load. @@ -118,15 +126,16 @@ type ClassifierDeps struct { // silently degrades that path rather than failing). func NewClassifierDeps(app *application.Application) ClassifierDeps { return ClassifierDeps{ - Scorer: app.Scorer, - Corpus: app.RouterCorpus(), - TokenCounter: app.TokenCounter, - Embedder: app.Embedder, - VectorStore: app.VectorStore, - Reranker: app.Reranker, - ModelLookup: app.ModelConfigLookup(), - Registry: app.RouterClassifierRegistry(), - Evaluator: app.TemplatesEvaluator(), + Scorer: app.Scorer, + Corpus: app.RouterCorpus(), + TokenCounter: app.TokenCounter, + Embedder: app.Embedder, + EmbedderFingerprint: app.EmbedderFingerprint, + VectorStore: app.VectorStore, + Reranker: app.Reranker, + ModelLookup: app.ModelConfigLookup(), + Registry: app.RouterClassifierRegistry(), + Evaluator: app.TemplatesEvaluator(), } } @@ -276,7 +285,11 @@ func GetOrBuildClassifier(registry *router.Registry, cfg *config.ModelConfig, de if deps.ModelLookup != nil { classifierCfg = deps.ModelLookup(cfg.Router.ClassifierModel) } - fp := routerConfigFingerprint(cfg.Router, classifierCfg) + embeddingFingerprint, err := resolvedKNNEmbeddingFingerprint(cfg, deps) + if err != nil { + return nil, err + } + fp := routerConfigFingerprint(cfg.Router, classifierCfg, embeddingFingerprint) if cached, ok := registry.Get(cfg.Name, fp); ok { return cached, nil } @@ -297,7 +310,7 @@ func GetOrBuildClassifier(registry *router.Registry, cfg *config.ModelConfig, de // unrelated changes (parameters, files, ...) don't burst the cache. // Pass classifierCfg=nil when no lookup is wired — the fingerprint // degenerates to the router-only form, matching pre-refactor behaviour. -func routerConfigFingerprint(rc config.RouterConfig, classifierCfg *config.ModelConfig) uint64 { +func routerConfigFingerprint(rc config.RouterConfig, classifierCfg *config.ModelConfig, additionalFingerprints ...string) uint64 { bytes, err := yaml.Marshal(rc) if err != nil { // Marshalling a value type can't fail in practice; fall @@ -326,9 +339,31 @@ func routerConfigFingerprint(rc config.RouterConfig, classifierCfg *config.Model h.Write([]byte(strconv.Itoa(*classifierCfg.ContextSize))) } } + for _, fingerprint := range additionalFingerprints { + h.Write([]byte{0}) + h.Write([]byte(fingerprint)) + } return h.Sum64() } +func resolvedKNNEmbeddingFingerprint(cfg *config.ModelConfig, deps ClassifierDeps) (string, error) { + name := cfg.Router.Classifier + if name == "" { + name = router.ClassifierScore + } + if name != router.ClassifierKNN || cfg.Router.KNN == nil || deps.EmbedderFingerprint == nil { + return "", nil + } + modelFingerprint, err := deps.EmbedderFingerprint(cfg.Router.KNN.EmbeddingModel) + if err != nil { + return "", fmt.Errorf("router classifier knn: fingerprint embedding_model %q: %w", cfg.Router.KNN.EmbeddingModel, err) + } + if modelFingerprint == "" { + return "", fmt.Errorf("router classifier knn: embedding_model %q returned an empty fingerprint", cfg.Router.KNN.EmbeddingModel) + } + return cfg.Router.KNN.ResolvedEmbeddingFingerprint(modelFingerprint), nil +} + func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Classifier, error) { rc := cfg.Router name := rc.Classifier @@ -423,6 +458,13 @@ func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Class if deps.Embedder == nil || deps.VectorStore == nil { return nil, fmt.Errorf("router classifier knn unavailable: embedder/vector-store factories not wired") } + embeddingFingerprint, err := resolvedKNNEmbeddingFingerprint(cfg, deps) + if err != nil { + return nil, err + } + if deps.Corpus != nil && embeddingFingerprint == "" { + return nil, fmt.Errorf("router classifier knn unavailable: embedder fingerprint factory not wired") + } embedder := deps.Embedder(rc.KNN.EmbeddingModel) if embedder == nil { return nil, fmt.Errorf("router classifier knn: embedding_model %q not loadable", rc.KNN.EmbeddingModel) @@ -433,12 +475,11 @@ func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Class return nil, fmt.Errorf("router classifier knn: vector store %q not loadable", storeName) } if deps.Corpus != nil { - // A failed load must not break routing — the classifier - // still works against whatever the index holds, and the - // epistemic gate falls back safely on an empty index. - if n, err := deps.Corpus.EnsureLoaded(context.Background(), storeName, rc.KNN.EmbeddingModel, embedder, vstore); err != nil { - xlog.Warn("router: knn corpus load failed; routing continues on the live index", - "router_model", cfg.Name, "store", storeName, "error", err) + // Loading fails closed: a live index from a different embedding + // space may have the same vector width and return plausible but + // incorrect routes. + if n, err := deps.Corpus.EnsureLoaded(context.Background(), storeName, rc.KNN.EmbeddingModel, embeddingFingerprint, embedder, vstore); err != nil { + return nil, fmt.Errorf("router classifier knn: load corpus %q: %w", storeName, err) } else if n > 0 { xlog.Info("router: knn corpus loaded", "router_model", cfg.Name, "store", storeName, "entries", n) diff --git a/core/http/middleware/route_model_test.go b/core/http/middleware/route_model_test.go index ea8dbe46ad01..8b63560cefd0 100644 --- a/core/http/middleware/route_model_test.go +++ b/core/http/middleware/route_model_test.go @@ -581,6 +581,12 @@ var ( errTestKNNInsert = errors.New("knn classifier must never insert into the corpus") ) +type failingCorpusLoader struct{ err error } + +func (f failingCorpusLoader) EnsureLoaded(context.Context, string, string, string, backend.Embedder, backend.VectorStore) (int, error) { + return 0, f.err +} + func corpusPayload(labels ...string) []byte { b, err := router.EncodeCorpusEntry(router.EntryID(strings.Join(labels, "+")), labels) Expect(err).NotTo(HaveOccurred()) @@ -704,6 +710,39 @@ var _ = Describe("RouteModel middleware (knn classifier)", func() { Expect(err.Error()).To(ContainSubstring("knn")) }) + It("fails closed when the persisted corpus cannot sync into the live index", func() { + routerCfg := newKNNRouterModel(modelDir, "smart-router") + writeCandidate(modelDir, "small-model") + writeCandidate(modelDir, "big-model") + deps := knnDeps() + deps.EmbedderFingerprint = func(string) (string, error) { return "fp", nil } + deps.Corpus = failingCorpusLoader{err: errors.New("incompatible persisted vectors")} + + _, err := runRouterWithDeps(loader, appConfig, store, routerCfg, + openAIChat("hello"), deps) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("incompatible persisted vectors")) + }) + + It("invalidates the classifier cache when the embedding fingerprint changes", func() { + routerCfg := newKNNRouterModel(modelDir, "smart-router") + fingerprint := "v1" + deps := knnDeps() + deps.EmbedderFingerprint = func(string) (string, error) { return fingerprint, nil } + registry := router.NewRegistry() + + first, err := GetOrBuildClassifier(registry, routerCfg, deps) + Expect(err).NotTo(HaveOccurred()) + cached, err := GetOrBuildClassifier(registry, routerCfg, deps) + Expect(err).NotTo(HaveOccurred()) + Expect(cached).To(BeIdenticalTo(first)) + + fingerprint = "v2" + rebuilt, err := GetOrBuildClassifier(registry, routerCfg, deps) + Expect(err).NotTo(HaveOccurred()) + Expect(rebuilt).NotTo(BeIdenticalTo(first)) + }) + It("ignores an embedding_cache block instead of double-embedding", func() { routerCfg := newKNNRouterModel(modelDir, "smart-router") routerCfg.Router.EmbeddingCache = &config.EmbeddingCacheConfig{EmbeddingModel: "embed-model"} diff --git a/core/services/routing/corpus/manager.go b/core/services/routing/corpus/manager.go index 3ace526e5466..bac177011870 100644 --- a/core/services/routing/corpus/manager.go +++ b/core/services/routing/corpus/manager.go @@ -3,13 +3,13 @@ // // The corpus FILE is the source of truth: one JSONL file per store // name under , each line an Entry {text, labels, vector, -// embedding_model}. The local-store vector backend is a pure in-memory -// index rebuilt from the file — it has no persistence of its own (its -// Load is an explicit no-op), and keeping it that way preserves the +// embedding_model, embedding_fingerprint}. The local-store vector backend is +// a pure in-memory index rebuilt from the file — it has no persistence of its +// own (its Load is an explicit no-op), and keeping it that way preserves the // documented swap point for external vector backends. Vectors are // cached in the file alongside the text so a restart re-indexes // without re-embedding; entries recorded under a different embedding -// model are re-embedded on load. +// fingerprint are re-embedded on load. // // Texts in the corpus file never leave the server: Stats exposes label // counts only, and there is deliberately no API that returns entries. @@ -18,7 +18,10 @@ package corpus import ( "bufio" "context" + "crypto/sha256" + "encoding/hex" "encoding/json" + "errors" "fmt" "os" "path/filepath" @@ -32,23 +35,30 @@ import ( "github.com/mudler/xlog" ) -// Entry is one labelled exemplar. Vector and EmbeddingModel are the -// embedding cache — absent (or stale) entries are re-embedded when the -// corpus is loaded or added to. +// Entry is one labelled exemplar. Vector, EmbeddingModel, and +// EmbeddingFingerprint are the embedding cache — absent (or stale) entries +// are re-embedded when the corpus is loaded or added to. type Entry struct { - Text string `json:"text"` - Labels []string `json:"labels"` - Vector []float32 `json:"vector,omitempty"` - EmbeddingModel string `json:"embedding_model,omitempty"` + Text string `json:"text"` + Labels []string `json:"labels"` + Vector []float32 `json:"vector,omitempty"` + EmbeddingModel string `json:"embedding_model,omitempty"` + EmbeddingFingerprint string `json:"embedding_fingerprint,omitempty"` } +// ErrLiveEmbeddingMismatch means a store already contains vectors from a +// different embedding space. Querying it with the new embedder can silently +// misroute when the dimensions happen to match, so callers must fail closed +// until LocalAI restarts with an empty live index. +var ErrLiveEmbeddingMismatch = errors.New("live corpus index uses a different embedding fingerprint") + // Stats is the inspection surface — counts only, never texts. The // router corpus API and the Routing tab render this. type Stats struct { StoreName string `json:"store_name"` Total int `json:"total"` LabelCounts map[string]int `json:"label_counts"` - // EmbeddingModels lists the distinct embedder fingerprints present + // EmbeddingModels lists the distinct embedding model names present // in the file. More than one entry means a re-embed is pending for // part of the corpus (it happens lazily on the next load). EmbeddingModels []string `json:"embedding_models,omitempty"` @@ -73,11 +83,11 @@ type Manager struct { dir string mu sync.Mutex - // loadedModel records which embedding model a store name was synced - // into the live index under. Guards double-loading and detects the - // embedding-model-changed-on-a-live-index case, which local-store - // cannot serve (it enforces one key length per store). - loadedModel map[string]string + // states records both embedding-space identity and which persisted + // file version reached the live index. needsSync survives a durable + // file write followed by a transient index failure, making an exact + // retry repair the index instead of skipping every duplicate entry. + states map[string]storeState // statsCache memoises Stats per store, keyed by the corpus file's // stat fingerprint — the file is the source of truth, so a matching @@ -89,6 +99,13 @@ type Manager struct { statsCache map[string]cachedStats } +type storeState struct { + embeddingFingerprint string + syncedFile fileFingerprint + needsSync bool + indexedEntries int +} + type cachedStats struct { key fileFingerprint stats Stats @@ -107,9 +124,9 @@ type fileFingerprint struct { // write). func NewManager(dir string) *Manager { return &Manager{ - dir: dir, - loadedModel: map[string]string{}, - statsCache: map[string]cachedStats{}, + dir: dir, + states: map[string]storeState{}, + statsCache: map[string]cachedStats{}, } } @@ -124,27 +141,39 @@ func (m *Manager) path(storeName string) string { } return '_' }, storeName) - return filepath.Join(m.dir, safe+".jsonl") + safe = strings.Trim(safe, ".") + if safe == "" { + safe = "store" + } + if len(safe) > 80 { + safe = safe[:80] + } + sum := sha256.Sum256([]byte(storeName)) + return filepath.Join(m.dir, safe+"-"+hex.EncodeToString(sum[:16])+".jsonl") } // EnsureLoaded syncs the persisted corpus for storeName into the live -// vector index, once per (store, embedding model) per process. Entries -// recorded under a different embedding model are re-embedded and the +// vector index, once per persisted file version and embedding fingerprint. +// Entries recorded under a different fingerprint are re-embedded and the // file rewritten. Returns the number of entries now indexed. A missing // file is an empty corpus, not an error. -func (m *Manager) EnsureLoaded(ctx context.Context, storeName, embeddingModel string, embedder backend.Embedder, store backend.VectorStore) (int, error) { +func (m *Manager) EnsureLoaded(ctx context.Context, storeName, embeddingModel, embeddingFingerprint string, embedder backend.Embedder, store backend.VectorStore) (int, error) { m.mu.Lock() defer m.mu.Unlock() - if prev, ok := m.loadedModel[storeName]; ok { - if prev == embeddingModel { + if embeddingFingerprint == "" { + return 0, fmt.Errorf("corpus %q: empty embedding fingerprint", storeName) + } + fileKey := m.fingerprint(storeName) + if state, ok := m.states[storeName]; ok { + if state.embeddingFingerprint != embeddingFingerprint { + if state.indexedEntries > 0 || state.needsSync { + return 0, fmt.Errorf("corpus %q %w; restart LocalAI to rebuild it", storeName, ErrLiveEmbeddingMismatch) + } + delete(m.states, storeName) + } else if !state.needsSync && state.syncedFile.equal(fileKey) { return 0, nil } - // The in-memory index already holds vectors from another - // embedder; local-store enforces a single key length, so mixing - // would fail on insert (or silently corrupt neighbourhoods when - // dimensions happen to match). A restart re-indexes cleanly. - return 0, fmt.Errorf("corpus %q was indexed with embedding model %q this process; restart LocalAI to re-index it with %q", storeName, prev, embeddingModel) } entries, _, err := m.readAll(storeName) @@ -152,13 +181,16 @@ func (m *Manager) EnsureLoaded(ctx context.Context, storeName, embeddingModel st return 0, err } if len(entries) == 0 { - m.loadedModel[storeName] = embeddingModel + m.states[storeName] = storeState{ + embeddingFingerprint: embeddingFingerprint, + syncedFile: fileKey, + } return 0, nil } dirty := false for i := range entries { - if len(entries[i].Vector) > 0 && entries[i].EmbeddingModel == embeddingModel { + if len(entries[i].Vector) > 0 && entries[i].EmbeddingFingerprint == embeddingFingerprint { continue } if embedder == nil { @@ -170,18 +202,29 @@ func (m *Manager) EnsureLoaded(ctx context.Context, storeName, embeddingModel st } entries[i].Vector = vec entries[i].EmbeddingModel = embeddingModel + entries[i].EmbeddingFingerprint = embeddingFingerprint dirty = true } if dirty { if err := m.write(storeName, entries); err != nil { return 0, err } + fileKey = m.fingerprint(storeName) } + if store == nil { + m.states[storeName] = storeState{embeddingFingerprint: embeddingFingerprint, needsSync: true, indexedEntries: len(entries)} + return 0, fmt.Errorf("corpus %q: no vector store available", storeName) + } if err := insertAll(ctx, store, entries); err != nil { + m.states[storeName] = storeState{embeddingFingerprint: embeddingFingerprint, needsSync: true, indexedEntries: len(entries)} return 0, err } - m.loadedModel[storeName] = embeddingModel + m.states[storeName] = storeState{ + embeddingFingerprint: embeddingFingerprint, + syncedFile: fileKey, + indexedEntries: len(entries), + } return len(entries), nil } @@ -197,16 +240,16 @@ func (m *Manager) EnsureLoaded(ctx context.Context, storeName, embeddingModel st // corpus operations for its whole duration. The commit re-checks for // duplicates under the lock, so a racing Add can at worst turn an // entry into a skip, never a double insert. -func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, embedder backend.Embedder, store backend.VectorStore, entries []Entry) (int, int, error) { +func (m *Manager) Add(ctx context.Context, storeName, embeddingModel, embeddingFingerprint string, embedder backend.Embedder, store backend.VectorStore, entries []Entry) (int, int, error) { if embedder == nil { return 0, 0, fmt.Errorf("corpus %q: no embedder available", storeName) } + if embeddingFingerprint == "" { + return 0, 0, fmt.Errorf("corpus %q: empty embedding fingerprint", storeName) + } for i, e := range entries { - if strings.TrimSpace(e.Text) == "" { - return 0, 0, fmt.Errorf("corpus entry %d: empty text", i) - } - if len(e.Labels) == 0 { - return 0, 0, fmt.Errorf("corpus entry %d: at least one label required", i) + if err := validateEntry(e); err != nil { + return 0, 0, fmt.Errorf("corpus entry %d: %w", i, err) } } @@ -218,12 +261,39 @@ func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, emb m.mu.Unlock() return 0, 0, err } + if state, ok := m.states[storeName]; ok && state.embeddingFingerprint != embeddingFingerprint && (state.indexedEntries > 0 || state.needsSync) { + m.mu.Unlock() + return 0, 0, fmt.Errorf("corpus %q %w; restart LocalAI to rebuild it", storeName, ErrLiveEmbeddingMismatch) + } + for _, e := range existing { + if e.EmbeddingFingerprint != embeddingFingerprint { + m.mu.Unlock() + return 0, 0, fmt.Errorf("corpus %q contains stale embedding vectors; load it before adding entries", storeName) + } + } + state, stateOK := m.states[storeName] + needsSync := len(existing) > 0 && (!stateOK || state.needsSync || !state.syncedFile.equal(m.fingerprint(storeName))) + if needsSync { + if store == nil { + m.states[storeName] = storeState{embeddingFingerprint: embeddingFingerprint, needsSync: true, indexedEntries: len(existing)} + m.mu.Unlock() + return 0, 0, fmt.Errorf("corpus %q: entries persisted but no vector store is available", storeName) + } + if err := insertAll(ctx, store, existing); err != nil { + m.states[storeName] = storeState{embeddingFingerprint: embeddingFingerprint, needsSync: true, indexedEntries: len(existing)} + m.mu.Unlock() + return 0, 0, fmt.Errorf("corpus %q: retrying live index sync: %w", storeName, err) + } + m.states[storeName] = storeState{ + embeddingFingerprint: embeddingFingerprint, + syncedFile: m.fingerprint(storeName), + indexedEntries: len(existing), + } + } seen := make(map[string]struct{}, len(existing)) for _, e := range existing { seen[e.Text] = struct{}{} } - m.mu.Unlock() - candidates := make([]Entry, 0, len(entries)) skipped := 0 for _, e := range entries { @@ -235,8 +305,10 @@ func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, emb candidates = append(candidates, e) } if len(candidates) == 0 { + m.mu.Unlock() return 0, skipped, nil } + m.mu.Unlock() for i := range candidates { vec, err := embedder.Embed(ctx, candidates[i].Text) @@ -245,6 +317,7 @@ func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, emb } candidates[i].Vector = vec candidates[i].EmbeddingModel = embeddingModel + candidates[i].EmbeddingFingerprint = embeddingFingerprint } m.mu.Lock() @@ -257,6 +330,14 @@ func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, emb if err != nil { return 0, skipped, err } + if state, ok := m.states[storeName]; ok && state.embeddingFingerprint != embeddingFingerprint && (state.indexedEntries > 0 || state.needsSync) { + return 0, skipped, fmt.Errorf("corpus %q %w; restart LocalAI to rebuild it", storeName, ErrLiveEmbeddingMismatch) + } + for _, e := range current { + if e.EmbeddingFingerprint != embeddingFingerprint { + return 0, skipped, fmt.Errorf("corpus %q contains stale embedding vectors; load it before adding entries", storeName) + } + } seen = make(map[string]struct{}, len(current)) for _, e := range current { seen[e.Text] = struct{}{} @@ -284,12 +365,21 @@ func (m *Manager) Add(ctx context.Context, storeName, embeddingModel string, emb if err != nil { return 0, skipped, err } + entryCount := len(current) + len(added) if store != nil { if err := insertAll(ctx, store, added); err != nil { // Durable but not indexed — routing won't see the new // entries until the next successful load. Surface loudly. - return len(added), skipped, fmt.Errorf("corpus %q: entries persisted but indexing failed (they will index on next load/restart): %w", storeName, err) + m.states[storeName] = storeState{embeddingFingerprint: embeddingFingerprint, needsSync: true, indexedEntries: entryCount} + return len(added), skipped, fmt.Errorf("corpus %q: entries persisted but indexing failed (retry the same seed request or restart): %w", storeName, err) + } + m.states[storeName] = storeState{ + embeddingFingerprint: embeddingFingerprint, + syncedFile: m.fingerprint(storeName), + indexedEntries: entryCount, } + } else { + m.states[storeName] = storeState{embeddingFingerprint: embeddingFingerprint, needsSync: true, indexedEntries: entryCount} } return len(added), skipped, nil } @@ -377,7 +467,7 @@ func (m *Manager) Clear(ctx context.Context, storeName string, store backend.Vec if err := os.Remove(m.path(storeName)); err != nil && !os.IsNotExist(err) { return 0, err } - delete(m.loadedModel, storeName) + delete(m.states, storeName) if len(entries) == 0 || store == nil { return len(entries), nil } @@ -400,6 +490,26 @@ func (m *Manager) Clear(ctx context.Context, storeName string, store backend.Vec return len(entries), nil } +func validateEntry(e Entry) error { + if strings.TrimSpace(e.Text) == "" { + return errors.New("empty text") + } + if len(e.Labels) == 0 { + return errors.New("at least one label required") + } + seen := make(map[string]struct{}, len(e.Labels)) + for _, label := range e.Labels { + if strings.TrimSpace(label) == "" { + return errors.New("labels must not be empty") + } + if _, ok := seen[label]; ok { + return fmt.Errorf("duplicate label %q", label) + } + seen[label] = struct{}{} + } + return nil +} + func insertAll(ctx context.Context, store backend.VectorStore, entries []Entry) error { vecs := make([][]float32, 0, len(entries)) payloads := make([][]byte, 0, len(entries)) diff --git a/core/services/routing/corpus/manager_test.go b/core/services/routing/corpus/manager_test.go index 0735d7e11ba3..6c50c25a6bdb 100644 --- a/core/services/routing/corpus/manager_test.go +++ b/core/services/routing/corpus/manager_test.go @@ -2,6 +2,7 @@ package corpus_test import ( "context" + "errors" "os" "path/filepath" "strings" @@ -33,10 +34,11 @@ func (e *countingEmbedder) Embed(_ context.Context, text string) ([]float32, err // capturingStore records index mutations. Search/SearchK are // irrelevant to the manager and return clean misses. type capturingStore struct { - mu sync.Mutex - payloads [][]byte - batches int - deleted [][]float32 + mu sync.Mutex + payloads [][]byte + batches int + deleted [][]float32 + failBatches int } func (s *capturingStore) Search(_ context.Context, _ []float32) (float64, []byte, bool, error) { @@ -58,11 +60,22 @@ func (s *capturingStore) InsertBatch(_ context.Context, vecs [][]float32, payloa s.mu.Lock() defer s.mu.Unlock() s.batches++ + if s.failBatches > 0 { + s.failBatches-- + return errors.New("transient batch failure") + } s.payloads = append(s.payloads, payloads...) _ = vecs return nil } +func onlyCorpusPath(dir string) string { + matches, err := filepath.Glob(filepath.Join(dir, "*.jsonl")) + Expect(err).NotTo(HaveOccurred()) + Expect(matches).To(HaveLen(1)) + return matches[0] +} + func (s *capturingStore) Delete(_ context.Context, vecs [][]float32) error { s.mu.Lock() defer s.mu.Unlock() @@ -79,7 +92,10 @@ var _ = Describe("corpus.Manager", func() { ctx context.Context ) - const storeName = "router-corpus-smart-router" + const ( + storeName = "router-corpus-smart-router" + fingerprint = "embed-1-fingerprint" + ) seed := []corpus.Entry{ {Text: "debug this Go null pointer", Labels: []string{"code-generation"}}, @@ -102,7 +118,7 @@ var _ = Describe("corpus.Manager", func() { }) It("adds entries: embeds, persists, and indexes them", func() { - added, skipped, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + added, skipped, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed) Expect(err).NotTo(HaveOccurred()) Expect(added).To(Equal(3)) Expect(skipped).To(BeZero()) @@ -114,54 +130,56 @@ var _ = Describe("corpus.Manager", func() { Expect(string(store.payloads[0])).To(ContainSubstring("code-generation")) // Persisted on disk under a sanitised name. - _, err = os.Stat(filepath.Join(dir, storeName+".jsonl")) + _, err = os.Stat(onlyCorpusPath(dir)) Expect(err).NotTo(HaveOccurred()) }) It("skips duplicate texts instead of double-weighting them", func() { - _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed) Expect(err).NotTo(HaveOccurred()) - added, skipped, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed[:2]) + added, skipped, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed[:2]) Expect(err).NotTo(HaveOccurred()) Expect(added).To(BeZero()) Expect(skipped).To(Equal(2)) }) It("rejects empty text and label-less entries", func() { - _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, []corpus.Entry{{Text: " ", Labels: []string{"x"}}}) + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, []corpus.Entry{{Text: " ", Labels: []string{"x"}}}) Expect(err).To(HaveOccurred()) - _, _, err = mgr.Add(ctx, storeName, "embed-1", embedder, store, []corpus.Entry{{Text: "hello", Labels: nil}}) + _, _, err = mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, []corpus.Entry{{Text: "hello", Labels: nil}}) Expect(err).To(HaveOccurred()) + _, _, err = mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, []corpus.Entry{{Text: "hello", Labels: []string{"x", "x"}}}) + Expect(err).To(MatchError(ContainSubstring("duplicate label"))) }) It("reloads a persisted corpus into a fresh index without re-embedding", func() { - _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed) Expect(err).NotTo(HaveOccurred()) // Simulate restart: fresh manager over the same dir, empty index. mgr2 := corpus.NewManager(dir) store2 := &capturingStore{} embedder2 := &countingEmbedder{model: 1} - n, err := mgr2.EnsureLoaded(ctx, storeName, "embed-1", embedder2, store2) + n, err := mgr2.EnsureLoaded(ctx, storeName, "embed-1", fingerprint, embedder2, store2) Expect(err).NotTo(HaveOccurred()) Expect(n).To(Equal(3)) Expect(store2.payloads).To(HaveLen(3)) Expect(embedder2.calls).To(BeZero(), "cached vectors must be reused") // Second call is a no-op — already synced this process. - n, err = mgr2.EnsureLoaded(ctx, storeName, "embed-1", embedder2, store2) + n, err = mgr2.EnsureLoaded(ctx, storeName, "embed-1", fingerprint, embedder2, store2) Expect(err).NotTo(HaveOccurred()) Expect(n).To(BeZero()) }) It("re-embeds entries recorded under a different embedding model", func() { - _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed) Expect(err).NotTo(HaveOccurred()) mgr2 := corpus.NewManager(dir) newEmbedder := &countingEmbedder{model: 2} store2 := &capturingStore{} - n, err := mgr2.EnsureLoaded(ctx, storeName, "embed-2", newEmbedder, store2) + n, err := mgr2.EnsureLoaded(ctx, storeName, "embed-2", "embed-2-fingerprint", newEmbedder, store2) Expect(err).NotTo(HaveOccurred()) Expect(n).To(Equal(3)) Expect(newEmbedder.calls).To(Equal(3), "fingerprint mismatch must re-embed") @@ -170,25 +188,25 @@ var _ = Describe("corpus.Manager", func() { // without touching the embedder again. mgr3 := corpus.NewManager(dir) embedder3 := &countingEmbedder{model: 2} - n, err = mgr3.EnsureLoaded(ctx, storeName, "embed-2", embedder3, &capturingStore{}) + n, err = mgr3.EnsureLoaded(ctx, storeName, "embed-2", "embed-2-fingerprint", embedder3, &capturingStore{}) Expect(err).NotTo(HaveOccurred()) Expect(n).To(Equal(3)) Expect(embedder3.calls).To(BeZero()) }) It("refuses to mix embedding models in a live index", func() { - _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed) Expect(err).NotTo(HaveOccurred()) - _, err = mgr.EnsureLoaded(ctx, storeName, "embed-1", embedder, store) + _, err = mgr.EnsureLoaded(ctx, storeName, "embed-1", fingerprint, embedder, store) Expect(err).NotTo(HaveOccurred()) - _, err = mgr.EnsureLoaded(ctx, storeName, "embed-2", &countingEmbedder{model: 2}, store) + _, err = mgr.EnsureLoaded(ctx, storeName, "embed-2", "embed-2-fingerprint", &countingEmbedder{model: 2}, store) Expect(err).To(HaveOccurred()) Expect(err.Error()).To(ContainSubstring("restart")) }) It("reports label counts and never texts", func() { - _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed) Expect(err).NotTo(HaveOccurred()) st, err := mgr.Stats(storeName) Expect(err).NotTo(HaveOccurred()) @@ -201,7 +219,7 @@ var _ = Describe("corpus.Manager", func() { }) It("clears the file and the live index", func() { - _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed) + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed) Expect(err).NotTo(HaveOccurred()) n, err := mgr.Clear(ctx, storeName, store) @@ -214,13 +232,13 @@ var _ = Describe("corpus.Manager", func() { Expect(st.Total).To(BeZero()) // And a load after clear indexes nothing. - loaded, err := corpus.NewManager(dir).EnsureLoaded(ctx, storeName, "embed-1", embedder, &capturingStore{}) + loaded, err := corpus.NewManager(dir).EnsureLoaded(ctx, storeName, "embed-1", fingerprint, embedder, &capturingStore{}) Expect(err).NotTo(HaveOccurred()) Expect(loaded).To(BeZero()) }) It("treats a missing file as an empty corpus", func() { - n, err := mgr.EnsureLoaded(ctx, "never-seeded", "embed-1", embedder, store) + n, err := mgr.EnsureLoaded(ctx, "never-seeded", "embed-1", fingerprint, embedder, store) Expect(err).NotTo(HaveOccurred()) Expect(n).To(BeZero()) st, err := mgr.Stats("never-seeded") @@ -229,11 +247,11 @@ var _ = Describe("corpus.Manager", func() { }) It("tolerates a torn final line and repairs it on the next add", func() { - _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed[:2]) + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed[:2]) Expect(err).NotTo(HaveOccurred()) // Simulate a crash mid-append: an unparseable tail. - path := filepath.Join(dir, storeName+".jsonl") + path := onlyCorpusPath(dir) f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o640) Expect(err).NotTo(HaveOccurred()) _, err = f.WriteString(`{"text":"torn`) @@ -246,7 +264,7 @@ var _ = Describe("corpus.Manager", func() { Expect(st.Total).To(Equal(2)) // The next add rewrites the file whole, clearing the damage. - added, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed[2:]) + added, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed[2:]) Expect(err).NotTo(HaveOccurred()) Expect(added).To(Equal(1)) raw, err := os.ReadFile(path) @@ -259,12 +277,12 @@ var _ = Describe("corpus.Manager", func() { }) It("fails on corruption that is not a torn tail", func() { - _, _, err := mgr.Add(ctx, storeName, "embed-1", embedder, store, seed[:1]) + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed[:1]) Expect(err).NotTo(HaveOccurred()) // Garbage with a valid successor line is real corruption, not // an interrupted append — reads must refuse to guess. - path := filepath.Join(dir, storeName+".jsonl") + path := onlyCorpusPath(dir) raw, err := os.ReadFile(path) Expect(err).NotTo(HaveOccurred()) Expect(os.WriteFile(path, append([]byte("garbage\n"), raw...), 0o640)).To(Succeed()) @@ -276,11 +294,54 @@ var _ = Describe("corpus.Manager", func() { It("sanitises hostile store names into the corpus dir", func() { hostile := "../../etc/passwd" - _, _, err := mgr.Add(ctx, hostile, "embed-1", embedder, store, seed[:1]) + _, _, err := mgr.Add(ctx, hostile, "embed-1", fingerprint, embedder, store, seed[:1]) Expect(err).NotTo(HaveOccurred()) entries, err := os.ReadDir(dir) Expect(err).NotTo(HaveOccurred()) Expect(entries).To(HaveLen(1)) Expect(strings.Contains(entries[0].Name(), "/")).To(BeFalse()) }) + + It("keeps store names distinct when their readable sanitised prefixes collide", func() { + _, _, err := mgr.Add(ctx, "team/a", "embed-1", fingerprint, embedder, store, seed[:1]) + Expect(err).NotTo(HaveOccurred()) + _, _, err = mgr.Add(ctx, "team?a", "embed-1", fingerprint, embedder, store, seed[1:2]) + Expect(err).NotTo(HaveOccurred()) + + files, err := filepath.Glob(filepath.Join(dir, "*.jsonl")) + Expect(err).NotTo(HaveOccurred()) + Expect(files).To(HaveLen(2)) + first, err := mgr.Stats("team/a") + Expect(err).NotTo(HaveOccurred()) + second, err := mgr.Stats("team?a") + Expect(err).NotTo(HaveOccurred()) + Expect(first.LabelCounts).To(HaveKey("code-generation")) + Expect(second.LabelCounts).To(HaveKey("math-reasoning")) + }) + + It("re-embeds after the fingerprint changes even when the model name does not", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + + fresh := corpus.NewManager(dir) + changed := &countingEmbedder{model: 9} + n, err := fresh.EnsureLoaded(ctx, storeName, "embed-1", "embed-1-upgraded", changed, &capturingStore{}) + Expect(err).NotTo(HaveOccurred()) + Expect(n).To(Equal(3)) + Expect(changed.calls).To(Equal(3)) + }) + + It("repairs a durable-file/live-index split on an exact retry", func() { + flaky := &capturingStore{failBatches: 1} + added, skipped, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, flaky, seed[:1]) + Expect(err).To(MatchError(ContainSubstring("retry the same seed request"))) + Expect(added).To(Equal(1)) + Expect(skipped).To(BeZero()) + + added, skipped, err = mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, flaky, seed[:1]) + Expect(err).NotTo(HaveOccurred()) + Expect(added).To(BeZero()) + Expect(skipped).To(Equal(1)) + Expect(flaky.payloads).To(HaveLen(1)) + }) }) diff --git a/core/services/routing/corpus/resolve.go b/core/services/routing/corpus/resolve.go index 13bf976dadf0..15763b488d2f 100644 --- a/core/services/routing/corpus/resolve.go +++ b/core/services/routing/corpus/resolve.go @@ -4,9 +4,11 @@ import ( "context" "errors" "fmt" + "strings" "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/routing/router" ) // The REST corpus endpoints and the assistant MCP tools are thin @@ -19,20 +21,23 @@ import ( var ( // ErrRouterNotFound: the named model doesn't exist (HTTP 404). ErrRouterNotFound = errors.New("not found") - // ErrNotKNNRouter: the model exists but declares no router.knn - // block (HTTP 400). + // ErrNotKNNRouter: the model exists but is not an active KNN router + // with a complete router.knn block (HTTP 400). ErrNotKNNRouter = errors.New("has no router.knn block (set classifier: knn and knn.embedding_model first)") // ErrUndeclaredLabel: a seed entry uses a label that is not a // declared router policy (HTTP 400). ErrUndeclaredLabel = errors.New("is not declared in router policies") + // ErrInvalidEntry: a seed entry is structurally invalid (HTTP 400). + ErrInvalidEntry = errors.New("invalid corpus entry") // ErrEmbedderUnavailable: the router's knn.embedding_model can't // be loaded (HTTP 400 — a config problem, not a server fault). ErrEmbedderUnavailable = errors.New("not loadable") ) // ResolveKNNRouter loads the named model config, requires it to declare -// a router.knn block, and resolves the corpus store name the classifier -// build uses — the single resolution path for every corpus surface. +// an active knn classifier with a router.knn block, and resolves the corpus +// store name the classifier build uses — the single resolution path for every +// corpus surface. // // LoadModelConfigFileByName returns a synthetic stub (empty Name) when // the model is unknown; that maps to ErrRouterNotFound so callers can @@ -45,7 +50,11 @@ func ResolveKNNRouter(loader *config.ModelConfigLoader, appConfig *config.Applic if cfg == nil || cfg.Name == "" { return nil, "", fmt.Errorf("model %q %w", name, ErrRouterNotFound) } - if cfg.Router.KNN == nil || cfg.Router.KNN.EmbeddingModel == "" { + classifier := cfg.Router.Classifier + if classifier == "" { + classifier = router.ClassifierScore + } + if !cfg.HasRouter() || classifier != router.ClassifierKNN || cfg.Router.KNN == nil || cfg.Router.KNN.EmbeddingModel == "" { return nil, "", fmt.Errorf("model %q %w", name, ErrNotKNNRouter) } return cfg, cfg.Router.KNN.ResolvedStoreName(cfg.Name), nil @@ -56,7 +65,7 @@ func ResolveKNNRouter(loader *config.ModelConfigLoader, appConfig *config.Applic // must not silently create an unroutable label), embeds and persists // them via the Manager, and returns the post-seed stats. func Seed(ctx context.Context, mgr *Manager, cfg *config.ModelConfig, storeName string, - embedderFor func(string) backend.Embedder, storeFor func(string) backend.VectorStore, + embedderFor func(string) backend.Embedder, fingerprintFor func(string) (string, error), storeFor func(string) backend.VectorStore, entries []Entry) (added, skipped int, stats Stats, err error) { declared := map[string]struct{}{} @@ -64,7 +73,18 @@ func Seed(ctx context.Context, mgr *Manager, cfg *config.ModelConfig, storeName declared[p.Label] = struct{}{} } for i, e := range entries { + if strings.TrimSpace(e.Text) == "" || len(e.Labels) == 0 { + return 0, 0, Stats{}, fmt.Errorf("entry %d: %w: text and at least one label are required", i, ErrInvalidEntry) + } + seen := make(map[string]struct{}, len(e.Labels)) for _, l := range e.Labels { + if strings.TrimSpace(l) == "" { + return 0, 0, Stats{}, fmt.Errorf("entry %d: %w: labels must not be empty", i, ErrInvalidEntry) + } + if _, duplicate := seen[l]; duplicate { + return 0, 0, Stats{}, fmt.Errorf("entry %d: %w: duplicate label %q", i, ErrInvalidEntry, l) + } + seen[l] = struct{}{} if _, ok := declared[l]; !ok { return 0, 0, Stats{}, fmt.Errorf("entry %d: label %q %w", i, l, ErrUndeclaredLabel) } @@ -79,12 +99,29 @@ func Seed(ctx context.Context, mgr *Manager, cfg *config.ModelConfig, storeName if embedder == nil { return 0, 0, Stats{}, fmt.Errorf("embedding_model %q %w", embeddingModel, ErrEmbedderUnavailable) } + if fingerprintFor == nil { + return 0, 0, Stats{}, errors.New("embedding fingerprint factory is unavailable") + } + modelFingerprint, err := fingerprintFor(embeddingModel) + if err != nil { + return 0, 0, Stats{}, fmt.Errorf("fingerprint embedding_model %q: %w", embeddingModel, err) + } + if modelFingerprint == "" { + return 0, 0, Stats{}, fmt.Errorf("fingerprint embedding_model %q: empty fingerprint", embeddingModel) + } + embeddingFingerprint := cfg.Router.KNN.ResolvedEmbeddingFingerprint(modelFingerprint) var store backend.VectorStore if storeFor != nil { store = storeFor(storeName) } + if store == nil { + return 0, 0, Stats{}, fmt.Errorf("vector store %q is unavailable", storeName) + } - added, skipped, err = mgr.Add(ctx, storeName, embeddingModel, embedder, store, entries) + if _, err = mgr.EnsureLoaded(ctx, storeName, embeddingModel, embeddingFingerprint, embedder, store); err != nil { + return 0, 0, Stats{}, err + } + added, skipped, err = mgr.Add(ctx, storeName, embeddingModel, embeddingFingerprint, embedder, store, entries) if err != nil { return added, skipped, Stats{}, err } diff --git a/core/services/routing/router/knn.go b/core/services/routing/router/knn.go index 2203278cbd12..529051e6fec7 100644 --- a/core/services/routing/router/knn.go +++ b/core/services/routing/router/knn.go @@ -7,6 +7,7 @@ import ( "encoding/json" "fmt" "sort" + "strings" "time" "github.com/mudler/LocalAI/core/backend" @@ -226,16 +227,33 @@ func EncodeCorpusEntry(id string, labels []string) ([]byte, error) { if id == "" { return nil, fmt.Errorf("corpus entry needs an id (EntryID of its text)") } - if len(labels) == 0 { - return nil, fmt.Errorf("corpus entry needs at least one label") + if err := validateCorpusLabels(labels); err != nil { + return nil, err } return json.Marshal(corpusEntry{ID: id, Labels: labels}) } func decodeCorpusEntry(b []byte) (corpusEntry, bool) { var e corpusEntry - if err := json.Unmarshal(b, &e); err != nil || len(e.Labels) == 0 { + if err := json.Unmarshal(b, &e); err != nil || e.ID == "" || validateCorpusLabels(e.Labels) != nil { return corpusEntry{}, false } return e, true } + +func validateCorpusLabels(labels []string) error { + if len(labels) == 0 { + return fmt.Errorf("corpus entry needs at least one label") + } + seen := make(map[string]struct{}, len(labels)) + for _, label := range labels { + if strings.TrimSpace(label) == "" { + return fmt.Errorf("corpus entry labels must not be empty") + } + if _, duplicate := seen[label]; duplicate { + return fmt.Errorf("corpus entry has duplicate label %q", label) + } + seen[label] = struct{}{} + } + return nil +} diff --git a/core/services/routing/router/knn_test.go b/core/services/routing/router/knn_test.go index 7b732c813f4f..9ec0bf5d9054 100644 --- a/core/services/routing/router/knn_test.go +++ b/core/services/routing/router/knn_test.go @@ -185,6 +185,21 @@ var _ = Describe("KNNClassifier", func() { Expect(err).To(HaveOccurred()) }) + It("rejects duplicate labels at encode time", func() { + _, err := router.EncodeCorpusEntry("e1", []string{"code", "code"}) + Expect(err).To(MatchError(ContainSubstring("duplicate label"))) + }) + + It("does not let duplicate labels in a corrupt payload inflate votes", func() { + store := &scriptedKNNStore{neighbors: []backend.Neighbor{ + {Similarity: 0.99, Payload: []byte(`{"id":"e1","labels":["code","code"]}`)}, + }} + c := router.NewKNNClassifier(embedder, store, router.KNNClassifierOptions{}) + d, err := c.Classify(ctx, probe) + Expect(err).NotTo(HaveOccurred()) + Expect(d.Labels).To(BeEmpty()) + }) + It("rejects corpus entries without an id at encode time", func() { _, err := router.EncodeCorpusEntry("", []string{"code"}) Expect(err).To(HaveOccurred()) diff --git a/docs/content/features/middleware.md b/docs/content/features/middleware.md index f4c735b10734..3903b79b9de1 100644 --- a/docs/content/features/middleware.md +++ b/docs/content/features/middleware.md @@ -443,6 +443,7 @@ router: fallback: gpt-4o-proxy # used whenever the prompt is unlike all corpus entries knn: embedding_model: nomic-embed-text-v1.5 + # embedding_revision: "2026-07" # bump for remote/in-place weight changes LocalAI cannot identify k: 3 # neighbours that vote (default 3) similarity_threshold: 0.80 # the epistemic gate (default 0.80) vote_threshold: 0.5 # weighted vote share a label needs (default 0.5) @@ -509,8 +510,9 @@ curl -X DELETE http://localhost:8080/api/router/smart-router/corpus ``` Entry labels must be declared in `policies` (same invariant as -candidate labels), and duplicate texts are skipped rather than -double-weighted. Label your exemplars with *outcomes*, not topics, +candidate labels). Empty and duplicate labels are rejected, and +duplicate texts are skipped rather than double-weighted. Label your exemplars +with *outcomes*, not topics, when routing for difficulty: an entry recording "the small model handled prompts like this" is exactly as useful as one recording that it failed — grade a sample of production traffic against your @@ -519,12 +521,19 @@ candidates and seed both. #### Persistence The corpus is persisted as one JSONL file per router under -`/router-corpus/` (text, labels, vector, embedding-model -fingerprint) — **the file is the source of truth** and survives -restarts; the local-store index is rebuilt from it at classifier build -time without re-embedding. Changing `knn.embedding_model` re-embeds -the corpus on the next load (restart LocalAI if the old index was -already live — mixed embedding spaces cannot be served). +`/router-corpus/` (text, labels, vector, embedding-model name, +and embedding fingerprint) — **the file is the source of truth** and +survives restarts; the local-store index is rebuilt from it at classifier +build time without re-embedding. The fingerprint follows the effective +embedding-model config and local artifact identity, so changing the model or +replacing its local weights re-embeds the corpus on the next process load. +For remote embedding services whose weights can change invisibly, bump +`knn.embedding_revision` explicitly. + +If an embedding fingerprint changes after that corpus is already present in +the live in-memory index, LocalAI fails the classifier build instead of +querying mixed embedding spaces. Restart LocalAI to rebuild the empty live +index and re-embed the persisted entries. #### Tuning notes diff --git a/pkg/mcp/localaitools/inproc/client.go b/pkg/mcp/localaitools/inproc/client.go index 3147ad9fa7ad..16e96a5482b3 100644 --- a/pkg/mcp/localaitools/inproc/client.go +++ b/pkg/mcp/localaitools/inproc/client.go @@ -75,9 +75,10 @@ type Client struct { // "unavailable" error. The factories mirror the middleware's // ClassifierDeps so the tools and the request path resolve models // and store namespaces identically. - RouterCorpus *corpus.Manager - RouterEmbedder func(modelName string) backend.Embedder - RouterVectorStore func(storeName string) backend.VectorStore + RouterCorpus *corpus.Manager + RouterEmbedder func(modelName string) backend.Embedder + RouterEmbedderFingerprint func(modelName string) (string, error) + RouterVectorStore func(storeName string) backend.VectorStore modelAdmin *modeladmin.ConfigService } @@ -919,7 +920,7 @@ func (c *Client) GetRouterCorpusStats(_ context.Context, routerModel string) (*l } func (c *Client) SeedRouterCorpus(ctx context.Context, req localaitools.RouterCorpusSeedRequest) (*localaitools.RouterCorpusSeedResult, error) { - if c.RouterCorpus == nil || c.RouterEmbedder == nil || c.RouterVectorStore == nil { + if c.RouterCorpus == nil || c.RouterEmbedder == nil || c.RouterEmbedderFingerprint == nil || c.RouterVectorStore == nil { return nil, errors.New("router corpus manager unavailable") } cfg, storeName, err := corpus.ResolveKNNRouter(c.ConfigLoader, c.AppConfig, req.Router) @@ -931,7 +932,7 @@ func (c *Client) SeedRouterCorpus(ctx context.Context, req localaitools.RouterCo entries = append(entries, corpus.Entry{Text: e.Text, Labels: e.Labels}) } added, skipped, stats, err := corpus.Seed(ctx, c.RouterCorpus, cfg, storeName, - c.RouterEmbedder, c.RouterVectorStore, entries) + c.RouterEmbedder, c.RouterEmbedderFingerprint, c.RouterVectorStore, entries) if err != nil { return nil, err } From 02a144a169425614ff745d56014ca81d91ee5ba7 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Wed, 15 Jul 2026 11:24:57 +0100 Subject: [PATCH 05/10] chore(mcp): align corpus tool prompts and the mutating-tool safety list Co-Authored-By: Claude Fable 5 --- pkg/mcp/localaitools/prompts/10_safety.md | 2 +- pkg/mcp/localaitools/prompts/20_tools.md | 11 +++++ .../prompts/skills/manage_router_corpus.md | 12 +++++ pkg/mcp/localaitools/prompts_test.go | 12 +++++ pkg/mcp/localaitools/server_test.go | 47 +++---------------- pkg/mcp/localaitools/tools.go | 27 ++++++++++- pkg/mcp/localaitools/tools_middleware.go | 2 +- 7 files changed, 68 insertions(+), 45 deletions(-) create mode 100644 pkg/mcp/localaitools/prompts/skills/manage_router_corpus.md diff --git a/pkg/mcp/localaitools/prompts/10_safety.md b/pkg/mcp/localaitools/prompts/10_safety.md index 3f338f32083c..0acea2d658f1 100644 --- a/pkg/mcp/localaitools/prompts/10_safety.md +++ b/pkg/mcp/localaitools/prompts/10_safety.md @@ -2,7 +2,7 @@ These rules are non-negotiable. The user trusts you to operate their server without unintended changes. -1. **Confirm before mutating.** Before calling any of these tools — `install_model`, `import_model_uri`, `delete_model`, `install_backend`, `upgrade_backend`, `edit_model_config`, `reload_models`, `load_model`, `toggle_model_state`, `toggle_model_pinned`, `create_voice_profile`, `delete_voice_profile` — first state in plain language what you are about to do (which tool, which target, which arguments) and wait for the user's explicit confirmation in the next turn. "Yes", "do it", "go ahead", "proceed" all count as confirmation. Anything else does not. +1. **Confirm before mutating.** Before calling any of these tools — `install_model`, `import_model_uri`, `delete_model`, `install_backend`, `upgrade_backend`, `edit_model_config`, `reload_models`, `load_model`, `toggle_model_state`, `toggle_model_pinned`, `set_branding`, `set_alias`, `seed_router_corpus`, `clear_router_corpus`, `create_voice_profile`, `delete_voice_profile`, `set_node_vram_budget` — first state in plain language what you are about to do (which tool, which target, which arguments) and wait for the user's explicit confirmation in the next turn. "Yes", "do it", "go ahead", "proceed" all count as confirmation. Anything else does not. 2. **Disambiguate before mutating.** If the user's request is ambiguous (several gallery candidates match, the model name has multiple installed versions, the backend has variants), present the candidates as a numbered list and ask the user to pick before calling any mutating tool. diff --git a/pkg/mcp/localaitools/prompts/20_tools.md b/pkg/mcp/localaitools/prompts/20_tools.md index a6cdc9b396f7..10a559dfa0e0 100644 --- a/pkg/mcp/localaitools/prompts/20_tools.md +++ b/pkg/mcp/localaitools/prompts/20_tools.md @@ -15,6 +15,13 @@ The MCP `tools/list` endpoint also exposes the full input schema for each of the - `system_info` — LocalAI version, paths, distributed flag, loaded models, installed backends. - `list_nodes` — List federated worker nodes (only useful in distributed mode). - `list_voice_profiles` — List reusable voice-cloning profiles and their stable TTS voice URIs. +- `get_branding` — Read the resolved instance name, tagline, and branding asset URLs. +- `get_usage_stats` — Read token/request usage aggregates when usage tracking is enabled. +- `get_pii_events` — Inspect recent PII-filter events without returning redacted request bodies. +- `get_middleware_status` — Inspect middleware configuration and health. +- `get_router_decisions` — Inspect recent router decisions and classifier signals. +- `get_router_corpus_stats` — Inspect a KNN router corpus by count and label only; exemplar texts are never returned. +- `list_aliases` — List configured model aliases and their targets. ## Mutating (require user confirmation per safety rule 1) @@ -28,5 +35,9 @@ The MCP `tools/list` endpoint also exposes the full input schema for each of the - `load_model` — Pre-load a model into memory so the first request pays no cold-start cost. For a realtime pipeline model, every sub-model (VAD, transcription, LLM, TTS, sound_detection, voice_recognition) is loaded. Inverse of stopping a model. - `toggle_model_state` — Enable or disable a model (`action`: `enable` or `disable`). - `toggle_model_pinned` — Pin or unpin a model (`action`: `pin` or `unpin`). +- `set_branding` — Change the instance name or tagline. +- `set_alias` — Create, update, or remove a model alias. +- `seed_router_corpus` — Add validated labelled exemplars to a KNN router corpus. +- `clear_router_corpus` — Permanently clear a KNN router corpus and its live index entries. - `create_voice_profile` — Save a consent-confirmed base64 PCM-WAV reference and exact transcript for reuse in TTS. - `delete_voice_profile` — Permanently delete a saved voice profile by UUID. diff --git a/pkg/mcp/localaitools/prompts/skills/manage_router_corpus.md b/pkg/mcp/localaitools/prompts/skills/manage_router_corpus.md new file mode 100644 index 000000000000..86f98d59efbe --- /dev/null +++ b/pkg/mcp/localaitools/prompts/skills/manage_router_corpus.md @@ -0,0 +1,12 @@ +# Skill: Manage a router corpus + +Use this when the user wants to inspect, seed, repair, or clear the labelled exemplar corpus behind a `knn` router. + +1. Call `get_model_config` for the router and verify that `router.classifier` is `knn`, `router.knn.embedding_model` is set, and the intended labels exist in `router.policies`. +2. Call `get_router_corpus_stats` and report the store name, embedding model, total entries, and per-label counts. Never claim to have inspected exemplar text: this interface deliberately returns counts only. +3. Before seeding, validate that every entry has non-empty text, at least one label, no duplicate labels, and only labels declared in `router.policies`. Summarise the number of entries and per-label additions, then ask for explicit confirmation. +4. On confirmation, call `seed_router_corpus`. If it reports a fingerprint or live-index mismatch, surface the error verbatim; do not retry automatically or guess that stale vectors are safe. +5. Verify a successful seed with `get_router_corpus_stats` and report added, skipped, and new totals. +6. Clearing is destructive. State the router and current entry count, ask for explicit confirmation, then call `clear_router_corpus`. Verify the total is zero with `get_router_corpus_stats`. + +Never call `seed_router_corpus` or `clear_router_corpus` without explicit confirmation in the immediately preceding user turn. diff --git a/pkg/mcp/localaitools/prompts_test.go b/pkg/mcp/localaitools/prompts_test.go index c85efdce38e9..1f0c762fdda3 100644 --- a/pkg/mcp/localaitools/prompts_test.go +++ b/pkg/mcp/localaitools/prompts_test.go @@ -49,9 +49,21 @@ var _ = Describe("SystemPrompt assembler", func() { "Skill: Upgrade a backend", "Skill: System status", "Skill: Safely edit a model config", + "Skill: Manage a router corpus", + "get_router_corpus_stats", } for _, s := range mustContain { Expect(out).To(ContainSubstring(s), "system prompt missing required anchor %q", s) } }) + + It("names every mutating tool in the confirmation safety rule", func() { + raw, err := promptsFS.ReadFile("prompts/10_safety.md") + Expect(err).NotTo(HaveOccurred()) + out := string(raw) + for _, name := range mutatingToolNames { + Expect(out).To(ContainSubstring("`"+name+"`"), + "system prompt confirmation rule is missing mutating tool %q", name) + } + }) }) diff --git a/pkg/mcp/localaitools/server_test.go b/pkg/mcp/localaitools/server_test.go index 99d51495e479..e8666739b71a 100644 --- a/pkg/mcp/localaitools/server_test.go +++ b/pkg/mcp/localaitools/server_test.go @@ -70,46 +70,6 @@ func resultText(res *mcp.CallToolResult) string { return b.String() } -// expectedFullCatalog is the tool set when DisableMutating=false. Sorted. -// References the Tool* constants so a rename can't drift code from tests. -var expectedFullCatalog = sortedStrings( - ToolDeleteModel, - ToolEditModelConfig, - ToolGallerySearch, - ToolGetBranding, - ToolGetJobStatus, - ToolGetMiddlewareStatus, - ToolGetModelConfig, - ToolGetPIIEvents, - ToolGetRouterCorpusStats, - ToolGetRouterDecisions, - ToolGetUsageStats, - ToolImportModelURI, - ToolInstallBackend, - ToolInstallModel, - ToolListBackends, - ToolListGalleries, - ToolListAliases, - ToolListInstalledModels, - ToolListKnownBackends, - ToolListNodes, - ToolListVoiceProfiles, - ToolLoadModel, - ToolReloadModels, - ToolSeedRouterCorpus, - ToolClearRouterCorpus, - ToolSetAlias, - ToolSetBranding, - ToolSystemInfo, - ToolToggleModelPinned, - ToolToggleModelState, - ToolUpgradeBackend, - ToolVRAMEstimate, - ToolCreateVoiceProfile, - ToolDeleteVoiceProfile, - ToolSetNodeVRAMBudget, -) - // expectedReadOnlyCatalog is the tool set when DisableMutating=true. Sorted. var expectedReadOnlyCatalog = sortedStrings( ToolGallerySearch, @@ -121,9 +81,9 @@ var expectedReadOnlyCatalog = sortedStrings( ToolGetRouterCorpusStats, ToolGetRouterDecisions, ToolGetUsageStats, - ToolListAliases, ToolListBackends, ToolListGalleries, + ToolListAliases, ToolListInstalledModels, ToolListKnownBackends, ToolListNodes, @@ -132,6 +92,11 @@ var expectedReadOnlyCatalog = sortedStrings( ToolVRAMEstimate, ) +// expectedFullCatalog derives from the read-only catalog plus the canonical +// mutating list used by the safety-prompt coverage test. Registering a new +// mutator now fails both catalog and prompt coverage until that list is updated. +var expectedFullCatalog = sortedStrings(append(append([]string(nil), expectedReadOnlyCatalog...), mutatingToolNames...)...) + func sortedStrings(in ...string) []string { out := append([]string(nil), in...) sort.Strings(out) diff --git a/pkg/mcp/localaitools/tools.go b/pkg/mcp/localaitools/tools.go index 7245eb0a2e61..4f06594ceefa 100644 --- a/pkg/mcp/localaitools/tools.go +++ b/pkg/mcp/localaitools/tools.go @@ -4,8 +4,8 @@ package localaitools // constants — never bare strings — when registering tools, asserting the // catalog in tests, or referencing tool names from other packages. The // embedded skill prompts under prompts/ keep the bare strings because -// go:embed-ed markdown can't reference Go constants; TestPromptsContain -// SafetyAnchors guards that those strings stay aligned. +// go:embed-ed markdown can't reference Go constants; prompts_test.go guards +// that the mutating names stay aligned with the confirmation rule. const ( // Read-only tools. ToolGallerySearch = "gallery_search" @@ -55,3 +55,26 @@ const ( // Options.ServerName is empty. Use the constant when you want a stable // reference across packages (e.g. test fixtures, CLI defaults). const DefaultServerName = "localai-admin" + +// mutatingToolNames is the canonical safety-prompt coverage list. Registration +// remains grouped by feature, while prompts_test.go mechanically ensures every +// state-changing tool is named in the confirmation rule. +var mutatingToolNames = []string{ + ToolInstallModel, + ToolImportModelURI, + ToolDeleteModel, + ToolEditModelConfig, + ToolReloadModels, + ToolLoadModel, + ToolInstallBackend, + ToolUpgradeBackend, + ToolToggleModelState, + ToolToggleModelPinned, + ToolSetBranding, + ToolSetAlias, + ToolSeedRouterCorpus, + ToolClearRouterCorpus, + ToolCreateVoiceProfile, + ToolDeleteVoiceProfile, + ToolSetNodeVRAMBudget, +} diff --git a/pkg/mcp/localaitools/tools_middleware.go b/pkg/mcp/localaitools/tools_middleware.go index d761768a80d6..528eb5811b9b 100644 --- a/pkg/mcp/localaitools/tools_middleware.go +++ b/pkg/mcp/localaitools/tools_middleware.go @@ -67,7 +67,7 @@ func registerMiddlewareTools(s *mcp.Server, client LocalAIClient, opts Options) mcp.AddTool(s, &mcp.Tool{ Name: ToolSeedRouterCorpus, - Description: "Add labelled example prompts to a knn router's corpus. Entries are embedded server-side with the router's knn.embedding_model, persisted, and indexed immediately — routing behaviour changes right away, so confirm with the user first (safety rule 1). Labels must be declared in the router's policies; duplicate texts are skipped.", + Description: "Add labelled example prompts to a knn router's corpus. Entries are embedded server-side with the router's knn.embedding_model, persisted, and indexed immediately — routing behaviour changes right away, so confirm with the user first (safety rule 1). Labels must be declared in the router's policies and unique per entry; duplicate texts are skipped.", }, func(ctx context.Context, _ *mcp.CallToolRequest, args RouterCorpusSeedRequest) (*mcp.CallToolResult, any, error) { if args.Router == "" { return errorResultf("router is required"), nil, nil From 94080276378781e820da5e78dffbded884605a03 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Wed, 15 Jul 2026 11:29:42 +0100 Subject: [PATCH 06/10] feat(proto,backend): report embedding shape from the llama-cpp backend Co-Authored-By: Claude Fable 5 --- backend/backend.proto | 9 +++++ backend/cpp/llama-cpp/grpc-server.cpp | 54 ++++++++++++--------------- tests/e2e/mock-backend/main.go | 33 +++++++++++++++- 3 files changed, 65 insertions(+), 31 deletions(-) diff --git a/backend/backend.proto b/backend/backend.proto index ad62c6df07ae..ca6d402afdbe 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -472,6 +472,15 @@ message Result { message EmbeddingResult { repeated float embeddings = 1; + // Shape of the payload above: dim is the embedding width, tokens is the + // number of vectors packed into `embeddings` (1 when the backend pooled + // server-side, N with pooling:none; total across prompts if a request + // carried several). tokens=0/dim=0 means the backend predates shape + // reporting. prompt_tokens is the number of prompt tokens evaluated, for + // usage accounting. + int32 tokens = 2; + int32 dim = 3; + int32 prompt_tokens = 4; } message TranscriptRequest { diff --git a/backend/cpp/llama-cpp/grpc-server.cpp b/backend/cpp/llama-cpp/grpc-server.cpp index cd4fb341d326..2c6b6686f105 100644 --- a/backend/cpp/llama-cpp/grpc-server.cpp +++ b/backend/cpp/llama-cpp/grpc-server.cpp @@ -2772,42 +2772,36 @@ class BackendServiceImpl final : public backend::Backend::Service { return grpc::Status(grpc::StatusCode::INTERNAL, all_results.error->to_json().value("message", "Error in receiving results")); } - // Collect responses - json responses = json::array(); + // Extract the embeddings typed, straight from the task results (no + // JSON round-trip), and report the payload shape alongside the same + // flat float array as before: dim is the embedding width, tokens the + // number of vectors packed into `embeddings` (1 per prompt when the + // server pooled, one per token with pooling:none; summed across + // prompts if the request carried several), prompt_tokens the prompt + // tokens evaluated, for usage accounting. Consumers seeing 0/0 know + // the backend predates shape reporting. + int32_t n_vectors = 0; + int32_t dim = 0; + int32_t prompt_tokens = 0; for (auto & res : all_results.results) { - GGML_ASSERT(dynamic_cast(res.get()) != nullptr); - responses.push_back(res->to_json()); - } - - std::cout << "[DEBUG] Responses size: " << responses.size() << std::endl; - - // Process the responses and extract embeddings - for (const auto & response_elem : responses) { - // Check if the response has an "embedding" field - if (response_elem.contains("embedding")) { - json embedding_data = json_value(response_elem, "embedding", json::array()); - - if (embedding_data.is_array() && !embedding_data.empty()) { - for (const auto & embedding_vector : embedding_data) { - if (embedding_vector.is_array()) { - for (const auto & embedding_value : embedding_vector) { - embeddingResult->add_embeddings(embedding_value.get()); - } - } - } + auto * embd_res = dynamic_cast(res.get()); + GGML_ASSERT(embd_res != nullptr); + prompt_tokens += embd_res->n_tokens; + for (const auto & vec : embd_res->embedding) { + for (const float value : vec) { + embeddingResult->add_embeddings(value); } - } else { - // Check if the response itself contains the embedding data directly - if (response_elem.is_array()) { - for (const auto & embedding_value : response_elem) { - embeddingResult->add_embeddings(embedding_value.get()); - } + if (!vec.empty()) { + n_vectors++; + dim = (int32_t) vec.size(); } } } + embeddingResult->set_tokens(n_vectors); + embeddingResult->set_dim(dim); + embeddingResult->set_prompt_tokens(prompt_tokens); - - + std::cout << "[DEBUG] Embedding vectors: " << n_vectors << " x " << dim << std::endl; return grpc::Status::OK; } diff --git a/tests/e2e/mock-backend/main.go b/tests/e2e/mock-backend/main.go index 10fb25296e24..c9ab48927eec 100644 --- a/tests/e2e/mock-backend/main.go +++ b/tests/e2e/mock-backend/main.go @@ -392,12 +392,43 @@ func mockToolNameFromRequest(in *pb.PredictOptions) string { func (m *MockBackend) Embedding(ctx context.Context, in *pb.PredictOptions) (*pb.EmbeddingResult, error) { xlog.Debug("Embedding called", "prompt", in.Prompt) + // Deterministic per-token mode for Go-side pooling tests: a prompt + // carrying the "per-token:" marker yields len(fields) vectors of dim 8 + // with vec[i][j] = (i+1)/(j+2), so endpoint tests can assert exact + // pooled goldens. The marker may sit mid-prompt (the embeddings + // messages[] path renders conversations as ": " lines), + // so match anywhere and tokenize what follows the first occurrence. + if idx := strings.Index(in.Prompt, "per-token:"); idx >= 0 { + fields := strings.Fields(in.Prompt[idx+len("per-token:"):]) + tokens := len(fields) + const dim = 8 + flat := make([]float32, 0, tokens*dim) + for i := 0; i < tokens; i++ { + for j := 0; j < dim; j++ { + flat = append(flat, float32(i+1)/float32(j+2)) + } + } + return &pb.EmbeddingResult{ + Embeddings: flat, + Tokens: int32(tokens), + Dim: dim, + PromptTokens: int32(tokens), + }, nil + } + // Legacy mode: a prompt carrying "no-shape:" simulates a backend built + // before EmbeddingResult carried shape fields (tokens/dim left 0), so + // tests can assert the fail-closed error when Go-side pooling is + // requested against such a backend. + legacyShape := strings.Contains(in.Prompt, "no-shape:") // Return a mock embedding vector of 768 dimensions embeddings := make([]float32, 768) for i := range embeddings { embeddings[i] = float32(i%100) / 100.0 // Pattern: 0.0, 0.01, 0.02, ..., 0.99, 0.0, ... } - return &pb.EmbeddingResult{Embeddings: embeddings}, nil + if legacyShape { + return &pb.EmbeddingResult{Embeddings: embeddings}, nil + } + return &pb.EmbeddingResult{Embeddings: embeddings, Tokens: 1, Dim: 768, PromptTokens: 1}, nil } func (m *MockBackend) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest) (*pb.Result, error) { From 067678fab5aa53a4d37e972971e5bb9fac8ff5a0 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Wed, 15 Jul 2026 11:38:19 +0100 Subject: [PATCH 07/10] =?UTF-8?q?feat(embeddings):=20Go-side=20pooling=20?= =?UTF-8?q?=E2=80=94=20mean/last/decayed=5Fmean=20with=20half-life?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 --- core/backend/embeddings.go | 43 ++++++- core/backend/pooling.go | 187 +++++++++++++++++++++++++++++ core/backend/pooling_test.go | 162 +++++++++++++++++++++++++ core/config/meta/registry.go | 24 ++++ core/config/model_config.go | 88 ++++++++++++++ core/config/pooling_config_test.go | 111 +++++++++++++++++ core/schema/prediction.go | 13 ++ 7 files changed, 624 insertions(+), 4 deletions(-) create mode 100644 core/backend/pooling.go create mode 100644 core/backend/pooling_test.go create mode 100644 core/config/pooling_config_test.go diff --git a/core/backend/embeddings.go b/core/backend/embeddings.go index eff88ef04b19..66300dc56971 100644 --- a/core/backend/embeddings.go +++ b/core/backend/embeddings.go @@ -9,9 +9,44 @@ import ( "github.com/mudler/LocalAI/core/trace" "github.com/mudler/LocalAI/pkg/grpc" + "github.com/mudler/LocalAI/pkg/grpc/proto" model "github.com/mudler/LocalAI/pkg/model" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) +// finishEmbeddingResult applies the model's Go-side pooling scheme to a raw +// per-token EmbeddingResult, or passes the backend's vector through +// untouched when pooling is delegated to the backend ("" / "backend" — +// today's exact behavior). Fails closed when a Go-side scheme is requested +// but the backend didn't report the payload shape (a build predating +// EmbeddingResult.tokens/dim), since silently treating a backend-pooled +// vector as per-token data would corrupt the embedding space. +func finishEmbeddingResult(res *proto.EmbeddingResult, modelConfig config.ModelConfig) ([]float32, error) { + scheme := modelConfig.Pooling + if scheme == "" || scheme == PoolingBackend { + return res.Embeddings, nil + } + if res.GetDim() == 0 { + return nil, fmt.Errorf( + "pooling %q needs per-token embeddings but the backend returned no shape: rebuild/update the backend so EmbeddingResult reports tokens/dim, and set options [\"pooling:none\"] on the model", + scheme) + } + return PoolEmbeddingResult(res, scheme, + float64(modelConfig.PoolingHalfLifeTokens), + embdNormalizeFromOptions(modelConfig.Options)) +} + +// mapEmbeddingGRPCError turns a gRPC ResourceExhausted — the per-token +// payload of a very long conversation exceeding the 50MB message cap — +// into an actionable message; everything else passes through unchanged. +func mapEmbeddingGRPCError(err error) error { + if status.Code(err) == codes.ResourceExhausted { + return fmt.Errorf("conversation too long for per-token embeddings (gRPC message limit exceeded): %w", err) + } + return err +} + // Embedder produces a fixed-dimension vector from a prompt. The // router's L2 embedding cache uses it to look up semantically-similar // past decisions. @@ -66,19 +101,19 @@ func ModelEmbedding(ctx context.Context, s string, tokens []int, loader *model.M res, err := model.Embeddings(appConfig.Context, predictOptions) if err != nil { - return nil, err + return nil, mapEmbeddingGRPCError(err) } - return res.Embeddings, nil + return finishEmbeddingResult(res, modelConfig) } predictOptions.Embeddings = s res, err := model.Embeddings(appConfig.Context, predictOptions) if err != nil { - return nil, err + return nil, mapEmbeddingGRPCError(err) } - return res.Embeddings, nil + return finishEmbeddingResult(res, modelConfig) } default: fn = func() ([]float32, error) { diff --git a/core/backend/pooling.go b/core/backend/pooling.go new file mode 100644 index 000000000000..f4f1799b8b07 --- /dev/null +++ b/core/backend/pooling.go @@ -0,0 +1,187 @@ +package backend + +import ( + "fmt" + "math" + "strconv" + "strings" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/grpc/proto" +) + +// Go-side embedding pooling schemes. The canonical strings live on the +// config package (config.Pooling* — mirroring the ScoreNormalization* +// pattern) so ModelConfig.Validate can reject unknown values without +// importing this package; these aliases are the API the backend layer +// (and the HTTP endpoints) program against. +const ( + // PoolingBackend leaves pooling to the inference backend — the exact + // pre-existing behavior of the embeddings path (also selected by ""). + PoolingBackend = config.PoolingBackend + // PoolingMean averages the per-token vectors. + PoolingMean = config.PoolingMean + // PoolingLast selects the last token's vector. + PoolingLast = config.PoolingLast + // PoolingDecayedMean is a mean weighted toward the most recent tokens: + // w_i = 2^(-(T-1-i)/H) for token i of T, with half-life H tokens. + PoolingDecayedMean = config.PoolingDecayedMean + + // DefaultPoolingHalfLifeTokens is the half-life used by + // PoolingDecayedMean when the model config / request doesn't set one. + DefaultPoolingHalfLifeTokens = 256 +) + +// reshapeEmbeddings views the flat float payload of an EmbeddingResult as +// tokens rows of dim columns. The gRPC contract packs vectors row-major +// (vector 0 first), so row i aliases flat[i*dim : (i+1)*dim]. +func reshapeEmbeddings(flat []float32, tokens, dim int) ([][]float32, error) { + if tokens <= 0 || dim <= 0 { + return nil, fmt.Errorf("invalid embedding shape: %d vectors x %d dims", tokens, dim) + } + if len(flat) != tokens*dim { + return nil, fmt.Errorf("embedding payload of %d floats does not match reported shape %d vectors x %d dims", len(flat), tokens, dim) + } + vecs := make([][]float32, tokens) + for i := range vecs { + vecs[i] = flat[i*dim : (i+1)*dim] + } + return vecs, nil +} + +// poolMean averages the per-token vectors with float64 accumulators. +func poolMean(vecs [][]float32) []float32 { + dim := len(vecs[0]) + acc := make([]float64, dim) + for _, v := range vecs { + for j, x := range v { + acc[j] += float64(x) + } + } + out := make([]float32, dim) + for j := range out { + out[j] = float32(acc[j] / float64(len(vecs))) + } + return out +} + +// poolLast returns (a copy of) the last token's vector. +func poolLast(vecs [][]float32) []float32 { + last := vecs[len(vecs)-1] + out := make([]float32, len(last)) + copy(out, last) + return out +} + +// poolDecayedMean computes a weighted mean over the per-token vectors with +// exponentially decaying weights anchored at the last token: token i of T +// gets w_i = 2^(-(T-1-i)/halfLife), so the last token always weighs 1 and a +// token halfLife positions earlier weighs 0.5. Accumulation is in float64. +func poolDecayedMean(vecs [][]float32, halfLife float64) []float32 { + dim := len(vecs[0]) + T := len(vecs) + acc := make([]float64, dim) + wsum := 0.0 + for i, v := range vecs { + w := math.Exp2(-float64(T-1-i) / halfLife) + wsum += w + for j, x := range v { + acc[j] += w * float64(x) + } + } + out := make([]float32, dim) + for j := range out { + out[j] = float32(acc[j] / wsum) + } + return out +} + +// normalizeEmbedding is an exact port of llama.cpp's common_embd_normalize +// (backend/cpp/llama-cpp/llama.cpp/common/common.cpp), applied after Go-side +// pooling because per-token vectors arrive RAW from llama.cpp with +// pooling:none (the server only normalizes vectors it pooled itself). +// embdNorm: <0 none, 0 max-abs scaled to int16 range (/32760.0), 2 L2 +// (llama.cpp default), anything else p-norm with p=embdNorm (1 = taxicab). +func normalizeEmbedding(v []float32, embdNorm int) []float32 { + sum := 0.0 + switch { + case embdNorm < 0: // no normalisation + sum = 1.0 + case embdNorm == 0: // max absolute + for _, x := range v { + if a := math.Abs(float64(x)); sum < a { + sum = a + } + } + sum /= 32760.0 // make an int16 range + case embdNorm == 2: // euclidean + for _, x := range v { + sum += float64(x) * float64(x) + } + sum = math.Sqrt(sum) + default: // p-norm (euclidean is p-norm p=2) + for _, x := range v { + sum += math.Pow(math.Abs(float64(x)), float64(embdNorm)) + } + sum = math.Pow(sum, 1.0/float64(embdNorm)) + } + + // llama.cpp computes the reciprocal as a float32 and multiplies in + // float32; mirror that so both paths yield bit-identical vectors. + var norm float32 + if sum > 0.0 { + norm = float32(1.0 / sum) + } + out := make([]float32, len(v)) + for i, x := range v { + out[i] = x * norm + } + return out +} + +// embdNormalizeFromOptions extracts the load-time embd_normalize backend +// option ("embd_normalize:", alias "embedding_normalize:") the same +// way the llama-cpp gRPC server parses it, so Go-side pooling normalizes +// with the exact norm the backend would have applied had it pooled +// server-side. Defaults to 2 (L2) like llama.cpp; unparsable values are +// ignored (llama.cpp swallows std::stoi failures). +func embdNormalizeFromOptions(options []string) int { + embdNorm := 2 + for _, opt := range options { + name, val, found := strings.Cut(opt, ":") + if !found || (name != "embd_normalize" && name != "embedding_normalize") { + continue + } + if n, err := strconv.Atoi(strings.TrimSpace(val)); err == nil { + embdNorm = n + } + } + return embdNorm +} + +// PoolEmbeddingResult reduces a per-token EmbeddingResult (the backend ran +// with pooling:none) to a single vector using scheme, then normalizes it +// with llama.cpp's common_embd_normalize semantics. halfLife only applies +// to PoolingDecayedMean; non-positive values fall back to +// DefaultPoolingHalfLifeTokens. +func PoolEmbeddingResult(res *proto.EmbeddingResult, scheme string, halfLife float64, embdNorm int) ([]float32, error) { + vecs, err := reshapeEmbeddings(res.GetEmbeddings(), int(res.GetTokens()), int(res.GetDim())) + if err != nil { + return nil, err + } + var pooled []float32 + switch scheme { + case PoolingMean: + pooled = poolMean(vecs) + case PoolingLast: + pooled = poolLast(vecs) + case PoolingDecayedMean: + if halfLife <= 0 { + halfLife = DefaultPoolingHalfLifeTokens + } + pooled = poolDecayedMean(vecs, halfLife) + default: + return nil, fmt.Errorf("unknown Go-side pooling scheme %q (expected %q, %q or %q)", scheme, PoolingMean, PoolingLast, PoolingDecayedMean) + } + return normalizeEmbedding(pooled, embdNorm), nil +} diff --git a/core/backend/pooling_test.go b/core/backend/pooling_test.go new file mode 100644 index 000000000000..78aa1734bb08 --- /dev/null +++ b/core/backend/pooling_test.go @@ -0,0 +1,162 @@ +package backend + +import ( + "math" + + "github.com/mudler/LocalAI/pkg/grpc/proto" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Go-side embedding pooling", func() { + Describe("reshapeEmbeddings", func() { + It("views a flat payload as tokens x dim rows", func() { + vecs, err := reshapeEmbeddings([]float32{1, 2, 3, 4, 5, 6}, 2, 3) + Expect(err).ToNot(HaveOccurred()) + Expect(vecs).To(Equal([][]float32{{1, 2, 3}, {4, 5, 6}})) + }) + + It("rejects a payload that does not match the reported shape", func() { + _, err := reshapeEmbeddings([]float32{1, 2, 3, 4, 5}, 2, 3) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("does not match reported shape")) + }) + + It("rejects a non-positive shape", func() { + _, err := reshapeEmbeddings(nil, 0, 3) + Expect(err).To(HaveOccurred()) + _, err = reshapeEmbeddings(nil, 3, 0) + Expect(err).To(HaveOccurred()) + }) + }) + + Describe("poolMean", func() { + It("averages the per-token vectors", func() { + Expect(poolMean([][]float32{{1, 2}, {3, 4}})).To(Equal([]float32{2, 3})) + }) + }) + + Describe("poolLast", func() { + It("returns the last token's vector", func() { + Expect(poolLast([][]float32{{1, 2}, {3, 4}})).To(Equal([]float32{3, 4})) + }) + }) + + Describe("poolDecayedMean", func() { + It("weights tokens by 2^(-(T-1-i)/H): H=1, T=3 gives [0.25 0.5 1]/1.75", func() { + vecs := [][]float32{{1, 0}, {0, 1}, {1, 1}} + got := poolDecayedMean(vecs, 1) + Expect(got[0]).To(BeNumerically("~", 1.25/1.75, 1e-6)) + Expect(got[1]).To(BeNumerically("~", 1.5/1.75, 1e-6)) + }) + + It("approaches the plain mean as the half-life grows", func() { + vecs := [][]float32{{1, 0}, {0, 1}, {1, 1}} + got := poolDecayedMean(vecs, 1e12) + want := poolMean(vecs) + Expect(got[0]).To(BeNumerically("~", want[0], 1e-6)) + Expect(got[1]).To(BeNumerically("~", want[1], 1e-6)) + }) + }) + + Describe("single-token conversations", func() { + It("agree across mean, last and decayed_mean", func() { + vecs := [][]float32{{3, 4}} + Expect(poolMean(vecs)).To(Equal([]float32{3, 4})) + Expect(poolLast(vecs)).To(Equal([]float32{3, 4})) + got := poolDecayedMean(vecs, 256) + Expect(got[0]).To(BeNumerically("~", 3, 1e-6)) + Expect(got[1]).To(BeNumerically("~", 4, 1e-6)) + }) + }) + + Describe("normalizeEmbedding (common_embd_normalize port)", func() { + v := []float32{3, -4} + + It("passes through untouched for negative embd_norm", func() { + Expect(normalizeEmbedding(v, -1)).To(Equal([]float32{3, -4})) + }) + + It("scales to the int16 range for embd_norm 0 (max-abs)", func() { + // max-abs = 4, sum = 4/32760, norm = 32760/4 = 8190 + got := normalizeEmbedding(v, 0) + Expect(got[0]).To(BeNumerically("~", 3*8190.0, 1e-2)) + Expect(got[1]).To(BeNumerically("~", -4*8190.0, 1e-2)) + }) + + It("applies the taxicab norm for embd_norm 1", func() { + got := normalizeEmbedding(v, 1) + Expect(got[0]).To(BeNumerically("~", 3.0/7.0, 1e-6)) + Expect(got[1]).To(BeNumerically("~", -4.0/7.0, 1e-6)) + }) + + It("applies the L2 norm for embd_norm 2", func() { + got := normalizeEmbedding(v, 2) + Expect(got[0]).To(BeNumerically("~", 0.6, 1e-6)) + Expect(got[1]).To(BeNumerically("~", -0.8, 1e-6)) + }) + + It("applies a p-norm for embd_norm > 2", func() { + p3 := math.Cbrt(27 + 64) // (|3|^3 + |-4|^3)^(1/3) + got := normalizeEmbedding(v, 3) + Expect(got[0]).To(BeNumerically("~", 3.0/p3, 1e-5)) + Expect(got[1]).To(BeNumerically("~", -4.0/p3, 1e-5)) + }) + + It("maps the all-zero vector to all zeros instead of dividing by zero", func() { + Expect(normalizeEmbedding([]float32{0, 0, 0}, 2)).To(Equal([]float32{0, 0, 0})) + }) + }) + + Describe("embdNormalizeFromOptions", func() { + It("defaults to 2 (L2) like llama.cpp", func() { + Expect(embdNormalizeFromOptions(nil)).To(Equal(2)) + Expect(embdNormalizeFromOptions([]string{"pooling:none", "gpu"})).To(Equal(2)) + }) + + It("parses embd_normalize and its embedding_normalize alias", func() { + Expect(embdNormalizeFromOptions([]string{"embd_normalize:0"})).To(Equal(0)) + Expect(embdNormalizeFromOptions([]string{"embedding_normalize:-1"})).To(Equal(-1)) + Expect(embdNormalizeFromOptions([]string{"embd_normalize: 3"})).To(Equal(3)) + }) + + It("keeps the default when the value does not parse", func() { + Expect(embdNormalizeFromOptions([]string{"embd_normalize:junk"})).To(Equal(2)) + Expect(embdNormalizeFromOptions([]string{"embd_normalize"})).To(Equal(2)) + }) + }) + + Describe("PoolEmbeddingResult", func() { + res := &proto.EmbeddingResult{ + Embeddings: []float32{1, 2, 3, 4}, + Tokens: 2, + Dim: 2, + } + + It("reshapes, pools and normalizes", func() { + got, err := PoolEmbeddingResult(res, PoolingMean, 0, -1) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal([]float32{2, 3})) + }) + + It("L2-normalizes by default norm 2", func() { + got, err := PoolEmbeddingResult(res, PoolingLast, 0, 2) + Expect(err).ToNot(HaveOccurred()) + Expect(got[0]).To(BeNumerically("~", 0.6, 1e-6)) + Expect(got[1]).To(BeNumerically("~", 0.8, 1e-6)) + }) + + It("rejects unknown pooling schemes", func() { + _, err := PoolEmbeddingResult(res, "sideways", 0, 2) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("unknown Go-side pooling scheme")) + }) + + It("propagates shape mismatches", func() { + bad := &proto.EmbeddingResult{Embeddings: []float32{1, 2, 3}, Tokens: 2, Dim: 2} + _, err := PoolEmbeddingResult(bad, PoolingMean, 0, 2) + Expect(err).To(HaveOccurred()) + }) + }) +}) diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index 027d5bc1ad3e..4443c1b13853 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -257,6 +257,30 @@ func DefaultRegistry() map[string]FieldMetaOverride { Advanced: true, Order: 35, }, + "parameters.pooling": { + Section: "parameters", + Label: "Embedding Pooling", + Description: "How per-token embedding vectors are reduced to one embedding: the backend pools itself (default), or LocalAI pools Go-side (mean, last, decayed_mean) from raw per-token vectors", + Component: "select", + Options: []FieldOption{ + {Value: "", Label: "Backend (default)"}, + {Value: "backend", Label: "Backend"}, + {Value: "mean", Label: "Mean"}, + {Value: "last", Label: "Last token"}, + {Value: "decayed_mean", Label: "Decayed mean"}, + }, + Advanced: true, + Order: 36, + }, + "parameters.pooling_half_life_tokens": { + Section: "parameters", + Label: "Pooling Half-Life (tokens)", + Description: "Half-life in tokens for decayed_mean pooling: a token's weight halves every this many positions counting back from the end of the conversation (default 256)", + Component: "number", + Min: f64(0), + Advanced: true, + Order: 37, + }, // --- Templates --- "template.chat": { diff --git a/core/config/model_config.go b/core/config/model_config.go index 4e714af47178..45b90c002a68 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -18,6 +18,7 @@ import ( "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/reasoning" "github.com/mudler/cogito" + "github.com/mudler/xlog" "gopkg.in/yaml.v3" ) @@ -1387,6 +1388,15 @@ func (cfg *ModelConfig) SetDefaults(opts ...ConfigLoaderOption) { } runBackendHooks(cfg, lo.modelPath) + // Go-side pooling schemes need the backend to hand back raw per-token + // vectors: unless the operator already pinned a pooling backend option + // explicitly, add "pooling:none" so llama.cpp skips server-side pooling. + if IsGoSidePooling(cfg.Pooling) && !hasBackendOption(cfg.Options, "pooling") { + cfg.Options = append(cfg.Options, "pooling:none") + xlog.Debug("Go-side pooling requested: auto-appending backend option", + "model", cfg.Name, "pooling", cfg.Pooling, "option", "pooling:none") + } + // Apply hardware-driven defaults (e.g. a larger physical batch on Blackwell) // LAST, after the context size is fully resolved (explicit config, LoadOptions, // then the GGUF guess inside runBackendHooks): the Blackwell batch guard sizes @@ -1550,6 +1560,24 @@ func (c *ModelConfig) Validate() (bool, error) { } } + // Go-side embedding pooling: reject unknown schemes and half-life + // misconfigurations at load time, and refuse contradictory setups where + // a Go-side scheme (which needs raw per-token vectors) is combined with + // a backend pooling option other than "none" — the backend would return + // a single pooled vector and the Go pooling would silently degenerate. + if err := ValidatePooling(c.Pooling, c.PoolingHalfLifeTokens); err != nil { + return false, fmt.Errorf("parameters: %w", err) + } + if IsGoSidePooling(c.Pooling) { + for _, o := range c.Options { + if name, val, ok := strings.Cut(o, ":"); ok && name == "pooling" && val != "none" { + return false, fmt.Errorf( + "parameters.pooling %q needs raw per-token vectors but backend option %q pools server-side: drop the option or set \"pooling:none\"", + c.Pooling, o) + } + } + } + return true, nil } @@ -1562,6 +1590,66 @@ const ( ScoreNormalizationMean = "mean" ) +// Embedding pooling schemes (parameters.pooling). Defined here so +// ModelConfig.Validate can reject unknown values without depending on +// core/backend (which already depends on config); core/backend/pooling.go +// aliases these for the pooling implementation. +const ( + // PoolingBackend (or the empty string) leaves pooling to the inference + // backend — the pre-existing behavior of the embeddings path. + PoolingBackend = "backend" + // PoolingMean averages the backend's raw per-token vectors Go-side. + PoolingMean = "mean" + // PoolingLast selects the last token's vector Go-side. + PoolingLast = "last" + // PoolingDecayedMean is a Go-side mean weighted toward the most recent + // tokens: w_i = 2^(-(T-1-i)/H) with half-life H tokens + // (parameters.pooling_half_life_tokens, default 256). + PoolingDecayedMean = "decayed_mean" +) + +// IsGoSidePooling reports whether scheme pools in LocalAI's Go layer, +// which requires raw per-token vectors from the backend (the "pooling:none" +// backend option) rather than a backend-pooled single vector. +func IsGoSidePooling(scheme string) bool { + switch scheme { + case PoolingMean, PoolingLast, PoolingDecayedMean: + return true + } + return false +} + +// ValidatePooling checks a pooling scheme + half-life pair. Shared by +// ModelConfig.Validate (model YAML) and the embeddings endpoint +// (per-request override) so both paths reject exactly the same shapes. +func ValidatePooling(scheme string, halfLifeTokens int) error { + switch scheme { + case "", PoolingBackend, PoolingMean, PoolingLast, PoolingDecayedMean: + default: + return fmt.Errorf("unknown pooling %q (expected %q, %q, %q or %q)", + scheme, PoolingBackend, PoolingMean, PoolingLast, PoolingDecayedMean) + } + if halfLifeTokens < 0 { + return fmt.Errorf("pooling_half_life_tokens must be non-negative, got %d", halfLifeTokens) + } + if halfLifeTokens > 0 && scheme != PoolingDecayedMean { + return fmt.Errorf("pooling_half_life_tokens only applies to pooling %q, not %q", + PoolingDecayedMean, scheme) + } + return nil +} + +// hasBackendOption reports whether options carries a "name:value" entry +// for the given option name. +func hasBackendOption(options []string, name string) bool { + for _, o := range options { + if n, _, ok := strings.Cut(o, ":"); ok && n == name { + return true + } + } + return false +} + func (c *ModelConfig) HasTemplate() bool { return c.TemplateConfig.Completion != "" || c.TemplateConfig.Edit != "" || c.TemplateConfig.Chat != "" || c.TemplateConfig.ChatMessage != "" || c.TemplateConfig.UseTokenizerTemplate } diff --git a/core/config/pooling_config_test.go b/core/config/pooling_config_test.go new file mode 100644 index 000000000000..0651d876e633 --- /dev/null +++ b/core/config/pooling_config_test.go @@ -0,0 +1,111 @@ +package config + +import ( + "strings" + + "github.com/mudler/LocalAI/core/schema" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func poolingTestConfig(pooling string, halfLife int, options ...string) *ModelConfig { + return &ModelConfig{ + Name: "pool-test", + PredictionOptions: schema.PredictionOptions{ + BasicModelRequest: schema.BasicModelRequest{Model: "foo.gguf"}, + Pooling: pooling, + PoolingHalfLifeTokens: halfLife, + }, + Options: options, + } +} + +var _ = Describe("Go-side pooling model config", func() { + Describe("SetDefaults", func() { + It("auto-appends pooling:none for a Go-side scheme with no explicit pooling option", func() { + cfg := poolingTestConfig(PoolingDecayedMean, 0) + cfg.SetDefaults() + Expect(cfg.Options).To(ContainElement("pooling:none")) + }) + + It("leaves an explicit pooling option alone", func() { + cfg := poolingTestConfig(PoolingMean, 0, "pooling:none") + cfg.SetDefaults() + // Other defaults (e.g. hardware-driven options) may append + // unrelated entries; the pooling option must stay single. + poolingOptions := []string{} + for _, o := range cfg.Options { + if strings.HasPrefix(o, "pooling:") { + poolingOptions = append(poolingOptions, o) + } + } + Expect(poolingOptions).To(Equal([]string{"pooling:none"})) + }) + + It("does not touch options when pooling is delegated to the backend", func() { + for _, scheme := range []string{"", PoolingBackend} { + cfg := poolingTestConfig(scheme, 0) + cfg.SetDefaults() + Expect(cfg.Options).ToNot(ContainElement("pooling:none")) + } + }) + }) + + Describe("Validate", func() { + It("accepts every known scheme", func() { + for _, scheme := range []string{"", PoolingBackend, PoolingMean, PoolingLast, PoolingDecayedMean} { + cfg := poolingTestConfig(scheme, 0) + valid, err := cfg.Validate() + Expect(err).ToNot(HaveOccurred(), "scheme %q", scheme) + Expect(valid).To(BeTrue(), "scheme %q", scheme) + } + }) + + It("rejects unknown schemes", func() { + cfg := poolingTestConfig("sideways", 0) + valid, err := cfg.Validate() + Expect(err).To(HaveOccurred()) + Expect(valid).To(BeFalse()) + Expect(err.Error()).To(ContainSubstring("pooling")) + }) + + It("rejects a negative half-life", func() { + cfg := poolingTestConfig(PoolingDecayedMean, -1) + valid, err := cfg.Validate() + Expect(err).To(HaveOccurred()) + Expect(valid).To(BeFalse()) + Expect(err.Error()).To(ContainSubstring("pooling_half_life_tokens")) + }) + + It("rejects a half-life on non-decayed schemes", func() { + cfg := poolingTestConfig(PoolingMean, 256) + valid, err := cfg.Validate() + Expect(err).To(HaveOccurred()) + Expect(valid).To(BeFalse()) + Expect(err.Error()).To(ContainSubstring("decayed_mean")) + }) + + It("accepts a half-life on decayed_mean", func() { + cfg := poolingTestConfig(PoolingDecayedMean, 256) + valid, err := cfg.Validate() + Expect(err).ToNot(HaveOccurred()) + Expect(valid).To(BeTrue()) + }) + + It("rejects Go-side pooling combined with a non-none backend pooling option", func() { + cfg := poolingTestConfig(PoolingMean, 0, "pooling:mean") + valid, err := cfg.Validate() + Expect(err).To(HaveOccurred()) + Expect(valid).To(BeFalse()) + Expect(err.Error()).To(ContainSubstring("pooling:none")) + }) + + It("accepts Go-side pooling with an explicit pooling:none option", func() { + cfg := poolingTestConfig(PoolingLast, 0, "pooling:none") + valid, err := cfg.Validate() + Expect(err).ToNot(HaveOccurred()) + Expect(valid).To(BeTrue()) + }) + }) +}) diff --git a/core/schema/prediction.go b/core/schema/prediction.go index d6e4fabf3a04..974850d0cd79 100644 --- a/core/schema/prediction.go +++ b/core/schema/prediction.go @@ -142,4 +142,17 @@ type PredictionOptions struct { // Embedding encoding format: "float" (default) or "base64" (OpenAI Node.js SDK default) EncodingFormat string `json:"encoding_format,omitempty" yaml:"encoding_format,omitempty"` + + // Pooling is a LocalAI extension for /v1/embeddings: how the backend's + // per-token vectors are reduced to a single embedding. "" or "backend" + // leaves pooling to the inference backend (the pre-existing behavior); + // "mean", "last" and "decayed_mean" pool Go-side from raw per-token + // vectors (the backend must run with the "pooling:none" option, which + // model configs get automatically when this is set). + Pooling string `json:"pooling,omitempty" yaml:"pooling,omitempty"` + // PoolingHalfLifeTokens is a LocalAI extension for /v1/embeddings: the + // half-life (in tokens) of the "decayed_mean" pooling scheme — a token's + // weight halves every this-many positions counting back from the end of + // the conversation. Defaults to 256 when unset. + PoolingHalfLifeTokens int `json:"pooling_half_life_tokens,omitempty" yaml:"pooling_half_life_tokens,omitempty"` } From f9869b9f2a097392e46bf8ce88851b4c3838e711 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Wed, 15 Jul 2026 11:54:59 +0100 Subject: [PATCH 08/10] feat(embeddings): accept chat messages[] and per-request pooling on /v1/embeddings Co-Authored-By: Claude Fable 5 --- core/http/embeddings_pooling_test.go | 209 +++++++++++++++++++++ core/http/endpoints/openai/embeddings.go | 47 ++++- core/http/middleware/request.go | 28 ++- core/http/routes/openai.go | 2 +- core/templates/evaluator.go | 27 +++ core/templates/evaluator_embedding_test.go | 88 +++++++++ docs/content/features/embeddings.md | 59 ++++++ tests/e2e/mock-backend/main.go | 15 +- tests/e2e/mock_backend_test.go | 30 +++ 9 files changed, 486 insertions(+), 19 deletions(-) create mode 100644 core/http/embeddings_pooling_test.go create mode 100644 core/templates/evaluator_embedding_test.go diff --git a/core/http/embeddings_pooling_test.go b/core/http/embeddings_pooling_test.go new file mode 100644 index 000000000000..39785f94a342 --- /dev/null +++ b/core/http/embeddings_pooling_test.go @@ -0,0 +1,209 @@ +package http_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "os" + "path/filepath" + "time" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/application" + "github.com/mudler/LocalAI/core/config" + . "github.com/mudler/LocalAI/core/http" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/xlog" +) + +// embedTestModel is served by the mock backend. embd_normalize:-1 disables +// Go-side normalization so the pooled goldens below stay hand-computable. +const embedTestModel = "embed-pooling-model" + +var _ = Describe("Embeddings chat messages[] and Go-side pooling", func() { + var app *echo.Echo + var localApp *application.Application + var localModelDir string + var c context.Context + var cancel context.CancelFunc + + const baseURL = "http://127.0.0.1:9092" + + postEmbeddings := func(body string) (int, map[string]any) { + resp, err := http.Post(baseURL+"/v1/embeddings", "application/json", bytes.NewBufferString(body)) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + payload, err := io.ReadAll(resp.Body) + Expect(err).ToNot(HaveOccurred()) + decoded := map[string]any{} + // Some error shapes are plain text; tolerate that and return nil map. + _ = json.Unmarshal(payload, &decoded) + if decoded == nil { + decoded = map[string]any{} + } + decoded["__raw"] = string(payload) + return resp.StatusCode, decoded + } + + embeddingOf := func(response map[string]any) []float64 { + data, ok := response["data"].([]any) + Expect(ok).To(BeTrue(), "response has no data array: %v", response["__raw"]) + Expect(data).To(HaveLen(1)) + item := data[0].(map[string]any) + raw := item["embedding"].([]any) + out := make([]float64, len(raw)) + for i, v := range raw { + out[i] = v.(float64) + } + return out + } + + BeforeEach(func() { + if mockBackendPath == "" { + Skip("mock-backend binary not built; run 'make build-mock-backend'") + } + + var err error + c, cancel = context.WithCancel(context.Background()) + + localModelDir, err = os.MkdirTemp("", "embeddings-pooling-models-") + Expect(err).ToNot(HaveOccurred()) + + mockModelYAML := "name: " + embedTestModel + "\n" + + "backend: mock-backend\n" + + "embeddings: true\n" + + "options:\n" + + "- embd_normalize:-1\n" + + "parameters:\n" + + " model: mock-model.bin\n" + Expect(os.WriteFile(filepath.Join(localModelDir, embedTestModel+".yaml"), []byte(mockModelYAML), 0644)).To(Succeed()) + + systemState, err := system.GetSystemState( + system.WithBackendPath(backendDir), + system.WithModelPath(localModelDir), + ) + Expect(err).ToNot(HaveOccurred()) + + localApp, err = application.New( + config.WithDebug(true), + config.WithContext(c), + config.WithSystemState(systemState), + ) + Expect(err).ToNot(HaveOccurred()) + localApp.ModelLoader().SetExternalBackend("mock-backend", mockBackendPath) + + app, err = API(localApp) + Expect(err).ToNot(HaveOccurred()) + + go func() { + if err := app.Start("127.0.0.1:9092"); err != nil && err != http.ErrServerClosed { + xlog.Error("server error", "error", err) + } + }() + + Eventually(func() error { + resp, err := http.Get(baseURL + "/healthz") + if err != nil { + return err + } + _ = resp.Body.Close() + return nil + }, "2m").ShouldNot(HaveOccurred()) + }) + + AfterEach(func() { + if localApp != nil { + _ = localApp.Shutdown() + localApp = nil + } + cancel() + if app != nil { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + Expect(app.Shutdown(ctx)).To(Succeed()) + app = nil + } + if localModelDir != "" { + _ = os.RemoveAll(localModelDir) + } + }) + + It("embeds a messages[] conversation with decayed_mean pooling to the exact golden", func() { + // The rendered fallback conversation is "user: per-token: alpha beta gamma", + // so the mock returns 3 per-token vectors of dim 8 with + // vec[i][j] = (i+1)/(j+2). decayed_mean with half-life 1 weighs them + // [0.25 0.5 1]/1.75; embd_normalize:-1 skips normalization, so the + // exact pooled value is hand-computable. + status, response := postEmbeddings(`{ + "model": "` + embedTestModel + `", + "messages": [{"role": "user", "content": "per-token: alpha beta gamma"}], + "pooling": "decayed_mean", + "pooling_half_life_tokens": 1 + }`) + Expect(status).To(Equal(200), "body: %v", response["__raw"]) + got := embeddingOf(response) + Expect(got).To(HaveLen(8)) + for j := 0; j < 8; j++ { + want := (0.25*1 + 0.5*2 + 1.0*3) / 1.75 / float64(j+2) + Expect(got[j]).To(BeNumerically("~", want, 1e-6), "component %d", j) + } + }) + + It("passes plain input through unchanged when no pooling is requested", func() { + status, response := postEmbeddings(`{"model": "` + embedTestModel + `", "input": "hello"}`) + Expect(status).To(Equal(200), "body: %v", response["__raw"]) + got := embeddingOf(response) + // The mock's default vector: index%100/100, 768 wide, passed through + // untouched (no Go-side pooling, no normalization). + Expect(got).To(HaveLen(768)) + Expect(got[0]).To(BeNumerically("~", 0.0, 1e-6)) + Expect(got[1]).To(BeNumerically("~", 0.01, 1e-6)) + Expect(got[99]).To(BeNumerically("~", 0.99, 1e-6)) + Expect(got[100]).To(BeNumerically("~", 0.0, 1e-6)) + }) + + It("embeds messages[] without pooling via the backend's own vector", func() { + status, response := postEmbeddings(`{ + "model": "` + embedTestModel + `", + "messages": [{"role": "user", "content": "hello"}] + }`) + Expect(status).To(Equal(200), "body: %v", response["__raw"]) + Expect(embeddingOf(response)).To(HaveLen(768)) + }) + + It("fails closed when pooling is requested but the backend reports no shape", func() { + // "no-shape:" makes the mock omit tokens/dim, simulating a backend + // built before EmbeddingResult carried shape fields. + status, response := postEmbeddings(`{ + "model": "` + embedTestModel + `", + "input": "no-shape: hello", + "pooling": "mean" + }`) + Expect(status).ToNot(Equal(200)) + Expect(response["__raw"]).To(ContainSubstring("no shape")) + }) + + It("rejects input combined with messages", func() { + status, response := postEmbeddings(`{ + "model": "` + embedTestModel + `", + "input": "hello", + "messages": [{"role": "user", "content": "hello"}] + }`) + Expect(status).To(Equal(400), "body: %v", response["__raw"]) + }) + + It("rejects an unknown pooling scheme", func() { + status, response := postEmbeddings(`{ + "model": "` + embedTestModel + `", + "input": "hello", + "pooling": "sideways" + }`) + Expect(status).To(Equal(400), "body: %v", response["__raw"]) + Expect(response["__raw"]).To(ContainSubstring("pooling")) + }) +}) diff --git a/core/http/endpoints/openai/embeddings.go b/core/http/endpoints/openai/embeddings.go index e6a3a353dbb5..b112153ef0b6 100644 --- a/core/http/endpoints/openai/embeddings.go +++ b/core/http/endpoints/openai/embeddings.go @@ -5,12 +5,14 @@ import ( "encoding/binary" "encoding/json" "math" + "net/http" "time" "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/templates" "github.com/mudler/LocalAI/pkg/model" "github.com/google/uuid" @@ -41,29 +43,58 @@ func embeddingItem(embeddings []float32, index int, encodingFormat string) schem } // EmbeddingsEndpoint is the OpenAI Embeddings API endpoint https://platform.openai.com/docs/api-reference/embeddings +// LocalAI extensions: a chat conversation can be embedded by sending +// `messages` (mutually exclusive with `input`; one conversation per request, +// one data item in the response), and `pooling`/`pooling_half_life_tokens` +// select a Go-side pooling scheme over the backend's per-token vectors. // @Summary Get a vector representation of a given input that can be easily consumed by machine learning models and algorithms. // @Tags embeddings // @Param request body schema.OpenAIRequest true "query params" // @Success 200 {object} schema.OpenAIResponse "Response" // @Router /v1/embeddings [post] -func EmbeddingsEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { +func EmbeddingsEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig) echo.HandlerFunc { return func(c echo.Context) error { input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest) if !ok || input.Model == "" { return echo.ErrBadRequest } - config, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) - if !ok || config == nil { + modelConfig, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + if !ok || modelConfig == nil { return echo.ErrBadRequest } - xlog.Debug("Parameter Config", "config", config) + // The middleware merged any per-request pooling override onto the + // per-request config copy; reject bad values before touching the + // model so the client gets a 400, not a load-time failure. + if err := config.ValidatePooling(modelConfig.Pooling, modelConfig.PoolingHalfLifeTokens); err != nil { + return echo.NewHTTPError(http.StatusBadRequest, err.Error()) + } + + if len(input.Messages) > 0 { + // One conversation per request: messages[] renders to a single + // input string, so it cannot be combined with input. + if len(modelConfig.InputStrings) > 0 || len(modelConfig.InputToken) > 0 { + return echo.NewHTTPError(http.StatusBadRequest, "input and messages are mutually exclusive: send the conversation via messages, or plain text/tokens via input") + } + // Non-text parts were parked in StringImages/StringVideos/ + // StringAudios by the request middleware; only text embeds. + for _, m := range input.Messages { + if len(m.StringImages) > 0 || len(m.StringVideos) > 0 || len(m.StringAudios) > 0 { + xlog.Debug("embeddings: ignoring non-text content parts in messages", "model", modelConfig.Name) + break + } + } + rendered := evaluator.RenderConversationForEmbedding(*input, input.Messages, modelConfig) + modelConfig.InputStrings = append(modelConfig.InputStrings, rendered) + } + + xlog.Debug("Parameter Config", "config", modelConfig) items := []schema.Item{} - for i, s := range config.InputToken { + for i, s := range modelConfig.InputToken { // get the model function to call for the result - embedFn, err := backend.ModelEmbedding(input.Context, "", s, ml, *config, appConfig) + embedFn, err := backend.ModelEmbedding(input.Context, "", s, ml, *modelConfig, appConfig) if err != nil { return err } @@ -75,9 +106,9 @@ func EmbeddingsEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, app items = append(items, embeddingItem(embeddings, i, input.EncodingFormat)) } - for i, s := range config.InputStrings { + for i, s := range modelConfig.InputStrings { // get the model function to call for the result - embedFn, err := backend.ModelEmbedding(input.Context, s, []int{}, ml, *config, appConfig) + embedFn, err := backend.ModelEmbedding(input.Context, s, []int{}, ml, *modelConfig, appConfig) if err != nil { return err } diff --git a/core/http/middleware/request.go b/core/http/middleware/request.go index 74f7e8565fb3..4840c8bbab69 100644 --- a/core/http/middleware/request.go +++ b/core/http/middleware/request.go @@ -591,6 +591,16 @@ func mergeOpenAIRequestAndModelConfig(config *config.ModelConfig, input *schema. } } + // Per-request Go-side pooling override for /v1/embeddings (LocalAI + // extension) — copy-if-set onto the per-request config like the Input + // handling above; the embeddings endpoint validates the merged values. + if input.Pooling != "" { + config.Pooling = input.Pooling + } + if input.PoolingHalfLifeTokens != 0 { + config.PoolingHalfLifeTokens = input.PoolingHalfLifeTokens + } + // Can be either a string or an object switch fnc := input.FunctionCall.(type) { case string: @@ -628,10 +638,13 @@ func mergeOpenAIRequestAndModelConfig(config *config.ModelConfig, input *schema. } } - if valid, _ := config.Validate(); valid { - return nil + if valid, err := config.Validate(); !valid { + // The base config validated at load time, so a post-merge failure + // can only come from request-supplied fields: surface it as a 400 + // with the underlying cause rather than an opaque 500. + return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("unable to validate configuration after merging: %v", err)) } - return fmt.Errorf("unable to validate configuration after merging") + return nil } func (re *RequestExtractor) SetOpenResponsesRequest(c echo.Context) error { @@ -753,8 +766,11 @@ func MergeOpenResponsesConfig(config *config.ModelConfig, input *schema.OpenResp } } - if valid, _ := config.Validate(); valid { - return nil + if valid, err := config.Validate(); !valid { + // The base config validated at load time, so a post-merge failure + // can only come from request-supplied fields: surface it as a 400 + // with the underlying cause rather than an opaque 500. + return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("unable to validate configuration after merging: %v", err)) } - return fmt.Errorf("unable to validate configuration after merging") + return nil } diff --git a/core/http/routes/openai.go b/core/http/routes/openai.go index 0820f4e31a49..0d602f5d3f85 100644 --- a/core/http/routes/openai.go +++ b/core/http/routes/openai.go @@ -133,7 +133,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, app.POST("/v1/engines/:model/completions", completionHandler, completionMiddleware...) // embeddings - embeddingHandler := openai.EmbeddingsEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()) + embeddingHandler := openai.EmbeddingsEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig()) embeddingMiddleware := []echo.MiddlewareFunc{ nodeHeaderMiddleware, usageMiddleware, diff --git a/core/templates/evaluator.go b/core/templates/evaluator.go index 8ee74e43d3e1..36e491b7de24 100644 --- a/core/templates/evaluator.go +++ b/core/templates/evaluator.go @@ -232,3 +232,30 @@ func (e *Evaluator) TemplateMessages(input schema.OpenAIRequest, messages []sche return predInput } + +// RenderConversationForEmbedding renders a chat conversation to the single +// string embedded for /v1/embeddings messages[] requests. When the model +// config carries full chat templates (both chat and chat_message), the +// conversation renders exactly like a chat prompt via TemplateMessages, so +// the embedding space matches what a chat model would actually see. +// +// Otherwise it falls back to a FROZEN role-prefixed rendering: one line per +// message, ": ", joined by "\n", skipping messages with +// empty text content (non-text parts — images/audio/video — are not +// embedded). This fallback DEFINES the embedding space for models without +// chat templates: any change to it — even whitespace — silently invalidates +// every vector clients have already stored, so it must never change. It is +// pinned by a golden test in evaluator_embedding_test.go. +func (e *Evaluator) RenderConversationForEmbedding(input schema.OpenAIRequest, messages []schema.Message, cfg *config.ModelConfig) string { + if cfg.TemplateConfig.Chat != "" && cfg.TemplateConfig.ChatMessage != "" { + return e.TemplateMessages(input, messages, cfg, nil, false) + } + lines := make([]string, 0, len(messages)) + for _, m := range messages { + if m.StringContent == "" { + continue + } + lines = append(lines, m.Role+": "+m.StringContent) + } + return strings.Join(lines, "\n") +} diff --git a/core/templates/evaluator_embedding_test.go b/core/templates/evaluator_embedding_test.go new file mode 100644 index 000000000000..186640be2153 --- /dev/null +++ b/core/templates/evaluator_embedding_test.go @@ -0,0 +1,88 @@ +package templates_test + +import ( + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/schema" + . "github.com/mudler/LocalAI/core/templates" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("RenderConversationForEmbedding", func() { + var evaluator *Evaluator + + BeforeEach(func() { + evaluator = NewEvaluator("") + }) + + Context("model without chat templates (frozen fallback)", func() { + It("renders the exact role-prefixed golden", func() { + messages := []schema.Message{ + {Role: "system", StringContent: "You are terse."}, + {Role: "user", StringContent: "What is the capital of France?"}, + {Role: "tool", StringContent: "Paris"}, + } + got := evaluator.RenderConversationForEmbedding(schema.OpenAIRequest{}, messages, &config.ModelConfig{}) + // GOLDEN: the fallback rendering defines the embedding space for + // template-less models — stored vectors embed this exact string. + // If this assertion ever fails, the change invalidates every + // previously stored conversation vector. Do not update the + // expectation; fix the code. + Expect(got).To(Equal("system: You are terse.\nuser: What is the capital of France?\ntool: Paris")) + }) + + It("skips messages with empty text content", func() { + messages := []schema.Message{ + {Role: "user", StringContent: "first"}, + {Role: "assistant", StringContent: ""}, + {Role: "user", StringContent: "second"}, + } + got := evaluator.RenderConversationForEmbedding(schema.OpenAIRequest{}, messages, &config.ModelConfig{}) + Expect(got).To(Equal("user: first\nuser: second")) + }) + + It("embeds the text parts of a content-part message and ignores parked media", func() { + // The request middleware decodes multi-part content into + // StringContent (text) and StringImages/... (media); only the + // text reaches the embedding. + messages := []schema.Message{ + { + Role: "user", + StringContent: "describe the image", + StringImages: []string{"aGVsbG8="}, + }, + } + got := evaluator.RenderConversationForEmbedding(schema.OpenAIRequest{}, messages, &config.ModelConfig{}) + Expect(got).To(Equal("user: describe the image")) + }) + }) + + Context("model with chat templates", func() { + It("renders through TemplateMessages so embeddings match the chat prompt", func() { + cfg := &config.ModelConfig{ + TemplateConfig: config.TemplateConfig{ + Chat: "{{.Input}} ", + ChatMessage: "[{{.RoleName}}] {{.Content}}", + }, + } + messages := []schema.Message{ + {Role: "system", StringContent: "sys"}, + {Role: "user", StringContent: "hi"}, + } + got := evaluator.RenderConversationForEmbedding(schema.OpenAIRequest{}, messages, cfg) + Expect(got).To(Equal("[system] sys\n[user] hi ")) + }) + + It("keeps the fallback when only one of the two templates is set", func() { + cfg := &config.ModelConfig{ + TemplateConfig: config.TemplateConfig{ + ChatMessage: "[{{.RoleName}}] {{.Content}}", + }, + } + messages := []schema.Message{{Role: "user", StringContent: "hi"}} + got := evaluator.RenderConversationForEmbedding(schema.OpenAIRequest{}, messages, cfg) + Expect(got).To(Equal("user: hi")) + }) + }) +}) diff --git a/docs/content/features/embeddings.md b/docs/content/features/embeddings.md index 707df2b83193..8bd83980f89e 100644 --- a/docs/content/features/embeddings.md +++ b/docs/content/features/embeddings.md @@ -96,6 +96,65 @@ curl http://localhost:8080/embeddings -X POST -H "Content-Type: application/json }' | jq "." ``` +## Embedding chat conversations and Go-side pooling + +`/v1/embeddings` also accepts a chat conversation via `messages` (a LocalAI +extension), plus a per-request `pooling` scheme that LocalAI applies itself to +the backend's raw per-token vectors: + +```bash +curl http://localhost:8080/v1/embeddings -X POST -H "Content-Type: application/json" -d '{ + "model": "my-awesome-model", + "messages": [ + {"role": "system", "content": "You are a support agent."}, + {"role": "user", "content": "My invoice is wrong."} + ], + "pooling": "decayed_mean", + "pooling_half_life_tokens": 256 +}' +``` + +- One conversation per request; the response is the standard OpenAI embeddings + shape with a single `data[0].embedding` item. +- `input` and `messages` are mutually exclusive (400 otherwise); an unknown + `pooling` value is also a 400. +- If the model config carries both `template.chat` and `template.chat_message`, + the conversation renders exactly like a chat prompt, so the embedding matches + what a chat model would actually see. Otherwise a frozen role-prefixed + fallback is used (`: ` lines joined by newlines, empty-content + messages skipped). Non-text content parts (images, audio, video) are ignored. + +`pooling` selects how the per-token vectors are reduced to one embedding: + +| Value | Meaning | +|-------|---------| +| _(empty)_ / `backend` | The backend pools by itself — the default, today's exact behavior. | +| `mean` | Average of all token vectors. | +| `last` | The last token's vector. | +| `decayed_mean` | Recency-weighted mean: token *i* of *T* weighs `2^(-(T-1-i)/H)` with half-life `H` = `pooling_half_life_tokens` (default 256) — recent turns dominate without erasing earlier context. | + +Go-side schemes need raw per-token vectors from the backend, so LocalAI +automatically adds the `pooling:none` backend option when `parameters.pooling` +is set in the model YAML. After pooling, the vector is normalized with +llama.cpp's `embd_normalize` rule (default L2; configurable through +`options: ["embd_normalize:"]`). + +Model-level defaults live under `parameters:`: + +```yaml +name: conversation-embedder +backend: llama-cpp +embeddings: true +parameters: + model: ggml-file.bin + pooling: decayed_mean + pooling_half_life_tokens: 256 +``` + +Go-side pooling requires an up-to-date llama-cpp backend: older builds don't +report the embedding shape, and the request fails closed with an error asking +you to rebuild the backend and set `options: ["pooling:none"]`. + ## 💡 Examples - Example that uses LLamaIndex and LocalAI as embedding: [here](https://github.com/mudler/LocalAI-examples/tree/main/query_data). diff --git a/tests/e2e/mock-backend/main.go b/tests/e2e/mock-backend/main.go index c9ab48927eec..c0956874f668 100644 --- a/tests/e2e/mock-backend/main.go +++ b/tests/e2e/mock-backend/main.go @@ -391,15 +391,22 @@ func mockToolNameFromRequest(in *pb.PredictOptions) string { } func (m *MockBackend) Embedding(ctx context.Context, in *pb.PredictOptions) (*pb.EmbeddingResult, error) { - xlog.Debug("Embedding called", "prompt", in.Prompt) + // The embeddings path ships the text in PredictOptions.Embeddings + // (see core/backend/embeddings.go), not Prompt; check both so the + // markers below work however the caller packed the request. + text := in.Embeddings + if text == "" { + text = in.Prompt + } + xlog.Debug("Embedding called", "text", text) // Deterministic per-token mode for Go-side pooling tests: a prompt // carrying the "per-token:" marker yields len(fields) vectors of dim 8 // with vec[i][j] = (i+1)/(j+2), so endpoint tests can assert exact // pooled goldens. The marker may sit mid-prompt (the embeddings // messages[] path renders conversations as ": " lines), // so match anywhere and tokenize what follows the first occurrence. - if idx := strings.Index(in.Prompt, "per-token:"); idx >= 0 { - fields := strings.Fields(in.Prompt[idx+len("per-token:"):]) + if idx := strings.Index(text, "per-token:"); idx >= 0 { + fields := strings.Fields(text[idx+len("per-token:"):]) tokens := len(fields) const dim = 8 flat := make([]float32, 0, tokens*dim) @@ -419,7 +426,7 @@ func (m *MockBackend) Embedding(ctx context.Context, in *pb.PredictOptions) (*pb // before EmbeddingResult carried shape fields (tokens/dim left 0), so // tests can assert the fail-closed error when Go-side pooling is // requested against such a backend. - legacyShape := strings.Contains(in.Prompt, "no-shape:") + legacyShape := strings.Contains(text, "no-shape:") // Return a mock embedding vector of 768 dimensions embeddings := make([]float32, 768) for i := range embeddings { diff --git a/tests/e2e/mock_backend_test.go b/tests/e2e/mock_backend_test.go index 24f4cbd94cf5..9f8e4a8511bf 100644 --- a/tests/e2e/mock_backend_test.go +++ b/tests/e2e/mock_backend_test.go @@ -132,6 +132,36 @@ var _ = Describe("Mock Backend E2E Tests", Label("MockBackend"), func() { Expect(len(resp.Data)).To(Equal(1)) Expect(len(resp.Data[0].Embedding)).To(Equal(768)) }) + + // LocalAI extension: a chat conversation can be embedded by sending + // messages[] instead of input — raw http.Post because the OpenAI SDK + // has no such parameter. + It("should embed a chat conversation sent via messages[]", func() { + body := `{"model":"mock-model","messages":[{"role":"system","content":"be brief"},{"role":"user","content":"hello"}]}` + resp, err := http.Post(apiURL+"/embeddings", "application/json", strings.NewReader(body)) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(200)) + var decoded struct { + Data []struct { + Embedding []float64 `json:"embedding"` + } `json:"data"` + } + payload, err := io.ReadAll(resp.Body) + Expect(err).ToNot(HaveOccurred()) + Expect(json.Unmarshal(payload, &decoded)).To(Succeed()) + // One conversation per request -> exactly one data item. + Expect(decoded.Data).To(HaveLen(1)) + Expect(decoded.Data[0].Embedding).To(HaveLen(768)) + }) + + It("should reject input combined with messages", func() { + body := `{"model":"mock-model","input":"x","messages":[{"role":"user","content":"hello"}]}` + resp, err := http.Post(apiURL+"/embeddings", "application/json", strings.NewReader(body)) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(400)) + }) }) Describe("TTS APIs", func() { From 48dbe64ca9c7cfccc4667d60cf6d275034f1e84f Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Wed, 15 Jul 2026 12:42:44 +0100 Subject: [PATCH 09/10] chore(middleware): name the failing fields when post-merge validation 400s An intermittent post-merge validation failure surfaced as an opaque 400 during integration (pooling scheme mismatch that no client had sent). Log the model, the request's pooling override, and the merged config's pooling fields at the failure point so the next occurrence identifies whether the request or the stored config carried the bad value. Co-Authored-By: Claude Fable 5 --- core/http/middleware/request.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/core/http/middleware/request.go b/core/http/middleware/request.go index 4840c8bbab69..1c44bc2c17bf 100644 --- a/core/http/middleware/request.go +++ b/core/http/middleware/request.go @@ -642,6 +642,10 @@ func mergeOpenAIRequestAndModelConfig(config *config.ModelConfig, input *schema. // The base config validated at load time, so a post-merge failure // can only come from request-supplied fields: surface it as a 400 // with the underlying cause rather than an opaque 500. + xlog.Warn("post-merge validation failed", + "model", config.Name, "err", err, + "inputPooling", input.Pooling, "cfgPooling", config.Pooling, + "cfgHalfLife", config.PoolingHalfLifeTokens, "options", config.Options) return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("unable to validate configuration after merging: %v", err)) } return nil From a074dcca2e27784ed4feecb83ea492d8efa4e171 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Wed, 15 Jul 2026 13:17:13 +0100 Subject: [PATCH 10/10] fix(embeddings): scheme override must not inherit the config's half-life A model config defaulting to decayed_mean pooling carries pooling_half_life_tokens; a request overriding the scheme to mean/last without its own half-life inherited that value, and post-merge validation rejected the pair the server itself had assembled. Zero the inherited half-life when the overridden scheme is not decayed_mean; a request that explicitly pairs a half-life with a non-decayed scheme still 400s. Co-Authored-By: Claude Fable 5 --- core/http/middleware/request.go | 9 +++ core/http/middleware/request_test.go | 86 ++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/core/http/middleware/request.go b/core/http/middleware/request.go index 1c44bc2c17bf..1599ef05c3ae 100644 --- a/core/http/middleware/request.go +++ b/core/http/middleware/request.go @@ -596,6 +596,15 @@ func mergeOpenAIRequestAndModelConfig(config *config.ModelConfig, input *schema. // handling above; the embeddings endpoint validates the merged values. if input.Pooling != "" { config.Pooling = input.Pooling + // The config's default half-life belongs to its own decayed_mean + // default. A request that overrides the scheme without sending its + // own half-life must not inherit it, or the merged pair (e.g. + // "mean" + 256) fails validation through no fault of the client. + // (literal because the `config` parameter shadows the package; + // this is config.PoolingDecayedMean) + if input.PoolingHalfLifeTokens == 0 && input.Pooling != "decayed_mean" { + config.PoolingHalfLifeTokens = 0 + } } if input.PoolingHalfLifeTokens != 0 { config.PoolingHalfLifeTokens = input.PoolingHalfLifeTokens diff --git a/core/http/middleware/request_test.go b/core/http/middleware/request_test.go index 010379714f6d..1b00c7f022e7 100644 --- a/core/http/middleware/request_test.go +++ b/core/http/middleware/request_test.go @@ -889,3 +889,89 @@ var _ = Describe("SetModelAndConfig metadata passthrough (chat completions)", fu Expect((*captured).RequestMetadata).To(HaveKeyWithValue("preserve_thinking", "true")) }) }) + +// A model config may default to decayed_mean pooling with a half-life +// (parameters: pooling / pooling_half_life_tokens). A request that +// overrides the scheme WITHOUT sending its own half-life must not inherit +// the config's half-life into an invalid pair — the middleware re-validates +// the merged config, and e.g. ("mean", 256) would 400 through no fault of +// the client. Found live: the first conv+mean A/B request against a +// decayed-default embedder config. +var _ = Describe("SetModelAndConfig per-request pooling override (embeddings)", func() { + var modelDir string + + BeforeEach(func() { + var err error + modelDir, err = os.MkdirTemp("", "localai-test-pooling-models-*") + Expect(err).ToNot(HaveOccurred()) + }) + + AfterEach(func() { + _ = os.RemoveAll(modelDir) + }) + + buildApp := func(cfgYAML string) (*echo.Echo, **config.ModelConfig) { + Expect(os.WriteFile(filepath.Join(modelDir, "test-model.yaml"), []byte(cfgYAML), 0644)).To(Succeed()) + + ss := &system.SystemState{Model: system.Model{ModelsPath: modelDir}} + appConfig := config.NewApplicationConfig() + appConfig.SystemState = ss + mcl := config.NewModelConfigLoader(modelDir) + ml := model.NewModelLoader(ss) + re := NewRequestExtractor(mcl, ml, appConfig) + + captured := new(*config.ModelConfig) + app := echo.New() + app.POST("/v1/embeddings", + func(c echo.Context) error { + if cfg, ok := c.Get(CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig); ok { + *captured = cfg + } + return c.String(http.StatusOK, "ok") + }, + re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }), + func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + if err := re.SetOpenAIRequest(c); err != nil { + return err + } + return next(c) + } + }, + ) + return app, captured + } + + decayedCfg := "name: test-model\nbackend: llama-cpp\nembeddings: true\n" + + "parameters:\n pooling: decayed_mean\n pooling_half_life_tokens: 256\n" + + embedReq := func(extra string) string { + return `{"model":"test-model","messages":[{"role":"user","content":"hi"}]` + extra + `}` + } + + It("drops the config half-life when the scheme override is not decayed_mean", func() { + app, captured := buildApp(decayedCfg) + rec := postJSON(app, "/v1/embeddings", embedReq(`,"pooling":"mean"`)) + + Expect(rec.Code).To(Equal(http.StatusOK), rec.Body.String()) + Expect(*captured).ToNot(BeNil()) + Expect((*captured).Pooling).To(Equal(config.PoolingMean)) + Expect((*captured).PoolingHalfLifeTokens).To(Equal(0)) + }) + + It("keeps the config half-life when overriding to decayed_mean explicitly", func() { + app, captured := buildApp(decayedCfg) + rec := postJSON(app, "/v1/embeddings", embedReq(`,"pooling":"decayed_mean"`)) + + Expect(rec.Code).To(Equal(http.StatusOK), rec.Body.String()) + Expect(*captured).ToNot(BeNil()) + Expect((*captured).PoolingHalfLifeTokens).To(Equal(256)) + }) + + It("still rejects a request pairing its own half-life with a non-decayed scheme", func() { + app, _ := buildApp(decayedCfg) + rec := postJSON(app, "/v1/embeddings", embedReq(`,"pooling":"mean","pooling_half_life_tokens":64`)) + + Expect(rec.Code).To(Equal(http.StatusBadRequest), rec.Body.String()) + }) +})