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

"""kb-postgres access for kb-query beyond what `kb_retrieval.retrieval` covers -- the envelope
join (`source` + `entities` for link/metadata building, app/links.py) is out of scope for the
shared retrieval package (module 5 phase 4 plan §4: "to jest nowy kod w warstwie
HTTP-handlera kb-query, nie zmiana w kb_retrieval")."""
from __future__ import annotations

import json

import asyncpg


async def create_pool(dsn: str) -> asyncpg.Pool:
    return await asyncpg.create_pool(dsn, min_size=1, max_size=5)


def _decode_jsonb(value: object) -> object:
    """asyncpg may return jsonb as a str or an already-decoded object."""
    if value is None:
        return None
    if isinstance(value, str):
        return json.loads(value)
    return value


async def fetch_envelopes(conn: asyncpg.Connection, envelope_ids: list[str]) -> dict[str, dict]:
    """Batch-fetch `source`/`entities` for a set of chunk hits, keyed by envelope id."""
    if not envelope_ids:
        return {}
    rows = await conn.fetch(
        "SELECT id, source, entities FROM envelope WHERE id = ANY($1::text[])",
        envelope_ids,
    )
    return {
        r["id"]: {"source": r["source"], "entities": _decode_jsonb(r["entities"]) or []}
        for r in rows
    }


async def fetch_summaries(
    conn: asyncpg.Connection, envelope_ids: list[str], summary_model: str
) -> dict[str, dict]:
    """Batch-fetch `document_summary.summary`/`tags` (the haiku compilation track) for the
    envelopes a search hit, keyed by envelope id -- frontend krok (plan §7): document summary
    as the per-envelope result header, chunks as its expandable fragments."""
    if not envelope_ids:
        return {}
    rows = await conn.fetch(
        "SELECT envelope_id, summary, tags FROM document_summary "
        "WHERE envelope_id = ANY($1::text[]) AND model = $2",
        envelope_ids, summary_model,
    )
    return {
        r["envelope_id"]: {"summary": r["summary"], "tags": _decode_jsonb(r["tags"]) or []}
        for r in rows
    }