diff --git a/server/server.go b/server/server.go index 473ce6c..e9f7521 100644 --- a/server/server.go +++ b/server/server.go @@ -888,124 +888,30 @@ func (s *Server) write(cs *connState, fc *plan9.Fcall) *plan9.Fcall { // /generate: stateless one-shot generation (per-fid state, like /net/dns) if f.path == "/generate" { cs.mu.Unlock() - var req struct { - Prompt string `json:"prompt"` - System string `json:"system"` - Backend string `json:"backend"` - Model string `json:"model"` - } + var req backend.GenerateRequest if err := json.Unmarshal(fc.Data, &req); err != nil { - // Plain text mode: treat entire write as prompt req.Prompt = strings.TrimSpace(string(fc.Data)) } - if req.Prompt == "" { - return errFcall(fc, "empty prompt") - } - be, err := backend.NewWithName(req.Backend) + result, err := backend.Generate(context.Background(), req) if err != nil { return errFcall(fc, err.Error()) } - if req.Model != "" { - be.SetModel(req.Model) - } - var msgs []backend.Message - if req.System != "" { - msgs = append(msgs, backend.Message{Role: "system", Content: req.System}) - } - msgs = append(msgs, backend.Message{Role: "user", Content: req.Prompt}) - ch, err := be.ChatStream(context.Background(), msgs, nil, backend.GenerationParams{}) - if err != nil { - return errFcall(fc, err.Error()) - } - var sb strings.Builder - for ev := range ch { - sb.WriteString(ev.Content) - } - f.generateResult = []byte(sb.String()) + f.generateResult = []byte(result) return &plan9.Fcall{Type: plan9.Rwrite, Tag: fc.Tag, Count: uint32(len(fc.Data))} } // /route: classify task and select appropriate backend+model (per-fid state, like /net/dns) - // Write a task description; a cheap model selects the best backend+model for it. if f.path == "/route" { cs.mu.Unlock() - var req struct { - Task string `json:"task"` - Backend string `json:"backend"` // optional: constrain to this backend - } + var req backend.RouteRequest if err := json.Unmarshal(fc.Data, &req); err != nil { - // Plain text mode: treat entire write as task description req.Task = strings.TrimSpace(string(fc.Data)) } - if req.Task == "" { - return errFcall(fc, "provide a task description") - } - - // Classify via LLM - classifyBackend := os.Getenv("OLLIE_ROUTE_BACKEND") - if classifyBackend == "" { - classifyBackend = "ollama" - } - classifyModel := os.Getenv("OLLIE_ROUTE_MODEL") - if classifyModel == "" { - classifyModel = "qwen3:8b" - } - cbe, err := backend.NewWithName(classifyBackend) + result, err := backend.Route(context.Background(), req) if err != nil { - return errFcall(fc, "route backend: "+err.Error()) + return errFcall(fc, err.Error()) } - cbe.SetModel(classifyModel) - - // Build the list of available models from env - backendName := req.Backend - if backendName == "" { - backendName = os.Getenv("OLLIE_BACKEND") - } - if backendName == "" { - backendName = "openrouter" - } - backendLower := strings.ToLower(backendName) - - // Gather configured models for this backend - 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 description, select the most appropriate model. - -Available models: -1. %s — fast, cheap, good for: simple edits, single-file changes, straightforward bugs, docs, tests, standard API usage -2. %s — powerful, expensive, good for: multi-file architecture, security-sensitive code, complex algorithms, system design, subtle debugging, performance optimization - -Reply with ONLY the model name, nothing else. - -Task: %s`, fastModel, powerModel, req.Task) - - msgs := []backend.Message{ - {Role: "user", Content: classifyPrompt}, - } - ch, err := cbe.ChatStream(context.Background(), msgs, nil, backend.GenerationParams{}) - if err != nil { - return errFcall(fc, "route classification failed: "+err.Error()) - } - var sb strings.Builder - for ev := range ch { - sb.WriteString(ev.Content) - } - result := strings.TrimSpace(sb.String()) - - // Match result to one of the configured models - model := fastModel - if strings.Contains(result, powerModel) { - model = powerModel - } - - f.routeResult = []byte(fmt.Sprintf("backend=%s\nmodel=%s\n", backendLower, model)) + f.routeResult = []byte(fmt.Sprintf("backend=%s\nmodel=%s\n", result.Backend, result.Model)) return &plan9.Fcall{Type: plan9.Rwrite, Tag: fc.Tag, Count: uint32(len(fc.Data))} }