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

"""Ollama embedding client -- module 5, phase 4, plan step 0 (kb/phases/kb-m5-faza4.md,
§3, decision 1). Moved 1:1 out of `jobs/documents_ingest/chunk_embed.py` (`embed_chunk`,
`_vector_literal`, `DEFAULT_MODEL`, `DEFAULT_OLLAMA_URL`) so both a venv-based host job
(`documents-ingest`) and a long-lived Docker service (`kb-query`) can depend on the same tested
client without the service image pulling in `jobs/`'s `anthropic` dependency and CLI scripts.

`check_ollama_health` is new here (not moved) -- a generic `/api/tags` probe kb-query's
health-check + circuit breaker (plan §2 decision 2) builds on, kept in this module so the HTTP
probing logic lives in exactly one place.

`embed_batch` -- module 5, faza mailowa, plan Krok 1 (kb/phases/kb-m5-faza-mailowa.md,
§4, decision 5): `/api/embed` with `input` as a list, batch 64 measured live at ~8-18 ms/chunk
vs ~150-200 ms/chunk sequential through `embed_chunk`'s `/api/embeddings` (plan §1.4). Used by
`jobs/mail-body-ingest` only -- `embed_chunk` stays the single-text path for kb-query (one query
= one text) and the paperless cyclic ingest.
"""
from __future__ import annotations

import time

import aiohttp

DEFAULT_OLLAMA_URL = "http://localhost:11434"
DEFAULT_MODEL = "bge-m3"
EXPECTED_DIM = 1024


class EmbeddingDimensionError(RuntimeError):
    """Ollama returned a vector of the wrong dimension for the target `document_chunk` schema.

    Abort-run class, mirroring `documents_ingest.chunk_embed.EmbeddingDimensionError` (a separate
    class in a separate module -- `embed_chunk`'s single-vector path still raises its own local
    one; this one is for `embed_batch` callers, currently only `jobs/mail-body-ingest`).
    """


def _vector_literal(embedding: list[float]) -> str:
    """Render a Python float list as a pgvector input literal, e.g. '[0.1,0.2,...]'."""
    return "[" + ",".join(repr(v) for v in embedding) + "]"


async def embed_chunk(
    session: aiohttp.ClientSession, base_url: str, model: str, text: str
) -> tuple[list[float], float]:
    """POST /api/embeddings on Ollama for one chunk. Returns (embedding, elapsed_seconds).
    No built-in timeout/retry -- inherits whatever `aiohttp.ClientSession(timeout=...)` the
    caller constructed; `raise_for_status()` propagates `aiohttp.ClientError` when Ollama is
    unreachable."""
    t0 = time.monotonic()
    async with session.post(f"{base_url}/api/embeddings", json={"model": model, "prompt": text}) as resp:
        resp.raise_for_status()
        data = await resp.json()
    elapsed = time.monotonic() - t0

    embedding = data.get("embedding")
    if not embedding:
        raise ValueError(f"ollama response missing 'embedding': {data!r}")
    return embedding, elapsed


async def embed_batch(
    session: aiohttp.ClientSession, base_url: str, model: str, texts: list[str]
) -> tuple[list[list[float]], float]:
    """POST /api/embed on Ollama with a batch of texts. Returns (embeddings, elapsed_seconds).

    Sequential batches (no HTTP parallelism) -- the GPU saturates on the batch itself (plan
    §1.4/§4 decision 5), so parallel requests would only add overhead and unpredictable error
    ordering in the caller's stats bilans.

    Raises `EmbeddingDimensionError` (abort-run) if the response's embedding count doesn't match
    `len(texts)`, or if any single embedding's dimension isn't `EXPECTED_DIM` -- never silently
    indexes a vector that doesn't match the `document_chunk.embedding VECTOR(1024)` column.
    Same `aiohttp.ClientError` propagation as `embed_chunk` -- no built-in timeout/retry.
    """
    t0 = time.monotonic()
    async with session.post(f"{base_url}/api/embed", json={"model": model, "input": texts}) as resp:
        resp.raise_for_status()
        data = await resp.json()
    elapsed = time.monotonic() - t0

    embeddings = data.get("embeddings")
    if embeddings is None:
        raise ValueError(f"ollama response missing 'embeddings': {data!r}")
    if len(embeddings) != len(texts):
        raise EmbeddingDimensionError(
            f"ollama returned {len(embeddings)} embeddings for {len(texts)} input texts"
        )
    for i, embedding in enumerate(embeddings):
        if len(embedding) != EXPECTED_DIM:
            raise EmbeddingDimensionError(
                f"ollama model={model!r} batch item {i} returned dim={len(embedding)}, "
                f"expected {EXPECTED_DIM} (document_chunk.embedding is VECTOR({EXPECTED_DIM}))"
            )
    return embeddings, elapsed


async def check_ollama_health(
    session: aiohttp.ClientSession, base_url: str, timeout_s: float
) -> bool:
    """Probe `GET {base_url}/api/tags`. True iff it answers within `timeout_s` with a non-error
    status; any exception (timeout, connection refused, 4xx/5xx) is treated as `down`, never
    raised -- callers use this for liveness decisions, not error propagation."""
    try:
        async with session.get(
            f"{base_url}/api/tags", timeout=aiohttp.ClientTimeout(total=timeout_s)
        ) as resp:
            return resp.status < 400
    except (aiohttp.ClientError, TimeoutError):
        return False