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:
parent
8a995c1b13
commit
b4c88a33d3
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
Reference in New Issue