mirror of
https://github.com/lukaszraczylo/llm-testbench.git
synced 2026-10-16 01:07:47 +00:00
feat(llm): make runs resilient and reproducible on slow backends
Retry with exponential backoff and jitter, and honour Retry-After. Retry empty-choice responses and HTTP 408. Add max_retries, retry_timeouts, model_timeouts and model_concurrency settings for slow local models. Add opt-in SSE streaming with an idle timeout. A hung server fails early, and a slow but active generation continues. Send an optional seed and record the backend fingerprint of each answer. Stamp saved artifacts with run metadata and write them atomically. Warn in compare when two runs differ. Save partial results on Ctrl-C. Turn an evaluator panic into an error row. Reject duplicate models and non-http endpoints in the config.
This commit is contained in:
@@ -37,6 +37,7 @@ type ToolCall struct {
|
||||
// non-empty the client advertises them with tool_choice "auto", so the
|
||||
// model decides whether and which to call.
|
||||
type Request struct {
|
||||
Seed *int
|
||||
Model string
|
||||
Messages []Message
|
||||
Tools []Tool
|
||||
@@ -48,6 +49,8 @@ type Request struct {
|
||||
type Response struct {
|
||||
Text string
|
||||
FinishReason string
|
||||
Fingerprint string
|
||||
ServedModel string
|
||||
ToolCalls []ToolCall
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
|
||||
+187
-36
@@ -7,22 +7,46 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/rand/v2"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// maxRetries is the number of retry attempts after the initial request on
|
||||
// transport errors or 5xx responses, per PLAN.md ("retry w/ backoff x2").
|
||||
// maxRetries is the default number of retry attempts after the initial
|
||||
// request on transient failures, per PLAN.md ("retry w/ backoff x2").
|
||||
const maxRetries = 2
|
||||
|
||||
// retryBaseDelay is the base backoff delay; attempt N waits N*retryBaseDelay.
|
||||
// retryBaseDelay is the base backoff delay; attempt N waits
|
||||
// retryBaseDelay*2^(N-1), capped at maxRetryDelay, plus jitter.
|
||||
const retryBaseDelay = 250 * time.Millisecond
|
||||
|
||||
// maxRetryDelay caps both the exponential backoff and a server-sent
|
||||
// Retry-After so one bad gateway cannot stall a worker for minutes.
|
||||
const maxRetryDelay = 30 * time.Second
|
||||
|
||||
// maxResponseBytes bounds how much of a response body is read; a runaway
|
||||
// or hostile upstream must not exhaust memory.
|
||||
const maxResponseBytes = 32 << 20
|
||||
|
||||
// jitterFraction is the share of the backoff delay added as random jitter,
|
||||
// so concurrent workers retrying one saturated gateway do not stampede it
|
||||
// in lockstep.
|
||||
const jitterFraction = 0.25
|
||||
|
||||
// OpenAIClient talks to an OpenAI-compatible /v1/chat/completions endpoint.
|
||||
type OpenAIClient struct {
|
||||
httpClient *http.Client
|
||||
endpoint string
|
||||
apiKey string
|
||||
httpClient *http.Client
|
||||
modelTimeouts map[string]time.Duration
|
||||
endpoint string
|
||||
apiKey string
|
||||
retries int
|
||||
baseDelay time.Duration
|
||||
timeout time.Duration
|
||||
noRetryTimeouts bool
|
||||
stream bool
|
||||
idleTimeout time.Duration
|
||||
}
|
||||
|
||||
// NewOpenAIClient builds a client against endpoint (e.g.
|
||||
@@ -30,21 +54,83 @@ type OpenAIClient struct {
|
||||
// per-request timeout.
|
||||
func NewOpenAIClient(endpoint, apiKey string, timeout time.Duration) *OpenAIClient {
|
||||
return &OpenAIClient{
|
||||
endpoint: endpoint,
|
||||
apiKey: apiKey,
|
||||
httpClient: &http.Client{
|
||||
Timeout: timeout,
|
||||
},
|
||||
endpoint: strings.TrimRight(endpoint, "/"),
|
||||
apiKey: apiKey,
|
||||
retries: maxRetries,
|
||||
baseDelay: retryBaseDelay,
|
||||
timeout: timeout,
|
||||
// The per-attempt deadline is applied through the context so it can
|
||||
// differ per model; the client itself has no overall timeout.
|
||||
httpClient: &http.Client{},
|
||||
}
|
||||
}
|
||||
|
||||
// WithModelTimeouts sets per-model request timeouts that override the
|
||||
// default; a slow local model can be given far longer than a hosted one.
|
||||
func (c *OpenAIClient) WithModelTimeouts(m map[string]time.Duration) *OpenAIClient {
|
||||
c.modelTimeouts = m
|
||||
return c
|
||||
}
|
||||
|
||||
// WithRetryTimeouts controls whether an attempt that hit its timeout is
|
||||
// retried. Retrying re-runs a full slow generation, so for slow models
|
||||
// disabling it avoids multiplying the wait by the attempt count.
|
||||
func (c *OpenAIClient) WithRetryTimeouts(retry bool) *OpenAIClient {
|
||||
c.noRetryTimeouts = !retry
|
||||
return c
|
||||
}
|
||||
|
||||
// WithStreaming makes the client request SSE streams. The per-request
|
||||
// timeout stays a hard ceiling; idle (default 5m when zero) is how long the
|
||||
// stream may go silent before the attempt is abandoned, so a hung server is
|
||||
// detected early while a slow-but-progressing generation is left to finish.
|
||||
func (c *OpenAIClient) WithStreaming(enabled bool, idle time.Duration) *OpenAIClient {
|
||||
c.stream = enabled
|
||||
c.idleTimeout = idle
|
||||
if enabled && idle <= 0 {
|
||||
c.idleTimeout = defaultIdleTimeout
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *OpenAIClient) timeoutFor(model string) time.Duration {
|
||||
if d, ok := c.modelTimeouts[model]; ok && d > 0 {
|
||||
return d
|
||||
}
|
||||
return c.timeout
|
||||
}
|
||||
|
||||
// WithRetries sets the number of retry attempts after the initial request
|
||||
// (0 disables retrying) and returns the client.
|
||||
func (c *OpenAIClient) WithRetries(n int) *OpenAIClient {
|
||||
c.retries = max(n, 0)
|
||||
return c
|
||||
}
|
||||
|
||||
// backoff returns the wait before retry number attempt (1-based): the
|
||||
// server's Retry-After when it sent one, else exponential with jitter.
|
||||
func (c *OpenAIClient) backoff(attempt int, retryAfter time.Duration) time.Duration {
|
||||
if retryAfter > 0 {
|
||||
return min(retryAfter, maxRetryDelay)
|
||||
}
|
||||
d := min(c.baseDelay<<(attempt-1), maxRetryDelay)
|
||||
return d + time.Duration(rand.Float64()*jitterFraction*float64(d)) //nolint:gosec // jitter needs no cryptographic randomness
|
||||
}
|
||||
|
||||
type chatCompletionRequest struct {
|
||||
Model string `json:"model"`
|
||||
ToolChoice string `json:"tool_choice,omitempty"`
|
||||
Messages []chatMsg `json:"messages"`
|
||||
Tools []apiTool `json:"tools,omitempty"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Temperature float64 `json:"temperature"`
|
||||
Seed *int `json:"seed,omitempty"`
|
||||
StreamOptions *streamOptions `json:"stream_options,omitempty"`
|
||||
Model string `json:"model"`
|
||||
ToolChoice string `json:"tool_choice,omitempty"`
|
||||
Messages []chatMsg `json:"messages"`
|
||||
Tools []apiTool `json:"tools,omitempty"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Temperature float64 `json:"temperature"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
}
|
||||
|
||||
type streamOptions struct {
|
||||
IncludeUsage bool `json:"include_usage"`
|
||||
}
|
||||
|
||||
type chatMsg struct {
|
||||
@@ -102,7 +188,9 @@ type chatCompletionResponse struct {
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
Choices []struct {
|
||||
Model string `json:"model"`
|
||||
SystemFingerprint string `json:"system_fingerprint"`
|
||||
Choices []struct {
|
||||
FinishReason string `json:"finish_reason"`
|
||||
Message responseMsg `json:"message"`
|
||||
} `json:"choices"`
|
||||
@@ -145,6 +233,7 @@ func (c *OpenAIClient) Complete(ctx context.Context, req Request) (Response, err
|
||||
Model: req.Model,
|
||||
MaxTokens: req.MaxTokens,
|
||||
Temperature: req.Temperature,
|
||||
Seed: req.Seed,
|
||||
}
|
||||
for _, m := range req.Messages {
|
||||
body.Messages = append(body.Messages, chatMsg(m))
|
||||
@@ -161,15 +250,21 @@ func (c *OpenAIClient) Complete(ctx context.Context, req Request) (Response, err
|
||||
}
|
||||
}
|
||||
|
||||
if c.stream {
|
||||
body.Stream = true
|
||||
body.StreamOptions = &streamOptions{IncludeUsage: true}
|
||||
}
|
||||
|
||||
payload, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return Response{}, fmt.Errorf("llm: marshal request: %w", err)
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for attempt := 0; attempt <= maxRetries; attempt++ {
|
||||
var retryAfter time.Duration
|
||||
for attempt := 0; attempt <= c.retries; attempt++ {
|
||||
if attempt > 0 {
|
||||
delay := time.Duration(attempt) * retryBaseDelay
|
||||
delay := c.backoff(attempt, retryAfter)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return Response{}, ctx.Err()
|
||||
@@ -181,16 +276,32 @@ func (c *OpenAIClient) Complete(ctx context.Context, req Request) (Response, err
|
||||
// reflects the successful attempt's own generation time, not the
|
||||
// cumulative wait across earlier failed/retried attempts (S8).
|
||||
attemptStart := time.Now()
|
||||
resp, err := c.doRequest(ctx, payload)
|
||||
attemptCtx, cancel := context.WithTimeout(ctx, c.timeoutFor(req.Model))
|
||||
resp, err := c.doRequest(attemptCtx, payload)
|
||||
timedOut := (attemptCtx.Err() == context.DeadlineExceeded && ctx.Err() == nil) || errors.Is(err, errIdleTimeout)
|
||||
cancel()
|
||||
retryAfter = 0
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
if isRetryable(err) {
|
||||
if timedOut && c.noRetryTimeouts {
|
||||
return Response{}, fmt.Errorf("llm: timed out after %s (not retried): %w", c.timeoutFor(req.Model), err)
|
||||
}
|
||||
var re *retryableError
|
||||
if errors.As(err, &re) {
|
||||
retryAfter = re.retryAfter
|
||||
continue
|
||||
}
|
||||
return Response{}, err
|
||||
}
|
||||
|
||||
latency := time.Since(attemptStart)
|
||||
if resp.Error == nil && len(resp.Choices) == 0 {
|
||||
// A 2xx with neither choices nor an error object is a gateway
|
||||
// glitch, not an answer; scoring it would record a spurious
|
||||
// wrong response.
|
||||
lastErr = errors.New("llm: response contained no choices")
|
||||
continue
|
||||
}
|
||||
var text, finishReason string
|
||||
var toolCalls []ToolCall
|
||||
if len(resp.Choices) > 0 {
|
||||
@@ -213,14 +324,17 @@ func (c *OpenAIClient) Complete(ctx context.Context, req Request) (Response, err
|
||||
PromptTokens: resp.Usage.PromptTokens,
|
||||
CompletionTokens: resp.Usage.CompletionTokens,
|
||||
Latency: latency,
|
||||
Fingerprint: resp.SystemFingerprint,
|
||||
ServedModel: resp.Model,
|
||||
}, nil
|
||||
}
|
||||
return Response{}, fmt.Errorf("llm: request failed after %d attempts: %w", maxRetries+1, lastErr)
|
||||
return Response{}, fmt.Errorf("llm: request failed after %d attempts: %w", c.retries+1, lastErr)
|
||||
}
|
||||
|
||||
type retryableError struct {
|
||||
err error
|
||||
status int
|
||||
err error
|
||||
status int
|
||||
retryAfter time.Duration
|
||||
}
|
||||
|
||||
func (e *retryableError) Error() string {
|
||||
@@ -232,18 +346,42 @@ func (e *retryableError) Error() string {
|
||||
|
||||
func (e *retryableError) Unwrap() error { return e.err }
|
||||
|
||||
func isRetryable(err error) bool {
|
||||
var re *retryableError
|
||||
return errors.As(err, &re)
|
||||
// parseRetryAfter reads a delta-seconds Retry-After header; HTTP-date form
|
||||
// and malformed values yield 0 (fall back to computed backoff).
|
||||
func parseRetryAfter(h string) time.Duration {
|
||||
secs, err := strconv.Atoi(strings.TrimSpace(h))
|
||||
if err != nil || secs <= 0 {
|
||||
return 0
|
||||
}
|
||||
return time.Duration(secs) * time.Second
|
||||
}
|
||||
|
||||
func (c *OpenAIClient) doRequest(ctx context.Context, payload []byte) (*chatCompletionResponse, error) {
|
||||
var watchdog *idleWatchdog
|
||||
if c.stream {
|
||||
var cancel context.CancelCauseFunc
|
||||
ctx, cancel = context.WithCancelCause(ctx)
|
||||
defer cancel(nil)
|
||||
watchdog = startIdleWatchdog(cancel, c.idleTimeout)
|
||||
defer watchdog.Stop()
|
||||
}
|
||||
out, err := c.send(ctx, payload, watchdog)
|
||||
if err != nil && c.stream && errors.Is(context.Cause(ctx), errIdleTimeout) {
|
||||
return nil, &retryableError{err: fmt.Errorf("no data for %s: %w", c.idleTimeout, errIdleTimeout)}
|
||||
}
|
||||
return out, err
|
||||
}
|
||||
|
||||
func (c *OpenAIClient) send(ctx context.Context, payload []byte, watchdog *idleWatchdog) (*chatCompletionResponse, error) {
|
||||
url := c.endpoint + "/chat/completions"
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("llm: build request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
if c.stream {
|
||||
httpReq.Header.Set("Accept", "text/event-stream")
|
||||
}
|
||||
if c.apiKey != "" {
|
||||
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||
}
|
||||
@@ -253,16 +391,29 @@ func (c *OpenAIClient) doRequest(ctx context.Context, payload []byte) (*chatComp
|
||||
return nil, &retryableError{err: err}
|
||||
}
|
||||
defer func() { _ = httpResp.Body.Close() }() // best-effort; body already fully read below
|
||||
|
||||
respBody, err := io.ReadAll(httpResp.Body)
|
||||
if err != nil {
|
||||
return nil, &retryableError{err: fmt.Errorf("read body: %w", err)}
|
||||
if watchdog != nil {
|
||||
watchdog.Touch()
|
||||
}
|
||||
|
||||
// 429 is as transient as a 5xx for a self-hosted gateway (rate limit or
|
||||
// a briefly saturated upstream); retry it with the same backoff budget.
|
||||
if httpResp.StatusCode >= 500 || httpResp.StatusCode == http.StatusTooManyRequests {
|
||||
return nil, &retryableError{status: httpResp.StatusCode}
|
||||
// 429 and 408 are as transient as a 5xx for a self-hosted gateway (rate
|
||||
// limit, briefly saturated upstream, slow proxy); retry them with the
|
||||
// same backoff budget.
|
||||
if httpResp.StatusCode >= 500 || httpResp.StatusCode == http.StatusTooManyRequests || httpResp.StatusCode == http.StatusRequestTimeout {
|
||||
return nil, &retryableError{
|
||||
status: httpResp.StatusCode,
|
||||
retryAfter: parseRetryAfter(httpResp.Header.Get("Retry-After")),
|
||||
}
|
||||
}
|
||||
|
||||
// A gateway may answer a streaming request with a plain JSON body (an
|
||||
// error, or no streaming support); only parse SSE when it sent SSE.
|
||||
if c.stream && httpResp.StatusCode < 400 && strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream") {
|
||||
return readStream(httpResp.Body, watchdog.Touch)
|
||||
}
|
||||
|
||||
respBody, err := io.ReadAll(io.LimitReader(httpResp.Body, maxResponseBytes))
|
||||
if err != nil {
|
||||
return nil, &retryableError{err: fmt.Errorf("read body: %w", err)}
|
||||
}
|
||||
if httpResp.StatusCode >= 400 {
|
||||
return nil, fmt.Errorf("llm: server returned status %d: %s", httpResp.StatusCode, string(respBody))
|
||||
|
||||
+213
-2
@@ -21,6 +21,12 @@ func newTestServer(t *testing.T, handler http.HandlerFunc) *httptest.Server {
|
||||
return srv
|
||||
}
|
||||
|
||||
func writeOK(w http.ResponseWriter, content string) {
|
||||
resp := chatCompletionResponse{}
|
||||
resp.Choices = append(resp.Choices, testChoice(content, "stop"))
|
||||
_ = json.NewEncoder(w).Encode(resp)
|
||||
}
|
||||
|
||||
// testChoice builds one response choice with a non-null content string, the
|
||||
// common case. Its return type is exactly chatCompletionResponse.Choices's
|
||||
// element type, so callers can append() it without repeating the anonymous
|
||||
@@ -137,7 +143,7 @@ func TestOpenAIClient_Complete_SendsTemperatureZero(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("read request body: %v", err)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(chatCompletionResponse{})
|
||||
writeOK(w, "ok")
|
||||
})
|
||||
|
||||
client := NewOpenAIClient(srv.URL, "", time.Second)
|
||||
@@ -320,7 +326,7 @@ func TestOpenAIClient_Complete_SendsAPIKey(t *testing.T) {
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer secret-key" {
|
||||
t.Errorf("Authorization = %q, want Bearer secret-key", got)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(chatCompletionResponse{})
|
||||
writeOK(w, "ok")
|
||||
})
|
||||
|
||||
client := NewOpenAIClient(srv.URL, "secret-key", time.Second)
|
||||
@@ -394,3 +400,208 @@ func TestOpenAIClient_Complete_MalformedToolArgs(t *testing.T) {
|
||||
t.Errorf("raw args = %q", resp.ToolCalls[0].RawArguments)
|
||||
}
|
||||
}
|
||||
|
||||
func fastClient(url string) *OpenAIClient {
|
||||
c := NewOpenAIClient(url, "", time.Second)
|
||||
c.baseDelay = time.Millisecond
|
||||
return c
|
||||
}
|
||||
|
||||
var oneMsg = Request{Model: "m", Messages: []Message{{Role: "user", Content: "x"}}}
|
||||
|
||||
func TestOpenAIClient_Complete_SendsSeedAndEchoesFingerprint(t *testing.T) {
|
||||
var decoded map[string]any
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
_ = json.Unmarshal(b, &decoded)
|
||||
_, _ = w.Write([]byte(`{"model":"served-1","system_fingerprint":"fp-9","choices":[{"finish_reason":"stop","message":{"content":"hi"}}]}`))
|
||||
})
|
||||
seed := 42
|
||||
req := oneMsg
|
||||
req.Seed = &seed
|
||||
resp, err := fastClient(srv.URL).Complete(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
if decoded["seed"] != float64(42) {
|
||||
t.Errorf("seed = %v, want 42", decoded["seed"])
|
||||
}
|
||||
if resp.Fingerprint != "fp-9" || resp.ServedModel != "served-1" {
|
||||
t.Errorf("fingerprint/model = %q/%q, want fp-9/served-1", resp.Fingerprint, resp.ServedModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Complete_OmitsSeedWhenUnset(t *testing.T) {
|
||||
var decoded map[string]any
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
_ = json.Unmarshal(b, &decoded)
|
||||
writeOK(w, "ok")
|
||||
})
|
||||
if _, err := fastClient(srv.URL).Complete(context.Background(), oneMsg); err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
if _, present := decoded["seed"]; present {
|
||||
t.Errorf("seed present in body without Request.Seed: %v", decoded)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Complete_RetriesEmptyChoicesThenSucceeds(t *testing.T) {
|
||||
var calls int32
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
if atomic.AddInt32(&calls, 1) == 1 {
|
||||
_, _ = w.Write([]byte(`{"choices":[]}`))
|
||||
return
|
||||
}
|
||||
writeOK(w, "ok")
|
||||
})
|
||||
resp, err := fastClient(srv.URL).Complete(context.Background(), oneMsg)
|
||||
if err != nil || resp.Text != "ok" {
|
||||
t.Fatalf("Complete() = %q, %v; want ok, nil", resp.Text, err)
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 2 {
|
||||
t.Errorf("calls = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Complete_RetriesOn408(t *testing.T) {
|
||||
var calls int32
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
if atomic.AddInt32(&calls, 1) == 1 {
|
||||
w.WriteHeader(http.StatusRequestTimeout)
|
||||
return
|
||||
}
|
||||
writeOK(w, "ok")
|
||||
})
|
||||
if _, err := fastClient(srv.URL).Complete(context.Background(), oneMsg); err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 2 {
|
||||
t.Errorf("calls = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClient_WithRetries(t *testing.T) {
|
||||
var calls int32
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&calls, 1)
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
})
|
||||
if _, err := fastClient(srv.URL).WithRetries(5).Complete(context.Background(), oneMsg); err == nil {
|
||||
t.Fatal("Complete() error = nil, want error")
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 6 {
|
||||
t.Errorf("calls = %d, want 6 (1 + 5 retries)", got)
|
||||
}
|
||||
|
||||
atomic.StoreInt32(&calls, 0)
|
||||
if _, err := fastClient(srv.URL).WithRetries(0).Complete(context.Background(), oneMsg); err == nil {
|
||||
t.Fatal("Complete() error = nil, want error")
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 1 {
|
||||
t.Errorf("calls = %d, want 1 (retries disabled)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Complete_HonorsRetryAfter(t *testing.T) {
|
||||
var calls int32
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
if atomic.AddInt32(&calls, 1) == 1 {
|
||||
w.Header().Set("Retry-After", "1")
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
writeOK(w, "ok")
|
||||
})
|
||||
start := time.Now()
|
||||
if _, err := fastClient(srv.URL).Complete(context.Background(), oneMsg); err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
if elapsed := time.Since(start); elapsed < 900*time.Millisecond {
|
||||
t.Errorf("elapsed = %v, want >= ~1s from Retry-After", elapsed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClient_backoff(t *testing.T) {
|
||||
c := NewOpenAIClient("http://x", "", time.Second)
|
||||
for attempt, base := range map[int]time.Duration{1: 250 * time.Millisecond, 2: 500 * time.Millisecond, 3: time.Second} {
|
||||
got := c.backoff(attempt, 0)
|
||||
if got < base || got > base+base/3 {
|
||||
t.Errorf("backoff(%d) = %v, want in [%v, %v]", attempt, got, base, base+base/3)
|
||||
}
|
||||
}
|
||||
if got := c.backoff(30, 0); got > maxRetryDelay*5/4 {
|
||||
t.Errorf("backoff(30) = %v, want capped near %v", got, maxRetryDelay)
|
||||
}
|
||||
if got := c.backoff(1, time.Hour); got != maxRetryDelay {
|
||||
t.Errorf("backoff with huge Retry-After = %v, want cap %v", got, maxRetryDelay)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewOpenAIClient_TrimsTrailingSlash(t *testing.T) {
|
||||
var path string
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
path = r.URL.Path
|
||||
writeOK(w, "ok")
|
||||
})
|
||||
if _, err := fastClient(srv.URL+"/v1/").Complete(context.Background(), oneMsg); err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
if path != "/v1/chat/completions" {
|
||||
t.Errorf("path = %q, want /v1/chat/completions", path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Complete_TimeoutNotRetriedWhenDisabled(t *testing.T) {
|
||||
var calls int32
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&calls, 1)
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
writeOK(w, "late")
|
||||
})
|
||||
c := NewOpenAIClient(srv.URL, "", 50*time.Millisecond).WithRetryTimeouts(false)
|
||||
c.baseDelay = time.Millisecond
|
||||
_, err := c.Complete(context.Background(), oneMsg)
|
||||
if err == nil || !strings.Contains(err.Error(), "not retried") {
|
||||
t.Fatalf("err = %v, want a timeout not retried", err)
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 1 {
|
||||
t.Errorf("calls = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Complete_TimeoutRetriedByDefault(t *testing.T) {
|
||||
var calls int32
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
if atomic.AddInt32(&calls, 1) == 1 {
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
}
|
||||
writeOK(w, "ok")
|
||||
})
|
||||
c := NewOpenAIClient(srv.URL, "", 100*time.Millisecond)
|
||||
c.baseDelay = time.Millisecond
|
||||
if _, err := c.Complete(context.Background(), oneMsg); err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 2 {
|
||||
t.Errorf("calls = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Complete_PerModelTimeout(t *testing.T) {
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
writeOK(w, "slow but fine")
|
||||
})
|
||||
c := NewOpenAIClient(srv.URL, "", 50*time.Millisecond).
|
||||
WithRetryTimeouts(false).
|
||||
WithModelTimeouts(map[string]time.Duration{"slow-local": 2 * time.Second})
|
||||
req := oneMsg
|
||||
req.Model = "slow-local"
|
||||
if resp, err := c.Complete(context.Background(), req); err != nil || resp.Text != "slow but fine" {
|
||||
t.Fatalf("slow-local = %q, %v; want success under its own timeout", resp.Text, err)
|
||||
}
|
||||
if _, err := c.Complete(context.Background(), oneMsg); err == nil {
|
||||
t.Error("default-timeout model succeeded, want timeout")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// defaultIdleTimeout is how long a stream may stay silent (no bytes at all,
|
||||
// including before the first token) before the attempt is abandoned. It
|
||||
// must exceed a slow local model's prompt-processing time.
|
||||
const defaultIdleTimeout = 5 * time.Minute
|
||||
|
||||
// maxStreamLine bounds one SSE line.
|
||||
const maxStreamLine = 1 << 20
|
||||
|
||||
const (
|
||||
sseDataPrefix = "data:"
|
||||
sseDone = "[DONE]"
|
||||
)
|
||||
|
||||
// errIdleTimeout marks an attempt abandoned because the stream went silent.
|
||||
var errIdleTimeout = errors.New("llm: stream idle timeout")
|
||||
|
||||
// errStreamTruncated marks a stream that ended without [DONE] or a
|
||||
// finish_reason, so the answer may be incomplete.
|
||||
var errStreamTruncated = errors.New("llm: stream ended before completion")
|
||||
|
||||
type streamToolCall struct {
|
||||
Function struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
Index int `json:"index"`
|
||||
}
|
||||
|
||||
type streamChunk struct {
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
Usage *struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
} `json:"usage"`
|
||||
Model string `json:"model"`
|
||||
SystemFingerprint string `json:"system_fingerprint"`
|
||||
Choices []struct {
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
Delta struct {
|
||||
Content *string `json:"content"`
|
||||
ToolCalls []streamToolCall `json:"tool_calls"`
|
||||
} `json:"delta"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
|
||||
// readStream folds an OpenAI-style SSE stream into the same response shape
|
||||
// the non-streaming path decodes. onActivity is called for every line read,
|
||||
// including keep-alive comments, so the caller can reset its idle timer.
|
||||
// reasoning_content deltas count as activity but are never accumulated,
|
||||
// matching the non-streaming path.
|
||||
func readStream(r io.Reader, onActivity func()) (*chatCompletionResponse, error) {
|
||||
sc := bufio.NewScanner(r)
|
||||
sc.Buffer(make([]byte, 0, 64<<10), maxStreamLine)
|
||||
|
||||
var (
|
||||
out chatCompletionResponse
|
||||
text strings.Builder
|
||||
sawText bool
|
||||
finish string
|
||||
done bool
|
||||
order []int
|
||||
calls = map[int]*apiToolCall{}
|
||||
tooLarge bool
|
||||
)
|
||||
|
||||
for sc.Scan() {
|
||||
onActivity()
|
||||
line := strings.TrimSpace(sc.Text())
|
||||
if !strings.HasPrefix(line, sseDataPrefix) {
|
||||
continue // blank separator, comment or other SSE field
|
||||
}
|
||||
data := strings.TrimSpace(strings.TrimPrefix(line, sseDataPrefix))
|
||||
if data == sseDone {
|
||||
done = true
|
||||
break
|
||||
}
|
||||
var ch streamChunk
|
||||
if err := json.Unmarshal([]byte(data), &ch); err != nil {
|
||||
return nil, &retryableError{err: fmt.Errorf("decode stream chunk: %w", err)}
|
||||
}
|
||||
if ch.Error != nil {
|
||||
out.Error = ch.Error
|
||||
return &out, nil
|
||||
}
|
||||
if ch.Model != "" {
|
||||
out.Model = ch.Model
|
||||
}
|
||||
if ch.SystemFingerprint != "" {
|
||||
out.SystemFingerprint = ch.SystemFingerprint
|
||||
}
|
||||
if ch.Usage != nil {
|
||||
out.Usage.PromptTokens = ch.Usage.PromptTokens
|
||||
out.Usage.CompletionTokens = ch.Usage.CompletionTokens
|
||||
}
|
||||
if len(ch.Choices) == 0 {
|
||||
continue
|
||||
}
|
||||
c := ch.Choices[0]
|
||||
if c.Delta.Content != nil {
|
||||
sawText = true
|
||||
if text.Len()+len(*c.Delta.Content) > maxResponseBytes {
|
||||
tooLarge = true
|
||||
} else {
|
||||
text.WriteString(*c.Delta.Content)
|
||||
}
|
||||
}
|
||||
for _, tc := range c.Delta.ToolCalls {
|
||||
acc, ok := calls[tc.Index]
|
||||
if !ok {
|
||||
acc = &apiToolCall{Type: "function"}
|
||||
calls[tc.Index] = acc
|
||||
order = append(order, tc.Index)
|
||||
}
|
||||
if tc.Function.Name != "" {
|
||||
acc.Function.Name = tc.Function.Name
|
||||
}
|
||||
acc.Function.Arguments += tc.Function.Arguments
|
||||
}
|
||||
if c.FinishReason != nil && *c.FinishReason != "" {
|
||||
finish = *c.FinishReason
|
||||
}
|
||||
}
|
||||
if err := sc.Err(); err != nil {
|
||||
return nil, &retryableError{err: fmt.Errorf("read stream: %w", err)}
|
||||
}
|
||||
if tooLarge {
|
||||
return nil, fmt.Errorf("llm: streamed response exceeded %d bytes", maxResponseBytes)
|
||||
}
|
||||
if !done && finish == "" {
|
||||
return nil, &retryableError{err: errStreamTruncated}
|
||||
}
|
||||
|
||||
msg := responseMsg{Role: "assistant"}
|
||||
if sawText {
|
||||
s := text.String()
|
||||
msg.Content = &s
|
||||
}
|
||||
for _, idx := range order {
|
||||
msg.ToolCalls = append(msg.ToolCalls, *calls[idx])
|
||||
}
|
||||
out.Choices = append(out.Choices, struct {
|
||||
FinishReason string `json:"finish_reason"`
|
||||
Message responseMsg `json:"message"`
|
||||
}{FinishReason: finish, Message: msg})
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
// idleWatchdog cancels an attempt when no activity is seen for idle.
|
||||
type idleWatchdog struct {
|
||||
timer *time.Timer
|
||||
idle time.Duration
|
||||
}
|
||||
|
||||
func startIdleWatchdog(cancel context.CancelCauseFunc, idle time.Duration) *idleWatchdog {
|
||||
return &idleWatchdog{idle: idle, timer: time.AfterFunc(idle, func() { cancel(errIdleTimeout) })}
|
||||
}
|
||||
|
||||
func (w *idleWatchdog) Touch() { w.timer.Reset(w.idle) }
|
||||
func (w *idleWatchdog) Stop() { w.timer.Stop() }
|
||||
@@ -0,0 +1,170 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func sseServer(t *testing.T, frames func(w http.ResponseWriter, fl http.Flusher)) string {
|
||||
t.Helper()
|
||||
return newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fl := w.(http.Flusher)
|
||||
frames(w, fl)
|
||||
}).URL
|
||||
}
|
||||
|
||||
func streamClient(url string, idle time.Duration) *OpenAIClient {
|
||||
c := fastClient(url).WithStreaming(true, idle)
|
||||
return c
|
||||
}
|
||||
|
||||
func TestStream_AggregatesContentUsageAndFingerprint(t *testing.T) {
|
||||
var body map[string]any
|
||||
url := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
_ = json.Unmarshal(b, &body)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
send(w, ": keepalive\n\n")
|
||||
send(w, `data: {"model":"srv","system_fingerprint":"fp","choices":[{"delta":{"reasoning_content":"thinking"}}]}`+"\n\n")
|
||||
send(w, `data: {"choices":[{"delta":{"content":"Hel"}}]}`+"\n\n")
|
||||
send(w, `data: {"choices":[{"delta":{"content":"lo"},"finish_reason":"stop"}]}`+"\n\n")
|
||||
send(w, `data: {"choices":[],"usage":{"prompt_tokens":5,"completion_tokens":7}}`+"\n\n")
|
||||
send(w, "data: [DONE]\n\n")
|
||||
}).URL
|
||||
resp, err := streamClient(url, time.Second).Complete(context.Background(), oneMsg)
|
||||
if err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
if resp.Text != "Hello" || resp.FinishReason != "stop" || resp.PromptTokens != 5 || resp.CompletionTokens != 7 {
|
||||
t.Errorf("resp = %+v", resp)
|
||||
}
|
||||
if resp.Fingerprint != "fp" || resp.ServedModel != "srv" {
|
||||
t.Errorf("fingerprint/model = %q/%q", resp.Fingerprint, resp.ServedModel)
|
||||
}
|
||||
if body["stream"] != true {
|
||||
t.Errorf("stream = %v, want true", body["stream"])
|
||||
}
|
||||
if opts, _ := body["stream_options"].(map[string]any); opts["include_usage"] != true {
|
||||
t.Errorf("stream_options = %v, want include_usage", body["stream_options"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestStream_AssemblesToolCallFragments(t *testing.T) {
|
||||
url := sseServer(t, func(w http.ResponseWriter, fl http.Flusher) {
|
||||
send(w, `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"get_w","arguments":"{\"ci"}}]}}]}`+"\n\n")
|
||||
send(w, `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"ty\":\"X\"}"}},{"index":1,"function":{"name":"b","arguments":"{}"}}]},"finish_reason":"tool_calls"}]}`+"\n\n")
|
||||
send(w, "data: [DONE]\n\n")
|
||||
})
|
||||
resp, err := streamClient(url, time.Second).Complete(context.Background(), oneMsg)
|
||||
if err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
if len(resp.ToolCalls) != 2 || resp.ToolCalls[0].Name != "get_w" || resp.ToolCalls[0].Arguments["city"] != "X" || !resp.ToolCalls[1].Decoded {
|
||||
t.Errorf("ToolCalls = %+v", resp.ToolCalls)
|
||||
}
|
||||
if resp.Text != "" {
|
||||
t.Errorf("Text = %q, want empty", resp.Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStream_ReasoningOnlyTruncationYieldsEmptyText(t *testing.T) {
|
||||
url := sseServer(t, func(w http.ResponseWriter, fl http.Flusher) {
|
||||
send(w, `data: {"choices":[{"delta":{"reasoning_content":"..."},"finish_reason":"length"}]}`+"\n\n")
|
||||
send(w, "data: [DONE]\n\n")
|
||||
})
|
||||
resp, err := streamClient(url, time.Second).Complete(context.Background(), oneMsg)
|
||||
if err != nil || resp.Text != "" || resp.FinishReason != FinishReasonLength {
|
||||
t.Fatalf("resp = %+v, err = %v; want empty text, length", resp, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStream_TruncatedStreamIsRetried(t *testing.T) {
|
||||
var calls int32
|
||||
url := sseServer(t, func(w http.ResponseWriter, fl http.Flusher) {
|
||||
if atomic.AddInt32(&calls, 1) == 1 {
|
||||
send(w, `data: {"choices":[{"delta":{"content":"par"}}]}`+"\n\n")
|
||||
return // connection ends with no [DONE] or finish_reason
|
||||
}
|
||||
send(w, `data: {"choices":[{"delta":{"content":"whole"},"finish_reason":"stop"}]}`+"\n\n")
|
||||
send(w, "data: [DONE]\n\n")
|
||||
})
|
||||
resp, err := streamClient(url, time.Second).Complete(context.Background(), oneMsg)
|
||||
if err != nil || resp.Text != "whole" {
|
||||
t.Fatalf("resp = %q, err = %v; want whole", resp.Text, err)
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 2 {
|
||||
t.Errorf("calls = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStream_IdleTimeoutAbandonsHungServer(t *testing.T) {
|
||||
var calls int32
|
||||
url := sseServer(t, func(w http.ResponseWriter, fl http.Flusher) {
|
||||
atomic.AddInt32(&calls, 1)
|
||||
send(w, `data: {"choices":[{"delta":{"content":"x"}}]}`+"\n\n")
|
||||
fl.Flush()
|
||||
time.Sleep(2 * time.Second) // goes silent
|
||||
})
|
||||
c := streamClient(url, 150*time.Millisecond).WithRetryTimeouts(false)
|
||||
start := time.Now()
|
||||
_, err := c.Complete(context.Background(), oneMsg)
|
||||
if err == nil || !strings.Contains(err.Error(), "not retried") {
|
||||
t.Fatalf("err = %v, want idle timeout not retried", err)
|
||||
}
|
||||
if elapsed := time.Since(start); elapsed > time.Second {
|
||||
t.Errorf("took %v, want abandoned near the 150ms idle timeout", elapsed)
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 1 {
|
||||
t.Errorf("calls = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStream_SlowButProgressingSurvivesPastIdleTimeout(t *testing.T) {
|
||||
url := sseServer(t, func(w http.ResponseWriter, fl http.Flusher) {
|
||||
for i := 0; i < 6; i++ {
|
||||
send(w, `data: {"choices":[{"delta":{"content":"a"}}]}`+"\n\n")
|
||||
fl.Flush()
|
||||
time.Sleep(80 * time.Millisecond)
|
||||
}
|
||||
send(w, `data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`+"\n\ndata: [DONE]\n\n")
|
||||
})
|
||||
// total ~480ms > idle 200ms, but never silent for 200ms.
|
||||
resp, err := streamClient(url, 200*time.Millisecond).Complete(context.Background(), oneMsg)
|
||||
if err != nil || resp.Text != "aaaaaa" {
|
||||
t.Fatalf("resp = %q, err = %v; want aaaaaa", resp.Text, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStream_ErrorChunkIsRetried(t *testing.T) {
|
||||
var calls int32
|
||||
url := sseServer(t, func(w http.ResponseWriter, fl http.Flusher) {
|
||||
if atomic.AddInt32(&calls, 1) == 1 {
|
||||
send(w, `data: {"error":{"message":"replica warming"}}`+"\n\n")
|
||||
return
|
||||
}
|
||||
send(w, `data: {"choices":[{"delta":{"content":"ok"},"finish_reason":"stop"}]}`+"\n\ndata: [DONE]\n\n")
|
||||
})
|
||||
resp, err := streamClient(url, time.Second).Complete(context.Background(), oneMsg)
|
||||
if err != nil || resp.Text != "ok" {
|
||||
t.Fatalf("resp = %q, err = %v", resp.Text, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStream_PlainJSONFallbackStillParsed(t *testing.T) {
|
||||
srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
writeOK(w, "plain")
|
||||
})
|
||||
resp, err := streamClient(srv.URL, time.Second).Complete(context.Background(), oneMsg)
|
||||
if err != nil || resp.Text != "plain" {
|
||||
t.Fatalf("resp = %q, err = %v; want plain JSON body accepted", resp.Text, err)
|
||||
}
|
||||
}
|
||||
|
||||
func send(w io.Writer, s string) { _, _ = io.WriteString(w, s) }
|
||||
Reference in New Issue
Block a user