fix(realtime): harden classifier slot completion

Reserve context for constrained slot filling, size completions from their encoded output, and encode enum grammar literals as valid JSON. Reject empty enum values and cover the failure modes with regression tests.

Assisted-by: Codex:gpt-5
Signed-off-by: Richard Palethorpe <io@richiejp.com>

Author: richiejp

core/http/endpoints/openai/realtime_classifier.go

@@ -429,7 +429,8 @@ func slotFillGrammar(slots []types.ClassifierSlot) string {
 				if vi > 0 {
 					rules.WriteString(" | ")
 				}
-				rules.WriteString(gbnfLiteral(`"` + v + `"`))
+				encoded, _ := json.Marshal(v) // validation rejects values JSON cannot encode
+				rules.WriteString(gbnfLiteral(string(encoded)))
 			}
 		default: // string
 			rules.WriteString("str")
@@ -446,6 +447,57 @@ func slotFillGrammar(slots []types.ClassifierSlot) string {
 	return root.String() + rules.String()
 }
 
+const (
+	// Free-form values need an explicit ceiling; forced enum values and field
+	// syntax are budgeted from their actual JSON encoding below.
+	slotFillStringTokens = 64
+	slotFillNumberTokens = 32
+)
+
+// slotFillMaxTokens conservatively budgets one token per output byte for the
+// forced JSON tail, plus explicit allowances for free-form values. This avoids
+// truncating long enum values or field names while keeping string generation
+// bounded.
+func slotFillMaxTokens(slots []types.ClassifierSlot) int {
+	tokens := 1 // closing brace
+	for i := range slots {
+		if i > 0 {
+			field, _ := json.Marshal(slots[i].Name)
+			tokens += len(field) + len(`, : `)
+		}
+		switch slots[i].Type {
+		case types.ClassifierSlotNumber:
+			tokens += slotFillNumberTokens
+		case types.ClassifierSlotString:
+			tokens += slotFillStringTokens
+		case types.ClassifierSlotEnum:
+			longest := 0
+			for _, value := range slots[i].Values {
+				encoded, _ := json.Marshal(value)
+				if len(encoded) > longest {
+					longest = len(encoded)
+				}
+			}
+			tokens += longest
+		}
+	}
+	return tokens
+}
+
+// slotFillContextReserve includes both the generated tail and the continuation
+// prefix appended after the scored prompt. It intentionally over-reserves by
+// counting bytes as tokens; preserving the identical scoring prompt is more
+// important than reclaiming a handful of context tokens.
+func slotFillContextReserve(option *types.ClassifierOption) int {
+	if option == nil || option.Tool == nil || len(option.Tool.Slots) == 0 {
+		return 0
+	}
+	route, _ := json.Marshal(option.ID)
+	field, _ := json.Marshal(option.Tool.Slots[0].Name)
+	prefixBytes := len(`{"route": , : `) + len(route) + len(field)
+	return prefixBytes + slotFillMaxTokens(option.Tool.Slots)
+}
+
 // parseSlotValues closes the completed route JSON and extracts each slot's
 // value as the string form SpliceArguments expects.
 func parseSlotValues(chosenID, firstSlot, generated string, slots []types.ClassifierSlot) (map[string]string, error) {

core/http/endpoints/openai/realtime_classifier_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"strings"
 
 	"github.com/mudler/LocalAI/core/config"
 	"github.com/mudler/LocalAI/core/http/endpoints/openai/types"
@@ -539,6 +540,24 @@ var _ = Describe("slotFillGrammar", func() {
 		Expect(g).To(ContainSubstring("num ::="))
 	})
 
+	It("JSON-encodes enum values before embedding them in the grammar", func() {
+		g := slotFillGrammar([]types.ClassifierSlot{
+			{Name: "units", Type: types.ClassifierSlotEnum, Values: []string{"quoted\"value", "line\nbreak", `back\slash`}},
+		})
+		Expect(g).To(ContainSubstring(gbnfLiteral(`"quoted\"value"`)))
+		Expect(g).To(ContainSubstring(gbnfLiteral(`"line\nbreak"`)))
+		Expect(g).To(ContainSubstring(gbnfLiteral(`"back\\slash"`)))
+	})
+
+	It("budgets forced enum and field text by encoded length", func() {
+		short := []types.ClassifierSlot{{Name: "value", Type: types.ClassifierSlotEnum, Values: []string{"m"}}}
+		long := []types.ClassifierSlot{
+			{Name: "value", Type: types.ClassifierSlotEnum, Values: []string{strings.Repeat("long-value-", 20)}},
+			{Name: strings.Repeat("field", 20), Type: types.ClassifierSlotNumber},
+		}
+		Expect(slotFillMaxTokens(long)).To(BeNumerically(">", slotFillMaxTokens(short)+200))
+	})
+
 	It("emits a string rule only when needed", func() {
 		g := slotFillGrammar([]types.ClassifierSlot{{Name: "what", Type: types.ClassifierSlotString}})
 		Expect(g).To(ContainSubstring("slot0 ::= str"))

core/http/endpoints/openai/realtime_model.go

@@ -456,6 +456,18 @@ func (m *wrappedModel) classifierFor(options []types.ClassifierOption, normaliza
 		CacheCap:      0,
 		Normalization: normalization,
 	}
+	if m.routerDeps != nil && m.routerDeps.TokenCounter != nil && cfg.ContextSize != nil {
+		opts.TokenCounter = m.routerDeps.TokenCounter(cfg.Name)
+		opts.MaxContextTokens = *cfg.ContextSize
+	}
+	for i := range options {
+		if options[i].Tool != nil && len(options[i].Tool.Slots) > 0 {
+			reserve := slotFillContextReserve(&options[i])
+			if reserve > opts.CompletionReserveTokens {
+				opts.CompletionReserveTokens = reserve
+			}
+		}
+	}
 	if m.evaluator != nil {
 		if renderer := middleware.NewTemplateRenderer(m.evaluator, cfg); renderer != nil {
 			opts.PromptRenderer = renderer
@@ -522,7 +534,7 @@ func (m *wrappedModel) FillToolArguments(ctx context.Context, messages schema.Me
 		return "", nil, fmt.Errorf("classifier: slot filling requires completion in the scoring model's known_usecases")
 	}
 	cfg.Grammar = slotFillGrammar(slots)
-	maxTokens := 16 + 16*len(slots)
+	maxTokens := slotFillMaxTokens(slots)
 	temperature := 0.0
 	cfg.Maxtokens = &maxTokens
 	cfg.Temperature = &temperature

core/http/endpoints/openai/types/classifier.go

@@ -390,6 +390,9 @@ func validateSlots(t *ClassifierTool) error {
 			if len(s.Values) == 0 {
 				return fmt.Errorf("slot %q: enum slots need values", s.Name)
 			}
+			if slices.Contains(s.Values, "") {
+				return fmt.Errorf("slot %q: enum values must be non-empty", s.Name)
+			}
 			if s.Default != "" && !slices.Contains(s.Values, s.Default) {
 				return fmt.Errorf("slot %q: default %q is not one of its values", s.Name, s.Default)
 			}

core/http/endpoints/openai/types/classifier_test.go

@@ -228,6 +228,12 @@ var _ = Describe("ClassifierTool slots", func() {
 			Expect(cfgWith(tool(bad, enumSlot)).Validate()).To(MatchError(ContainSubstring("does not parse")))
 		})
 
+		It("rejects empty enum values that cannot be spliced", func() {
+			bad := enumSlot
+			bad.Values = []string{"m", ""}
+			Expect(cfgWith(tool(numberSlot, bad)).Validate()).To(MatchError(ContainSubstring("must be non-empty")))
+		})
+
 		It("rejects slots the template never references", func() {
 			t := tool(numberSlot, enumSlot, types.ClassifierSlot{Name: "speed", Type: types.ClassifierSlotNumber})
 			Expect(cfgWith(t).Validate()).To(MatchError(ContainSubstring("{{speed}}")))

core/services/routing/router/score.go

@@ -98,6 +98,11 @@ type ScoreClassifierOptions struct {
 	// sends Probe.Prompt as-is and relies on the backend's n_ctx guard.
 	TokenCounter     func(string) (int, error)
 	MaxContextTokens int
+
+	// CompletionReserveTokens reserves additional context beyond the longest
+	// scoring candidate. Classifier slot filling uses this to ensure the prompt
+	// scored here can be continued without overflowing the model context.
+	CompletionReserveTokens int
 }
 
 // ScoreClassifier scores every policy label as the model's actual
@@ -202,8 +207,13 @@ func NewScoreClassifier(policies []ScorePolicy, scorer backend.Scorer, opts Scor
 		systemPrompt:        systemPrompt,
 		labelOrder:          labels,
 		candidates:          candidates,
-		budget:              &lazyBudget{tokenize: opts.TokenCounter, maxContext: opts.MaxContextTokens, extras: candidates},
-		cache:               newLabelSetCache(opts.CacheCap),
+		budget: &lazyBudget{
+			tokenize:   opts.TokenCounter,
+			maxContext: opts.MaxContextTokens,
+			extras:     candidates,
+			reserve:    opts.CompletionReserveTokens,
+		},
+		cache: newLabelSetCache(opts.CacheCap),
 	}
 }
 

core/services/routing/router/score_test.go

@@ -368,6 +368,19 @@ var _ = Describe("ScoreClassifier conversation trimming", func() {
 		Expect(len(strings.Fields(s.lastP))).To(BeNumerically("<", 20000), "must be trimmed, not the full transcript")
 	})
 
+	It("reserves context for a completion after scoring", func() {
+		without := NewScoreClassifier(testPolicies(), &stubScorer{}, ScoreClassifierOptions{
+			TokenCounter:     wordCount,
+			MaxContextTokens: 10000,
+		})
+		with := NewScoreClassifier(testPolicies(), &stubScorer{}, ScoreClassifierOptions{
+			TokenCounter:            wordCount,
+			MaxContextTokens:        10000,
+			CompletionReserveTokens: 257,
+		})
+		Expect(with.probeTokenBudget()).To(Equal(without.probeTokenBudget() - 257))
+	})
+
 	It("keeps the newest turn whole even when it alone exceeds the budget", func() {
 		s := &stubScorer{results: threeScores}
 		c := NewScoreClassifier(testPolicies(), s, ScoreClassifierOptions{

core/services/routing/router/trim.go

@@ -123,6 +123,7 @@ type lazyBudget struct {
 	tokenize   func(string) (int, error)
 	maxContext int
 	extras     []string
+	reserve    int
 
 	mu    sync.Mutex
 	value atomic.Int64 // 0=unset, >0=budget, -1=disabled
@@ -156,7 +157,7 @@ func (l *lazyBudget) get() int {
 			longest = n
 		}
 	}
-	b := l.maxContext - longest - tokenBudgetMargin
+	b := l.maxContext - longest - l.reserve - tokenBudgetMargin
 	if b <= 0 {
 		l.value.Store(-1)
 		return 0