142 lines
3.7 KiB
Go
142 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"` // optional: only consider models from this 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
|
|
// 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")
|
|
}
|
|
|
|
// 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)
|
|
|
|
// Fetch available models
|
|
type modelEntry struct {
|
|
backend string
|
|
model string
|
|
}
|
|
var entries []modelEntry
|
|
var modelList strings.Builder
|
|
|
|
backends := Backends()
|
|
if req.Backend != "" {
|
|
backends = []string{req.Backend}
|
|
}
|
|
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 and a numbered list of available models, select the single best model for the task.
|
|
|
|
Available models:
|
|
%s
|
|
Reply with ONLY the number of the best model. Nothing else.
|
|
|
|
Task: %s`, modelList.String(), 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())
|
|
|
|
// 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: picked.backend, Model: picked.model}, nil
|
|
}
|