homelab-codex-ws/services/kb-query/app/fallback.py

"""Active embed fallback state machine -- module 5 phase 4 plan §2 decision 2 / §5
(docs/kb/modules/05-faza4-plan.md). One process-global circuit breaker (`SolCircuitBreaker`,
held in `app.state`, not a module global -- keeps tests isolated) shared by every `/search`
request.

State machine (plan §2 decision 2 table):
  1. Cached status still fresh (< CACHE_TTL_S old) -> use it, no network call at all.
  2. Cache expired -> probe `GET {solaria}/api/tags` with a short timeout (HEALTH_PROBE_TIMEOUT_S);
     cache the result (up or down) for CACHE_TTL_S. Same TTL for both outcomes -- plan explicitly
     defers a separate down-backoff until flapping is observed.
  3. status == "up" -> embed on SOLARIA with a hard per-call timeout (EMBED_TIMEOUT_S), distinct
     from the probe timeout: a slow-but-technically-reachable SOLARIA must not hang the user's
     query past this bound.
  3a. Success -> return, sol_status="up".
  3b. Timeout/error *on the real embed call* (not just the probe) -> flip the breaker to "down"
      immediately (one-shot switch) and fall through to step 4 **in the same request** -- the
      user never sees an error for this, only the first unlucky hit in a 30s window pays one
      extra EMBED_TIMEOUT_S before falling through.
  4. status == "down" -> embed locally via `piha_url`, sol_status="down".

Both branches call `embed_chunk(..., model=embed_model)` with the exact same `embed_model`
argument -- there is no separate "PIHA model" and no per-request model choice (plan §2: "nie ma
dziś ... per-request wyboru modelu"), so the hard model invariant (app/startup.py, enforced once
at process startup against document_chunk/document_summary) already covers both paths by
construction. Adding a second, redundant per-request DB check here would guard against a
scenario that cannot occur while `embed_model` stays a single constant threaded through one
function -- instead, `tests/test_fallback.py` proves the invariant holds by asserting both the
failed SOLARIA attempt and the successful PIHA attempt in the same request carry an identical
`model` field.
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Callable, Optional

import aiohttp

from kb_retrieval.embed import check_ollama_health, embed_chunk

CACHE_TTL_S = 30.0
HEALTH_PROBE_TIMEOUT_S = 0.5
EMBED_TIMEOUT_S = 3.0


@dataclass
class SolCircuitBreaker:
    """Process-global cached `sol_status` ("up"/"down"), one instance per running app (held in
    `app.state.sol_breaker`, created fresh in `lifespan` -- never a module-level singleton, so
    tests get a clean breaker per instance). `clock` is injectable (default `time.monotonic`,
    never wall-clock/datetime) purely to test TTL expiry without real sleeps."""

    cache_ttl_s: float = CACHE_TTL_S
    clock: Callable[[], float] = field(default=time.monotonic)
    _status: Optional[str] = field(default=None, init=False)
    _checked_at: Optional[float] = field(default=None, init=False)

    @property
    def status(self) -> Optional[str]:
        """Cached status if still within `cache_ttl_s`, else `None` (caller must (re)probe)."""
        if self._status is None or self._checked_at is None:
            return None
        if self.clock() - self._checked_at >= self.cache_ttl_s:
            return None
        return self._status

    def set(self, status: str) -> None:
        self._status = status
        self._checked_at = self.clock()


async def resolve_sol_status(
    breaker: SolCircuitBreaker, session: aiohttp.ClientSession, solaria_url: str
) -> str:
    """Cache-or-probe: returns "up"/"down". Shared by `/search` (embed_with_fallback) and
    `/healthz` (plan: "sol_status w stopce UI ... podepnij realny stan") so both endpoints
    agree on the same cached view instead of running independent probes."""
    cached = breaker.status
    if cached is not None:
        return cached
    up = await check_ollama_health(session, solaria_url, HEALTH_PROBE_TIMEOUT_S)
    status = "up" if up else "down"
    breaker.set(status)
    return status


async def embed_with_fallback(
    breaker: SolCircuitBreaker,
    session: aiohttp.ClientSession,
    solaria_url: str,
    piha_url: str,
    embed_model: str,
    query_text: str,
) -> tuple[list[float], str]:
    """query_text -> (embedding, sol_status_used). `sol_status_used` is "up" (SOLARIA answered
    the real embed call) or "down" (PIHA answered -- either the breaker already cached "down",
    or SOLARIA just failed mid-request and this call transparently fell through, plan §2 step
    3b's one-shot switch)."""
    status = await resolve_sol_status(breaker, session, solaria_url)

    if status == "up":
        try:
            embedding, _elapsed = await embed_chunk(
                session, solaria_url, embed_model, query_text, timeout_s=EMBED_TIMEOUT_S
            )
            return embedding, "up"
        except (aiohttp.ClientError, TimeoutError):
            # One-shot switch (plan §2 step 3b): don't fail this request, fall through to PIHA
            # below and remember "down" for the rest of the cache window.
            breaker.set("down")

    embedding, _elapsed = await embed_chunk(session, piha_url, embed_model, query_text)
    return embedding, "down"