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

"""`/search` core -- module 5 phase 4 (docs/kb/modules/05-faza4-plan.md §4). Kept decoupled
from FastAPI so it can be unit-tested with fake `conn`/`session` 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. Purely additive: does not change any field already covered by
the phase-4 gate's HTTP-equivalence check (plan §9).
"""
from __future__ import annotations

import aiohttp
import asyncpg

from kb_retrieval.retrieval import cascade_query, flat_query, hybrid_query

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


async def run_search(
    conn: asyncpg.Connection,
    session: aiohttp.ClientSession,
    ollama_url: str,
    query_text: str,
    mode: str,
    embed_model: str,
    summary_model: str,
) -> dict:
    if mode == "flat":
        retrieval = await flat_query(conn, session, ollama_url, query_text, embed_model=embed_model)
    elif mode == "hybrid":
        retrieval = await hybrid_query(
            conn, session, ollama_url, query_text,
            summary_model=summary_model, embed_model=embed_model,
        )
    else:
        retrieval = await cascade_query(
            conn, session, ollama_url, query_text,
            summary_model=summary_model, embed_model=embed_model,
        )

    chunks = retrieval["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,
        "sol_status": "up",  # reaching this point means the embed call above succeeded
        "results": results,
    }