10 Commits

Author SHA1 Message Date
steve 1206261e6a feat(failover): package-level default observer for transparently-built chains
CI / Build, Test & Lint (push) Successful in 10m43s
The transparent comma-Parse path builds failover chains via NewFailoverModel
with no options, so defaultFailoverConfig() left the observer nil and observers
only fired when a caller passed WithFailoverObserver explicitly. Add a
package-level default observer (SetFailoverObserver / DefaultFailoverObserver),
guarded by the existing defaultsMu, and seed it in defaultFailoverConfig() so
chains built transparently still notify it. An explicit WithFailoverObserver
still overrides the default per-chain. mort sets this at boot to persist
failover events.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-01 00:43:24 +02:00
steve 361999550e fix(failover): preserve manual bench against automatic cooldown downgrade
CI / Build, Test & Lint (push) Successful in 10m48s
recordTransientFailure and benchNow unconditionally set manual=false and
reset until to now+cooldown. When the best-effort all-benched failover path
re-tries a model an operator manually benched via BenchModel, a subsequent
failure downgraded manual=true -> false and could shorten the operator's
window to the short auto cooldown.

Both functions now read existing state: if it is an active manual bench
(manual && now.Before(until)) they bump consecutiveFails but keep manual=true
and the later until. Non-manual or expired-manual state still gets the
automatic cooldown. Adds TestFailover_ManualBenchSurvivesAutomaticDowngrade
covering no-prior, prior-auto, active-manual, and expired-manual cases.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-01 00:37:38 +02:00
steve ae8e194fad feat(failover): model failover chains via comma-separated specs
Parse("a,b,c") now returns one composite *llm.Model that tries each model
in order, retrying transient failures, benching dead models, and failing
over to the next. Comma-free specs are completely unchanged.

- classify.go: Classify(err) ErrKind + IsTransient(err) error classifier
  mapping anthropic (typed Is*Err helpers + RequestError status),
  openai-go (*openai.Error status), openaicompat.FeatureUnsupportedError,
  context errors, and ollama "HTTP <code>" strings to
  transient/auth-dead/request-specific/unknown.
- failover.go: failoverProvider (satisfies provider.Provider) wrapped into a
  *Model via NewClient. Process-wide mutex-guarded modelHealth bench
  registry keyed by concrete spec, with cooldowns and a control API
  (ListBenched/BenchModel/UnbenchModel/IsBenched). NewFailoverModel +
  ParseChain constructors, FailoverOption config, FailoverObserver (carries
  the full request), and configurable package-level defaults.
- parse.go: comma-aware Parse splits into a failover chain; alias/resolver
  targets that expand to comma chains are routed through the comma-aware
  path and flattened.

All access to global health is mutex-guarded; tests reset it via
resetHealthForTest and pass under go test -race.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-01 00:30:08 +02:00
steve 67c3ebe067 feat(ollama): add automatic retry with exponential backoff for transient HTTP errors
CI / Build, Test & Lint (push) Successful in 10m50s
Ollama Cloud returns HTTP 503 when the model is temporarily overloaded,
429 on rate limit, and 502 on upstream failures. These are transient
conditions that resolve on retry. Previously they bubbled up as hard
errors, forcing callers to implement their own retry logic.

