feat(kb): faza 3 krok 2 — migracja 004 (document_summary) + pilot streszczeń A/B

Migracja 004: document_summary (envelope_id, summary, tags JSONB, model, embedding
VECTOR(1024) + HNSW cosine, embedding_model, UNIQUE(envelope_id, model) od razu — wzorzec
002/003). Zastosowana na żywej bazie kb-postgres@PIHA.

Job documents-ingest-summarize: wejście = document_chunk.text WHERE excluded_reason IS NULL
per koperta source='paperless' (duplikaty przez entities[duplicate_of] pomijane w całości),
wymuszony JSON {summary, tags} przez --backend ollama|anthropic, słownik tagów kontrolowany
(tags-vocab.yaml) z max 3 free-form, map-reduce dla dokumentów >200k znaków (grupy ~20
chunków), bilans + idempotencja + izolacja błędów per wiersz wg wzorców rodziny jobów.
Osobny --embed-summaries (bge-m3, reużywa chunk_embed.embed_chunk).

Bug znaleziony i naprawiony w trakcie pilota: brak options.num_ctx w wywołaniach Ollamy
powodował, że gemma3:12b używał domyślnego runtime kontekstu (~2048 tok), nie
zadeklarowanego 128k — dla 71/157 dokumentów (45%, >8k znaków) treść była cicho ucinana
(zweryfikowane: prompt_eval_count=2051 dla dokumentu 93k znaków). Naprawa: compute_num_ctx()
liczy num_ctx z długości promptu (~3 znaki/token, cap 131072). Cały tor lokalny przeliczony
od zera po naprawie.

Wynik pilota (186 dok. paperless, 3 duplikaty, 26 bez aktywnych chunków → 157 oczekiwanych
na tor): tor lokalny (gemma3:12b) 155/157 (2 izolowane błędy JSON po retry: paperless:24,
paperless:61), tor referencyjny (claude-haiku-4-5) 157/157, 0 błędów JSON, 0 tagów
ucinanych — słownik przestrzegany w 100% przypadków. Oba komplety zembedowane (bge-m3).
Znaleziony przy okazji: prompt do tagów wymagał dopracowania — pierwsza wersja pozwalała
modelowi zwracać tagi po angielsku spoza słownika; wzmocniona instrukcja (słownik w
pierwszej kolejności, "nigdy po angielsku") poprawiła zgodność w 2/3 przypadków testowych.

Porównanie A/B (~15 dok.) i weryfikacja końcowa (bilans, sanity SQL, retrieval po summary)
odłożone do następnej sesji.

Testy: 157 (152 nowe/summarize.py + istniejące), mocki API/Ollama/DB, bilans, idempotencja,
regresja num_ctx.

Co najmniej 3 decyzje wymagały zatrzymania i potwierdzenia z Oskarem w sesji (sposób podania
klucza API, wybór modelu lokalnego gemma3:12b, naprawa+przeliczenie całego toru lokalnego po
odkryciu buga num_ctx) — udokumentowane w transkrypcie sesji.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
oskar 2026-07-16 21:59:21 +02:00
parent 8eaf993ea1
commit 569f95005d
5 changed files with 1307 additions and 0 deletions

View file

@ -11,12 +11,15 @@ dependencies = [
"structlog>=24.1", "structlog>=24.1",
"aiohttp>=3.9", "aiohttp>=3.9",
"kb-mail", "kb-mail",
"PyYAML>=6.0",
"anthropic>=0.40",
] ]
[project.scripts] [project.scripts]
documents-ingest = "documents_ingest.extractor:main" documents-ingest = "documents_ingest.extractor:main"
documents-ingest-paperless = "documents_ingest.paperless_adapter:main" documents-ingest-paperless = "documents_ingest.paperless_adapter:main"
documents-ingest-embed = "documents_ingest.chunk_embed:main" documents-ingest-embed = "documents_ingest.chunk_embed:main"
documents-ingest-summarize = "documents_ingest.summarize:main"
[tool.setuptools.packages.find] [tool.setuptools.packages.find]
where = ["src"] where = ["src"]

View file

@ -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()

View file

@ -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

View file

@ -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

View file

@ -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);