feat: add configurable AgentRouter AI provider

2026-08-22 21:47:41 +08:00
parent 4f91726edb
commit cf2f114dc8
20 changed files with 1204 additions and 79 deletions
--- a/internal/aiprovider/agentrouter.go
+++ b/internal/aiprovider/agentrouter.go
@@ -0,0 +1,145 @@
+package aiprovider
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+)
+
+const maxAgentRouterRequestTimeout = 90 * time.Second
+
+const agentRouterUserAgent = "claude-cli/1.0.0 (external, cli)"
+
+type AgentRouterClient struct {
+	config Config
+	client *http.Client
+}
+
+var _ Generator = (*AgentRouterClient)(nil)
+
+func NewAgentRouter(config Config, client *http.Client) *AgentRouterClient {
+	if config.BaseURL == "" {
+		config.BaseURL = "https://agentrouter.org/v1"
+	}
+	config.BaseURL = strings.TrimRight(config.BaseURL, "/")
+	if !strings.HasSuffix(config.BaseURL, "/v1") {
+		config.BaseURL += "/v1"
+	}
+	if config.Timeout <= 0 || config.Timeout > maxAgentRouterRequestTimeout {
+		config.Timeout = maxAgentRouterRequestTimeout
+	}
+	if client == nil {
+		client = &http.Client{}
+	}
+	return &AgentRouterClient{config: config, client: client}
+}
+
+func (c *AgentRouterClient) Model() string { return c.config.Model }
+
+func (c *AgentRouterClient) Generate(ctx context.Context, request Request) (json.RawMessage, Usage, error) {
+	if c.config.APIKey == "" || c.config.Model == "" {
+		return nil, Usage{}, errors.New("agentrouter: disabled")
+	}
+	if request.MaxOutputTokens <= 0 {
+		request.MaxOutputTokens = 1600
+	}
+	body := map[string]any{
+		"model":               c.config.Model,
+		"messages":            []map[string]any{{"role": "system", "content": request.Instructions}, {"role": "user", "content": string(request.Input)}},
+		"max_tokens":          request.MaxOutputTokens,
+		"reasoning_effort":    "medium",
+		"parallel_tool_calls": false,
+		"response_format": map[string]any{
+			"type":        "json_schema",
+			"json_schema": map[string]any{"name": request.SchemaName, "strict": true, "schema": request.Schema},
+		},
+	}
+	var total Usage
+	var last error
+	for attempt := 0; attempt < 2; attempt++ {
+		output, usage, transient, err := c.do(ctx, body)
+		total = mergeUsage(total, usage)
+		if err == nil {
+			return output, total, nil
+		}
+		last = err
+		if !transient || ctx.Err() != nil {
+			break
+		}
+	}
+	return nil, total, last
+}
+
+type agentRouterEnvelope struct {
+	ID      string `json:"id"`
+	Choices []struct {
+		FinishReason string `json:"finish_reason"`
+		Message      struct {
+			Content string `json:"content"`
+			Refusal string `json:"refusal"`
+		} `json:"message"`
+	} `json:"choices"`
+	Usage struct {
+		PromptTokens     int64 `json:"prompt_tokens"`
+		CompletionTokens int64 `json:"completion_tokens"`
+	} `json:"usage"`
+}
+
+func (c *AgentRouterClient) do(parent context.Context, value map[string]any) (json.RawMessage, Usage, bool, error) {
+	body, err := json.Marshal(value)
+	if err != nil {
+		return nil, Usage{}, false, fmt.Errorf("encoding agentrouter request: %w", err)
+	}
+	ctx, cancel := context.WithTimeout(parent, c.config.Timeout)
+	defer cancel()
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.config.BaseURL+"/chat/completions", bytes.NewReader(body))
+	if err != nil {
+		return nil, Usage{}, false, fmt.Errorf("creating agentrouter request: %w", err)
+	}
+	req.Header.Set("Authorization", "Bearer "+c.config.APIKey)
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("User-Agent", agentRouterUserAgent)
+	started := time.Now()
+	resp, err := c.client.Do(req)
+	if err != nil {
+		return nil, Usage{ProviderRequests: 1, LatencyMS: time.Since(started).Milliseconds()}, true, fmt.Errorf("calling agentrouter: %w", err)
+	}
+	defer resp.Body.Close()
+	responseBody, readErr := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes+1))
+	usage := Usage{ProviderRequests: 1, LatencyMS: time.Since(started).Milliseconds(), ProviderRequestID: resp.Header.Get("X-Request-ID")}
+	if readErr != nil {
+		return nil, usage, true, fmt.Errorf("reading agentrouter response: %w", readErr)
+	}
+	if len(responseBody) > maxResponseBytes {
+		return nil, usage, false, errors.New("agentrouter: response too large")
+	}
+	var result agentRouterEnvelope
+	decodeErr := json.Unmarshal(responseBody, &result)
+	usage.InputTokens, usage.OutputTokens = result.Usage.PromptTokens, result.Usage.CompletionTokens
+	if usage.ProviderRequestID == "" {
+		usage.ProviderRequestID = result.ID
+	}
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return nil, usage, resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500, fmt.Errorf("agentrouter: http %d", resp.StatusCode)
+	}
+	if decodeErr != nil {
+		return nil, usage, false, fmt.Errorf("decoding agentrouter response: %w", decodeErr)
+	}
+	if len(result.Choices) != 1 || result.Choices[0].FinishReason != "stop" {
+		return nil, usage, false, errors.New("agentrouter: incomplete chat completion")
+	}
+	message := result.Choices[0].Message
+	if message.Refusal != "" {
+		return nil, usage, false, errors.New("agentrouter: refused")
+	}
+	if !json.Valid([]byte(message.Content)) {
+		return nil, usage, false, errors.New("agentrouter: invalid structured output")
+	}
+	return json.RawMessage(message.Content), usage, false, nil
+}