NexusOS/tests/test_documents.py

"""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_conversation_recall_uses_vec_and_matches_brute_force():
    s = _store()
    if not s.vec_enabled:
        import pytest
        pytest.skip("sqlite-vec not loadable on this host")

    async def run():
        s.create_conversation("c1")
        s.add_message("c1", "user", "tell me about lego star wars")
        s.add_message("c1", "assistant", "lego star wars is a fun game")
        s.create_conversation("c2")
        s.add_message("c2", "user", "gpu vega vram notes")
        s.add_message("c2", "assistant", "vega has 4gb")

        hits = await s.semantic_search_conversations("lego star wars", _fake_embed, limit=2, min_score=0.1)
        assert hits and hits[0]["id"] == "c1"

        conn = s._connect()
        n = conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0]
        conn.close()
        assert n >= 2  # dual-write populated the message vec index

        s.vec_enabled = False
        bf = await s.semantic_search_conversations("lego star wars", _fake_embed, limit=2, min_score=0.1)
        assert bf[0]["id"] == hits[0]["id"]

    asyncio.run(run())


def test_projects_scope_documents_and_survive_delete():
    s = _store()

    async def run():
        p = s.create_project("Star Wars")
        assert [x["name"] for x in s.list_projects()] == ["Star Wars"]

        await s.add_document("Lego", "lego star wars boss tips", _fake_embed, project_id=p["id"])
        await s.add_document("GPU", "gpu vega notes", _fake_embed)  # unscoped

        # project sees only its own; unscoped/all sees both
        assert [d["title"] for d in s.list_documents(p["id"])] == ["Lego"]
        assert len(s.list_documents(None)) == 2

        # scoped search only returns the project's docs
        scoped = await s.search_documents("star wars", _fake_embed, min_score=0.1, project_id=p["id"])
        assert scoped and all(h["title"] == "Lego" for h in scoped)

        # deleting the project keeps the docs but unscopes them
        assert s.delete_project(p["id"]) is True
        assert s.list_projects() == []
        assert len(s.list_documents(None)) == 2
        assert len(s.list_documents(p["id"])) == 0  # nothing left in that project

    asyncio.run(run())


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