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

"""Embed-backend router -- module 5 phase 4, Krok 2 (docs/kb/modules/05-faza4-plan.md §2
Decyzja 2, §5): the active SOLARIA->PIHA fallback state machine `/search` embeds through.
One instance per process (`app.state.embed_router`), holding the single global `sol_status`
cache the plan describes:

| state           | behaviour                                                            |
|-----------------|----------------------------------------------------------------------|
| cache valid     | use remembered up/down, no probe                                     |
| cache expired   | GET {primary}/api/tags, short timeout -> up/down, cached for the TTL |
| up -> embed     | POST {primary}/api/embeddings under a hard timeout                   |
| embed fails     | one-shot flip (plan step 3b): mark down for one TTL window and serve |
|                 | THE SAME query from the fallback -- the user never sees a SOLARIA    |
|                 | error unless there is no fallback configured                         |
| down -> embed   | straight to the fallback, no primary attempt until the TTL expires   |

Both legs call `kb_retrieval.embed.embed_chunk(..., model=embed_model)` with the same
configured model (the plan's "twardy inwariant"). On top of the DB-side startup check
(app/startup.py pins EMBED_MODEL to `document_chunk.model`/`document_summary.embedding_model`),
each backend is additionally verified ONCE per process lifetime, lazily at its first use:
its `/api/tags` must list `embed_model`. A backend that answers but lacks bge-m3 raises
`ModelMismatchError` (ERROR log, 500 upstream) -- never a silent embed in a vector space
that doesn't match the index. Verification is lazy rather than at startup because either
backend may legitimately be offline when kb-query boots (SOLARIA sleeps ~16 h/day, plan
§1) and must not block the service.

Timeouts: probe 1.5 s (task spec "krótki, np. 1-2 s"; the plan's 500 ms proposal was for
LAN-adjacent Tailscale -- 1.5 s absorbs LTE-grade jitter without hurting the 30 s cache
economics), primary embed 3 s hard (plan step 3). The fallback embed gets NO hard timeout:
Pi-5 CPU embed plus a cold model load (OLLAMA_KEEP_ALIVE=0 on ollama-piha) is legitimately
slow (single seconds to low tens) and it is the last resort -- cutting it off would turn
"slow but alive" back into "dead".
"""
from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field

import aiohttp

from kb_retrieval.embed import check_ollama_health, embed_chunk

logger = logging.getLogger("kb-query.embed")


class EmbedBackendError(RuntimeError):
    """No configured backend could embed the query (primary down and the fallback is
    unconfigured or unreachable). Maps to 503 in the HTTP layer."""


class ModelMismatchError(RuntimeError):
    """A backend is reachable but its /api/tags does not list the configured embed model.
    Config/ops error (wrong URL, model not pulled) -- deliberately loud, maps to 500."""


@dataclass
class _Backend:
    name: str
    url: str
    model_verified: bool = field(default=False)


