Run the routing heads in Go and route with them #206

Merged
claude merged 6 commits from task/664-routing-heads-in-go into master 2026-08-08 23:21:41 +02:00
13 changed files with 138 additions and 15 deletions
Showing only changes of commit a4abcdefa3 - Show all commits
+1 -1
View File
@@ -19,7 +19,7 @@ func TestChatAnswersWithNoLlamaServer(t *testing.T) {
dead := llm.New("http://127.0.0.1:1", 500*time.Millisecond)
emb := router.NewHashEmbedder(1024)
h.recall.embedder = emb
h.router = buildRouter(emb, h.matcher, 0.55, pickLLMRouter(true, dead))
h.router = buildRouter(emb, h.matcher, 0.55, pickLLMRouter(true, dead), nil)
h.replier = newLLMReplier(dead, nil)
ctx := withDialogueID(context.Background(), dialogueIDFor(sourceText, "web"))
+2 -2
View File
@@ -317,7 +317,7 @@ func TestClarifyExpiryIsAnnouncedAndWordsStillRoute(t *testing.T) {
h, _, now := newClarifyHandler(t)
emb := router.NewHashEmbedder(1024)
h.recall.embedder = emb
h.router = buildRouter(emb, h.matcher, 0.55, nil)
h.router = buildRouter(emb, h.matcher, 0.55, nil, nil)
if _, asked := h.askClarify(ctx, clarifyDec(router.IntentReminder, router.Slots{Text: "напомни"}, "напомни")); !asked {
t.Fatal("expected a question")
@@ -671,7 +671,7 @@ func TestUnresolvedActSaysItDoesNotKnowTheCommand(t *testing.T) {
func newRoutingClarifyHandler(t *testing.T) (*reactiveHandler, *store.Store) {
t.Helper()
h, st, _ := newClarifyHandler(t)
h.router = buildRouter(router.NewHashEmbedder(1024), h.matcher, 0.55, nil)
h.router = buildRouter(router.NewHashEmbedder(1024), h.matcher, 0.55, nil, nil)
h.recall = recallWiring{embedder: router.NewHashEmbedder(1024), memStore: memory.NewInMemoryStore()}
return h, st
}
+1 -1
View File
@@ -25,7 +25,7 @@ func traceHandler(t *testing.T, ring *decision.Ring) *reactiveHandler {
return &reactiveHandler{
api: api,
recall: recallWiring{embedder: emb, memStore: memory.NewInMemoryStore()},
router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil),
router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil, nil),
replier: voice.NewStubReplier(),
now: func() time.Time { return now },
dataStore: st,
+1 -1
View File
@@ -171,7 +171,7 @@ func newDialogueHandler(t *testing.T) (*reactiveHandler, *store.Store, *time.Tim
// and never a coincidence (V-577, V-579). checkEnd refuses any reminder
// landing on it, and at 09:00 the row that answers "на 9" would trip that.
*now = time.Date(2026, 7, 31, 9, 17, 0, 0, time.UTC)
h.router = buildRouter(router.NewHashEmbedder(1024), h.matcher, 0.55, nil)
h.router = buildRouter(router.NewHashEmbedder(1024), h.matcher, 0.55, nil, nil)
h.recall = recallWiring{embedder: router.NewHashEmbedder(1024), memStore: memory.NewInMemoryStore()}
return h, st, now
}
+1 -1
View File
@@ -24,7 +24,7 @@ func TestApplyAction_FactCapture_QueuesEntityResolution(t *testing.T) {
emb := router.NewHashEmbedder(1024)
matcher := tool.NewMatcher(api)
rtr := buildRouter(emb, matcher, 0.55, nil)
rtr := buildRouter(emb, matcher, 0.55, nil, nil)
h := &reactiveHandler{
api: api,
+1 -1
View File
@@ -20,7 +20,7 @@ func newFactGateHandler(t *testing.T, now time.Time) (*reactiveHandler, ipc.Core
h := &reactiveHandler{
api: api,
recall: recallWiring{embedder: emb, memStore: memory.NewInMemoryStore()},
router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil),
router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil, nil),
replier: voice.NewStubReplier(),
now: func() time.Time { return now },
dataStore: st,
+1 -1
View File
@@ -45,7 +45,7 @@ func newNoteHandler(t *testing.T) (*reactiveHandler, *store.Store) {
h := &reactiveHandler{
api: api,
recall: recallWiring{embedder: emb, memStore: memory.NewInMemoryStore()},
router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil),
router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil, nil),
replier: voice.NewStubReplier(),
now: func() time.Time { return now },
dataStore: st,
+2 -2
View File
@@ -22,7 +22,7 @@ func TestReactiveNotesReminders(t *testing.T) {
emb := router.NewHashEmbedder(1024)
matcher := tool.NewMatcher(api)
rtr := buildRouter(emb, matcher, 0.55, nil)
rtr := buildRouter(emb, matcher, 0.55, nil, nil)
h := &reactiveHandler{
api: api,
@@ -104,7 +104,7 @@ func TestSpokenTaskCaptureFilesATask(t *testing.T) {
h := &reactiveHandler{
api: api,
recall: recallWiring{embedder: emb, memStore: memory.NewInMemoryStore()},
router: buildRouter(emb, matcher, 0.55, nil),
router: buildRouter(emb, matcher, 0.55, nil, nil),
replier: voice.NewStubReplier(),
now: func() time.Time { return now },
dataStore: st,
+1 -1
View File
@@ -474,7 +474,7 @@ func newSimWorld(t *testing.T, sc scenario) *simWorld {
// used to be built on a nil API, which meant any scenario that produced an
// act panicked the moment the matcher was consulted.
matcher := tool.NewMatcher(api)
rtr := buildRouter(emb, matcher, config.DefaultRouterThreshold, router.NewLLMRouter(scripted))
rtr := buildRouter(emb, matcher, config.DefaultRouterThreshold, router.NewLLMRouter(scripted), nil)
w.handler = &reactiveHandler{
stt: simTranscriber{},
+29 -3
View File
@@ -36,7 +36,9 @@ type voiceWiring struct {
sessions *voice.Sessions
voiceSink delivery.Sink
embedder router.Embedder
handler *reactiveHandler // the reactive handler for IPC Chat
// heads — the routing heads, nil unless embedder.heads_path is set.
heads *router.RouterHeads
handler *reactiveHandler // the reactive handler for IPC Chat
// worker clients (set when configured as Remote): closed on shutdown so
// mavsttd / mavttsd don't keep a stale conn into a restarting daemon.
sttClient *worker.Client
@@ -72,6 +74,9 @@ func (w *voiceWiring) close() {
if w.embedder != nil {
_ = w.embedder.Close()
}
if w.heads != nil {
_ = w.heads.Close()
}
if w.server != nil {
_ = w.server.Close()
}
@@ -147,6 +152,24 @@ func wireVoice(cfg *config.Config, coreAPI ipc.CoreAPI, phr phraser.Phraser, mem
emb = router.NewHashEmbedder(1024)
}
w.embedder = emb
// ----- router: routing heads (only when configured, and never fatal) -----
// A missing or broken weights file logs and leaves w.heads nil, which is
// byte-for-byte the cascade that shipped before V-664. Refusing to start
// over a routing accelerator would trade a working box for a better one.
if cfg.Voice.Embedder != nil && cfg.Voice.Embedder.HeadsPath != "" {
h, err := router.NewRouterHeads(
cfg.Voice.Embedder.HeadsPath,
cfg.Voice.Embedder.TokenizerPath,
)
if err != nil {
log.Printf("voice: routing heads unavailable, cascade unchanged: %v", err)
} else {
log.Printf("voice: routing heads loaded from %s", cfg.Voice.Embedder.HeadsPath)
w.heads = h
}
}
repairFactVectors(dataStore, emb)
checkStoredEmbedder(dataStore, emb)
// Retention is enforced on write, which is not enough on its own: a box that
@@ -223,7 +246,8 @@ func wireVoice(cfg *config.Config, coreAPI ipc.CoreAPI, phr phraser.Phraser, mem
// against the classifier's 50.0%, at about 1s a turn instead of 30ms (see
// config.VoiceConfig.LLMRouter). The classifier always stays wired as the
// fallback, so a model error never breaks a turn.
rtr := buildRouter(emb, matcher, threshold, pickLLMRouter(cfg.Voice.UseLLMRouter(), hot))
rtr := buildRouter(emb, matcher, threshold,
pickLLMRouter(cfg.Voice.UseLLMRouter(), hot), w.heads)
// ----- sessions registry (shared with voicesink) -----
sessions := voice.NewSessions()
@@ -390,7 +414,8 @@ func pickLLMRouter(enabled bool, c router.Completer) *router.LLMRouter {
// intent from seedDir (models/seeds/<intent>.txt) — see seedClassifier
// below for the current intent list and file names.
// - Threshold is from voice.router_threshold config (default 0.55).
func buildRouter(emb router.Embedder, acts router.ActMatcher, threshold float64, llmR *router.LLMRouter) *router.Router {
func buildRouter(emb router.Embedder, acts router.ActMatcher, threshold float64,
llmR *router.LLMRouter, heads *router.RouterHeads) *router.Router {
cls := router.NewClassifier(emb)
seedClassifier(cls)
grammars := router.DefaultGrammars(acts)
@@ -442,6 +467,7 @@ func buildRouter(emb router.Embedder, acts router.ActMatcher, threshold float64,
},
Threshold: threshold,
LLM: llmR,
Heads: heads,
})
}
+2 -1
View File
@@ -231,7 +231,8 @@
"embedder": {
"model_path": "/opt/maven/models/embedder/multilingual-e5-small/model_quantized.onnx",
"tokenizer_path": "/opt/maven/models/embedder/multilingual-e5-small/tokenizer.json",
"lib_path": "/opt/maven/lib/libonnxruntime.so"
"lib_path": "/opt/maven/lib/libonnxruntime.so",
"heads_path": "/opt/maven/models/embedder/router-heads/router_heads.onnx"
},
"llm_router": true,
"query_min_score": 0.55,
+10
View File
@@ -37,6 +37,16 @@ type EmbedderConfig struct {
ModelPath string `json:"model_path,omitempty"`
TokenizerPath string `json:"tokenizer_path,omitempty"`
LibPath string `json:"lib_path,omitempty"`
// HeadsPath — the routing heads graph, which is a fine-tuned COPY of the
// model above with four linear heads on its pooled output (V-664). Empty
// means no heads, and the cascade runs exactly as it did before they
// existed. It shares LibPath and TokenizerPath, and router_heads.json is
// read from the same directory.
//
// It must never be pointed at ModelPath. Memory recall depends on the
// resident copy scoring what it scored, and the fine-tuned one does not.
HeadsPath string `json:"heads_path,omitempty"`
}
// WeatherConfig configures the weather provider for voice queries.
+86
View File
@@ -0,0 +1,86 @@
package eval
import (
"context"
"os"
"path/filepath"
"testing"
"github.com/kami/maven/internal/config"
"github.com/kami/maven/internal/router"
)
// TestONNXRoutingHeads — the cascade with the routing heads wired, which is
// what V-664 deploys. Opt-in via MAVEN_ONNX_LIB, same as TestONNXBaseline, and
// one TestONNX* per process.
//
// The comparison worth reading is against TestONNXBaseline, which is the same
// cascade with the same grammars and the same classifier floor and no heads.
// Only the middle arm varies.
//
// It also checks the Go unigram tokenizer against the Python one, because the
// heads were trained through transformers and are read through a hand-written
// tokenizer. A mismatch shows up here as a score below what Python measured on
// the same weights, and nowhere else.
func TestONNXRoutingHeads(t *testing.T) {
lib := os.Getenv("MAVEN_ONNX_LIB")
if lib == "" {
t.Skip("MAVEN_ONNX_LIB unset — see AGENTS.md § Embedder model for intent routing")
}
// Absolute, because onnxruntime resolves a graph's external weights file
// against the model path it was given, and a relative one lands in the
// test's working directory.
root, err := filepath.Abs("../../..")
if err != nil {
t.Fatal(err)
}
model := filepath.Join(root, "models/embedder/multilingual-e5-small/model_quantized.onnx")
tok := filepath.Join(root, "models/embedder/multilingual-e5-small/tokenizer.json")
heads := filepath.Join(root, "models/embedder/router-heads/router_heads.onnx")
for _, p := range []string{lib, model, tok, heads} {
if _, err := os.Stat(p); err != nil {
t.Skipf("missing %s: %v", p, err)
}
}
emb, err2 := router.NewONNXEmbedder(model, tok, lib)
if err2 != nil {
t.Skipf("onnx embedder unavailable: %v", err2)
}
err = nil
defer emb.Close()
h, err := router.NewRouterHeads(heads, tok)
if err != nil {
t.Skipf("routing heads unavailable: %v", err)
}
defer h.Close()
f, err := Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
rep, err := Score(context.Background(), "heads+classifier", withHeads(t, emb, h), f)
if err != nil {
t.Fatalf("Score: %v", err)
}
t.Log("\n" + rep.String() + rep.Failures())
}
// withHeads mirrors newBaselineRouter and adds the one arm under test. It is a
// separate function rather than a parameter so the baseline's signature stays
// the shape every other test calls it with.
func withHeads(t *testing.T, emb router.Embedder, h *router.RouterHeads) *router.Router {
t.Helper()
acts := router.DefaultActMatcher{Fns: actFns}
return router.New(router.Config{
Grammars: baselineGrammars(acts),
Classifier: newBaselineClassifier(t, emb),
Extractor: router.Extractor{
Time: router.StubDateTimeParser{},
Acts: acts,
Facts: router.DefaultFactParser{},
},
Threshold: config.DefaultRouterThreshold,
Heads: h,
})
}