The retry is implemented at the HTTP transport level in doChatRequest,
so both Complete and Stream benefit transparently. Strategy: up to 3
retries with exponential backoff (1s, 2s, 4s), Retry-After header
respected for 429, context cancellation checked between retries.
Non-transient errors (400, 401, 403, 404, 500) are never retried.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-25 11:58:25 -04:00
steve 6bac4cb3ed Merge pull request 'feat(v2): add Parse() and extensible Registry' (#5) from feat/parse-registry into main
CI / Build, Test & Lint (push) Successful in 11m58s
2026-05-24 15:14:45 +00:00
steve a1f405de74 Merge branch 'main' into feat/parse-registry
CI / Build, Test & Lint (pull_request) Successful in 12m0s
2026-05-24 15:14:17 +00:00
steve 72e281f82b ci: optimize workflow by merging jobs and adding module cache
CI / Build, Test & Lint (pull_request) Successful in 11m33s
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-05-23 23:28:20 -04:00
steve db03e8d4ed fix(ci): write correct YAML workflow (was double-encoded as base64)
CI / Lint (pull_request) Successful in 9s
CI / V2 Module (pull_request) Successful in 1m50s
The previous CI fix commits stored the workflow file as a literal base64
string instead of decoded YAML, so Gitea could not parse it as a valid
workflow and no CI runs triggered. This writes the correct YAML content.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-23 23:24:19 -04:00
steve ddf1a9b556 ci: fix workflow to use v2/ module directory 2026-05-24 03:02:55 +00:00
steve 4522310f5a feat(v2): add Parse() function and extensible Registry for model string resolution
CI / Root Module (pull_request) Failing after 3s
CI / Lint (pull_request) Failing after 3s
CI / V2 Module (pull_request) Successful in 1m25s
Introduces llm.Parse(spec) backed by an extensible Registry that resolves
model strings like "openai/gpt-4o", aliases like "fast", and named targets
like "m5/qwen3:30b" (via LLM_M5 env var DSNs) into ready-to-use *Model
objects. Extension points: RegisterProvider, RegisterAlias, RegisterResolver.
Adds Foreman constructor and sentinel errors ErrAliasLoop, ErrUnknownProvider,
ErrInvalidDSN.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-23 22:58:14 -04:00
12 changed files with 3084 additions and 22 deletions
+12
View File
@@ -32,3 +32,15 @@
6. Streaming via pull-based `StreamReader.Next()`
7. Middleware for logging, retry, timeout, usage tracking
8. Ollama uses the native `/api/chat` API rather than the OpenAI-compat `/v1` endpoint. Native API supports `think: false` for thinking-capable models, has more reliable tool calling, and is approximately 15-20% lower latency. Both local and cloud share the same provider; only the apiKey/baseURL differ. `llm.Ollama()` targets `http://localhost:11434` with no Authorization header; `llm.OllamaCloud(key)` targets `https://ollama.com` with `Authorization: Bearer <key>`.
### DD#9 — Parse() function and extensible Registry (2026-05-23)
**Context:** mort's ParseModelRequest resolves "provider/model" strings but is
mort-specific. Multi-instance providers (foreman) need named targets.
**Decision:** Add `llm.Parse(spec)` backed by an extensible `Registry`. Supports
aliases, dynamic resolvers, and `LLM_X` env var DSNs for named targets.
Provider/model syntax: `"openai/gpt-4o"`, aliases: `"fast"`, named targets:
`"m5/qwen3:30b"` (reads `LLM_M5` env var). Registry is extensible via
`RegisterProvider`, `RegisterAlias`, `RegisterResolver`.
**Consequence:** Any go-llm consumer gets model-string parsing. mort migrates by
registering its tier aliases as resolvers. Foreman instances are addressed via
`LLM_X` DSN env vars without code changes.
+207
View File
@@ -0,0 +1,207 @@
package llm
import (
"context"
"errors"
"strings"
anth "github.com/liushuangls/go-anthropic/v2"
"github.com/openai/openai-go"
"gitea.stevedudenhoeffer.com/steve/go-llm/v2/openaicompat"
)
// ErrKind classifies a provider error for failover decision-making.
//
// Why: failover must decide, per error, whether to retry the same model
// (transient), bench it as broken (auth/model dead), or fail over without
// benching (this request's fault). Without a classifier every error looks
// the same and we'd either thrash a dead model or bench a healthy one.
// What: an enum of the four outcomes the failover algorithm distinguishes.
// Test: see classify_test.go — every branch is table-tested with faked SDK errors.
type ErrKind int
const (
// ErrUnknown is an unrecognized error. Failover treats it as transient
// (conservative — retry then fail over), EXCEPT context.Canceled which
// the caller special-cases as an abort.
ErrUnknown ErrKind = iota
// ErrTransient is a temporary failure (429/5xx/timeout): retry, then
// bench-and-fail-over if retries are exhausted.
ErrTransient
// ErrAuthDead is an auth failure or model-not-found (401/403/404): the
// model is unusable; bench immediately and fail over.
ErrAuthDead
// ErrRequestSpecific is the caller's fault for THIS request (400/413/422,
// unsupported feature): fail over to try a more capable model, but do NOT
// bench — the model itself is healthy.
ErrRequestSpecific
)
// classifyStatus maps an HTTP status code to an ErrKind.
//
// Why: openai-go and anthropic RequestError both expose a numeric StatusCode;
// centralizing the mapping keeps the per-SDK branches thin and consistent.
// What: 408/409/429/5xx → transient, 401/403/404 → auth-dead, 400/413/422 →
// request-specific, anything else → unknown.
// Test: covered indirectly via Classify table tests for each SDK.
func classifyStatus(code int) ErrKind {
switch code {
case 408, 409, 429, 500, 502, 503, 504:
return ErrTransient
case 401, 403, 404:
return ErrAuthDead
case 400, 413, 422:
return ErrRequestSpecific
default:
return ErrUnknown
}
}
// Classify inspects a provider error and returns its ErrKind.
//
// Why: the failover composite needs typed, status-code-aware classification to
// retry/bench/skip correctly across the anthropic, openai-compat, and ollama
// providers, each of which surfaces errors differently.
// What: prefers anthropic's typed Is*Err helpers, falls back to numeric status
// codes (openai-go, anthropic RequestError), then the openaicompat
// FeatureUnsupportedError, context errors, and finally an ollama HTTP-string
// fallback; unrecognized errors are ErrUnknown.
// Test: classify_test.go faked SDK errors exercise every branch.
func Classify(err error) ErrKind {
if err == nil {
return ErrUnknown
}
// context.Canceled is reported as ErrUnknown here; the failover algorithm
// special-cases it as an abort before consulting the kind.
if errors.Is(err, context.Canceled) {
return ErrUnknown
}
if errors.Is(err, context.DeadlineExceeded) {
return ErrTransient
}
// FeatureUnsupportedError is a permanent, request-shaped failure.
var featErr *openaicompat.FeatureUnsupportedError
if errors.As(err, &featErr) {
return ErrRequestSpecific
}
// Anthropic APIError: prefer the typed helpers (no StatusCode available).
var apiErr *anth.APIError
if errors.As(err, &apiErr) {
switch {
case apiErr.IsRateLimitErr(), apiErr.IsOverloadedErr(), apiErr.IsApiErr():
return ErrTransient
case apiErr.IsAuthenticationErr(), apiErr.IsPermissionErr(), apiErr.IsNotFoundErr():
return ErrAuthDead
case apiErr.IsTooLargeErr(), apiErr.IsInvalidRequestErr():
return ErrRequestSpecific
default:
return ErrUnknown
}
}
// Anthropic RequestError: status-code based.
var anthReqErr *anth.RequestError
if errors.As(err, &anthReqErr) {
return classifyStatus(anthReqErr.StatusCode)
}
// openai-go (openai/deepseek/moonshot/xai/groq): status-code based.
var oaiErr *openai.Error
if errors.As(err, &oaiErr) {
return classifyStatus(oaiErr.StatusCode)
}
// Ollama: no typed status — fall back to its "ollama: HTTP <code>:" string.
if k := classifyOllamaString(err.Error()); k != ErrUnknown {
return k
}
return ErrUnknown
}
// classifyOllamaString extracts an HTTP status from ollama's error string
// format ("ollama: HTTP <code>: ...") and classifies it.
//
// Why: the ollama provider stringifies errors without a typed status code, so
// failover can only classify by parsing the message.
// What: looks for "HTTP <code>" in the message and maps the code; returns
// ErrUnknown when no recognizable status is present.
// Test: classify_test.go ollama cases cover 5xx/429/401/404/400.
func classifyOllamaString(msg string) ErrKind {
const marker = "HTTP "
idx := strings.Index(msg, marker)
if idx < 0 {
return ErrUnknown
}
rest := msg[idx+len(marker):]
// Read up to 3 leading digits.
end := 0
for end < len(rest) && end < 3 && rest[end] >= '0' && rest[end] <= '9' {
end++
}
if end == 0 {
return ErrUnknown
}
code := 0
for i := 0; i < end; i++ {
code = code*10 + int(rest[i]-'0')
}
return classifyStatus(code)
}
// extractStatus best-effort pulls an HTTP status code out of a provider error
// for structured logging. Returns 0 when none is available.
//
// Why: log lines benefit from the numeric status even though classification
// may use typed helpers; this keeps that detail out of the hot path.
// What: checks anthropic RequestError and openai-go Error StatusCode fields,
// then parses ollama's "HTTP <code>" string; returns 0 otherwise.
// Test: covered indirectly via failover log assertions / manual inspection.
func extractStatus(err error) int {
if err == nil {
return 0
}
var anthReqErr *anth.RequestError
if errors.As(err, &anthReqErr) {
return anthReqErr.StatusCode
}
var oaiErr *openai.Error
if errors.As(err, &oaiErr) {
return oaiErr.StatusCode
}
const marker = "HTTP "
msg := err.Error()
if idx := strings.Index(msg, marker); idx >= 0 {
rest := msg[idx+len(marker):]
code, n := 0, 0
for n < len(rest) && n < 3 && rest[n] >= '0' && rest[n] <= '9' {
code = code*10 + int(rest[n]-'0')
n++
}
if n > 0 {
return code
}
}
return 0
}
// IsTransient reports whether an error should be retried/failed-over rather
// than treated as a hard, model-specific failure.
//
// Why: callers (and failover) want a one-call "is this worth retrying?" check
// that is conservative about unknown errors.
// What: returns true for ErrTransient and ErrUnknown (conservative), false for
// ErrAuthDead and ErrRequestSpecific.
// Test: TestIsTransient asserts 503→true, unknown→true, 401→false, 400→false.
func IsTransient(err error) bool {
switch Classify(err) {
case ErrTransient, ErrUnknown:
return true
default:
return false
}
}
+115
View File
@@ -0,0 +1,115 @@
package llm
import (
"context"
"errors"
"fmt"
"testing"
anth "github.com/liushuangls/go-anthropic/v2"
"github.com/openai/openai-go"
"gitea.stevedudenhoeffer.com/steve/go-llm/v2/openaicompat"
)
func TestClassify(t *testing.T) {
tests := []struct {
name string
err error
want ErrKind
}{
// nil
{"nil", nil, ErrUnknown},
// openai-go status codes (transient)
{"openai 408", &openai.Error{StatusCode: 408}, ErrTransient},
{"openai 409", &openai.Error{StatusCode: 409}, ErrTransient},
{"openai 429", &openai.Error{StatusCode: 429}, ErrTransient},
{"openai 500", &openai.Error{StatusCode: 500}, ErrTransient},
{"openai 502", &openai.Error{StatusCode: 502}, ErrTransient},
{"openai 503", &openai.Error{StatusCode: 503}, ErrTransient},
{"openai 504", &openai.Error{StatusCode: 504}, ErrTransient},
// openai-go status codes (auth dead)
{"openai 401", &openai.Error{StatusCode: 401}, ErrAuthDead},
{"openai 403", &openai.Error{StatusCode: 403}, ErrAuthDead},
{"openai 404", &openai.Error{StatusCode: 404}, ErrAuthDead},
// openai-go status codes (request specific)
{"openai 400", &openai.Error{StatusCode: 400}, ErrRequestSpecific},
{"openai 413", &openai.Error{StatusCode: 413}, ErrRequestSpecific},
{"openai 422", &openai.Error{StatusCode: 422}, ErrRequestSpecific},
// openai unrecognized status -> unknown
{"openai 418", &openai.Error{StatusCode: 418}, ErrUnknown},
// wrapped openai error (providers wrap with %w)
{"wrapped openai 503", fmt.Errorf("openai completion error: %w", &openai.Error{StatusCode: 503}), ErrTransient},
// FeatureUnsupportedError -> request specific
{"feature unsupported", &openaicompat.FeatureUnsupportedError{Feature: "tools", Model: "m"}, ErrRequestSpecific},
{"wrapped feature unsupported", fmt.Errorf("x: %w", &openaicompat.FeatureUnsupportedError{Feature: "vision", Model: "m"}), ErrRequestSpecific},
// anthropic RequestError (status-code based)
{"anth req 503", &anth.RequestError{StatusCode: 503}, ErrTransient},
{"anth req 429", &anth.RequestError{StatusCode: 429}, ErrTransient},
{"anth req 401", &anth.RequestError{StatusCode: 401}, ErrAuthDead},
{"anth req 400", &anth.RequestError{StatusCode: 400}, ErrRequestSpecific},
{"wrapped anth req 502", fmt.Errorf("anthropic completion error: %w", &anth.RequestError{StatusCode: 502}), ErrTransient},
// anthropic APIError (helper based)
{"anth rate limit", &anth.APIError{Type: anth.ErrTypeRateLimit}, ErrTransient},
{"anth overloaded", &anth.APIError{Type: anth.ErrTypeOverloaded}, ErrTransient},
{"anth api", &anth.APIError{Type: anth.ErrTypeApi}, ErrTransient},
{"anth auth", &anth.APIError{Type: anth.ErrTypeAuthentication}, ErrAuthDead},
{"anth permission", &anth.APIError{Type: anth.ErrTypePermission}, ErrAuthDead},
{"anth not found", &anth.APIError{Type: anth.ErrTypeNotFound}, ErrAuthDead},
{"anth too large", &anth.APIError{Type: anth.ErrTypeTooLarge}, ErrRequestSpecific},
{"anth invalid request", &anth.APIError{Type: anth.ErrTypeInvalidRequest}, ErrRequestSpecific},
{"wrapped anth api error", fmt.Errorf("error, status code: 529, message: %w", &anth.APIError{Type: anth.ErrTypeOverloaded}), ErrTransient},
// context errors
{"context canceled", context.Canceled, ErrUnknown},
{"context deadline", context.DeadlineExceeded, ErrTransient},
{"wrapped deadline", fmt.Errorf("call failed: %w", context.DeadlineExceeded), ErrTransient},
// ollama string-based
{"ollama HTTP 503", errors.New("ollama: HTTP 503: service unavailable"), ErrTransient},
{"ollama HTTP 500", errors.New("ollama: HTTP 500: internal"), ErrTransient},
{"ollama HTTP 502", errors.New("ollama: HTTP 502: bad gateway"), ErrTransient},
{"ollama HTTP 504", errors.New("ollama: HTTP 504: timeout"), ErrTransient},
{"ollama HTTP 429", errors.New("ollama: HTTP 429: too many requests"), ErrTransient},
{"ollama HTTP 401", errors.New("ollama: HTTP 401: unauthorized"), ErrAuthDead},
{"ollama HTTP 404", errors.New("ollama: HTTP 404: not found"), ErrAuthDead},
{"ollama HTTP 400", errors.New("ollama: HTTP 400: bad request"), ErrRequestSpecific},
// unknown
{"random error", errors.New("something weird"), ErrUnknown},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := Classify(tt.err)
if got != tt.want {
t.Errorf("Classify(%v) = %v, want %v", tt.err, got, tt.want)
}
})
}
}
func TestIsTransient(t *testing.T) {
// IsTransient treats both ErrTransient and ErrUnknown as "should retry"
// (conservative). Auth/request-specific are not transient.
if !IsTransient(&openai.Error{StatusCode: 503}) {
t.Error("503 should be transient")
}
if !IsTransient(errors.New("mystery")) {
t.Error("unknown should be treated as transient (conservative)")
}
if IsTransient(&openai.Error{StatusCode: 401}) {
t.Error("401 (auth dead) should NOT be transient")
}
if IsTransient(&openai.Error{StatusCode: 400}) {
t.Error("400 (request specific) should NOT be transient")
}
}
+16
View File
@@ -119,6 +119,22 @@ func Ollama(opts ...ClientOption) *Client {
return NewClient(ollamaProvider.New("", cfg.baseURL))
}
// Foreman creates a client targeting a foreman daemon (a private, authenticated
// Ollama endpoint with queuing and observability). The token is sent as a Bearer
// token; pass "" for unauthenticated (network-trusted) deployments. Use
// WithBaseURL to set the foreman host URL.
//
// Example:
//
// model := llm.Foreman("my-token", llm.WithBaseURL("https://foreman.local")).Model("qwen3:30b")
func Foreman(token string, opts ...ClientOption) *Client {
cfg := &clientConfig{}
for _, opt := range opts {
opt(cfg)
}
return NewClient(ollamaProvider.New(token, cfg.baseURL))
}
// OllamaCloud creates a client targeting Ollama Cloud (https://ollama.com).
// The apiKey is required and is sent as `Authorization: Bearer <key>`. Use
// WithBaseURL to point at a private Ollama deployment that requires auth.
+12
View File
@@ -17,4 +17,16 @@ var (
// ErrNoStructuredOutput is returned when the model did not return a structured output tool call.
ErrNoStructuredOutput = errors.New("model did not return structured output")
// ErrAliasLoop is returned when alias resolution exceeds the maximum depth (10),
// indicating a cycle such as "a" → "b" → "a".
ErrAliasLoop = errors.New("alias resolution loop detected (depth > 10)")
// ErrUnknownProvider is returned when a spec references a provider name that
// is not registered and has no corresponding LLM_X environment variable.
ErrUnknownProvider = errors.New("unknown provider")
// ErrInvalidDSN is returned when a DSN string (from an LLM_X env var) cannot
// be parsed. Expected format: scheme://[token@]host[/path].
ErrInvalidDSN = errors.New("invalid DSN")
)
+638
View File
@@ -0,0 +1,638 @@
package llm
import (
"context"
"errors"
"fmt"
"log/slog"
"math/rand"
"strings"
"sync"
"time"
"gitea.stevedudenhoeffer.com/steve/go-llm/v2/provider"
)
// ---------------------------------------------------------------------------
// Package-level defaults (mort configures these at boot via SetFailoverDefaults)
// ---------------------------------------------------------------------------
var (
// DefaultFailoverMaxRetries is the number of attempts per chain entry on
// transient errors before benching and moving to the next entry.
DefaultFailoverMaxRetries = 3
// DefaultFailoverCooldown is how long a model stays benched after a
// qualifying failure.
DefaultFailoverCooldown = 5 * time.Minute
// DefaultFailoverBackoff is the default exponential-with-jitter backoff.
DefaultFailoverBackoff = defaultBackoff
// defaultFailoverObserver is the package-level observer applied to chains
// built without an explicit WithFailoverObserver (e.g. the transparent
// comma-Parse path). Kept unexported and behind defaultsMu so reads/writes
// are race-safe under -race. mort sets this at boot to persist failover events.
defaultFailoverObserver FailoverObserver
defaultsMu sync.Mutex
)
// defaultBackoff returns an exponential backoff with full jitter.
//
// Why: spreads retries to avoid thundering-herd against a recovering provider.
// What: base 200ms doubling per attempt, capped at 10s, with uniform jitter.
// Test: failover retry tests inject a fast backoff; this is the production default.
func defaultBackoff(attempt int) time.Duration {
if attempt < 1 {
attempt = 1
}
base := 200 * time.Millisecond
d := base << (attempt - 1)
if d > 10*time.Second {
d = 10 * time.Second
}
// Full jitter in [0, d].
return time.Duration(rand.Int63n(int64(d) + 1))
}
// SetFailoverDefaults overrides the package-level failover defaults used when
// no per-model options are supplied (e.g. comma-spec Parse).
//
// Why: mort wants to tune retries/cooldown once at boot without threading
// options through every Parse call.
// What: sets DefaultFailoverMaxRetries and DefaultFailoverCooldown under a lock.
// Test: set defaults, build a comma model, assert its cfg reflects them.
func SetFailoverDefaults(maxRetries int, cooldown time.Duration) {
defaultsMu.Lock()
defer defaultsMu.Unlock()
DefaultFailoverMaxRetries = maxRetries
DefaultFailoverCooldown = cooldown
}
// SetFailoverObserver sets the package-level default observer notified on
// failover decisions for chains built without an explicit WithFailoverObserver.
//
// Why: the transparent comma-Parse path builds chains via NewFailoverModel with
// no options, so without a package default no observer ever fires; mort sets
// this once at boot to persist failover events from every chain.
// What: stores the observer under defaultsMu; pass nil to disable.
// Test: set an observer, build a no-option chain, assert it fires on failover.
func SetFailoverObserver(obs FailoverObserver) {
defaultsMu.Lock()
defer defaultsMu.Unlock()
defaultFailoverObserver = obs
}
// DefaultFailoverObserver returns the current package-level default observer.
//
// Why: lets tests assert/restore the default without reaching into the unexported var.
// What: reads defaultFailoverObserver under defaultsMu.
// Test: set via SetFailoverObserver, assert this returns a non-nil func.
func DefaultFailoverObserver() FailoverObserver {
defaultsMu.Lock()
defer defaultsMu.Unlock()
return defaultFailoverObserver
}
// ---------------------------------------------------------------------------
// Global model health (process-wide bench registry)
// ---------------------------------------------------------------------------
// modelHealth tracks which concrete models are temporarily disabled (benched).
//
// Why: bench decisions must persist across requests and across all failover
// chains in the process, so a model that's down isn't retried by every chain.
// What: a mutex-guarded map keyed by specKey to its disabled state.
// Test: failover tests reset it via resetHealthForTest and assert via IsBenched.
type modelHealth struct {
mu sync.Mutex
disabled map[string]disabledState
}
type disabledState struct {
until time.Time
consecutiveFails int
manual bool
}
// globalHealth is the process-wide singleton shared by every failover chain.
var globalHealth = &modelHealth{disabled: map[string]disabledState{}}
// benchThreshold is the number of consecutive transient failures (each after
// exhausting retries) required before a model is benched. Auth-dead benches
// immediately regardless.
const benchThreshold = 1
// resetHealthForTest clears all bench state. Test-only.
func resetHealthForTest() {
globalHealth.mu.Lock()
defer globalHealth.mu.Unlock()
globalHealth.disabled = map[string]disabledState{}
}
// isBenched reports whether key is currently benched (and not expired).
func (h *modelHealth) isBenched(key string, now time.Time) bool {
h.mu.Lock()
defer h.mu.Unlock()
st, ok := h.disabled[key]
if !ok {
return false
}
if now.After(st.until) {
delete(h.disabled, key)
return false
}
return true
}
// recordSuccess clears any failure state for key.
func (h *modelHealth) recordSuccess(key string) {
h.mu.Lock()
defer h.mu.Unlock()
delete(h.disabled, key)
}
// recordTransientFailure increments the consecutive failure count and benches
// the model once the threshold is reached. Returns whether it is now benched
// and for how long.
func (h *modelHealth) recordTransientFailure(key string, cooldown time.Duration, now time.Time) (benched bool, until time.Time) {
h.mu.Lock()
defer h.mu.Unlock()
st := h.disabled[key]
// Preserve an active manual bench: automatic logic must not clear the
// operator's manual flag or shorten their window. Still count the failure.
manualActive := st.manual && now.Before(st.until)
st.consecutiveFails++
if manualActive {
h.disabled[key] = st
return true, st.until
}
if st.consecutiveFails >= benchThreshold {
st.until = now.Add(cooldown)
st.manual = false
h.disabled[key] = st
return true, st.until
}
h.disabled[key] = st
return false, time.Time{}
}
// benchNow benches a model immediately (used for auth-dead errors).
func (h *modelHealth) benchNow(key string, cooldown time.Duration, now time.Time) time.Time {
h.mu.Lock()
defer h.mu.Unlock()
st := h.disabled[key]
// Preserve an active manual bench: automatic logic must not clear the
// operator's manual flag or shorten their window. Still count the failure.
manualActive := st.manual && now.Before(st.until)
st.consecutiveFails++
if manualActive {
h.disabled[key] = st
return st.until
}
st.until = now.Add(cooldown)
st.manual = false
h.disabled[key] = st
return st.until
}
// benchManual benches a model until the given time, marking it manual.
func (h *modelHealth) benchManual(key string, until time.Time) {
h.mu.Lock()
defer h.mu.Unlock()
st := h.disabled[key]
st.until = until
st.manual = true
h.disabled[key] = st
}
// unbench removes a model's bench state, reporting whether it was benched.
func (h *modelHealth) unbench(key string, now time.Time) bool {
h.mu.Lock()
defer h.mu.Unlock()
st, ok := h.disabled[key]
if !ok || now.After(st.until) {
delete(h.disabled, key)
return false
}
delete(h.disabled, key)
return true
}
// list returns a snapshot of all currently-benched (non-expired) models.
func (h *modelHealth) list(now time.Time) []BenchedModel {
h.mu.Lock()
defer h.mu.Unlock()
var out []BenchedModel
for k, st := range h.disabled {
if now.After(st.until) {
delete(h.disabled, k)
continue
}
out = append(out, BenchedModel{
Model: k,
Until: st.until,
ConsecutiveFails: st.consecutiveFails,
Manual: st.manual,
})
}
return out
}
// ---------------------------------------------------------------------------
// Control API (admin commands / UI drive these)
// ---------------------------------------------------------------------------
// BenchedModel is a snapshot of a benched model's state.
type BenchedModel struct {
Model string
Until time.Time
ConsecutiveFails int
Manual bool
}
// ListBenched returns all currently-benched models across the process.
//
// Why: admin tooling needs to display which models are sidelined and why.
// What: snapshots the global health map, pruning expired entries.
// Test: BenchModel then ListBenched returns it with Manual=true.
func ListBenched() []BenchedModel { return globalHealth.list(time.Now()) }
// BenchModel manually benches a model until the given time.
//
// Why: operators sometimes need to force a model offline (incident, cost).
// What: records a manual bench in the global health registry.
// Test: BenchModel then IsBenched returns true and ListBenched shows Manual.
func BenchModel(spec string, until time.Time) { globalHealth.benchManual(spec, until) }
// UnbenchModel clears a model's bench state, returning whether it was benched.
//
// Why: operators need to bring a model back early after manual or auto bench.
// What: deletes the global health entry, reporting prior benched state.
// Test: bench then UnbenchModel returns true; a second call returns false.
func UnbenchModel(spec string) bool { return globalHealth.unbench(spec, time.Now()) }
// IsBenched reports whether a model is currently benched.
//
// Why: callers/tests want a quick health check for a concrete model.
// What: consults the global health registry (expired benches read as false).
// Test: BenchModel makes it true; an expired bench reads false.
func IsBenched(spec string) bool { return globalHealth.isBenched(spec, time.Now()) }
// ---------------------------------------------------------------------------
// Observer
// ---------------------------------------------------------------------------
// FailoverEvent describes a single failover decision for an observer.
type FailoverEvent struct {
Model string
Err error
Kind ErrKind
Attempt int
Benched bool
BenchedFor time.Duration
NextModel string
Request provider.Request
}
// FailoverObserver receives a FailoverEvent for each failover decision. mort
// uses this to persist the full prompt chain on failover.
type FailoverObserver func(ctx context.Context, ev FailoverEvent)
// ---------------------------------------------------------------------------
// Config + options
// ---------------------------------------------------------------------------
type failoverConfig struct {
maxRetries int
cooldown time.Duration
backoff func(attempt int) time.Duration
observer FailoverObserver
}
func defaultFailoverConfig() failoverConfig {
defaultsMu.Lock()
defer defaultsMu.Unlock()
return failoverConfig{
maxRetries: DefaultFailoverMaxRetries,
cooldown: DefaultFailoverCooldown,
backoff: DefaultFailoverBackoff,
// Seed the package-level default observer. An explicit
// WithFailoverObserver applied after this in NewFailoverModel/ParseChain
// overrides it for that chain. Read under the same defaultsMu we already
// hold (a single Lock above), so no re-lock / deadlock.
observer: defaultFailoverObserver,
}
}
// FailoverOption configures a failover model.
type FailoverOption func(*failoverConfig)
// WithFailoverMaxRetries sets attempts per entry on transient errors.
func WithFailoverMaxRetries(n int) FailoverOption {
return func(c *failoverConfig) {
if n < 1 {
n = 1
}
c.maxRetries = n
}
}
// WithFailoverCooldown sets how long a model stays benched after failure.
func WithFailoverCooldown(d time.Duration) FailoverOption {
return func(c *failoverConfig) { c.cooldown = d }
}
// WithFailoverBackoff sets the retry backoff function.
func WithFailoverBackoff(fn func(attempt int) time.Duration) FailoverOption {
return func(c *failoverConfig) {
if fn != nil {
c.backoff = fn
}
}
}
// WithFailoverObserver sets an observer notified on every failover decision.
func WithFailoverObserver(obs FailoverObserver) FailoverOption {
return func(c *failoverConfig) { c.observer = obs }
}
// ---------------------------------------------------------------------------
// Composite provider
// ---------------------------------------------------------------------------
type failoverEntry struct {
provider provider.Provider
model string // bare model name sent to the provider
specKey string // global health key (full concrete spec)
}
type failoverProvider struct {
entries []failoverEntry
cfg failoverConfig
health *modelHealth
}
// bareModel strips a leading "provider/" prefix, returning the model name the
// underlying provider expects. Specs without a slash are returned unchanged.
func bareModel(spec string) string {
if i := strings.Index(spec, "/"); i >= 0 {
return spec[i+1:]
}
return spec
}
// NewFailoverModel builds a composite *Model that tries each sub-model in order,
// retrying/benching per the configured policy and failing over on error.
//
// Why: callers hold *Model and the base Complete handler is hardwired to one
// provider, so failover (which must switch providers) is implemented as a
// composite provider wrapped back into a *Model.
// What: flattens any nested failover sub-models, derives a specKey per entry
// from its model string, and returns NewClient(fp).Model("failover").
// Test: failover_test.go exercises success, failover, bench, abort, and flatten.
func NewFailoverModel(models []*Model, opts ...FailoverOption) *Model {
cfg := defaultFailoverConfig()
for _, opt := range opts {
opt(&cfg)
}
var entries []failoverEntry
for _, m := range models {
if m == nil {
continue
}
// Flatten nested failover models so cooldowns/keys stay flat.
if fp, ok := m.provider.(*failoverProvider); ok {
entries = append(entries, fp.entries...)
continue
}
entries = append(entries, failoverEntry{
provider: m.provider,
model: bareModel(m.model),
specKey: m.model,
})
}
fp := &failoverProvider{
entries: entries,
cfg: cfg,
health: globalHealth,
}
return NewClient(fp).Model("failover")
}
// ParseChain parses each spec and combines them into one failover model.
//
// Why: lets callers build a failover chain from a slice of specs (full
// resolution per entry) without manually wiring providers.
// What: Parse each spec, preserve the original spec string as the bench key,
// flatten nested failover models, and return a composite *Model.
// Test: parse_test.go covers comma-spec parsing through the registry.
func ParseChain(specs []string, opts ...FailoverOption) (*Model, error) {
return DefaultRegistry.ParseChain(specs, opts...)
}
// ParseChain is the registry-scoped form of ParseChain.
func (r *Registry) ParseChain(specs []string, opts ...FailoverOption) (*Model, error) {
cfg := defaultFailoverConfig()
for _, opt := range opts {
opt(&cfg)
}
var entries []failoverEntry
for _, spec := range specs {
spec = strings.TrimSpace(spec)
if spec == "" {
continue
}
m, err := r.Parse(spec)
if err != nil {
return nil, fmt.Errorf("failover chain: parse %q: %w", spec, err)
}
if fp, ok := m.provider.(*failoverProvider); ok {
// A sub-spec was itself a comma/failover spec — splice its entries.
entries = append(entries, fp.entries...)
continue
}
entries = append(entries, failoverEntry{
provider: m.provider,
model: m.model,
specKey: spec,
})
}
if len(entries) == 0 {
return nil, fmt.Errorf("failover chain: no valid specs")
}
fp := &failoverProvider{entries: entries, cfg: cfg, health: globalHealth}
return NewClient(fp).Model("failover"), nil
}
// reqWithModel returns a shallow copy of req with its Model set to model.
//
// Why: each sub-provider must receive its own bare model name; the incoming
// req carries the placeholder "failover" model from the composite *Model.
// What: copies the struct (slices/pointers are shared, which is safe here since
// providers treat the request as read-only) and overrides Model.
// Test: TestFailover_PassesModelNameToProvider asserts the provider sees the bare name.
func reqWithModel(req provider.Request, model string) provider.Request {
req.Model = model
return req
}
// Complete implements provider.Provider with ordered failover.
func (f *failoverProvider) Complete(ctx context.Context, req provider.Request) (provider.Response, error) {
now := time.Now()
// 1. Build the live set (not currently benched). Best-effort: if all are
// benched, ignore cooldowns rather than hard-fail.
var live []failoverEntry
for _, e := range f.entries {
if !f.health.isBenched(e.specKey, now) {
live = append(live, e)
}
}
if len(live) == 0 {
live = f.entries
}
var causes []error
for i, entry := range live {
nextModel := ""
if i+1 < len(live) {
nextModel = live[i+1].specKey
}
for attempt := 1; attempt <= f.cfg.maxRetries; attempt++ {
resp, err := entry.provider.Complete(ctx, reqWithModel(req, entry.model))
if err == nil {
f.health.recordSuccess(entry.specKey)
return resp, nil
}
// Caller aborted: stop everything, no failover, no bench.
if errors.Is(err, context.Canceled) {
return provider.Response{}, err
}
kind := Classify(err)
switch kind {
case ErrRequestSpecific:
f.emit(ctx, FailoverEvent{
Model: entry.specKey, Err: err, Kind: kind, Attempt: attempt,
NextModel: nextModel, Request: req,
})
slog.Warn("failover: request-specific error, trying next model",
"model", entry.specKey, "kind", "request_specific",
"status", statusOf(err), "attempt", attempt, "next", nextModel)
causes = append(causes, fmt.Errorf("%s: %w", entry.specKey, err))
goto nextEntry
case ErrAuthDead:
until := f.health.benchNow(entry.specKey, f.cfg.cooldown, time.Now())
f.emit(ctx, FailoverEvent{
Model: entry.specKey, Err: err, Kind: kind, Attempt: attempt,
Benched: true, BenchedFor: f.cfg.cooldown, NextModel: nextModel, Request: req,
})
slog.Warn("failover: auth/model-dead error, benching model",
"model", entry.specKey, "kind", "auth_dead", "status", statusOf(err),
"attempt", attempt, "benched", true, "cooldown", f.cfg.cooldown,
"until", until, "next", nextModel)
causes = append(causes, fmt.Errorf("%s: %w", entry.specKey, err))
goto nextEntry
default: // ErrTransient or ErrUnknown -> retry, then bench.
if attempt >= f.cfg.maxRetries {
benched, until := f.health.recordTransientFailure(entry.specKey, f.cfg.cooldown, time.Now())
f.emit(ctx, FailoverEvent{
Model: entry.specKey, Err: err, Kind: kind, Attempt: attempt,
Benched: benched, BenchedFor: f.cfg.cooldown, NextModel: nextModel, Request: req,
})
slog.Warn("failover: transient error, retries exhausted",
"model", entry.specKey, "kind", kindString(kind), "status", statusOf(err),
"attempt", attempt, "benched", benched, "cooldown", f.cfg.cooldown,
"until", until, "next", nextModel)
causes = append(causes, fmt.Errorf("%s: %w", entry.specKey, err))
goto nextEntry
}
// Sleep before retrying (respect ctx).
select {
case <-ctx.Done():
return provider.Response{}, ctx.Err()
case <-time.After(f.cfg.backoff(attempt)):
}
}
}
nextEntry:
}
return provider.Response{}, fmt.Errorf("failover: all %d models in chain failed: %w",
len(live), errors.Join(causes...))
}
// Stream implements provider.Provider. It fails over only on the INITIAL Stream
// call error (before any event). Once a stream begins, mid-stream failures are
// surfaced as-is — failover does not replay a partially-consumed stream.
func (f *failoverProvider) Stream(ctx context.Context, req provider.Request, events chan<- provider.StreamEvent) error {
now := time.Now()
var live []failoverEntry
for _, e := range f.entries {
if !f.health.isBenched(e.specKey, now) {
live = append(live, e)
}
}
if len(live) == 0 {
live = f.entries
}
var causes []error
for i, entry := range live {
nextModel := ""
if i+1 < len(live) {
nextModel = live[i+1].specKey
}
err := entry.provider.Stream(ctx, reqWithModel(req, entry.model), events)
if err == nil {
f.health.recordSuccess(entry.specKey)
return nil
}
if errors.Is(err, context.Canceled) {
return err
}
kind := Classify(err)
switch kind {
case ErrAuthDead:
f.health.benchNow(entry.specKey, f.cfg.cooldown, time.Now())
case ErrTransient, ErrUnknown:
f.health.recordTransientFailure(entry.specKey, f.cfg.cooldown, time.Now())
}
slog.Warn("failover(stream): error, trying next model",
"model", entry.specKey, "kind", kindString(kind), "status", statusOf(err), "next", nextModel)
causes = append(causes, fmt.Errorf("%s: %w", entry.specKey, err))
}
return fmt.Errorf("failover(stream): all %d models in chain failed: %w", len(live), errors.Join(causes...))
}
func (f *failoverProvider) emit(ctx context.Context, ev FailoverEvent) {
if f.cfg.observer != nil {
f.cfg.observer(ctx, ev)
}
}
func kindString(k ErrKind) string {
switch k {
case ErrTransient:
return "transient"
case ErrAuthDead:
return "auth_dead"
case ErrRequestSpecific:
return "request_specific"
default:
return "unknown"
}
}
// statusOf best-effort extracts an HTTP status code for logging, or 0.
func statusOf(err error) int { return extractStatus(err) }
+509
View File
@@ -0,0 +1,509 @@
package llm
import (
"context"
"errors"
"strings"
"sync"
"testing"
"time"
"github.com/openai/openai-go"
"gitea.stevedudenhoeffer.com/steve/go-llm/v2/provider"
)
// fastBackoff is a near-zero backoff so retry tests don't sleep.
func fastBackoff(int) time.Duration { return time.Microsecond }
func testFailoverOpts(extra ...FailoverOption) []FailoverOption {
base := []FailoverOption{
WithFailoverMaxRetries(2),
WithFailoverBackoff(fastBackoff),
WithFailoverCooldown(time.Minute),
}
return append(base, extra...)
}
// modelFor builds a *Model around a mock provider with a concrete model name,
// mimicking what Parse produces (so specKey resolution works).
func modelFor(p provider.Provider, name string) *Model {
return &Model{provider: p, model: name}
}
func TestFailover_FirstSucceeds(t *testing.T) {
resetHealthForTest()
a := newMockProvider(provider.Response{Text: "from-a"})
b := newMockProvider(provider.Response{Text: "from-b"})
fo := NewFailoverModel([]*Model{modelFor(a, "anthropic/a"), modelFor(b, "openai/b")}, testFailoverOpts()...)
resp, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if resp.Text != "from-a" {
t.Errorf("expected from-a, got %q", resp.Text)
}
// b must not have been called.
b.mu.Lock()
n := len(b.Requests)
b.mu.Unlock()
if n != 0 {
t.Errorf("expected b untouched, got %d calls", n)
}
}
func TestFailover_FailsOverToSecond(t *testing.T) {
resetHealthForTest()
// a always returns a request-specific error (400) -> fail over, no retry-bench loop noise.
a := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
return provider.Response{}, &openai.Error{StatusCode: 400}
})
b := newMockProvider(provider.Response{Text: "from-b"})
fo := NewFailoverModel([]*Model{modelFor(a, "p/a"), modelFor(b, "p/b")}, testFailoverOpts()...)
resp, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if resp.Text != "from-b" {
t.Errorf("expected from-b, got %q", resp.Text)
}
// 400 is request-specific: a must NOT be benched.
if IsBenched("p/a") {
t.Error("p/a should not be benched on a 400")
}
}
func TestFailover_PassesModelNameToProvider(t *testing.T) {
resetHealthForTest()
a := newMockProvider(provider.Response{Text: "ok"})
fo := NewFailoverModel([]*Model{modelFor(a, "anthropic/claude-x")}, testFailoverOpts()...)
_, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatal(err)
}
if got := a.lastRequest().Model; got != "claude-x" {
t.Errorf("provider received model %q, want bare model name claude-x", got)
}
}
func TestFailover_AuthDeadBenchesImmediately(t *testing.T) {
resetHealthForTest()
a := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
return provider.Response{}, &openai.Error{StatusCode: 401}
})
b := newMockProvider(provider.Response{Text: "from-b"})
fo := NewFailoverModel([]*Model{modelFor(a, "p/a"), modelFor(b, "p/b")}, testFailoverOpts()...)
resp, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatal(err)
}
if resp.Text != "from-b" {
t.Errorf("expected from-b, got %q", resp.Text)
}
if !IsBenched("p/a") {
t.Error("p/a should be benched after auth-dead error")
}
// a should have been called exactly once (no retries on auth-dead).
a.mu.Lock()
n := len(a.Requests)
a.mu.Unlock()
if n != 1 {
t.Errorf("auth-dead should not retry; a called %d times", n)
}
}
func TestFailover_TransientRetriesThenBenches(t *testing.T) {
resetHealthForTest()
var calls int
var mu sync.Mutex
a := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
mu.Lock()
calls++
mu.Unlock()
return provider.Response{}, &openai.Error{StatusCode: 503}
})
b := newMockProvider(provider.Response{Text: "from-b"})
// maxRetries=2 means 2 attempts total per entry.
fo := NewFailoverModel([]*Model{modelFor(a, "p/a"), modelFor(b, "p/b")},
WithFailoverMaxRetries(2), WithFailoverBackoff(fastBackoff), WithFailoverCooldown(time.Minute))
resp, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatal(err)
}
if resp.Text != "from-b" {
t.Errorf("expected from-b, got %q", resp.Text)
}
mu.Lock()
n := calls
mu.Unlock()
if n != 2 {
t.Errorf("expected 2 attempts on transient model, got %d", n)
}
if !IsBenched("p/a") {
t.Error("p/a should be benched after exhausting retries")
}
}
func TestFailover_AllFail(t *testing.T) {
resetHealthForTest()
a := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
return provider.Response{}, &openai.Error{StatusCode: 400}
})
b := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
return provider.Response{}, &openai.Error{StatusCode: 400}
})
fo := NewFailoverModel([]*Model{modelFor(a, "p/a"), modelFor(b, "p/b")}, testFailoverOpts()...)
_, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err == nil {
t.Fatal("expected error when all models fail")
}
if !strings.Contains(err.Error(), "2") {
t.Errorf("error should mention all 2 models failed: %v", err)
}
}
func TestFailover_ContextCanceledAborts(t *testing.T) {
resetHealthForTest()
a := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
return provider.Response{}, context.Canceled
})
b := newMockProvider(provider.Response{Text: "from-b"})
fo := NewFailoverModel([]*Model{modelFor(a, "p/a"), modelFor(b, "p/b")}, testFailoverOpts()...)
_, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if !errors.Is(err, context.Canceled) {
t.Errorf("expected context.Canceled to abort, got %v", err)
}
// b must not be tried.
b.mu.Lock()
n := len(b.Requests)
b.mu.Unlock()
if n != 0 {
t.Errorf("canceled should not fail over; b called %d times", n)
}
if IsBenched("p/a") {
t.Error("canceled should not bench")
}
}
func TestFailover_AllBenchedBestEffort(t *testing.T) {
resetHealthForTest()
// Manually bench both, then ensure Complete still tries (best-effort) and succeeds.
BenchModel("p/a", time.Now().Add(time.Hour))
BenchModel("p/b", time.Now().Add(time.Hour))
a := newMockProvider(provider.Response{Text: "from-a"})
b := newMockProvider(provider.Response{Text: "from-b"})
fo := NewFailoverModel([]*Model{modelFor(a, "p/a"), modelFor(b, "p/b")}, testFailoverOpts()...)
resp, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatalf("best-effort should still try benched models: %v", err)
}
if resp.Text != "from-a" {
t.Errorf("expected from-a, got %q", resp.Text)
}
}
func TestFailover_Observer(t *testing.T) {
resetHealthForTest()
a := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
return provider.Response{}, &openai.Error{StatusCode: 401}
})
b := newMockProvider(provider.Response{Text: "from-b"})
var mu sync.Mutex
var events []FailoverEvent
obs := func(ctx context.Context, ev FailoverEvent) {
mu.Lock()
events = append(events, ev)
mu.Unlock()
}
fo := NewFailoverModel([]*Model{modelFor(a, "p/a"), modelFor(b, "p/b")},
append(testFailoverOpts(), WithFailoverObserver(obs))...)
_, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatal(err)
}
mu.Lock()
defer mu.Unlock()
if len(events) == 0 {
t.Fatal("expected observer to be called")
}
ev := events[0]
if ev.Model != "p/a" || ev.Kind != ErrAuthDead || !ev.Benched {
t.Errorf("unexpected event: %+v", ev)
}
if ev.NextModel != "p/b" {
t.Errorf("expected NextModel p/b, got %q", ev.NextModel)
}
if len(ev.Request.Messages) == 0 {
t.Error("observer event should carry the full request")
}
}
// TestFailover_DefaultObserverFiresOnTransparentChain verifies that a chain
// built via NewFailoverModel with NO options still notifies a package-level
// default observer set via SetFailoverObserver. This is the transparent
// comma-Parse path: defaultFailoverConfig() must seed the default observer.
func TestFailover_DefaultObserverFiresOnTransparentChain(t *testing.T) {
resetHealthForTest()
t.Cleanup(func() {
SetFailoverObserver(nil)
resetHealthForTest()
})
var mu sync.Mutex
var events []FailoverEvent
SetFailoverObserver(func(ctx context.Context, ev FailoverEvent) {
mu.Lock()
events = append(events, ev)
mu.Unlock()
})
a := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
return provider.Response{}, &openai.Error{StatusCode: 401}
})
b := newMockProvider(provider.Response{Text: "from-b"})
// NO options: the only way the observer can fire is via the package default.
fo := NewFailoverModel([]*Model{modelFor(a, "p/a"), modelFor(b, "p/b")})
_, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatal(err)
}
mu.Lock()
defer mu.Unlock()
if len(events) == 0 {
t.Fatal("expected default observer to fire on a transparently-built chain")
}
if events[0].Model != "p/a" || events[0].Kind != ErrAuthDead {
t.Errorf("unexpected event: %+v", events[0])
}
}
// TestFailover_DefaultObserverFiresOnParseChain verifies the comma-Parse seam:
// a chain built through the registry's ParseChain (no per-call observer) fires
// the package-level default observer.
func TestFailover_DefaultObserverFiresOnParseChain(t *testing.T) {
resetHealthForTest()
t.Cleanup(func() {
SetFailoverObserver(nil)
resetHealthForTest()
})
var mu sync.Mutex
var events []FailoverEvent
SetFailoverObserver(func(ctx context.Context, ev FailoverEvent) {
mu.Lock()
events = append(events, ev)
mu.Unlock()
})
r, alpha, _ := testRegistry(nil)
// alpha returns a 401 (auth-dead) so the chain fails over and emits an event.
alpha.err = &openai.Error{StatusCode: 401}
m, err := r.Parse("alpha/model-a,beta/model-b")
if err != nil {
t.Fatalf("parse failed: %v", err)
}
if _, ok := m.provider.(*failoverProvider); !ok {
t.Fatalf("expected a failover provider, got %T", m.provider)
}
_, err = m.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatal(err)
}
mu.Lock()
defer mu.Unlock()
if len(events) == 0 {
t.Fatal("expected default observer to fire on a comma-Parse'd chain")
}
}
// TestFailover_ExplicitObserverOverridesDefault verifies WithFailoverObserver
// still wins: when both a package default and an explicit observer are present,
// only the explicit one fires for that chain.
func TestFailover_ExplicitObserverOverridesDefault(t *testing.T) {
resetHealthForTest()
t.Cleanup(func() {
SetFailoverObserver(nil)
resetHealthForTest()
})
var mu sync.Mutex
var defaultCalls, explicitCalls int
SetFailoverObserver(func(ctx context.Context, ev FailoverEvent) {
mu.Lock()
defaultCalls++
mu.Unlock()
})
explicit := func(ctx context.Context, ev FailoverEvent) {
mu.Lock()
explicitCalls++
mu.Unlock()
}
a := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
return provider.Response{}, &openai.Error{StatusCode: 401}
})
b := newMockProvider(provider.Response{Text: "from-b"})
fo := NewFailoverModel([]*Model{modelFor(a, "p/a"), modelFor(b, "p/b")},
WithFailoverObserver(explicit))
_, err := fo.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatal(err)
}
mu.Lock()
defer mu.Unlock()
if explicitCalls == 0 {
t.Error("explicit observer should fire")
}
if defaultCalls != 0 {
t.Errorf("default observer must NOT fire when an explicit one is set; got %d calls", defaultCalls)
}
}
func TestFailover_ControlAPI(t *testing.T) {
resetHealthForTest()
if IsBenched("x/y") {
t.Error("should start unbenched")
}
until := time.Now().Add(time.Hour)
BenchModel("x/y", until)
if !IsBenched("x/y") {
t.Error("should be benched")
}
list := ListBenched()
if len(list) != 1 || list[0].Model != "x/y" || !list[0].Manual {
t.Errorf("unexpected list: %+v", list)
}
if !UnbenchModel("x/y") {
t.Error("UnbenchModel should report it was benched")
}
if IsBenched("x/y") {
t.Error("should be unbenched now")
}
if UnbenchModel("x/y") {
t.Error("UnbenchModel on non-benched should return false")
}
}
func TestFailover_ExpiredBenchIsLive(t *testing.T) {
resetHealthForTest()
// Bench in the past -> should be considered live again.
BenchModel("p/a", time.Now().Add(-time.Hour))
if IsBenched("p/a") {
t.Error("expired bench should not count as benched")
}
}
// TestParseChain exercises ParseChain via a registry-backed seam is covered in
// parse_test.go; here we verify NewFailoverModel flattens nested failover models.
func TestNewFailoverModel_Flattens(t *testing.T) {
resetHealthForTest()
a := newMockProvider(provider.Response{Text: "a"})
b := newMockProvider(provider.Response{Text: "b"})
c := newMockProvider(provider.Response{Text: "c"})
inner := NewFailoverModel([]*Model{modelFor(b, "p/b"), modelFor(c, "p/c")}, testFailoverOpts()...)
outer := NewFailoverModel([]*Model{modelFor(a, "p/a"), inner}, testFailoverOpts()...)
fp, ok := outer.provider.(*failoverProvider)
if !ok {
t.Fatalf("expected *failoverProvider, got %T", outer.provider)
}
if len(fp.entries) != 3 {
t.Errorf("expected flattened 3 entries, got %d", len(fp.entries))
}
keys := []string{fp.entries[0].specKey, fp.entries[1].specKey, fp.entries[2].specKey}
want := []string{"p/a", "p/b", "p/c"}
for i := range want {
if keys[i] != want[i] {
t.Errorf("entry %d specKey = %q, want %q", i, keys[i], want[i])
}
}
}
// TestFailover_ManualBenchSurvivesAutomaticDowngrade is a regression test:
// an active manual bench (long window) must NOT be cleared or shortened by the
// automatic recordTransientFailure/benchNow paths. The best-effort all-benched
// failover loop can re-try a manually benched model; if it fails, the automatic
// logic previously overwrote manual=true -> false and shortened until to the
// short auto cooldown. The operator's intent must win.
func TestFailover_ManualBenchSurvivesAutomaticDowngrade(t *testing.T) {
resetHealthForTest()
const key = "p/manual"
longUntil := time.Now().Add(time.Hour)
shortCooldown := time.Minute
now := time.Now()
// Operator manually benches the model for a long window.
BenchModel(key, longUntil)
assertManualPreserved := func(stage string) {
t.Helper()
list := ListBenched()
var got *BenchedModel
for i := range list {
if list[i].Model == key {
got = &list[i]
break
}
}
if got == nil {
t.Fatalf("%s: %q missing from ListBenched; manual bench was cleared", stage, key)
}
if !got.Manual {
t.Errorf("%s: Manual=false, want true (automatic logic downgraded manual bench)", stage)
}
if !got.Until.Equal(longUntil) {
t.Errorf("%s: Until=%v, want %v (automatic logic shortened operator window)", stage, got.Until, longUntil)
}
}
// (c) Active manual bench hit by the automatic transient path.
benched, until := globalHealth.recordTransientFailure(key, shortCooldown, now)
if !benched {
t.Errorf("recordTransientFailure: benched=false, want true (model stays benched)")
}
if !until.Equal(longUntil) {
t.Errorf("recordTransientFailure: until=%v, want %v (long manual window preserved)", until, longUntil)
}
assertManualPreserved("after recordTransientFailure")
// (c) Active manual bench hit by the automatic auth-dead path.
until = globalHealth.benchNow(key, shortCooldown, now)
if !until.Equal(longUntil) {
t.Errorf("benchNow: until=%v, want %v (long manual window preserved)", until, longUntil)
}
assertManualPreserved("after benchNow")
// (d) Expired manual bench: automatic logic IS allowed to take over.
resetHealthForTest()
expired := time.Now().Add(-time.Hour)
BenchModel(key, expired)
until = globalHealth.benchNow(key, shortCooldown, now)
if want := now.Add(shortCooldown); !until.Equal(want) {
t.Errorf("benchNow on expired manual: until=%v, want auto cooldown %v", until, want)
}
// (a) No prior state and (b) prior automatic bench: automatic cooldown applies.
resetHealthForTest()
benched, until = globalHealth.recordTransientFailure(key, shortCooldown, now) // (a)
if !benched || !until.Equal(now.Add(shortCooldown)) {
t.Errorf("recordTransientFailure(no prior): benched=%v until=%v, want true %v", benched, until, now.Add(shortCooldown))
}
until = globalHealth.benchNow(key, shortCooldown, now) // (b) prior automatic state
if want := now.Add(shortCooldown); !until.Equal(want) {
t.Errorf("benchNow(prior auto): until=%v, want %v", until, want)
}
if list := ListBenched(); len(list) != 1 || list[0].Manual {
t.Errorf("automatic bench should not be Manual: %+v", list)
}
}
+93 -11
View File
@@ -12,8 +12,11 @@ import (
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"strconv"
"strings"
"time"
"gitea.stevedudenhoeffer.com/steve/go-llm/v2/provider"
)
@@ -25,6 +28,22 @@ const DefaultLocalBaseURL = "http://localhost:11434"
// DefaultCloudBaseURL is the default base URL for Ollama Cloud.
const DefaultCloudBaseURL = "https://ollama.com"
// retryMaxAttempts is the maximum number of retry attempts for transient HTTP
// errors (503, 429, 502). Total attempts = 1 initial + retryMaxAttempts.
const retryMaxAttempts = 3
// retryBaseDelay is the base delay for exponential backoff between retries.
// Actual delays: 1s, 2s, 4s (base * 2^attempt).
const retryBaseDelay = 1 * time.Second
// isTransientHTTPStatus reports whether the HTTP status code indicates a
// transient server-side condition that may resolve on retry.
func isTransientHTTPStatus(code int) bool {
return code == http.StatusBadGateway || // 502
code == http.StatusServiceUnavailable || // 503
code == http.StatusTooManyRequests // 429
}
// Provider implements provider.Provider over Ollama's native /api/chat
// endpoint. An empty apiKey means local-mode (no Authorization header sent);
// a non-empty apiKey is sent as a Bearer token (cloud-mode).
@@ -32,6 +51,10 @@ type Provider struct {
apiKey string
baseURL string
client *http.Client
// retryBaseDelayOverride, when non-zero, replaces retryBaseDelay for
// testing. Production code leaves this at the zero value.
retryBaseDelayOverride time.Duration
}
// newNative constructs a native Ollama provider. Callers should use the
@@ -420,22 +443,81 @@ func (p *Provider) buildChatRequest(req provider.Request, stream bool) ([]byte,
}
// doChatRequest POSTs the wire body to /api/chat and returns the raw HTTP
// response. The caller is responsible for closing the response body.
// response. Transient HTTP errors (502, 503, 429) are retried with exponential
// backoff up to retryMaxAttempts times. The caller is responsible for closing
// the response body.
func (p *Provider) doChatRequest(ctx context.Context, body []byte) (*http.Response, error) {
url := strings.TrimRight(p.baseURL, "/") + "/api/chat"
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("ollama: build request: %w", err)
for attempt := 0; ; attempt++ {
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("ollama: build request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
if p.apiKey != "" {
httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
}
resp, err := p.client.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("ollama: HTTP request: %w", err)
}
// On success or non-transient error, return immediately.
if !isTransientHTTPStatus(resp.StatusCode) || attempt >= retryMaxAttempts {
return resp, nil
}
// Transient error — drain and close the body before retrying.
respBody, _ := io.ReadAll(resp.Body)
resp.Body.Close()
delay := retryBackoff(attempt, resp.Header, p.retryBaseDelayOverride)
slog.Info("ollama: retrying after transient HTTP error",
"status", resp.StatusCode,
"attempt", attempt+1,
"max_attempts", retryMaxAttempts,
"delay", delay,
"body", truncateBody(respBody, 200),
)
// Wait for backoff or context cancellation.
timer := time.NewTimer(delay)
select {
case <-ctx.Done():
timer.Stop()
return nil, ctx.Err()
case <-timer.C:
}
}
httpReq.Header.Set("Content-Type", "application/json")
if p.apiKey != "" {
httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
}
// retryBackoff computes the delay before the next retry attempt. It uses
// exponential backoff (base * 2^attempt), but respects the Retry-After header
// when present (for 429 responses). baseOverride, when non-zero, replaces the
// package-level retryBaseDelay constant (used by tests to avoid real waits).
func retryBackoff(attempt int, header http.Header, baseOverride time.Duration) time.Duration {
// Check Retry-After header (seconds value or HTTP-date; we only parse seconds).
if ra := header.Get("Retry-After"); ra != "" {
if secs, err := strconv.Atoi(ra); err == nil && secs > 0 {
return time.Duration(secs) * time.Second
}
}
resp, err := p.client.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("ollama: HTTP request: %w", err)
base := retryBaseDelay
if baseOverride > 0 {
base = baseOverride
}
return resp, nil
return base * (1 << attempt)
}
// truncateBody returns a string of at most maxLen bytes from b, appending
// "..." when truncated. Used for readable log output of error response bodies.
func truncateBody(b []byte, maxLen int) string {
if len(b) <= maxLen {
return string(b)
}
return string(b[:maxLen]) + "..."
}
// convertMessage maps a provider.Message into a native wire message.
+314
View File
@@ -571,3 +571,317 @@ func equalStrings(a, b []string) bool {
}
return true
}
// --- Retry tests ---
// newTestNative creates a Provider with a minimal retry delay so tests run fast.
func newTestNative(apiKey, baseURL string) *Provider {
p := newNative(apiKey, baseURL)
p.retryBaseDelayOverride = 1 * time.Millisecond
return p
}
func TestRetryOnTransientHTTPError(t *testing.T) {
t.Run("503 retries then succeeds", func(t *testing.T) {
var attempts int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Drain the request body so the client can reuse the connection.
_, _ = io.ReadAll(r.Body)
attempts++
if attempts <= 2 {
w.WriteHeader(503)
_, _ = w.Write([]byte(`{"error":"model is temporarily overloaded"}`))
return
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(200)
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"ok"},"done":true,"prompt_eval_count":1,"eval_count":1}`))
}))
t.Cleanup(srv.Close)
p := newTestNative("key", srv.URL)
resp, err := p.Complete(context.Background(), provider.Request{
Model: "test",
Messages: []provider.Message{{Role: "user", Content: "hi"}},
})
if err != nil {
t.Fatalf("expected success after retries, got error: %v", err)
}
if resp.Text != "ok" {
t.Errorf("Text: want %q, got %q", "ok", resp.Text)
}
if attempts != 3 {
t.Errorf("attempts: want 3, got %d", attempts)
}
})
t.Run("429 retries then succeeds", func(t *testing.T) {
var attempts int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.ReadAll(r.Body)
attempts++
if attempts == 1 {
w.WriteHeader(429)
_, _ = w.Write([]byte(`{"error":"rate limited"}`))
return
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(200)
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"done"},"done":true}`))
}))
t.Cleanup(srv.Close)
p := newTestNative("", srv.URL)
resp, err := p.Complete(context.Background(), provider.Request{
Model: "test",
Messages: []provider.Message{{Role: "user", Content: "hi"}},
})
if err != nil {
t.Fatalf("expected success after retry, got error: %v", err)
}
if resp.Text != "done" {
t.Errorf("Text: want %q, got %q", "done", resp.Text)
}
if attempts != 2 {
t.Errorf("attempts: want 2, got %d", attempts)
}
})
t.Run("502 retries then succeeds", func(t *testing.T) {
var attempts int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.ReadAll(r.Body)
attempts++
if attempts == 1 {
w.WriteHeader(502)
_, _ = w.Write([]byte(`Bad Gateway`))
return
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(200)
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"yes"},"done":true}`))
}))
t.Cleanup(srv.Close)
p := newTestNative("", srv.URL)
resp, err := p.Complete(context.Background(), provider.Request{
Model: "test",
Messages: []provider.Message{{Role: "user", Content: "hi"}},
})
if err != nil {
t.Fatalf("expected success after retry, got error: %v", err)
}
if resp.Text != "yes" {
t.Errorf("Text: want %q, got %q", "yes", resp.Text)
}
})
t.Run("exhausts retries and returns error", func(t *testing.T) {
var attempts int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.ReadAll(r.Body)
attempts++
w.WriteHeader(503)
_, _ = w.Write([]byte(`{"error":"overloaded"}`))
}))
t.Cleanup(srv.Close)
p := newTestNative("", srv.URL)
_, err := p.Complete(context.Background(), provider.Request{
Model: "test",
Messages: []provider.Message{{Role: "user", Content: "hi"}},
})
if err == nil {
t.Fatal("expected error after exhausting retries")
}
if !strings.Contains(err.Error(), "503") {
t.Errorf("error should mention status 503, got: %v", err)
}
// 1 initial + 3 retries = 4 total
if attempts != 4 {
t.Errorf("attempts: want 4, got %d", attempts)
}
})
t.Run("400 is not retried", func(t *testing.T) {
var attempts int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.ReadAll(r.Body)
attempts++
w.WriteHeader(400)
_, _ = w.Write([]byte(`{"error":"bad request"}`))
}))
t.Cleanup(srv.Close)
p := newTestNative("", srv.URL)
_, err := p.Complete(context.Background(), provider.Request{
Model: "test",
Messages: []provider.Message{{Role: "user", Content: "hi"}},
})
if err == nil {
t.Fatal("expected error for 400")
}
if attempts != 1 {
t.Errorf("attempts: want 1 (no retries for 400), got %d", attempts)
}
})
t.Run("context cancellation during backoff aborts retry", func(t *testing.T) {
var attempts int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.ReadAll(r.Body)
attempts++
w.WriteHeader(503)
_, _ = w.Write([]byte(`{"error":"overloaded"}`))
}))
t.Cleanup(srv.Close)
ctx, cancel := context.WithCancel(context.Background())
// Cancel shortly after the first attempt so the backoff wait is interrupted.
go func() {
time.Sleep(50 * time.Millisecond)
cancel()
}()
p := newTestNative("", srv.URL)
// Use a longer backoff so the context cancel fires during the wait.
p.retryBaseDelayOverride = 2 * time.Second
_, err := p.Complete(ctx, provider.Request{
Model: "test",
Messages: []provider.Message{{Role: "user", Content: "hi"}},
})
if err == nil {
t.Fatal("expected error from cancelled context")
}
if !strings.Contains(err.Error(), "canceled") && !strings.Contains(err.Error(), "context") {
t.Errorf("expected context error, got: %v", err)
}
// Should have made only 1 attempt before the context cancelled during backoff.
if attempts != 1 {
t.Errorf("attempts: want 1, got %d", attempts)
}
})
t.Run("stream retries on 503 then succeeds", func(t *testing.T) {
var attempts int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.ReadAll(r.Body)
attempts++
if attempts == 1 {
w.WriteHeader(503)
_, _ = w.Write([]byte(`{"error":"overloaded"}`))
return
}
w.Header().Set("Content-Type", "application/x-ndjson")
w.WriteHeader(200)
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"streamed"},"done":true,"prompt_eval_count":1,"eval_count":1}` + "\n"))
}))
t.Cleanup(srv.Close)
p := newTestNative("", srv.URL)
events := collectStream(t, p, provider.Request{
Model: "test",
Messages: []provider.Message{{Role: "user", Content: "hi"}},
})
var gotDone bool
for _, ev := range events {
if ev.Type == provider.StreamEventDone && ev.Response != nil {
if ev.Response.Text != "streamed" {
t.Errorf("Response.Text: want %q, got %q", "streamed", ev.Response.Text)
}
gotDone = true
}
}
if !gotDone {
t.Fatal("expected StreamEventDone")
}
if attempts != 2 {
t.Errorf("attempts: want 2, got %d", attempts)
}
})
}
func TestRetryBackoff(t *testing.T) {
t.Run("exponential backoff without Retry-After", func(t *testing.T) {
h := http.Header{}
d0 := retryBackoff(0, h, 0)
d1 := retryBackoff(1, h, 0)
d2 := retryBackoff(2, h, 0)
if d0 != 1*time.Second {
t.Errorf("attempt 0: want 1s, got %v", d0)
}
if d1 != 2*time.Second {
t.Errorf("attempt 1: want 2s, got %v", d1)
}
if d2 != 4*time.Second {
t.Errorf("attempt 2: want 4s, got %v", d2)
}
})
t.Run("Retry-After header overrides backoff", func(t *testing.T) {
h := http.Header{}
h.Set("Retry-After", "5")
d := retryBackoff(0, h, 0)
if d != 5*time.Second {
t.Errorf("want 5s from Retry-After, got %v", d)
}
})
t.Run("invalid Retry-After falls back to exponential", func(t *testing.T) {
h := http.Header{}
h.Set("Retry-After", "not-a-number")
d := retryBackoff(1, h, 0)
if d != 2*time.Second {
t.Errorf("want 2s fallback, got %v", d)
}
})
t.Run("baseOverride replaces default base delay", func(t *testing.T) {
h := http.Header{}
d := retryBackoff(0, h, 500*time.Millisecond)
if d != 500*time.Millisecond {
t.Errorf("attempt 0 with 500ms override: want 500ms, got %v", d)
}
d1 := retryBackoff(2, h, 500*time.Millisecond)
if d1 != 2*time.Second {
t.Errorf("attempt 2 with 500ms override: want 2s, got %v", d1)
}
})
}
func TestIsTransientHTTPStatus(t *testing.T) {
cases := []struct {
code int
want bool
}{
{200, false},
{400, false},
{401, false},
{403, false},
{404, false},
{429, true},
{500, false},
{502, true},
{503, true},
}
for _, c := range cases {
if got := isTransientHTTPStatus(c.code); got != c.want {
t.Errorf("isTransientHTTPStatus(%d): want %v, got %v", c.code, c.want, got)
}
}
}
func TestTruncateBody(t *testing.T) {
short := "hello"
if got := truncateBody([]byte(short), 10); got != short {
t.Errorf("short: want %q, got %q", short, got)
}
long := "hello world this is a very long string"
got := truncateBody([]byte(long), 10)
if got != "hello worl..." {
t.Errorf("long: want %q, got %q", "hello worl...", got)
}
}
+266
View File
@@ -0,0 +1,266 @@
package llm
import (
"fmt"
"os"
"strings"
)
// DSN represents a parsed Data Source Name for an LLM provider endpoint.
// Format: scheme://[token@]host[/path]
//
// Why: multi-instance providers (e.g., multiple foreman daemons) need a compact
// way to encode scheme, credentials, and host in a single env var.
// What: holds the three components after parsing.
// Test: ParseDSN("foreman://tok@host") → {Scheme:"foreman", Token:"tok", Host:"host"}.
type DSN struct {
Scheme string // provider type: "foreman", "ollama", etc.
Token string // API key / bearer token; empty = none
Host string // hostname[:port][/path], no scheme prefix
}
// ParseDSN parses a raw DSN string into its components.
// Expected format: scheme://[token@]host[/path]
//
// Why: LLM_X env vars encode provider type, optional credentials, and host
// in a single string; this function decodes that.
// What: splits on "://", then optional "@" for token, remainder is host.
// Test: valid DSNs parse correctly, missing scheme or host returns ErrInvalidDSN.
func ParseDSN(raw string) (DSN, error) {
schemeEnd := strings.Index(raw, "://")
if schemeEnd < 0 {
return DSN{}, fmt.Errorf("%w: missing scheme://: %q", ErrInvalidDSN, raw)
}
scheme := raw[:schemeEnd]
rest := raw[schemeEnd+3:]
var token, host string
if atIdx := strings.Index(rest, "@"); atIdx >= 0 {
token = rest[:atIdx]
host = rest[atIdx+1:]
} else {
host = rest
}
host = strings.TrimRight(host, "/")
if host == "" {
return DSN{}, fmt.Errorf("%w: missing host: %q", ErrInvalidDSN, raw)
}
return DSN{Scheme: scheme, Token: token, Host: host}, nil
}
// splitReasoning strips a trailing ":low", ":medium", or ":high" reasoning
// suffix from a spec string. Only the last colon-delimited segment is
// considered, and only if it matches a known ReasoningLevel. This preserves
// Ollama-style tags like ":30b" or ":14b" which are not reasoning levels.
//
// Why: reasoning level is encoded as a suffix in the spec grammar, but
// colons also appear in Ollama model tags — this function disambiguates.
// What: returns the base string and the extracted level (empty if none).
// Test: "gpt-4o:high" → ("gpt-4o", high); "qwen3:30b" → ("qwen3:30b", "").
func splitReasoning(s string) (string, ReasoningLevel) {
idx := strings.LastIndex(s, ":")
if idx < 0 || idx == len(s)-1 {
return s, ""
}
suffix := ReasoningLevel(s[idx+1:])
switch suffix {
case ReasoningLow, ReasoningMedium, ReasoningHigh:
return s[:idx], suffix
}
return s, ""
}
// Parse resolves a spec string to a ready-to-use *Model using the
// DefaultRegistry. See Registry.Parse for the full grammar and resolution
// order.
//
// Why: provides a one-call entry point for resolving model strings without
// requiring callers to interact with the Registry directly.
// What: delegates to DefaultRegistry.Parse.
// Test: Parse("openai/gpt-4o") returns a non-nil *Model.
func Parse(spec string) (*Model, error) {
return DefaultRegistry.Parse(spec)
}
// Parse resolves a spec string to a ready-to-use *Model.
//
// Spec grammar:
//
// spec = alias | provider "/" model | envname "/" model
// alias = name (registered via RegisterAlias or matched by a Resolver)
// provider = registered-name (e.g., "openai", "foreman")
// envname = name (resolved via LLM_{UPPER(name)} env var containing a DSN)
// model = everything after the first "/"
//
// Any spec may carry a ":reasoning" suffix (":low", ":medium", ":high") after
// the last colon. Ollama-style tags like ":30b" are NOT consumed as reasoning.
//
// Resolution order:
// 1. Strip reasoning suffix
// 2. Check static aliases → recurse
// 3. Check dynamic resolvers → recurse
// 4. Split on first "/" → provider/model
// 5. Look up provider in registry
// 6. Look up LLM_{UPPER(left)} env var → parse DSN → create client
// 7. Return ErrUnknownProvider
//
// Why: consumers need a single function to go from a user-supplied string
// (CLI flag, config file, database row) to a ready-to-use Model.
// What: walks the resolution chain and returns the final *Model with reasoning applied.
// Test: see parse_test.go for comprehensive table-driven tests.
func (r *Registry) Parse(spec string) (*Model, error) {
// Comma-separated specs become an ordered failover chain. A single part
// (after trimming/dropping empties) falls through to normal single-model
// parsing, preserving exact existing behavior for comma-free specs.
if strings.Contains(spec, ",") {
var parts []string
for _, p := range strings.Split(spec, ",") {
if p = strings.TrimSpace(p); p != "" {
parts = append(parts, p)
}
}
if len(parts) == 0 {
return nil, fmt.Errorf("%w: empty failover spec %q", ErrUnknownProvider, spec)
}
if len(parts) > 1 {
return r.ParseChain(parts)
}
spec = parts[0]
}
m, level, err := r.parse(spec, 0)
if err != nil {
return nil, err
}
if level != "" {
m = m.WithReasoning(level)
}
return m, nil
}
// parse is the internal recursive resolver. depth is bounded to prevent
// alias loops.
func (r *Registry) parse(spec string, depth int) (*Model, ReasoningLevel, error) {
if depth > 10 {
return nil, "", ErrAliasLoop
}
// 1. Strip reasoning suffix.
base, userLevel := splitReasoning(spec)
// 2. Check static aliases.
r.mu.RLock()
target, isAlias := r.aliases[base]
r.mu.RUnlock()
if isAlias {
// An alias may expand to a comma-separated failover chain; route those
// through the comma-aware public Parse so the chain is built correctly.
if strings.Contains(target, ",") {
m, err := r.Parse(target)
if err != nil {
return nil, "", err
}
return m, userLevel, nil
}
m, aliasLevel, err := r.parse(target, depth+1)
if err != nil {
return nil, "", err
}
if userLevel != "" {
return m, userLevel, nil
}
return m, aliasLevel, nil
}
// 3. Check dynamic resolvers (copy slice under lock to avoid holding
// the lock while calling back — resolvers may access the registry).
r.mu.RLock()
resolvers := make([]Resolver, len(r.resolvers))
copy(resolvers, r.resolvers)
r.mu.RUnlock()
for _, res := range resolvers {
if resolved, defaultLevel, ok := res.Resolve(base); ok {
if resolved == "" {
return nil, "", fmt.Errorf("resolver returned empty spec for %q", base)
}
// A resolver may return a comma-separated failover chain.
if strings.Contains(resolved, ",") {
m, err := r.Parse(resolved)
if err != nil {
return nil, "", err
}
level := defaultLevel
if userLevel != "" {
level = userLevel
}
return m, level, nil
}
m, resolvedLevel, err := r.parse(resolved, depth+1)
if err != nil {
return nil, "", err
}
level := resolvedLevel
if defaultLevel != "" && level == "" {
level = defaultLevel
}
if userLevel != "" {
level = userLevel
}
return m, level, nil
}
}
// 4. Split on first "/".
slashIdx := strings.Index(base, "/")
if slashIdx < 0 {
return nil, "", fmt.Errorf("%w: %q is not an alias and has no provider/ prefix", ErrUnknownProvider, spec)
}
left := base[:slashIdx]
right := base[slashIdx+1:]
// 5. Look up in provider registry.
if info := r.ProviderByName(left); info != nil {
client := r.createClient(info)
return client.Model(right), userLevel, nil
}
// 6. Check LLM_{UPPER(left)} env var.
envKey := "LLM_" + strings.ToUpper(strings.ReplaceAll(left, "-", "_"))
lookup := r.envLookup
if lookup == nil {
lookup = os.Getenv
}
envVal := lookup(envKey)
if envVal == "" {
return nil, "", fmt.Errorf("%w: %q (checked registry and %s env var)", ErrUnknownProvider, left, envKey)
}
dsn, err := ParseDSN(envVal)
if err != nil {
return nil, "", fmt.Errorf("parse %s: %w", envKey, err)
}
schemeInfo := r.ProviderByName(dsn.Scheme)
if schemeInfo == nil {
return nil, "", fmt.Errorf("%w: DSN scheme %q in %s is not a registered provider", ErrUnknownProvider, dsn.Scheme, envKey)
}
url := "https://" + dsn.Host
client := schemeInfo.New(dsn.Token, WithBaseURL(url))
return client.Model(right), userLevel, nil
}
// createClient builds a Client from a ProviderInfo, reading the API key from
// the environment when the provider specifies an EnvKey.
//
// Why: avoids duplicating env-var lookup logic across parse paths.
// What: reads the env var (if any), calls info.New with the key.
// Test: indirectly tested via Parse("openai/gpt-4o") with injected envLookup.
func (r *Registry) createClient(info *ProviderInfo) *Client {
apiKey := ""
if info.EnvKey != "" {
lookup := r.envLookup
if lookup == nil {
lookup = os.Getenv
}
apiKey = lookup(info.EnvKey)
}
return info.New(apiKey)
}
+723
View File
@@ -0,0 +1,723 @@
package llm
import (
"context"
"errors"
"testing"
"gitea.stevedudenhoeffer.com/steve/go-llm/v2/provider"
)
// recordingProvider captures the model name passed to Complete so tests can
// verify that Parse resolved to the correct model without network calls.
type recordingProvider struct {
lastModel string
// err, when non-nil, is returned from Complete so failover tests can drive
// a comma-Parse'd chain through a failover decision. Defaults to nil (success).
err error
}
func (p *recordingProvider) Complete(_ context.Context, req provider.Request) (provider.Response, error) {
p.lastModel = req.Model
if p.err != nil {
return provider.Response{}, p.err
}
return provider.Response{Text: "ok"}, nil
}
func (p *recordingProvider) Stream(_ context.Context, _ provider.Request, events chan<- provider.StreamEvent) error {
close(events)
return nil
}
// testRegistry builds a Registry with two mock providers ("alpha" and "beta")
// and an injectable envLookup. No real API keys or network access required.
func testRegistry(envFn func(string) string) (*Registry, *recordingProvider, *recordingProvider) {
alpha := &recordingProvider{}
beta := &recordingProvider{}
r := &Registry{
providers: map[string]ProviderInfo{
"alpha": {
Name: "alpha",
DisplayName: "Alpha",
EnvKey: "ALPHA_API_KEY",
Models: []string{"model-a"},
New: func(apiKey string, opts ...ClientOption) *Client {
return NewClient(alpha)
},
},
"beta": {
Name: "beta",
DisplayName: "Beta",
EnvKey: "",
Models: []string{"model-b"},
New: func(_ string, opts ...ClientOption) *Client {
return NewClient(beta)
},
},
},
order: []string{"alpha", "beta"},
aliases: make(map[string]string),
envLookup: envFn,
}
return r, alpha, beta
}
// ---------------------------------------------------------------------------
// splitReasoning
// ---------------------------------------------------------------------------
func TestSplitReasoning(t *testing.T) {
tests := []struct {
input string
wantBase string
wantLevel ReasoningLevel
}{
{"gpt-4o", "gpt-4o", ""},
{"gpt-4o:high", "gpt-4o", ReasoningHigh},
{"gpt-4o:low", "gpt-4o", ReasoningLow},
{"gpt-4o:medium", "gpt-4o", ReasoningMedium},
{"qwen3:30b", "qwen3:30b", ""}, // Ollama tag, not a level
{"qwen3:30b:high", "qwen3:30b", ReasoningHigh}, // tag + reasoning
{"model:", "model:", ""}, // trailing colon, empty suffix
{"", "", ""}, // empty string
{"a:b:c:high", "a:b:c", ReasoningHigh}, // multiple colons
{"a:b:c:30b", "a:b:c:30b", ""}, // multiple colons, non-level
}
for _, tt := range tests {
t.Run(tt.input, func(t *testing.T) {
base, level := splitReasoning(tt.input)
if base != tt.wantBase {
t.Errorf("splitReasoning(%q) base = %q, want %q", tt.input, base, tt.wantBase)
}
if level != tt.wantLevel {
t.Errorf("splitReasoning(%q) level = %q, want %q", tt.input, level, tt.wantLevel)
}
})
}
}
// ---------------------------------------------------------------------------
// ParseDSN
// ---------------------------------------------------------------------------
func TestParseDSN(t *testing.T) {
tests := []struct {
name string
raw string
want DSN
wantErr bool
}{
{
name: "full DSN with token",
raw: "foreman://test-token@foreman-m5.orgrimmar.dudenhoeffer.casa",
want: DSN{Scheme: "foreman", Token: "test-token", Host: "foreman-m5.orgrimmar.dudenhoeffer.casa"},
},
{
name: "no token",
raw: "ollama://localhost:11434",
want: DSN{Scheme: "ollama", Token: "", Host: "localhost:11434"},
},
{
name: "trailing slash stripped",
raw: "foreman://tok@host.com/",
want: DSN{Scheme: "foreman", Token: "tok", Host: "host.com"},
},
{
name: "with path",
raw: "foreman://tok@host.com/v1/api",
want: DSN{Scheme: "foreman", Token: "tok", Host: "host.com/v1/api"},
},
{
name: "missing scheme",
raw: "no-scheme-here",
wantErr: true,
},
{
name: "missing host",
raw: "foreman://token@",
wantErr: true,
},
{
name: "empty string",
raw: "",
wantErr: true,
},
{
name: "scheme only",
raw: "foreman://",
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := ParseDSN(tt.raw)
if tt.wantErr {
if err == nil {
t.Fatalf("ParseDSN(%q) expected error, got nil", tt.raw)
}
if !errors.Is(err, ErrInvalidDSN) {
t.Errorf("ParseDSN(%q) error = %v, want ErrInvalidDSN", tt.raw, err)
}
return
}
if err != nil {
t.Fatalf("ParseDSN(%q) unexpected error: %v", tt.raw, err)
}
if got != tt.want {
t.Errorf("ParseDSN(%q) = %+v, want %+v", tt.raw, got, tt.want)
}
})
}
}
// ---------------------------------------------------------------------------
// Registry.Parse
// ---------------------------------------------------------------------------
func TestRegistryParse(t *testing.T) {
t.Run("basic provider/model", func(t *testing.T) {
r, alpha, _ := testRegistry(nil)
m, err := r.Parse("alpha/model-a")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if m == nil {
t.Fatal("expected non-nil model")
}
// Verify the model name by exercising it.
_, _ = m.Complete(context.Background(), nil)
if alpha.lastModel != "model-a" {
t.Errorf("model = %q, want %q", alpha.lastModel, "model-a")
}
})
t.Run("provider/model with reasoning suffix", func(t *testing.T) {
r, alpha, _ := testRegistry(nil)
m, err := r.Parse("alpha/model-a:high")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if m.defaultReasoning != ReasoningHigh {
t.Errorf("reasoning = %q, want %q", m.defaultReasoning, ReasoningHigh)
}
_, _ = m.Complete(context.Background(), nil)
if alpha.lastModel != "model-a" {
t.Errorf("model = %q, want %q", alpha.lastModel, "model-a")
}
})
t.Run("ollama-style tag preserved", func(t *testing.T) {
r, _, beta := testRegistry(nil)
m, err := r.Parse("beta/qwen3:30b")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if m.defaultReasoning != "" {
t.Errorf("reasoning = %q, want empty (30b is not a level)", m.defaultReasoning)
}
_, _ = m.Complete(context.Background(), nil)
if beta.lastModel != "qwen3:30b" {
t.Errorf("model = %q, want %q", beta.lastModel, "qwen3:30b")
}
})
t.Run("ollama-style tag with reasoning", func(t *testing.T) {
r, _, beta := testRegistry(nil)
m, err := r.Parse("beta/qwen3:30b:high")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if m.defaultReasoning != ReasoningHigh {
t.Errorf("reasoning = %q, want %q", m.defaultReasoning, ReasoningHigh)
}
_, _ = m.Complete(context.Background(), nil)
if beta.lastModel != "qwen3:30b" {
t.Errorf("model = %q, want %q", beta.lastModel, "qwen3:30b")
}
})
t.Run("static alias", func(t *testing.T) {
r, alpha, _ := testRegistry(nil)
r.RegisterAlias("fast", "alpha/model-a")
m, err := r.Parse("fast")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
_, _ = m.Complete(context.Background(), nil)
if alpha.lastModel != "model-a" {
t.Errorf("model = %q, want %q", alpha.lastModel, "model-a")
}
})
t.Run("alias with embedded reasoning", func(t *testing.T) {
r, alpha, _ := testRegistry(nil)
r.RegisterAlias("thinking", "alpha/model-a:high")
m, err := r.Parse("thinking")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if m.defaultReasoning != ReasoningHigh {
t.Errorf("reasoning = %q, want %q", m.defaultReasoning, ReasoningHigh)
}
_, _ = m.Complete(context.Background(), nil)
if alpha.lastModel != "model-a" {
t.Errorf("model = %q, want %q", alpha.lastModel, "model-a")
}
})
t.Run("user reasoning overrides alias default", func(t *testing.T) {
r, _, _ := testRegistry(nil)
r.RegisterAlias("thinking", "alpha/model-a:high")
m, err := r.Parse("thinking:low")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if m.defaultReasoning != ReasoningLow {
t.Errorf("reasoning = %q, want %q (user override)", m.defaultReasoning, ReasoningLow)
}
})
t.Run("dynamic resolver", func(t *testing.T) {
r, alpha, _ := testRegistry(nil)
r.RegisterResolver(ResolverFunc(func(name string) (string, ReasoningLevel, bool) {
if name == "custom" {
return "alpha/resolved-model", "", true
}
return "", "", false
}))
m, err := r.Parse("custom")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
_, _ = m.Complete(context.Background(), nil)
if alpha.lastModel != "resolved-model" {
t.Errorf("model = %q, want %q", alpha.lastModel, "resolved-model")
}
})
t.Run("resolver with default reasoning", func(t *testing.T) {
r, _, _ := testRegistry(nil)
r.RegisterResolver(ResolverFunc(func(name string) (string, ReasoningLevel, bool) {
if name == "smart" {
return "alpha/model-a", ReasoningHigh, true
}
return "", "", false
}))
m, err := r.Parse("smart")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if m.defaultReasoning != ReasoningHigh {
t.Errorf("reasoning = %q, want %q", m.defaultReasoning, ReasoningHigh)
}
})
t.Run("resolver default reasoning overridden by user", func(t *testing.T) {
r, _, _ := testRegistry(nil)
r.RegisterResolver(ResolverFunc(func(name string) (string, ReasoningLevel, bool) {
if name == "smart" {
return "alpha/model-a", ReasoningHigh, true
}
return "", "", false
}))
m, err := r.Parse("smart:low")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if m.defaultReasoning != ReasoningLow {
t.Errorf("reasoning = %q, want %q (user override)", m.defaultReasoning, ReasoningLow)
}
})
t.Run("resolver returns empty spec", func(t *testing.T) {
r, _, _ := testRegistry(nil)
r.RegisterResolver(ResolverFunc(func(name string) (string, ReasoningLevel, bool) {
if name == "bad" {
return "", "", true
}
return "", "", false
}))
_, err := r.Parse("bad")
if err == nil {
t.Fatal("expected error for empty resolver spec")
}
})
t.Run("LLM_X env var with token", func(t *testing.T) {
envFn := func(key string) string {
if key == "LLM_M5" {
return "alpha://mytoken@myhost.com"
}
return ""
}
r, alpha, _ := testRegistry(envFn)
m, err := r.Parse("m5/qwen3:30b")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
_, _ = m.Complete(context.Background(), nil)
if alpha.lastModel != "qwen3:30b" {
t.Errorf("model = %q, want %q", alpha.lastModel, "qwen3:30b")
}
})
t.Run("LLM_X env var without token", func(t *testing.T) {
envFn := func(key string) string {
if key == "LLM_LOCAL" {
return "beta://localhost:11434"
}
return ""
}
r, _, beta := testRegistry(envFn)
m, err := r.Parse("local/llama3.2")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
_, _ = m.Complete(context.Background(), nil)
if beta.lastModel != "llama3.2" {
t.Errorf("model = %q, want %q", beta.lastModel, "llama3.2")
}
})
t.Run("LLM_X with hyphenated name", func(t *testing.T) {
envFn := func(key string) string {
if key == "LLM_MY_BOX" {
return "alpha://tok@host.com"
}
return ""
}
r, alpha, _ := testRegistry(envFn)
m, err := r.Parse("my-box/some-model")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
_, _ = m.Complete(context.Background(), nil)
if alpha.lastModel != "some-model" {
t.Errorf("model = %q, want %q", alpha.lastModel, "some-model")
}
})
t.Run("unknown provider no env var", func(t *testing.T) {
r, _, _ := testRegistry(func(string) string { return "" })
_, err := r.Parse("unknown/model")
if err == nil {
t.Fatal("expected error for unknown provider")
}
if !errors.Is(err, ErrUnknownProvider) {
t.Errorf("error = %v, want ErrUnknownProvider", err)
}
})
t.Run("bare unknown name no slash", func(t *testing.T) {
r, _, _ := testRegistry(nil)
_, err := r.Parse("nonexistent")
if err == nil {
t.Fatal("expected error for bare unknown name")
}
if !errors.Is(err, ErrUnknownProvider) {
t.Errorf("error = %v, want ErrUnknownProvider", err)
}
})
t.Run("alias loop", func(t *testing.T) {
r, _, _ := testRegistry(nil)
r.RegisterAlias("a", "b")
r.RegisterAlias("b", "a")
_, err := r.Parse("a")
if err == nil {
t.Fatal("expected error for alias loop")
}
if !errors.Is(err, ErrAliasLoop) {
t.Errorf("error = %v, want ErrAliasLoop", err)
}
})
t.Run("deep alias chain within limit", func(t *testing.T) {
r, alpha, _ := testRegistry(nil)
// chain: a0 → a1 → a2 → ... → a9 → alpha/model-a (depth 10 is fine, >10 is not)
r.RegisterAlias("a9", "alpha/model-a")
for i := 8; i >= 0; i-- {
r.RegisterAlias(
"a"+string(rune('0'+i)),
"a"+string(rune('0'+i+1)),
)
}
m, err := r.Parse("a0")
if err != nil {
t.Fatalf("unexpected error for deep but valid chain: %v", err)
}
_, _ = m.Complete(context.Background(), nil)
if alpha.lastModel != "model-a" {
t.Errorf("model = %q, want %q", alpha.lastModel, "model-a")
}
})
t.Run("RegisterProvider replaces existing", func(t *testing.T) {
replaced := &recordingProvider{}
r, _, _ := testRegistry(nil)
r.RegisterProvider(ProviderInfo{
Name: "alpha",
DisplayName: "Alpha Replaced",
New: func(_ string, _ ...ClientOption) *Client {
return NewClient(replaced)
},
})
m, err := r.Parse("alpha/new-model")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
_, _ = m.Complete(context.Background(), nil)
if replaced.lastModel != "new-model" {
t.Errorf("model = %q, want %q", replaced.lastModel, "new-model")
}
// Verify order is preserved (alpha is still at index 0).
providers := r.Providers()
if providers[0].Name != "alpha" {
t.Errorf("first provider = %q, want %q (order preserved)", providers[0].Name, "alpha")
}
if providers[0].DisplayName != "Alpha Replaced" {
t.Errorf("display name = %q, want %q", providers[0].DisplayName, "Alpha Replaced")
}
})
t.Run("RegisterProvider adds new provider", func(t *testing.T) {
gamma := &recordingProvider{}
r, _, _ := testRegistry(nil)
r.RegisterProvider(ProviderInfo{
Name: "gamma",
DisplayName: "Gamma",
New: func(_ string, _ ...ClientOption) *Client {
return NewClient(gamma)
},
})
m, err := r.Parse("gamma/g-model")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
_, _ = m.Complete(context.Background(), nil)
if gamma.lastModel != "g-model" {
t.Errorf("model = %q, want %q", gamma.lastModel, "g-model")
}
// New provider should be last.
providers := r.Providers()
last := providers[len(providers)-1]
if last.Name != "gamma" {
t.Errorf("last provider = %q, want %q", last.Name, "gamma")
}
})
t.Run("DSN with unknown scheme", func(t *testing.T) {
envFn := func(key string) string {
if key == "LLM_X" {
return "nope://tok@host.com"
}
return ""
}
r, _, _ := testRegistry(envFn)
_, err := r.Parse("x/model")
if err == nil {
t.Fatal("expected error for unknown DSN scheme")
}
if !errors.Is(err, ErrUnknownProvider) {
t.Errorf("error = %v, want ErrUnknownProvider", err)
}
})
t.Run("DSN with invalid format", func(t *testing.T) {
envFn := func(key string) string {
if key == "LLM_BAD" {
return "no-scheme"
}
return ""
}
r, _, _ := testRegistry(envFn)
_, err := r.Parse("bad/model")
if err == nil {
t.Fatal("expected error for invalid DSN")
}
if !errors.Is(err, ErrInvalidDSN) {
t.Errorf("error = %v, want ErrInvalidDSN", err)
}
})
t.Run("createClient reads env key", func(t *testing.T) {
var capturedKey string
r, _, _ := testRegistry(func(key string) string {
if key == "MY_KEY" {
return "secret-123"
}
return ""
})
rec := &recordingProvider{}
r.RegisterProvider(ProviderInfo{
Name: "keyed",
EnvKey: "MY_KEY",
New: func(apiKey string, opts ...ClientOption) *Client {
capturedKey = apiKey
return NewClient(rec)
},
})
_, err := r.Parse("keyed/m")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if capturedKey != "secret-123" {
t.Errorf("apiKey = %q, want %q", capturedKey, "secret-123")
}
})
}
// ---------------------------------------------------------------------------
// Backward compatibility: package-level Providers() and ProviderByName()
// ---------------------------------------------------------------------------
func TestBackwardCompat(t *testing.T) {
t.Run("Providers includes foreman", func(t *testing.T) {
providers := Providers()
found := false
for _, p := range providers {
if p.Name == "foreman" {
found = true
break
}
}
if !found {
t.Error("Providers() does not include foreman")
}
})
t.Run("ProviderByName openai", func(t *testing.T) {
p := ProviderByName("openai")
if p == nil {
t.Fatal("ProviderByName(\"openai\") returned nil")
}
if p.DisplayName != "OpenAI" {
t.Errorf("DisplayName = %q, want %q", p.DisplayName, "OpenAI")
}
})
t.Run("ProviderByName foreman", func(t *testing.T) {
p := ProviderByName("foreman")
if p == nil {
t.Fatal("ProviderByName(\"foreman\") returned nil")
}
if p.DisplayName != "Foreman" {
t.Errorf("DisplayName = %q, want %q", p.DisplayName, "Foreman")
}
})
t.Run("ProviderByName unknown", func(t *testing.T) {
p := ProviderByName("does-not-exist")
if p != nil {
t.Errorf("expected nil for unknown provider, got %+v", p)
}
})
t.Run("Providers count", func(t *testing.T) {
providers := Providers()
// providerRegistry has 9 entries (openai, anthropic, google, deepseek,
// moonshot, xai, groq, ollama, ollama-cloud) + foreman = 10
if len(providers) < 10 {
t.Errorf("len(Providers()) = %d, want >= 10", len(providers))
}
})
}
// ---------------------------------------------------------------------------
// NewRegistry isolation
// ---------------------------------------------------------------------------
func TestNewRegistryIsolation(t *testing.T) {
r1 := NewRegistry()
r2 := NewRegistry()
r1.RegisterAlias("only-in-r1", "openai/gpt-4o")
// r2 should not see r1's alias
r2.mu.RLock()
_, found := r2.aliases["only-in-r1"]
r2.mu.RUnlock()
if found {
t.Error("alias registered in r1 should not appear in r2")
}
}
// ---------------------------------------------------------------------------
// Comma-separated failover chains
// ---------------------------------------------------------------------------
func TestParse_CommaProducesFailover(t *testing.T) {
resetHealthForTest()
r, alpha, beta := testRegistry(func(string) string { return "" })
m, err := r.Parse("alpha/model-a,beta/model-b")
if err != nil {
t.Fatalf("Parse failover spec: %v", err)
}
fp, ok := m.provider.(*failoverProvider)
if !ok {
t.Fatalf("expected *failoverProvider, got %T", m.provider)
}
if len(fp.entries) != 2 {
t.Fatalf("expected 2 entries, got %d", len(fp.entries))
}
if fp.entries[0].specKey != "alpha/model-a" || fp.entries[1].specKey != "beta/model-b" {
t.Errorf("unexpected specKeys: %q, %q", fp.entries[0].specKey, fp.entries[1].specKey)
}
// Complete routes to the first provider and passes the bare model name.
_, err = m.Complete(context.Background(), []Message{{Role: RoleUser, Content: Content{Text: "hi"}}})
if err != nil {
t.Fatal(err)
}
if alpha.lastModel != "model-a" {
t.Errorf("alpha got model %q, want model-a", alpha.lastModel)
}
if beta.lastModel != "" {
t.Errorf("beta should not have been called, got model %q", beta.lastModel)
}
}
func TestParse_NoCommaUnchanged(t *testing.T) {
resetHealthForTest()
r, _, _ := testRegistry(func(string) string { return "" })
m, err := r.Parse("alpha/model-a")
if err != nil {
t.Fatal(err)
}
if _, ok := m.provider.(*failoverProvider); ok {
t.Error("single (comma-free) spec must NOT produce a failover provider")
}
}
func TestParse_CommaSinglePartFallsThrough(t *testing.T) {
resetHealthForTest()
r, _, _ := testRegistry(func(string) string { return "" })
// Trailing comma / whitespace collapses to a single real part.
m, err := r.Parse("alpha/model-a, ")
if err != nil {
t.Fatal(err)
}
if _, ok := m.provider.(*failoverProvider); ok {
t.Error("a single effective part must not produce a failover provider")
}
}
func TestParse_CommaFlattensNested(t *testing.T) {
resetHealthForTest()
r, _, _ := testRegistry(func(string) string { return "" })
// Register an alias that is itself a comma chain.
r.RegisterAlias("pair", "alpha/model-a,beta/model-b")
m, err := r.Parse("pair,beta/model-b2")
if err != nil {
t.Fatal(err)
}
fp, ok := m.provider.(*failoverProvider)
if !ok {
t.Fatalf("expected *failoverProvider, got %T", m.provider)
}
if len(fp.entries) != 3 {
t.Errorf("expected flattened 3 entries, got %d", len(fp.entries))
}
}
+179 -11
View File
@@ -1,6 +1,9 @@
package llm
import (
"os"
"sync"
"gitea.stevedudenhoeffer.com/steve/go-llm/v2/deepseek"
"gitea.stevedudenhoeffer.com/steve/go-llm/v2/groq"
"gitea.stevedudenhoeffer.com/steve/go-llm/v2/moonshot"
@@ -38,7 +41,7 @@ type ProviderInfo struct {
// providerRegistry is the in-process list of known providers. Order is
// intentional: the three original providers first, then OpenAI-compatible
// additions in the order they were added.
// additions in the order they were added. This slice seeds NewRegistry().
var providerRegistry = []ProviderInfo{
{
Name: "openai",
@@ -151,24 +154,189 @@ var providerRegistry = []ProviderInfo{
},
New: OllamaCloud,
},
{
Name: "foreman",
DisplayName: "Foreman",
EnvKey: "", // no single env key; discovered via LLM_* DSNs
DefaultURL: "", // always requires a URL
Models: []string{},
New: Foreman,
},
}
// ---------------------------------------------------------------------------
// Resolver — dynamic model-spec resolution
// ---------------------------------------------------------------------------
// Resolver resolves an alias or short name to a full spec string. Consumers
// register resolvers for dynamic lookups (e.g., database-backed tier aliases).
type Resolver interface {
// Resolve returns the resolved spec and an optional default reasoning
// level. ok is false when the resolver does not handle this name.
Resolve(name string) (spec string, defaultReasoning ReasoningLevel, ok bool)
}
// ResolverFunc adapts a plain function to the Resolver interface.
type ResolverFunc func(name string) (string, ReasoningLevel, bool)
// Resolve implements Resolver by calling the underlying function.
func (f ResolverFunc) Resolve(name string) (string, ReasoningLevel, bool) {
return f(name)
}
// ---------------------------------------------------------------------------
// Registry — extensible provider/alias/resolver store
// ---------------------------------------------------------------------------
// Registry holds providers, static aliases, and dynamic resolvers. Use
// NewRegistry to create one pre-populated with the built-in providers, or
// use the package-level DefaultRegistry.
type Registry struct {
mu sync.RWMutex
providers map[string]ProviderInfo
order []string // insertion order for Providers()
aliases map[string]string
resolvers []Resolver
envLookup func(string) string // defaults to os.Getenv
}
// NewRegistry creates a Registry pre-populated with all built-in providers
// (the same set returned by the providerRegistry package variable).
//
// Why: provides a fresh, isolated registry for testing or multi-tenant
// scenarios while reusing the canonical provider list.
// What: copies every entry from providerRegistry into a new Registry.
// Test: call NewRegistry(), verify Providers() length matches providerRegistry
// and ProviderByName("openai") is non-nil.
func NewRegistry() *Registry {
r := &Registry{
providers: make(map[string]ProviderInfo, len(providerRegistry)),
order: make([]string, 0, len(providerRegistry)),
aliases: make(map[string]string),
envLookup: os.Getenv,
}
for _, info := range providerRegistry {
r.providers[info.Name] = info
r.order = append(r.order, info.Name)
}
return r
}
// DefaultRegistry is the package-level registry used by the convenience
// functions Parse, Providers, ProviderByName, RegisterProvider, RegisterAlias,
// and RegisterResolver. Initialized in init() with all built-in providers.
var DefaultRegistry *Registry
func init() {
DefaultRegistry = NewRegistry()
}
// RegisterProvider adds or replaces a provider in the registry. When
// replacing, the provider keeps its original position in the ordered list.
//
// Why: allows consumers to override built-in factories (e.g., wrapping with
// middleware) or add entirely new providers at runtime.
// What: upserts info by Name into the provider map and order slice.
// Test: register a custom "openai" factory, verify ProviderByName returns it.
func (r *Registry) RegisterProvider(info ProviderInfo) {
r.mu.Lock()
defer r.mu.Unlock()
if _, exists := r.providers[info.Name]; !exists {
r.order = append(r.order, info.Name)
}
r.providers[info.Name] = info
}
// RegisterAlias maps a short name to a full spec string. The spec is resolved
// recursively by Parse, so an alias can point to another alias or to a
// "provider/model" string.
//
// Why: lets consumers define convenient shortcuts like "fast" → "openai/gpt-4o-mini".
// What: stores name→spec in the alias map.
// Test: register "fast" → "openai/gpt-4o-mini", parse "fast", verify model.
func (r *Registry) RegisterAlias(name, spec string) {
r.mu.Lock()
defer r.mu.Unlock()
r.aliases[name] = spec
}
// RegisterResolver appends a dynamic resolver. Resolvers are checked in
// registration order after static aliases. A resolver may return a spec
// string that is itself an alias or "provider/model" — it will be recursed.
//
// Why: supports dynamic alias sources (databases, remote config) without
// requiring static registration of every possible name.
// What: appends res to the resolver list.
// Test: register a ResolverFunc, parse a name it handles, verify resolution.
func (r *Registry) RegisterResolver(res Resolver) {
r.mu.Lock()
defer r.mu.Unlock()
r.resolvers = append(r.resolvers, res)
}
// ProviderByName returns the registered ProviderInfo with the given name, or
// nil if no such provider is registered. Name matching is exact.
//
// Why: callers need to look up provider metadata by name for factory calls,
// discovery, and DSN scheme resolution.
// What: returns a copy of the ProviderInfo or nil.
// Test: verify ProviderByName("openai") is non-nil, ProviderByName("nope") is nil.
func (r *Registry) ProviderByName(name string) *ProviderInfo {
r.mu.RLock()
defer r.mu.RUnlock()
if info, ok := r.providers[name]; ok {
return &info
}
return nil
}
// Providers returns a copy of all registered providers in insertion order.
//
// Why: CLI pickers and admin tools need the full list for display.
// What: returns a freshly allocated slice of ProviderInfo copies.
// Test: verify length matches expected count after registration.
func (r *Registry) Providers() []ProviderInfo {
r.mu.RLock()
defer r.mu.RUnlock()
out := make([]ProviderInfo, 0, len(r.order))
for _, name := range r.order {
if info, ok := r.providers[name]; ok {
out = append(out, info)
}
}
return out
}
// ---------------------------------------------------------------------------
// Package-level convenience functions — delegate to DefaultRegistry
// ---------------------------------------------------------------------------
// Providers returns a copy of the registered provider list so callers cannot
// mutate library state.
func Providers() []ProviderInfo {
out := make([]ProviderInfo, len(providerRegistry))
copy(out, providerRegistry)
return out
return DefaultRegistry.Providers()
}
// ProviderByName returns the registered ProviderInfo with the given name, or
// nil if no such provider is registered. Name matching is exact.
func ProviderByName(name string) *ProviderInfo {
for i := range providerRegistry {
if providerRegistry[i].Name == name {
p := providerRegistry[i]
return &p
}
}
return nil
return DefaultRegistry.ProviderByName(name)
}
// RegisterProvider adds or replaces a provider in the DefaultRegistry.
func RegisterProvider(info ProviderInfo) {
DefaultRegistry.RegisterProvider(info)
}
// RegisterAlias maps a short name to a full spec in the DefaultRegistry.
func RegisterAlias(name, spec string) {
DefaultRegistry.RegisterAlias(name, spec)
}
// RegisterResolver appends a dynamic resolver to the DefaultRegistry.
func RegisterResolver(res Resolver) {
DefaultRegistry.RegisterResolver(res)
}