feat(router): make KNN a first-class classifier with a persisted, curated corpus

backend/backend.proto

@@ -472,6 +472,15 @@ message Result {
 
 message EmbeddingResult {
   repeated float embeddings = 1;
+  // Shape of the payload above: dim is the embedding width, tokens is the
+  // number of vectors packed into `embeddings` (1 when the backend pooled
+  // server-side, N with pooling:none; total across prompts if a request
+  // carried several). tokens=0/dim=0 means the backend predates shape
+  // reporting. prompt_tokens is the number of prompt tokens evaluated, for
+  // usage accounting.
+  int32 tokens = 2;
+  int32 dim = 3;
+  int32 prompt_tokens = 4;
 }
 
 message TranscriptRequest {

backend/cpp/llama-cpp/grpc-server.cpp

@@ -2772,42 +2772,36 @@ class BackendServiceImpl final : public backend::Backend::Service {
             return grpc::Status(grpc::StatusCode::INTERNAL, all_results.error->to_json().value("message", "Error in receiving results"));
         }
 
-        // Collect responses
-        json responses = json::array();
+        // Extract the embeddings typed, straight from the task results (no
+        // JSON round-trip), and report the payload shape alongside the same
+        // flat float array as before: dim is the embedding width, tokens the
+        // number of vectors packed into `embeddings` (1 per prompt when the
+        // server pooled, one per token with pooling:none; summed across
+        // prompts if the request carried several), prompt_tokens the prompt
+        // tokens evaluated, for usage accounting. Consumers seeing 0/0 know
+        // the backend predates shape reporting.
+        int32_t n_vectors = 0;
+        int32_t dim = 0;
+        int32_t prompt_tokens = 0;
         for (auto & res : all_results.results) {
-            GGML_ASSERT(dynamic_cast<server_task_result_embd*>(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<float>());
-                            }
-                        }
-                    }
+            auto * embd_res = dynamic_cast<server_task_result_embd*>(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<float>());
-                    }
+                if (!vec.empty()) {
+                    n_vectors++;
+                    dim = (int32_t) vec.size();
                 }
             }
         }
+        embeddingResult->set_tokens(n_vectors);
+        embeddingResult->set_dim(dim);
+        embeddingResult->set_prompt_tokens(prompt_tokens);
 
-
-
+        std::cout << "[DEBUG] Embedding vectors: " << n_vectors << " x " << dim << std::endl;
 
         return grpc::Status::OK;
     }

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"
@@ -76,6 +79,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
@@ -125,6 +129,10 @@ func newApplication(appConfig *config.ApplicationConfig) *Application {
 		applicationConfig:  appConfig,
 		templatesEvaluator: templates.NewEvaluator(appConfig.SystemState.Model.ModelsPath),
 		voiceProfileStore:  voiceprofile.NewStore(appConfig.DataPath),
+		// KNN corpus files live under <state dir>/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.
@@ -549,6 +557,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.

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, storeName)
 }
+
+// RouterCorpus returns the process-wide KNN corpus manager, built in
+// newApplication.
+func (a *Application) RouterCorpus() *corpus.Manager {
+	return a.routerCorpus
+}

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())

core/backend/embeddings.go

@@ -9,9 +9,44 @@ import (
 	"github.com/mudler/LocalAI/core/trace"
 
 	"github.com/mudler/LocalAI/pkg/grpc"
+	"github.com/mudler/LocalAI/pkg/grpc/proto"
 	model "github.com/mudler/LocalAI/pkg/model"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
+// finishEmbeddingResult applies the model's Go-side pooling scheme to a raw
+// per-token EmbeddingResult, or passes the backend's vector through
+// untouched when pooling is delegated to the backend ("" / "backend" —
+// today's exact behavior). Fails closed when a Go-side scheme is requested
+// but the backend didn't report the payload shape (a build predating
+// EmbeddingResult.tokens/dim), since silently treating a backend-pooled
+// vector as per-token data would corrupt the embedding space.
+func finishEmbeddingResult(res *proto.EmbeddingResult, modelConfig config.ModelConfig) ([]float32, error) {
+	scheme := modelConfig.Pooling
+	if scheme == "" || scheme == PoolingBackend {
+		return res.Embeddings, nil
+	}
+	if res.GetDim() == 0 {
+		return nil, fmt.Errorf(
+			"pooling %q needs per-token embeddings but the backend returned no shape: rebuild/update the backend so EmbeddingResult reports tokens/dim, and set options [\"pooling:none\"] on the model",
+			scheme)
+	}
+	return PoolEmbeddingResult(res, scheme,
+		float64(modelConfig.PoolingHalfLifeTokens),
+		embdNormalizeFromOptions(modelConfig.Options))
+}
+
+// mapEmbeddingGRPCError turns a gRPC ResourceExhausted — the per-token
+// payload of a very long conversation exceeding the 50MB message cap —
+// into an actionable message; everything else passes through unchanged.
+func mapEmbeddingGRPCError(err error) error {
+	if status.Code(err) == codes.ResourceExhausted {
+		return fmt.Errorf("conversation too long for per-token embeddings (gRPC message limit exceeded): %w", err)
+	}
+	return err
+}
+
 // Embedder produces a fixed-dimension vector from a prompt. The
 // router's L2 embedding cache uses it to look up semantically-similar
 // past decisions.
@@ -66,19 +101,19 @@ func ModelEmbedding(ctx context.Context, s string, tokens []int, loader *model.M
 
 				res, err := model.Embeddings(appConfig.Context, predictOptions)
 				if err != nil {
-					return nil, err
+					return nil, mapEmbeddingGRPCError(err)
 				}
 
-				return res.Embeddings, nil
+				return finishEmbeddingResult(res, modelConfig)
 			}
 			predictOptions.Embeddings = s
 
 			res, err := model.Embeddings(appConfig.Context, predictOptions)
 			if err != nil {
-				return nil, err
+				return nil, mapEmbeddingGRPCError(err)
 			}
 
-			return res.Embeddings, nil
+			return finishEmbeddingResult(res, modelConfig)
 		}
 	default:
 		fn = func() ([]float32, error) {