diff --git a/ingestion/cmd/server/main.go b/ingestion/cmd/server/main.go index b3e4e1a..7e15ae4 100644 --- a/ingestion/cmd/server/main.go +++ b/ingestion/cmd/server/main.go @@ -236,6 +236,11 @@ func main() { fallbackModel := envOr("BRAIN_LLM_FALLBACK_MODEL", "gemma4:31b") router.Fallback = llm.New(fallbackURL, "", fallbackModel, timeout) } + if tertiaryURL := os.Getenv("BRAIN_LLM_TERTIARY_URL"); tertiaryURL != "" { + tertiaryModel := envOr("BRAIN_LLM_TERTIARY_MODEL", "gemini-2.5-flash") + tertiaryKey := os.Getenv("GEMINI_API_KEY") + router.Tertiary = llm.New(tertiaryURL, tertiaryKey, tertiaryModel, timeout) + } answerComplete = router.Complete logger.Info("brain answer LLM configured", "primary", primaryURL, "model", primaryModel) } diff --git a/ingestion/internal/llm/router.go b/ingestion/internal/llm/router.go index ab61f0f..c3e8d75 100644 --- a/ingestion/internal/llm/router.go +++ b/ingestion/internal/llm/router.go @@ -5,14 +5,18 @@ import ( "fmt" ) -// Router calls Primary first; on any error falls back to Fallback. -// Fallback may be nil, in which case primary errors are returned directly. +// Router calls Primary first; on any error falls back to Fallback, then to +// Tertiary. Fallback and Tertiary may each be nil — a nil tier is skipped and +// the chain returns the errors accumulated so far. type Router struct { Primary *Client Fallback *Client + Tertiary *Client } -// Complete implements pipeline.CompleteFunc, routing through Primary then Fallback. +// Complete implements pipeline.CompleteFunc, routing through Primary, then +// Fallback, then Tertiary — trying each configured tier in order and +// returning the first success. func (r *Router) Complete(ctx context.Context, system, user string) (string, error) { out, err := r.Primary.Complete(ctx, system, user) if err == nil { @@ -22,8 +26,15 @@ func (r *Router) Complete(ctx context.Context, system, user string) (string, err return "", fmt.Errorf("primary llm: %w", err) } out, err2 := r.Fallback.Complete(ctx, system, user) - if err2 != nil { + if err2 == nil { + return out, nil + } + if r.Tertiary == nil { return "", fmt.Errorf("primary llm: %w; fallback llm: %v", err, err2) } + out, err3 := r.Tertiary.Complete(ctx, system, user) + if err3 != nil { + return "", fmt.Errorf("primary llm: %w; fallback llm: %v; tertiary llm: %v", err, err2, err3) + } return out, nil } diff --git a/ingestion/internal/llm/router_test.go b/ingestion/internal/llm/router_test.go index 04aaa19..4a75653 100644 --- a/ingestion/internal/llm/router_test.go +++ b/ingestion/internal/llm/router_test.go @@ -69,3 +69,78 @@ func TestRouter_NilFallback(t *testing.T) { _, err := r.Complete(context.Background(), "sys", "user") assert.Error(t, err) } + +func TestRouter_FallsBackToTertiaryWhenPrimaryAndFallbackFail(t *testing.T) { + fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "unavailable", http.StatusServiceUnavailable) + })) + defer fail.Close() + tertiary := mockServer(t, "from-tertiary") + defer tertiary.Close() + + r := &Router{ + Primary: New(fail.URL, "", "m", time.Second), + Fallback: New(fail.URL, "", "m", time.Second), + Tertiary: New(tertiary.URL, "", "m", time.Second), + } + out, err := r.Complete(context.Background(), "sys", "user") + require.NoError(t, err) + assert.Equal(t, "from-tertiary", out) +} + +func TestRouter_TertiaryNotCalledWhenFallbackSucceeds(t *testing.T) { + fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "unavailable", http.StatusServiceUnavailable) + })) + defer fail.Close() + fallback := mockServer(t, "from-fallback") + defer fallback.Close() + tertiary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + t.Error("tertiary must not be called when fallback succeeds") + })) + defer tertiary.Close() + + r := &Router{ + Primary: New(fail.URL, "", "m", time.Second), + Fallback: New(fallback.URL, "", "m", time.Second), + Tertiary: New(tertiary.URL, "", "m", time.Second), + } + out, err := r.Complete(context.Background(), "sys", "user") + require.NoError(t, err) + assert.Equal(t, "from-fallback", out) +} + +func TestRouter_AllThreeFail(t *testing.T) { + fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "err", http.StatusBadGateway) + })) + defer fail.Close() + + r := &Router{ + Primary: New(fail.URL, "", "m", time.Second), + Fallback: New(fail.URL, "", "m", time.Second), + Tertiary: New(fail.URL, "", "m", time.Second), + } + _, err := r.Complete(context.Background(), "sys", "user") + assert.Error(t, err) + assert.Contains(t, err.Error(), "primary llm") + assert.Contains(t, err.Error(), "fallback llm") + assert.Contains(t, err.Error(), "tertiary llm") +} + +func TestRouter_NilTertiary(t *testing.T) { + // Existing 2-tier deployments (no BRAIN_LLM_TERTIARY_URL) must keep working + // exactly as before — Tertiary is opt-in, not required. + fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "err", http.StatusBadGateway) + })) + defer fail.Close() + + r := &Router{ + Primary: New(fail.URL, "", "m", time.Second), + Fallback: New(fail.URL, "", "m", time.Second), + } + _, err := r.Complete(context.Background(), "sys", "user") + assert.Error(t, err) + assert.NotContains(t, err.Error(), "tertiary") +}