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

"""`/search` core -- module 5 phase 4 (kb/phases/kb-m5-faza4.md §4). Kept decoupled
from FastAPI so it can be unit-tested with fake `conn`/`session`/`router` objects, the same
style as `kb_retrieval`'s own tests, instead of needing a live DB/Ollama behind a TestClient.

Response shape (plan §4 exactly): `{"query", "mode", "sol_status", "results": [...]}`, each
result carrying `dist` un-filtered -- the 0.45/0.55 colour thresholds (plan §7) are a frontend
concern (Krok 4, out of this step's scope), never applied server-side.

`summary`/`summary_tags` (document_summary, haiku track) are an additive field added in Krok 4
for the frontend's per-envelope result header (plan §7) -- `None`/`[]` when the envelope has no
summary for `summary_model` yet.

Krok 2 (aktywny fallback, plan §5): the query embedding no longer happens inside
`kb_retrieval.cascade_query`/`flat_query`/`hybrid_query` -- it goes through
`app.embed_router.EmbedRouter` (SOLARIA primary -> PIHA fallback state machine) and the
resulting vector feeds the same `*_retrieve` functions those wrappers call. Two additive
consequences for the response: `sol_status` now reports the router's real world view ("down"
while serving from the fallback -- the frontend already renders that as "offline (fallback
embed)"), and `embed_backend` names which backend actually embedded THIS query (task spec:
needed to debug result quality per backend).
"""
from __future__ import annotations

import aiohttp
import asyncpg

from kb_retrieval.embed import _vector_literal  # same private-import convention as kb_retrieval.retrieval
from kb_retrieval.retrieval import cascade_retrieve, flat_retrieve, hybrid_retrieve

from app.db import fetch_envelopes, fetch_summaries
from app.embed_router import EmbedRouter
from app.links import build_result


async def run_search(
    conn: asyncpg.Connection,
    session: aiohttp.ClientSession,
    router: EmbedRouter,
    query_text: str,
    mode: str,
    summary_model: str,
) -> dict:
    embedding, backend = await router.embed(session, query_text)
    query_vector = _vector_literal(embedding)

    if mode == "flat":
        chunks = await flat_retrieve(conn, query_vector)
    elif mode == "hybrid":
        chunks = (await hybrid_retrieve(conn, query_vector, summary_model))["chunks"]
    else:
        chunks = (await cascade_retrieve(conn, query_vector, summary_model))["chunks"]

    envelope_ids = sorted({c["envelope_id"] for c in chunks})
    envelopes = await fetch_envelopes(conn, envelope_ids)
    summaries = await fetch_summaries(conn, envelope_ids, summary_model)

    results = []
    for chunk in chunks:
        result = build_result(chunk, envelopes.get(chunk["envelope_id"]))
        summary = summaries.get(chunk["envelope_id"])
        result["summary"] = summary["summary"] if summary else None
        result["summary_tags"] = summary["tags"] if summary else []
        results.append(result)

    return {
        "query": query_text,
        "mode": mode,
        # "up" iff the primary (SOLARIA) embedded this very query -- the fallback path
        # implies the router just observed the primary down (probe or mid-embed failure).
        "sol_status": "up" if backend == router.primary.name else "down",
        "embed_backend": backend,
        "results": results,
    }