Compare commits
10 Commits
8d4e99be1a
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| 1206261e6a | |||
| 361999550e | |||
| ae8e194fad | |||
| 67c3ebe067 | |||
| 6bac4cb3ed | |||
| a1f405de74 | |||
| 72e281f82b | |||
| db03e8d4ed | |||
| ddf1a9b556 | |||
| 4522310f5a |
@@ -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
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -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
@@ -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) }
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user