class EmbedRouter:
    def __init__(
        self,
        primary_url: str,
        fallback_url: str | None = None,
        *,
        embed_model: str,
        primary_name: str = "solaria",
        fallback_name: str = "piha",
        health_ttl_s: float = 30.0,
        health_timeout_s: float = 1.5,
        primary_embed_timeout_s: float = 3.0,
        clock=time.monotonic,
    ):
        self.primary = _Backend(primary_name, primary_url.rstrip("/"))
        self.fallback = _Backend(fallback_name, fallback_url.rstrip("/")) if fallback_url else None
        self.embed_model = embed_model
        self.health_ttl_s = health_ttl_s
        self.health_timeout_s = health_timeout_s
        self.primary_embed_timeout_s = primary_embed_timeout_s
        self._clock = clock
        self._sol_status: str | None = None  # "up" | "down" | None (never probed yet)
        self._checked_at: float = 0.0

    def _cache_valid(self) -> bool:
        return self._sol_status is not None and (self._clock() - self._checked_at) < self.health_ttl_s

    def _set_status(self, status: str) -> None:
        self._sol_status = status
        self._checked_at = self._clock()

    async def primary_status(self, session: aiohttp.ClientSession) -> str:
        """Cached primary liveness ("up"/"down"), probing /api/tags only when the TTL expired.
        Shared by /search routing and /healthz so both report the same world view."""
        if not self._cache_valid():
            up = await check_ollama_health(session, self.primary.url, self.health_timeout_s)
            self._set_status("up" if up else "down")
        return self._sol_status

    async def fallback_status(self, session: aiohttp.ClientSession) -> str:
        """Live fallback liveness for /healthz: "up"/"down"/"unconfigured". Uncached -- the
        probe is a cheap /api/tags (does not load the model on ollama-piha) and monitoring
        calls this every 30 s anyway."""
        if self.fallback is None:
            return "unconfigured"
        up = await check_ollama_health(session, self.fallback.url, self.health_timeout_s)
        return "up" if up else "down"

    async def _verify_model(self, session: aiohttp.ClientSession, backend: _Backend) -> None:
        """First-use guard: `/api/tags` of `backend` must list `embed_model` (exact name or
        same base, so a pulled "bge-m3:latest" satisfies EMBED_MODEL=bge-m3). Raises
        ModelMismatchError when reachable-but-missing; connection errors propagate to the
        caller's failover handling. Passed once -> never re-checked for this process."""
        if backend.model_verified:
            return
        async with session.get(
            f"{backend.url}/api/tags",
            timeout=aiohttp.ClientTimeout(total=self.health_timeout_s),
        ) as resp:
            resp.raise_for_status()
            data = await resp.json()
        names = {m.get("name", "") for m in data.get("models", [])}
        base = self.embed_model.split(":")[0]
        if not any(n == self.embed_model or n.split(":")[0] == base for n in names):
            raise ModelMismatchError(
                f"embed backend {backend.name} ({backend.url}) does not serve "
                f"model {self.embed_model!r} (available: {sorted(names) or 'none'}) -- "
                f"refusing to embed in a vector space that doesn't match document_chunk"
            )
        backend.model_verified = True

    async def embed(self, session: aiohttp.ClientSession, text: str) -> tuple[list[float], str]:
        """Embed one query text, returning (embedding, backend_name). Implements the state
        machine from the module docstring; raises EmbedBackendError (503) when no backend
        can serve, ModelMismatchError (500) when the fallback serves the wrong model."""
        if await self.primary_status(session) == "up":
            try:
                await self._verify_model(session, self.primary)
                embedding, elapsed = await asyncio.wait_for(
                    embed_chunk(session, self.primary.url, self.embed_model, text),
                    timeout=self.primary_embed_timeout_s,
                )
                logger.info(
                    "query embedded backend=%s url=%s elapsed_ms=%.0f",
                    self.primary.name, self.primary.url, elapsed * 1000,
                )
                return embedding, self.primary.name
            except ModelMismatchError as exc:
                # Reachable but serving the wrong model set -- loud, then try the fallback
                # (which is verified independently below, so this cannot cascade silently).
                logger.error("primary embed backend rejected: %s", exc)
                self._set_status("down")
            except (aiohttp.ClientError, TimeoutError, ValueError) as exc:
                # Plan step 3b: one-shot flip -- the SAME query continues to the fallback.
                logger.warning(
                    "primary embed failed backend=%s url=%s (%s: %s); circuit open for %.0fs",
                    self.primary.name, self.primary.url, type(exc).__name__, exc, self.health_ttl_s,
                )
                self._set_status("down")

        if self.fallback is None:
            raise EmbedBackendError(
                f"primary embed backend {self.primary.name} ({self.primary.url}) is down "
                f"and no EMBED_FALLBACK_URL is configured"
            )

        try:
            await self._verify_model(session, self.fallback)
        except ModelMismatchError:
            raise  # loud config error -- never silently embed in a mismatched space
        except (aiohttp.ClientError, TimeoutError) as exc:
            raise EmbedBackendError(
                f"primary embed backend {self.primary.name} is down and fallback "
                f"{self.fallback.name} ({self.fallback.url}) is unreachable: {exc}"
            ) from exc

        embedding, elapsed = await embed_chunk(session, self.fallback.url, self.embed_model, text)
        logger.info(
            "query embedded backend=%s url=%s elapsed_ms=%.0f (fallback; primary=%s down)",
            self.fallback.name, self.fallback.url, elapsed * 1000, self.primary.name,
        )
        return embedding, self.fallback.name