route: fetch real models from backends, pick by number

Route() now queries all configured backends for their model lists,
presents them numbered to the classifier, and parses the number.
No more hardcoded OLLIE_ROUTE_MODEL_FAST/POWER env vars.
This commit is contained in:
ollie 2026-07-18 17:09:55 +02:00
parent 8a995c1b13
commit b4c88a33d3
1 changed files with 38 additions and 30 deletions

View File

@ -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
}