134 lines
3.7 KiB
Go
134 lines
3.7 KiB
Go
// Package backend — stateless one-shot operations.
|
|
// These functions provide request/response LLM access without session state.
|
|
package backend
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
)
|
|
|
|
// GenerateRequest is the input for a stateless one-shot generation.
|
|
type GenerateRequest struct {
|
|
Prompt string `json:"prompt"`
|
|
System string `json:"system"`
|
|
Backend string `json:"backend"`
|
|
Model string `json:"model"`
|
|
}
|
|
|
|
// Generate performs a stateless one-shot LLM generation with no session context.
|
|
// Returns the full model response text.
|
|
func Generate(ctx context.Context, req GenerateRequest) (string, error) {
|
|
if req.Prompt == "" {
|
|
return "", fmt.Errorf("empty prompt")
|
|
}
|
|
be, err := NewWithName(req.Backend)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if req.Model != "" {
|
|
be.SetModel(req.Model)
|
|
}
|
|
var msgs []Message
|
|
if req.System != "" {
|
|
msgs = append(msgs, Message{Role: "system", Content: req.System})
|
|
}
|
|
msgs = append(msgs, Message{Role: "user", Content: req.Prompt})
|
|
ch, err := be.ChatStream(ctx, msgs, nil, GenerationParams{})
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
var sb strings.Builder
|
|
for ev := range ch {
|
|
sb.WriteString(ev.Content)
|
|
}
|
|
return sb.String(), nil
|
|
}
|
|
|
|
// RouteRequest is the input for model routing.
|
|
type RouteRequest struct {
|
|
Task string `json:"task"`
|
|
Backend string `json:"backend"`
|
|
}
|
|
|
|
// RouteResult is the output of model routing.
|
|
type RouteResult struct {
|
|
Backend string
|
|
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.
|
|
func Route(ctx context.Context, req RouteRequest) (RouteResult, error) {
|
|
if req.Task == "" {
|
|
return RouteResult{}, fmt.Errorf("provide a task description")
|
|
}
|
|
|
|
// Classifier backend/model
|
|
classifyBackend := os.Getenv("OLLIE_ROUTE_BACKEND")
|
|
if classifyBackend == "" {
|
|
classifyBackend = "ollama"
|
|
}
|
|
classifyModel := os.Getenv("OLLIE_ROUTE_MODEL")
|
|
if classifyModel == "" {
|
|
classifyModel = "qwen3:8b"
|
|
}
|
|
cbe, err := NewWithName(classifyBackend)
|
|
if err != nil {
|
|
return RouteResult{}, fmt.Errorf("route backend: %w", err)
|
|
}
|
|
cbe.SetModel(classifyModel)
|
|
|
|
// Target backend
|
|
backendName := req.Backend
|
|
if backendName == "" {
|
|
backendName = os.Getenv("OLLIE_BACKEND")
|
|
}
|
|
if backendName == "" {
|
|
backendName = "openrouter"
|
|
}
|
|
backendLower := strings.ToLower(backendName)
|
|
|
|
// Candidate models
|
|
fastModel := os.Getenv("OLLIE_ROUTE_MODEL_FAST")
|
|
powerModel := os.Getenv("OLLIE_ROUTE_MODEL_POWER")
|
|
if fastModel == "" {
|
|
fastModel = "mistralai/devstral-small"
|
|
}
|
|
if powerModel == "" {
|
|
powerModel = "anthropic/claude-sonnet-4-20250514"
|
|
}
|
|
|
|
classifyPrompt := fmt.Sprintf(`You are a task router. Given a task, reply with ONLY the number 1 or 2.
|
|
|
|
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
|
|
|
|
Reply with ONLY the digit 1 or 2. Nothing else.
|
|
|
|
Task: %s`, req.Task)
|
|
|
|
msgs := []Message{
|
|
{Role: "user", Content: classifyPrompt},
|
|
}
|
|
ch, err := cbe.ChatStream(ctx, msgs, nil, GenerationParams{})
|
|
if err != nil {
|
|
return RouteResult{}, fmt.Errorf("route classification failed: %w", err)
|
|
}
|
|
var sb strings.Builder
|
|
for ev := range ch {
|
|
sb.WriteString(ev.Content)
|
|
}
|
|
result := strings.TrimSpace(sb.String())
|
|
|
|
// Match: if the response contains "2", pick the powerful model
|
|
model := fastModel
|
|
if strings.Contains(result, "2") {
|
|
model = powerModel
|
|
}
|
|
|
|
return RouteResult{Backend: backendLower, Model: model}, nil
|
|
}
|