f964631d02
Module 5 phase 4 step 0 (docs/kb/modules/05-faza4-plan.md, §3, decision 1): kb-query is a long-lived Docker service, documents-ingest is a venv job with an `anthropic` dependency and CLI scripts it doesn't need. Move embed_chunk/_vector_literal/cascade_query/flat_query into a shared package with minimal deps (asyncpg, aiohttp only) so both can depend on the same tested retrieval code without the service image pulling in the job's extras. documents_ingest.chunk_embed/retrieval keep thin re-exports so nothing importing the old paths breaks. Pure refactor: retrieval_eval.py run live against kb-postgres@PIHA + Ollama@SOLARIA before/after gives byte-identical `dist`/hit@3/gate results (still PASS) — zero behavior change in the cascade. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
672 lines
27 KiB
Python
672 lines
27 KiB
Python
"""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
|
|
`kb_retrieval.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
|
|
from kb_retrieval.embed import _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()
|