diff --git a/jobs/documents-ingest/pyproject.toml b/jobs/documents-ingest/pyproject.toml index 9cead27..b930283 100644 --- a/jobs/documents-ingest/pyproject.toml +++ b/jobs/documents-ingest/pyproject.toml @@ -11,12 +11,15 @@ dependencies = [ "structlog>=24.1", "aiohttp>=3.9", "kb-mail", + "PyYAML>=6.0", + "anthropic>=0.40", ] [project.scripts] documents-ingest = "documents_ingest.extractor:main" documents-ingest-paperless = "documents_ingest.paperless_adapter:main" documents-ingest-embed = "documents_ingest.chunk_embed:main" +documents-ingest-summarize = "documents_ingest.summarize:main" [tool.setuptools.packages.find] where = ["src"] diff --git a/jobs/documents-ingest/src/documents_ingest/summarize.py b/jobs/documents-ingest/src/documents_ingest/summarize.py new file mode 100644 index 0000000..f41530c --- /dev/null +++ b/jobs/documents-ingest/src/documents_ingest/summarize.py @@ -0,0 +1,675 @@ +"""Summarize + tag job — module 5 phase 3, plan step 3 (docs/kb/modules/05-faza3-plan.md, +§5 step 3, §2 decision 3: two-track A/B pilot). + +Pipeline: `document_chunk.text WHERE excluded_reason IS NULL ORDER BY chunk_index` per +`source='paperless'` envelope (duplicates via `entities[type=duplicate_of]` skipped whole) +-> LLM (`--backend ollama|anthropic`) forced-JSON `{"summary": ..., "tags": [...]}` +-> `INSERT document_summary` (`services/kb-postgres/init/004_summaries.sql`). `envelope` and +`document_chunk` are read-only here; this job only ever `INSERT`s into `document_summary`. + +Two backends write the same table under different `model` values — `UNIQUE (envelope_id, +model)` exists precisely so both A/B tracks coexist (plan §2 decision 3). A second, separate +mode (`--embed-summaries`) embeds existing summaries with bge-m3 via Ollama, reusing +`chunk_embed.embed_chunk` 1:1 — writing and embedding are split so either can be re-run +without redoing the other (SOLARIA asleep blocks embedding, not writing; API downtime blocks +writing, not embedding). + +Long documents (plan §1.2 outlier: ~360k chars) exceed practical LLM context: chunks are +grouped ~20-at-a-time into partial summaries (plain text, no JSON), then a final pass +synthesizes those partials with the normal JSON prompt. Counted separately +(`documents_mapreduce`), never silently skipped. + +Runs on PIHA (rsync `src/` + `tags-vocab.yaml` to `/tmp`, per plan §1.4) against +kb-postgres@PIHA and OLLAMA_URL=http://solaria:11434, or on SOLARIA directly for the local +backend. Anthropic backend needs `ANTHROPIC_API_KEY` in the environment for the run only — +never logged, never written to disk. + +Usage: + # Dry run (default) — fetch, chunk-concat, count; no LLM calls, no DB writes: + documents-ingest-summarize --dsn ... --backend ollama --model gemma3:12b + + # Real run, local track: + documents-ingest-summarize --dsn ... --backend ollama --model gemma3:12b --apply + + # Real run, API track (ANTHROPIC_API_KEY must be set): + documents-ingest-summarize --dsn ... --backend anthropic --model claude-haiku-4-5 --apply + + # Embed existing summaries (bge-m3, second pass): + documents-ingest-summarize --dsn ... --embed-summaries --embed-model bge-m3 --apply +""" +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import re +import sys +import unicodedata +from pathlib import Path +from typing import Optional, Protocol + +import aiohttp +import asyncpg +import structlog +import yaml +from anthropic import AsyncAnthropic + +from documents_ingest.chunk_embed import ( + EmbeddingDimensionError, + EXPECTED_DIM, + _vector_literal, + embed_chunk, +) + +_log = structlog.get_logger(__name__) + +DEFAULT_OLLAMA_URL = "http://localhost:11434" +DEFAULT_EMBED_MODEL = "bge-m3" +DEFAULT_MAX_TOKENS = 2048 +DEFAULT_CHUNKS_PER_GROUP = 20 + +# Ollama's *runtime* context window defaults to a few thousand tokens regardless of a +# model's advertised max (gemma3:12b advertises 131072) -- confirmed empirically during the +# pilot: a 93k-char OWU document was silently truncated to prompt_eval_count=2051 tokens with +# no num_ctx set, and the resulting summary described the document's tail (a RODO clause) +# instead of its actual insurance terms. `options.num_ctx` must be set explicitly on every +# call, sized to the prompt -- never left to Ollama's default. +OLLAMA_CHARS_PER_TOKEN = 3.0 # conservative for diacritic-heavy Polish text +OLLAMA_MIN_NUM_CTX = 4096 +OLLAMA_MAX_NUM_CTX = 131072 # gemma3:12b's advertised context_length +OLLAMA_CTX_RESPONSE_RESERVE = 4096 # headroom for template + generated output + + +def compute_num_ctx(prompt_chars: int) -> int: + """Size Ollama's num_ctx to the actual prompt, rounded up to a 1024-token step (stable + KV-cache allocation), clamped to [OLLAMA_MIN_NUM_CTX, OLLAMA_MAX_NUM_CTX].""" + needed = int(prompt_chars / OLLAMA_CHARS_PER_TOKEN) + OLLAMA_CTX_RESPONSE_RESERVE + needed = ((needed + 1023) // 1024) * 1024 + return max(OLLAMA_MIN_NUM_CTX, min(OLLAMA_MAX_NUM_CTX, needed)) +# Plan §1.2: the pilot corpus has one ~360k-char outlier; everything else is far smaller. +# This threshold only exists to catch that shape of document, not to tune routine chunking. +DEFAULT_MAPREDUCE_THRESHOLD_CHARS = 200_000 +DEFAULT_TAGS_VOCAB_PATH = Path(__file__).resolve().parent.parent.parent / "tags-vocab.yaml" + +SYSTEM_PROMPT = ( + "Jesteś archiwistą domowej bazy wiedzy. Streszczasz dokumenty po polsku, zwięźle i " + "faktograficznie: kto, co, kiedy, kwoty, numery umów/polis, terminy. Nie zgadujesz — " + "czego nie ma w tekście, tego nie piszesz. Zwracasz wyłącznie JSON." +) + +PARTIAL_SYSTEM_PROMPT = ( + "Jesteś archiwistą domowej bazy wiedzy. Streszczasz fragment dłuższego dokumentu po " + "polsku w 2-4 zdaniach, wypisując wszystkie fakty: kwoty, daty, numery, strony. To " + "streszczenie częściowe zostanie później połączone z innymi częściami tego samego " + "dokumentu — nie pomijaj konkretów. Odpowiadasz czystym tekstem, bez JSON." +) + +SUMMARY_JSON_SCHEMA = { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["summary", "tags"], + "additionalProperties": False, +} + +_INSERT_SQL = """ +INSERT INTO document_summary (envelope_id, summary, tags, model) +VALUES ($1, $2, $3::jsonb, $4) +ON CONFLICT (envelope_id, model) DO NOTHING +""" + +_TAG_CLEAN_RE = re.compile(r"[^\w-]+", re.UNICODE) +_TAG_DASH_COLLAPSE_RE = re.compile(r"-{2,}") + + +class LLMBackend(Protocol): + """A backend answers one prompt at a time. `json_mode=True` requests forced-JSON output + (final summary+tags call); `json_mode=False` is a plain-text call (map-reduce partials).""" + + async def complete(self, system: str, user: str, json_mode: bool = True) -> str: ... + + +class OllamaBackend: + def __init__(self, session: aiohttp.ClientSession, base_url: str, model: str): + self.session = session + self.base_url = base_url + self.model = model + + async def complete(self, system: str, user: str, json_mode: bool = True) -> str: + num_ctx = compute_num_ctx(len(system) + len(user)) + payload = { + "model": self.model, + "messages": [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ], + "stream": False, + "options": {"temperature": 0, "num_ctx": num_ctx}, + } + if json_mode: + payload["format"] = "json" + async with self.session.post(f"{self.base_url}/api/chat", json=payload) as resp: + resp.raise_for_status() + data = await resp.json() + return data["message"]["content"] + + +class AnthropicBackend: + def __init__(self, client, model: str, max_tokens: int = DEFAULT_MAX_TOKENS): + self.client = client + self.model = model + self.max_tokens = max_tokens + + async def complete(self, system: str, user: str, json_mode: bool = True) -> str: + kwargs = {} + if json_mode: + kwargs["output_config"] = { + "format": {"type": "json_schema", "schema": SUMMARY_JSON_SCHEMA} + } + response = await self.client.messages.create( + model=self.model, + max_tokens=self.max_tokens, + temperature=0, + system=system, + messages=[{"role": "user", "content": user}], + **kwargs, + ) + return next((b.text for b in response.content if b.type == "text"), "") + + +def load_tags_vocab(path: Path) -> list[str]: + """Controlled tag vocabulary (plan §2 decision 4) — a versioned YAML list.""" + with open(path, encoding="utf-8") as f: + data = yaml.safe_load(f) + return list(data.get("tags", [])) + + +def build_user_prompt(vocab: list[str], content: str) -> str: + return ( + "Słownik kontrolowany tagów (WYBIERZ Z TEJ LISTY w pierwszej kolejności, maks. 3 " + "najtrafniejsze pozycje, dokładnie w tym zapisie):\n" + f"{', '.join(vocab)}\n\n" + "Dopiero jeśli żaden z powyższych tagów nie pasuje do dokumentu, możesz dodać maks. " + "3 dodatkowe tagi spoza listy — one też MUSZĄ być PO POLSKU, małymi literami, " + "kebab-case (np. \"wspolnota-mieszkaniowa\", nie \"community\" ani \"HOA\"). " + "Nigdy nie zwracaj tagów po angielsku.\n\n" + "Treść dokumentu:\n" + f"{content}\n\n" + 'Zwróć JSON: {"summary": "<3-8 zdań PO POLSKU>", "tags": ["tag1", "tag2", ...]}' + ) + + +def build_partial_prompt(group_text: str, idx: int, total: int) -> str: + return ( + f"Fragment dokumentu, część {idx}/{total}:\n\n{group_text}\n\n" + "Streść ten fragment w 2-4 zdaniach po polsku, z konkretami (kwoty, daty, numery)." + ) + + +def is_duplicate(entities: list) -> bool: + """`entities[type=duplicate_of]` (plan §3.2 / §5.1) — duplicate envelopes are skipped + whole, never summarized even partially.""" + return any( + isinstance(e, dict) and e.get("type") == "duplicate_of" for e in (entities or []) + ) + + +def normalize_tag(tag: str) -> str: + """lowercase, kebab-case, NFC (plan §2 decision 4).""" + t = unicodedata.normalize("NFC", tag.strip().lower()) + t = t.replace(" ", "-") + t = _TAG_CLEAN_RE.sub("-", t) + t = _TAG_DASH_COLLAPSE_RE.sub("-", t) + return t.strip("-") + + +def normalize_and_validate_tags(raw_tags: list, vocab: list[str]) -> tuple[list[str], int]: + """Normalizes + dedupes tags; vocab tags are kept in full, free-form tags capped at 3 + (plan §2 decision 4). Returns (final_tags, truncated_count) — truncated_count is + informational only, not part of the stats balance invariant.""" + vocab_set = set(vocab) + seen: list[str] = [] + for raw in raw_tags or []: + if not isinstance(raw, str): + continue + norm = normalize_tag(raw) + if norm and norm not in seen: + seen.append(norm) + + vocab_tags = [t for t in seen if t in vocab_set] + freeform = [t for t in seen if t not in vocab_set] + truncated = max(0, len(freeform) - 3) + return vocab_tags + freeform[:3], truncated + + +async def fetch_documents(conn: asyncpg.Connection, limit: Optional[int], offset: Optional[int]) -> list: + """`source='paperless'` envelopes, ordered by id for stable --limit/--offset slicing.""" + query = "SELECT id, entities FROM envelope WHERE source = 'paperless' ORDER BY id" + params: list = [] + if limit is not None: + params.append(limit) + query += f" LIMIT ${len(params)}" + if offset: + params.append(offset) + query += f" OFFSET ${len(params)}" + return await conn.fetch(query, *params) + + +async def fetch_active_chunk_texts(conn: asyncpg.Connection, envelope_id: str) -> list[str]: + """`document_chunk.text WHERE excluded_reason IS NULL` (plan §5.1) — junk and duplicate + chunks never enter a summary.""" + rows = await conn.fetch( + "SELECT text FROM document_chunk WHERE envelope_id = $1 AND excluded_reason IS NULL " + "ORDER BY chunk_index", + envelope_id, + ) + return [r["text"] for r in rows] + + +async def fetch_existing_summary_keys(conn: asyncpg.Connection, model: str) -> set[str]: + """envelope_ids already summarized with this `model` — idempotency + dry-run preview.""" + rows = await conn.fetch("SELECT envelope_id FROM document_summary WHERE model = $1", model) + return {r["envelope_id"] for r in rows} + + +def _decode_jsonb(value: object) -> object: + if value is None: + return None + if isinstance(value, str): + return json.loads(value) + return value + + +async def get_summary_and_tags( + backend: LLMBackend, system: str, user: str, retries: int = 1 +) -> Optional[dict]: + """Forced-JSON call with one retry on invalid JSON/schema (plan §5.1). Returns None + (caller counts `llm_errors`) after `retries` failed attempts.""" + for attempt in range(retries + 1): + try: + raw = await backend.complete(system, user, json_mode=True) + data = json.loads(raw) + except Exception: + data = None + + if ( + isinstance(data, dict) + and isinstance(data.get("summary"), str) + and isinstance(data.get("tags"), list) + ): + return data + + if attempt < retries: + user = ( + user + + "\n\nUWAGA: Poprzednia odpowiedź nie była poprawnym JSON zgodnym ze " + "schematem. Zwróć WYŁĄCZNIE poprawny JSON, bez dodatkowego tekstu." + ) + return None + + +async def summarize_mapreduce( + backend: LLMBackend, chunks: list[str], group_size: int +) -> Optional[str]: + """Groups chunks ~`group_size`-at-a-time into partial summaries (plain text), returns + their concatenation as the "content" fed into the final JSON summarization pass. Returns + None if any partial call fails — caller counts the whole document as `llm_errors`, never + silently drops a partial.""" + groups = [chunks[i : i + group_size] for i in range(0, len(chunks), group_size)] + partials: list[str] = [] + for idx, group in enumerate(groups, start=1): + group_text = "\n\n".join(group) + try: + partial = await backend.complete( + PARTIAL_SYSTEM_PROMPT, + build_partial_prompt(group_text, idx, len(groups)), + json_mode=False, + ) + except Exception: + _log.warning("skip.mapreduce_partial_error", group_index=idx, group_total=len(groups)) + return None + partial = partial.strip() + if not partial: + _log.warning("skip.mapreduce_partial_empty", group_index=idx, group_total=len(groups)) + return None + partials.append(partial) + return "\n\n".join(partials) + + +async def insert_summary( + conn: asyncpg.Connection, envelope_id: str, summary: str, tags: list[str], model: str +) -> str: + """Returns asyncpg's command tag so the caller can tell a real insert from an + ON CONFLICT no-op (second line of defense behind the pre-fetched `existing` set).""" + return await conn.execute(_INSERT_SQL, envelope_id, summary, json.dumps(tags), model) + + +def _rows_affected(command_tag: str) -> int: + return int(command_tag.rsplit(" ", 1)[-1]) + + +async def run_summarize( + dsn: str, + backend_name: str, + model: str, + ollama_url: str = DEFAULT_OLLAMA_URL, + anthropic_api_key: Optional[str] = None, + tags_vocab_path: Path = DEFAULT_TAGS_VOCAB_PATH, + limit: Optional[int] = None, + offset: Optional[int] = None, + apply: bool = False, + max_tokens: int = DEFAULT_MAX_TOKENS, + chunks_per_group: int = DEFAULT_CHUNKS_PER_GROUP, + mapreduce_threshold_chars: int = DEFAULT_MAPREDUCE_THRESHOLD_CHARS, +) -> dict: + """Summarize+tag one --limit/--offset slice of `source='paperless'` envelopes into + `document_summary` under `model`. + + Stats must balance: + documents_fetched = duplicates_skipped + no_active_chunks + already_summarized + + summarized + llm_errors + + `duplicates_skipped` (envelope has `entities[type=duplicate_of]`) and `no_active_chunks` + (zero `document_chunk` rows with `excluded_reason IS NULL`) are both counted before any + LLM call — dry-run and apply agree on this split even though dry-run never calls the LLM. + A failed DB insert (rare — the pre-fetched `existing` set is the first line of defense + against duplicate work) is folded into `llm_errors`, same as `chunk_embed.py`'s + `chunks_errors` covers both embed and insert failures. + """ + stats = { + "documents_fetched": 0, + "duplicates_skipped": 0, + "no_active_chunks": 0, + "already_summarized": 0, + "summarized": 0, + "llm_errors": 0, + "documents_mapreduce": 0, + "tags_truncated": 0, + } + + conn = await asyncpg.connect(dsn) + try: + docs = await fetch_documents(conn, limit, offset) + existing = await fetch_existing_summary_keys(conn, model) + vocab = load_tags_vocab(tags_vocab_path) + + session: Optional[aiohttp.ClientSession] = None + client = None + backend: Optional[LLMBackend] = None + if apply: + if backend_name == "ollama": + session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=300)) + backend = OllamaBackend(session, ollama_url, model) + else: + client = AsyncAnthropic(api_key=anthropic_api_key) + backend = AnthropicBackend(client, model, max_tokens) + + try: + for row in docs: + stats["documents_fetched"] += 1 + envelope_id = row["id"] + entities = _decode_jsonb(row["entities"]) or [] + + if is_duplicate(entities): + stats["duplicates_skipped"] += 1 + continue + + if envelope_id in existing: + stats["already_summarized"] += 1 + continue + + chunks = await fetch_active_chunk_texts(conn, envelope_id) + content = "\n\n".join(chunks).strip() + if not content: + stats["no_active_chunks"] += 1 + continue + + if not apply: + stats["summarized"] += 1 + continue + + assert backend is not None + if len(content) > mapreduce_threshold_chars: + stats["documents_mapreduce"] += 1 + content = await summarize_mapreduce(backend, chunks, chunks_per_group) + if content is None: + stats["llm_errors"] += 1 + continue + + user_prompt = build_user_prompt(vocab, content) + data = await get_summary_and_tags(backend, SYSTEM_PROMPT, user_prompt) + if data is None: + stats["llm_errors"] += 1 + _log.warning("skip.llm_error", envelope_id=envelope_id) + continue + + tags, truncated = normalize_and_validate_tags(data["tags"], vocab) + stats["tags_truncated"] += truncated + + try: + command_tag = await insert_summary(conn, envelope_id, data["summary"], tags, model) + except Exception: + _log.warning("skip.insert_error", envelope_id=envelope_id, exc_info=True) + stats["llm_errors"] += 1 + continue + + existing.add(envelope_id) + if _rows_affected(command_tag) == 0: + # ON CONFLICT no-op: the pre-fetched `existing` set should make this rare; + # treat it the same as already-summarized rather than adding a bucket the + # plan's balance formula doesn't have. + stats["already_summarized"] += 1 + else: + stats["summarized"] += 1 + finally: + if session is not None: + await session.close() + if client is not None: + await client.close() + finally: + await conn.close() + + balance = ( + stats["duplicates_skipped"] + + stats["no_active_chunks"] + + stats["already_summarized"] + + stats["summarized"] + + stats["llm_errors"] + ) + if balance != stats["documents_fetched"]: + _log.error("stats_mismatch", **stats) + + _log.info("run_complete", apply=apply, backend=backend_name, model=model, **stats) + return stats + + +async def run_embed_summaries( + dsn: str, + ollama_url: str = DEFAULT_OLLAMA_URL, + embed_model: str = DEFAULT_EMBED_MODEL, + model_filter: Optional[str] = None, + limit: Optional[int] = None, + offset: Optional[int] = None, + apply: bool = False, +) -> dict: + """Second, separate pass: embeds `document_summary.summary` (bge-m3, via Ollama) for + rows with `embedding IS NULL`. Reuses `chunk_embed.embed_chunk` — same Ollama + `/api/embeddings` call, same dimension guard. `--model` (the summarizing model) filters + which track to embed; omit to embed both tracks' pending rows in one run. + + Stats must balance: summaries_fetched = embedded + errors + """ + stats = {"summaries_fetched": 0, "embedded": 0, "errors": 0} + + query = "SELECT id, summary FROM document_summary WHERE embedding IS NULL" + params: list = [] + if model_filter: + params.append(model_filter) + query += f" AND model = ${len(params)}" + query += " ORDER BY id" + if limit is not None: + params.append(limit) + query += f" LIMIT ${len(params)}" + if offset: + params.append(offset) + query += f" OFFSET ${len(params)}" + + conn = await asyncpg.connect(dsn) + try: + rows = await conn.fetch(query, *params) + + session: Optional[aiohttp.ClientSession] = None + if apply: + session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120)) + + try: + for row in rows: + stats["summaries_fetched"] += 1 + if not apply: + stats["embedded"] += 1 + continue + + assert session is not None + try: + embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, row["summary"]) + except Exception: + _log.warning("skip.embed_error", summary_id=row["id"], exc_info=True) + stats["errors"] += 1 + continue + + if len(embedding) != EXPECTED_DIM: + raise EmbeddingDimensionError( + f"ollama model={embed_model!r} returned dim={len(embedding)}, " + f"expected {EXPECTED_DIM}" + ) + + try: + await conn.execute( + "UPDATE document_summary SET embedding = $1::vector, embedding_model = $2 " + "WHERE id = $3", + _vector_literal(embedding), + embed_model, + row["id"], + ) + except Exception: + _log.warning("skip.update_error", summary_id=row["id"], exc_info=True) + stats["errors"] += 1 + continue + + stats["embedded"] += 1 + finally: + if session is not None: + await session.close() + finally: + await conn.close() + + if stats["embedded"] + stats["errors"] != stats["summaries_fetched"]: + _log.error("stats_mismatch", **stats) + + _log.info("run_complete", apply=apply, embed_model=embed_model, **stats) + return stats + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Summarize+tag source='paperless' documents into document_summary, or " + "embed existing summaries with --embed-summaries (module 5, phase 3, plan §5)." + ) + parser.add_argument("--dsn", default=os.environ.get("KB_DSN"), + help="asyncpg DSN for kb-postgres (or set KB_DSN env var)") + parser.add_argument("--limit", type=int, default=None, help="Max documents/summaries to process") + parser.add_argument("--offset", type=int, default=0, help="Slice offset, ordered by id") + parser.add_argument("--apply", action="store_true", + help="Actually call the LLM/Ollama and write rows. Default is dry-run.") + + parser.add_argument("--embed-summaries", action="store_true", + help="Embed existing document_summary rows (bge-m3) instead of writing new ones.") + parser.add_argument("--embed-model", default=os.environ.get("SUMMARY_EMBED_MODEL", DEFAULT_EMBED_MODEL), + help=f"Ollama embedding model for --embed-summaries (default: {DEFAULT_EMBED_MODEL})") + + parser.add_argument("--backend", choices=["ollama", "anthropic"], + help="LLM backend for writing summaries (required unless --embed-summaries)") + parser.add_argument("--model", help="Backend model name, e.g. gemma3:12b or claude-haiku-4-5 " + "(writing mode); or a filter on document_summary.model (--embed-summaries)") + parser.add_argument("--ollama-url", default=os.environ.get("OLLAMA_URL", DEFAULT_OLLAMA_URL), + help=f"Ollama base URL (default: {DEFAULT_OLLAMA_URL}, or set OLLAMA_URL)") + parser.add_argument("--anthropic-api-key", default=os.environ.get("ANTHROPIC_API_KEY"), + help="Anthropic API key (or set ANTHROPIC_API_KEY env var) — never logged") + parser.add_argument("--tags-vocab", type=Path, default=DEFAULT_TAGS_VOCAB_PATH, + help=f"Path to tags-vocab.yaml (default: {DEFAULT_TAGS_VOCAB_PATH})") + parser.add_argument("--max-tokens", type=int, default=DEFAULT_MAX_TOKENS, + help=f"Anthropic max_tokens (default: {DEFAULT_MAX_TOKENS})") + parser.add_argument("--chunks-per-group", type=int, default=DEFAULT_CHUNKS_PER_GROUP, + help=f"Map-reduce group size in chunks (default: {DEFAULT_CHUNKS_PER_GROUP})") + parser.add_argument("--mapreduce-threshold-chars", type=int, default=DEFAULT_MAPREDUCE_THRESHOLD_CHARS, + help=f"Content length above which map-reduce kicks in (default: {DEFAULT_MAPREDUCE_THRESHOLD_CHARS})") + args = parser.parse_args() + + if not args.dsn: + _log.error("missing_dsn", hint="pass --dsn or set KB_DSN") + sys.exit(1) + + if args.embed_summaries: + stats = asyncio.run( + run_embed_summaries( + dsn=args.dsn, + ollama_url=args.ollama_url, + embed_model=args.embed_model, + model_filter=args.model, + limit=args.limit, + offset=args.offset, + apply=args.apply, + ) + ) + mode = "APPLY" if args.apply else "DRY-RUN" + _log.info("summary", mode=mode, **stats) + balanced = stats["embedded"] + stats["errors"] == stats["summaries_fetched"] + sys.exit(1 if (stats["errors"] > 0 or not balanced) else 0) + + if not args.backend or not args.model: + _log.error("missing_backend_or_model", hint="--backend and --model are required unless --embed-summaries") + sys.exit(1) + if args.backend == "anthropic" and args.apply and not args.anthropic_api_key: + _log.error("missing_anthropic_api_key", hint="pass --anthropic-api-key or set ANTHROPIC_API_KEY") + sys.exit(1) + if not args.tags_vocab.exists(): + _log.error("missing_tags_vocab", path=str(args.tags_vocab)) + sys.exit(1) + + try: + stats = asyncio.run( + run_summarize( + dsn=args.dsn, + backend_name=args.backend, + model=args.model, + ollama_url=args.ollama_url, + anthropic_api_key=args.anthropic_api_key, + tags_vocab_path=args.tags_vocab, + limit=args.limit, + offset=args.offset, + apply=args.apply, + max_tokens=args.max_tokens, + chunks_per_group=args.chunks_per_group, + mapreduce_threshold_chars=args.mapreduce_threshold_chars, + ) + ) + except EmbeddingDimensionError as exc: + _log.error("dim_mismatch_abort", error=str(exc)) + sys.exit(1) + + mode = "APPLY" if args.apply else "DRY-RUN" + _log.info("summary", mode=mode, backend=args.backend, model=args.model, **stats) + + balanced = stats["documents_fetched"] == ( + stats["duplicates_skipped"] + stats["no_active_chunks"] + stats["already_summarized"] + + stats["summarized"] + stats["llm_errors"] + ) + failed = stats["llm_errors"] > 0 or not balanced + sys.exit(1 if failed else 0) + + +if __name__ == "__main__": + main() diff --git a/jobs/documents-ingest/tags-vocab.yaml b/jobs/documents-ingest/tags-vocab.yaml new file mode 100644 index 0000000..1d4ad1d --- /dev/null +++ b/jobs/documents-ingest/tags-vocab.yaml @@ -0,0 +1,20 @@ +# Controlled tag vocabulary for document_summary.tags (module 5, phase 3, plan §2 decision 4). +# Prompt forces the model to pick from this list; up to 3 extra free-form tags are allowed +# per document (normalized lowercase/kebab-case/NFC), the rest truncated (`tags_truncated`). +# After the pilot: review free-form tags via `jsonb_array_elements` + count, promote frequent +# ones here. Versioned — changes to this list are a deliberate, reviewable diff. +tags: + - ubezpieczenie + - bank + - kredyt + - faktura + - umowa + - urzad + - auto + - nieruchomosc + - zdrowie + - szkola + - fll + - praca + - subskrypcja + - regulamin diff --git a/jobs/documents-ingest/tests/test_summarize.py b/jobs/documents-ingest/tests/test_summarize.py new file mode 100644 index 0000000..03c2011 --- /dev/null +++ b/jobs/documents-ingest/tests/test_summarize.py @@ -0,0 +1,586 @@ +"""Unit tests for the summarize+tag job — no DB, no real HTTP, no real Ollama/Anthropic.""" +from __future__ import annotations + +import json + +import pytest + +from documents_ingest.summarize import ( + DEFAULT_TAGS_VOCAB_PATH, + OLLAMA_MAX_NUM_CTX, + OLLAMA_MIN_NUM_CTX, + build_partial_prompt, + build_user_prompt, + compute_num_ctx, + get_summary_and_tags, + insert_summary, + is_duplicate, + load_tags_vocab, + normalize_and_validate_tags, + normalize_tag, + run_embed_summaries, + run_summarize, + summarize_mapreduce, +) + +VOCAB = ["ubezpieczenie", "bank", "faktura"] + + +def _envelope_row(envelope_id, entities=None): + return {"id": envelope_id, "entities": json.dumps(entities or [])} + + +class TestComputeNumCtx: + def test_short_prompt_at_least_floor(self): + assert compute_num_ctx(100) >= OLLAMA_MIN_NUM_CTX + + def test_zero_length_prompt_uses_floor(self): + assert compute_num_ctx(0) == OLLAMA_MIN_NUM_CTX + + def test_scales_with_prompt_length(self): + # Regression case from the pilot: a 93k-char document was silently truncated to + # prompt_eval_count=2051 tokens by Ollama's runtime default -- num_ctx must scale + # well past that for a prompt this size. + ctx = compute_num_ctx(93_471) + assert ctx > 30_000 + + def test_huge_prompt_clamped_to_model_max(self): + assert compute_num_ctx(10_000_000) == OLLAMA_MAX_NUM_CTX + + def test_rounds_up_to_1024_step(self): + assert compute_num_ctx(1000) % 1024 == 0 + + +class TestTagsVocab: + def test_default_vocab_file_loads_and_is_nonempty(self): + vocab = load_tags_vocab(DEFAULT_TAGS_VOCAB_PATH) + assert "ubezpieczenie" in vocab + assert len(vocab) >= 5 + + +class TestIsDuplicate: + def test_no_entities_not_duplicate(self): + assert is_duplicate([]) is False + assert is_duplicate(None) is False + + def test_duplicate_of_entity_detected(self): + entities = [{"type": "content", "text": "x"}, {"type": "duplicate_of", "envelope_id": "paperless:14"}] + assert is_duplicate(entities) is True + + def test_unrelated_entities_not_duplicate(self): + entities = [{"type": "content", "text": "x"}, {"type": "tag", "name": "faktura"}] + assert is_duplicate(entities) is False + + +class TestNormalizeTag: + def test_lowercases_and_kebab_cases(self): + assert normalize_tag(" Ubezpieczenie Auto ") == "ubezpieczenie-auto" + + def test_strips_punctuation(self): + assert normalize_tag("PZU!!") == "pzu" + + def test_collapses_repeated_dashes(self): + assert normalize_tag("a b") == "a-b" + + def test_nfc_normalizes_diacritics(self): + # decomposed "s" + combining cedilla vs precomposed - should normalize the same way + import unicodedata + decomposed = unicodedata.normalize("NFD", "ś") + "rodowisko" # "ś" decomposed + assert normalize_tag(decomposed) == normalize_tag("środowisko") + + +class TestNormalizeAndValidateTags: + def test_keeps_all_vocab_tags(self): + tags, truncated = normalize_and_validate_tags(["bank", "faktura"], VOCAB) + assert tags == ["bank", "faktura"] + assert truncated == 0 + + def test_caps_freeform_at_three(self): + raw = ["bank", "polisa", "pzu", "auto-osobowe", "extra-one"] + tags, truncated = normalize_and_validate_tags(raw, VOCAB) + assert tags[0] == "bank" + assert len(tags) == 1 + 3 # 1 vocab + 3 freeform kept + assert truncated == 1 # "extra-one" dropped + + def test_dedupes_after_normalization(self): + tags, truncated = normalize_and_validate_tags(["Bank", "bank", " bank "], VOCAB) + assert tags == ["bank"] + assert truncated == 0 + + def test_ignores_non_string_entries(self): + tags, truncated = normalize_and_validate_tags(["bank", 123, None], VOCAB) + assert tags == ["bank"] + + +class TestPromptBuilders: + def test_user_prompt_includes_vocab_and_content(self): + prompt = build_user_prompt(VOCAB, "treść dokumentu") + assert "ubezpieczenie" in prompt + assert "treść dokumentu" in prompt + assert "JSON" in prompt + + def test_partial_prompt_includes_index_and_total(self): + prompt = build_partial_prompt("fragment", 2, 5) + assert "2/5" in prompt + assert "fragment" in prompt + + +class _FakeChatResponse: + def __init__(self, payload, status=200): + self._payload = payload + self._status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def json(self): + return self._payload + + def raise_for_status(self): + if self._status >= 400: + raise RuntimeError(f"HTTP {self._status}") + + +class _FakeOllamaChatSession: + """Serves a queue of /api/chat responses, one per call (in order).""" + + def __init__(self, contents: list[str]): + self._contents = list(contents) + self.requests: list[dict] = [] + + def post(self, url, json): + self.requests.append({"url": url, "json": json}) + content = self._contents.pop(0) + return _FakeChatResponse({"message": {"content": content}}) + + async def close(self): + pass + + +class _FakeBackend: + """A scripted LLMBackend: returns queued responses in order, regardless of prompt.""" + + def __init__(self, responses: list[str]): + self._responses = list(responses) + self.calls: list[tuple] = [] + + async def complete(self, system, user, json_mode=True): + self.calls.append((system, user, json_mode)) + return self._responses.pop(0) + + +class _FailingBackend: + async def complete(self, system, user, json_mode=True): + raise RuntimeError("boom") + + +class TestGetSummaryAndTags: + async def test_valid_json_first_try(self): + backend = _FakeBackend([json.dumps({"summary": "s", "tags": ["bank"]})]) + data = await get_summary_and_tags(backend, "sys", "user") + assert data == {"summary": "s", "tags": ["bank"]} + assert len(backend.calls) == 1 + + async def test_invalid_json_then_valid_on_retry(self): + backend = _FakeBackend(["not json", json.dumps({"summary": "s", "tags": []})]) + data = await get_summary_and_tags(backend, "sys", "user") + assert data == {"summary": "s", "tags": []} + assert len(backend.calls) == 2 + + async def test_missing_keys_counts_as_invalid(self): + backend = _FakeBackend([json.dumps({"summary": "s"}), json.dumps({"summary": "s", "tags": []})]) + data = await get_summary_and_tags(backend, "sys", "user") + assert data == {"summary": "s", "tags": []} + + async def test_gives_up_after_retries_exhausted(self): + backend = _FakeBackend(["not json", "still not json"]) + data = await get_summary_and_tags(backend, "sys", "user", retries=1) + assert data is None + + async def test_backend_exception_is_treated_as_invalid(self): + data = await get_summary_and_tags(_FailingBackend(), "sys", "user", retries=0) + assert data is None + + +class TestSummarizeMapreduce: + async def test_combines_partials_in_order(self): + backend = _FakeBackend(["partial one", "partial two"]) + result = await summarize_mapreduce(backend, ["c1", "c2", "c3"], group_size=2) + assert result == "partial one\n\npartial two" + # two groups of size 2 and 1 -> two calls, both plain-text (json_mode=False) + assert len(backend.calls) == 2 + assert all(call[2] is False for call in backend.calls) + + async def test_partial_failure_aborts_whole_document(self): + result = await summarize_mapreduce(_FailingBackend(), ["c1", "c2"], group_size=2) + assert result is None + + async def test_empty_partial_aborts(self): + backend = _FakeBackend([" "]) + result = await summarize_mapreduce(backend, ["c1"], group_size=2) + assert result is None + + +class TestInsertSql: + def test_on_conflict_target_includes_model(self): + from documents_ingest.summarize import _INSERT_SQL + assert "ON CONFLICT (envelope_id, model)" in _INSERT_SQL + + +class _FakeConn: + """Fakes the three query shapes this job issues: envelope fetch, active-chunk-text + fetch, and existing-summary-keys fetch; plus execute() for INSERT/UPDATE.""" + + def __init__(self, docs=None, chunks_by_envelope=None, existing_summaries=None, + existing_model="bge-m3", execute_results=None): + self._docs = docs or [] + self._chunks_by_envelope = chunks_by_envelope or {} + self._existing_summaries = list(existing_summaries or []) + self._existing_model = existing_model + self._execute_results = list(execute_results) if execute_results is not None else None + self.execute_calls: list[tuple] = [] + + async def fetch(self, query, *params): + if "FROM document_chunk" in query: + envelope_id = params[0] + return [{"text": t} for t in self._chunks_by_envelope.get(envelope_id, [])] + if "FROM document_summary" in query: + if params and params[0] != self._existing_model: + return [] + return [{"envelope_id": eid} for eid in self._existing_summaries] + return self._docs + + async def execute(self, query, *params): + self.execute_calls.append(params) + if self._execute_results is not None: + return self._execute_results.pop(0) + return "INSERT 0 1" + + async def close(self): + pass + + +class TestRunSummarize: + def _patch(self, monkeypatch, conn, session=None): + async def _fake_connect(dsn): + return conn + monkeypatch.setattr("documents_ingest.summarize.asyncpg.connect", _fake_connect) + if session is not None: + def _fake_session_factory(*args, **kwargs): + return session + monkeypatch.setattr("documents_ingest.summarize.aiohttp.ClientSession", _fake_session_factory) + + async def test_dry_run_counts_without_calling_backend_or_db(self, monkeypatch): + conn = _FakeConn( + docs=[_envelope_row("paperless:1")], + chunks_by_envelope={"paperless:1": ["treść dokumentu"]}, + ) + self._patch(monkeypatch, conn) + stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False) + assert stats["documents_fetched"] == 1 + assert stats["summarized"] == 1 + assert stats["llm_errors"] == 0 + assert conn.execute_calls == [] + + async def test_duplicate_envelope_skipped_whole(self, monkeypatch): + conn = _FakeConn( + docs=[_envelope_row("paperless:74", entities=[{"type": "duplicate_of", "envelope_id": "paperless:14"}])], + ) + self._patch(monkeypatch, conn) + stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False) + assert stats["duplicates_skipped"] == 1 + assert stats["summarized"] == 0 + + async def test_no_active_chunks_counted_separately(self, monkeypatch): + conn = _FakeConn(docs=[_envelope_row("paperless:2")], chunks_by_envelope={}) + self._patch(monkeypatch, conn) + stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False) + assert stats["no_active_chunks"] == 1 + assert stats["summarized"] == 0 + + async def test_already_summarized_skipped(self, monkeypatch): + conn = _FakeConn( + docs=[_envelope_row("paperless:1")], + chunks_by_envelope={"paperless:1": ["text"]}, + existing_summaries=["paperless:1"], + existing_model="gemma3:12b", + ) + self._patch(monkeypatch, conn) + stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False) + assert stats["already_summarized"] == 1 + assert stats["summarized"] == 0 + + async def test_apply_writes_summary_and_tags(self, monkeypatch): + conn = _FakeConn( + docs=[_envelope_row("paperless:1")], + chunks_by_envelope={"paperless:1": ["treść z kwotą 100 zł"]}, + ) + session = _FakeOllamaChatSession([json.dumps({"summary": "Streszczenie.", "tags": ["bank", "extra"]})]) + self._patch(monkeypatch, conn, session) + + stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True) + + assert stats["summarized"] == 1 + assert stats["llm_errors"] == 0 + assert len(conn.execute_calls) == 1 + envelope_id, summary, tags_json, model = conn.execute_calls[0] + assert envelope_id == "paperless:1" + assert summary == "Streszczenie." + assert json.loads(tags_json) == ["bank", "extra"] + assert model == "gemma3:12b" + # Regression guard: every Ollama call must carry an explicit num_ctx (plan §5.1 + # limitation) -- Ollama's runtime default silently truncates long documents otherwise. + assert session.requests[0]["json"]["options"]["num_ctx"] >= OLLAMA_MIN_NUM_CTX + + async def test_apply_with_anthropic_backend(self, monkeypatch): + conn = _FakeConn( + docs=[_envelope_row("paperless:1")], + chunks_by_envelope={"paperless:1": ["treść"]}, + ) + + class _TextBlock: + type = "text" + text = json.dumps({"summary": "Streszczenie API.", "tags": ["faktura"]}) + + class _FakeResponse: + content = [_TextBlock()] + + class _FakeMessages: + def __init__(self): + self.create_calls = [] + + async def create(self, **kwargs): + self.create_calls.append(kwargs) + return _FakeResponse() + + class _FakeAsyncAnthropic: + def __init__(self, api_key=None): + self.api_key = api_key + self.messages = _FakeMessages() + + async def close(self): + pass + + fake_client_holder = {} + + def _factory(api_key=None): + client = _FakeAsyncAnthropic(api_key=api_key) + fake_client_holder["client"] = client + return client + + async def _fake_connect(dsn): + return conn + monkeypatch.setattr("documents_ingest.summarize.asyncpg.connect", _fake_connect) + monkeypatch.setattr("documents_ingest.summarize.AsyncAnthropic", _factory) + + stats = await run_summarize( + dsn="x", backend_name="anthropic", model="claude-haiku-4-5", + anthropic_api_key="sk-fake", apply=True, + ) + + assert stats["summarized"] == 1 + client = fake_client_holder["client"] + assert client.api_key == "sk-fake" + assert len(client.messages.create_calls) == 1 + call = client.messages.create_calls[0] + assert call["model"] == "claude-haiku-4-5" + assert call["temperature"] == 0 + assert call["output_config"]["format"]["type"] == "json_schema" + + async def test_llm_error_isolated_and_counted(self, monkeypatch): + conn = _FakeConn( + docs=[_envelope_row("paperless:1"), _envelope_row("paperless:2")], + chunks_by_envelope={"paperless:1": ["a"], "paperless:2": ["b"]}, + ) + session = _FakeOllamaChatSession([ + "not json", "still not json", # paperless:1 exhausts retry -> llm_error + json.dumps({"summary": "ok", "tags": []}), # paperless:2 succeeds + ]) + self._patch(monkeypatch, conn, session) + + stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True) + + assert stats["documents_fetched"] == 2 + assert stats["llm_errors"] == 1 + assert stats["summarized"] == 1 + assert len(conn.execute_calls) == 1 # only the successful one got inserted + + async def test_conflict_skip_counted_as_already_summarized(self, monkeypatch): + conn = _FakeConn( + docs=[_envelope_row("paperless:1")], + chunks_by_envelope={"paperless:1": ["a"]}, + execute_results=["INSERT 0 0"], # ON CONFLICT DO NOTHING no-op + ) + session = _FakeOllamaChatSession([json.dumps({"summary": "s", "tags": []})]) + self._patch(monkeypatch, conn, session) + + stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True) + + assert stats["already_summarized"] == 1 + assert stats["summarized"] == 0 + + async def test_stats_balance_invariant(self, monkeypatch): + conn = _FakeConn( + docs=[ + _envelope_row("paperless:1", entities=[{"type": "duplicate_of", "envelope_id": "paperless:0"}]), + _envelope_row("paperless:2"), + _envelope_row("paperless:3"), + ], + chunks_by_envelope={"paperless:3": ["content"]}, + ) + self._patch(monkeypatch, conn) + stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False) + balance = ( + stats["duplicates_skipped"] + stats["no_active_chunks"] + + stats["already_summarized"] + stats["summarized"] + stats["llm_errors"] + ) + assert balance == stats["documents_fetched"] == 3 + + async def test_limit_and_offset_passed_through(self, monkeypatch): + class _Conn(_FakeConn): + async def fetch(self, query, *params): + if "FROM envelope" in query: + self.captured = (query, params) + return [] + return await super().fetch(query, *params) + + conn = _Conn() + self._patch(monkeypatch, conn) + await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", limit=5, offset=10, apply=False) + query, params = conn.captured + assert "LIMIT" in query and "OFFSET" in query + assert params == (5, 10) + + async def test_rerun_after_apply_writes_nothing_new(self, monkeypatch): + conn1 = _FakeConn( + docs=[_envelope_row("paperless:1")], + chunks_by_envelope={"paperless:1": ["a"]}, + ) + session1 = _FakeOllamaChatSession([json.dumps({"summary": "s", "tags": []})]) + self._patch(monkeypatch, conn1, session1) + await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True) + + conn2 = _FakeConn( + docs=[_envelope_row("paperless:1")], + chunks_by_envelope={"paperless:1": ["a"]}, + existing_summaries=["paperless:1"], + existing_model="gemma3:12b", + ) + session2 = _FakeOllamaChatSession([]) + self._patch(monkeypatch, conn2, session2) + stats2 = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True) + + assert stats2["already_summarized"] == 1 + assert stats2["summarized"] == 0 + + +class TestMapreduceIntegration: + def _patch(self, monkeypatch, conn, session=None): + async def _fake_connect(dsn): + return conn + monkeypatch.setattr("documents_ingest.summarize.asyncpg.connect", _fake_connect) + if session is not None: + monkeypatch.setattr("documents_ingest.summarize.aiohttp.ClientSession", lambda *a, **k: session) + + async def test_long_document_triggers_mapreduce(self, monkeypatch): + long_chunks = ["x" * 100_000, "y" * 150_000] # combined > threshold + conn = _FakeConn( + docs=[_envelope_row("paperless:1")], + chunks_by_envelope={"paperless:1": long_chunks}, + ) + session = _FakeOllamaChatSession([ + "partial summary 1", # map-reduce partial (plain text, one group covers both chunks) + json.dumps({"summary": "final", "tags": []}), # final synthesis (JSON) + ]) + self._patch(monkeypatch, conn, session) + + stats = await run_summarize( + dsn="x", backend_name="ollama", model="gemma3:12b", apply=True, + mapreduce_threshold_chars=200_000, chunks_per_group=20, + ) + + assert stats["documents_mapreduce"] == 1 + assert stats["summarized"] == 1 + # first call has no "format" key (plain text), second has format=json + assert "format" not in session.requests[0]["json"] + assert session.requests[1]["json"]["format"] == "json" + + +class TestRunEmbedSummaries: + class _EmbedConn: + def __init__(self, rows): + self._rows = rows + self.update_calls: list[tuple] = [] + + async def fetch(self, query, *params): + return self._rows + + async def execute(self, query, *params): + self.update_calls.append(params) + return "UPDATE 1" + + async def close(self): + pass + + class _EmbedSession: + def __init__(self, dim=1024): + self._dim = dim + + def post(self, url, json): + return _FakeEmbedResponseForEmbed({"embedding": [0.1] * self._dim}) + + async def close(self): + pass + + def _patch(self, monkeypatch, conn, session=None): + async def _fake_connect(dsn): + return conn + monkeypatch.setattr("documents_ingest.summarize.asyncpg.connect", _fake_connect) + if session is not None: + monkeypatch.setattr("documents_ingest.summarize.aiohttp.ClientSession", lambda *a, **k: session) + + async def test_dry_run_counts_without_writing(self, monkeypatch): + conn = self._EmbedConn(rows=[{"id": 1, "summary": "s"}]) + self._patch(monkeypatch, conn) + stats = await run_embed_summaries(dsn="x", apply=False) + assert stats["summaries_fetched"] == 1 + assert stats["embedded"] == 1 + assert conn.update_calls == [] + + async def test_apply_embeds_and_updates(self, monkeypatch): + conn = self._EmbedConn(rows=[{"id": 1, "summary": "s"}]) + session = self._EmbedSession() + self._patch(monkeypatch, conn, session) + stats = await run_embed_summaries(dsn="x", embed_model="bge-m3", apply=True) + assert stats["embedded"] == 1 + assert stats["errors"] == 0 + assert len(conn.update_calls) == 1 + vector_literal, embed_model, summary_id = conn.update_calls[0] + assert embed_model == "bge-m3" + assert summary_id == 1 + + async def test_balance_invariant(self, monkeypatch): + conn = self._EmbedConn(rows=[{"id": 1, "summary": "s"}, {"id": 2, "summary": "t"}]) + self._patch(monkeypatch, conn) + stats = await run_embed_summaries(dsn="x", apply=False) + assert stats["embedded"] + stats["errors"] == stats["summaries_fetched"] + + +class _FakeEmbedResponseForEmbed: + def __init__(self, payload): + self._payload = payload + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def json(self): + return self._payload + + def raise_for_status(self): + pass diff --git a/services/kb-postgres/init/004_summaries.sql b/services/kb-postgres/init/004_summaries.sql new file mode 100644 index 0000000..c82a2ca --- /dev/null +++ b/services/kb-postgres/init/004_summaries.sql @@ -0,0 +1,23 @@ +-- KB spine: document_summary — summaries + tags per document, half-product of compilation +-- Additive: does not modify 001_envelope.sql, 002_chunks.sql, or 003_chunk_model_key.sql. +-- Version: 004 — document_summary (module 5, phase 3, plan §4) +-- +-- 1:N to envelope by model: the same envelope can carry summaries from multiple models +-- (pilot A/B: API vs local GPU); retrieval selects a model via config. + +CREATE TABLE IF NOT EXISTS document_summary ( + id BIGSERIAL PRIMARY KEY, + envelope_id TEXT NOT NULL REFERENCES envelope(id) ON DELETE CASCADE, + summary TEXT NOT NULL, + tags JSONB NOT NULL DEFAULT '[]', + model TEXT NOT NULL, -- model that WROTE the summary + embedding VECTOR(1024), -- NULL until embedded + embedding_model TEXT, -- embedding model (bge-m3); NULL with embedding + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE (envelope_id, model) +); + +CREATE INDEX IF NOT EXISTS document_summary_envelope_idx + ON document_summary (envelope_id); +CREATE INDEX IF NOT EXISTS document_summary_embedding_hnsw_idx + ON document_summary USING hnsw (embedding vector_cosine_ops);