fix: deliver validated AI insight responses

2026-08-22 15:47:38 +08:00
parent ca55c2fb3d
commit 3fba4261ee
68 changed files with 6870 additions and 493 deletions
--- a/internal/aiprovider/client.go
+++ b/internal/aiprovider/client.go
@@ -0,0 +1,196 @@
+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
+}