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

package aiprovider

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

const maxResponseBytes = 1 << 20

type Config struct {
	BaseURL string
	APIKey  string
	Model   string
	Timeout time.Duration
}

type Request struct {
	SchemaName      string
	Schema          map[string]any
	Instructions    string
	Input           json.RawMessage
	MaxOutputTokens int
}

type Usage struct {
	InputTokens       int64
	OutputTokens      int64
	ProviderRequests  int64
	LatencyMS         int64
	ProviderRequestID string
	EstimatedCost     string
}

type Client struct {
	config Config
	client *http.Client
}

func New(config Config, client *http.Client) *Client {
	if config.BaseURL == "" {
		config.BaseURL = "https://api.avalai.ir"
	}
	config.BaseURL = strings.TrimSuffix(strings.TrimRight(config.BaseURL, "/"), "/v1")
	if config.Timeout <= 0 || config.Timeout > 30*time.Second {
		config.Timeout = 30 * time.Second
	}
	if client == nil {
		client = &http.Client{}
	}
	return &Client{config: config, client: client}
}

func (c *Client) Model() string { return c.config.Model }

func (c *Client) Generate(ctx context.Context, request Request) (json.RawMessage, Usage, error) {
	if c.config.APIKey == "" || c.config.Model == "" {
		return nil, Usage{}, errors.New("avalai: disabled")
	}
	if request.MaxOutputTokens <= 0 {
		request.MaxOutputTokens = 1600
	}
	bodyValue := map[string]any{
		"model": c.config.Model, "store": false, "parallel_tool_calls": false,
		"max_output_tokens": request.MaxOutputTokens,
		"instructions":      request.Instructions,
		"input":             string(request.Input),
		"text": map[string]any{"format": map[string]any{
			"type": "json_schema", "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, bodyValue)
		total = mergeUsage(total, usage)
		if err == nil {
			return output, total, nil
		}
		last = err
		if !transient || ctx.Err() != nil {
			break
		}
	}
	return nil, total, last
}

type envelope struct {
	Status string `json:"status"`
	Output []struct {
		Type    string `json:"type"`
		Content []struct {
			Type    string `json:"type"`
			Text    string `json:"text"`
			Refusal string `json:"refusal"`
		} `json:"content"`
	} `json:"output"`
	Usage struct {
		InputTokens  int64 `json:"input_tokens"`
		OutputTokens int64 `json:"output_tokens"`
	} `json:"usage"`
	EstimatedCost json.RawMessage `json:"estimated_cost"`
}

func (c *Client) 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 AvalAI request: %w", err)
	}
	ctx, cancel := context.WithTimeout(parent, c.config.Timeout)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.config.BaseURL+"/v1/responses", bytes.NewReader(body))
	if err != nil {
		return nil, Usage{}, false, fmt.Errorf("creating AvalAI request: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+c.config.APIKey)
	req.Header.Set("Content-Type", "application/json")
	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 AvalAI: %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 AvalAI response: %w", readErr)
	}
	if len(responseBody) > maxResponseBytes {
		return nil, usage, false, errors.New("avalai: response too large")
	}
	var result envelope
	decodeErr := json.Unmarshal(responseBody, &result)
	usage.InputTokens, usage.OutputTokens = result.Usage.InputTokens, result.Usage.OutputTokens
	usage.EstimatedCost = safeCost(result.EstimatedCost)
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, usage, resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500, fmt.Errorf("avalai: HTTP %d", resp.StatusCode)
	}
	if decodeErr != nil {
		return nil, usage, false, fmt.Errorf("decoding AvalAI response: %w", decodeErr)
	}
	if result.Status != "completed" {
		return nil, usage, false, fmt.Errorf("avalai: incomplete response status %q", result.Status)
	}
	var text strings.Builder
	for _, item := range result.Output {
		if item.Type != "message" {
			continue
		}
		for _, content := range item.Content {
			if content.Refusal != "" {
				return nil, usage, false, errors.New("avalai: refused")
			}
			if content.Type == "output_text" {
				text.WriteString(content.Text)
			}
		}
	}
	if !json.Valid([]byte(text.String())) {
		return nil, usage, false, errors.New("avalai: invalid structured output")
	}
	return json.RawMessage(text.String()), usage, false, nil
}

func mergeUsage(left, right Usage) Usage {
	left.InputTokens += right.InputTokens
	left.OutputTokens += right.OutputTokens
	left.ProviderRequests += right.ProviderRequests
	left.LatencyMS += right.LatencyMS
	if right.ProviderRequestID != "" {
		left.ProviderRequestID = right.ProviderRequestID
	}
	if right.EstimatedCost != "" {
		left.EstimatedCost = right.EstimatedCost
	}
	return left
}

func safeCost(raw json.RawMessage) string {
	value := strings.Trim(strings.TrimSpace(string(raw)), `"`)
	if len(value) > 64 {
		return ""
	}
	for _, r := range value {
		if (r < '0' || r > '9') && r != '.' && r != '-' && r != 'e' && r != 'E' {
			return ""
		}
	}
	return value
}