From d10374f84905f13ce236f2db29bc1de7b11b97dd Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Tue, 18 Aug 2026 08:37:43 +0100 Subject: [PATCH] feat(router): make KNN a first-class classifier with a persisted, curated corpus (#10652) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(router): make KNN a first-class classifier with a persisted, curated corpus 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] Signed-off-by: Richard Palethorpe * 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] Signed-off-by: Richard Palethorpe * 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] Signed-off-by: Richard Palethorpe * feat(router): bind knn corpus vectors to an embedder fingerprint and fail closed on mismatch Assisted-by: Claude:claude-fable-5 [Claude Code] Signed-off-by: Richard Palethorpe * chore(mcp): align corpus tool prompts and the mutating-tool safety list Assisted-by: Claude:claude-fable-5 [Claude Code] Signed-off-by: Richard Palethorpe * feat(proto,backend): report embedding shape from the llama-cpp backend Assisted-by: Claude:claude-fable-5 [Claude Code] Signed-off-by: Richard Palethorpe * feat(embeddings): Go-side pooling — mean/last/decayed_mean with half-life Assisted-by: Claude:claude-fable-5 [Claude Code] Signed-off-by: Richard Palethorpe * feat(embeddings): accept chat messages[] and per-request pooling on /v1/embeddings Assisted-by: Claude:claude-fable-5 [Claude Code] Signed-off-by: Richard Palethorpe * 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. Assisted-by: Claude:claude-fable-5 [Claude Code] Signed-off-by: Richard Palethorpe * 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. Assisted-by: Claude:claude-fable-5 [Claude Code] Signed-off-by: Richard Palethorpe * fix embedding pooling validation and router bounds Declare backend embedding layouts and reject incompatible pooling modes. Reset local-store dimensions after a full clear, validate KNN thresholds, and add real backend and store integration coverage. Assisted-by: Codex:gpt-5 Signed-off-by: Richard Palethorpe * ci: run local-store integration tests Build and install the local-store backend in the Linux test job, then run the existing store integration suite so new specs are discovered automatically. Assisted-by: Codex:gpt-5 Signed-off-by: Richard Palethorpe --------- Signed-off-by: Richard Palethorpe --- .github/workflows/test-extra.yml | 1 + .github/workflows/test.yml | 6 + Makefile | 12 + backend/backend.proto | 20 + backend/cpp/ik-llama-cpp/grpc-server.cpp | 1 + backend/cpp/llama-cpp/grpc-server.cpp | 58 +- backend/go/local-store/store.go | 9 +- backend/go/local-store/store_test.go | 40 ++ backend/python/insightface/backend.py | 5 +- backend/python/tinygrad/backend.py | 5 +- backend/python/transformers/backend.py | 5 +- backend/python/vllm/backend.py | 5 +- core/application/application.go | 15 + core/application/router_factories.go | 67 ++ core/application/router_factories_test.go | 31 + core/backend/embeddings.go | 80 ++- core/backend/pooling.go | 187 +++++ core/backend/pooling_test.go | 229 +++++++ core/backend/stores.go | 113 ++- core/config/backend_capabilities.go | 12 +- core/config/backend_capabilities_test.go | 11 + core/config/meta/registry.go | 86 ++- core/config/model_config.go | 211 ++++++ core/config/pooling_config_test.go | 158 +++++ core/config/router_knn_config_test.go | 48 ++ core/http/embeddings_pooling_test.go | 229 +++++++ .../endpoints/localai/api_instructions.go | 2 +- core/http/endpoints/localai/router_corpus.go | 168 +++++ .../endpoints/localai/router_corpus_test.go | 230 +++++++ core/http/endpoints/localai/router_decide.go | 29 +- .../endpoints/mcp/localai_assistant_test.go | 12 + core/http/endpoints/openai/embeddings.go | 53 +- core/http/endpoints/openai/realtime_model.go | 13 +- core/http/middleware/request.go | 41 +- core/http/middleware/request_test.go | 86 +++ core/http/middleware/route_model.go | 156 ++++- core/http/middleware/route_model_test.go | 222 ++++++ .../http/react-ui/e2e/middleware-page.spec.js | 51 +- core/http/react-ui/src/App.css | 16 + core/http/react-ui/src/pages/Middleware.jsx | 49 +- core/http/routes/anthropic.go | 11 +- core/http/routes/middleware.go | 90 ++- core/http/routes/openai.go | 13 +- core/schema/localai.go | 142 +++- core/schema/prediction.go | 13 + .../routing/corpus/corpus_suite_test.go | 13 + core/services/routing/corpus/manager.go | 648 ++++++++++++++++++ core/services/routing/corpus/manager_test.go | 347 ++++++++++ core/services/routing/corpus/resolve.go | 130 ++++ core/services/routing/router/decisions.go | 34 +- .../routing/router/embedding_cache_test.go | 26 + core/services/routing/router/knn.go | 259 +++++++ core/services/routing/router/knn_test.go | 219 ++++++ core/services/routing/router/resolve.go | 2 + core/services/routing/router/types.go | 42 ++ core/templates/evaluator.go | 27 + core/templates/evaluator_embedding_test.go | 88 +++ docs/content/features/embeddings.md | 71 ++ docs/content/operations/middleware.md | 177 ++++- pkg/grpc/server.go | 5 +- pkg/mcp/localaitools/client.go | 12 + pkg/mcp/localaitools/coverage_test.go | 41 +- 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 | 79 +++ 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 | 49 +- pkg/mcp/localaitools/tools.go | 68 +- pkg/mcp/localaitools/tools_middleware.go | 57 +- swagger/docs.go | 296 ++++++++ swagger/swagger.json | 296 ++++++++ swagger/swagger.yaml | 208 ++++++ tests/e2e-backends/backend_test.go | 27 +- tests/e2e/mock-backend/main.go | 48 +- tests/e2e/mock_backend_test.go | 30 + tests/integration/stores_test.go | 17 + 80 files changed, 6165 insertions(+), 319 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 create mode 100644 core/config/router_knn_config_test.go create mode 100644 core/http/embeddings_pooling_test.go 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/corpus/resolve.go create mode 100644 core/services/routing/router/knn.go create mode 100644 core/services/routing/router/knn_test.go create mode 100644 core/templates/evaluator_embedding_test.go create mode 100644 pkg/mcp/localaitools/prompts/skills/manage_router_corpus.md diff --git a/.github/workflows/test-extra.yml b/.github/workflows/test-extra.yml index e4e00f7d1..8087a9413 100644 --- a/.github/workflows/test-extra.yml +++ b/.github/workflows/test-extra.yml @@ -526,6 +526,7 @@ jobs: - name: Build llama-cpp backend image and run gRPC e2e tests run: | make test-extra-backend-llama-cpp + make test-extra-backend-llama-cpp-embeddings tests-llama-cpp-grpc-transcription: needs: detect-changes if: needs.detect-changes.outputs.llama-cpp == 'true' || needs.detect-changes.outputs.run-all == 'true' diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7081d630c..7929f6b65 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -65,6 +65,12 @@ jobs: - name: Test (with coverage gate) run: | PATH="$PATH:/root/go/bin" make --jobs 5 --output-sync=target test-coverage-check + # tests/integration is outside the coverage roots because its store specs + # need a live backend. test-stores builds and installs local-store before + # running the complete suite, so new local-store specs are collected + # automatically without adding another workflow entry. + - name: Test local-store integration + run: PATH="$PATH:$HOME/go/bin" make test-stores - name: Upload coverage report if: ${{ always() }} uses: actions/upload-artifact@v4 diff --git a/Makefile b/Makefile index 1f25e246c..6b813586e 100644 --- a/Makefile +++ b/Makefile @@ -676,6 +676,7 @@ test-extra: prepare-test-extra ## BACKEND_TEST_PROMPT Override the prompt used in predict/stream specs. ## BACKEND_TEST_OPTIONS Comma-separated Options[] entries forwarded to LoadModel, ## e.g. "tool_parser:hermes,reasoning_parser:qwen3". +## BACKEND_TEST_EMBEDDING_LAYOUT Expected EmbeddingResult layout: "final" or "per_token". ## ## Direct usage (image already built, no docker-build-* dependency): ## @@ -705,6 +706,7 @@ test-extra-backend: protogen-go BACKEND_TEST_CAPS="$$BACKEND_TEST_CAPS" \ BACKEND_TEST_PROMPT="$$BACKEND_TEST_PROMPT" \ BACKEND_TEST_OPTIONS="$$BACKEND_TEST_OPTIONS" \ + BACKEND_TEST_EMBEDDING_LAYOUT="$$BACKEND_TEST_EMBEDDING_LAYOUT" \ BACKEND_TEST_TOOL_PROMPT="$$BACKEND_TEST_TOOL_PROMPT" \ BACKEND_TEST_TOOL_NAME="$$BACKEND_TEST_TOOL_NAME" \ BACKEND_TEST_CACHE_TYPE_K="$$BACKEND_TEST_CACHE_TYPE_K" \ @@ -724,6 +726,15 @@ test-extra-backend-llama-cpp: docker-build-llama-cpp BACKEND_TEST_CAPS=health,load,predict,stream,logprobs,logit_bias \ $(MAKE) test-extra-backend +## Raw llama.cpp embeddings are required by Go-side pooling. This exercises the +## real C++ backend and verifies that it marks the flattened matrix per-token. +test-extra-backend-llama-cpp-embeddings: docker-build-llama-cpp + BACKEND_IMAGE=local-ai-backend:llama-cpp \ + BACKEND_TEST_CAPS=health,load,embeddings \ + BACKEND_TEST_OPTIONS=pooling:none \ + BACKEND_TEST_EMBEDDING_LAYOUT=per_token \ + $(MAKE) test-extra-backend + test-extra-backend-ik-llama-cpp: docker-build-ik-llama-cpp BACKEND_IMAGE=local-ai-backend:ik-llama-cpp $(MAKE) test-extra-backend @@ -813,6 +824,7 @@ test-extra-backend-tinygrad-embeddings: docker-build-tinygrad BACKEND_IMAGE=local-ai-backend:tinygrad \ BACKEND_TEST_MODEL_NAME=Qwen/Qwen3-0.6B \ BACKEND_TEST_CAPS=health,load,embeddings \ + BACKEND_TEST_EMBEDDING_LAYOUT=final \ $(MAKE) test-extra-backend ## tinygrad — Stable Diffusion 1.5. The original CompVis/runwayml repos have diff --git a/backend/backend.proto b/backend/backend.proto index 2de6c8711..7be3d5a4c 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -536,8 +536,28 @@ message Result { bool success = 2; } +// EmbeddingLayout describes whether embeddings contains one final vector or +// a matrix of per-token vectors. Go-side pooling must never infer this from +// tokens/dim alone: a one-token raw matrix and a final vector have the same +// shape. +enum EmbeddingLayout { + EMBEDDING_LAYOUT_UNSPECIFIED = 0; + EMBEDDING_LAYOUT_FINAL = 1; + EMBEDDING_LAYOUT_PER_TOKEN = 2; +} + 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; + EmbeddingLayout layout = 5; } message TranscriptRequest { diff --git a/backend/cpp/ik-llama-cpp/grpc-server.cpp b/backend/cpp/ik-llama-cpp/grpc-server.cpp index ce43d6d01..600e50e8a 100644 --- a/backend/cpp/ik-llama-cpp/grpc-server.cpp +++ b/backend/cpp/ik-llama-cpp/grpc-server.cpp @@ -2565,6 +2565,7 @@ public: grpc::Status Embedding(ServerContext* context, const backend::PredictOptions* request, backend::EmbeddingResult* embeddingResult) { auto identity = checkModelIdentity(request); if (!identity.ok()) return identity; + embeddingResult->set_layout(backend::EMBEDDING_LAYOUT_FINAL); json data = parse_options(false, request, llama); const int task_id = llama.queue_tasks.get_new_id(); llama.queue_results.add_waiting_task_id(task_id); diff --git a/backend/cpp/llama-cpp/grpc-server.cpp b/backend/cpp/llama-cpp/grpc-server.cpp index 1091a786d..82080c644 100644 --- a/backend/cpp/llama-cpp/grpc-server.cpp +++ b/backend/cpp/llama-cpp/grpc-server.cpp @@ -2911,42 +2911,40 @@ public: 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); + embeddingResult->set_layout( + llama_pooling_type(ctx_server.get_llama_context()) == LLAMA_POOLING_TYPE_NONE + ? backend::EMBEDDING_LAYOUT_PER_TOKEN + : backend::EMBEDDING_LAYOUT_FINAL); - - + std::cout << "[DEBUG] Embedding vectors: " << n_vectors << " x " << dim << std::endl; return grpc::Status::OK; } diff --git a/backend/go/local-store/store.go b/backend/go/local-store/store.go index a46c52691..e7fd1d091 100644 --- a/backend/go/local-store/store.go +++ b/backend/go/local-store/store.go @@ -38,8 +38,9 @@ type Store struct { // keysAreNormalized stays true until any non-unit-magnitude key // is added; once false, the magnitude-aware fallback path is // used by Find. Re-evaluated only at Set time, never again on - // its own — a deletion of the offending key does NOT flip it - // back to true (the bookkeeping cost would dominate the gain). + // its own — a partial deletion of the offending key does NOT flip + // it back to true (the bookkeeping cost would dominate the gain). + // An empty store returns to its initial state. keysAreNormalized bool // keyLen is the dimension of every stored key. -1 means "no @@ -142,6 +143,10 @@ func (s *Store) StoresDelete(opts *pb.StoresDeleteOptions) error { mergedV = append(mergedV, tailV...) s.keys = mergedK s.values = mergedV + if len(s.keys) == 0 { + s.keyLen = -1 + s.keysAreNormalized = true + } assert(slices.IsSortedFunc(s.keys, slices.Compare[[]float32]), "Delete: s.keys not sorted post-merge") assert(len(s.keys) == len(s.values), "Delete: keys/values length skew") return nil diff --git a/backend/go/local-store/store_test.go b/backend/go/local-store/store_test.go index a7f4d7e2e..23ed9b172 100644 --- a/backend/go/local-store/store_test.go +++ b/backend/go/local-store/store_test.go @@ -105,6 +105,46 @@ var _ = Describe("StoresDelete", func() { })).To(Succeed(), "delete of missing key should succeed") Expect(s.keys).To(HaveLen(1)) }) + + It("reopens the dimension after deleting every key", func() { + s := NewStore() + oldKey := []float32{2, 0, 0} + mustSet(s, [][]float32{oldKey}, [][]byte{[]byte("3d")}) + Expect(s.keysAreNormalized).To(BeFalse()) + + Expect(s.StoresDelete(&pb.StoresDeleteOptions{ + Keys: wrapKeys([][]float32{oldKey}), + })).To(Succeed()) + Expect(s.keys).To(BeEmpty()) + Expect(s.keyLen).To(Equal(-1)) + Expect(s.keysAreNormalized).To(BeTrue()) + + newKey := normalizeVec([]float32{1, 1}) + mustSet(s, [][]float32{newKey}, [][]byte{[]byte("2d")}) + res, err := s.StoresFind(&pb.StoresFindOptions{ + Key: &pb.StoresKey{Floats: newKey}, + TopK: 1, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(res.Values).To(HaveLen(1)) + Expect(string(res.Values[0].Bytes)).To(Equal("2d")) + }) + + It("retains the dimension after a partial delete", func() { + s := NewStore() + mustSet(s, + [][]float32{{1, 0, 0}, {0, 1, 0}}, + [][]byte{[]byte("x"), []byte("y")}, + ) + Expect(s.StoresDelete(&pb.StoresDeleteOptions{ + Keys: wrapKeys([][]float32{{1, 0, 0}}), + })).To(Succeed()) + Expect(s.keyLen).To(Equal(3)) + Expect(s.StoresSet(&pb.StoresSetOptions{ + Keys: wrapKeys([][]float32{{1, 0}}), + Values: wrapValues([][]byte{[]byte("2d")}), + })).NotTo(Succeed()) + }) }) var _ = Describe("StoresFind", func() { diff --git a/backend/python/insightface/backend.py b/backend/python/insightface/backend.py index 60b89ef0c..763e8d0ce 100644 --- a/backend/python/insightface/backend.py +++ b/backend/python/insightface/backend.py @@ -127,7 +127,10 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): context.set_code(grpc.StatusCode.NOT_FOUND) context.set_details("no face detected") return backend_pb2.EmbeddingResult() - return backend_pb2.EmbeddingResult(embeddings=[float(x) for x in vec]) + return backend_pb2.EmbeddingResult( + embeddings=[float(x) for x in vec], + layout=backend_pb2.EMBEDDING_LAYOUT_FINAL, + ) def Detect(self, request, context): if self.engine is None: diff --git a/backend/python/tinygrad/backend.py b/backend/python/tinygrad/backend.py index fa643568c..3f2d5da50 100644 --- a/backend/python/tinygrad/backend.py +++ b/backend/python/tinygrad/backend.py @@ -638,7 +638,10 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): normalized = (pooled / (norm + 1e-12)) vec = normalized.cast(dtypes.float32).tolist() - return backend_pb2.EmbeddingResult(embeddings=[float(x) for x in vec]) + return backend_pb2.EmbeddingResult( + embeddings=[float(x) for x in vec], + layout=backend_pb2.EMBEDDING_LAYOUT_FINAL, + ) except Exception as exc: import traceback traceback.print_exc() diff --git a/backend/python/transformers/backend.py b/backend/python/transformers/backend.py index 30d6334ba..aa427267a 100644 --- a/backend/python/transformers/backend.py +++ b/backend/python/transformers/backend.py @@ -375,7 +375,10 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # Pool to get sentence embeddings; i.e. generate one 1024 vector for the entire sentence sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask']) embeds = sentence_embeddings[0] - return backend_pb2.EmbeddingResult(embeddings=embeds) + return backend_pb2.EmbeddingResult( + embeddings=embeds, + layout=backend_pb2.EMBEDDING_LAYOUT_FINAL, + ) async def _predict(self, request, context, streaming=False): set_seed(request.Seed) diff --git a/backend/python/vllm/backend.py b/backend/python/vllm/backend.py index a5a85041b..8fd3c2dc1 100644 --- a/backend/python/vllm/backend.py +++ b/backend/python/vllm/backend.py @@ -336,7 +336,10 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): context.set_code(grpc.StatusCode.INVALID_ARGUMENT) context.set_details("No embeddings were calculated.") return backend_pb2.EmbeddingResult() - return backend_pb2.EmbeddingResult(embeddings=outputs[0].outputs.embedding) + return backend_pb2.EmbeddingResult( + embeddings=outputs[0].outputs.embedding, + layout=backend_pb2.EMBEDDING_LAYOUT_FINAL, + ) async def PredictStream(self, request, context): """ diff --git a/core/application/application.go b/core/application/application.go index fd2f1e404..b49851d47 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" @@ -19,6 +21,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" @@ -77,6 +80,7 @@ type Application struct { mitmHostConflicts atomic.Pointer[map[string][]string] routerDecisions router.DecisionStore routerRegistry *router.Registry + routerCorpus *corpus.Manager admissionLimiter *admission.Limiter watchdogMutex sync.Mutex watchdogStop chan bool @@ -149,6 +153,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. @@ -578,6 +586,13 @@ 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.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 // 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 523c4465e..dff2d5798 100644 --- a/core/application/router_factories.go +++ b/core/application/router_factories.go @@ -2,10 +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 @@ -100,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 @@ -118,3 +179,9 @@ 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, a.backendLoader, storeName) } + +// RouterCorpus returns the process-wide KNN corpus manager, built in +// newApplication. +func (a *Application) RouterCorpus() *corpus.Manager { + return a.routerCorpus +} diff --git a/core/application/router_factories_test.go b/core/application/router_factories_test.go index 91b26d491..66ee0578b 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/backend/embeddings.go b/core/backend/embeddings.go index d68a334fb..55abb76f5 100644 --- a/core/backend/embeddings.go +++ b/core/backend/embeddings.go @@ -2,6 +2,7 @@ package backend import ( "context" + "errors" "fmt" "time" @@ -9,9 +10,80 @@ 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" ) +type embeddingPoolingCompatibilityError struct { + message string +} + +func (e *embeddingPoolingCompatibilityError) Error() string { + return e.message +} + +func poolingCompatibilityErrorf(format string, args ...any) error { + return &embeddingPoolingCompatibilityError{message: fmt.Sprintf(format, args...)} +} + +// IsEmbeddingPoolingCompatibilityError reports errors caused by a requested +// pooling scheme disagreeing with the layout declared by the loaded backend. +// HTTP callers map these client-selectable incompatibilities to status 400. +func IsEmbeddingPoolingCompatibilityError(err error) bool { + var target *embeddingPoolingCompatibilityError + return errors.As(err, &target) +} + +// finishEmbeddingResult applies the model's Go-side pooling scheme only when +// the backend declares that it returned per-token vectors. Shape alone is not +// sufficient: one raw token and one final vector are both reported as 1 x dim. +// Legacy backends remain compatible with backend pooling, but cannot opt in to +// Go-side pooling until they declare their result layout. +func finishEmbeddingResult(res *proto.EmbeddingResult, modelConfig config.ModelConfig) ([]float32, error) { + scheme := modelConfig.Pooling + if scheme == "" || scheme == PoolingBackend { + switch res.GetLayout() { + case proto.EmbeddingLayout_EMBEDDING_LAYOUT_UNSPECIFIED, proto.EmbeddingLayout_EMBEDDING_LAYOUT_FINAL: + return res.Embeddings, nil + case proto.EmbeddingLayout_EMBEDDING_LAYOUT_PER_TOKEN: + return nil, poolingCompatibilityErrorf( + "pooling %q cannot pass through per-token embeddings: choose %q, %q or %q, or load the backend with pooling enabled", + PoolingBackend, PoolingMean, PoolingLast, PoolingDecayedMean) + default: + return nil, poolingCompatibilityErrorf("pooling %q cannot use unknown embedding layout %d", PoolingBackend, res.GetLayout()) + } + } + switch res.GetLayout() { + case proto.EmbeddingLayout_EMBEDDING_LAYOUT_PER_TOKEN: + // Pool below after validating the reported matrix shape. + case proto.EmbeddingLayout_EMBEDDING_LAYOUT_FINAL: + return nil, poolingCompatibilityErrorf( + "pooling %q needs per-token embeddings but this backend returned a final vector; configure raw per-token output if the backend supports it (llama.cpp: options [\"pooling:none\"])", + scheme) + case proto.EmbeddingLayout_EMBEDDING_LAYOUT_UNSPECIFIED: + return nil, poolingCompatibilityErrorf( + "pooling %q needs per-token embeddings but the backend did not declare its embedding layout: rebuild/update the backend to report EmbeddingResult.layout", + scheme) + default: + return nil, poolingCompatibilityErrorf("pooling %q cannot use unknown embedding layout %d", scheme, res.GetLayout()) + } + 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 +138,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 000000000..f4f1799b8 --- /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 000000000..eb3720518 --- /dev/null +++ b/core/backend/pooling_test.go @@ -0,0 +1,229 @@ +package backend + +import ( + "math" + + "github.com/mudler/LocalAI/core/config" + "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()) + }) + }) + + Describe("finishEmbeddingResult", func() { + finish := func(res *proto.EmbeddingResult, scheme string) ([]float32, error) { + cfg := config.ModelConfig{} + cfg.Pooling = scheme + cfg.Options = []string{"embd_normalize:-1"} + return finishEmbeddingResult(res, cfg) + } + final := func() *proto.EmbeddingResult { + return &proto.EmbeddingResult{ + Embeddings: []float32{1, 2}, + Tokens: 1, + Dim: 2, + Layout: proto.EmbeddingLayout_EMBEDDING_LAYOUT_FINAL, + } + } + perToken := func() *proto.EmbeddingResult { + return &proto.EmbeddingResult{ + Embeddings: []float32{1, 2, 3, 4}, + Tokens: 2, + Dim: 2, + Layout: proto.EmbeddingLayout_EMBEDDING_LAYOUT_PER_TOKEN, + } + } + + It("passes a final vector through to backend pooling", func() { + got, err := finish(final(), PoolingBackend) + Expect(err).NotTo(HaveOccurred()) + Expect(got).To(Equal([]float32{1, 2})) + }) + + It("rejects Go-side pooling after the backend returned a final vector", func() { + _, err := finish(final(), PoolingMean) + Expect(err).To(MatchError(ContainSubstring("final vector"))) + }) + + It("pools a backend-declared per-token matrix", func() { + got, err := finish(perToken(), PoolingMean) + Expect(err).NotTo(HaveOccurred()) + Expect(got).To(Equal([]float32{2, 3})) + }) + + It("rejects backend pass-through of a per-token matrix", func() { + _, err := finish(perToken(), PoolingBackend) + Expect(err).To(MatchError(ContainSubstring("cannot pass through per-token"))) + }) + + It("allows legacy layout only for backend pooling", func() { + legacy := &proto.EmbeddingResult{Embeddings: []float32{1, 2}, Tokens: 1, Dim: 2} + got, err := finish(legacy, PoolingBackend) + Expect(err).NotTo(HaveOccurred()) + Expect(got).To(Equal([]float32{1, 2})) + + _, err = finish(legacy, PoolingMean) + Expect(err).To(MatchError(ContainSubstring("did not declare"))) + }) + + It("fails closed for an unknown layout value", func() { + res := perToken() + res.Layout = proto.EmbeddingLayout(99) + _, err := finish(res, PoolingMean) + Expect(err).To(MatchError(ContainSubstring("unknown embedding layout"))) + _, err = finish(res, PoolingBackend) + Expect(err).To(MatchError(ContainSubstring("unknown embedding layout"))) + }) + }) +}) diff --git a/core/backend/stores.go b/core/backend/stores.go index 5c97e9792..6a566af53 100644 --- a/core/backend/stores.go +++ b/core/backend/stores.go @@ -15,13 +15,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. // cl resolves the per-store model config (backend + options); it may be nil, @@ -45,18 +55,30 @@ func (s *localVectorStore) backend(_ context.Context) (grpc.Backend, error) { return StoreBackend(s.loader, s.appConfig, s.cl, s.storeName, "") } -func (s *localVectorStore) Search(ctx context.Context, vec []float32) (sim float64, payload []byte, ok bool, err error) { +// 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) { outcome := "hit" + sim := 0.0 be, berr := s.backend(ctx) if berr != nil { outcome = "backend_load_error" err = fmt.Errorf("vector store load: %w", berr) s.recordTrace("", time.Now(), "search", len(vec), 0, outcome, err) - return 0, nil, false, err + return nil, err } release, err := AcquireGlobalBackendSlot() if err != nil { - return 0, nil, false, err + return nil, err } defer release() start := time.Now() @@ -64,16 +86,21 @@ func (s *localVectorStore) Search(ctx context.Context, vec []float32) (sim float defer func() { s.recordTrace(traceID, start, "search", len(vec), sim, outcome, err) }() - _, values, similarities, ferr := store.Find(ctx, be, vec, 1) + _, values, similarities, ferr := store.Find(ctx, be, vec, k) if ferr != nil { outcome = "find_error" - return 0, nil, false, fmt.Errorf("vector store find: %w", ferr) + return nil, fmt.Errorf("vector store find: %w", ferr) } - if len(values) == 0 || len(similarities) == 0 { + if len(values) == 0 { outcome = "miss" - return 0, nil, false, nil + return nil, nil } - return float64(similarities[0]), values[0], true, 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) { @@ -102,6 +129,72 @@ 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) { + outcome := "ok" + dim := 0 + if len(vecs) > 0 { + dim = len(vecs[0]) + } + be, berr := s.backend(ctx) + if berr != nil { + outcome = "backend_load_error" + err = fmt.Errorf("vector store load: %w", berr) + s.recordTrace("", time.Now(), "insert_batch", dim, 0, outcome, err) + return err + } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return err + } + defer release() + start := time.Now() + traceID := s.beginTrace(start, "insert_batch") + defer func() { + s.recordTrace(traceID, start, "insert_batch", dim, 0, outcome, err) + }() + 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) { + outcome := "ok" + dim := 0 + if len(vecs) > 0 { + dim = len(vecs[0]) + } + be, berr := s.backend(ctx) + if berr != nil { + outcome = "backend_load_error" + err = fmt.Errorf("vector store load: %w", berr) + s.recordTrace("", time.Now(), "delete", dim, 0, outcome, err) + return err + } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return err + } + defer release() + start := time.Now() + traceID := s.beginTrace(start, "delete") + defer func() { + s.recordTrace(traceID, start, "delete", dim, 0, outcome, err) + }() + 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/backend_capabilities.go b/core/config/backend_capabilities.go index 50ace9b36..c0693d048 100644 --- a/core/config/backend_capabilities.go +++ b/core/config/backend_capabilities.go @@ -341,20 +341,20 @@ var BackendCapabilities = map[string]BackendCapability{ Description: "HuggingFace transformers — general-purpose Python inference", }, "mlx": { - GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream, MethodEmbedding}, - PossibleUsecases: []string{UsecaseChat, UsecaseCompletion, UsecaseEmbeddings}, + GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream}, + PossibleUsecases: []string{UsecaseChat, UsecaseCompletion}, DefaultUsecases: []string{UsecaseChat}, Description: "Apple MLX framework — optimized for Apple Silicon", }, "mlx-distributed": { - GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream, MethodEmbedding}, - PossibleUsecases: []string{UsecaseChat, UsecaseCompletion, UsecaseEmbeddings}, + GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream}, + PossibleUsecases: []string{UsecaseChat, UsecaseCompletion}, DefaultUsecases: []string{UsecaseChat}, Description: "MLX distributed inference across multiple Apple Silicon devices", }, "mlx-vlm": { - GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream, MethodEmbedding}, - PossibleUsecases: []string{UsecaseChat, UsecaseCompletion, UsecaseEmbeddings, UsecaseVision}, + GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream}, + PossibleUsecases: []string{UsecaseChat, UsecaseCompletion, UsecaseVision}, DefaultUsecases: []string{UsecaseChat, UsecaseVision}, AcceptsImages: true, AcceptsAudios: true, diff --git a/core/config/backend_capabilities_test.go b/core/config/backend_capabilities_test.go index b18c41069..eeb6c3135 100644 --- a/core/config/backend_capabilities_test.go +++ b/core/config/backend_capabilities_test.go @@ -157,6 +157,17 @@ var _ = Describe("nemo-speech-cpp capabilities", func() { }) }) +var _ = Describe("MLX embedding capabilities", func() { + It("does not advertise RPCs that the backends return as unimplemented", func() { + for _, name := range []string{"mlx", "mlx-distributed", "mlx-vlm"} { + capability := GetBackendCapability(name) + Expect(capability).NotTo(BeNil()) + Expect(capability.GRPCMethods).ToNot(ContainElement(MethodEmbedding), name) + Expect(capability.PossibleUsecases).ToNot(ContainElement(UsecaseEmbeddings), name) + } + }) +}) + // audio-cpp advertises voice cloning from the backend itself and ships // audio-cpp-chatterbox, whose family serves cloning and NOT plain TTS, so a // reference clip is the only way to use it. Without a capability entry diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index aa25e57ad..9e3711af2 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -1,6 +1,9 @@ package meta -import "github.com/mudler/LocalAI/core/services/routing/piipattern" +import ( + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/routing/piipattern" +) // builtinPatternOptions turns the piipattern built-in catalogue into select // options for the editor's built-in-patterns checklist, keeping the catalogue @@ -257,6 +260,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": { @@ -991,17 +1018,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, @@ -1093,5 +1122,56 @@ 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.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 (1-1024). 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), + Max: f64(float64(config.RouterKNNMaxK)), + Order: 243, + }, + "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: 244, + }, + "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: 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: 246, + }, } } diff --git a/core/config/model_config.go b/core/config/model_config.go index 9a43bd539..76509598a 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -1,8 +1,11 @@ package config import ( + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" + "math" "os" "path/filepath" "regexp" @@ -17,6 +20,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" ) @@ -358,7 +362,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 @@ -392,6 +406,97 @@ type EmbeddingCacheConfig struct { StoreName string `yaml:"store_name,omitempty" json:"store_name,omitempty"` } +// RouterKNNMaxK bounds the local priority queue and the number of neighbors +// returned across the store boundary for one routing decision. +const RouterKNNMaxK = 1024 + +// 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"` + + // 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 + // 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"` +} + +// Validate rejects values that would make weighted cosine voting invalid or +// allow one request to allocate an unbounded neighbor result. Zero retains the +// documented default sentinel for backward compatibility. +func (k *RouterKNNConfig) Validate() error { + if k.K < 0 || k.K > RouterKNNMaxK { + return fmt.Errorf("router.knn.k must be 0 (default) or between 1 and %d, got %d", RouterKNNMaxK, k.K) + } + if math.IsNaN(k.SimilarityThreshold) || math.IsInf(k.SimilarityThreshold, 0) { + return fmt.Errorf("router.knn.similarity_threshold must be finite, got %v", k.SimilarityThreshold) + } + if k.SimilarityThreshold < 0 || k.SimilarityThreshold > 1 { + return fmt.Errorf("router.knn.similarity_threshold must be 0 (default) or greater than 0 and at most 1, got %v", k.SimilarityThreshold) + } + if math.IsNaN(k.VoteThreshold) || math.IsInf(k.VoteThreshold, 0) { + return fmt.Errorf("router.knn.vote_threshold must be finite, got %v", k.VoteThreshold) + } + if k.VoteThreshold < 0 || k.VoteThreshold > 1 { + return fmt.Errorf("router.knn.vote_threshold must be 0 (default) or greater than 0 and at most 1, got %v", k.VoteThreshold) + } + return nil +} + +// 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 +} + +// 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 @@ -1411,6 +1516,15 @@ func (cfg *ModelConfig) SetDefaults(opts ...ConfigLoaderOption) { } runBackendHooks(cfg, lo.modelPath) + // llama.cpp selects raw-vs-final embedding layout at model load. Other + // backends may expose per-token embeddings through their own contracts, so + // do not inject a llama.cpp-specific option into them. + if IsGoSidePooling(cfg.Pooling) && IsLlamaCppBackend(cfg.Backend) && !hasLlamaPoolingOption(cfg.Options) { + 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 @@ -1565,6 +1679,11 @@ func (c *ModelConfig) Validate() (bool, error) { return false, fmt.Errorf("router: unknown score_normalization %q (expected %q or %q)", c.Router.ScoreNormalization, ScoreNormalizationRaw, ScoreNormalizationMean) } + if c.Router.KNN != nil { + if err := c.Router.KNN.Validate(); err != nil { + return false, err + } + } // router.classifier_system_template parses as Go text/template // (Sprig funcs available at execution time). Reject malformed @@ -1576,6 +1695,29 @@ func (c *ModelConfig) Validate() (bool, error) { } } + // Go-side embedding pooling: reject unknown schemes and half-life + // misconfigurations at load time. llama.cpp chooses its result layout at + // model load, so its configured layout must agree with the model-level + // pooling scheme in both directions. Other backends declare the actual + // result layout on EmbeddingResult and are checked at request time. + if err := ValidatePooling(c.Pooling, c.PoolingHalfLifeTokens); err != nil { + return false, fmt.Errorf("parameters: %w", err) + } + if IsLlamaCppBackend(c.Backend) { + backendPooling, configured := llamaPoolingOption(c.Options) + raw := configured && backendPooling == "none" + if IsGoSidePooling(c.Pooling) && configured && !raw { + return false, fmt.Errorf( + "parameters.pooling %q needs raw per-token vectors but llama.cpp pooling %q returns a final vector: drop the option or set \"pooling:none\"", + c.Pooling, backendPooling) + } + if !IsGoSidePooling(c.Pooling) && raw { + return false, fmt.Errorf( + "parameters.pooling %q cannot pass through llama.cpp per-token vectors from \"pooling:none\": choose %q, %q or %q, or enable backend pooling", + c.Pooling, PoolingMean, PoolingLast, PoolingDecayedMean) + } + } + return true, nil } @@ -1588,6 +1730,75 @@ 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 +} + +// llamaPoolingOption returns the last load-time llama.cpp pooling option. The +// server accepts both names and processes options in order, so mirroring that +// behavior avoids disagreeing with the layout the backend actually loads. +func llamaPoolingOption(options []string) (string, bool) { + var value string + found := false + for _, o := range options { + if name, candidate, ok := strings.Cut(o, ":"); ok && (name == "pooling" || name == "pooling_type") { + value = candidate + found = true + } + } + return value, found +} + +func hasLlamaPoolingOption(options []string) bool { + _, found := llamaPoolingOption(options) + return found +} + 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 000000000..1f09770d4 --- /dev/null +++ b/core/config/pooling_config_test.go @@ -0,0 +1,158 @@ +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")) + } + }) + + It("does not inject llama.cpp options into another backend", func() { + for _, backend := range []string{"transformers", "ik-llama-cpp"} { + cfg := poolingTestConfig(PoolingMean, 0) + cfg.Backend = backend + cfg.SetDefaults() + Expect(cfg.Options).ToNot(ContainElement("pooling:none"), backend) + } + }) + + It("configures pinned llama.cpp variants for raw output", func() { + cfg := poolingTestConfig(PoolingMean, 0) + cfg.Backend = "cuda12-llama-cpp" + cfg.SetDefaults() + Expect(cfg.Options).To(ContainElement("pooling:none")) + }) + + It("recognizes llama.cpp's pooling_type alias", func() { + cfg := poolingTestConfig(PoolingMean, 0, "pooling_type:none") + 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()) + }) + + It("rejects backend pooling with a raw llama.cpp instance", func() { + for _, scheme := range []string{"", PoolingBackend} { + cfg := poolingTestConfig(scheme, 0, "pooling:none") + valid, err := cfg.Validate() + Expect(err).To(HaveOccurred(), "scheme %q", scheme) + Expect(valid).To(BeFalse(), "scheme %q", scheme) + Expect(err.Error()).To(ContainSubstring("per-token")) + } + }) + + It("uses the last llama.cpp pooling option like the backend", func() { + cfg := poolingTestConfig(PoolingMean, 0, "pooling:mean", "pooling_type:none") + valid, err := cfg.Validate() + Expect(err).ToNot(HaveOccurred()) + Expect(valid).To(BeTrue()) + }) + + It("leaves non-llama pooling compatibility to the backend result", func() { + cfg := poolingTestConfig(PoolingMean, 0, "pooling:custom") + cfg.Backend = "future-embedding-backend" + valid, err := cfg.Validate() + Expect(err).ToNot(HaveOccurred()) + Expect(valid).To(BeTrue()) + }) + }) +}) diff --git a/core/config/router_knn_config_test.go b/core/config/router_knn_config_test.go new file mode 100644 index 000000000..9c3b4de40 --- /dev/null +++ b/core/config/router_knn_config_test.go @@ -0,0 +1,48 @@ +package config + +import ( + "math" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("RouterKNNConfig validation", func() { + validate := func(knn RouterKNNConfig) error { + cfg := poolingTestConfig("", 0) + cfg.Backend = "transformers" + cfg.Router.KNN = &knn + _, err := cfg.Validate() + return err + } + + It("accepts defaults and documented boundaries", func() { + for _, knn := range []RouterKNNConfig{ + {}, + {K: 1, SimilarityThreshold: math.SmallestNonzeroFloat64, VoteThreshold: math.SmallestNonzeroFloat64}, + {K: RouterKNNMaxK, SimilarityThreshold: 1, VoteThreshold: 1}, + } { + Expect(validate(knn)).To(Succeed()) + } + }) + + DescribeTable("rejects an invalid bound", + func(knn RouterKNNConfig, field string) { + err := validate(knn) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(field)) + }, + Entry("negative k", RouterKNNConfig{K: -1}, "router.knn.k"), + Entry("k above the operational cap", RouterKNNConfig{K: RouterKNNMaxK + 1}, "router.knn.k"), + Entry("negative similarity", RouterKNNConfig{SimilarityThreshold: -math.SmallestNonzeroFloat64}, "router.knn.similarity_threshold"), + Entry("similarity above one", RouterKNNConfig{SimilarityThreshold: math.Nextafter(1, 2)}, "router.knn.similarity_threshold"), + Entry("NaN similarity", RouterKNNConfig{SimilarityThreshold: math.NaN()}, "router.knn.similarity_threshold"), + Entry("positive infinite similarity", RouterKNNConfig{SimilarityThreshold: math.Inf(1)}, "router.knn.similarity_threshold"), + Entry("negative infinite similarity", RouterKNNConfig{SimilarityThreshold: math.Inf(-1)}, "router.knn.similarity_threshold"), + Entry("negative vote", RouterKNNConfig{VoteThreshold: -math.SmallestNonzeroFloat64}, "router.knn.vote_threshold"), + Entry("vote above one", RouterKNNConfig{VoteThreshold: math.Nextafter(1, 2)}, "router.knn.vote_threshold"), + Entry("NaN vote", RouterKNNConfig{VoteThreshold: math.NaN()}, "router.knn.vote_threshold"), + Entry("positive infinite vote", RouterKNNConfig{VoteThreshold: math.Inf(1)}, "router.knn.vote_threshold"), + Entry("negative infinite vote", RouterKNNConfig{VoteThreshold: math.Inf(-1)}, "router.knn.vote_threshold"), + ) +}) diff --git a/core/http/embeddings_pooling_test.go b/core/http/embeddings_pooling_test.go new file mode 100644 index 000000000..454cce5b4 --- /dev/null +++ b/core/http/embeddings_pooling_test.go @@ -0,0 +1,229 @@ +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 layout", func() { + // "no-shape:" makes the mock omit layout, simulating a backend built + // before EmbeddingResult declared whether its vector was final or raw. + status, response := postEmbeddings(`{ + "model": "` + embedTestModel + `", + "input": "no-shape: hello", + "pooling": "mean" + }`) + Expect(status).To(Equal(400)) + Expect(response["__raw"]).To(ContainSubstring("did not declare")) + }) + + It("rejects Go-side pooling when the backend returns a final vector", func() { + status, response := postEmbeddings(`{ + "model": "` + embedTestModel + `", + "input": "hello", + "pooling": "mean" + }`) + Expect(status).To(Equal(400)) + Expect(response["__raw"]).To(ContainSubstring("final vector")) + }) + + It("rejects backend pooling when the backend returns per-token vectors", func() { + status, response := postEmbeddings(`{ + "model": "` + embedTestModel + `", + "input": "per-token: alpha beta", + "pooling": "backend" + }`) + Expect(status).To(Equal(400)) + Expect(response["__raw"]).To(ContainSubstring("per-token")) + }) + + 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/localai/api_instructions.go b/core/http/endpoints/localai/api_instructions.go index c8a5fb74b..d4574d2e0 100644 --- a/core/http/endpoints/localai/api_instructions.go +++ b/core/http/endpoints/localai/api_instructions.go @@ -133,7 +133,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 000000000..a5bde7899 --- /dev/null +++ b/core/http/endpoints/localai/router_corpus.go @@ -0,0 +1,168 @@ +package localai + +import ( + "errors" + "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" + "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. +// +// 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 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, 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()) + } +} + +// 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") + } + entries := make([]corpus.Entry, 0, len(req.Entries)) + for _, e := range req.Entries { + entries = append(entries, corpus.Entry{Text: e.Text, Labels: e.Labels}) + } + + added, skipped, stats, err := corpus.Seed(c.Request().Context(), mgr, cfg, storeName, + 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.ErrInvalidEntry) || 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{ + 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 + } + 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()) + } + 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 000000000..010b1f73d --- /dev/null +++ b/core/http/endpoints/localai/router_corpus_test.go @@ -0,0 +1,230 @@ +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{} }, + 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)) + 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 { + req := httptest.NewRequest(method, path, strings.NewReader(body)) + 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 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":[]}`) + 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 11fddcf57..ca24e69e6 100644 --- a/core/http/endpoints/localai/router_decide.go +++ b/core/http/endpoints/localai/router_decide.go @@ -94,16 +94,27 @@ 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, - 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, + Neighbors: neighbors, }) } } diff --git a/core/http/endpoints/mcp/localai_assistant_test.go b/core/http/endpoints/mcp/localai_assistant_test.go index 213f6fa0c..71ce38a7a 100644 --- a/core/http/endpoints/mcp/localai_assistant_test.go +++ b/core/http/endpoints/mcp/localai_assistant_test.go @@ -202,3 +202,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/embeddings.go b/core/http/endpoints/openai/embeddings.go index e6a3a353d..9026d52f3 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,49 +43,84 @@ 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 } embeddings, err := embedFn() if err != nil { + if backend.IsEmbeddingPoolingCompatibilityError(err) { + return echo.NewHTTPError(http.StatusBadRequest, err.Error()) + } return err } 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 } embeddings, err := embedFn() if err != nil { + if backend.IsEmbeddingPoolingCompatibilityError(err) { + return echo.NewHTTPError(http.StatusBadRequest, err.Error()) + } return err } items = append(items, embeddingItem(embeddings, i, input.EncodingFormat)) diff --git a/core/http/endpoints/openai/realtime_model.go b/core/http/endpoints/openai/realtime_model.go index 030a0c914..68dd61997 100644 --- a/core/http/endpoints/openai/realtime_model.go +++ b/core/http/endpoints/openai/realtime_model.go @@ -833,22 +833,13 @@ func buildRealtimeRoutingContext(a *application.Application, sessionID string) * if a == nil { return nil } - deps := &middleware.ClassifierDeps{ - Scorer: a.Scorer, - 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/request.go b/core/http/middleware/request.go index 74f7e8565..1599ef05c 100644 --- a/core/http/middleware/request.go +++ b/core/http/middleware/request.go @@ -591,6 +591,25 @@ 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 + // 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 + } + // Can be either a string or an object switch fnc := input.FunctionCall.(type) { case string: @@ -628,10 +647,17 @@ 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. + 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 fmt.Errorf("unable to validate configuration after merging") + return nil } func (re *RequestExtractor) SetOpenResponsesRequest(c echo.Context) error { @@ -753,8 +779,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/middleware/request_test.go b/core/http/middleware/request_test.go index 010379714..1b00c7f02 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()) + }) +}) diff --git a/core/http/middleware/route_model.go b/core/http/middleware/route_model.go index 643c6fa12..cac68ae9d 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" @@ -32,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. @@ -49,6 +55,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, embeddingFingerprint 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 @@ -61,10 +77,17 @@ type ModelConfigLookup func(modelName string) *config.ModelConfig // 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. + Corpus CorpusLoader // ModelLookup resolves the classifier_model name to its config so // buildClassifier can reject misconfigurations that would @@ -95,6 +118,27 @@ 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, + EmbedderFingerprint: app.EmbedderFingerprint, + 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 @@ -241,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 } @@ -262,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 @@ -291,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 @@ -312,6 +382,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") } @@ -363,6 +436,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") } @@ -375,8 +451,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") + } + 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) + } + 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) + } + if deps.Corpus != nil { + // 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) + } + } + 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) + } + 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}, ", ")) + return nil, fmt.Errorf("router: unknown classifier %q (supported: %s)", name, strings.Join(router.AllClassifiers, ", ")) } if rc.EmbeddingCache == nil { @@ -422,15 +551,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) { - if rc.ClassifierModel == "" { - 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/middleware/route_model_test.go b/core/http/middleware/route_model_test.go index 0496ed9dd..97aa58580 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,224 @@ 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") +) + +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()) + 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("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"} + 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 98e011c1e..a12733a99 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/App.css b/core/http/react-ui/src/App.css index 463ad8c84..a6c590096 100644 --- a/core/http/react-ui/src/App.css +++ b/core/http/react-ui/src/App.css @@ -11073,6 +11073,22 @@ button.collapsible-header:focus-visible { .mw-hist__bar--hit { background: var(--color-success); opacity: 1; } .mw-hist__bar--miss { background: var(--color-warning); opacity: 1; } +.mw-knn { + font-family: var(--font-mono); + font-size: 0.6875rem; + line-height: 1.3; +} +.mw-knn__model { font-weight: var(--font-weight-semibold); } +.mw-knn__meta { color: var(--color-text-muted); } +.mw-knn__labels { + display: flex; + flex-wrap: wrap; + gap: 6px; + margin-top: 2px; + color: var(--color-text-muted); +} +.mw-knn__label { color: var(--color-primary); } + .mw-alert { padding: var(--spacing-md); border-left: 3px solid var(--color-error); } .mw-list { margin: 0; padding-left: 20px; } .mw-list--mono { font-family: var(--font-mono); } diff --git a/core/http/react-ui/src/pages/Middleware.jsx b/core/http/react-ui/src/pages/Middleware.jsx index b325d1539..34d55d049 100644 --- a/core/http/react-ui/src/pages/Middleware.jsx +++ b/core/http/react-ui/src/pages/Middleware.jsx @@ -461,7 +461,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).'}
) } @@ -559,7 +561,7 @@ function RoutingTab({ status, decisions }) { Model Classifier Candidates - Embedding cache + Cache / corpus Fallback @@ -580,7 +582,7 @@ function RoutingTab({ status, decisions }) { ))} - + {m.knn ? : } {m.fallback || '—'} @@ -1040,6 +1042,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 288aa6f57..124557655 100644 --- a/core/http/routes/anthropic.go +++ b/core/http/routes/anthropic.go @@ -55,16 +55,7 @@ func RegisterAnthropicRoutes(app *echo.Echo, application.FallbackUser(), middleware.AnthropicProbe, router.SourceAnthropic, - middleware.ClassifierDeps{ - Scorer: application.Scorer, - 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 6d130863a..0195f445e 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,18 +118,31 @@ func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { decideHandler := localai.RouterDecideEndpoint( app.ModelConfigLoader(), app.ApplicationConfig(), - middleware.ClassifierDeps{ - Scorer: app.Scorer, - 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 { + 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 + // 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.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) + 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"}) @@ -153,8 +150,8 @@ func RegisterMiddlewareRoutes(e *echo.Echo, app *application.Application) { if viewer.Role != auth.RoleAdmin { return c.JSON(http.StatusForbidden, map[string]string{"error": "admin access required"}) } - return decideHandler(c) - }) + return h(c) + } } // buildRouterStatus inventories every model that declares a Router @@ -210,6 +207,25 @@ func buildRouterStatus(app *application.Application) map[string]any { } entry["embedding_cache"] = cacheEntry } + if kc := cfg.Router.KNN; kc != nil { + storeName := kc.ResolvedStoreName(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 +240,7 @@ func buildRouterStatus(app *application.Application) map[string]any { "configured": hasAny, "models": models, "recent_decision_count": recentCount, - "available_classifiers": []string{router.ClassifierScore}, + "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 7f3caf072..24ee4e1e3 100644 --- a/core/http/routes/openai.go +++ b/core/http/routes/openai.go @@ -71,16 +71,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, application.FallbackUser(), middleware.OpenAIProbe, router.SourceChat, - middleware.ClassifierDeps{ - Scorer: application.Scorer, - 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 @@ -156,7 +147,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, app.POST("/moderations", moderationHandler, moderationMiddleware...) // 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/schema/localai.go b/core/schema/localai.go index cecbc6d80..e160badad 100644 --- a/core/schema/localai.go +++ b/core/schema/localai.go @@ -51,23 +51,23 @@ 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 3D asset generation request body. Generation is image-conditioned @@ -96,11 +96,11 @@ type Model3DRemeshRequest struct { // @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 // Speed is the OpenAI `speed` field (0.25-4.0). It is a pointer so an @@ -214,10 +214,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 { @@ -289,14 +289,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" @@ -577,6 +577,80 @@ 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"` + // 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 +// 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/schema/prediction.go b/core/schema/prediction.go index d6e4fabf3..974850d0c 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"` } 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 000000000..36e2abea4 --- /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 000000000..bac177011 --- /dev/null +++ b/core/services/routing/corpus/manager.go @@ -0,0 +1,648 @@ +// 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, 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 +// 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. +package corpus + +import ( + "bufio" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "time" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/services/routing/router" + "github.com/mudler/xlog" +) + +// 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"` + 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 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"` +} + +// 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 + // 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 + // 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 storeState struct { + embeddingFingerprint string + syncedFile fileFingerprint + needsSync bool + indexedEntries int +} + +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, + states: map[string]storeState{}, + statsCache: map[string]cachedStats{}, + } +} + +// 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) + 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 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, embeddingFingerprint string, embedder backend.Embedder, store backend.VectorStore) (int, error) { + m.mu.Lock() + defer m.mu.Unlock() + + 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 + } + } + + entries, _, err := m.readAll(storeName) + if err != nil { + return 0, err + } + if len(entries) == 0 { + 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].EmbeddingFingerprint == embeddingFingerprint { + 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 + 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.states[storeName] = storeState{ + embeddingFingerprint: embeddingFingerprint, + syncedFile: fileKey, + indexedEntries: len(entries), + } + 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. +// +// 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, 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 err := validateEntry(e); err != nil { + return 0, 0, fmt.Errorf("corpus entry %d: %w", i, err) + } + } + + // First pass: dedupe against the persisted corpus (and within the + // request) so known texts aren't re-embedded. + m.mu.Lock() + existing, _, err := m.readAll(storeName) + if err != nil { + 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{}{} + } + candidates := make([]Entry, 0, len(entries)) + skipped := 0 + for _, e := range entries { + if _, dup := seen[e.Text]; dup { + skipped++ + continue + } + seen[e.Text] = struct{}{} + 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) + if err != nil { + 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 + candidates[i].EmbeddingFingerprint = embeddingFingerprint + } + + 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 + } + 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{}{} + } + added := candidates[:0] + for _, e := range candidates { + if _, dup := seen[e.Text]; dup { + skipped++ + continue + } + added = append(added, e) + } + if len(added) == 0 { + return 0, skipped, 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 + } + 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. + 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 +} + +// 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) { + 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.readAll(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) + + 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 +// 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.readAll(storeName) + if err != nil { + return 0, err + } + if err := os.Remove(m.path(storeName)); err != nil && !os.IsNotExist(err) { + return 0, err + } + delete(m.states, 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 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)) + for _, e := range entries { + payload, err := router.EncodeCorpusEntry(router.EntryID(e.Text), 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 +} + +// 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, false, nil + } + if err != nil { + return nil, false, err + } + defer func() { _ = f.Close() }() + 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 { + 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, false, err + } + 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 + 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 + } + 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.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 new file mode 100644 index 000000000..6c50c25a6 --- /dev/null +++ b/core/services/routing/corpus/manager_test.go @@ -0,0 +1,347 @@ +package corpus_test + +import ( + "context" + "errors" + "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 + failBatches int +} + +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++ + 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() + 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" + fingerprint = "embed-1-fingerprint" + ) + + 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", fingerprint, 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(onlyCorpusPath(dir)) + Expect(err).NotTo(HaveOccurred()) + }) + + It("skips duplicate texts instead of double-weighting them", func() { + _, _, err := mgr.Add(ctx, storeName, "embed-1", fingerprint, embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + 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", fingerprint, embedder, store, []corpus.Entry{{Text: " ", Labels: []string{"x"}}}) + Expect(err).To(HaveOccurred()) + _, _, 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", 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", 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", 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", 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", "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") + + // 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", "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", fingerprint, embedder, store, seed) + Expect(err).NotTo(HaveOccurred()) + _, err = mgr.EnsureLoaded(ctx, storeName, "embed-1", fingerprint, embedder, store) + Expect(err).NotTo(HaveOccurred()) + + _, 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", fingerprint, 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", fingerprint, 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", 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", fingerprint, 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("tolerates a torn final line and repairs it on the next add", func() { + _, _, 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 := onlyCorpusPath(dir) + 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", fingerprint, 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", 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 := onlyCorpusPath(dir) + 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", 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 new file mode 100644 index 000000000..15763b488 --- /dev/null +++ b/core/services/routing/corpus/resolve.go @@ -0,0 +1,130 @@ +package corpus + +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 +// 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 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 +// 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 +// 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) + } + 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 +} + +// 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, fingerprintFor func(string) (string, error), 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 { + 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) + } + } + } + + 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) + } + 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) + } + + 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 + } + stats, err = mgr.Stats(storeName) + return added, skipped, stats, err +} diff --git a/core/services/routing/router/decisions.go b/core/services/routing/router/decisions.go index d446ac29a..492da3af8 100644 --- a/core/services/routing/router/decisions.go +++ b/core/services/routing/router/decisions.go @@ -11,29 +11,35 @@ 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). 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 // 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 e36b049c3..41be408a0 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 000000000..529051e6f --- /dev/null +++ b/core/services/routing/router/knn.go @@ -0,0 +1,259 @@ +package router + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "sort" + "strings" + "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)) + } + + // 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} + entry, ok := decodeCorpusEntry(n.Payload) + if ok { + ref.ID = entry.ID + ref.Labels = entry.Labels + } + refs = append(refs, ref) + if n.Similarity > best { + best = n.Similarity + } + // 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 + for _, l := range entry.Labels { + votes[l] += n.Similarity + } + } + 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, + 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, + Neighbors: refs, + 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 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"` +} + +// 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 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 || 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 new file mode 100644 index 000000000..9ec0bf5d9 --- /dev/null +++ b/core/services/routing/router/knn_test.go @@ -0,0 +1,219 @@ +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(id string, labels ...string) []byte { + b, err := router.EncodeCorpusEntry(id, 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("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"})) + 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)) + // 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("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"})) + 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("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() { + // 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("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"})) + 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("e1", "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)) + // 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() { + 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("e1", nil) + 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()) + }) + + 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 a474d6d48..6950bf8ca 100644 --- a/core/services/routing/router/resolve.go +++ b/core/services/routing/router/resolve.go @@ -165,8 +165,10 @@ 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, + 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 05efdc349..02cc67256 100644 --- a/core/services/routing/router/types.go +++ b/core/services/routing/router/types.go @@ -71,6 +71,34 @@ 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"` + + // 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 @@ -127,8 +155,22 @@ 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" ) +// 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/core/templates/evaluator.go b/core/templates/evaluator.go index 8ee74e43d..36e491b7d 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 000000000..186640be2 --- /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 a7f948d90..8c5dc8fc3 100644 --- a/docs/content/features/embeddings.md +++ b/docs/content/features/embeddings.md @@ -96,6 +96,77 @@ 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. Each backend +declares whether an embedding result is final or per-token; LocalAI rejects a +Go-side scheme for a final vector and rejects `backend` pass-through for a +per-token matrix instead of guessing from its shape. Older backends that do +not declare a layout remain compatible with `backend` pooling only. + +llama.cpp chooses this layout when the model is loaded. LocalAI automatically +adds the `pooling:none` backend option when a llama.cpp model sets a Go-side +`parameters.pooling` scheme. That raw-loaded instance can switch between +`mean`, `last`, and `decayed_mean` per request, but it cannot switch back to +`backend` pooling without reloading. Conversely, a backend-pooled llama.cpp +instance rejects per-request Go pooling. Other backends may support Go-side +pooling when they explicitly return per-token vectors. + +After Go-side 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 backend that reports its embedding +layout. A legacy backend fails closed for Go-side schemes with an error asking +you to rebuild or update it. + ## 💡 Examples - Example that uses LLamaIndex and LocalAI as embedding: [here](https://github.com/mudler/LocalAI-examples/tree/main/query_data). diff --git a/docs/content/operations/middleware.md b/docs/content/operations/middleware.md index 6f61ac548..fc2f2a3c9 100644 --- a/docs/content/operations/middleware.md +++ b/docs/content/operations/middleware.md @@ -353,17 +353,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 @@ -444,6 +447,143 @@ 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 + # 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) + # 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". + +The numeric configuration is bounded so one request cannot allocate an +unbounded neighbour result and the cosine-weighted vote remains meaningful: + +| Field | Accepted values | +|-------|-----------------| +| `k` | `0` for the default (`3`), otherwise `1` through `1024` | +| `similarity_threshold` | `0` for the default (`0.80`), otherwise greater than `0` and at most `1` | +| `vote_threshold` | `0` for the default (`0.5`), otherwise greater than `0` and at most `1` | + +Negative, non-finite, and above-range values are rejected when the model +configuration is loaded. The `k` cap bounds the store priority queue and +returned neighbour arrays. Similarity is restricted to non-negative cosine +weights because negative weights would make vote totals and shares invalid; +vote share itself is necessarily between zero and one. + +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 +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). 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 +candidates and seed both. + +#### Persistence + +The corpus is persisted as one JSONL file per router under +`/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 + +- **`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. The maximum is `1024`. +- **`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 @@ -451,8 +591,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 @@ -571,9 +712,12 @@ For each request: paraphrases. 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. +default — each router gets its own collection so two routers can't +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 @@ -614,7 +758,10 @@ 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/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`. | +| 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 @@ -622,10 +769,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. | +| `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 routing config - adding a candidate, changing the classifier -model - is YAML-only today; reload with `POST /models/reload` to pick -up edits without restarting. +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/grpc/server.go b/pkg/grpc/server.go index 528ba0a36..679777d6c 100644 --- a/pkg/grpc/server.go +++ b/pkg/grpc/server.go @@ -96,7 +96,10 @@ func (s *server) Embedding(ctx context.Context, in *pb.PredictOptions) (*pb.Embe return nil, err } - return &pb.EmbeddingResult{Embeddings: embeds}, nil + return &pb.EmbeddingResult{ + Embeddings: embeds, + Layout: pb.EmbeddingLayout_EMBEDDING_LAYOUT_FINAL, + }, nil } func (s *server) LoadModel(ctx context.Context, in *pb.ModelOptions) (*pb.Result, error) { diff --git a/pkg/mcp/localaitools/client.go b/pkg/mcp/localaitools/client.go index e8b7b2ef5..bf683a508 100644 --- a/pkg/mcp/localaitools/client.go +++ b/pkg/mcp/localaitools/client.go @@ -118,4 +118,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 dbe132ba1..a802226f0 100644 --- a/pkg/mcp/localaitools/coverage_test.go +++ b/pkg/mcp/localaitools/coverage_test.go @@ -26,25 +26,26 @@ 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", - ToolListScheduling: "GET /api/nodes/scheduling", - ToolGetScheduling: "GET /api/nodes/scheduling/:model", - 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", + ToolListScheduling: "GET /api/nodes/scheduling", + ToolGetScheduling: "GET /api/nodes/scheduling/:model", + 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", @@ -59,6 +60,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 51938b786..a8e734c84 100644 --- a/pkg/mcp/localaitools/dto.go +++ b/pkg/mcp/localaitools/dto.go @@ -347,6 +347,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 932f429a8..b0697b272 100644 --- a/pkg/mcp/localaitools/fakes_test.go +++ b/pkg/mcp/localaitools/fakes_test.go @@ -378,3 +378,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 550c1ef01..923f35eba 100644 --- a/pkg/mcp/localaitools/httpapi/client.go +++ b/pkg/mcp/localaitools/httpapi/client.go @@ -798,3 +798,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 759cd5bba..1d80c30df 100644 --- a/pkg/mcp/localaitools/inproc/client.go +++ b/pkg/mcp/localaitools/inproc/client.go @@ -26,6 +26,7 @@ import ( "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "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" @@ -71,6 +72,17 @@ 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 + RouterEmbedderFingerprint func(modelName string) (string, error) + RouterVectorStore func(storeName string) backend.VectorStore + modelAdmin *modeladmin.ConfigService } @@ -1008,3 +1020,70 @@ func capabilityFlagsOf(m *config.ModelConfig) []string { } return out } + +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 := corpus.ResolveKNNRouter(c.ConfigLoader, c.AppConfig, 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.RouterEmbedderFingerprint == nil || c.RouterVectorStore == nil { + return nil, errors.New("router corpus manager unavailable") + } + cfg, storeName, err := corpus.ResolveKNNRouter(c.ConfigLoader, c.AppConfig, req.Router) + if err != nil { + return nil, err + } + entries := make([]corpus.Entry, 0, len(req.Entries)) + for _, e := range req.Entries { + 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.RouterEmbedderFingerprint, c.RouterVectorStore, entries) + 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 := corpus.ResolveKNNRouter(c.ConfigLoader, c.AppConfig, 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/prompts/10_safety.md b/pkg/mcp/localaitools/prompts/10_safety.md index bd5941808..15671fdc4 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`, `set_node_vram_budget`, `set_scheduling`, `delete_scheduling` — 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`, `set_scheduling`, `delete_scheduling` — 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 0a2d7dddc..2a2830417 100644 --- a/pkg/mcp/localaitools/prompts/20_tools.md +++ b/pkg/mcp/localaitools/prompts/20_tools.md @@ -17,6 +17,13 @@ The MCP `tools/list` endpoint also exposes the full input schema for each of the - `list_scheduling` — List distributed per-model scheduling configs. - `get_scheduling` — Read the distributed scheduling config for one model. - `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) @@ -30,6 +37,10 @@ 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. - `set_node_vram_budget` — Set or clear a federated node's VRAM budget override. 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 000000000..86f98d59e --- /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 c85efdce3..1f0c762fd 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 ce1ede915..c18a2301c 100644 --- a/pkg/mcp/localaitools/server_test.go +++ b/pkg/mcp/localaitools/server_test.go @@ -70,47 +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, - ToolGetRouterDecisions, - ToolGetUsageStats, - ToolImportModelURI, - ToolInstallBackend, - ToolInstallModel, - ToolListBackends, - ToolListGalleries, - ToolListAliases, - ToolListInstalledModels, - ToolListKnownBackends, - ToolListNodes, - ToolListScheduling, - ToolListVoiceProfiles, - ToolLoadModel, - ToolReloadModels, - ToolSetAlias, - ToolSetBranding, - ToolSystemInfo, - ToolToggleModelPinned, - ToolToggleModelState, - ToolUpgradeBackend, - ToolVRAMEstimate, - ToolCreateVoiceProfile, - ToolDeleteVoiceProfile, - ToolSetNodeVRAMBudget, - ToolSetScheduling, - ToolDeleteScheduling, - ToolGetScheduling, -) - // expectedReadOnlyCatalog is the tool set when DisableMutating=true. Sorted. var expectedReadOnlyCatalog = sortedStrings( ToolGallerySearch, @@ -119,11 +78,12 @@ var expectedReadOnlyCatalog = sortedStrings( ToolGetMiddlewareStatus, ToolGetModelConfig, ToolGetPIIEvents, + ToolGetRouterCorpusStats, ToolGetRouterDecisions, ToolGetUsageStats, - ToolListAliases, ToolListBackends, ToolListGalleries, + ToolListAliases, ToolListInstalledModels, ToolListKnownBackends, ToolListNodes, @@ -134,6 +94,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 89d99432b..cfc7f0bfd 100644 --- a/pkg/mcp/localaitools/tools.go +++ b/pkg/mcp/localaitools/tools.go @@ -4,28 +4,29 @@ 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" - 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" - ToolListScheduling = "list_scheduling" - ToolGetScheduling = "get_scheduling" - 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" + ToolListScheduling = "list_scheduling" + ToolGetScheduling = "get_scheduling" + 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). @@ -41,6 +42,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" @@ -56,3 +59,28 @@ 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, + ToolSetScheduling, + ToolDeleteScheduling, +} diff --git a/pkg/mcp/localaitools/tools_middleware.go b/pkg/mcp/localaitools/tools_middleware.go index 5dd8066fd..528eb5811 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 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 + } + 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 d785faa68..ecfe71cc2 100644 --- a/swagger/docs.go +++ b/swagger/docs.go @@ -1530,6 +1530,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 a bounded, newest-first page of captured API exchange traces. Request and response bodies plus headers are omitted unless full=true; fetch them per-trace from /api/traces/{id}. Paging metadata is returned in the X-Total-Count, X-Trace-Offset and X-Trace-Limit headers.", @@ -7200,6 +7375,116 @@ 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.RouterDecideNeighbor": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "labels": { + "type": "array", + "items": { + "type": "string" + } + }, + "similarity": { + "type": "number" + } + } + }, "schema.RouterDecideRequest": { "type": "object", "properties": { @@ -7247,6 +7532,17 @@ 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" + }, + "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 8b7f4803a..5f20e5955 100644 --- a/swagger/swagger.json +++ b/swagger/swagger.json @@ -1527,6 +1527,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 a bounded, newest-first page of captured API exchange traces. Request and response bodies plus headers are omitted unless full=true; fetch them per-trace from /api/traces/{id}. Paging metadata is returned in the X-Total-Count, X-Trace-Offset and X-Trace-Limit headers.", @@ -7197,6 +7372,116 @@ } } }, + "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.RouterDecideNeighbor": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "labels": { + "type": "array", + "items": { + "type": "string" + } + }, + "similarity": { + "type": "number" + } + } + }, "schema.RouterDecideRequest": { "type": "object", "properties": { @@ -7244,6 +7529,17 @@ "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" + }, + "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 f944d88f8..36cc627f5 100644 --- a/swagger/swagger.yaml +++ b/swagger/swagger.yaml @@ -2351,6 +2351,84 @@ 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.RouterDecideNeighbor: + properties: + id: + type: string + labels: + items: + type: string + type: array + similarity: + type: number + type: object schema.RouterDecideRequest: properties: input: @@ -2402,6 +2480,21 @@ 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 + 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 @@ -3859,6 +3952,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: diff --git a/tests/e2e-backends/backend_test.go b/tests/e2e-backends/backend_test.go index 7ac3ad410..a86a2a088 100644 --- a/tests/e2e-backends/backend_test.go +++ b/tests/e2e-backends/backend_test.go @@ -71,6 +71,9 @@ import ( // BACKEND_TEST_THREADS Override Threads passed to LoadModel (default 4). // BACKEND_TEST_OPTIONS Comma-separated Options[] entries passed to LoadModel, // e.g. "tool_parser:hermes,reasoning_parser:qwen3". +// BACKEND_TEST_EMBEDDING_LAYOUT Expected EmbeddingResult layout: "final" or +// "per_token". When set, the embeddings spec also +// validates tokens/dim against the returned payload. // BACKEND_TEST_CACHE_TYPE_K Sets ModelOptions.CacheTypeKey (llama.cpp -ctk), // e.g. "q8_0" — exercises KV-cache quantization code paths. // BACKEND_TEST_CACHE_TYPE_V Sets ModelOptions.CacheTypeValue (llama.cpp -ctv). @@ -401,6 +404,7 @@ var _ = Describe("Backend container", Ordered, func() { NGPULayers: 0, MMap: true, NBatch: 128, + Embeddings: caps[capEmbeddings], Options: options, MMProj: mmprojFile, CacheTypeKey: os.Getenv("BACKEND_TEST_CACHE_TYPE_K"), @@ -549,7 +553,28 @@ var _ = Describe("Backend container", Ordered, func() { }) Expect(err).NotTo(HaveOccurred()) Expect(res.GetEmbeddings()).NotTo(BeEmpty(), "Embedding returned empty vector") - GinkgoWriter.Printf("Embedding: %d dims\n", len(res.GetEmbeddings())) + + expectedLayout := strings.ToLower(strings.TrimSpace(os.Getenv("BACKEND_TEST_EMBEDDING_LAYOUT"))) + if expectedLayout != "" { + var layout pb.EmbeddingLayout + switch expectedLayout { + case "final": + layout = pb.EmbeddingLayout_EMBEDDING_LAYOUT_FINAL + Expect(res.GetTokens()).To(Equal(int32(1)), "a final embedding must contain one vector") + case "per_token": + layout = pb.EmbeddingLayout_EMBEDDING_LAYOUT_PER_TOKEN + Expect(res.GetTokens()).To(BeNumerically(">", 1), "the test prompt should produce multiple token vectors") + default: + Fail(fmt.Sprintf("unsupported BACKEND_TEST_EMBEDDING_LAYOUT %q", expectedLayout)) + } + + Expect(res.GetLayout()).To(Equal(layout)) + Expect(res.GetDim()).To(BeNumerically(">", 0)) + Expect(res.GetEmbeddings()).To(HaveLen(int(res.GetTokens() * res.GetDim()))) + } + + GinkgoWriter.Printf("Embedding: layout=%s vectors=%d dim=%d\n", + res.GetLayout(), res.GetTokens(), res.GetDim()) }) // TokenClassify is the PII-NER RPC (privacy-filter backend). The crown-jewel diff --git a/tests/e2e/mock-backend/main.go b/tests/e2e/mock-backend/main.go index 290205d5e..578482d0c 100644 --- a/tests/e2e/mock-backend/main.go +++ b/tests/e2e/mock-backend/main.go @@ -420,13 +420,57 @@ func (m *MockBackend) Embedding(ctx context.Context, in *pb.PredictOptions) (*pb if err := checkModelIdentity(in); err != nil { return nil, err } - 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(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) + 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), + Layout: pb.EmbeddingLayout_EMBEDDING_LAYOUT_PER_TOKEN, + }, nil + } + // Legacy mode: a prompt carrying "no-shape:" simulates a backend built + // before EmbeddingResult carried shape and layout fields, so tests can + // assert the fail-closed error when Go-side pooling is requested. + legacyShape := strings.Contains(text, "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, + Layout: pb.EmbeddingLayout_EMBEDDING_LAYOUT_FINAL, + }, nil } func (m *MockBackend) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest) (*pb.Result, error) { diff --git a/tests/e2e/mock_backend_test.go b/tests/e2e/mock_backend_test.go index e95dd875e..b12dff4d9 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() { diff --git a/tests/integration/stores_test.go b/tests/integration/stores_test.go index 849b4b944..ad9b51b7a 100644 --- a/tests/integration/stores_test.go +++ b/tests/integration/stores_test.go @@ -176,6 +176,23 @@ var _ = Describe("Integration tests for the stores backend(s) and internal APIs" Expect(vals).To(HaveLen(0)) }) + It("should accept a new embedding dimension after deleting every key", func() { + ctx := context.Background() + oldKey := []float32{0.1, 0.2, 0.3} + newKey := []float32{0.6, 0.8} + + Expect(store.SetSingle(ctx, sc, oldKey, []byte("old"))).To(Succeed()) + Expect(store.DeleteSingle(ctx, sc, oldKey)).To(Succeed()) + Expect(store.SetSingle(ctx, sc, newKey, []byte("new"))).To(Succeed()) + + keys, vals, sims, err := store.Find(ctx, sc, newKey, 1) + Expect(err).ToNot(HaveOccurred()) + Expect(keys).To(Equal([][]float32{newKey})) + Expect(vals).To(Equal([][]byte{[]byte("new")})) + Expect(sims).To(HaveLen(1)) + Expect(sims[0]).To(BeNumerically("~", 1, 0.0001)) + }) + It("should be able to find smilar keys", func() { // set 3 vectors that are at varying angles to {0.5, 0.5, 0.5} err := store.SetCols(context.Background(), sc, [][]float32{{0.5, 0.5, 0.5}, {0.6, 0.6, -0.6}, {0.7, -0.7, -0.7}}, [][]byte{[]byte("test1"), []byte("test2"), []byte("test3")})