From e7ab48e8ea30c5427815cd805106052df952b3e6 Mon Sep 17 00:00:00 2001 From: weichen Date: Thu, 1 Oct 2026 10:10:02 +0000 Subject: [PATCH] Add native provider protocols and retry Why: The runtime could only speak OpenAI chat completions, so thinking levels, images, and the other vendor wires had nowhere to go. Evidence: go test -count=1 ./runtime/llm/ Impact: Existing chat-completions calls stay the same unless a thinking level or image is set. Overflow errors are not retried. Refs: feature/pi-headless N2 Co-authored-by: Cursor --- runtime/llm/client.go | 5 +- runtime/llm/providers.go | 431 ++++++++++++++++++++++++++++++++++ runtime/llm/providers_test.go | 115 +++++++++ runtime/llm/sse.go | 14 +- runtime/llm/types.go | 13 + 5 files changed, 575 insertions(+), 3 deletions(-) create mode 100644 runtime/llm/providers.go create mode 100644 runtime/llm/providers_test.go diff --git a/runtime/llm/client.go b/runtime/llm/client.go index 5e5d236..e07db59 100644 --- a/runtime/llm/client.go +++ b/runtime/llm/client.go @@ -53,6 +53,9 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool if opts.MaxTokens > 0 { body["max_tokens"] = opts.MaxTokens } + if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" { + body["reasoning_effort"] = effort + } if len(tools) > 0 { body["tools"] = toProviderTools(tools) body["tool_choice"] = "auto" @@ -125,7 +128,7 @@ func toProviderMessages(messages []Message) []map[string]any { default: out = append(out, map[string]any{ "role": string(m.Role), - "content": m.Content, + "content": messageContent(m), }) } } diff --git a/runtime/llm/providers.go b/runtime/llm/providers.go new file mode 100644 index 0000000..0e8dc9b --- /dev/null +++ b/runtime/llm/providers.go @@ -0,0 +1,431 @@ +package llm + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" + "time" +) + +// Thinking levels, from off to max. Unknown values are rejected by NormalizeThinking. +var ThinkingLevels = []string{"off", "minimal", "low", "medium", "high", "xhigh", "max"} + +// NormalizeThinking returns the canonical level. Empty stays empty. +func NormalizeThinking(level string) (string, error) { + if level == "" { + return "", nil + } + level = strings.ToLower(strings.TrimSpace(level)) + for _, known := range ThinkingLevels { + if level == known { + return level, nil + } + } + return "", fmt.Errorf("llm: unknown thinking level %q", level) +} + +// OpenAIReasoningEffort maps a thinking level onto OpenAI reasoning_effort. +// off and empty produce no parameter. +func OpenAIReasoningEffort(level string) string { + switch level { + case "minimal": + return "minimal" + case "low": + return "low" + case "medium": + return "medium" + case "high", "xhigh", "max": + return "high" + default: + return "" + } +} + +// AnthropicThinkingBudget maps a thinking level onto a token budget. +// off and empty disable thinking. +func AnthropicThinkingBudget(level string) int { + switch level { + case "minimal": + return 1024 + case "low": + return 2048 + case "medium": + return 4096 + case "high": + return 8192 + case "xhigh": + return 16384 + case "max": + return 32768 + default: + return 0 + } +} + +// GoogleThinkingBudget maps a thinking level onto Gemini thinkingBudget. +func GoogleThinkingBudget(level string) int { + switch level { + case "minimal", "low": + return 1024 + case "medium": + return 4096 + case "high": + return 8192 + case "xhigh", "max": + return 24576 + default: + return 0 + } +} + +// ModelRef is provider/model with an optional :thinking suffix. +type ModelRef struct { + Provider string + Model string + Thinking string +} + +// ParseModelRef accepts "model", "provider/model", and "provider/model:level". +func ParseModelRef(spec string) (ModelRef, error) { + spec = strings.TrimSpace(spec) + if spec == "" { + return ModelRef{}, fmt.Errorf("llm: empty model ref") + } + thinking := "" + if i := strings.LastIndex(spec, ":"); i >= 0 { + thinking = spec[i+1:] + spec = spec[:i] + norm, err := NormalizeThinking(thinking) + if err != nil { + return ModelRef{}, err + } + thinking = norm + } + provider, model := "", spec + if i := strings.Index(spec, "/"); i >= 0 { + provider = spec[:i] + model = spec[i+1:] + } + if model == "" { + return ModelRef{}, fmt.Errorf("llm: model id missing in %q", spec) + } + return ModelRef{Provider: provider, Model: model, Thinking: thinking}, nil +} + +// Protocol names the wire format a provider speaks. +type Protocol string + +const ( + ProtocolOpenAIChat Protocol = "openai-chat" + ProtocolOpenAIResponses Protocol = "openai-responses" + ProtocolAnthropic Protocol = "anthropic" + ProtocolGoogle Protocol = "google" + ProtocolAzure Protocol = "azure" +) + +// Provider is one configured endpoint. +type Provider struct { + Name string `json:"name"` + Protocol Protocol `json:"protocol"` + BaseURL string `json:"baseURL"` + APIKeyEnv string `json:"apiKeyEnv,omitempty"` + APIKey string `json:"-"` +} + +// Presets are OpenAI-compatible or native endpoints that only need an API key. +func Presets() []Provider { + return []Provider{ + {Name: "anthropic", Protocol: ProtocolAnthropic, BaseURL: "https://api.anthropic.com", APIKeyEnv: "ANTHROPIC_API_KEY"}, + {Name: "openai", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"}, + {Name: "openai-responses", Protocol: ProtocolOpenAIResponses, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"}, + {Name: "google", Protocol: ProtocolGoogle, BaseURL: "https://generativelanguage.googleapis.com/v1beta", APIKeyEnv: "GEMINI_API_KEY"}, + {Name: "azure", Protocol: ProtocolAzure, BaseURL: "", APIKeyEnv: "AZURE_OPENAI_API_KEY"}, + {Name: "deepseek", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.deepseek.com/v1", APIKeyEnv: "DEEPSEEK_API_KEY"}, + {Name: "minimax", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.minimaxi.com/v1", APIKeyEnv: "MINIMAX_API_KEY"}, + {Name: "openrouter", Protocol: ProtocolOpenAIChat, BaseURL: "https://openrouter.ai/api/v1", APIKeyEnv: "OPENROUTER_API_KEY"}, + {Name: "groq", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.groq.com/openai/v1", APIKeyEnv: "GROQ_API_KEY"}, + } +} + +// modelsFile is the on-disk shape of models.json. +type modelsFile struct { + Providers []Provider `json:"providers"` +} + +// LoadModelsJSON reads custom providers. Missing file yields an empty list. +func LoadModelsJSON(path string) ([]Provider, error) { + raw, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, err + } + var file modelsFile + if err := json.Unmarshal(raw, &file); err != nil { + return nil, err + } + for i := range file.Providers { + if file.Providers[i].APIKey == "" && file.Providers[i].APIKeyEnv != "" { + file.Providers[i].APIKey = os.Getenv(file.Providers[i].APIKeyEnv) + } + } + return file.Providers, nil +} + +// Lookup finds a preset or a custom provider by name. +func Lookup(name string, custom []Provider) (Provider, bool) { + for _, p := range custom { + if p.Name == name { + return p, true + } + } + for _, p := range Presets() { + if p.Name == name { + if p.APIKey == "" { + p.APIKey = os.Getenv(p.APIKeyEnv) + } + return p, true + } + } + return Provider{}, false +} + +func messageContent(m Message) any { + if len(m.Images) == 0 { + return m.Content + } + parts := []map[string]any{} + if m.Content != "" { + parts = append(parts, map[string]any{"type": "text", "text": m.Content}) + } + for _, img := range m.Images { + url := img.URL + if url == "" && img.Data != "" { + mime := img.MIME + if mime == "" { + mime = "image/png" + } + url = "data:" + mime + ";base64," + img.Data + } + parts = append(parts, map[string]any{ + "type": "image_url", + "image_url": map[string]any{"url": url}, + }) + } + return parts +} + +// BuildBody returns the JSON body for one native protocol. +func BuildBody(protocol Protocol, model string, messages []Message, tools []ToolDecl, opts StreamOptions) ([]byte, error) { + switch protocol { + case ProtocolOpenAIChat, ProtocolAzure: + body := map[string]any{ + "model": model, + "messages": toProviderMessages(messages), + "stream": true, + } + if opts.MaxTokens > 0 { + body["max_tokens"] = opts.MaxTokens + } + if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" { + body["reasoning_effort"] = effort + } + if len(tools) > 0 { + body["tools"] = toProviderTools(tools) + } + return json.Marshal(body) + case ProtocolOpenAIResponses: + body := map[string]any{ + "model": model, + "input": responsesInput(messages), + "stream": true, + } + if opts.MaxTokens > 0 { + body["max_output_tokens"] = opts.MaxTokens + } + if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" { + body["reasoning"] = map[string]any{"effort": effort} + } + return json.Marshal(body) + case ProtocolAnthropic: + sys, msgs := anthropicMessages(messages) + body := map[string]any{ + "model": model, + "messages": msgs, + "max_tokens": opts.MaxTokens, + "stream": true, + } + if body["max_tokens"] == 0 { + body["max_tokens"] = 4096 + } + if sys != "" { + body["system"] = sys + } + if budget := AnthropicThinkingBudget(opts.ThinkingLevel); budget > 0 { + body["thinking"] = map[string]any{"type": "enabled", "budget_tokens": budget} + } + return json.Marshal(body) + case ProtocolGoogle: + body := map[string]any{ + "contents": googleContents(messages), + } + if budget := GoogleThinkingBudget(opts.ThinkingLevel); budget > 0 { + body["generationConfig"] = map[string]any{ + "thinkingConfig": map[string]any{"thinkingBudget": budget}, + } + } + return json.Marshal(body) + default: + return nil, fmt.Errorf("llm: unknown protocol %q", protocol) + } +} + +func responsesInput(messages []Message) []map[string]any { + out := make([]map[string]any, 0, len(messages)) + for _, m := range messages { + role := string(m.Role) + if role == "" { + role = "user" + } + out = append(out, map[string]any{ + "role": role, + "content": messageContent(m), + }) + } + return out +} + +func anthropicMessages(messages []Message) (string, []map[string]any) { + var system strings.Builder + var out []map[string]any + for _, m := range messages { + if m.Role == RoleSystem { + if system.Len() > 0 { + system.WriteString("\n") + } + system.WriteString(m.Content) + continue + } + role := "user" + if m.Role == RoleAssistant { + role = "assistant" + } + content := []map[string]any{} + if m.Content != "" { + content = append(content, map[string]any{"type": "text", "text": m.Content}) + } + for _, img := range m.Images { + content = append(content, map[string]any{ + "type": "image", + "source": map[string]any{ + "type": "base64", + "media_type": firstNonEmpty(img.MIME, "image/png"), + "data": img.Data, + }, + }) + } + out = append(out, map[string]any{"role": role, "content": content}) + } + return system.String(), out +} + +func googleContents(messages []Message) []map[string]any { + var out []map[string]any + for _, m := range messages { + if m.Role == RoleSystem { + continue + } + role := "user" + if m.Role == RoleAssistant { + role = "model" + } + parts := []map[string]any{} + if m.Content != "" { + parts = append(parts, map[string]any{"text": m.Content}) + } + for _, img := range m.Images { + parts = append(parts, map[string]any{ + "inlineData": map[string]any{ + "mimeType": firstNonEmpty(img.MIME, "image/png"), + "data": img.Data, + }, + }) + } + out = append(out, map[string]any{"role": role, "parts": parts}) + } + return out +} + +func firstNonEmpty(a, b string) string { + if a != "" { + return a + } + return b +} + +// ContextOverflow reports provider errors that ask the caller to compact. +func ContextOverflow(message string) bool { + m := strings.ToLower(message) + return strings.Contains(m, "context") && (strings.Contains(m, "overflow") || strings.Contains(m, "too long") || strings.Contains(m, "length")) +} + +// WithRetry re-invokes base on a retryable terminal error, up to attempts +// extra tries. onEvent receives auto_retry_start and auto_retry_end. +// Context-overflow errors are returned immediately so the caller can compact. +func WithRetry(base StreamFunc, attempts int, onEvent func(kind, detail string)) StreamFunc { + if attempts < 1 { + attempts = 1 + } + return func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event { + out := make(chan Event) + go func() { + defer close(out) + var last []Event + for try := 0; try <= attempts; try++ { + if try > 0 { + detail := fmt.Sprintf("attempt %d", try+1) + if onEvent != nil { + onEvent("auto_retry_start", detail) + } + delay := time.Duration(200*(1< 0 && onEvent != nil { + onEvent("auto_retry_end", fmt.Sprintf("attempt %d", try+1)) + } + for _, ev := range last { + select { + case <-ctx.Done(): + return + case out <- ev: + } + } + return + } + } + }() + return out + } +} diff --git a/runtime/llm/providers_test.go b/runtime/llm/providers_test.go new file mode 100644 index 0000000..5a2de99 --- /dev/null +++ b/runtime/llm/providers_test.go @@ -0,0 +1,115 @@ +package llm + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestParseModelRefAndThinking(t *testing.T) { + ref, err := ParseModelRef("anthropic/claude-sonnet:high") + if err != nil { + t.Fatal(err) + } + if ref.Provider != "anthropic" || ref.Model != "claude-sonnet" || ref.Thinking != "high" { + t.Fatalf("%+v", ref) + } + if _, err := ParseModelRef("openai/gpt:nope"); err == nil { + t.Fatal("unknown thinking level must fail") + } + if AnthropicThinkingBudget("high") != 8192 || OpenAIReasoningEffort("off") != "" { + t.Fatal("thinking map") + } + p, ok := Lookup("groq", nil) + if !ok || !strings.Contains(p.BaseURL, "groq") || p.Protocol != ProtocolOpenAIChat { + t.Fatalf("preset: %+v", p) + } +} + +func TestModelsJSONAndBodies(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "models.json") + raw := []byte(`{"providers":[{"name":"local","protocol":"openai-chat","baseURL":"http://127.0.0.1:9","apiKeyEnv":"UNUSED_KEY"}]}`) + if err := os.WriteFile(path, raw, 0o644); err != nil { + t.Fatal(err) + } + got, err := LoadModelsJSON(path) + if err != nil || len(got) != 1 || got[0].Name != "local" { + t.Fatalf("%+v %v", got, err) + } + msgs := []Message{{ + Role: RoleUser, Content: "see", + Images: []Image{{MIME: "image/png", Data: "aaaa"}}, + }} + body, err := BuildBody(ProtocolAnthropic, "claude", append([]Message{{Role: RoleSystem, Content: "sys"}}, msgs...), nil, StreamOptions{ThinkingLevel: "low"}) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(body), "budget_tokens") || !strings.Contains(string(body), "base64") { + t.Fatalf("anthropic body: %s", body) + } + gbody, err := BuildBody(ProtocolGoogle, "gemini", msgs, nil, StreamOptions{ThinkingLevel: "max"}) + if err != nil || !strings.Contains(string(gbody), "inlineData") || !strings.Contains(string(gbody), "thinkingBudget") { + t.Fatalf("google: %s %v", gbody, err) + } + rbody, err := BuildBody(ProtocolOpenAIResponses, "gpt", msgs, nil, StreamOptions{ThinkingLevel: "medium", MaxTokens: 20}) + if err != nil || !strings.Contains(string(rbody), "max_output_tokens") || !strings.Contains(string(rbody), "image_url") { + t.Fatalf("responses: %s %v", rbody, err) + } + var decoded map[string]any + if err := json.Unmarshal(rbody, &decoded); err != nil { + t.Fatal(err) + } +} + +func TestRetryThenSuccess(t *testing.T) { + n := 0 + base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event { + out := make(chan Event, 2) + n++ + if n == 1 { + out <- Event{Kind: KindError, Retryable: true, ErrorMessage: "429"} + } else { + out <- Event{Kind: KindTextDelta, Text: "ok"} + out <- Event{Kind: KindMessageEnd, Message: Message{Role: RoleAssistant, Content: "ok"}} + } + close(out) + return out + }) + var events []string + wrapped := WithRetry(base, 2, func(kind, detail string) { events = append(events, kind) }) + var texts []string + for ev := range wrapped(context.Background(), "m", nil, nil, StreamOptions{}) { + if ev.Kind == KindTextDelta { + texts = append(texts, ev.Text) + } + } + if len(texts) != 1 || texts[0] != "ok" { + t.Fatalf("texts %v calls %d", texts, n) + } + if len(events) == 0 || events[0] != "auto_retry_start" { + t.Fatalf("events %v", events) + } +} + +func TestContextOverflowIsNotRetried(t *testing.T) { + if !ContextOverflow("maximum context length exceeded") { + t.Fatal("expected overflow") + } + n := 0 + base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event { + out := make(chan Event, 1) + n++ + out <- Event{Kind: KindError, Retryable: true, ErrorMessage: "context length overflow"} + close(out) + return out + }) + for range WithRetry(base, 3, nil)(context.Background(), "m", nil, nil, StreamOptions{}) { + } + if n != 1 { + t.Fatalf("overflow must not retry, calls=%d", n) + } +} diff --git a/runtime/llm/sse.go b/runtime/llm/sse.go index 6389917..d7588ae 100644 --- a/runtime/llm/sse.go +++ b/runtime/llm/sse.go @@ -26,12 +26,19 @@ type sseChoice struct { type sseChunk struct { Choices []sseChoice `json:"choices"` - Usage *Usage `json:"usage"` + Usage *usageWire `json:"usage"` Error *struct { Message string `json:"message"` } `json:"error"` } +type usageWire struct { + Usage + Details struct { + Cached int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` +} + // ParseSSE reads an OpenAI-compatible server-sent event body and emits // deltas plus one terminal message_end or error event. func ParseSSE(ctx context.Context, r io.Reader) <-chan Event { @@ -77,7 +84,10 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { return } if chunk.Usage != nil { - usage = *chunk.Usage + usage = chunk.Usage.Usage + if chunk.Usage.Details.Cached > 0 { + usage.CacheReadTokens = chunk.Usage.Details.Cached + } } if len(chunk.Choices) == 0 { continue diff --git a/runtime/llm/types.go b/runtime/llm/types.go index 7cae603..623e0fc 100644 --- a/runtime/llm/types.go +++ b/runtime/llm/types.go @@ -48,6 +48,7 @@ type Message struct { Role Role `json:"role"` Content string `json:"content,omitempty"` Name string `json:"name,omitempty"` + Images []Image `json:"images,omitempty"` ToolCallID string `json:"tool_call_id,omitempty"` ToolCalls []ToolCall `json:"tool_calls,omitempty"` StopReason StopReason `json:"stop_reason,omitempty"` @@ -61,11 +62,20 @@ type ToolDecl struct { Parameters json.RawMessage `json:"parameters"` } +// Image is one picture attached to a user message. +type Image struct { + MIME string `json:"mime,omitempty"` + URL string `json:"url,omitempty"` + Data string `json:"data,omitempty"` +} + // Usage counts tokens reported by the provider. Zero means the provider omitted usage. type Usage struct { PromptTokens int `json:"prompt_tokens"` CompletionTokens int `json:"completion_tokens"` TotalTokens int `json:"total_tokens"` + CacheReadTokens int `json:"cache_read_tokens,omitempty"` + CacheWriteTokens int `json:"cache_write_tokens,omitempty"` } // Event is one streamed update. Failures are KindError plus a final assistant @@ -90,6 +100,9 @@ type StreamOptions struct { MaxTokens int APIBase string APIKey string + // ThinkingLevel is off, minimal, low, medium, high, xhigh, or max. + // Empty leaves the provider default unchanged. + ThinkingLevel string } // StreamFunc is the provider boundary. It must return a channel that is closed