From c45198f39fc63b692598341eb3ec2d588082dade Mon Sep 17 00:00:00 2001 From: BMAD CI Fix Agent Date: Wed, 16 Sep 2026 16:12:09 -0500 Subject: [PATCH] Add LLM classifier routing --- README.md | 14 ++- configs/router.example.yaml | 15 +-- internal/config/config.go | 66 ++++++++++- internal/server/server.go | 198 ++++++++++++++++++++++++++++++++- internal/server/server_test.go | 151 +++++++++++++++++++++++++ 5 files changed, 424 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 632c05b..2cd0020 100644 --- a/README.md +++ b/README.md @@ -136,12 +136,14 @@ models: - id: large-prompt target_model: qwen/qwen3.6-35b-a3b min_prompt_chars: 12000 - - id: hard-work-keywords - target_model: qwen/qwen3.6-35b-a3b - any_keywords: - - architecture - - refactor - - security + classifier: + backend: lmstudio + model: qwen3-coder-30b-a3b-instruct + target_models: + - qwen3-coder-30b-a3b-instruct + - qwen/qwen3.6-35b-a3b + timeout: 15s + max_tokens: 64 ensure: mode: command command: diff --git a/configs/router.example.yaml b/configs/router.example.yaml index 69a8efe..5744f61 100644 --- a/configs/router.example.yaml +++ b/configs/router.example.yaml @@ -49,13 +49,14 @@ models: - id: large-prompt target_model: qwen/qwen3.6-35b-a3b min_prompt_chars: 12000 - - id: hard-work-keywords - target_model: qwen/qwen3.6-35b-a3b - any_keywords: - - architecture - - migration - - refactor - - security + classifier: + backend: lmstudio + model: qwen3-coder-30b-a3b-instruct + target_models: + - qwen3-coder-30b-a3b-instruct + - qwen/qwen3.6-35b-a3b + timeout: 15s + max_tokens: 64 backends: - id: lmstudio diff --git a/internal/config/config.go b/internal/config/config.go index 739be7c..1ba0a80 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -52,7 +52,8 @@ type EnsureConfig struct { } type RoutingConfig struct { - Rules []RoutingRuleConfig `yaml:"rules"` + Rules []RoutingRuleConfig `yaml:"rules"` + Classifier RoutingClassifierConfig `yaml:"classifier"` } type RoutingRuleConfig struct { @@ -65,6 +66,15 @@ type RoutingRuleConfig struct { AnyKeywords []string `yaml:"any_keywords"` } +type RoutingClassifierConfig struct { + Backend string `yaml:"backend"` + Model string `yaml:"model"` + TargetModels []string `yaml:"target_models"` + SystemPrompt string `yaml:"system_prompt"` + Timeout string `yaml:"timeout"` + MaxTokens int `yaml:"max_tokens"` +} + type CommandArgs []string func (args *CommandArgs) UnmarshalYAML(value *yaml.Node) error { @@ -171,6 +181,11 @@ func (cfg Config) Validate() error { if err := model.Routing.Validate(); err != nil { return fmt.Errorf("model %q routing is invalid: %w", model.ID, err) } + if model.Routing.Classifier.Enabled() { + if _, ok := backends[model.Routing.Classifier.Backend]; !ok { + return fmt.Errorf("model %q routing classifier references unknown backend %q", model.ID, model.Routing.Classifier.Backend) + } + } if _, ok := models[model.ID]; ok { return fmt.Errorf("model %q is duplicated", model.ID) } @@ -218,9 +233,58 @@ func (routing RoutingConfig) Validate() error { } } + if err := routing.Classifier.Validate(); err != nil { + return fmt.Errorf("classifier is invalid: %w", err) + } + return nil } +func (classifier RoutingClassifierConfig) Validate() error { + if classifier.Backend == "" && classifier.Model == "" && len(classifier.TargetModels) == 0 && classifier.SystemPrompt == "" && classifier.Timeout == "" && classifier.MaxTokens == 0 { + return nil + } + if classifier.Backend == "" { + return errors.New("backend is required") + } + if classifier.Model == "" { + return errors.New("model is required") + } + if len(classifier.TargetModels) == 0 { + return errors.New("target_models is required") + } + for index, target := range classifier.TargetModels { + if strings.TrimSpace(target) == "" { + return fmt.Errorf("target_models[%d] must not be empty", index) + } + } + if classifier.MaxTokens < 0 { + return errors.New("max_tokens must be non-negative") + } + _, err := classifier.TimeoutDuration() + return err +} + +func (classifier RoutingClassifierConfig) Enabled() bool { + return classifier.Backend != "" || classifier.Model != "" || len(classifier.TargetModels) > 0 || classifier.SystemPrompt != "" || classifier.Timeout != "" || classifier.MaxTokens != 0 +} + +func (classifier RoutingClassifierConfig) TimeoutDuration() (time.Duration, error) { + if classifier.Timeout == "" { + return 15 * time.Second, nil + } + + duration, err := time.ParseDuration(classifier.Timeout) + if err != nil { + return 0, err + } + if duration <= 0 { + return 0, errors.New("duration must be positive") + } + + return duration, nil +} + func (cfg Config) Model(id string) (ModelConfig, bool) { for _, model := range cfg.Models { if model.ID == id { diff --git a/internal/server/server.go b/internal/server/server.go index 58cc460..5b74f31 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -151,7 +151,7 @@ func (s *Server) proxyOpenAI(w http.ResponseWriter, r *http.Request) { return } - targetModel, routeRule := selectTargetModel(body, model) + targetModel, routeRule := s.selectTargetModel(r.Context(), body, model, requestID) targetMaxOutputTokens := 0 if targetAlias, ok := s.cfg.Model(targetModel); ok { targetMaxOutputTokens = targetAlias.MaxOutputTokens @@ -881,11 +881,7 @@ type routingFeatures struct { RequestedMaxToken int } -func selectTargetModel(body []byte, model config.ModelConfig) (string, string) { - if len(model.Routing.Rules) == 0 { - return model.TargetModel, "default" - } - +func (s *Server) selectTargetModel(ctx context.Context, body []byte, model config.ModelConfig, requestID string) (string, string) { features := extractRoutingFeatures(body) for _, rule := range model.Routing.Rules { if routingRuleMatches(rule, features) { @@ -896,9 +892,89 @@ func selectTargetModel(body []byte, model config.ModelConfig) (string, string) { } } + if model.Routing.Classifier.Enabled() { + target, ok := s.classifyTargetModel(ctx, model, features, requestID) + if ok { + return target, "classifier" + } + } + return model.TargetModel, "default" } +func (s *Server) classifyTargetModel(ctx context.Context, model config.ModelConfig, features routingFeatures, requestID string) (string, bool) { + classifier := model.Routing.Classifier + backend, ok := s.cfg.Backend(classifier.Backend) + if !ok { + slog.Warn("routing classifier backend is missing", "request_id", requestID, "alias", model.ID, "backend", classifier.Backend) + return "", false + } + + timeout, err := classifier.TimeoutDuration() + if err != nil { + slog.Warn("routing classifier timeout is invalid", "request_id", requestID, "alias", model.ID, "error", err) + return "", false + } + classifierCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + requestPayload := map[string]any{ + "model": classifier.Model, + "messages": classifierMessages(classifier, features), + "temperature": 0, + "stream": false, + "max_tokens": classifierMaxTokens(classifier), + } + requestBody, err := json.Marshal(requestPayload) + if err != nil { + slog.Warn("routing classifier request could not be encoded", "request_id", requestID, "alias", model.ID, "error", err) + return "", false + } + + endpoint, err := backendChatCompletionsURL(backend) + if err != nil { + slog.Warn("routing classifier backend URL is invalid", "request_id", requestID, "alias", model.ID, "backend", backend.ID, "error", err) + return "", false + } + + req, err := http.NewRequestWithContext(classifierCtx, http.MethodPost, endpoint, bytes.NewReader(requestBody)) + if err != nil { + slog.Warn("routing classifier request could not be created", "request_id", requestID, "alias", model.ID, "error", err) + return "", false + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Request-ID", requestID) + setBackendAuth(req, backend) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + slog.Warn("routing classifier request failed", "request_id", requestID, "alias", model.ID, "error", err) + return "", false + } + defer func() { + _ = resp.Body.Close() + }() + + responseBody, err := io.ReadAll(io.LimitReader(resp.Body, 64*1024)) + if err != nil { + slog.Warn("routing classifier response could not be read", "request_id", requestID, "alias", model.ID, "status", resp.StatusCode, "error", err) + return "", false + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + slog.Warn("routing classifier returned non-success status", "request_id", requestID, "alias", model.ID, "status", resp.StatusCode, "body", trimForLog(string(responseBody), 2048)) + return "", false + } + + target, ok := classifierTargetFromResponse(responseBody, classifier.TargetModels) + if !ok { + slog.Warn("routing classifier returned no allowed target", "request_id", requestID, "alias", model.ID, "allowed_targets", strings.Join(classifier.TargetModels, ","), "body", trimForLog(string(responseBody), 2048)) + return "", false + } + + slog.Info("routing classifier selected target", "request_id", requestID, "alias", model.ID, "target_model", target) + return target, true +} + func extractRoutingFeatures(body []byte) routingFeatures { var payload map[string]any if err := json.Unmarshal(body, &payload); err != nil { @@ -982,6 +1058,116 @@ func numericJSONValue(value any) (int, bool) { } } +func classifierMessages(classifier config.RoutingClassifierConfig, features routingFeatures) []map[string]string { + systemPrompt := strings.TrimSpace(classifier.SystemPrompt) + if systemPrompt == "" { + systemPrompt = defaultClassifierSystemPrompt(classifier.TargetModels) + } + + return []map[string]string{ + {"role": "system", "content": systemPrompt}, + {"role": "user", "content": classifierUserPrompt(features, classifier.TargetModels)}, + } +} + +func defaultClassifierSystemPrompt(targets []string) string { + return "You route coding-agent requests to the cheapest adequate model. Return only JSON with target_model set to one of: " + strings.Join(targets, ", ") + ". Choose the stronger model for architecture, debugging, security, migrations, production risk, broad refactors, ambiguous planning, long-context synthesis, or high requested output. Choose the faster model for routine edits, short questions, status checks, formatting, simple commands, and low-risk local changes." +} + +func classifierUserPrompt(features routingFeatures, targets []string) string { + text := features.Text + if len(text) > 6000 { + head := text[:3000] + tail := text[len(text)-3000:] + text = head + "\n...[middle omitted]...\n" + tail + } + + return fmt.Sprintf( + "/no_think\nAllowed targets: %s\nPrompt chars: %d\nRequested max output tokens: %d\nRequest text:\n%s\n\nReturn only: {\"target_model\":\"\"}", + strings.Join(targets, ", "), + features.PromptChars, + features.RequestedMaxToken, + text, + ) +} + +func classifierMaxTokens(classifier config.RoutingClassifierConfig) int { + if classifier.MaxTokens > 0 { + return classifier.MaxTokens + } + return 64 +} + +func backendChatCompletionsURL(backend config.BackendConfig) (string, error) { + parsed, err := url.Parse(backend.BaseURL) + if err != nil { + return "", err + } + parsed.Path = joinOpenAIPath(parsed.Path, "/v1/chat/completions") + return parsed.String(), nil +} + +func classifierTargetFromResponse(body []byte, allowed []string) (string, bool) { + text := classifierResponseText(body) + if target, ok := classifierTargetFromText(text, allowed); ok { + return target, true + } + return classifierTargetFromText(string(body), allowed) +} + +func classifierResponseText(body []byte) string { + var payload struct { + Choices []struct { + Message map[string]any `json:"message"` + Text string `json:"text"` + } `json:"choices"` + } + if err := json.Unmarshal(body, &payload); err != nil || len(payload.Choices) == 0 { + return "" + } + + var builder strings.Builder + for _, choice := range payload.Choices { + builder.WriteString(choice.Text) + builder.WriteByte('\n') + for _, key := range []string{"content", "reasoning_content"} { + if value, ok := choice.Message[key].(string); ok { + builder.WriteString(value) + builder.WriteByte('\n') + } + } + } + return builder.String() +} + +func classifierTargetFromText(text string, allowed []string) (string, bool) { + var payload struct { + TargetModel string `json:"target_model"` + } + if err := json.Unmarshal([]byte(strings.TrimSpace(text)), &payload); err == nil { + for _, target := range allowed { + if payload.TargetModel == target { + return target, true + } + } + } + + for _, target := range allowed { + if strings.Contains(text, target) { + return target, true + } + } + return "", false +} + +func trimForLog(value string, limit int) string { + value = strings.TrimSpace(value) + if len(value) <= limit { + return value + } + return value[:limit] + "...[truncated]" +} + func rewriteModel(body []byte, targetModel string, maxOutputTokens int) ([]byte, error) { var payload map[string]any if err := json.Unmarshal(body, &payload); err != nil { diff --git a/internal/server/server_test.go b/internal/server/server_test.go index cb2efaf..6801e75 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -218,6 +218,157 @@ func TestRoutingRuleFallsBackToDefaultTarget(t *testing.T) { } } +func TestRoutingClassifierSelectsTarget(t *testing.T) { + t.Parallel() + + var classifierModel string + classifier := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload struct { + Model string `json:"model"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Errorf("decode classifier request: %v", err) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + classifierModel = payload.Model + writeJSON(w, http.StatusOK, map[string]any{ + "choices": []map[string]any{{ + "message": map[string]string{ + "role": "assistant", + "content": `{"target_model":"deep-model"}`, + }, + }}, + }) + })) + t.Cleanup(classifier.Close) + + var backendModel string + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload struct { + Model string `json:"model"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Errorf("decode backend request: %v", err) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + backendModel = payload.Model + writeJSON(w, http.StatusOK, map[string]string{"model": payload.Model}) + })) + t.Cleanup(backend.Close) + + srv, err := New(config.Config{ + Server: config.ServerConfig{Address: "127.0.0.1:0"}, + Models: []config.ModelConfig{{ + ID: "local-coder-auto", + Backend: "lmstudio", + TargetModel: "fast-model", + Routing: config.RoutingConfig{ + Classifier: config.RoutingClassifierConfig{ + Backend: "classifier", + Model: "classifier-model", + TargetModels: []string{"fast-model", "deep-model"}, + }, + }, + }}, + Backends: []config.BackendConfig{ + {ID: "lmstudio", BaseURL: backend.URL}, + {ID: "classifier", BaseURL: classifier.URL}, + }, + }) + if err != nil { + t.Fatalf("create server: %v", err) + } + + req := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + strings.NewReader(`{"model":"local-coder-auto","messages":[{"role":"user","content":"plan a tricky migration"}]}`), + ) + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("unexpected status: %d", rec.Code) + } + if classifierModel != "classifier-model" { + t.Fatalf("unexpected classifier model: %q", classifierModel) + } + if backendModel != "deep-model" { + t.Fatalf("unexpected backend model: %q", backendModel) + } +} + +func TestRoutingClassifierFallsBackToDefaultTarget(t *testing.T) { + t.Parallel() + + classifier := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + writeJSON(w, http.StatusOK, map[string]any{ + "choices": []map[string]any{{ + "message": map[string]string{ + "role": "assistant", + "content": `{"target_model":"unknown-model"}`, + }, + }}, + }) + })) + t.Cleanup(classifier.Close) + + var backendModel string + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload struct { + Model string `json:"model"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Errorf("decode backend request: %v", err) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + backendModel = payload.Model + writeJSON(w, http.StatusOK, map[string]string{"model": payload.Model}) + })) + t.Cleanup(backend.Close) + + srv, err := New(config.Config{ + Server: config.ServerConfig{Address: "127.0.0.1:0"}, + Models: []config.ModelConfig{{ + ID: "local-coder-auto", + Backend: "lmstudio", + TargetModel: "fast-model", + Routing: config.RoutingConfig{ + Classifier: config.RoutingClassifierConfig{ + Backend: "classifier", + Model: "classifier-model", + TargetModels: []string{"fast-model", "deep-model"}, + }, + }, + }}, + Backends: []config.BackendConfig{ + {ID: "lmstudio", BaseURL: backend.URL}, + {ID: "classifier", BaseURL: classifier.URL}, + }, + }) + if err != nil { + t.Fatalf("create server: %v", err) + } + + req := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + strings.NewReader(`{"model":"local-coder-auto","messages":[{"role":"user","content":"say ok"}]}`), + ) + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("unexpected status: %d", rec.Code) + } + if backendModel != "fast-model" { + t.Fatalf("unexpected backend model: %q", backendModel) + } +} + func TestRewriteModelClampsOversizedTargetAliasOutput(t *testing.T) { t.Parallel()