fix(router): production-ready request router + auto-size batch for embedding/rerank (#10104)

* fix(router): score classifier production-readiness

Conversation trimming runs through the classifier model's chat template
and trims by exact token count, sized to the model's n_batch which is
now scaled to context so long probes can't crash the backend. Missing
chat_message templates are a hard error at router build time. Router-
facing factories (Embedder/Scorer/Reranker/TokenCounter) re-resolve
ModelConfig per call so a model installed post-startup doesn't bind a
stub Backend="" config and silently fall into the loader's auto-
iterate path.

New 'vector_store' backend trace recorded inside localVectorStore on
every Search/Insert — including the backend-load-failure path that
previously vanished into an xlog.Warn — with outcome tagging
(hit/miss/empty_store/backend_load_error/find_error/insert_error/ok).
Companion cleanup drops misleading similarity:0 and input_tokens_count:0
from non-hit and text-mode traces.

Gallery local-store-development aliases to 'local-store' so the master
image satisfies pkg/model.LocalStoreBackend lookups from the embedding
cache.

Misc: llama-cpp TokenizeString reads the correct 'prompt' JSON key
(the original bug); ModelTokenize nil-guard; non-fatal mitm proxy
startup; PII 'route_local' renamed to 'allow' with docs/UI in sync;
model-editor footer no longer eats the edit area on small screens;
several config-editor template/dropdown/section fixes.

Tests: e2e router specs (casual/code-hint + long-conversation trim),
vector_store trace specs, lazy-factory specs, gallery dev-alias
resolution, Playwright trace badge + scroll regression.

Assisted-by: Claude:claude-opus-4-7 [Claude Code]
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* feat(backend): auto-size batch to context for embedding and rerank models

Embedding and rerank models pool over the whole input in a single physical batch (n_ubatch). With batch left at the 512 default, the backend rejects longer inputs with "input is too large to process", silently capping a large-context embedder (e.g. 8k/32k) at 512 tokens. Size n_batch to the context for these single-pass usecases, mirroring the existing FLAG_SCORE behaviour; an explicit batch: still wins.

Extracts EffectiveContextSize/EffectiveBatchSize from grpcModelOpts so the effective decode window has one home for other callers to reuse.

Adds an e2e-aio regression test that embeds a >512-token input. The AIO embedding model is switched to nomic-embed-text-v1.5 (2048 context) because the previous granite model was capped at 512 tokens and could not exercise the larger batch.

Assisted-by: claude-code:claude-opus-4-8 [Claude Code]
Signed-off-by: Richard Palethorpe <io@richiejp.com>

* fix(gallery): raise arch-router scoring output cap via parallel:64

Scoring decodes the whole prompt+candidate in a single llama_decode and
reads one logit row per candidate token. The vendored llama.cpp server
caps causal output rows at n_parallel, so the default of 1 aborts with
GGML_ASSERT(n_outputs_max <= cparams.n_outputs_max) on multi-token route
labels. Set options: [parallel:64] on both arch-router quant entries to
lift the cap; kv_unified (the grpc-server default) keeps the full context
per sequence, so this does not split the KV cache.

Assisted-by: claude-code:claude-opus-4-8 [Claude Code]
Signed-off-by: Richard Palethorpe <io@richiejp.com>

---------

Signed-off-by: Richard Palethorpe <io@richiejp.com>

core/backend/embeddings.go

@@ -100,8 +100,13 @@ func ModelEmbedding(ctx context.Context, s string, tokens []int, loader *model.M
 		trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes)
 
 		traceData := map[string]any{
-			"input_text":         trace.TruncateString(s, 1000),
-			"input_tokens_count": len(tokens),
+			"input_text": trace.TruncateString(s, 1000),
+		}
+		// Only present for token-mode callers (pre-tokenized override);
+		// emitting "0" alongside input_text would read as "consumed zero
+		// tokens", which is wrong.
+		if len(tokens) > 0 {
+			traceData["input_tokens_count"] = len(tokens)
 		}
 
 		startTime := time.Now()

core/backend/options.go

@@ -87,11 +87,47 @@ func getSeed(c config.ModelConfig) int32 {
 	return seed
 }
 
-func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {
-	b := 512
-	if c.Batch != 0 {
-		b = c.Batch
+// DefaultContextSize and DefaultBatchSize are the backend's fallbacks when a
+// model config leaves them unset. Exported so callers that must respect the
+// effective decode window — notably the router's prompt trimmer — resolve the
+// same numbers grpcModelOpts does instead of guessing.
+const (
+	DefaultContextSize = 4096
+	DefaultBatchSize   = 512
+)
+
+// EffectiveContextSize is the context window the backend will run with: the
+// configured value, or DefaultContextSize when unset.
+func EffectiveContextSize(c config.ModelConfig) int {
+	if c.ContextSize != nil {
+		return *c.ContextSize
 	}
+	return DefaultContextSize
+}
+
+// EffectiveBatchSize is the single-decode batch the backend will run with.
+// Score, embedding and rerank all process the whole input in one pass: score
+// decodes prompt+candidate (asserts n_tokens <= n_batch), and embedding/rerank
+// pool over the full sequence in one physical batch (n_ubatch). So the batch
+// is sized to the context — anything that fits the context fits one pass,
+// avoiding both the GGML_ASSERT crash and the "input is too large to process"
+// error. Explicit `batch:` always wins.
+func EffectiveBatchSize(c config.ModelConfig) int {
+	if c.Batch != 0 {
+		return c.Batch
+	}
+	singlePass := c.HasUsecases(config.FLAG_SCORE) ||
+		c.HasUsecases(config.FLAG_EMBEDDINGS) ||
+		c.HasUsecases(config.FLAG_RERANK)
+	if ctx := EffectiveContextSize(c); singlePass && ctx > DefaultBatchSize {
+		return ctx
+	}
+	return DefaultBatchSize
+}
+
+func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {
+	ctxSize := EffectiveContextSize(c)
+	b := EffectiveBatchSize(c)
 
 	flashAttention := "auto"
 
@@ -134,11 +170,6 @@ func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {
 		}
 	}
 
-	ctxSize := 4096
-	if c.ContextSize != nil {
-		ctxSize = *c.ContextSize
-	}
-
 	mmlock := false
 	if c.MMlock != nil {
 		mmlock = *c.MMlock

core/backend/options_internal_test.go

@@ -97,3 +97,67 @@ var _ = Describe("gRPCPredictOpts reasoning_effort metadata", func() {
 		Expect(opts.Metadata).ToNot(HaveKey("reasoning_effort"))
 	})
 })
+
+var _ = Describe("grpcModelOpts NBatch", func() {
+	scoreUsecase := config.FLAG_SCORE
+	threads := 1
+	ctx := 4096
+
+	It("defaults to 512 for an ordinary model", func() {
+		cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}}
+		opts := grpcModelOpts(cfg, "/tmp/models")
+		Expect(opts.NBatch).To(BeEquivalentTo(512))
+	})
+
+	It("sizes the batch to the context window for score models", func() {
+		// Score models decode the whole prompt+candidate in one
+		// llama_decode; n_batch must cover it or the backend aborts.
+		cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}, KnownUsecases: &scoreUsecase}
+		opts := grpcModelOpts(cfg, "/tmp/models")
+		Expect(opts.NBatch).To(BeEquivalentTo(4096))
+	})
+
+	It("keeps an explicit batch over the score default", func() {
+		cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}, KnownUsecases: &scoreUsecase}
+		cfg.Batch = 1024
+		opts := grpcModelOpts(cfg, "/tmp/models")
+		Expect(opts.NBatch).To(BeEquivalentTo(1024))
+	})
+
+	It("sizes the batch to the context window for embedding models", func() {
+		// Embedding/rerank pool over the whole sequence in one physical batch
+		// (n_ubatch); without this the input is capped at the 512 default and
+		// the backend returns "input is too large to process".
+		embeddings := true
+		cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}}
+		cfg.Embeddings = &embeddings
+		opts := grpcModelOpts(cfg, "/tmp/models")
+		Expect(opts.NBatch).To(BeEquivalentTo(4096))
+	})
+
+	It("sizes the batch to the context window for rerank models", func() {
+		reranking := true
+		cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}}
+		cfg.Reranking = &reranking
+		opts := grpcModelOpts(cfg, "/tmp/models")
+		Expect(opts.NBatch).To(BeEquivalentTo(4096))
+	})
+
+	It("does not raise the batch when a score model's context is below the default", func() {
+		small := 256
+		cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &small}, KnownUsecases: &scoreUsecase}
+		opts := grpcModelOpts(cfg, "/tmp/models")
+		Expect(opts.NBatch).To(BeEquivalentTo(512))
+	})
+
+	It("sizes the batch to the effective 4096 default for a score model with no explicit context_size", func() {
+		// The crash case: the backend defaults n_ctx to 4096, so n_batch must
+		// follow even when context_size is unset — otherwise n_batch stays 512
+		// against a 4096 window and the score decode hits the GGML_ASSERT.
+		cfg := config.ModelConfig{Threads: &threads, KnownUsecases: &scoreUsecase}
+		Expect(cfg.ContextSize).To(BeNil())
+		opts := grpcModelOpts(cfg, "/tmp/models")
+		Expect(opts.NBatch).To(BeEquivalentTo(4096))
+		Expect(opts.ContextSize).To(BeEquivalentTo(4096), "n_batch must match the effective n_ctx the backend receives")
+	})
+})

