- Chunker: char overlap across boundaries + hard-split of oversized paragraphs. - Retrieval knobs: rag_top_k / rag_min_score in settings + Settings UI. - Vector index: sqlite-vec ANN over document embeddings, dual-written and backfilled, with brute-force cosine as the guaranteed fallback. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
114 lines
4.3 KiB
Python
114 lines
4.3 KiB
Python
"""Document ingest / RAG store — hermetic (fake embeddings, temp DB)."""
|
|
import asyncio
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from synapse.memory.store import PersistentMemoryStore
|
|
|
|
|
|
def _store():
|
|
return PersistentMemoryStore(Path(tempfile.mkdtemp()) / "t.db")
|
|
|
|
|
|
def test_chunker_packs_and_splits():
|
|
s = _store()
|
|
one = s._chunk_text("short one.\n\nshort two.")
|
|
assert one == ["short one.\n\nshort two."] # both fit one chunk
|
|
many = s._chunk_text("a" * 700 + "\n\n" + "b" * 700)
|
|
assert len(many) == 2 # each paragraph near the size cap -> own chunk
|
|
|
|
|
|
async def _fake_embed(text):
|
|
kws = ["lego", "star", "wars", "gpu", "vega"]
|
|
v = [float(text.lower().count(k)) for k in kws]
|
|
return v if any(v) else None
|
|
|
|
|
|
def test_add_list_search_delete_roundtrip():
|
|
async def run():
|
|
s = _store()
|
|
r = await s.add_document(
|
|
"Guide",
|
|
"Beat the lego star wars boss with the force.\n\nUnrelated gpu vega notes.",
|
|
_fake_embed,
|
|
)
|
|
assert r["chunks"] >= 1
|
|
assert [d["title"] for d in s.list_documents()] == ["Guide"]
|
|
|
|
hits = await s.search_documents("lego star wars", _fake_embed, limit=2, min_score=0.1)
|
|
assert hits and "lego" in hits[0]["text"].lower()
|
|
# scores are sorted descending
|
|
assert all(hits[i]["score"] >= hits[i + 1]["score"] for i in range(len(hits) - 1))
|
|
|
|
# get_document returns ordered chunks for the viewer
|
|
chunks = s.get_document(r["doc_id"])
|
|
assert [c["chunk_idx"] for c in chunks] == list(range(len(chunks)))
|
|
|
|
assert s.delete_document(r["doc_id"]) is True
|
|
assert s.list_documents() == []
|
|
assert s.get_document(r["doc_id"]) == [] # gone -> no chunks
|
|
assert s.delete_document(r["doc_id"]) is False # already gone
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_search_empty_query_returns_nothing():
|
|
s = _store()
|
|
assert asyncio.run(s.search_documents("", _fake_embed)) == []
|
|
|
|
|
|
def test_chunker_overlap_and_hard_split():
|
|
s = _store()
|
|
# a single oversized paragraph (no blank lines, as in PDF text) is split
|
|
big = "x" * 2000
|
|
parts = s._chunk_text(big, size=800, overlap=120)
|
|
# each chunk is one <=size unit, plus at most an overlap tail (+separator)
|
|
assert len(parts) >= 3 and all(len(p) <= 800 + 120 + 2 for p in parts)
|
|
# consecutive chunks share an overlap tail
|
|
two = s._chunk_text("A" * 700 + "\n\n" + "B" * 700, size=800, overlap=120)
|
|
assert len(two) == 2 and two[1].startswith("A" * 120)
|
|
|
|
|
|
def test_vec_index_used_and_matches_brute_force():
|
|
# On a host that can load sqlite-vec, the fast path must be exercised (not a
|
|
# silent fallback) and agree with brute force on the top hit.
|
|
s = _store()
|
|
if not s.vec_enabled:
|
|
import pytest
|
|
pytest.skip("sqlite-vec not loadable on this host")
|
|
|
|
async def run():
|
|
await s.add_document("Lego", "lego star wars boss fight tips", _fake_embed)
|
|
await s.add_document("GPU", "gpu vega vram notes", _fake_embed)
|
|
# vec table populated by the dual-write
|
|
conn = s._connect()
|
|
n = conn.execute("SELECT COUNT(*) FROM vec_documents").fetchone()[0]
|
|
conn.close()
|
|
assert n == 2
|
|
|
|
vec_hits = await s.search_documents("lego star wars", _fake_embed, limit=2, min_score=0.1)
|
|
assert vec_hits and vec_hits[0]["title"] == "Lego"
|
|
|
|
# force the brute-force path and compare the top title
|
|
s.vec_enabled = False
|
|
bf_hits = await s.search_documents("lego star wars", _fake_embed, limit=2, min_score=0.1)
|
|
assert bf_hits[0]["title"] == vec_hits[0]["title"]
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_extract_text_by_type():
|
|
from synapse.main import _extract_text
|
|
# plain text / markdown -> UTF-8 decode
|
|
assert _extract_text("notes.md", b"# Title\n\nbody") == "# Title\n\nbody"
|
|
assert _extract_text("x.txt", "café".encode("utf-8")) == "café"
|
|
# a real (tiny) PDF built with pypdf -> text extracted back out
|
|
from pypdf import PdfWriter, PdfReader
|
|
import io
|
|
w = PdfWriter()
|
|
w.add_blank_page(width=200, height=200)
|
|
buf = io.BytesIO(); w.write(buf)
|
|
out = _extract_text("blank.pdf", buf.getvalue())
|
|
assert isinstance(out, str) # blank page -> "" or whitespace, never raises
|
|
assert PdfReader(io.BytesIO(buf.getvalue())).pages # sanity: it was a valid PDF
|