zarinpal-analytics-backend/internal/aiprovider/agentrouter.go

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
}