""" Tests for citation-aware ingestion (M2 Phase E). lambdagent.ingest — page-preserving chunker + Obsidian import + the rag.RAGTool `[Source N | file p.12]` citation format (FR-005). PDF extraction itself (extract_pdf_pages) needs pdfplumber/PyPDF2 + a real PDF, so we test the chunker against synthetic page tuples (the part that actually carries provenance) and exercise Obsidian import with real temp markdown files. """ from __future__ import annotations import os import pytest from lambdagent.ingest import ( CitedChunk, chunk_text, chunk_pages, import_obsidian_vault, ) from lambdagent.rag import RAGTool, SimpleVectorStore # ── chunk_text ────────────────────────────────────────────────────────────── def test_chunk_text_short_returns_single(): assert chunk_text("short text", chunk_size=800) == ["short text"] def test_chunk_text_empty_returns_empty(): assert chunk_text("") == [] assert chunk_text(" \n ") == [] def test_chunk_text_splits_long_with_overlap(): text = "para one.\n\n" + ("word " * 400) # well over 800 chars chunks = chunk_text(text, chunk_size=300, overlap=50) assert len(chunks) > 1 # No chunk wildly exceeds chunk_size (allow some boundary slack). assert all(len(c) <= 300 + 50 for c in chunks) # ── chunk_pages: provenance preserved ─────────────────────────────────────── def test_chunk_pages_preserves_page_numbers(): pages = [ (1, "Introduction. " + "alpha " * 200), (2, "Methods. " + "beta " * 200), (3, ""), # blank page — produces no chunks but doesn't break numbering (4, "Results. " + "gamma " * 200), ] chunks = chunk_pages(pages, source="paper.pdf", chunk_size=300, overlap=40) assert chunks, "expected chunks" # Every chunk knows its source + a real page number. assert all(c.source == "paper.pdf" for c in chunks) assert all(c.page in (1, 2, 4) for c in chunks) # A chunk never spans pages: page-1 chunks only contain alpha, etc. for c in chunks: if c.page == 1: assert "alpha" in c.text and "beta" not in c.text if c.page == 2: assert "beta" in c.text and "gamma" not in c.text # chunk_index is globally increasing. idxs = [c.chunk_index for c in chunks] assert idxs == sorted(idxs) assert len(set(idxs)) == len(idxs) def test_cited_chunk_to_metadata(): c = CitedChunk(text="x", source="paper.pdf", page=12, chunk_index=3) md = c.to_metadata() assert md["source"] == "paper.pdf" assert md["page"] == 12 assert md["chunk_index"] == 3 # Non-paged source (markdown) omits page. c2 = CitedChunk(text="y", source="note.md", page=None) assert "page" not in c2.to_metadata() # ── Obsidian vault import ──────────────────────────────────────────────────── def test_import_obsidian_vault(tmp_path): vault = tmp_path / "vault" (vault / "notes").mkdir(parents=True) (vault / "index.md").write_text("# Index\nSee [[notes/qgt|the QGT note]] for details.") (vault / "notes" / "qgt.md").write_text("# QGT\nQuantum geometric tensor basics. " * 5) # Obsidian config dir + a non-md file must be ignored. (vault / ".obsidian").mkdir() (vault / ".obsidian" / "app.json").write_text("{}") (vault / "image.png").write_bytes(b"\x89PNG") chunks = import_obsidian_vault(str(vault)) sources = {c.source for c in chunks} # Only the two .md files, with vault-relative paths. assert "index.md" in sources assert os.path.join("notes", "qgt.md") in sources assert not any(".obsidian" in s for s in sources) assert not any(s.endswith(".png") for s in sources) # Markdown chunks have no page. assert all(c.page is None for c in chunks) # Wikilink was flattened to its alias text + recorded in extra.links. index_chunks = [c for c in chunks if c.source == "index.md"] assert index_chunks assert "the QGT note" in index_chunks[0].text assert "[[" not in index_chunks[0].text assert "notes/qgt" in index_chunks[0].extra.get("links", []) def test_import_obsidian_missing_dir_raises(tmp_path): with pytest.raises(NotADirectoryError): import_obsidian_vault(str(tmp_path / "nope")) # ── RAGTool citation format (FR-005) ───────────────────────────────────────── def test_ragtool_renders_pdf_citation(): """A chunk ingested from a PDF (metadata source+page) must render as `[Source N | paper.pdf p.12]` — the headline FR-005 deliverable.""" store = SimpleVectorStore() store.add( "The quantum geometric tensor encodes the metric and Berry curvature.", metadata=CitedChunk(text="", source="qgt.pdf", page=12).to_metadata(), ) store.add( "Unrelated content about cats and dogs.", metadata=CitedChunk(text="", source="pets.pdf", page=3).to_metadata(), ) rag = RAGTool(store, top_k=1, format="numbered") out = rag.apply("quantum geometric tensor metric Berry curvature") assert "| qgt.pdf p.12" in out, out def test_ragtool_renders_source_without_page(): """Markdown chunk (no page) renders `| note.md`, no `p.None`.""" store = SimpleVectorStore() store.add( "Obsidian note about lambda calculus and Y combinators.", metadata=CitedChunk(text="", source="lambda.md", page=None).to_metadata(), ) rag = RAGTool(store, top_k=1, format="numbered") out = rag.apply("lambda calculus Y combinator") assert "| lambda.md" in out assert "p.None" not in out