core/backend/stores.go

@@ -3,9 +3,10 @@ package backend
 import (
 	"context"
 	"fmt"
-	"strings"
+	"time"
 
 	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/trace"
 
 	"github.com/mudler/LocalAI/pkg/grpc"
 	"github.com/mudler/LocalAI/pkg/model"
@@ -39,34 +40,85 @@ func (s *localVectorStore) backend(_ context.Context) (grpc.Backend, error) {
 	return StoreBackend(s.loader, s.appConfig, s.storeName, "")
 }
 
-func (s *localVectorStore) Search(ctx context.Context, vec []float32) (float64, []byte, bool, error) {
-	be, err := s.backend(ctx)
-	if err != nil {
-		return 0, nil, false, fmt.Errorf("vector store load: %w", err)
+func (s *localVectorStore) Search(ctx context.Context, vec []float32) (sim float64, payload []byte, ok bool, err error) {
+	start := time.Now()
+	outcome := "hit"
+	defer func() {
+		s.recordTrace(start, "search", len(vec), sim, outcome, err)
+	}()
+	be, berr := s.backend(ctx)
+	if berr != nil {
+		outcome = "backend_load_error"
+		return 0, nil, false, fmt.Errorf("vector store load: %w", berr)
 	}
-	_, values, similarities, err := store.Find(ctx, be, vec, 1)
-	if err != nil {
-		// local-store's Find returns "existing length is -1" before
-		// any keys are inserted. Surface that as a clean miss so the
-		// cache layer treats it as an empty store and proceeds to
-		// Insert rather than skipping.
-		if strings.Contains(err.Error(), "existing length is -1") {
-			return 0, nil, false, nil
-		}
-		return 0, nil, false, fmt.Errorf("vector store find: %w", err)
+	_, values, similarities, ferr := store.Find(ctx, be, vec, 1)
+	if ferr != nil {
+		outcome = "find_error"
+		return 0, nil, false, fmt.Errorf("vector store find: %w", ferr)
 	}
 	if len(values) == 0 || len(similarities) == 0 {
+		outcome = "miss"
 		return 0, nil, false, nil
 	}
 	return float64(similarities[0]), values[0], true, nil
 }
 
-func (s *localVectorStore) Insert(ctx context.Context, vec []float32, payload []byte) error {
-	be, err := s.backend(ctx)
-	if err != nil {
-		return fmt.Errorf("vector store load: %w", err)
+func (s *localVectorStore) Insert(ctx context.Context, vec []float32, payload []byte) (err error) {
+	start := time.Now()
+	outcome := "ok"
+	defer func() {
+		s.recordTrace(start, "insert", len(vec), 0, outcome, err)
+	}()
+	be, berr := s.backend(ctx)
+	if berr != nil {
+		outcome = "backend_load_error"
+		return fmt.Errorf("vector store load: %w", berr)
 	}
-	return store.SetSingle(ctx, be, vec, payload)
+	if serr := store.SetSingle(ctx, be, vec, payload); serr != nil {
+		outcome = "insert_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
+// admins can tell which router's cache misbehaved; the backend is always
+// "local-store" and can't disambiguate.
+func (s *localVectorStore) recordTrace(start time.Time, op string, vecDim int, sim float64, outcome string, err error) {
+	if s.appConfig == nil || !s.appConfig.EnableTracing {
+		return
+	}
+	trace.InitBackendTracingIfEnabled(s.appConfig.TracingMaxItems, s.appConfig.TracingMaxBodyBytes)
+	errStr := ""
+	if err != nil {
+		errStr = err.Error()
+	}
+	summary := op + " " + outcome
+	if op == "search" && outcome == "hit" {
+		summary = fmt.Sprintf("search hit (sim=%.3f)", sim)
+	}
+	data := map[string]any{
+		"op":         op,
+		"outcome":    outcome,
+		"vector_dim": vecDim,
+	}
+	// Only include similarity for a real neighbor — miss/empty_store would
+	// otherwise render "similarity: 0" and read as a measured value.
+	if op == "search" && outcome == "hit" {
+		data["similarity"] = sim
+	}
+	trace.RecordBackendTrace(trace.BackendTrace{
+		Timestamp: start,
+		Duration:  time.Since(start),
+		Type:      trace.BackendTraceVectorStore,
+		ModelName: s.storeName,
+		Backend:   model.LocalStoreBackend,
+		Summary:   summary,
+		Error:     errStr,
+		Data:      data,
+	})
 }
 
 func StoreBackend(sl *model.ModelLoader, appConfig *config.ApplicationConfig, storeName string, backend string) (grpc.Backend, error) {

core/backend/stores_test.go

@@ -0,0 +1,88 @@
+package backend
+
+import (
+	"context"
+
+	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/trace"
+	"github.com/mudler/LocalAI/pkg/model"
+	"github.com/mudler/LocalAI/pkg/system"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+// findVectorStoreTrace returns the most recent vector_store trace whose
+// model_name matches storeName, or nil if none was recorded. Used by
+// the specs below to assert the trace landed without relying on
+// ring-buffer ordering across other tests in the suite.
+func findVectorStoreTrace(storeName string) *trace.BackendTrace {
+	traces := trace.GetBackendTraces()
+	for i := range traces {
+		bt := &traces[i]
+		if bt.Type == trace.BackendTraceVectorStore && bt.ModelName == storeName {
+			return bt
+		}
+	}
+	return nil
+}
+
+var _ = Describe("localVectorStore tracing", func() {
+	// Pin the trace surface admins read from /api/backend-traces.
+	// The original failure mode that motivated these specs — the
+	// local-store backend not installed — was silent on every surface
+	// except a per-call xlog.Warn. With tracing wired in, the row
+	// appears next to the embedder/score traces for the same request.
+	BeforeEach(func() {
+		trace.ClearBackendTraces()
+	})
+
+	It("records a vector_store trace with outcome=backend_load_error when the backend can't be loaded", func() {
+		// nil ModelLoader → s.backend → StoreBackend → panics on load.
+		// Use a real-but-empty loader so the failure surfaces as an
+		// error instead, exercising the load-failure trace path the
+		// admin would hit when local-store isn't installed.
+		appCfg := &config.ApplicationConfig{
+			EnableTracing:       true,
+			TracingMaxItems:     16,
+			TracingMaxBodyBytes: 1024,
+		}
+		s := &localVectorStore{
+			loader:    model.NewModelLoader(&system.SystemState{}),
+			appConfig: appCfg,
+			storeName: "router-cache-test",
+		}
+
+		// Search must surface the error AND record a trace describing it.
+		_, _, _, err := s.Search(context.Background(), []float32{0.1, 0.2, 0.3})
+		Expect(err).To(HaveOccurred())
+
+		Eventually(func() *trace.BackendTrace {
+			return findVectorStoreTrace("router-cache-test")
+		}).ShouldNot(BeNil())
+
+		bt := findVectorStoreTrace("router-cache-test")
+		Expect(bt.Backend).To(Equal(model.LocalStoreBackend))
+		Expect(bt.Data["op"]).To(Equal("search"))
+		Expect(bt.Data["outcome"]).To(Equal("backend_load_error"))
+		Expect(bt.Data["vector_dim"]).To(Equal(3))
+		// Error is the wrapped "vector store load: …" surfaced to the caller.
+		Expect(bt.Error).To(ContainSubstring("vector store load"))
+	})
+
+	It("does not record a trace when tracing is disabled", func() {
+		// Opt-out path: appConfig.EnableTracing=false must short-circuit
+		// before InitBackendTracingIfEnabled, so a workload with tracing
+		// turned off doesn't pay the channel-send cost per cache call.
+		appCfg := &config.ApplicationConfig{EnableTracing: false}
+		s := &localVectorStore{
+			loader:    model.NewModelLoader(&system.SystemState{}),
+			appConfig: appCfg,
+			storeName: "router-cache-disabled",
+		}
+		_, _, _, _ = s.Search(context.Background(), []float32{1})
+		Consistently(func() *trace.BackendTrace {
+			return findVectorStoreTrace("router-cache-disabled")
+		}).Should(BeNil())
+	})
+})

core/backend/tokenize.go

@@ -7,9 +7,23 @@ import (
 	"github.com/mudler/LocalAI/core/schema"
 	"github.com/mudler/LocalAI/core/trace"
 	"github.com/mudler/LocalAI/pkg/grpc"
+	pb "github.com/mudler/LocalAI/pkg/grpc/proto"
 	"github.com/mudler/LocalAI/pkg/model"
 )
 
+// tokenizeTokenCount returns the number of tokens in a backend response,
+// treating a nil response as zero. The gRPC client returns (nil, err) on
+// failure, and the tracing block below runs before that error is returned —
+// so the count must be read nil-safely here. Reading resp.Tokens on a nil
+// resp previously panicked the whole HTTP handler when tracing was enabled
+// (e.g. a transient tokenize failure during router probe-budget sizing).
+func tokenizeTokenCount(resp *pb.TokenizationResponse) int {
+	if resp == nil {
+		return 0
+	}
+	return len(resp.Tokens)
+}
+
 func ModelTokenize(s string, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (schema.TokenizeResponse, error) {
 
 	var inferenceModel grpc.Backend
@@ -40,10 +54,7 @@ func ModelTokenize(s string, loader *model.ModelLoader, modelConfig config.Model
 			errStr = err.Error()
 		}
 
-		tokenCount := 0
-		if resp.Tokens != nil {
-			tokenCount = len(resp.Tokens)
-		}
+		tokenCount := tokenizeTokenCount(resp)
 
 		trace.RecordBackendTrace(trace.BackendTrace{
 			Timestamp: startTime,
@@ -64,8 +75,8 @@ func ModelTokenize(s string, loader *model.ModelLoader, modelConfig config.Model
 		return schema.TokenizeResponse{}, err
 	}
 
-	if resp.Tokens == nil {
-		resp.Tokens = make([]int32, 0)
+	if resp == nil || resp.Tokens == nil {
+		return schema.TokenizeResponse{Tokens: make([]int32, 0)}, nil
 	}
 
 	return schema.TokenizeResponse{

core/backend/tokenize_test.go

@@ -0,0 +1,27 @@
+package backend
+
+import (
+	pb "github.com/mudler/LocalAI/pkg/grpc/proto"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("tokenizeTokenCount", func() {
+	// Regression: the gRPC client returns (nil, err) when a tokenize call
+	// fails, and ModelTokenize's tracing block reads the token count before
+	// the error is returned. Dereferencing a nil response there panicked the
+	// HTTP handler (nil pointer dereference) — e.g. a transient tokenize
+	// failure while the router sized its probe-token budget.
+	It("returns zero for a nil response instead of panicking", func() {
+		Expect(tokenizeTokenCount(nil)).To(Equal(0))
+	})
+
+	It("returns zero when the response carries no tokens", func() {
+		Expect(tokenizeTokenCount(&pb.TokenizationResponse{})).To(Equal(0))
+	})
+
+	It("counts the tokens present on the response", func() {
+		Expect(tokenizeTokenCount(&pb.TokenizationResponse{Tokens: []int32{1, 2, 3}})).To(Equal(3))
+	})
+})