diff --git a/README.md b/README.md index e79f40d..b5c47e6 100644 --- a/README.md +++ b/README.md @@ -75,6 +75,8 @@ Both `list` and `run` accept: `60s`). - `--out FILE` - write the run as a JSON artifact (raw per-attempt results plus the per-test discrimination rollup) for later `llmtest compare`. +- `--seed N` - send `seed: N` with every request (overrides the config's + `seed`). Whether the backend honors it is backend-specific. - `--quiet` - suppress progress output on stderr. Progress is on by default: a run can issue hundreds of requests at 30-50s each, and a silent CLI reads as a hang. diff --git a/cmd/llmtest/main.go b/cmd/llmtest/main.go index a92d4d7..f93d224 100644 --- a/cmd/llmtest/main.go +++ b/cmd/llmtest/main.go @@ -4,11 +4,15 @@ package main import ( "context" + "errors" "flag" "fmt" "os" + "os/signal" "sort" "strings" + "syscall" + "time" "github.com/lukaszraczylo/llm-testbench/internal/config" "github.com/lukaszraczylo/llm-testbench/internal/llm" @@ -121,6 +125,7 @@ func runCommand(args []string) error { repeat := fs.Int("repeat", 1, "samples per (model, test); >1 exposes response instability at temperature 0") out := fs.String("out", "", "write the run as a JSON artifact to this file (for llmtest compare)") quiet := fs.Bool("quiet", false, "suppress progress output on stderr") + seedFlag := fs.Int("seed", 0, "send this seed with every request (overrides config's seed)") if err := fs.Parse(args); err != nil { return err } @@ -132,7 +137,7 @@ func runCommand(args []string) error { models := cfg.Models if *modelsCSV != "" { - models = splitCSV(*modelsCSV) + models = dedupe(splitCSV(*modelsCSV)) } reqTimeout := cfg.RequestTimeout @@ -145,7 +150,18 @@ func runCommand(args []string) error { conc = *concurrency } - client := llm.NewOpenAIClient(cfg.Endpoint, cfg.APIKey, reqTimeout) + seed := cfg.Seed + fs.Visit(func(f *flag.Flag) { + if f.Name == "seed" { + seed = seedFlag + } + }) + + client := llm.NewOpenAIClient(cfg.Endpoint, cfg.APIKey, reqTimeout). + WithRetries(cfg.MaxRetries). + WithRetryTimeouts(cfg.RetryTimeouts). + WithModelTimeouts(cfg.ModelTimeouts). + WithStreaming(cfg.Stream, cfg.IdleTimeout) registry := tests.All() selected, err := selectTests(registry, *testsCSV, shared.category, shared.subcategory) @@ -164,10 +180,23 @@ func runCommand(args []string) error { MaxTokensDefault: cfg.MaxTokensDefault, Reporter: reporter, Repeat: *repeat, + Seed: seed, + ModelConcurrency: cfg.ModelConcurrency, }) - ctx := context.Background() + // First Ctrl-C stops dispatching and lets in-flight calls drain so the + // partial run can still be saved; stop() restores default handling, so + // a second Ctrl-C kills the process. + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + startedAt := time.Now().UTC() results := r.Run(ctx, models, selected) + interrupted := ctx.Err() != nil + if interrupted { + stop() + results = dropCancelled(results) + fmt.Fprintf(os.Stderr, "llmtest: interrupted; reporting %d completed results\n", len(results)) + } f, err := validateFormat(*format) if err != nil { @@ -175,7 +204,17 @@ func runCommand(args []string) error { } if *out != "" { - if err := report.WriteArtifact(*out, selected, models, results); err != nil { + meta := report.RunMeta{ + Version: version, + StartedAt: startedAt.Format(time.RFC3339), + Endpoint: cfg.Endpoint, + Seed: seed, + Temperature: requestTemperature, + MaxTokensDefault: cfg.MaxTokensDefault, + Repeat: max(*repeat, 1), + Interrupted: interrupted, + } + if err := report.WriteArtifact(*out, meta, selected, models, results); err != nil { return fmt.Errorf("write --out: %w", err) } } @@ -183,6 +222,19 @@ func runCommand(args []string) error { return report.Render(os.Stdout, f, selected, models, results) } +// dropCancelled removes results that never ran because the run was +// interrupted, so they are not saved as failures. +func dropCancelled(results []runner.Result) []runner.Result { + kept := results[:0:0] + for _, r := range results { + if r.Err != nil && errors.Is(r.Err, context.Canceled) { + continue + } + kept = append(kept, r) + } + return kept +} + // selectTests resolves the test subset: --tests (exact IDs, error on any // unknown ID) overrides --category/--subcategory filtering. func selectTests(registry *testkit.Registry, testsCSV, category, subcategory string) ([]testkit.Test, error) { @@ -234,7 +286,7 @@ func compareCommand(args []string) error { return err } - return report.RenderCompare(os.Stdout, report.CompareArtifacts(baseline, current)) + return report.RenderCompare(os.Stdout, report.CompareArtifacts(baseline, current), report.CompareWarnings(baseline, current)...) } // healthCommand audits one or more saved artifacts for suite health: @@ -271,6 +323,20 @@ func validateFormat(format string) (report.Format, error) { } } +// dedupe drops repeated entries, keeping first-seen order; a model listed +// twice would otherwise run, and be averaged, twice. +func dedupe(in []string) []string { + seen := make(map[string]bool, len(in)) + out := make([]string, 0, len(in)) + for _, v := range in { + if !seen[v] { + seen[v] = true + out = append(out, v) + } + } + return out +} + func splitCSV(s string) []string { parts := strings.Split(s, ",") out := make([]string, 0, len(parts)) diff --git a/cmd/llmtest/main_test.go b/cmd/llmtest/main_test.go index 50c66e5..c6632d8 100644 --- a/cmd/llmtest/main_test.go +++ b/cmd/llmtest/main_test.go @@ -1,11 +1,15 @@ package main import ( + "context" + "errors" + "fmt" "strings" "testing" "github.com/lukaszraczylo/llm-testbench/internal/eval" "github.com/lukaszraczylo/llm-testbench/internal/report" + "github.com/lukaszraczylo/llm-testbench/internal/runner" "github.com/lukaszraczylo/llm-testbench/internal/testkit" ) @@ -101,3 +105,22 @@ func TestHealthCommand_RequiresArtifacts(t *testing.T) { t.Errorf("healthCommand(missing file) error = %v, want load error naming the file", err) } } + +func TestDedupe(t *testing.T) { + got := dedupe([]string{"a", "b", "a", "c", "b"}) + if strings.Join(got, ",") != "a,b,c" { + t.Errorf("dedupe = %v, want [a b c]", got) + } +} + +func TestDropCancelled(t *testing.T) { + in := []runner.Result{ + {TestID: "ok"}, + {TestID: "cancelled", Err: fmt.Errorf("wrapped: %w", context.Canceled)}, + {TestID: "real-error", Err: errors.New("boom")}, + } + got := dropCancelled(in) + if len(got) != 2 || got[0].TestID != "ok" || got[1].TestID != "real-error" { + t.Errorf("dropCancelled = %+v, want ok and real-error kept", got) + } +} diff --git a/config.example.yaml b/config.example.yaml index c6e1754..3567dd1 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -5,7 +5,19 @@ models: - uni/qwen3.6-27b - uni/btl-4 concurrency: 8 # default when omitted; matches internal/config's default -request_timeout: 900s # default when omitted -max_tokens_default: 12000 # default when omitted; a reasoning model can spend +request_timeout: 900s # default when omitted is 300s +max_tokens_default: 40000 # default when omitted; a reasoning model can spend # thousands of tokens on reasoning_content before # writing an answer, so this must stay generous +max_retries: 2 # retries per request after the first try (default 2; 0 disables). + # Raise for flaky gateways: backoff is exponential and + # honors Retry-After. +# seed: 1234 # optional; sent with every request for backends that honor it +# Slow local models: give them their own timeout and queue depth. +# retry_timeouts: false # default true; false stops a timed-out call re-running a full slow generation +# model_timeouts: # overrides request_timeout per model +# local/big-model: 45m +# model_concurrency: # caps in-flight calls per model (global concurrency still applies) +# local/big-model: 1 +# stream: false # true = SSE; detects a hung server after idle_timeout instead of the full request timeout +# idle_timeout: 5m # max silence in a stream (must exceed prompt-processing time); default 5m diff --git a/internal/config/config.go b/internal/config/config.go index 9f82339..90fbb01 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -3,6 +3,7 @@ package config import ( "fmt" + "net/url" "os" "time" @@ -27,27 +28,42 @@ const ( defaultConcurrency = 8 defaultRequestTimeout = 300 * time.Second defaultMaxTokens = 40000 + defaultMaxRetries = 2 ) // Config is the top-level configuration loaded from config.yaml. type Config struct { - Endpoint string `yaml:"endpoint"` - APIKey string `yaml:"api_key"` - Models []string `yaml:"models"` - Concurrency int `yaml:"concurrency"` - RequestTimeout time.Duration `yaml:"request_timeout"` - MaxTokensDefault int `yaml:"max_tokens_default"` + Seed *int `yaml:"seed"` + ModelTimeouts map[string]time.Duration `yaml:"-"` + ModelConcurrency map[string]int `yaml:"model_concurrency"` + Endpoint string `yaml:"endpoint"` + APIKey string `yaml:"api_key"` + Models []string `yaml:"models"` + Concurrency int `yaml:"concurrency"` + RequestTimeout time.Duration `yaml:"request_timeout"` + MaxTokensDefault int `yaml:"max_tokens_default"` + MaxRetries int `yaml:"max_retries"` + IdleTimeout time.Duration `yaml:"-"` + RetryTimeouts bool `yaml:"retry_timeouts"` + Stream bool `yaml:"stream"` } // rawConfig mirrors Config but with RequestTimeout as a string, since // yaml.v3 does not natively decode Go duration strings. type rawConfig struct { - Endpoint string `yaml:"endpoint"` - APIKey string `yaml:"api_key"` - RequestTimeout string `yaml:"request_timeout"` - Models []string `yaml:"models"` - Concurrency int `yaml:"concurrency"` - MaxTokensDefault int `yaml:"max_tokens_default"` + MaxRetries *int `yaml:"max_retries"` + Seed *int `yaml:"seed"` + RetryTimeouts *bool `yaml:"retry_timeouts"` + ModelTimeouts map[string]string `yaml:"model_timeouts"` + ModelConcurrency map[string]int `yaml:"model_concurrency"` + APIKey string `yaml:"api_key"` + Endpoint string `yaml:"endpoint"` + RequestTimeout string `yaml:"request_timeout"` + IdleTimeout string `yaml:"idle_timeout"` + Models []string `yaml:"models"` + Concurrency int `yaml:"concurrency"` + MaxTokensDefault int `yaml:"max_tokens_default"` + Stream bool `yaml:"stream"` } // Load reads and parses the YAML config at path, applies the LLMTB_API_KEY @@ -78,6 +94,31 @@ func parse(data []byte) (Config, error) { Models: raw.Models, Concurrency: raw.Concurrency, MaxTokensDefault: raw.MaxTokensDefault, + MaxRetries: defaultMaxRetries, + Seed: raw.Seed, + } + if raw.MaxRetries != nil { + cfg.MaxRetries = *raw.MaxRetries + } + cfg.RetryTimeouts = raw.RetryTimeouts == nil || *raw.RetryTimeouts + cfg.ModelConcurrency = raw.ModelConcurrency + cfg.Stream = raw.Stream + if raw.IdleTimeout != "" { + d, err := time.ParseDuration(raw.IdleTimeout) + if err != nil { + return Config{}, fmt.Errorf("config: parse idle_timeout %q: %w", raw.IdleTimeout, err) + } + cfg.IdleTimeout = d + } + if len(raw.ModelTimeouts) > 0 { + cfg.ModelTimeouts = make(map[string]time.Duration, len(raw.ModelTimeouts)) + for m, v := range raw.ModelTimeouts { + d, err := time.ParseDuration(v) + if err != nil { + return Config{}, fmt.Errorf("config: parse model_timeouts[%q] %q: %w", m, v, err) + } + cfg.ModelTimeouts[m] = d + } } if raw.RequestTimeout != "" { @@ -123,10 +164,35 @@ func (c Config) Validate() error { if len(c.Models) == 0 { return fmt.Errorf("config: models must list at least one model") } + u, err := url.Parse(c.Endpoint) + if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" { + return fmt.Errorf("config: endpoint %q must be an http(s) URL", c.Endpoint) + } + seen := make(map[string]bool, len(c.Models)) for i, m := range c.Models { if m == "" { return fmt.Errorf("config: models[%d] must not be empty", i) } + if seen[m] { + return fmt.Errorf("config: models[%d] duplicates %q", i, m) + } + seen[m] = true + } + for m, d := range c.ModelTimeouts { + if d <= 0 { + return fmt.Errorf("config: model_timeouts[%q] must be > 0, got %s", m, d) + } + } + for m, n := range c.ModelConcurrency { + if n <= 0 { + return fmt.Errorf("config: model_concurrency[%q] must be > 0, got %d", m, n) + } + } + if c.IdleTimeout < 0 { + return fmt.Errorf("config: idle_timeout must be >= 0, got %s", c.IdleTimeout) + } + if c.MaxRetries < 0 { + return fmt.Errorf("config: max_retries must be >= 0, got %d", c.MaxRetries) } if c.Concurrency <= 0 { return fmt.Errorf("config: concurrency must be > 0, got %d", c.Concurrency) diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 2221843..f2072a7 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -160,6 +160,9 @@ func TestValidate(t *testing.T) { {name: "zero concurrency", mutate: func(c Config) Config { c.Concurrency = 0; return c }}, {name: "negative concurrency", mutate: func(c Config) Config { c.Concurrency = -1; return c }}, {name: "zero timeout", mutate: func(c Config) Config { c.RequestTimeout = 0; return c }}, + {name: "non-http endpoint", mutate: func(c Config) Config { c.Endpoint = "llm-gateway.example.com/v1"; return c }}, + {name: "duplicate model", mutate: func(c Config) Config { c.Models = []string{"m", "m"}; return c }}, + {name: "negative retries", mutate: func(c Config) Config { c.MaxRetries = -1; return c }}, {name: "zero max tokens", mutate: func(c Config) Config { c.MaxTokensDefault = 0; return c }}, } for _, tt := range tests { @@ -194,3 +197,65 @@ func TestLoad_MissingFile(t *testing.T) { t.Fatal("Load() error = nil, want error for missing file") } } + +func TestParse_RetriesAndSeed(t *testing.T) { + base := "endpoint: https://x\nmodels: [m]\n" + + cfg, err := parse([]byte(base)) + if err != nil { + t.Fatalf("parse() error = %v", err) + } + if cfg.MaxRetries != defaultMaxRetries || cfg.Seed != nil { + t.Errorf("defaults = retries %d seed %v, want %d nil", cfg.MaxRetries, cfg.Seed, defaultMaxRetries) + } + + cfg, err = parse([]byte(base + "max_retries: 0\nseed: 7\n")) + if err != nil { + t.Fatalf("parse() error = %v", err) + } + if cfg.MaxRetries != 0 { + t.Errorf("explicit max_retries 0 = %d, want 0 (retrying disabled)", cfg.MaxRetries) + } + if cfg.Seed == nil || *cfg.Seed != 7 { + t.Errorf("Seed = %v, want 7", cfg.Seed) + } +} + +func TestParse_ModelOverrides(t *testing.T) { + cfg, err := parse([]byte(` +endpoint: https://x +models: [a, b] +retry_timeouts: false +model_timeouts: {a: 45m} +model_concurrency: {a: 1} +`)) + if err != nil { + t.Fatalf("parse() error = %v", err) + } + if cfg.RetryTimeouts || cfg.ModelTimeouts["a"] != 45*time.Minute || cfg.ModelConcurrency["a"] != 1 { + t.Errorf("got %+v", cfg) + } + def, _ := parse([]byte("endpoint: https://x\nmodels: [a]\n")) + if !def.RetryTimeouts { + t.Error("retry_timeouts default = false, want true (legacy behavior)") + } + for _, bad := range []string{"model_timeouts: {a: nope}", "model_timeouts: {a: 0s}", "model_concurrency: {a: 0}"} { + if _, err := parse([]byte("endpoint: https://x\nmodels: [a]\n" + bad)); err == nil { + t.Errorf("parse(%q) error = nil, want error", bad) + } + } +} + +func TestParse_Streaming(t *testing.T) { + cfg, err := parse([]byte("endpoint: https://x\nmodels: [a]\nstream: true\nidle_timeout: 10m\n")) + if err != nil || !cfg.Stream || cfg.IdleTimeout != 10*time.Minute { + t.Fatalf("cfg = %+v, err = %v", cfg, err) + } + def, _ := parse([]byte("endpoint: https://x\nmodels: [a]\n")) + if def.Stream { + t.Error("stream default = true, want false") + } + if _, err := parse([]byte("endpoint: https://x\nmodels: [a]\nidle_timeout: nope\n")); err == nil { + t.Error("bad idle_timeout accepted") + } +} diff --git a/internal/llm/llm.go b/internal/llm/llm.go index 263a8b5..18e84fd 100644 --- a/internal/llm/llm.go +++ b/internal/llm/llm.go @@ -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 diff --git a/internal/llm/openai.go b/internal/llm/openai.go index 2f6b7f0..77f1aad 100644 --- a/internal/llm/openai.go +++ b/internal/llm/openai.go @@ -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)) diff --git a/internal/llm/openai_test.go b/internal/llm/openai_test.go index def77df..e0696b3 100644 --- a/internal/llm/openai_test.go +++ b/internal/llm/openai_test.go @@ -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") + } +} diff --git a/internal/llm/stream.go b/internal/llm/stream.go new file mode 100644 index 0000000..5e380a9 --- /dev/null +++ b/internal/llm/stream.go @@ -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() } diff --git a/internal/llm/stream_test.go b/internal/llm/stream_test.go new file mode 100644 index 0000000..d127b2d --- /dev/null +++ b/internal/llm/stream_test.go @@ -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) } diff --git a/internal/report/compare.go b/internal/report/compare.go index 4a5a830..ff63e31 100644 --- a/internal/report/compare.go +++ b/internal/report/compare.go @@ -170,8 +170,11 @@ func CompareArtifacts(baseline, current Artifact) []Comparison { // RenderCompare writes a comparison report: headline counts, then only // changed/missing rows (stable rows are summarized, not listed). -func RenderCompare(w io.Writer, comparisons []Comparison) error { +func RenderCompare(w io.Writer, comparisons []Comparison, warnings ...string) error { ew := &errWriter{w: w} + for _, warn := range warnings { + ew.println("WARNING: " + warn) + } var reg, imp, stable, added, removed, errored int for _, c := range comparisons { @@ -219,3 +222,9 @@ func deltaText(v float64) string { } return fmt.Sprintf("%+.2f", v) } + +// CompareWarnings reports reasons the two artifacts may not be comparable +// like-for-like (catalog, sampling or backend-build drift). +func CompareWarnings(baseline, current Artifact) []string { + return compareWarnings(baseline.Meta, current.Meta) +} diff --git a/internal/report/json.go b/internal/report/json.go index a64e5d7..11b5384 100644 --- a/internal/report/json.go +++ b/internal/report/json.go @@ -6,6 +6,7 @@ import ( "io" "math" "os" + "path/filepath" "github.com/lukaszraczylo/llm-testbench/internal/runner" "github.com/lukaszraczylo/llm-testbench/internal/testkit" @@ -22,6 +23,7 @@ type jsonResult struct { Error string `json:"error,omitempty"` ResponseText string `json:"response_text"` FinishReason string `json:"finish_reason,omitempty"` + Fingerprint string `json:"fingerprint,omitempty"` Score float64 `json:"score"` LatencyMS int64 `json:"latency_ms"` PromptTokens int `json:"prompt_tokens"` @@ -36,6 +38,7 @@ type jsonResult struct { // plus the per-test discrimination rollup, so downstream tooling // (compare) can re-derive stability without re-running the suite. type jsonArtifact struct { + Meta *RunMeta `json:"meta,omitempty"` Results []jsonResult `json:"results"` Stats []jsonStatRow `json:"stats,omitempty"` } @@ -77,6 +80,7 @@ func buildArtifact(tests []testkit.Test, models []string, results []runner.Resul Detail: r.Score.Detail, ResponseText: r.ResponseText, FinishReason: r.FinishReason, + Fingerprint: r.Fingerprint, Skipped: r.Score.Skipped, Truncated: r.Truncated(), LatencyMS: r.Latency.Milliseconds(), @@ -150,12 +154,37 @@ func LoadArtifact(path string) (jsonArtifact, error) { return jsonArtifact{Results: legacy}, nil } -// WriteArtifact saves the run as a JSON artifact at path, for later -// `llmtest compare` runs, regardless of the stdout report format. -func WriteArtifact(path string, tests []testkit.Test, models []string, results []runner.Result) error { - b, err := json.MarshalIndent(buildArtifact(tests, models, results), "", " ") +// WriteArtifact saves the run, stamped with meta, as a JSON artifact at +// path for later `llmtest compare` runs, regardless of the stdout report +// format. The write is atomic (temp file + rename), so an interrupted +// process never leaves a truncated artifact that replaced a good one. +func WriteArtifact(path string, meta RunMeta, tests []testkit.Test, models []string, results []runner.Result) error { + art := buildArtifact(tests, models, results) + meta.CatalogHash = CatalogHash(tests) + meta.TestCount = len(tests) + meta.Fingerprints = servedFingerprints(results) + art.Meta = &meta + + b, err := json.MarshalIndent(art, "", " ") if err != nil { return err } - return os.WriteFile(path, append(b, '\n'), 0o644) //nolint:gosec // 0o644: the artifact is a report meant to be shared/committed, not a secret file + + tmp, err := os.CreateTemp(filepath.Dir(path), ".llmtest-artifact-*") + if err != nil { + return err + } + defer func() { _ = os.Remove(tmp.Name()) }() // no-op after a successful rename + if _, err := tmp.Write(append(b, '\n')); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + // 0o644: the artifact is a report meant to be shared/committed, not a secret file. + if err := os.Chmod(tmp.Name(), 0o644); err != nil { //nolint:gosec // see comment above + return err + } + return os.Rename(tmp.Name(), path) } diff --git a/internal/report/meta.go b/internal/report/meta.go new file mode 100644 index 0000000..15b00a8 --- /dev/null +++ b/internal/report/meta.go @@ -0,0 +1,109 @@ +package report + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "sort" + "strconv" + "strings" + + "github.com/lukaszraczylo/llm-testbench/internal/runner" + "github.com/lukaszraczylo/llm-testbench/internal/testkit" +) + +// RunMeta records the settings and environment a run used, so two +// artifacts can be judged comparable and a result reproduced. +type RunMeta struct { + Fingerprints map[string][]string `json:"fingerprints,omitempty"` + Seed *int `json:"seed,omitempty"` + Version string `json:"version"` + StartedAt string `json:"started_at"` + Endpoint string `json:"endpoint,omitempty"` + CatalogHash string `json:"catalog_hash"` + Temperature float64 `json:"temperature"` + MaxTokensDefault int `json:"max_tokens_default"` + Repeat int `json:"repeat"` + TestCount int `json:"test_count"` + Interrupted bool `json:"interrupted,omitempty"` +} + +// CatalogHash returns a stable digest of the test inputs that determine a +// model's answer (id, system prompt, prompt, token floor, tool names). +func CatalogHash(tests []testkit.Test) string { + sorted := append([]testkit.Test(nil), tests...) + sort.Slice(sorted, func(i, j int) bool { return sorted[i].ID < sorted[j].ID }) + h := sha256.New() + for _, t := range sorted { + names := make([]string, 0, len(t.Tools)) + for _, tl := range t.Tools { + names = append(names, tl.Name) + } + _, _ = fmt.Fprintf(h, "%s\x00%s\x00%s\x00%d\x00%s\x01", t.ID, t.System, t.Prompt, t.MaxTokens, strings.Join(names, ",")) + } + return hex.EncodeToString(h.Sum(nil))[:16] +} + +// servedFingerprints collects the distinct backend fingerprints each model +// answered with; more than one means the serving build changed mid-run. +func servedFingerprints(results []runner.Result) map[string][]string { + seen := make(map[string]map[string]bool) + for _, r := range results { + if r.Fingerprint == "" { + continue + } + if seen[r.Model] == nil { + seen[r.Model] = make(map[string]bool) + } + seen[r.Model][r.Fingerprint] = true + } + if len(seen) == 0 { + return nil + } + out := make(map[string][]string, len(seen)) + for m, fps := range seen { + for fp := range fps { + out[m] = append(out[m], fp) + } + sort.Strings(out[m]) + } + return out +} + +// compareWarnings lists reasons two artifacts may not be like-for-like. +// Artifacts without metadata (older runs) produce no warnings. +func compareWarnings(base, cur *RunMeta) []string { + if base == nil || cur == nil { + return nil + } + var w []string + if base.CatalogHash != cur.CatalogHash { + w = append(w, fmt.Sprintf("catalog differs (%s vs %s): prompts changed between runs", base.CatalogHash, cur.CatalogHash)) + } + if base.Temperature != cur.Temperature { + w = append(w, fmt.Sprintf("temperature differs (%g vs %g)", base.Temperature, cur.Temperature)) + } + if base.MaxTokensDefault != cur.MaxTokensDefault { + w = append(w, fmt.Sprintf("max_tokens_default differs (%d vs %d)", base.MaxTokensDefault, cur.MaxTokensDefault)) + } + if fmtSeed(base.Seed) != fmtSeed(cur.Seed) { + w = append(w, fmt.Sprintf("seed differs (%s vs %s)", fmtSeed(base.Seed), fmtSeed(cur.Seed))) + } + if base.Interrupted || cur.Interrupted { + w = append(w, "an artifact is from an interrupted run: missing results are not regressions") + } + for m, cfps := range cur.Fingerprints { + if bfps, ok := base.Fingerprints[m]; ok && strings.Join(bfps, ",") != strings.Join(cfps, ",") { + w = append(w, fmt.Sprintf("%s served by a different backend build (%s vs %s)", m, strings.Join(bfps, "+"), strings.Join(cfps, "+"))) + } + } + sort.Strings(w) + return w +} + +func fmtSeed(s *int) string { + if s == nil { + return "unset" + } + return strconv.Itoa(*s) +} diff --git a/internal/report/meta_test.go b/internal/report/meta_test.go new file mode 100644 index 0000000..01929a9 --- /dev/null +++ b/internal/report/meta_test.go @@ -0,0 +1,71 @@ +package report + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/lukaszraczylo/llm-testbench/internal/runner" + "github.com/lukaszraczylo/llm-testbench/internal/testkit" +) + +func TestCatalogHash_StableAndSensitive(t *testing.T) { + a := testkit.Test{ID: "a", Prompt: "p1"} + b := testkit.Test{ID: "b", Prompt: "p2"} + if CatalogHash([]testkit.Test{a, b}) != CatalogHash([]testkit.Test{b, a}) { + t.Error("hash depends on input order, want stable") + } + b2 := b + b2.Prompt = "p2 changed" + if CatalogHash([]testkit.Test{a, b}) == CatalogHash([]testkit.Test{a, b2}) { + t.Error("hash ignores a prompt change") + } +} + +func TestCompareWarnings(t *testing.T) { + seed := 1 + base := &RunMeta{CatalogHash: "x", Temperature: 0, MaxTokensDefault: 100, Fingerprints: map[string][]string{"m": {"fp1"}}} + if w := compareWarnings(base, base); len(w) != 0 { + t.Errorf("identical meta warnings = %v, want none", w) + } + if w := compareWarnings(nil, base); len(w) != 0 { + t.Errorf("legacy artifact warnings = %v, want none", w) + } + cur := &RunMeta{CatalogHash: "y", Temperature: 0, MaxTokensDefault: 200, Seed: &seed, Interrupted: true, + Fingerprints: map[string][]string{"m": {"fp2"}}} + got := strings.Join(compareWarnings(base, cur), "\n") + for _, want := range []string{"catalog differs", "max_tokens_default differs", "seed differs", "interrupted", "different backend build"} { + if !strings.Contains(got, want) { + t.Errorf("warnings missing %q:\n%s", want, got) + } + } +} + +func TestWriteArtifact_AtomicWithMetaRoundTrip(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "run.json") + if err := os.WriteFile(path, []byte("previous"), 0o600); err != nil { + t.Fatal(err) + } + tests := []testkit.Test{{ID: "t1", Category: "c", Prompt: "p"}} + results := []runner.Result{{Model: "m", TestID: "t1", Fingerprint: "fp"}} + + if err := WriteArtifact(path, RunMeta{Version: "v1", Repeat: 1}, tests, []string{"m"}, results); err != nil { + t.Fatalf("WriteArtifact() error = %v", err) + } + art, err := LoadArtifact(path) + if err != nil { + t.Fatalf("LoadArtifact() error = %v", err) + } + if art.Meta == nil || art.Meta.Version != "v1" || art.Meta.TestCount != 1 || art.Meta.CatalogHash == "" { + t.Errorf("Meta = %+v, want version, test count and catalog hash set", art.Meta) + } + if got := art.Meta.Fingerprints["m"]; len(got) != 1 || got[0] != "fp" { + t.Errorf("Fingerprints = %v, want m:[fp]", art.Meta.Fingerprints) + } + entries, _ := os.ReadDir(dir) + if len(entries) != 1 { + t.Errorf("dir has %d entries after write, want only the artifact (no temp leftovers)", len(entries)) + } +} diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 0b1441b..54c0533 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -5,6 +5,7 @@ package runner import ( "context" "encoding/json" + "fmt" "sync/atomic" "time" @@ -59,6 +60,8 @@ type Result struct { TestID string ResponseText string FinishReason string + Fingerprint string + ServedModel string Score eval.Score Latency time.Duration PromptTokens int @@ -83,7 +86,13 @@ func (r Result) TotalTokens() int { // Config controls how the Runner fans work out. type Config struct { - Reporter ProgressReporter + Reporter ProgressReporter + // Seed, when non-nil, is sent with every request for backends that + // honor it. + Seed *int + // ModelConcurrency caps in-flight calls per model (on top of the global + // Concurrency), so one slow local model is not flooded into timeouts. + ModelConcurrency map[string]int Concurrency int Temperature float64 MaxTokensDefault int @@ -138,15 +147,32 @@ func (r *Runner) Run(ctx context.Context, models []string, tests []testkit.Test) if concurrency <= 0 { concurrency = 1 } - g.SetLimit(concurrency) + // Jobs wait for their per-model slot first, then a global slot, so a + // capped slow model's queue never occupies global slots other models + // need. + global := make(chan struct{}, concurrency) r.cfg.Reporter.ReportStart(len(tests), len(models), concurrency) var completed int32 // go.mod's go directive is 1.22+, so loop variables are already // per-iteration scoped; no manual i, j := i, j shadowing needed (N4). + sems := make(map[string]chan struct{}, len(r.cfg.ModelConcurrency)) + for m, n := range r.cfg.ModelConcurrency { + if n > 0 { + sems[m] = make(chan struct{}, n) + } + } for i, j := range jobs { g.Go(func() error { + if sem, ok := sems[j.model]; ok { + if acquire(gctx, sem) { + defer func() { <-sem }() + } + } + if acquire(gctx, global) { + defer func() { <-global }() + } res := r.runOne(gctx, j.model, j.test) res.Attempt = j.attempt results[i] = res @@ -167,8 +193,32 @@ func (r *Runner) Run(ctx context.Context, models []string, tests []testkit.Test) return results } -// runOne executes a single (model, test) combination. -func (r *Runner) runOne(ctx context.Context, model string, test testkit.Test) Result { +// acquire takes a slot from sem, reporting false if ctx ended first. +func acquire(ctx context.Context, sem chan struct{}) bool { + select { + case sem <- struct{}{}: + return true + case <-ctx.Done(): + return false + } +} + +// runOne executes a single (model, test) combination. A cancelled context +// short-circuits to an error result, and a panic in an evaluator becomes an +// error result instead of killing a multi-hour run. +func (r *Runner) runOne(ctx context.Context, model string, test testkit.Test) (res Result) { + if err := ctx.Err(); err != nil { + return Result{Model: model, TestID: test.ID, Err: err} + } + defer func() { + if p := recover(); p != nil { + res = Result{Model: model, TestID: test.ID, Err: fmt.Errorf("runner: panic scoring test: %v", p)} + } + }() + return r.execute(ctx, model, test) +} + +func (r *Runner) execute(ctx context.Context, model string, test testkit.Test) Result { // cfg.MaxTokensDefault is a floor, not a fallback used only when the // test leaves MaxTokens unset: a per-test value only ever raises the // budget, it never lowers it below the default. A reasoning model can @@ -190,6 +240,7 @@ func (r *Runner) runOne(ctx context.Context, model string, test testkit.Test) Re Tools: test.Tools, MaxTokens: maxTokens, Temperature: r.cfg.Temperature, + Seed: r.cfg.Seed, } // Wall-clock across all attempts, recorded even on error: an error row @@ -223,6 +274,8 @@ func (r *Runner) runOne(ctx context.Context, model string, test testkit.Test) Re Score: score, ResponseText: scored, FinishReason: resp.FinishReason, + Fingerprint: resp.Fingerprint, + ServedModel: resp.ServedModel, Latency: resp.Latency, PromptTokens: resp.PromptTokens, CompletionTokens: resp.CompletionTokens, diff --git a/internal/runner/runner_test.go b/internal/runner/runner_test.go index 49743f2..7e26208 100644 --- a/internal/runner/runner_test.go +++ b/internal/runner/runner_test.go @@ -410,3 +410,89 @@ func TestRunner_Run_RepeatSamplesEachCombination(t *testing.T) { } } } + +func TestRunner_Run_PassesSeedAndRecordsFingerprint(t *testing.T) { + var gotSeed *int + client := &mockClient{respond: func(req llm.Request) (llm.Response, error) { + gotSeed = req.Seed + return llm.Response{Text: "OK", Fingerprint: "fp-1", ServedModel: "served"}, nil + }} + seed := 7 + r := New(client, Config{Concurrency: 1, MaxTokensDefault: 100, Seed: &seed}) + res := r.Run(context.Background(), []string{"m"}, []testkit.Test{{ID: "t", Category: "c", Prompt: "p", Eval: echoEval()}}) + if gotSeed == nil || *gotSeed != 7 { + t.Errorf("request seed = %v, want 7", gotSeed) + } + if res[0].Fingerprint != "fp-1" || res[0].ServedModel != "served" { + t.Errorf("fingerprint/model = %q/%q, want fp-1/served", res[0].Fingerprint, res[0].ServedModel) + } +} + +func TestRunner_Run_EvaluatorPanicBecomesErrorResult(t *testing.T) { + boom := eval.EvaluatorFunc(func(context.Context, string) eval.Score { panic("boom") }) + r := New(&mockClient{}, Config{Concurrency: 2, MaxTokensDefault: 100}) + res := r.Run(context.Background(), []string{"m"}, []testkit.Test{ + {ID: "bad", Category: "c", Prompt: "p", Eval: boom}, + {ID: "good", Category: "c", Prompt: "p", Eval: echoEval()}, + }) + if res[0].Err == nil || !strings.Contains(res[0].Err.Error(), "boom") { + t.Errorf("panicking test Err = %v, want panic error", res[0].Err) + } + if res[1].Err != nil || res[1].Score.Value != 1 { + t.Errorf("sibling test = %+v, want unaffected pass", res[1]) + } +} + +func TestRunner_Run_CancelledContextSkipsCalls(t *testing.T) { + var calls int32 + client := &mockClient{respond: func(llm.Request) (llm.Response, error) { + atomic.AddInt32(&calls, 1) + return llm.Response{Text: "OK"}, nil + }} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + r := New(client, Config{Concurrency: 2, MaxTokensDefault: 100}) + res := r.Run(ctx, []string{"m"}, []testkit.Test{{ID: "t", Category: "c", Prompt: "p", Eval: echoEval()}}) + if atomic.LoadInt32(&calls) != 0 { + t.Errorf("client called %d times after cancel, want 0", calls) + } + if !errors.Is(res[0].Err, context.Canceled) { + t.Errorf("Err = %v, want context.Canceled", res[0].Err) + } +} + +func TestRunner_Run_ModelConcurrencyCapDoesNotStarveOthers(t *testing.T) { + var slowIn, slowMax int32 + fastDone := make(chan struct{}, 8) + client := &mockClient{respond: func(req llm.Request) (llm.Response, error) { + if req.Model == "slow" { + cur := atomic.AddInt32(&slowIn, 1) + for { + m := atomic.LoadInt32(&slowMax) + if cur <= m || atomic.CompareAndSwapInt32(&slowMax, m, cur) { + break + } + } + time.Sleep(30 * time.Millisecond) + atomic.AddInt32(&slowIn, -1) + } else { + fastDone <- struct{}{} + } + return llm.Response{Text: "OK"}, nil + }} + r := New(client, Config{Concurrency: 2, MaxTokensDefault: 100, ModelConcurrency: map[string]int{"slow": 1}}) + tests := make([]testkit.Test, 6) + for i := range tests { + tests[i] = testkit.Test{ID: string(rune('a' + i)), Category: "c", Prompt: "p", Eval: echoEval()} + } + res := r.Run(context.Background(), []string{"slow", "fast"}, tests) + if len(res) != 12 { + t.Fatalf("results = %d, want 12", len(res)) + } + if got := atomic.LoadInt32(&slowMax); got != 1 { + t.Errorf("slow model max in-flight = %d, want 1", got) + } + if len(fastDone) != 6 { + t.Errorf("fast calls = %d, want 6", len(fastDone)) + } +}