homelab-codex-ws/jobs/mail-body-ingest/src/mail_body_ingest/ingest.py

645 lines
26 KiB
Python
Raw Normal View History

feat(mail-body-ingest): new job to chunk+embed gmail body content (faza mailowa Krok 2) Second full pass over the gmail .eml archive (gmail-bulk-import's first pass skipped inline text/plain and text/html on purpose). Per envelope: typed parse with compat32 fallback -> body extraction (inline text/plain preferred, HTML->text via a small stdlib HTMLParser otherwise) -> quote-strip (reply markers + `>`-quoted lines, EN/PL/Outlook patterns) -> newsletter classification (List-Unsubscribe/List-Id/Precedence, chunked but not embedded, excluded_reason='newsletter') -> Temat/Od/Data prefix from the already-backfilled entities[type=headers] (zero header re-parse) -> chunk via kb_mail.chunking -> batched embed_batch (64) -> INSERT document_chunk. In-Reply-To/References are appended as entities[type=threading] during the same read (idempotent WHERE NOT EXISTS append, 1:1 with gmail-header-backfill) -- the only DB writes are document_chunk INSERTs and an additive envelope.entities UPDATE; the .eml archive stays read-only. --dsn/KB_DSN, --archive-root, --since/--limit/--offset, --batch-size, --apply (dry-run default). Idempotency keys on (envelope_id, chunk_index) pre-fetched scoped to --model, built correctly from the start per the plan's flagged chunk_embed.py precedent. A failed embed batch is isolated (chunks_errors, no abort) for Ollama's documented instability; a wrong embedding dimension aborts the whole run. 48 tests, DoD smoke run against live kb-postgres@PIHA confirmed wiring (archive not yet rsync'd to SOLARIA, so all 5 rows correctly reported missing_file). docs/kb/modules/05-faza-mailowa-plan.md, §5.
2026-07-22 19:01:14 +02:00
"""Mail body ingest job — module 5, faza mailowa, plan Krok 2
(docs/kb/modules/05-faza-mailowa-plan.md, §5). Second full pass over the gmail .eml archive
(the first was gmail-bulk-import's manifest-only import): extracts inline body text that
`_parse_attachments` deliberately skipped, chunks it, embeds it (batched), and appends
`entities[type=threading]` the only new writes are `document_chunk` INSERTs and an additive
`envelope.entities` UPDATE. The archive itself is read-only; `envelope.id`/`raw_ref`/`ts` are
never touched.
Runs on SOLARIA (needs Ollama on localhost) against kb-postgres@PIHA over Tailscale the .eml
archive is rsync'd PIHA -> SOLARIA once (plan §7 Krok 4), not read live over the network.
Install (from repo root):
pip install -e packages/kb-mail/
pip install -e packages/kb-retrieval/
pip install -e jobs/mail-body-ingest/
Usage:
# Dry run (default) — parse, quote-strip, classify, chunk, count. Zero Ollama calls, zero
# DB writes (including the threading UPDATE):
mail-body-ingest --dsn postgresql://kb:<pw>@piha:5433/kb --archive-root /path/to/archive
# Etap A pilot — last 12 months only:
mail-body-ingest --dsn ... --since 2025-07-01 --apply
# Smoke-test slice:
mail-body-ingest --dsn ... --apply --limit 10
DSN can come from KB_DSN, Ollama URL from OLLAMA_URL (default http://localhost:11434 this job
is meant to run where Ollama lives, not call it over Tailscale per-chunk).
Idempotency: a pre-fetched set of existing (envelope_id, chunk_index) pairs for this `model`
skips chunks already inserted (the fetch itself filters `WHERE model = $1`, so the pair is
already scoped to the model without needing it redundantly in the tuple same pattern as
`documents_ingest.chunk_embed`, built correctly from the start here per plan §1.5). The
`ON CONFLICT (envelope_id, chunk_index, model) DO NOTHING` insert is the second line of defense;
its command tag is checked so a silently-skipped row counts as `chunks_conflict_skipped`, never
`chunks_inserted`. Threading is separately idempotent via `WHERE NOT EXISTS (... type=threading)`
the same 1:1 pattern as `gmail_header_backfill`.
Newsletter chunks (`List-Unsubscribe`/`List-Id`/`Precedence: bulk|list`) are inserted with
`excluded_reason='newsletter'`, `embedding=NULL` no Ollama call, no HNSW entry, but the text
stays in the DB (reversible: `UPDATE ... SET excluded_reason=NULL WHERE excluded_reason=
'newsletter'` + a re-embed run un-flags them later, per plan Decyzja 4).
Ollama-offline tolerance: a failed `embed_batch()` call is caught per-batch (`chunks_errors +=
len(batch)`), never aborting the run those chunks never enter the idempotency set, so a later
re-run naturally retries them. Only a wrong embedding dimension aborts the whole run
(`EmbeddingDimensionError`) never silently indexes a mismatched vector.
"""
from __future__ import annotations
import argparse
import asyncio
import email
import email.policy
import json
import os
import re
import sys
from datetime import datetime, timezone
from html.parser import HTMLParser
from pathlib import Path
from typing import Optional
import aiohttp
import asyncpg
import structlog
from kb_mail.chunking import OVERLAP_CHARS, TARGET_CHARS, chunk_text
from kb_mail.text import sanitize_surrogates as _sanitize
from kb_retrieval.embed import (
DEFAULT_MODEL,
DEFAULT_OLLAMA_URL,
EmbeddingDimensionError,
_vector_literal,
check_ollama_health,
embed_batch,
)
_log = structlog.get_logger(__name__)
DEFAULT_ARCHIVE_ROOT = Path("/home/oskar/kb/mail/archive")
DEFAULT_BATCH_SIZE = 64
THREADING_UPDATE_BATCH_SIZE = 500
_CHUNK_INSERT_SQL = """
INSERT INTO document_chunk (envelope_id, chunk_index, text, embedding, model, excluded_reason)
VALUES ($1, $2, $3, $4::vector, $5, $6)
ON CONFLICT (envelope_id, chunk_index, model) DO NOTHING
"""
_THREADING_UPDATE_SQL = """
UPDATE envelope
SET entities = entities || $2::jsonb
WHERE id = $1
AND NOT EXISTS (
SELECT 1 FROM jsonb_array_elements(entities) e WHERE e->>'type' = 'threading'
)
"""
# Decision 2b: reply markers -- earliest match of any of these truncates the rest of the body.
# `.` doesn't cross newlines (no re.DOTALL) -- these are meant to match a single header line,
# same shape gmail/outlook/pl clients actually emit.
_REPLY_MARKER_PATTERNS = [
re.compile(r"^On .+ wrote:\s*$", re.IGNORECASE | re.MULTILINE),
re.compile(r"^Dnia .+ napisał\(a\):\s*$", re.IGNORECASE | re.MULTILINE),
re.compile(r"^W dniu .+ pisze:\s*$", re.IGNORECASE | re.MULTILINE),
re.compile(r"^-{5,}\s*Original Message\s*-{5,}\s*$", re.IGNORECASE | re.MULTILINE),
re.compile(r"^_{10,}\s*$", re.MULTILINE), # Outlook's horizontal-rule separator
]
# Decision 2a: quoted lines -- applied to whatever remains after marker truncation.
_QUOTE_LINE_RE = re.compile(r"^\s*>")
_ID_TOKEN_RE = re.compile(r"<[^<>]+>")
class _MailHTMLTextExtractor(HTMLParser):
"""HTML -> text: skips `style`/`script`/`head` content and `blockquote`/`div.gmail_quote`
subtrees (Decision 2c quoted reply chains in HTML mail), inserts newlines at block-level
boundaries so `kb_mail.chunking.chunk_text`'s paragraph splitter has something to work with.
Tag-name-matched stack (search backward on close, truncate from the match): real-world mail
HTML frequently writes void elements (`<meta>`, `<br>`, `<img>`, ...) without a self-closing
slash. A naive LIFO push/pop-on-any-endtag desyncs on these -- e.g. `<head><meta><meta>
</head>` pops a `meta` "frame" for the literal `</head>`, leaving `skip_depth` stuck at 1 for
the rest of the document (confirmed live: a real HTML-only mail with a `<head>` full of
`<meta>` tags extracted to '' entirely). Void elements are never pushed at all; every other
close searches backward for its matching open and discards anything opened-but-never-closed
after it, which also self-heals other malformed nesting instead of just the void-tag case.
feat(mail-body-ingest): new job to chunk+embed gmail body content (faza mailowa Krok 2) Second full pass over the gmail .eml archive (gmail-bulk-import's first pass skipped inline text/plain and text/html on purpose). Per envelope: typed parse with compat32 fallback -> body extraction (inline text/plain preferred, HTML->text via a small stdlib HTMLParser otherwise) -> quote-strip (reply markers + `>`-quoted lines, EN/PL/Outlook patterns) -> newsletter classification (List-Unsubscribe/List-Id/Precedence, chunked but not embedded, excluded_reason='newsletter') -> Temat/Od/Data prefix from the already-backfilled entities[type=headers] (zero header re-parse) -> chunk via kb_mail.chunking -> batched embed_batch (64) -> INSERT document_chunk. In-Reply-To/References are appended as entities[type=threading] during the same read (idempotent WHERE NOT EXISTS append, 1:1 with gmail-header-backfill) -- the only DB writes are document_chunk INSERTs and an additive envelope.entities UPDATE; the .eml archive stays read-only. --dsn/KB_DSN, --archive-root, --since/--limit/--offset, --batch-size, --apply (dry-run default). Idempotency keys on (envelope_id, chunk_index) pre-fetched scoped to --model, built correctly from the start per the plan's flagged chunk_embed.py precedent. A failed embed batch is isolated (chunks_errors, no abort) for Ollama's documented instability; a wrong embedding dimension aborts the whole run. 48 tests, DoD smoke run against live kb-postgres@PIHA confirmed wiring (archive not yet rsync'd to SOLARIA, so all 5 rows correctly reported missing_file). docs/kb/modules/05-faza-mailowa-plan.md, §5.
2026-07-22 19:01:14 +02:00
"""
_SKIP_CONTENT_TAGS = {"style", "script", "head"}
_BLOCK_TAGS = {"p", "div", "tr", "li", "h1", "h2", "h3", "h4", "h5", "h6", "br"}
_VOID_ELEMENTS = {
"area", "base", "br", "col", "embed", "hr", "img", "input",
"link", "meta", "param", "source", "track", "wbr",
}
feat(mail-body-ingest): new job to chunk+embed gmail body content (faza mailowa Krok 2) Second full pass over the gmail .eml archive (gmail-bulk-import's first pass skipped inline text/plain and text/html on purpose). Per envelope: typed parse with compat32 fallback -> body extraction (inline text/plain preferred, HTML->text via a small stdlib HTMLParser otherwise) -> quote-strip (reply markers + `>`-quoted lines, EN/PL/Outlook patterns) -> newsletter classification (List-Unsubscribe/List-Id/Precedence, chunked but not embedded, excluded_reason='newsletter') -> Temat/Od/Data prefix from the already-backfilled entities[type=headers] (zero header re-parse) -> chunk via kb_mail.chunking -> batched embed_batch (64) -> INSERT document_chunk. In-Reply-To/References are appended as entities[type=threading] during the same read (idempotent WHERE NOT EXISTS append, 1:1 with gmail-header-backfill) -- the only DB writes are document_chunk INSERTs and an additive envelope.entities UPDATE; the .eml archive stays read-only. --dsn/KB_DSN, --archive-root, --since/--limit/--offset, --batch-size, --apply (dry-run default). Idempotency keys on (envelope_id, chunk_index) pre-fetched scoped to --model, built correctly from the start per the plan's flagged chunk_embed.py precedent. A failed embed batch is isolated (chunks_errors, no abort) for Ollama's documented instability; a wrong embedding dimension aborts the whole run. 48 tests, DoD smoke run against live kb-postgres@PIHA confirmed wiring (archive not yet rsync'd to SOLARIA, so all 5 rows correctly reported missing_file). docs/kb/modules/05-faza-mailowa-plan.md, §5.
2026-07-22 19:01:14 +02:00
def __init__(self) -> None:
super().__init__(convert_charrefs=True)
self._tag_stack: list[tuple[str, bool, bool]] = []
feat(mail-body-ingest): new job to chunk+embed gmail body content (faza mailowa Krok 2) Second full pass over the gmail .eml archive (gmail-bulk-import's first pass skipped inline text/plain and text/html on purpose). Per envelope: typed parse with compat32 fallback -> body extraction (inline text/plain preferred, HTML->text via a small stdlib HTMLParser otherwise) -> quote-strip (reply markers + `>`-quoted lines, EN/PL/Outlook patterns) -> newsletter classification (List-Unsubscribe/List-Id/Precedence, chunked but not embedded, excluded_reason='newsletter') -> Temat/Od/Data prefix from the already-backfilled entities[type=headers] (zero header re-parse) -> chunk via kb_mail.chunking -> batched embed_batch (64) -> INSERT document_chunk. In-Reply-To/References are appended as entities[type=threading] during the same read (idempotent WHERE NOT EXISTS append, 1:1 with gmail-header-backfill) -- the only DB writes are document_chunk INSERTs and an additive envelope.entities UPDATE; the .eml archive stays read-only. --dsn/KB_DSN, --archive-root, --since/--limit/--offset, --batch-size, --apply (dry-run default). Idempotency keys on (envelope_id, chunk_index) pre-fetched scoped to --model, built correctly from the start per the plan's flagged chunk_embed.py precedent. A failed embed batch is isolated (chunks_errors, no abort) for Ollama's documented instability; a wrong embedding dimension aborts the whole run. 48 tests, DoD smoke run against live kb-postgres@PIHA confirmed wiring (archive not yet rsync'd to SOLARIA, so all 5 rows correctly reported missing_file). docs/kb/modules/05-faza-mailowa-plan.md, §5.
2026-07-22 19:01:14 +02:00
self._skip_depth = 0
self._quote_depth = 0
self._parts: list[str] = []
def _visible(self) -> bool:
return self._skip_depth == 0 and self._quote_depth == 0
def handle_starttag(self, tag: str, attrs: list) -> None:
if tag in self._VOID_ELEMENTS:
if self._visible() and tag in self._BLOCK_TAGS:
self._parts.append("\n")
return
feat(mail-body-ingest): new job to chunk+embed gmail body content (faza mailowa Krok 2) Second full pass over the gmail .eml archive (gmail-bulk-import's first pass skipped inline text/plain and text/html on purpose). Per envelope: typed parse with compat32 fallback -> body extraction (inline text/plain preferred, HTML->text via a small stdlib HTMLParser otherwise) -> quote-strip (reply markers + `>`-quoted lines, EN/PL/Outlook patterns) -> newsletter classification (List-Unsubscribe/List-Id/Precedence, chunked but not embedded, excluded_reason='newsletter') -> Temat/Od/Data prefix from the already-backfilled entities[type=headers] (zero header re-parse) -> chunk via kb_mail.chunking -> batched embed_batch (64) -> INSERT document_chunk. In-Reply-To/References are appended as entities[type=threading] during the same read (idempotent WHERE NOT EXISTS append, 1:1 with gmail-header-backfill) -- the only DB writes are document_chunk INSERTs and an additive envelope.entities UPDATE; the .eml archive stays read-only. --dsn/KB_DSN, --archive-root, --since/--limit/--offset, --batch-size, --apply (dry-run default). Idempotency keys on (envelope_id, chunk_index) pre-fetched scoped to --model, built correctly from the start per the plan's flagged chunk_embed.py precedent. A failed embed batch is isolated (chunks_errors, no abort) for Ollama's documented instability; a wrong embedding dimension aborts the whole run. 48 tests, DoD smoke run against live kb-postgres@PIHA confirmed wiring (archive not yet rsync'd to SOLARIA, so all 5 rows correctly reported missing_file). docs/kb/modules/05-faza-mailowa-plan.md, §5.
2026-07-22 19:01:14 +02:00
was_visible = self._visible()
starts_skip = tag in self._SKIP_CONTENT_TAGS
starts_quote = tag == "blockquote" or (
tag == "div" and "gmail_quote" in (dict(attrs).get("class") or "").split()
)
if starts_skip:
self._skip_depth += 1
if starts_quote:
self._quote_depth += 1
self._tag_stack.append((tag, starts_skip, starts_quote))
feat(mail-body-ingest): new job to chunk+embed gmail body content (faza mailowa Krok 2) Second full pass over the gmail .eml archive (gmail-bulk-import's first pass skipped inline text/plain and text/html on purpose). Per envelope: typed parse with compat32 fallback -> body extraction (inline text/plain preferred, HTML->text via a small stdlib HTMLParser otherwise) -> quote-strip (reply markers + `>`-quoted lines, EN/PL/Outlook patterns) -> newsletter classification (List-Unsubscribe/List-Id/Precedence, chunked but not embedded, excluded_reason='newsletter') -> Temat/Od/Data prefix from the already-backfilled entities[type=headers] (zero header re-parse) -> chunk via kb_mail.chunking -> batched embed_batch (64) -> INSERT document_chunk. In-Reply-To/References are appended as entities[type=threading] during the same read (idempotent WHERE NOT EXISTS append, 1:1 with gmail-header-backfill) -- the only DB writes are document_chunk INSERTs and an additive envelope.entities UPDATE; the .eml archive stays read-only. --dsn/KB_DSN, --archive-root, --since/--limit/--offset, --batch-size, --apply (dry-run default). Idempotency keys on (envelope_id, chunk_index) pre-fetched scoped to --model, built correctly from the start per the plan's flagged chunk_embed.py precedent. A failed embed batch is isolated (chunks_errors, no abort) for Ollama's documented instability; a wrong embedding dimension aborts the whole run. 48 tests, DoD smoke run against live kb-postgres@PIHA confirmed wiring (archive not yet rsync'd to SOLARIA, so all 5 rows correctly reported missing_file). docs/kb/modules/05-faza-mailowa-plan.md, §5.
2026-07-22 19:01:14 +02:00
if was_visible and tag in self._BLOCK_TAGS:
self._parts.append("\n")
def handle_endtag(self, tag: str) -> None:
for i in range(len(self._tag_stack) - 1, -1, -1):
if self._tag_stack[i][0] == tag:
_, starts_skip, starts_quote = self._tag_stack[i]
del self._tag_stack[i:] # also drops anything opened-but-never-closed after it
if starts_skip and self._skip_depth > 0:
self._skip_depth -= 1
if starts_quote and self._quote_depth > 0:
self._quote_depth -= 1
if self._visible() and tag in self._BLOCK_TAGS:
self._parts.append("\n")
return
# no matching open tag on the stack -- stray/unbalanced closing tag, ignore
feat(mail-body-ingest): new job to chunk+embed gmail body content (faza mailowa Krok 2) Second full pass over the gmail .eml archive (gmail-bulk-import's first pass skipped inline text/plain and text/html on purpose). Per envelope: typed parse with compat32 fallback -> body extraction (inline text/plain preferred, HTML->text via a small stdlib HTMLParser otherwise) -> quote-strip (reply markers + `>`-quoted lines, EN/PL/Outlook patterns) -> newsletter classification (List-Unsubscribe/List-Id/Precedence, chunked but not embedded, excluded_reason='newsletter') -> Temat/Od/Data prefix from the already-backfilled entities[type=headers] (zero header re-parse) -> chunk via kb_mail.chunking -> batched embed_batch (64) -> INSERT document_chunk. In-Reply-To/References are appended as entities[type=threading] during the same read (idempotent WHERE NOT EXISTS append, 1:1 with gmail-header-backfill) -- the only DB writes are document_chunk INSERTs and an additive envelope.entities UPDATE; the .eml archive stays read-only. --dsn/KB_DSN, --archive-root, --since/--limit/--offset, --batch-size, --apply (dry-run default). Idempotency keys on (envelope_id, chunk_index) pre-fetched scoped to --model, built correctly from the start per the plan's flagged chunk_embed.py precedent. A failed embed batch is isolated (chunks_errors, no abort) for Ollama's documented instability; a wrong embedding dimension aborts the whole run. 48 tests, DoD smoke run against live kb-postgres@PIHA confirmed wiring (archive not yet rsync'd to SOLARIA, so all 5 rows correctly reported missing_file). docs/kb/modules/05-faza-mailowa-plan.md, §5.
2026-07-22 19:01:14 +02:00
def handle_data(self, data: str) -> None:
if self._visible():
self._parts.append(data)
def text(self) -> str:
raw = "".join(self._parts)
raw = re.sub(r"[ \t]+", " ", raw)
raw = re.sub(r"\n[ \t]*\n+", "\n\n", raw)
return raw.strip()
def html_to_text(html: str) -> str:
parser = _MailHTMLTextExtractor()
parser.feed(html)
parser.close()
return parser.text()
def strip_quotes(text: str) -> tuple[str, int]:
"""Decision 2: truncate at the earliest reply marker, then drop remaining `>`-quoted lines.
Returns (clean_text, chars_stripped)."""
original_len = len(text)
earliest: Optional[int] = None
for pattern in _REPLY_MARKER_PATTERNS:
m = pattern.search(text)
if m and (earliest is None or m.start() < earliest):
earliest = m.start()
if earliest is not None:
text = text[:earliest]
lines = [line for line in text.split("\n") if not _QUOTE_LINE_RE.match(line)]
text = "\n".join(lines).strip()
return text, original_len - len(text)
def parse_message(raw: bytes) -> email.message.Message:
"""Typed parse (RFC 2047 decoding) with a compat32 fallback for the ~9/225030 messages the
typed parser rejects (same failure mode `gmail_header_backfill.parse_headers_fallback`
documents) both give a usable `Message`/`EmailMessage` for body walking and header access."""
try:
return email.message_from_bytes(raw, policy=email.policy.default)
except Exception:
return email.message_from_bytes(raw, policy=email.policy.compat32)
def _inline_text_parts(msg: email.message.Message, content_type: str):
for part in msg.walk():
if part.get_content_maintype() == "multipart":
continue
if part.get_content_type() != content_type:
continue
disposition = part.get_content_disposition() or ""
if disposition == "attachment" or part.get_filename():
continue
yield part
def _decode_payload(part: email.message.Message) -> Optional[str]:
payload = part.get_payload(decode=True)
if payload is None:
return None
charset = part.get_content_charset() or "utf-8"
try:
text = payload.decode(charset, errors="replace")
except (LookupError, UnicodeDecodeError):
text = payload.decode("utf-8", errors="replace")
return _sanitize(text)
def extract_body(msg: email.message.Message) -> str:
"""Inline text/plain preferred; HTML->text fallback when the mail is HTML-only (plan §1.3:
15% of the corpus). Neither present (attachment-only mail) -> ''."""
for part in _inline_text_parts(msg, "text/plain"):
text = _decode_payload(part)
if text is not None:
return text
for part in _inline_text_parts(msg, "text/html"):
html = _decode_payload(part)
if html is not None:
return html_to_text(html)
return ""
def is_newsletter(msg: email.message.Message) -> bool:
"""Decision 4: List-Unsubscribe or List-Id present, or Precedence: bulk|list."""
if msg.get("List-Unsubscribe") is not None or msg.get("List-Id") is not None:
return True
precedence = msg.get("Precedence")
if precedence is not None and str(precedence).strip().lower() in ("bulk", "list"):
return True
return False
def extract_threading(msg: email.message.Message) -> dict:
"""Decision 10: raw In-Reply-To / References, angle brackets stripped (matches
`envelope.id`'s bare-Message-ID convention from `gmail_bulk_import._message_id`)."""
in_reply_to_raw = msg.get("In-Reply-To")
references_raw = msg.get("References")
in_reply_to = None
if in_reply_to_raw:
m = _ID_TOKEN_RE.search(_sanitize(str(in_reply_to_raw)) or "")
if m:
in_reply_to = m.group(0)[1:-1]
references: list[str] = []
if references_raw:
for m in _ID_TOKEN_RE.finditer(_sanitize(str(references_raw)) or ""):
references.append(m.group(0)[1:-1])
return {"type": "threading", "in_reply_to": in_reply_to, "references": references}
def build_prefix(entities: list, ts: datetime) -> str:
"""Decision 3: 'Temat: ... | Od: ... | Data: YYYY-MM-DD' from entities[type=headers]
(already in DB zero header re-parse) + envelope.ts for the date."""
headers = next(
(e for e in (entities or []) if isinstance(e, dict) and e.get("type") == "headers"), {}
)
subject = headers.get("subject") or "(brak tematu)"
from_obj = headers.get("from")
from_display = (from_obj.get("name") or from_obj.get("address") or "?") if from_obj else "?"
return f"Temat: {subject} | Od: {from_display} | Data: {ts:%Y-%m-%d}"
def _decode_jsonb(value: object) -> object:
if value is None:
return None
if isinstance(value, str):
return json.loads(value)
return value
def _has_threading(entities: list) -> bool:
return any(isinstance(e, dict) and e.get("type") == "threading" for e in entities or [])
async def fetch_envelopes(
conn: asyncpg.Connection, since: Optional[datetime], limit: Optional[int], offset: Optional[int]
) -> list:
"""`source='gmail'` envelopes, ordered by id for stable --limit/--offset slicing."""
query = "SELECT id, ts, raw_ref, entities FROM envelope WHERE source = 'gmail'"
params: list = []
if since is not None:
params.append(since)
query += f" AND ts >= ${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)}"
return await conn.fetch(query, *params)
async def fetch_existing_chunk_keys(conn: asyncpg.Connection, model: str) -> set[tuple[str, int]]:
"""(envelope_id, chunk_index) pairs already inserted with this model — idempotency + dry-run
preview. The query filters `WHERE model = $1`, so the pair need not redundantly carry model."""
rows = await conn.fetch(
"SELECT envelope_id, chunk_index FROM document_chunk WHERE model = $1", model
)
return {(r["envelope_id"], r["chunk_index"]) for r in rows}
async def insert_chunk(
conn: asyncpg.Connection, envelope_id: str, chunk_index: int, text: str,
embedding: Optional[list[float]], model: str, excluded_reason: Optional[str] = None,
) -> str:
vector_literal = _vector_literal(embedding) if embedding is not None else None
return await conn.execute(
_CHUNK_INSERT_SQL, envelope_id, chunk_index, text, vector_literal, model, excluded_reason
)
def _rows_affected(command_tag: str) -> int:
return int(command_tag.rsplit(" ", 1)[-1])
def _new_stats() -> dict:
return {
"mails_scanned": 0,
"missing_file": 0,
"read_errors": 0,
"parse_errors": 0,
"body_empty": 0,
"mails_chunked": 0,
"chunks_total": 0,
"chunks_already_embedded": 0,
"chunks_inserted": 0,
"chunks_newsletter_flagged": 0,
"chunks_conflict_skipped": 0,
"chunks_errors": 0,
"quoted_chars_stripped_total": 0,
"embed_calls": 0,
"embed_seconds_total": 0.0,
"threading_updated": 0,
"threading_already_present": 0,
}
async def run(
dsn: str,
archive_root: Path = DEFAULT_ARCHIVE_ROOT,
ollama_url: str = DEFAULT_OLLAMA_URL,
model: str = DEFAULT_MODEL,
since: Optional[datetime] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
apply: bool = False,
batch_size: int = DEFAULT_BATCH_SIZE,
) -> dict:
"""Process one --limit/--offset (optionally --since-filtered) slice of `source='gmail'`
envelopes. dry-run (apply=False): parse + quote-strip + classify + chunk + count, zero
Ollama calls, zero writes (including the threading UPDATE). Stats must always balance:
mails_scanned = missing_file + read_errors + parse_errors + body_empty + mails_chunked
chunks_total = chunks_inserted + chunks_newsletter_flagged + chunks_already_embedded
+ chunks_conflict_skipped + chunks_errors
"""
stats = _new_stats()
conn = await asyncpg.connect(dsn)
try:
envelopes = await fetch_envelopes(conn, since, limit, offset)
existing = await fetch_existing_chunk_keys(conn, model)
session: Optional[aiohttp.ClientSession] = None
if apply:
session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120))
if not await check_ollama_health(session, ollama_url, 3.0):
_log.warning("ollama_down_at_start", ollama_url=ollama_url)
embed_buffer: list[tuple[str, int, str]] = []
threading_pending: list[tuple[str, str]] = []
async def flush_embed_buffer() -> None:
if not embed_buffer:
return
texts = [t for (_eid, _idx, t) in embed_buffer]
try:
embeddings, elapsed = await embed_batch(session, ollama_url, model, texts)
except EmbeddingDimensionError:
raise
except aiohttp.ClientError:
_log.warning("skip.embed_batch_error", count=len(embed_buffer), exc_info=True)
stats["chunks_errors"] += len(embed_buffer)
embed_buffer.clear()
return
stats["embed_calls"] += 1
stats["embed_seconds_total"] += elapsed
for (eid, idx, chunk), embedding in zip(embed_buffer, embeddings):
try:
command_tag = await insert_chunk(conn, eid, idx, chunk, embedding, model)
except Exception:
_log.warning("skip.insert_error", envelope_id=eid, chunk_index=idx, exc_info=True)
stats["chunks_errors"] += 1
continue
existing.add((eid, idx))
if _rows_affected(command_tag) == 0:
_log.warning("chunk.conflict_skipped", envelope_id=eid, chunk_index=idx, model=model)
stats["chunks_conflict_skipped"] += 1
else:
stats["chunks_inserted"] += 1
embed_buffer.clear()
async def flush_threading() -> None:
if apply and threading_pending:
await conn.executemany(_THREADING_UPDATE_SQL, threading_pending)
threading_pending.clear()
try:
for row in envelopes:
stats["mails_scanned"] += 1
envelope_id = row["id"]
ts = row["ts"]
entities = _decode_jsonb(row["entities"]) or []
eml_path = archive_root / row["raw_ref"]
try:
raw = eml_path.read_bytes()
except FileNotFoundError:
_log.info("skip.missing_file", envelope_id=envelope_id, expected_path=str(eml_path))
stats["missing_file"] += 1
continue
except OSError:
_log.warning("skip.read_error", envelope_id=envelope_id, path=str(eml_path))
stats["read_errors"] += 1
continue
try:
msg = parse_message(raw)
body = extract_body(msg)
clean_body, stripped = strip_quotes(body)
stats["quoted_chars_stripped_total"] += stripped
newsletter = is_newsletter(msg)
threading_patch = extract_threading(msg)
except Exception:
_log.warning("skip.parse_error", envelope_id=envelope_id, exc_info=True)
stats["parse_errors"] += 1
continue
if not _has_threading(entities):
stats["threading_updated"] += 1
if apply:
threading_pending.append((envelope_id, json.dumps([threading_patch])))
if len(threading_pending) >= THREADING_UPDATE_BATCH_SIZE:
await flush_threading()
else:
stats["threading_already_present"] += 1
if not clean_body.strip():
stats["body_empty"] += 1
continue
stats["mails_chunked"] += 1
prefix = build_prefix(entities, ts)
chunks = chunk_text(f"{prefix}\n\n{clean_body}", TARGET_CHARS, OVERLAP_CHARS)
for idx, chunk in enumerate(chunks):
stats["chunks_total"] += 1
key = (envelope_id, idx)
if key in existing:
stats["chunks_already_embedded"] += 1
continue
if newsletter:
if not apply:
stats["chunks_newsletter_flagged"] += 1
continue
try:
command_tag = await insert_chunk(
conn, envelope_id, idx, chunk, None, model, excluded_reason="newsletter"
)
except Exception:
_log.warning(
"skip.insert_error", envelope_id=envelope_id, chunk_index=idx, exc_info=True
)
stats["chunks_errors"] += 1
continue
existing.add(key)
if _rows_affected(command_tag) == 0:
stats["chunks_conflict_skipped"] += 1
else:
stats["chunks_newsletter_flagged"] += 1
continue
if not apply:
stats["chunks_inserted"] += 1
continue
embed_buffer.append((envelope_id, idx, chunk))
if len(embed_buffer) >= batch_size:
await flush_embed_buffer()
if stats["mails_scanned"] % 500 == 0:
_log.info("progress", **stats)
if apply:
await flush_embed_buffer()
await flush_threading()
finally:
if session is not None:
await session.close()
finally:
await conn.close()
balance_mails = (
stats["missing_file"] + stats["read_errors"] + stats["parse_errors"]
+ stats["body_empty"] + stats["mails_chunked"]
)
balance_chunks = (
stats["chunks_inserted"] + stats["chunks_newsletter_flagged"] + stats["chunks_already_embedded"]
+ stats["chunks_conflict_skipped"] + stats["chunks_errors"]
)
if balance_mails != stats["mails_scanned"] or balance_chunks != stats["chunks_total"]:
_log.error("stats_mismatch", **stats)
_log.info("run_complete", apply=apply, model=model, **stats)
return stats
def _parse_since(value: str) -> datetime:
return datetime.strptime(value, "%Y-%m-%d").replace(tzinfo=timezone.utc)
def main() -> None:
parser = argparse.ArgumentParser(
description="Chunk + embed source='gmail' envelope body content into document_chunk "
"(module 5, faza mailowa — plan §5, Krok 2)."
)
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("--archive-root", type=Path, default=DEFAULT_ARCHIVE_ROOT,
help=f"Mail .eml archive root (default: {DEFAULT_ARCHIVE_ROOT})")
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("--model", default=os.environ.get("OLLAMA_EMBED_MODEL", DEFAULT_MODEL),
help=f"Ollama embedding model (default: {DEFAULT_MODEL})")
parser.add_argument("--since", type=_parse_since, default=None, metavar="YYYY-MM-DD",
help="Only envelopes with ts >= this date (for staged runs, plan Decyzja 9)")
parser.add_argument("--limit", type=int, default=None,
help="Max envelopes to process (default: all)")
parser.add_argument("--offset", type=int, default=0,
help="Slice offset, ordered by envelope id (default: 0)")
parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE,
help=f"Ollama /api/embed batch size (default: {DEFAULT_BATCH_SIZE})")
parser.add_argument("--apply", action="store_true",
help="Actually call Ollama, insert chunks, and update threading entities. "
"Default is dry-run (parse + classify + chunk + count only).")
args = parser.parse_args()
if not args.dsn:
_log.error("missing_dsn", hint="pass --dsn or set KB_DSN")
sys.exit(1)
if not args.archive_root.is_dir():
_log.error("archive_root_not_found", path=str(args.archive_root))
sys.exit(1)
try:
stats = asyncio.run(
run(
dsn=args.dsn,
archive_root=args.archive_root,
ollama_url=args.ollama_url,
model=args.model,
since=args.since,
limit=args.limit,
offset=args.offset,
apply=args.apply,
batch_size=args.batch_size,
)
)
except EmbeddingDimensionError as exc:
_log.error("dim_mismatch_abort", error=str(exc))
sys.exit(1)
mode = "APPLY" if args.apply else "DRY-RUN"
avg_embed_ms = (
1000 * stats["embed_seconds_total"] / stats["embed_calls"] if stats["embed_calls"] else 0.0
)
_log.info("summary", mode=mode, avg_embed_ms_per_batch=round(avg_embed_ms, 2), **stats)
balanced = (
stats["mails_scanned"] == (
stats["missing_file"] + stats["read_errors"] + stats["parse_errors"]
+ stats["body_empty"] + stats["mails_chunked"]
)
and stats["chunks_total"] == (
stats["chunks_inserted"] + stats["chunks_newsletter_flagged"] + stats["chunks_already_embedded"]
+ stats["chunks_conflict_skipped"] + stats["chunks_errors"]
)
)
failed = (
stats["read_errors"] > 0 or stats["parse_errors"] > 0 or stats["missing_file"] > 0
or stats["chunks_errors"] > 0 or stats["chunks_conflict_skipped"] > 0
or not balanced
)
sys.exit(1 if failed else 0)
if __name__ == "__main__":
main()