diff --git a/README.md b/README.md index 7a53ee2..632c05b 100644 --- a/README.md +++ b/README.md @@ -131,6 +131,17 @@ models: max_concurrent_requests: 2 max_queue_size: 4 queue_timeout: 30s + routing: + rules: + - 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 ensure: mode: command command: diff --git a/configs/router.example.yaml b/configs/router.example.yaml index 487cb72..69a8efe 100644 --- a/configs/router.example.yaml +++ b/configs/router.example.yaml @@ -34,6 +34,28 @@ models: - /usr/local/bin/lmstudio-load-profile - local-coder-large timeout: 45s + - id: local-coder-auto + name: Local Coder Auto + backend: lmstudio + target_model: qwen3-coder-30b-a3b-instruct + context_window: 65536 + max_output_tokens: 4096 + tool_calls: true + max_concurrent_requests: 1 + max_queue_size: 4 + queue_timeout: 2m + routing: + rules: + - 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 backends: - id: lmstudio diff --git a/docs/architecture.md b/docs/architecture.md index d9d192e..dbf6b91 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -36,8 +36,34 @@ The first router is deliberately simple: - The router rewrites the request to the configured backend model. - The backend handles inference. -Future routing can add deterministic policy, queueing, health-aware selection, -RouteLLM-style strong/weak model routing, and second-pass review workflows. +Model aliases can also use deterministic request routing rules. The alias still +has a default `target_model`, but ordered rules can choose a different target +from cheap request metadata before the request is proxied: + +```yaml +models: + - id: local-coder-auto + backend: lmstudio + target_model: qwen3-coder-30b-a3b-instruct + routing: + rules: + - 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 +``` + +Rules are evaluated in order. A rule matches only when all configured numeric +conditions pass; `any_keywords` is an OR condition across the lower-cased +request text. If no rule matches, the alias uses its default `target_model`. +Future routing can add health-aware selection, RouteLLM-style strong/weak model +routing, and second-pass review workflows. ## Request Limits diff --git a/internal/config/config.go b/internal/config/config.go index c186f6c..739be7c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -5,6 +5,7 @@ import ( "fmt" "net/url" "os" + "strings" "time" "gopkg.in/yaml.v3" @@ -34,6 +35,7 @@ type ModelConfig struct { MaxQueueSize int `yaml:"max_queue_size"` QueueTimeout string `yaml:"queue_timeout"` Ensure EnsureConfig + Routing RoutingConfig `yaml:"routing"` } type BackendConfig struct { @@ -49,6 +51,20 @@ type EnsureConfig struct { Timeout string `yaml:"timeout"` } +type RoutingConfig struct { + Rules []RoutingRuleConfig `yaml:"rules"` +} + +type RoutingRuleConfig struct { + ID string `yaml:"id"` + TargetModel string `yaml:"target_model"` + MinPromptChars int `yaml:"min_prompt_chars"` + MaxPromptChars int `yaml:"max_prompt_chars"` + MinOutputTokens int `yaml:"min_output_tokens"` + MaxOutputTokens int `yaml:"max_output_tokens"` + AnyKeywords []string `yaml:"any_keywords"` +} + type CommandArgs []string func (args *CommandArgs) UnmarshalYAML(value *yaml.Node) error { @@ -152,6 +168,9 @@ func (cfg Config) Validate() error { if err := model.Ensure.Validate(); err != nil { return fmt.Errorf("model %q ensure is invalid: %w", model.ID, err) } + if err := model.Routing.Validate(); err != nil { + return fmt.Errorf("model %q routing is invalid: %w", model.ID, err) + } if _, ok := models[model.ID]; ok { return fmt.Errorf("model %q is duplicated", model.ID) } @@ -161,6 +180,47 @@ func (cfg Config) Validate() error { return nil } +func (routing RoutingConfig) Validate() error { + for index, rule := range routing.Rules { + if rule.TargetModel == "" { + return fmt.Errorf("rule %d target_model is required", index) + } + if rule.MinPromptChars < 0 { + return fmt.Errorf("rule %d min_prompt_chars must be non-negative", index) + } + if rule.MaxPromptChars < 0 { + return fmt.Errorf("rule %d max_prompt_chars must be non-negative", index) + } + if rule.MinPromptChars > 0 && rule.MaxPromptChars > 0 && rule.MinPromptChars > rule.MaxPromptChars { + return fmt.Errorf("rule %d min_prompt_chars must be less than or equal to max_prompt_chars", index) + } + if rule.MinOutputTokens < 0 { + return fmt.Errorf("rule %d min_output_tokens must be non-negative", index) + } + if rule.MaxOutputTokens < 0 { + return fmt.Errorf("rule %d max_output_tokens must be non-negative", index) + } + if rule.MinOutputTokens > 0 && rule.MaxOutputTokens > 0 && rule.MinOutputTokens > rule.MaxOutputTokens { + return fmt.Errorf("rule %d min_output_tokens must be less than or equal to max_output_tokens", index) + } + hasCondition := rule.MinPromptChars > 0 || + rule.MaxPromptChars > 0 || + rule.MinOutputTokens > 0 || + rule.MaxOutputTokens > 0 || + len(rule.AnyKeywords) > 0 + if !hasCondition { + return fmt.Errorf("rule %d must define at least one condition", index) + } + for keywordIndex, keyword := range rule.AnyKeywords { + if strings.TrimSpace(keyword) == "" { + return fmt.Errorf("rule %d any_keywords[%d] must not be empty", index, keywordIndex) + } + } + } + + return nil +} + func (cfg Config) Model(id string) (ModelConfig, bool) { for _, model := range cfg.Models { if model.ID == id { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 0f1dfa0..4f66ed2 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -70,6 +70,63 @@ func TestValidateQueueSettings(t *testing.T) { } } +func TestValidateRoutingRules(t *testing.T) { + t.Parallel() + + cfg := Config{ + Models: []ModelConfig{{ + ID: "local-coder", + Backend: "lmstudio", + TargetModel: "fast-model", + Routing: RoutingConfig{ + Rules: []RoutingRuleConfig{{ + ID: "large-prompt", + TargetModel: "deep-model", + MinPromptChars: 12000, + AnyKeywords: []string{"refactor"}, + }}, + }, + }}, + Backends: []BackendConfig{{ + ID: "lmstudio", + BaseURL: "http://127.0.0.1:1234/v1", + }}, + } + + if err := cfg.Validate(); err != nil { + t.Fatalf("validate config: %v", err) + } +} + +func TestValidateRoutingRuleRequiresCondition(t *testing.T) { + t.Parallel() + + cfg := Config{ + Models: []ModelConfig{{ + ID: "local-coder", + Backend: "lmstudio", + TargetModel: "fast-model", + Routing: RoutingConfig{ + Rules: []RoutingRuleConfig{{ + TargetModel: "deep-model", + }}, + }, + }}, + Backends: []BackendConfig{{ + ID: "lmstudio", + BaseURL: "http://127.0.0.1:1234/v1", + }}, + } + + err := cfg.Validate() + if err == nil { + t.Fatal("expected validation error") + } + if !strings.Contains(err.Error(), "routing") { + t.Fatalf("expected routing error, got: %v", err) + } +} + func TestLoadEnsureCommandString(t *testing.T) { t.Parallel() diff --git a/internal/server/server.go b/internal/server/server.go index 904488a..58cc460 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -88,6 +88,8 @@ func (s *Server) handleModels(w http.ResponseWriter, _ *http.Request) { OwnedBy string `json:"owned_by"` Name string `json:"name,omitempty"` ContextWindow int `json:"context_window,omitempty"` + MaxOutput int `json:"max_output_tokens,omitempty"` + ToolCall bool `json:"tool_call,omitempty"` TargetModel string `json:"target_model,omitempty"` } @@ -99,6 +101,8 @@ func (s *Server) handleModels(w http.ResponseWriter, _ *http.Request) { OwnedBy: "devrail-router", Name: model.Name, ContextWindow: model.ContextWindow, + MaxOutput: model.MaxOutputTokens, + ToolCall: model.ToolCalls, TargetModel: model.TargetModel, }) } @@ -147,12 +151,19 @@ func (s *Server) proxyOpenAI(w http.ResponseWriter, r *http.Request) { return } - body, err = rewriteModel(body, model.TargetModel) + targetModel, routeRule := selectTargetModel(body, model) + targetMaxOutputTokens := 0 + if targetAlias, ok := s.cfg.Model(targetModel); ok { + targetMaxOutputTokens = targetAlias.MaxOutputTokens + } + body, err = rewriteModel(body, targetModel, targetMaxOutputTokens) if err != nil { writeOpenAIError(w, http.StatusBadRequest, err.Error(), "invalid_request_error", "invalid_request") s.metrics.record(requestMetricsFromModel(model, backend, http.StatusBadRequest, time.Now())) return } + routedModel := model + routedModel.TargetModel = targetModel r.Body = io.NopCloser(bytes.NewReader(body)) r.ContentLength = int64(len(body)) @@ -175,7 +186,7 @@ func (s *Server) proxyOpenAI(w http.ResponseWriter, r *http.Request) { setBackendAuth(req, backend) } proxy.ModifyResponse = func(resp *http.Response) error { - instrumentBackendResponse(resp, started, waited, model, backend, requestID, s.metrics) + instrumentBackendResponse(resp, started, waited, routedModel, backend, requestID, s.metrics) return nil } proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, proxyErr error) { @@ -185,16 +196,17 @@ func (s *Server) proxyOpenAI(w http.ResponseWriter, r *http.Request) { "path", req.URL.Path, "request_id", requestID, "alias", model.ID, - "target_model", model.TargetModel, + "target_model", targetModel, + "route_rule", routeRule, "backend", backend.ID, "duration_ms", time.Since(started).Milliseconds(), "error", proxyErr, ) - s.metrics.record(requestMetricsFromModel(model, backend, http.StatusBadGateway, started)) + s.metrics.record(requestMetricsFromModel(routedModel, backend, http.StatusBadGateway, started)) writeOpenAIError(rw, http.StatusBadGateway, "backend request failed", "devrail_backend_error", "backend_request_failed") } - slog.Info("routing request", "request_id", requestID, "alias", model.ID, "target_model", model.TargetModel, "backend", backend.ID) + slog.Info("routing request", "request_id", requestID, "alias", model.ID, "target_model", targetModel, "route_rule", routeRule, "backend", backend.ID) proxy.ServeHTTP(w, r) } @@ -863,13 +875,121 @@ func requestModel(r *http.Request) (string, []byte, error) { return model, body, nil } -func rewriteModel(body []byte, targetModel string) ([]byte, error) { +type routingFeatures struct { + PromptChars int + Text string + RequestedMaxToken int +} + +func selectTargetModel(body []byte, model config.ModelConfig) (string, string) { + if len(model.Routing.Rules) == 0 { + return model.TargetModel, "default" + } + + features := extractRoutingFeatures(body) + for _, rule := range model.Routing.Rules { + if routingRuleMatches(rule, features) { + if rule.ID != "" { + return rule.TargetModel, rule.ID + } + return rule.TargetModel, rule.TargetModel + } + } + + return model.TargetModel, "default" +} + +func extractRoutingFeatures(body []byte) routingFeatures { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + return routingFeatures{} + } + + text := requestText(payload) + return routingFeatures{ + PromptChars: len(text), + Text: strings.ToLower(text), + RequestedMaxToken: requestMaxTokens(payload), + } +} + +func routingRuleMatches(rule config.RoutingRuleConfig, features routingFeatures) bool { + if rule.MinPromptChars > 0 && features.PromptChars < rule.MinPromptChars { + return false + } + if rule.MaxPromptChars > 0 && features.PromptChars > rule.MaxPromptChars { + return false + } + if rule.MinOutputTokens > 0 && features.RequestedMaxToken < rule.MinOutputTokens { + return false + } + if rule.MaxOutputTokens > 0 && features.RequestedMaxToken > 0 && features.RequestedMaxToken > rule.MaxOutputTokens { + return false + } + if len(rule.AnyKeywords) > 0 { + for _, keyword := range rule.AnyKeywords { + if strings.Contains(features.Text, strings.ToLower(strings.TrimSpace(keyword))) { + return true + } + } + return false + } + + return true +} + +func requestText(payload map[string]any) string { + var builder strings.Builder + appendTextValue(&builder, payload["messages"]) + appendTextValue(&builder, payload["input"]) + appendTextValue(&builder, payload["prompt"]) + return builder.String() +} + +func appendTextValue(builder *strings.Builder, value any) { + switch typed := value.(type) { + case string: + builder.WriteString(typed) + builder.WriteByte('\n') + case []any: + for _, item := range typed { + appendTextValue(builder, item) + } + case map[string]any: + for _, key := range []string{"role", "content", "text", "input_text", "prompt"} { + appendTextValue(builder, typed[key]) + } + } +} + +func requestMaxTokens(payload map[string]any) int { + for _, key := range []string{"max_completion_tokens", "max_tokens"} { + if value, ok := numericJSONValue(payload[key]); ok { + return value + } + } + return 0 +} + +func numericJSONValue(value any) (int, bool) { + switch typed := value.(type) { + case float64: + return int(typed), true + case int: + return typed, true + default: + return 0, false + } +} + +func rewriteModel(body []byte, targetModel string, maxOutputTokens int) ([]byte, error) { var payload map[string]any if err := json.Unmarshal(body, &payload); err != nil { return nil, fmt.Errorf("parse request body: %w", err) } payload["model"] = targetModel + clampRequestMaxTokens(payload, maxOutputTokens) body, err := json.Marshal(payload) if err != nil { return nil, fmt.Errorf("encode request body: %w", err) @@ -878,6 +998,19 @@ func rewriteModel(body []byte, targetModel string) ([]byte, error) { return body, nil } +func clampRequestMaxTokens(payload map[string]any, maxOutputTokens int) { + if maxOutputTokens <= 0 { + return + } + for _, key := range []string{"max_completion_tokens", "max_tokens"} { + value, ok := numericJSONValue(payload[key]) + if !ok || value <= maxOutputTokens { + continue + } + payload[key] = maxOutputTokens + } +} + func setBackendAuth(req *http.Request, backend config.BackendConfig) { if backend.APIKeyEnv == "" { return diff --git a/internal/server/server_test.go b/internal/server/server_test.go index fb945e4..cb2efaf 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -124,6 +124,123 @@ func TestRequestIDHeaderIsPropagated(t *testing.T) { } } +func TestRoutingRuleSelectsTargetByPromptSize(t *testing.T) { + t.Parallel() + + 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 := testServerWithBackend(t, backend.URL, config.ModelConfig{ + ID: "local-coder-auto", + Backend: "lmstudio", + TargetModel: "fast-model", + Routing: config.RoutingConfig{ + Rules: []config.RoutingRuleConfig{{ + ID: "large-prompt", + TargetModel: "deep-model", + MinPromptChars: 20, + }}, + }, + }) + req := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + strings.NewReader(`{"model":"local-coder-auto","messages":[{"role":"user","content":"please analyze this larger prompt"}]}`), + ) + rec := httptest.NewRecorder() + + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("unexpected status: %d", rec.Code) + } + if backendModel != "deep-model" { + t.Fatalf("unexpected backend model: %q", backendModel) + } +} + +func TestRoutingRuleFallsBackToDefaultTarget(t *testing.T) { + t.Parallel() + + 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 := testServerWithBackend(t, backend.URL, config.ModelConfig{ + ID: "local-coder-auto", + Backend: "lmstudio", + TargetModel: "fast-model", + Routing: config.RoutingConfig{ + Rules: []config.RoutingRuleConfig{{ + ID: "hard-work", + TargetModel: "deep-model", + AnyKeywords: []string{"architecture"}, + }}, + }, + }) + 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() + + body, err := rewriteModel([]byte(`{"model":"local-coder-auto","messages":[],"max_tokens":32768}`), "local-coder-fast", 8192) + if err != nil { + t.Fatalf("rewrite model: %v", err) + } + + var payload struct { + Model string `json:"model"` + MaxToken int `json:"max_tokens"` + } + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("decode rewritten payload: %v", err) + } + if payload.Model != "local-coder-fast" { + t.Fatalf("unexpected model: %q", payload.Model) + } + if payload.MaxToken != 8192 { + t.Fatalf("unexpected max_tokens: %d", payload.MaxToken) + } +} + func TestJoinOpenAIPathAvoidsDuplicateVersionPrefix(t *testing.T) { t.Parallel()