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:
parent
8eaf993ea1
commit
569f95005d
|
|
@ -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"]
|
||||
|
|
|
|||
675
jobs/documents-ingest/src/documents_ingest/summarize.py
Normal file
675
jobs/documents-ingest/src/documents_ingest/summarize.py
Normal 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()
|
||||
20
jobs/documents-ingest/tags-vocab.yaml
Normal file
20
jobs/documents-ingest/tags-vocab.yaml
Normal 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
|
||||
586
jobs/documents-ingest/tests/test_summarize.py
Normal file
586
jobs/documents-ingest/tests/test_summarize.py
Normal 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
|
||||
23
services/kb-postgres/init/004_summaries.sql
Normal file
23
services/kb-postgres/init/004_summaries.sql
Normal 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);
|
||||
Loading…
Reference in a new issue