homelab-codex-ws/packages/kb-retrieval/src/kb_retrieval/retrieval.py

"""Retrieval module -- module 5, phase 3, plan step 4 (docs/kb/modules/05-faza3-plan.md, §6).
Moved 1:1 into `packages/kb-retrieval` in phase 4 (docs/kb/modules/05-faza4-plan.md, §3,
decision 1) so both `documents-ingest` (venv job) and `kb-query` (Docker service) share one
tested module instead of the service image needing to pull in all of `jobs/documents-ingest`.

Three retrieval paths over the same corpus, sharing one query embedding (bge-m3, via Ollama):

- `flat_retrieve`: baseline -- ranks every active `document_chunk` row directly. This
  formalizes the pilot's ad hoc `/tmp/kbq.sh` query (docs/kb/eval/retrieval-pilot-2026-07-16.md)
  into a tested, versioned module instead of a script living only in a session transcript.
- `cascade_retrieve`: pre-filters to the top-N `document_summary` matches for one configured
  `model` (plan §2 decision 3, resolved 2026-07-17 as D3: `claude-haiku-4-5` is the compilation
  track; `gemma3:12b` stays in reserve as the local track) before ranking `document_chunk`
  within just those envelopes. At the 186-document pilot scale this doesn't speed anything up
  -- it is an architecture test for the mail-scale corpus (225k envelopes, plan §1.1) where a
  flat chunk scan stops being cheap. `eval/retrieval_eval.py` runs the quality gate (plan §6.2)
  that decides whether it becomes the default path.
- `hybrid_retrieve` -- module 5, faza mailowa, plan Krok 3 (docs/kb/modules/05-faza-mailowa-plan.md,
  §6, decision 6): mail (gmail) envelopes never get a `document_summary` (streszczenie maila
  would usually be longer than the mail itself -- decision 6's rejected-summaries reasoning), so
  they are invisible to the cascade's stage 1 pre-filter. Hybrid runs the cascade for sources that
  DO have summaries (paperless) and, in parallel, a direct flat-style HNSW scan restricted to
  `summaryless_sources` (gmail), then merges both result sets by `dist` (same embedder, same
  cosine space -- merge is a plain sort, no re-normalization needed). This is the only path that
  can surface mail content until/unless mail gets its own summaries later (decision 6, deferred).

`cascade_query` / `flat_query` / `hybrid_query` are the intended clean entry points for kb-query:
query_text -> chunk hits with `dist` (cosine distance) and `source` ("flat" | "cascade" |
"hybrid"). `envelope`, `document_chunk`, and `document_summary` are read-only here -- this module
only ever `SELECT`s.
"""
from __future__ import annotations

import aiohttp
import asyncpg

from kb_retrieval.embed import _vector_literal, embed_chunk

DEFAULT_SUMMARY_MODEL = "claude-haiku-4-5"  # plan §2 D3 resolution 2026-07-17: compilation track
DEFAULT_EMBED_MODEL = "bge-m3"
DEFAULT_N = 10  # plan §6.1 start value
DEFAULT_K = 5  # plan §6.1 start value
DEFAULT_SUMMARYLESS_SOURCES = ("gmail",)  # faza mailowa plan §6 Krok 3: sources with no summaries


async def flat_retrieve(conn: asyncpg.Connection, query_vector: str, k: int = DEFAULT_K) -> list[dict]:
    """Baseline: rank every active chunk directly against the query embedding, no pre-filter."""
    rows = await conn.fetch(
        "SELECT envelope_id, chunk_index, text, embedding <=> $1::vector AS dist "
        "FROM document_chunk WHERE excluded_reason IS NULL AND embedding IS NOT NULL "
        "ORDER BY embedding <=> $1::vector LIMIT $2",
        query_vector, k,
    )
    return [
        {
            "envelope_id": r["envelope_id"],
            "chunk_index": r["chunk_index"],
            "text": r["text"],
            "dist": r["dist"],
            "source": "flat",
        }
        for r in rows
    ]


async def cascade_retrieve(
    conn: asyncpg.Connection,
    query_vector: str,
    summary_model: str = DEFAULT_SUMMARY_MODEL,
    n: int = DEFAULT_N,
    k: int = DEFAULT_K,
) -> dict:
    """Stage 1: top-N `document_summary` envelopes for `summary_model`. Stage 2: top-k
    `document_chunk` ranked within just those envelopes.

    `n` exceeding the number of summarized envelopes is not an error -- plain SQL `LIMIT`
    semantics just return all of them. An empty stage 1 (no summaries for this model, or a
    corpus that hasn't been summarized yet) short-circuits before stage 2 runs at all: a
    cascade can never rank chunks in envelopes it didn't pre-filter into, so there is nothing
    for stage 2 to query.
    """
    stage1 = await conn.fetch(
        "SELECT envelope_id, embedding <=> $1::vector AS dist FROM document_summary "
        "WHERE model = $2 AND embedding IS NOT NULL ORDER BY embedding <=> $1::vector LIMIT $3",
        query_vector, summary_model, n,
    )
    stage1_summaries = [{"envelope_id": r["envelope_id"], "dist": r["dist"]} for r in stage1]

    envelope_ids = [s["envelope_id"] for s in stage1_summaries]
    if not envelope_ids:
        return {"stage1_summaries": stage1_summaries, "chunks": []}

    stage2 = await conn.fetch(
        "SELECT envelope_id, chunk_index, text, embedding <=> $1::vector AS dist "
        "FROM document_chunk WHERE envelope_id = ANY($2::text[]) AND excluded_reason IS NULL "
        "AND embedding IS NOT NULL ORDER BY embedding <=> $1::vector LIMIT $3",
        query_vector, envelope_ids, k,
    )
    chunks = [
        {
            "envelope_id": r["envelope_id"],
            "chunk_index": r["chunk_index"],
            "text": r["text"],
            "dist": r["dist"],
            "source": "cascade",
        }
        for r in stage2
    ]
    return {"stage1_summaries": stage1_summaries, "chunks": chunks}


