fix: deliver validated AI insight responses
This commit is contained in:
196
internal/aiprovider/client.go
Normal file
196
internal/aiprovider/client.go
Normal file
@@ -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
|
||||
}
|
||||
59
internal/aiprovider/client_test.go
Normal file
59
internal/aiprovider/client_test.go
Normal file
@@ -0,0 +1,59 @@
|
||||
package aiprovider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestClientNormalizesV1AndUsesPrivateStrictResponsesRequest(t *testing.T) {
|
||||
var received map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/responses" {
|
||||
t.Fatalf("path = %q", r.URL.Path)
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&received); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w.Header().Set("X-Request-ID", "aval-request-1")
|
||||
_, _ = w.Write([]byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"{\"answer\":\"ok\"}"}]}],"usage":{"input_tokens":12,"output_tokens":4}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := New(Config{BaseURL: server.URL + "/v1/", APIKey: "secret", Model: "test-model", Timeout: time.Second}, server.Client())
|
||||
result, usage, err := client.Generate(context.Background(), Request{SchemaName: "test_v1", Schema: map[string]any{"type": "object", "additionalProperties": false}, Input: json.RawMessage(`{"safe":"aggregate"}`)})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(result) != `{"answer":"ok"}` || usage.ProviderRequestID != "aval-request-1" {
|
||||
t.Fatalf("result=%s usage=%+v", result, usage)
|
||||
}
|
||||
if received["store"] != false {
|
||||
t.Fatalf("store = %#v, want false", received["store"])
|
||||
}
|
||||
format := received["text"].(map[string]any)["format"].(map[string]any)
|
||||
if format["strict"] != true || format["name"] != "test_v1" {
|
||||
t.Fatalf("format = %#v", format)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientRetriesOnlyTransientFailure(t *testing.T) {
|
||||
calls := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
calls++
|
||||
if calls == 1 {
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"{}"}]}]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
client := New(Config{BaseURL: server.URL, APIKey: "secret", Model: "model"}, server.Client())
|
||||
_, usage, err := client.Generate(t.Context(), Request{SchemaName: "test", Schema: map[string]any{"type": "object"}, Input: json.RawMessage(`{}`)})
|
||||
if err != nil || calls != 2 || usage.ProviderRequests != 2 {
|
||||
t.Fatalf("err=%v calls=%d usage=%+v", err, calls, usage)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user