Agent-Native/runtime/llm/providers.go

package llm

import (
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"strconv"
	"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: ProtocolOpenAIChat, BaseURL: "https://generativelanguage.googleapis.com/v1beta/openai", 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
}

// ClientFor 返回会说该协议的客户端。responses、azure 以及原生 Google
// 协议当前没有接到客户端上，返回明确错误而不是发到错误的路径。
func ClientFor(p Provider) (*Client, error) {
	switch p.Protocol {
	case ProtocolOpenAIChat:
		return &Client{BaseURL: p.BaseURL, APIKey: p.APIKey}, nil
	case ProtocolAnthropic:
		return &Client{BaseURL: p.BaseURL, APIKey: p.APIKey, Protocol: ProtocolAnthropic}, nil
	default:
		return nil, fmt.Errorf("llm: protocol %q is not implemented", p.Protocol)
	}
}

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)
	// ProtocolGoogle 只供离线组包（客户端未接）；Anthropic 与线上客户端
	// 共用 anthropicBody（同一组包，消灭离线/线上两张皮）。
	case ProtocolAnthropic:
		return anthropicBody(model, messages, tools, opts)
	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 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
}

const (
	retryBaseDelay = 200 * time.Millisecond
	maxRetryAfter  = 30 * time.Second
)

// overflowNeedles are substrings that mean the prompt does not fit.
// Matching is case-insensitive. A bare "context" is not enough.
var overflowNeedles = []string{
	"status 413",
	"payload too large",
	"context_length_exceeded",
	"context length",
	"maximum context",
	"context window",
	"prompt is too long",
	"exceeds the maximum",
	"exceeds max",
	"too many tokens",
	"input is too long",
	"input length",
	"tokens exceed",
	"上下文",
	"超出最大",
	"输入过长",
	"提示过长",
	"请求过长",
	"请求体过大",
}

// ContextOverflow reports provider errors that ask the caller to compact.
func ContextOverflow(message string) bool {
	m := strings.ToLower(message)
	for _, needle := range overflowNeedles {
		if strings.Contains(m, strings.ToLower(needle)) {
			return true
		}
	}
	return false
}

// CapacityLimited reports an in-stream overload that is worth retrying
// before any delta has been forwarded.
func CapacityLimited(message string) bool {
	m := strings.ToLower(message)
	return strings.Contains(m, "overloaded") || strings.Contains(m, "at capacity") || strings.Contains(m, "too many requests") || strings.Contains(m, "rate limit")
}

// StatusRetryable is true for timeouts, rate limits, server errors, and 529.
func StatusRetryable(code int) bool {
	return code == http.StatusRequestTimeout || code == http.StatusTooManyRequests || code == 529 || code >= 500
}

// ParseRetryAfter reads a delay in seconds or an HTTP date.
// The returned duration is capped. ok is false when the header is missing or garbage.
func ParseRetryAfter(raw string, now time.Time) (time.Duration, bool) {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return 0, false
	}
	if secs, err := strconv.Atoi(raw); err == nil {
		if secs < 0 {
			return 0, false
		}
		return capRetryAfter(time.Duration(secs) * time.Second), true
	}
	when, err := http.ParseTime(raw)
	if err != nil {
		return 0, false
	}
	d := when.Sub(now)
	if d < 0 {
		d = 0
	}
	return capRetryAfter(d), true
}

func capRetryAfter(d time.Duration) time.Duration {
	if d > maxRetryAfter {
		return maxRetryAfter
	}
	return d
}

func retryBackoff(extraTry int) time.Duration {
	if extraTry < 1 {
		extraTry = 1
	}
	shift := extraTry - 1
	if shift > 4 {
		shift = 4
	}
	delay := retryBaseDelay * time.Duration(1<<uint(shift))
	jitter := time.Duration(rand.Int63n(int64(delay/2) + 1))
	return delay + jitter
}

func forwardedDelta(ev Event) bool {
	switch ev.Kind {
	case KindTextDelta, KindReasoningDelta, KindToolCallDelta:
		return true
	default:
		return false
	}
}

// WithRetry re-invokes base on a retryable terminal error, up to attempts
// extra tries. Events are forwarded as they arrive. A retry happens only when
// no delta has been forwarded yet. onEvent receives auto_retry_start and
// auto_retry_end. Context-overflow errors are returned immediately.
// attempts <= 0 disables extra tries.
func WithRetry(base StreamFunc, attempts int, onEvent func(kind, detail string)) StreamFunc {
	if attempts < 0 {
		attempts = 0
	}
	return func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
		out := make(chan Event)
		go func() {
			defer close(out)
			headerWait := time.Duration(0)
			useHeader := false
			for try := 0; try <= attempts; try++ {
				if try > 0 {
					delay := retryBackoff(try)
					if useHeader {
						delay = headerWait
					}
					detail := fmt.Sprintf("attempt %d", try+1)
					if onEvent != nil {
						onEvent("auto_retry_start", detail)
					}
					select {
					case <-ctx.Done():
						emitAbort(out, "", ctx.Err())
						return
					case <-time.After(delay):
					}
				}
				useHeader = false
				forwarded := false
				var terminal *Event
				for ev := range base(ctx, model, messages, tools, opts) {
					if terminal != nil {
						continue
					}
					if ev.Kind == KindError {
						cp := ev
						terminal = &cp
						continue
					}
					if forwardedDelta(ev) {
						forwarded = true
					}
					select {
					case <-ctx.Done():
						return
					case out <- ev:
					}
					if ev.Kind == KindMessageEnd {
						if try > 0 && onEvent != nil {
							onEvent("auto_retry_end", fmt.Sprintf("attempt %d", try+1))
						}
						return
					}
				}
				if terminal == nil {
					return
				}
				overflow := ContextOverflow(terminal.ErrorMessage)
				if forwarded || overflow || !terminal.Retryable || try == attempts {
					if try > 0 && onEvent != nil {
						onEvent("auto_retry_end", fmt.Sprintf("attempt %d", try+1))
					}
					select {
					case <-ctx.Done():
						return
					case out <- *terminal:
					}
					return
				}
				if terminal.HasRetryAfter {
					useHeader = true
					headerWait = terminal.RetryAfter
				}
			}
		}()
		return out
	}
}