async def hybrid_retrieve(
    conn: asyncpg.Connection,
    query_vector: str,
    summary_model: str = DEFAULT_SUMMARY_MODEL,
    summaryless_sources: tuple[str, ...] = DEFAULT_SUMMARYLESS_SOURCES,
    n: int = DEFAULT_N,
    k: int = DEFAULT_K,
) -> dict:
    """Cascade (for summarized sources) merged with a direct chunk scan over
    `summaryless_sources` (plan §6 decision 6) -- mail envelopes never reach the cascade's
    stage 1, so this is the only path today that can surface them.

    Both branches share the same query embedding/cosine space, so the merge is a plain sort by
    `dist`, then truncate to top-k -- no re-normalization needed (plan §6 decision 6 rationale).
    Each returned chunk's `source` is overwritten to `"hybrid"` regardless of which branch it
    came from, matching `flat`/`cascade`'s convention that `source` names the retrieval mode,
    not an internal sub-path.
    """
    cascade = await cascade_retrieve(conn, query_vector, summary_model, n=n, k=k)

    mail_rows = await conn.fetch(
        "SELECT c.envelope_id, c.chunk_index, c.text, c.embedding <=> $1::vector AS dist "
        "FROM document_chunk c JOIN envelope e ON e.id = c.envelope_id "
        "WHERE e.source = ANY($2::text[]) AND c.excluded_reason IS NULL "
        "AND c.embedding IS NOT NULL ORDER BY c.embedding <=> $1::vector LIMIT $3",
        query_vector, list(summaryless_sources), k,
    )
    mail_chunks = [
        {
            "envelope_id": r["envelope_id"],
            "chunk_index": r["chunk_index"],
            "text": r["text"],
            "dist": r["dist"],
            "source": "hybrid",
        }
        for r in mail_rows
    ]

    merged = sorted(
        [{**c, "source": "hybrid"} for c in cascade["chunks"]] + mail_chunks,
        key=lambda c: c["dist"],
    )[:k]
    return {"stage1_summaries": cascade["stage1_summaries"], "chunks": merged}


async def flat_query(
    conn: asyncpg.Connection,
    session: aiohttp.ClientSession,
    ollama_url: str,
    query_text: str,
    embed_model: str = DEFAULT_EMBED_MODEL,
    k: int = DEFAULT_K,
) -> dict:
    """query_text -> flat chunk hits. One Ollama embed call, one SQL query."""
    embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text)
    chunks = await flat_retrieve(conn, _vector_literal(embedding), k=k)
    return {"query": query_text, "k": k, "chunks": chunks}


async def cascade_query(
    conn: asyncpg.Connection,
    session: aiohttp.ClientSession,
    ollama_url: str,
    query_text: str,
    summary_model: str = DEFAULT_SUMMARY_MODEL,
    embed_model: str = DEFAULT_EMBED_MODEL,
    n: int = DEFAULT_N,
    k: int = DEFAULT_K,
) -> dict:
    """query_text -> cascade chunk hits, plus the stage-1 envelope pre-filter (needed for the
    gate's diagnosis order, plan §6.2: "N too small" is diagnosed by looking at stage 1).

    Same query embedding as `flat_query`'s single Ollama call -- the cascade's only added
    cost over the flat path is one extra SQL query (stage 1), never an extra embed.
    """
    embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text)
    result = await cascade_retrieve(conn, _vector_literal(embedding), summary_model, n=n, k=k)
    result["query"] = query_text
    result["n"] = n
    result["k"] = k
    return result


async def hybrid_query(
    conn: asyncpg.Connection,
    session: aiohttp.ClientSession,
    ollama_url: str,
    query_text: str,
    summary_model: str = DEFAULT_SUMMARY_MODEL,
    embed_model: str = DEFAULT_EMBED_MODEL,
    summaryless_sources: tuple[str, ...] = DEFAULT_SUMMARYLESS_SOURCES,
    n: int = DEFAULT_N,
    k: int = DEFAULT_K,
) -> dict:
    """query_text -> hybrid chunk hits (cascade + direct mail scan, merged by dist).

    Same single query embedding as `flat_query`/`cascade_query` -- hybrid's only added cost
    over cascade is one extra SQL query (the mail branch), never an extra Ollama call.
    """
    embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text)
    result = await hybrid_retrieve(
        conn, _vector_literal(embedding), summary_model, summaryless_sources, n=n, k=k
    )
    result["query"] = query_text
    result["n"] = n
    result["k"] = k
    return result