lambdagent/tests/test_ingest.py

"""
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