348ce10b1d
chunk_text/hard_split/split_paragraphs + 2400/600-char constants move out of documents_ingest.chunk_embed into kb_mail.chunking so the upcoming jobs/mail-body-ingest shares the exact same chunker instead of a copy-pasted drift (the phase-4 lesson for retrieval.py -> packages/kb-retrieval, applied again). chunk_embed.py re-exports unchanged; zero behavior change, tests moved 1:1 alongside the code (docs/kb/modules/05-faza-mailowa-plan.md, §3).
487 lines
19 KiB
Python
487 lines
19 KiB
Python
"""Unit tests for the chunk + embed job — no DB, no real HTTP, no real Ollama.
|
|
|
|
`chunk_text`/`hard_split`/`split_paragraphs` tests moved to
|
|
packages/kb-mail/tests/test_chunking.py (module 5, faza mailowa, Krok 0) alongside the code."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from documents_ingest.chunk_embed import (
|
|
EmbeddingDimensionError,
|
|
embed_chunk,
|
|
extract_content,
|
|
fetch_documents,
|
|
fetch_existing_chunk_keys,
|
|
insert_chunk,
|
|
is_ocr_junk,
|
|
run,
|
|
_vector_literal,
|
|
)
|
|
|
|
|
|
def _paragraph(char: str, length: int) -> str:
|
|
return char * length
|
|
|
|
|
|
class TestExtractContent:
|
|
def test_finds_content_entity(self):
|
|
entities = [{"type": "filename", "value": "a.pdf"}, {"type": "content", "text": "hello"}]
|
|
assert extract_content(entities) == "hello"
|
|
|
|
def test_missing_content_entity_returns_empty_string(self):
|
|
entities = [{"type": "filename", "value": "a.pdf"}]
|
|
assert extract_content(entities) == ""
|
|
|
|
def test_none_text_returns_empty_string(self):
|
|
entities = [{"type": "content", "text": None}]
|
|
assert extract_content(entities) == ""
|
|
|
|
def test_empty_entities_list(self):
|
|
assert extract_content([]) == ""
|
|
assert extract_content(None) == ""
|
|
|
|
|
|
class TestVectorLiteral:
|
|
def test_formats_as_bracketed_csv(self):
|
|
assert _vector_literal([0.1, 0.2, -0.3]) == "[0.1,0.2,-0.3]"
|
|
|
|
|
|
class TestIsOcrJunk:
|
|
"""Pilot cases from docs/kb/modules/05-faza3-plan.md §3.1 and the 2026-07-16 calibration
|
|
session (docs/kb/eval/retrieval-pilot-2026-07-16.md)."""
|
|
|
|
def test_clean_text_is_not_junk(self):
|
|
text = (
|
|
"Ogólne Warunki Ubezpieczenia PZU Auto ustalają zasady odpowiedzialności "
|
|
"ubezpieczyciela za szkody powstałe w związku z ruchem pojazdów mechanicznych."
|
|
)
|
|
assert is_ocr_junk(text) is False
|
|
|
|
def test_ground_barcode_binary_noise_is_junk(self):
|
|
# paperless:39/96/128/186/192 pattern: dense run of C0 control chars (scrambled
|
|
# barcode), well above the calibrated raw-count threshold of 5.
|
|
text = "\x01\x02\x03\x05\x06\x07\x0f\x10\x11\x12\x14\x15\x16\x17\x18\x19qQQQ\x11q\x01\x11!!AA"
|
|
assert is_ocr_junk(text) is True
|
|
|
|
def test_partial_mojibake_is_not_junk(self):
|
|
# paperless:119: a confirmed retrieval hit (query 4, eval-set) despite partial
|
|
# mojibake and a couple of stray control bytes — mojibake alone must not exclude it.
|
|
text = (
|
|
"Wygenerowano z systemu wydarzeń FIRST Israel | Turniej Regionalny FIRST LEGO "
|
|
"League\n\x01integraln¹ czêœæ dokumentacji zawodów\nARTEFAKTY Z WĘDKI"
|
|
)
|
|
assert is_ocr_junk(text) is False
|
|
|
|
def test_stray_single_control_char_in_legible_text_is_not_junk(self):
|
|
# paperless:141 pattern: one stray control byte in an otherwise clean English
|
|
# paragraph (a PDF-extraction artifact, not OCR noise).
|
|
text = "separable\x0fconvolution is equal to the combination of a self-attention layer"
|
|
assert is_ocr_junk(text) is False
|
|
|
|
def test_dot_leader_toc_is_junk(self):
|
|
text = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
|
|
assert is_ocr_junk(text) is True
|
|
|
|
def test_empty_text_is_not_junk(self):
|
|
assert is_ocr_junk("") is False
|
|
|
|
|
|
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 RuntimeError(f"HTTP {self._status}")
|
|
|
|
|
|
class _FakeOllamaSession:
|
|
"""Serves a fixed embedding vector for every /api/embeddings POST, or errors by filename."""
|
|
|
|
def __init__(self, dim=1024, fail_for=None):
|
|
self._dim = dim
|
|
self._fail_for = fail_for or set()
|
|
self.requests: list[dict] = []
|
|
|
|
def post(self, url, json):
|
|
self.requests.append({"url": url, "json": json})
|
|
if json["prompt"] in self._fail_for:
|
|
return _FakeEmbedResponse({}, status=500)
|
|
return _FakeEmbedResponse({"embedding": [0.01] * self._dim})
|
|
|
|
async def close(self):
|
|
pass
|
|
|
|
|
|
class TestEmbedChunk:
|
|
async def test_returns_embedding_and_elapsed(self):
|
|
session = _FakeOllamaSession(dim=1024)
|
|
embedding, elapsed = await embed_chunk(session, "http://fake-ollama", "bge-m3", "hello world")
|
|
assert len(embedding) == 1024
|
|
assert elapsed >= 0
|
|
assert session.requests == [
|
|
{"url": "http://fake-ollama/api/embeddings", "json": {"model": "bge-m3", "prompt": "hello world"}}
|
|
]
|
|
|
|
async def test_missing_embedding_key_raises(self):
|
|
class _EmptyResponse(_FakeEmbedResponse):
|
|
pass
|
|
|
|
class _Session:
|
|
def post(self, url, json):
|
|
return _EmptyResponse({})
|
|
|
|
with pytest.raises(ValueError):
|
|
await embed_chunk(_Session(), "http://fake-ollama", "bge-m3", "text")
|
|
|
|
|
|
class TestInsertSql:
|
|
def test_on_conflict_target_includes_model(self):
|
|
# Plan §3.3 fix: the old UNIQUE(envelope_id, chunk_index) silently no-opped a
|
|
# second embedding model via ON CONFLICT DO NOTHING (review finding, phase 2 step
|
|
# 6). The target must now include `model`.
|
|
from documents_ingest.chunk_embed import _INSERT_SQL
|
|
assert "ON CONFLICT (envelope_id, chunk_index, model)" in _INSERT_SQL
|
|
|
|
|
|
class TestInsertChunk:
|
|
async def test_junk_chunk_inserts_null_embedding_with_reason(self):
|
|
conn = _FakeConn()
|
|
await insert_chunk(conn, "paperless:1", 0, "junk text", None, "bge-m3", excluded_reason="ocr_junk")
|
|
|
|
assert conn.execute_calls == [("paperless:1", 0, "junk text", None, "bge-m3", "ocr_junk")]
|
|
|
|
async def test_normal_chunk_inserts_vector_with_no_reason(self):
|
|
conn = _FakeConn()
|
|
await insert_chunk(conn, "paperless:1", 0, "real text", [0.1, 0.2], "bge-m3")
|
|
|
|
assert conn.execute_calls == [("paperless:1", 0, "real text", "[0.1,0.2]", "bge-m3", None)]
|
|
|
|
|
|
class _FakeConn:
|
|
def __init__(self, docs=None, existing_keys=None, existing_model="bge-m3", execute_results=None):
|
|
self._docs = docs or []
|
|
self._existing_keys = list(existing_keys or [])
|
|
self._existing_model = existing_model
|
|
# Optional queue of command tags returned by successive execute() calls, in order
|
|
# (e.g. ["INSERT 0 0"] to simulate an ON CONFLICT DO NOTHING no-op). Defaults to a
|
|
# real insert every time.
|
|
self._execute_results = list(execute_results) if execute_results is not None else None
|
|
self.execute_calls: list[tuple] = []
|
|
|
|
async def fetch(self, query, *params):
|
|
if "FROM document_chunk" in query:
|
|
# Mirrors the real `WHERE model = $1` filter: existing_keys were "written" under
|
|
# existing_model, so a query for a different model must not see them.
|
|
if params and params[0] != self._existing_model:
|
|
return []
|
|
return [{"envelope_id": eid, "chunk_index": idx} for eid, idx in self._existing_keys]
|
|
return self._docs
|
|
|
|
async def execute(self, query, *params):
|
|
self.execute_calls.append(params)
|
|
if self._execute_results is not None:
|
|
return self._execute_results.pop(0)
|
|
return "INSERT 0 1"
|
|
|
|
async def close(self):
|
|
pass
|
|
|
|
|
|
def _row(envelope_id, content):
|
|
return {"id": envelope_id, "entities": json.dumps([{"type": "content", "text": content}])}
|
|
|
|
|
|
class TestFetchHelpers:
|
|
async def test_fetch_documents_no_limit_offset(self):
|
|
conn = _FakeConn(docs=[_row("paperless:1", "hi")])
|
|
docs = await fetch_documents(conn, limit=None, offset=None)
|
|
assert docs == [_row("paperless:1", "hi")]
|
|
|
|
async def test_fetch_existing_chunk_keys(self):
|
|
conn = _FakeConn(existing_keys=[("paperless:1", 0), ("paperless:1", 1)])
|
|
keys = await fetch_existing_chunk_keys(conn, model="bge-m3")
|
|
assert keys == {("paperless:1", 0), ("paperless:1", 1)}
|
|
|
|
|
|
class TestRun:
|
|
def _patch(self, monkeypatch, conn, ollama_session):
|
|
async def _fake_connect(dsn):
|
|
return conn
|
|
monkeypatch.setattr("documents_ingest.chunk_embed.asyncpg.connect", _fake_connect)
|
|
|
|
def _fake_session_factory(*args, **kwargs):
|
|
return ollama_session
|
|
monkeypatch.setattr("documents_ingest.chunk_embed.aiohttp.ClientSession", _fake_session_factory)
|
|
|
|
async def test_dry_run_counts_without_calling_ollama_or_db(self, monkeypatch):
|
|
conn = _FakeConn(docs=[_row("paperless:1", "short doc")])
|
|
ollama = _FakeOllamaSession()
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=False)
|
|
|
|
assert stats["documents_fetched"] == 1
|
|
assert stats["documents_chunked"] == 1
|
|
assert stats["chunks_total"] == 1
|
|
assert stats["chunks_inserted"] == 1
|
|
assert ollama.requests == []
|
|
assert conn.execute_calls == []
|
|
|
|
async def test_empty_content_counted_separately_not_as_error(self, monkeypatch):
|
|
conn = _FakeConn(docs=[_row("paperless:1", ""), _row("paperless:2", "some text")])
|
|
ollama = _FakeOllamaSession()
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=False)
|
|
|
|
assert stats["empty_content"] == 1
|
|
assert stats["documents_chunked"] == 1
|
|
assert stats["documents_fetched"] == 2
|
|
assert stats["chunks_errors"] == 0
|
|
|
|
async def test_apply_embeds_and_inserts(self, monkeypatch):
|
|
conn = _FakeConn(docs=[_row("paperless:1", "some real content")])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True)
|
|
|
|
assert stats["chunks_inserted"] == 1
|
|
assert stats["embed_calls"] == 1
|
|
assert len(conn.execute_calls) == 1
|
|
assert len(ollama.requests) == 1
|
|
|
|
async def test_idempotent_skips_already_embedded_chunks(self, monkeypatch):
|
|
conn = _FakeConn(
|
|
docs=[_row("paperless:1", "some real content")],
|
|
existing_keys=[("paperless:1", 0)],
|
|
)
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True)
|
|
|
|
assert stats["chunks_already_embedded"] == 1
|
|
assert stats["chunks_inserted"] == 0
|
|
assert ollama.requests == []
|
|
assert conn.execute_calls == []
|
|
|
|
async def test_rerun_after_apply_inserts_nothing_new(self, monkeypatch):
|
|
doc = _row("paperless:1", "some real content")
|
|
|
|
conn1 = _FakeConn(docs=[doc])
|
|
ollama1 = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn1, ollama1)
|
|
first = await run(dsn="postgresql://fake", apply=True)
|
|
assert first["chunks_inserted"] == 1
|
|
|
|
conn2 = _FakeConn(docs=[doc], existing_keys=[("paperless:1", 0)])
|
|
ollama2 = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn2, ollama2)
|
|
second = await run(dsn="postgresql://fake", apply=True)
|
|
|
|
assert second["chunks_inserted"] == 0
|
|
assert second["chunks_already_embedded"] == 1
|
|
assert ollama2.requests == []
|
|
|
|
async def test_dimension_mismatch_aborts(self, monkeypatch):
|
|
conn = _FakeConn(docs=[_row("paperless:1", "some real content")])
|
|
ollama = _FakeOllamaSession(dim=768) # wrong dim
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
with pytest.raises(EmbeddingDimensionError):
|
|
await run(dsn="postgresql://fake", apply=True)
|
|
|
|
async def test_embed_error_is_isolated_and_counted(self, monkeypatch):
|
|
# Two chunks: force one to fail via a document long enough to produce 2 chunks
|
|
# (chunk_size/overlap kept small to make the test fast and explicit). Real
|
|
# word-like tokens (not a single 100-char run) so neither chunk trips is_ocr_junk.
|
|
para1 = " ".join(["word"] * 20)
|
|
para2 = " ".join(["term"] * 20)
|
|
big = para1 + "\n\n" + para2
|
|
conn = _FakeConn(docs=[_row("paperless:1", big)])
|
|
ollama = _FakeOllamaSession(dim=1024, fail_for={para1})
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True, chunk_size=100, chunk_overlap=20)
|
|
|
|
assert stats["chunks_errors"] == 1
|
|
assert stats["chunks_inserted"] == 1
|
|
assert stats["chunks_total"] == (
|
|
stats["chunks_already_embedded"] + stats["chunks_inserted"]
|
|
+ stats["chunks_conflict_skipped"] + stats["chunks_errors"]
|
|
)
|
|
|
|
async def test_insert_error_is_isolated_and_counted(self, monkeypatch):
|
|
# Two chunks; the first chunk's INSERT raises (simulated DB blip) but the run
|
|
# continues and the second chunk still gets embedded and inserted normally.
|
|
# Real word-like tokens so neither chunk trips is_ocr_junk.
|
|
big = " ".join(["word"] * 20) + "\n\n" + " ".join(["term"] * 20)
|
|
conn = _FakeConn(docs=[_row("paperless:1", big)])
|
|
|
|
real_execute = conn.execute
|
|
|
|
async def _flaky_execute(query, *params):
|
|
if params[1] == 0: # chunk_index 0
|
|
conn.execute_calls.append(params)
|
|
raise RuntimeError("simulated db error")
|
|
return await real_execute(query, *params)
|
|
|
|
conn.execute = _flaky_execute
|
|
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True, chunk_size=100, chunk_overlap=20)
|
|
|
|
assert stats["chunks_errors"] == 1
|
|
assert stats["chunks_inserted"] == 1
|
|
assert stats["chunks_total"] == (
|
|
stats["chunks_already_embedded"] + stats["chunks_inserted"]
|
|
+ stats["chunks_conflict_skipped"] + stats["chunks_errors"]
|
|
)
|
|
|
|
async def test_conflict_skipped_counted_separately_from_inserted(self, monkeypatch):
|
|
# ON CONFLICT DO NOTHING no-op: command tag reports 0 rows affected.
|
|
conn = _FakeConn(
|
|
docs=[_row("paperless:1", "some real content")], execute_results=["INSERT 0 0"]
|
|
)
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True)
|
|
|
|
assert stats["chunks_conflict_skipped"] == 1
|
|
assert stats["chunks_inserted"] == 0
|
|
assert stats["chunks_total"] == (
|
|
stats["chunks_already_embedded"] + stats["chunks_inserted"]
|
|
+ stats["chunks_conflict_skipped"] + stats["chunks_errors"]
|
|
)
|
|
|
|
async def test_existing_keys_scoped_to_model(self, monkeypatch):
|
|
# existing_keys were "written" under a different model — the WHERE model = $1
|
|
# filter must not treat them as already-embedded for the current --model.
|
|
conn = _FakeConn(
|
|
docs=[_row("paperless:1", "some real content")],
|
|
existing_keys=[("paperless:1", 0)],
|
|
existing_model="some-other-model",
|
|
)
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True, model="bge-m3")
|
|
|
|
assert stats["chunks_already_embedded"] == 0
|
|
assert stats["chunks_inserted"] == 1
|
|
|
|
async def test_junk_chunk_flagged_without_calling_ollama(self, monkeypatch):
|
|
junk = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
|
|
conn = _FakeConn(docs=[_row("paperless:1", junk)])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True)
|
|
|
|
assert stats["chunks_junk_flagged"] == 1
|
|
assert stats["chunks_inserted"] == 0
|
|
assert ollama.requests == []
|
|
assert conn.execute_calls == [("paperless:1", 0, junk, None, "bge-m3", "ocr_junk")]
|
|
assert stats["chunks_total"] == (
|
|
stats["chunks_already_embedded"] + stats["chunks_inserted"] + stats["chunks_junk_flagged"]
|
|
+ stats["chunks_conflict_skipped"] + stats["chunks_errors"]
|
|
)
|
|
|
|
async def test_junk_chunk_dry_run_counted_without_db_or_ollama(self, monkeypatch):
|
|
junk = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
|
|
conn = _FakeConn(docs=[_row("paperless:1", junk)])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=False)
|
|
|
|
assert stats["chunks_junk_flagged"] == 1
|
|
assert stats["chunks_inserted"] == 0
|
|
assert ollama.requests == []
|
|
assert conn.execute_calls == []
|
|
|
|
async def test_junk_chunk_rerun_is_idempotent(self, monkeypatch):
|
|
# Same (envelope_id, chunk_index) already flagged ocr_junk under this model — a
|
|
# rerun must skip it, same as a normal embedded chunk (plan §3.1: backfill is a
|
|
# one-shot, but chunk_embed.py itself must stay idempotent for new junk too).
|
|
junk = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
|
|
conn = _FakeConn(
|
|
docs=[_row("paperless:1", junk)],
|
|
existing_keys=[("paperless:1", 0)],
|
|
)
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True)
|
|
|
|
assert stats["chunks_already_embedded"] == 1
|
|
assert stats["chunks_junk_flagged"] == 0
|
|
assert conn.execute_calls == []
|
|
|
|
async def test_junk_conflict_skip_not_double_counted(self, monkeypatch):
|
|
junk = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
|
|
conn = _FakeConn(docs=[_row("paperless:1", junk)], execute_results=["INSERT 0 0"])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True)
|
|
|
|
assert stats["chunks_conflict_skipped"] == 1
|
|
assert stats["chunks_junk_flagged"] == 0
|
|
assert ollama.requests == []
|
|
|
|
async def test_stats_balance_documents_and_chunks(self, monkeypatch):
|
|
conn = _FakeConn(docs=[
|
|
_row("paperless:1", "content one"),
|
|
_row("paperless:2", ""),
|
|
_row("paperless:3", "content three"),
|
|
])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", apply=True)
|
|
|
|
assert stats["documents_fetched"] == stats["empty_content"] + stats["documents_chunked"]
|
|
assert stats["chunks_total"] == (
|
|
stats["chunks_already_embedded"] + stats["chunks_inserted"]
|
|
+ stats["chunks_conflict_skipped"] + stats["chunks_errors"]
|
|
)
|
|
|
|
async def test_limit_and_offset_passed_through_to_query(self, monkeypatch):
|
|
captured = {}
|
|
|
|
class _Conn(_FakeConn):
|
|
async def fetch(self, query, *params):
|
|
if "FROM document_chunk" in query:
|
|
return []
|
|
captured["query"] = query
|
|
captured["params"] = params
|
|
return []
|
|
|
|
conn = _Conn()
|
|
ollama = _FakeOllamaSession()
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
await run(dsn="postgresql://fake", apply=False, limit=10, offset=5)
|
|
|
|
assert "LIMIT $1" in captured["query"]
|
|
assert "OFFSET $2" in captured["query"]
|
|
assert captured["params"] == (10, 5)
|