From e06cc105711e4453264d7100a4d4a182026cb56c Mon Sep 17 00:00:00 2001 From: oskar Date: Wed, 22 Jul 2026 19:01:14 +0200 Subject: [PATCH] feat(mail-body-ingest): new job to chunk+embed gmail body content (faza mailowa Krok 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- jobs/mail-body-ingest/pyproject.toml | 31 + .../src/mail_body_ingest/ingest.py | 627 ++++++++++++++++++ jobs/mail-body-ingest/tests/test_ingest.py | 577 ++++++++++++++++ 3 files changed, 1235 insertions(+) create mode 100644 jobs/mail-body-ingest/pyproject.toml create mode 100644 jobs/mail-body-ingest/src/mail_body_ingest/ingest.py create mode 100644 jobs/mail-body-ingest/tests/test_ingest.py diff --git a/jobs/mail-body-ingest/pyproject.toml b/jobs/mail-body-ingest/pyproject.toml new file mode 100644 index 0000000..9ed1ead --- /dev/null +++ b/jobs/mail-body-ingest/pyproject.toml @@ -0,0 +1,31 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "mail-body-ingest" +version = "0.1.0" +requires-python = ">=3.11" +dependencies = [ + "asyncpg>=0.29", + "aiohttp>=3.9", + "structlog>=24.1", + "kb-mail", + "kb-retrieval", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.1", + "pytest-asyncio>=0.23", +] + +[project.scripts] +mail-body-ingest = "mail_body_ingest.ingest:main" + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] diff --git a/jobs/mail-body-ingest/src/mail_body_ingest/ingest.py b/jobs/mail-body-ingest/src/mail_body_ingest/ingest.py new file mode 100644 index 0000000..e903adc --- /dev/null +++ b/jobs/mail-body-ingest/src/mail_body_ingest/ingest.py @@ -0,0 +1,627 @@ +"""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:@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. + + LIFO tag-stack assumption (no tag-name matching on close): correct for well-nested HTML, + which is what mail clients emit in practice; a truly malformed document may under/over-skip + at the margins -- acceptable for a best-effort heuristic with zero new dependencies. + """ + + _SKIP_CONTENT_TAGS = {"style", "script", "head"} + _BLOCK_TAGS = {"p", "div", "tr", "li", "h1", "h2", "h3", "h4", "h5", "h6", "br"} + + def __init__(self) -> None: + super().__init__(convert_charrefs=True) + self._tag_stack: list[tuple[bool, bool]] = [] + 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: + 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((starts_skip, starts_quote)) + if was_visible and tag in self._BLOCK_TAGS: + self._parts.append("\n") + + def handle_endtag(self, tag: str) -> None: + if not self._tag_stack: + return + starts_skip, starts_quote = self._tag_stack.pop() + 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") + + 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() diff --git a/jobs/mail-body-ingest/tests/test_ingest.py b/jobs/mail-body-ingest/tests/test_ingest.py new file mode 100644 index 0000000..48e4bba --- /dev/null +++ b/jobs/mail-body-ingest/tests/test_ingest.py @@ -0,0 +1,577 @@ +"""Unit tests for the mail body ingest job — no DB, no real HTTP, no real Ollama. + +Fixture .eml bytes are built inline (email.mime helpers / raw byte literals), matching the +convention in jobs/gmail-header-backfill/tests/test_backfill.py rather than separate fixture +files -- easier to see exactly what each test is asserting against.""" +from __future__ import annotations + +import json +from datetime import datetime, timezone +from email.mime.application import MIMEApplication +from email.mime.multipart import MIMEMultipart +from email.mime.text import MIMEText + +import aiohttp +import pytest + +from kb_retrieval.embed import EmbeddingDimensionError +from mail_body_ingest.ingest import ( + _CHUNK_INSERT_SQL, + _decode_jsonb, + _has_threading, + build_prefix, + extract_body, + extract_threading, + fetch_envelopes, + fetch_existing_chunk_keys, + html_to_text, + insert_chunk, + is_newsletter, + parse_message, + run, + strip_quotes, +) + + +def _plain_eml(headers: dict, body: str) -> bytes: + msg = MIMEText(body, "plain", "utf-8") + for k, v in headers.items(): + msg[k] = v + return msg.as_bytes() + + +def _html_eml(headers: dict, html_body: str) -> bytes: + msg = MIMEText(html_body, "html", "utf-8") + for k, v in headers.items(): + msg[k] = v + return msg.as_bytes() + + +def _multipart_with_attachment_eml(headers: dict, body: str) -> bytes: + msg = MIMEMultipart() + for k, v in headers.items(): + msg[k] = v + msg.attach(MIMEText(body, "plain", "utf-8")) + attachment = MIMEApplication(b"%PDF-fake", Name="doc.pdf") + attachment["Content-Disposition"] = 'attachment; filename="doc.pdf"' + msg.attach(attachment) + return msg.as_bytes() + + +# Pattern 1 from the gmail-header-backfill diagnosis (2026-07-14): RFC 2047 encoded-word whose +# decoded text contains a newline in a display name -- policy.default raises ValueError. +_CRLF_ENCODED_WORD_EML = ( + b"From: Rekrutacja z =?utf-8?Q?ExampleCorp=0A?= \r\n" + b"To: Jan Kowalski \r\n" + b"Subject: Nowe oferty\r\n" + b"Date: Mon, 15 Sep 2014 13:24:18 +0000\r\n" + b"\r\nbody" +) + + +class TestStripQuotes: + def test_no_marker_no_quotes_returns_unchanged(self): + clean, stripped = strip_quotes("Hello,\n\nJust checking in.\n\nRegards,\nA") + assert clean == "Hello,\n\nJust checking in.\n\nRegards,\nA" + assert stripped == 0 + + def test_gmail_wrote_marker_truncates_everything_after(self): + text = "Sure, sounds good.\n\nOn Tue, 10 Jun 2025 at 12:00, Alice wrote:\n> old text\n> more old" + clean, stripped = strip_quotes(text) + assert clean == "Sure, sounds good." + assert stripped > 0 + + def test_polish_dnia_napisal_marker(self): + text = "OK, dzieki.\n\nDnia 10 czerwca 2025 Jan napisał(a):\n> stara tresc" + clean, _ = strip_quotes(text) + assert clean == "OK, dzieki." + + def test_polish_w_dniu_pisze_marker(self): + text = "Dobrze.\n\nW dniu 10.06.2025 Anna pisze:\n> cos tam" + clean, _ = strip_quotes(text) + assert clean == "Dobrze." + + def test_outlook_original_message_marker(self): + text = "Ok, dzieki za info.\n\n-----Original Message-----\nFrom: bob@example.com\nold content" + clean, _ = strip_quotes(text) + assert clean == "Ok, dzieki za info." + + def test_outlook_underscore_separator(self): + text = "Widziane, dzieki.\n\n________________________________\nFrom: c@example.com\nstuff" + clean, _ = strip_quotes(text) + assert clean == "Widziane, dzieki." + + def test_bare_quote_lines_without_marker_are_dropped(self): + text = "My reply here.\n> quoted line one\n> quoted line two\nNot quoted trailing line" + clean, stripped = strip_quotes(text) + assert ">" not in clean + assert "My reply here." in clean + assert stripped > 0 + + def test_chars_stripped_counts_removed_characters(self): + text = "short reply\n\nOn 1 Jan 2020 X wrote:\n> a very long quoted paragraph indeed" + _clean, stripped = strip_quotes(text) + assert stripped == len(text) - len("short reply") + + +class TestHtmlToText: + def test_skips_style_and_script(self): + html = "

Real text

" + assert "color:red" not in html_to_text(html) + assert "evil()" not in html_to_text(html) + assert "Real text" in html_to_text(html) + + def test_skips_blockquote_subtree(self): + html = "
Reply text
quoted old text
" + text = html_to_text(html) + assert "Reply text" in text + assert "quoted old text" not in text + + def test_skips_gmail_quote_div(self): + html = ( + '
My reply
' + '
On Tue, X wrote:
nested quoted content
' + ) + text = html_to_text(html) + assert "My reply" in text + assert "nested quoted content" not in text + + def test_plain_div_without_gmail_quote_class_is_not_skipped(self): + html = '
Still visible
' + assert "Still visible" in html_to_text(html) + + +class TestIsNewsletter: + def test_list_unsubscribe_present(self): + msg = parse_message(_plain_eml({"From": "a@b.com", "List-Unsubscribe": ""}, "hi")) + assert is_newsletter(msg) is True + + def test_list_id_present(self): + msg = parse_message(_plain_eml({"From": "a@b.com", "List-Id": "newsletter.example.com"}, "hi")) + assert is_newsletter(msg) is True + + def test_precedence_bulk(self): + msg = parse_message(_plain_eml({"From": "a@b.com", "Precedence": "bulk"}, "hi")) + assert is_newsletter(msg) is True + + def test_precedence_list_case_insensitive(self): + msg = parse_message(_plain_eml({"From": "a@b.com", "Precedence": "LIST"}, "hi")) + assert is_newsletter(msg) is True + + def test_ordinary_mail_is_not_newsletter(self): + msg = parse_message(_plain_eml({"From": "a@b.com", "Subject": "hi"}, "hi")) + assert is_newsletter(msg) is False + + +class TestExtractThreading: + def test_in_reply_to_and_references_stripped_of_brackets(self): + msg = parse_message(_plain_eml( + {"From": "a@b.com", "In-Reply-To": "", + "References": " "}, "hi", + )) + threading = extract_threading(msg) + assert threading == { + "type": "threading", "in_reply_to": "msg1@x.com", + "references": ["msg0@x.com", "msg1@x.com"], + } + + def test_missing_headers_yield_none_and_empty_list(self): + msg = parse_message(_plain_eml({"From": "a@b.com"}, "hi")) + threading = extract_threading(msg) + assert threading == {"type": "threading", "in_reply_to": None, "references": []} + + +class TestBuildPrefix: + def test_full_headers(self): + entities = [{"type": "headers", "subject": "Twoja polisa", "from": {"name": "PZU", "address": "x@pzu.pl"}}] + ts = datetime(2025, 8, 3, tzinfo=timezone.utc) + assert build_prefix(entities, ts) == "Temat: Twoja polisa | Od: PZU | Data: 2025-08-03" + + def test_missing_subject_and_from_fallback(self): + entities = [{"type": "headers", "subject": None, "from": None}] + ts = datetime(2025, 1, 1, tzinfo=timezone.utc) + prefix = build_prefix(entities, ts) + assert "(brak tematu)" in prefix + assert "Od: ?" in prefix + + def test_from_without_name_uses_address(self): + entities = [{"type": "headers", "subject": "s", "from": {"name": None, "address": "a@b.com"}}] + ts = datetime(2025, 1, 1, tzinfo=timezone.utc) + assert "Od: a@b.com" in build_prefix(entities, ts) + + def test_no_headers_entity_at_all(self): + ts = datetime(2025, 1, 1, tzinfo=timezone.utc) + prefix = build_prefix([], ts) + assert "(brak tematu)" in prefix + assert "Od: ?" in prefix + + +class TestExtractBody: + def test_plain_text_preferred(self): + msg = parse_message(_plain_eml({"From": "a@b.com"}, "hello plain body")) + assert extract_body(msg) == "hello plain body" + + def test_html_only_falls_back_to_html_to_text(self): + msg = parse_message(_html_eml({"From": "a@b.com"}, "

hello html body

")) + body = extract_body(msg) + assert "hello" in body and "html" in body and "body" in body + assert "<" not in body + + def test_attachment_only_yields_empty_body(self): + msg = MIMEApplication(b"%PDF-fake", Name="doc.pdf") + msg["From"] = "a@b.com" + assert extract_body(msg) == "" + + def test_multipart_with_attachment_extracts_only_the_text_part(self): + msg = parse_message(_multipart_with_attachment_eml({"From": "a@b.com"}, "the actual message")) + assert extract_body(msg) == "the actual message" + + +class TestParseMessage: + def test_typed_parse_succeeds_for_ordinary_mail(self): + msg = parse_message(_plain_eml({"From": "a@b.com", "Subject": "hi"}, "body")) + assert msg.get("Subject") == "hi" + + def test_falls_back_to_compat32_on_typed_parse_failure(self): + # This exact byte pattern raises ValueError under policy.default (see backfill's + # diagnosis) -- parse_message must still return a usable Message via the fallback. + msg = parse_message(_CRLF_ENCODED_WORD_EML) + assert msg.get("Subject") is not None + + +class TestDecodeJsonbAndThreadingHelpers: + def test_decode_jsonb_handles_str_and_object(self): + assert _decode_jsonb('[{"type": "headers"}]') == [{"type": "headers"}] + assert _decode_jsonb([{"type": "headers"}]) == [{"type": "headers"}] + assert _decode_jsonb(None) is None + + def test_has_threading(self): + assert _has_threading([{"type": "threading"}]) is True + assert _has_threading([{"type": "headers"}]) is False + assert _has_threading([]) is False + + +class TestInsertSql: + def test_on_conflict_target_includes_model(self): + assert "ON CONFLICT (envelope_id, chunk_index, model)" in _CHUNK_INSERT_SQL + + +class TestInsertChunk: + async def test_newsletter_chunk_inserts_null_embedding_with_reason(self): + conn = _FakeConn() + await insert_chunk(conn, "msgid@x", 0, "junk text", None, "bge-m3", excluded_reason="newsletter") + assert conn.execute_calls == [("msgid@x", 0, "junk text", None, "bge-m3", "newsletter")] + + async def test_normal_chunk_inserts_vector_with_no_reason(self): + conn = _FakeConn() + await insert_chunk(conn, "msgid@x", 0, "real text", [0.1, 0.2], "bge-m3") + assert conn.execute_calls == [("msgid@x", 0, "real text", "[0.1,0.2]", "bge-m3", None)] + + +class TestFetchHelpers: + async def test_fetch_envelopes_no_filters(self): + conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")]) + rows = await fetch_envelopes(conn, since=None, limit=None, offset=None) + assert rows == [_env_row("m1@x", "hi")] + query, params = conn.queries[-1] + assert "ORDER BY id" in query + assert "LIMIT" not in query and "OFFSET" not in query + + async def test_fetch_envelopes_with_since_limit_offset(self): + conn = _FakeConn(envelopes=[]) + await fetch_envelopes(conn, since=datetime(2025, 7, 1, tzinfo=timezone.utc), limit=10, offset=5) + query, params = conn.queries[-1] + assert "ts >= $1" in query + assert "LIMIT $2" in query + assert "OFFSET $3" in query + assert params == (datetime(2025, 7, 1, tzinfo=timezone.utc), 10, 5) + + async def test_fetch_existing_chunk_keys(self): + conn = _FakeConn(existing_keys=[("m1@x", 0), ("m1@x", 1)]) + keys = await fetch_existing_chunk_keys(conn, model="bge-m3") + assert keys == {("m1@x", 0), ("m1@x", 1)} + + +def _env_row(envelope_id: str, subject: str, ts: datetime = datetime(2025, 8, 1, tzinfo=timezone.utc)): + entities = [{"type": "headers", "subject": subject, "from": {"name": "A", "address": "a@b.com"}}] + return {"id": envelope_id, "ts": ts, "raw_ref": f"gmail/2025/08/{envelope_id}.eml", "entities": json.dumps(entities)} + + +class _FakeConn: + """Mirrors documents_ingest's/backfill's `_FakeConn` test style. `existing_keys` were + "written" under `existing_model` -- a fetch for a different model must not see them, + same as the real `WHERE model = $1` filter.""" + + def __init__(self, envelopes=None, existing_keys=None, existing_model="bge-m3", execute_results=None): + self._envelopes = envelopes or [] + self._existing_keys = list(existing_keys 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] = [] + self.executemany_calls: list[tuple] = [] + self.queries: list[tuple] = [] + + async def fetch(self, query, *params): + self.queries.append((query, params)) + if "FROM document_chunk" in query: + if params and params[0] != self._existing_model: + return [] + return [{"envelope_id": eid, "chunk_index": idx} for eid, idx in self._existing_keys] + if "FROM envelope" in query: + return self._envelopes + raise AssertionError(f"unexpected query: {query}") + + 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 executemany(self, query, rows): + self.executemany_calls.append((query, list(rows))) + + async def close(self): + pass + + +class _FakeEmbedResponse: + 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 aiohttp.ClientConnectionError(f"simulated HTTP {self._status}") + + +class _FakeTagsResponse: + def __init__(self, status=200): + self.status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + +class _FakeOllamaSession: + def __init__(self, dim=1024, fail_batches=False, health_up=True): + self._dim = dim + self._fail_batches = fail_batches + self._health_up = health_up + self.requests: list[dict] = [] + self.closed = False + + def post(self, url, json): + assert url.endswith("/api/embed") + self.requests.append({"url": url, "json": json}) + if self._fail_batches: + return _FakeEmbedResponse({}, status=500) + n = len(json["input"]) + return _FakeEmbedResponse({"embeddings": [[0.01] * self._dim for _ in range(n)]}) + + def get(self, url, timeout=None): + return _FakeTagsResponse(200 if self._health_up else 500) + + async def close(self): + self.closed = True + + +def _write_eml(archive_root, rel_path: str, raw: bytes) -> None: + dest = archive_root / rel_path + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(raw) + + +class TestRun: + def _patch(self, monkeypatch, conn, ollama_session): + async def _fake_connect(dsn): + return conn + monkeypatch.setattr("mail_body_ingest.ingest.asyncpg.connect", _fake_connect) + + def _fake_session_factory(*args, **kwargs): + return ollama_session + monkeypatch.setattr("mail_body_ingest.ingest.aiohttp.ClientSession", _fake_session_factory) + + async def test_dry_run_counts_without_calling_ollama_or_db(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" + )) + conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")]) + ollama = _FakeOllamaSession() + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=False) + + assert stats["mails_scanned"] == 1 + assert stats["mails_chunked"] == 1 + assert stats["chunks_total"] == 1 + assert stats["chunks_inserted"] == 1 + assert stats["threading_updated"] == 1 # would-update, not applied + assert ollama.requests == [] + assert conn.execute_calls == [] + assert conn.executemany_calls == [] + + async def test_apply_embeds_inserts_and_updates_threading(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" + )) + conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")]) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + + assert stats["chunks_inserted"] == 1 + assert stats["embed_calls"] == 1 + assert len(conn.execute_calls) == 1 + assert len(ollama.requests) == 1 + assert len(conn.executemany_calls) == 1 # threading update batch flushed + + async def test_idempotent_skips_already_inserted_chunks(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" + )) + conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")], existing_keys=[("m1@x", 0)]) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + + assert stats["chunks_already_embedded"] == 1 + assert stats["chunks_inserted"] == 0 + assert ollama.requests == [] + + async def test_rerun_after_apply_inserts_nothing_new(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" + )) + conn1 = _FakeConn(envelopes=[_env_row("m1@x", "hi")]) + ollama1 = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn1, ollama1) + first = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + assert first["chunks_inserted"] == 1 + + conn2 = _FakeConn(envelopes=[_env_row("m1@x", "hi")], existing_keys=[("m1@x", 0)]) + ollama2 = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn2, ollama2) + second = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + + assert second["chunks_inserted"] == 0 + assert second["chunks_already_embedded"] == 1 + assert ollama2.requests == [] + + async def test_newsletter_chunk_flagged_and_not_embedded(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/nl@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "Promo", "List-Unsubscribe": "", + "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, + "Kup teraz z rabatem!", + )) + conn = _FakeConn(envelopes=[_env_row("nl@x", "Promo")]) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + + assert stats["chunks_newsletter_flagged"] == 1 + assert stats["chunks_inserted"] == 0 + assert ollama.requests == [] # no Ollama call for newsletter chunks + assert conn.execute_calls[0][3] is None # embedding column is NULL + assert conn.execute_calls[0][5] == "newsletter" + + async def test_body_empty_after_quote_strip_still_updates_threading(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/empty@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "Re: x", "In-Reply-To": "", + "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, + "On Tue wrote:\n> everything is quoted", + )) + conn = _FakeConn(envelopes=[_env_row("empty@x", "Re: x")]) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + + assert stats["body_empty"] == 1 + assert stats["chunks_total"] == 0 + assert stats["threading_updated"] == 1 + assert len(conn.executemany_calls) == 1 + patch = json.loads(conn.executemany_calls[0][1][0][1]) + assert patch[0]["in_reply_to"] == "orig@x.com" + + async def test_missing_file_counted(self, tmp_path, monkeypatch): + conn = _FakeConn(envelopes=[_env_row("missing@x", "hi")]) + ollama = _FakeOllamaSession() + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=False) + + assert stats["missing_file"] == 1 + assert stats["mails_scanned"] == 1 + + async def test_ollama_offline_batch_error_isolated_not_aborting(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" + )) + conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")]) + ollama = _FakeOllamaSession(dim=1024, fail_batches=True) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + + assert stats["chunks_errors"] == 1 + assert stats["chunks_inserted"] == 0 + assert stats["chunks_total"] == ( + stats["chunks_inserted"] + stats["chunks_newsletter_flagged"] + stats["chunks_already_embedded"] + + stats["chunks_conflict_skipped"] + stats["chunks_errors"] + ) + + async def test_dimension_mismatch_aborts(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" + )) + conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")]) + ollama = _FakeOllamaSession(dim=768) # wrong dim + self._patch(monkeypatch, conn, ollama) + + with pytest.raises(EmbeddingDimensionError): + await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + + async def test_conflict_skipped_counted_separately(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" + )) + conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")], execute_results=["INSERT 0 0"]) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + + assert stats["chunks_conflict_skipped"] == 1 + assert stats["chunks_inserted"] == 0 + + async def test_threading_already_present_is_skipped(self, tmp_path, monkeypatch): + _write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml( + {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" + )) + row = _env_row("m1@x", "hi") + entities = json.loads(row["entities"]) + [{"type": "threading", "in_reply_to": None, "references": []}] + row["entities"] = json.dumps(entities) + conn = _FakeConn(envelopes=[row]) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + + assert stats["threading_already_present"] == 1 + assert stats["threading_updated"] == 0 + assert conn.executemany_calls == []