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 <cursoragent@cursor.com>
This commit is contained in:
@@ -53,6 +53,9 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool
|
|||||||
if opts.MaxTokens > 0 {
|
if opts.MaxTokens > 0 {
|
||||||
body["max_tokens"] = opts.MaxTokens
|
body["max_tokens"] = opts.MaxTokens
|
||||||
}
|
}
|
||||||
|
if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" {
|
||||||
|
body["reasoning_effort"] = effort
|
||||||
|
}
|
||||||
if len(tools) > 0 {
|
if len(tools) > 0 {
|
||||||
body["tools"] = toProviderTools(tools)
|
body["tools"] = toProviderTools(tools)
|
||||||
body["tool_choice"] = "auto"
|
body["tool_choice"] = "auto"
|
||||||
@@ -125,7 +128,7 @@ func toProviderMessages(messages []Message) []map[string]any {
|
|||||||
default:
|
default:
|
||||||
out = append(out, map[string]any{
|
out = append(out, map[string]any{
|
||||||
"role": string(m.Role),
|
"role": string(m.Role),
|
||||||
"content": m.Content,
|
"content": messageContent(m),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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<<uint(try-1))) * time.Millisecond
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
emitAbort(out, "", ctx.Err())
|
||||||
|
return
|
||||||
|
case <-time.After(delay):
|
||||||
|
}
|
||||||
|
}
|
||||||
|
last = nil
|
||||||
|
for ev := range base(ctx, model, messages, tools, opts) {
|
||||||
|
last = append(last, ev)
|
||||||
|
}
|
||||||
|
overflow, retryable := false, false
|
||||||
|
for _, ev := range last {
|
||||||
|
if ev.Kind == KindError && ContextOverflow(ev.ErrorMessage) {
|
||||||
|
overflow = true
|
||||||
|
}
|
||||||
|
if ev.Kind == KindError && ev.Retryable {
|
||||||
|
retryable = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if overflow || !retryable || try == attempts {
|
||||||
|
if try > 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
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
+12
-2
@@ -26,12 +26,19 @@ type sseChoice struct {
|
|||||||
|
|
||||||
type sseChunk struct {
|
type sseChunk struct {
|
||||||
Choices []sseChoice `json:"choices"`
|
Choices []sseChoice `json:"choices"`
|
||||||
Usage *Usage `json:"usage"`
|
Usage *usageWire `json:"usage"`
|
||||||
Error *struct {
|
Error *struct {
|
||||||
Message string `json:"message"`
|
Message string `json:"message"`
|
||||||
} `json:"error"`
|
} `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
|
// ParseSSE reads an OpenAI-compatible server-sent event body and emits
|
||||||
// deltas plus one terminal message_end or error event.
|
// deltas plus one terminal message_end or error event.
|
||||||
func ParseSSE(ctx context.Context, r io.Reader) <-chan 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
|
return
|
||||||
}
|
}
|
||||||
if chunk.Usage != nil {
|
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 {
|
if len(chunk.Choices) == 0 {
|
||||||
continue
|
continue
|
||||||
|
|||||||
@@ -48,6 +48,7 @@ type Message struct {
|
|||||||
Role Role `json:"role"`
|
Role Role `json:"role"`
|
||||||
Content string `json:"content,omitempty"`
|
Content string `json:"content,omitempty"`
|
||||||
Name string `json:"name,omitempty"`
|
Name string `json:"name,omitempty"`
|
||||||
|
Images []Image `json:"images,omitempty"`
|
||||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||||
StopReason StopReason `json:"stop_reason,omitempty"`
|
StopReason StopReason `json:"stop_reason,omitempty"`
|
||||||
@@ -61,11 +62,20 @@ type ToolDecl struct {
|
|||||||
Parameters json.RawMessage `json:"parameters"`
|
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.
|
// Usage counts tokens reported by the provider. Zero means the provider omitted usage.
|
||||||
type Usage struct {
|
type Usage struct {
|
||||||
PromptTokens int `json:"prompt_tokens"`
|
PromptTokens int `json:"prompt_tokens"`
|
||||||
CompletionTokens int `json:"completion_tokens"`
|
CompletionTokens int `json:"completion_tokens"`
|
||||||
TotalTokens int `json:"total_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
|
// Event is one streamed update. Failures are KindError plus a final assistant
|
||||||
@@ -90,6 +100,9 @@ type StreamOptions struct {
|
|||||||
MaxTokens int
|
MaxTokens int
|
||||||
APIBase string
|
APIBase string
|
||||||
APIKey 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
|
// StreamFunc is the provider boundary. It must return a channel that is closed
|
||||||
|
|||||||
Reference in New Issue
Block a user