diff --git a/pkg/backend/oneshot.go b/pkg/backend/oneshot.go index f295438..8a1c16c 100644 --- a/pkg/backend/oneshot.go +++ b/pkg/backend/oneshot.go @@ -49,7 +49,7 @@ func Generate(ctx context.Context, req GenerateRequest) (string, error) { // RouteRequest is the input for model routing. type RouteRequest struct { Task string `json:"task"` - Backend string `json:"backend"` + Backend string `json:"backend"` // optional: only consider models from this backend } // RouteResult is the output of model routing. @@ -58,9 +58,9 @@ type RouteResult struct { Model string } -// Route classifies a task description and selects the most appropriate model. -// Uses a cheap classifier model (OLLIE_ROUTE_BACKEND/OLLIE_ROUTE_MODEL) to pick -// between OLLIE_ROUTE_MODEL_FAST and OLLIE_ROUTE_MODEL_POWER. +// Route classifies a task description and selects the most appropriate model +// from the actually available models. Fetches the model list from configured +// backends, presents them to a cheap classifier, and returns the selection. func Route(ctx context.Context, req RouteRequest) (RouteResult, error) { if req.Task == "" { return RouteResult{}, fmt.Errorf("provide a task description") @@ -81,34 +81,39 @@ func Route(ctx context.Context, req RouteRequest) (RouteResult, error) { } cbe.SetModel(classifyModel) - // Target backend - backendName := req.Backend - if backendName == "" { - backendName = os.Getenv("OLLIE_BACKEND") + // Fetch available models + type modelEntry struct { + backend string + model string } - if backendName == "" { - backendName = "openrouter" - } - backendLower := strings.ToLower(backendName) + var entries []modelEntry + var modelList strings.Builder - // Candidate models - fastModel := os.Getenv("OLLIE_ROUTE_MODEL_FAST") - powerModel := os.Getenv("OLLIE_ROUTE_MODEL_POWER") - if fastModel == "" { - fastModel = "mistralai/devstral-small" + backends := Backends() + if req.Backend != "" { + backends = []string{req.Backend} } - if powerModel == "" { - powerModel = "anthropic/claude-sonnet-4-20250514" + for _, name := range backends { + b, err := NewWithName(name) + if err != nil { + continue + } + for _, m := range b.Models(ctx) { + entries = append(entries, modelEntry{backend: name, model: m}) + fmt.Fprintf(&modelList, "%d. %s/%s\n", len(entries), name, m) + } + } + if len(entries) == 0 { + return RouteResult{}, fmt.Errorf("no models available") } - classifyPrompt := fmt.Sprintf(`You are a task router. Given a task, reply with ONLY the number 1 or 2. + classifyPrompt := fmt.Sprintf(`You are a task router. Given a task and a numbered list of available models, select the single best model for the task. -1 = fast/cheap model, good for: simple edits, single-file changes, straightforward bugs, docs, tests, standard API usage -2 = powerful/expensive model, good for: multi-file architecture, security-sensitive code, complex algorithms, system design, subtle debugging, performance optimization +Available models: +%s +Reply with ONLY the number of the best model. Nothing else. -Reply with ONLY the digit 1 or 2. Nothing else. - -Task: %s`, req.Task) +Task: %s`, modelList.String(), req.Task) msgs := []Message{ {Role: "user", Content: classifyPrompt}, @@ -123,11 +128,14 @@ Task: %s`, req.Task) } result := strings.TrimSpace(sb.String()) - // Match: if the response contains "2", pick the powerful model - model := fastModel - if strings.Contains(result, "2") { - model = powerModel + // Parse the number from the response + var choice int + fmt.Sscanf(result, "%d", &choice) + if choice < 1 || choice > len(entries) { + // Default to first model if parsing fails + choice = 1 } + picked := entries[choice-1] - return RouteResult{Backend: backendLower, Model: model}, nil + return RouteResult{Backend: picked.backend, Model: picked.model}, nil }