tests/test_chroma_rag.py

from __future__ import annotations

import asyncio
import pickle
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

import chromadb
import tiktoken
from data.scraping_scripts.add_context_to_nodes import process
from data.scraping_scripts.create_vector_stores import write_retrieval_artifacts
from llama_index.core import Document
from app.chroma_rag import (
    BM25Index,
    ChunkRecord,
    LocalChromaRetriever,
    build_chunk_records,
    heading_aware_markdown_chunks,
    reciprocal_rank_fusion,
    load_bm25_index,
    rerank_results,
    SearchResult,
)


class ChromaRagTestCase(unittest.TestCase):
    def test_heading_aware_chunks_keep_code_blocks_intact(self) -> None:
        code_lines = "\n".join(f"print({index})" for index in range(120))
        markdown = f"""# Guide

## Install

Use `pip install`.

## Example

```python
{code_lines}
```

After the example.
"""

        chunks = heading_aware_markdown_chunks(
            markdown,
            title="Guide",
            chunk_size=80,
        )

        code_chunks = [chunk for chunk in chunks if "print(0)" in chunk.text]
        self.assertEqual(len(code_chunks), 1)
        self.assertIn("print(119)", code_chunks[0].text)
        self.assertIn("Example", code_chunks[0].heading_path)

    def test_build_chunk_records_adds_heading_metadata(self) -> None:
        records = build_chunk_records(
            [
                {
                    "doc_id": "doc-1",
                    "name": "Guide",
                    "url": "https://example.com/guide",
                    "source": "transformers",
                    "retrieve_doc": False,
                    "tokens": 1000,
                    "content": "# Guide\n\n## Install\n\nUse `AutoModel`.",
                }
            ]
        )

        self.assertEqual(records[0].metadata["heading_path"], "Guide")
        self.assertEqual(records[1].metadata["heading_path"], "Guide > Install")
        self.assertIn("source_version", records[0].metadata)

    def test_bm25_search_finds_keywords_and_filters_sources(self) -> None:
        records = [
            ChunkRecord(
                chunk_id="a",
                doc_id="doc-a",
                text="Use AutoModel.from_pretrained for model loading.",
                metadata={"doc_id": "doc-a", "source": "transformers"},
            ),
            ChunkRecord(
                chunk_id="b",
                doc_id="doc-b",
                text="Create a prompt template for chains.",
                metadata={"doc_id": "doc-b", "source": "langchain"},
            ),
        ]
        index = BM25Index.build(records)

        hits = index.search(
            "AutoModel.from_pretrained", allowed_sources=["transformers"]
        )

        self.assertEqual([record.chunk_id for record, _score in hits], ["a"])
        self.assertEqual(index.search("AutoModel", allowed_sources=["langchain"]), [])

    def test_retrieval_artifact_writer_persists_bm25_and_document_dict(self) -> None:
        document_rows = [
            {
                "doc_id": "doc-1",
                "name": "Transformers Loading",
                "url": "https://example.com/loading",
                "source": "transformers",
                "retrieve_doc": False,
                "tokens": 1200,
                "content": "# Loading\n\n## AutoModel\n\nUse `AutoModel.from_pretrained`.",
            }
        ]

        with tempfile.TemporaryDirectory() as temp_dir:
            db_path = Path(temp_dir)
            count = write_retrieval_artifacts(
                config={
                    "document_dict_file": "document_dict_test.pkl",
                    "bm25_index_file": "bm25_index_test.json.gz",
                },
                document_rows=document_rows,
                db_path=str(db_path),
            )

            document_dict_path = db_path / "document_dict_test.pkl"
            bm25_path = db_path / "bm25_index_test.json.gz"

            self.assertGreaterEqual(count, 1)
            self.assertTrue(document_dict_path.exists())
            self.assertTrue(bm25_path.exists())

            with open(document_dict_path, "rb") as handle:
                document_dict = pickle.load(handle)
            self.assertEqual(document_dict["doc-1"]["name"], "Transformers Loading")

            index = load_bm25_index(str(bm25_path))
            self.assertIsNotNone(index)
            assert index is not None
            hits = index.search("AutoModel.from_pretrained")
            self.assertEqual(hits[0][0].doc_id, "doc-1")
            self.assertTrue(
                any(record.metadata["heading_path"] for record in index.records)
            )

    def test_context_processing_uses_heading_chunks_and_raw_text_metadata(self) -> None:
        async def fake_situate_context(_doc: str, chunk: str) -> str:
            return f"Situated {chunk.splitlines()[0]}"

        document = Document(
            doc_id="doc-1",
            text="# Guide\n\n## Setup\n\nUse `AutoModel.from_pretrained`.",
            metadata={
                "title": "Guide",
                "url": "https://example.com/guide",
                "tokens": 1000,
                "retrieve_doc": False,
                "source": "transformers",
            },
        )

        with patch(
            "data.scraping_scripts.add_context_to_nodes.situate_context",
            fake_situate_context,
        ):
            records = asyncio.run(process([document], semaphore_limit=1))

        self.assertGreaterEqual(len(records), 1)
        setup_record = next(
            record
            for record in records
            if record.metadata["heading_path"] == "Guide > Setup"
        )
        self.assertIn("raw_text", setup_record.metadata)
        self.assertIn("Title: Guide", setup_record.text)
        self.assertIn("Heading path: Guide > Setup", setup_record.text)
        self.assertIn("Context: Situated", setup_record.text)

    def test_rerank_scores_matched_chunk_for_retrieve_doc_results(self) -> None:
        # retrieve_doc results carry the whole document in `content`; the
        # reranker must score the matched chunk (`chunk_content`) instead, so
        # relevance is not diluted toward the doc average and the payload stays
        # within Cohere's per-document token limit.
        full_doc = "Intro paragraph.\n" * 500
        results = [
            SearchResult(
                chunk_id="doc-chunk",
                doc_id="doc-1",
                title="Doc",
                url="",
                source="test",
                retrieve_doc=True,
                tokens=4000,
                score=0.5,
                content=full_doc,
                chunk_content="the matched chunk about AutoModel",
                heading_path="section",
                retrieval_method="dense",
            ),
            SearchResult(
                chunk_id="plain-chunk",
                doc_id="doc-2",
                title="Plain",
                url="",
                source="test",
                retrieve_doc=False,
                tokens=100,
                score=0.4,
                content="formatted chunk body",
                chunk_content="raw chunk body",
                heading_path="section",
                retrieval_method="dense",
            ),
        ]

        captured: dict[str, list[str]] = {}

        class _FakeItem:
            def __init__(self, index: int, score: float) -> None:
                self.index = index
                self.relevance_score = score

        class _FakeResponse:
            def __init__(self, items: list["_FakeItem"]) -> None:
                self.results = items

        class _FakeCohere:
            def rerank(self, *, model, query, documents, top_n):  # type: ignore[no-untyped-def]
                captured["documents"] = list(documents)
                return _FakeResponse(
                    [
                        _FakeItem(i, 1.0 - i * 0.1)
                        for i in range(min(top_n, len(documents)))
                    ]
                )

        reranked = rerank_results(_FakeCohere(), "AutoModel", results)

        # The full document never reaches the reranker; the matched chunk does.
        self.assertEqual(
            captured["documents"],
            ["the matched chunk about AutoModel", "formatted chunk body"],
        )
        # The returned result still carries the full document for the answer.
        self.assertEqual(reranked[0].content, full_doc)

    def test_rrf_prefers_overlap_across_ranked_lists(self) -> None:
        dense_only = self._result("dense-only", 0.9, "dense")
        overlap_dense = self._result("overlap", 0.7, "dense")
        overlap_bm25 = self._result("overlap", 4.0, "bm25")
        bm25_only = self._result("bm25-only", 5.0, "bm25")

        fused = reciprocal_rank_fusion(
            [[dense_only, overlap_dense], [bm25_only, overlap_bm25]],
            top_k=4,
        )

        self.assertEqual(fused[0].chunk_id, "overlap")
        self.assertEqual(fused[0].retrieval_method, "hybrid")

    def test_rrf_counts_each_key_once_per_ranked_list(self) -> None:
        # A section split into several chunks can land at multiple ranks of ONE
        # retriever's list (same dedupe key). Standard RRF scores a key once per
        # list, at its best rank; per-occurrence accumulation would let one
        # retriever's duplicates masquerade as cross-retriever consensus.
        dup_top = self._result("dup:0", 0.9, "dense", doc_id="dup-doc")
        dup_mid = self._result("dup:1", 0.8, "dense", doc_id="dup-doc")
        dup_low = self._result("dup:2", 0.7, "dense", doc_id="dup-doc")
        consensus_dense = self._result("uni:0", 0.6, "dense", doc_id="consensus-doc")
        consensus_bm25 = self._result("uni:0", 5.0, "bm25", doc_id="consensus-doc")

        fused = reciprocal_rank_fusion(
            [[dup_top, dup_mid, dup_low, consensus_dense], [consensus_bm25]],
            top_k=5,
        )

        by_doc = {result.doc_id: result for result in fused}
        # One contribution at the best rank (1), nothing from ranks 2-3.
        self.assertAlmostEqual(by_doc["dup-doc"].score, 1.0 / 61)
        # Rank 4 in dense + rank 1 in bm25.
        self.assertAlmostEqual(by_doc["consensus-doc"].score, 1.0 / 64 + 1.0 / 61)
        # Genuine cross-retriever consensus outranks single-list duplication.
        self.assertEqual(fused[0].doc_id, "consensus-doc")
        # Representative selection still works: best-scoring dense duplicate.
        self.assertEqual(by_doc["dup-doc"].chunk_id, "dup:0")

    def _result(
        self,
        chunk_id: str,
        score: float,
        method: str,
        *,
        doc_id: str | None = None,
        content: str | None = None,
        retrieve_doc: bool = False,
    ) -> SearchResult:
        return SearchResult(
            chunk_id=chunk_id,
            doc_id=doc_id if doc_id is not None else chunk_id,
            title=chunk_id,
            url="",
            source="test",
            retrieve_doc=retrieve_doc,
            tokens=10,
            score=score,
            content=content if content is not None else chunk_id,
            chunk_content=content if content is not None else chunk_id,
            heading_path="section",
            retrieval_method=method,
        )


class TokenBudgetTestCase(unittest.TestCase):
    def _retriever(self) -> LocalChromaRetriever:
        retriever = LocalChromaRetriever.__new__(LocalChromaRetriever)
        retriever._encoding = tiktoken.get_encoding("cl100k_base")
        retriever._token_budget = 100_000
        return retriever

    def _result(self, chunk_id: str, score: float, content: str) -> SearchResult:
        return SearchResult(
            chunk_id=chunk_id,
            doc_id=chunk_id,
            title=chunk_id,
            url="",
            source="test",
            retrieve_doc=False,
            tokens=10,
            score=score,
            content=content,
            chunk_content=content,
            heading_path="section",
            retrieval_method="dense",
        )

    def test_budget_skips_oversized_result_and_fills_with_smaller(self) -> None:
        # A rank-1 retrieve_doc result whose full document exceeds a small
        # per-request budget must not empty the whole result list: it is
        # skipped and the budget is filled with lower-ranked results that fit,
        # in rank order.
        retriever = self._retriever()
        oversized = self._result("big", 0.9, "word " * 500)
        small_one = self._result("small-1", 0.8, "word " * 15)
        small_two = self._result("small-2", 0.7, "word " * 15)

        kept = retriever._apply_token_budget(
            [oversized, small_one, small_two], token_budget=50
        )
        self.assertEqual([result.chunk_id for result in kept], ["small-1", "small-2"])

        # With the default (large) budget everything still fits.
        kept_all = retriever._apply_token_budget([oversized, small_one, small_two])
        self.assertEqual(len(kept_all), 3)


class CollectionOpenTestCase(unittest.TestCase):
    def _write_document_dict(self, directory: str) -> str:
        path = Path(directory) / "document_dict_test.pkl"
        with open(path, "wb") as handle:
            pickle.dump({}, handle)
        return str(path)

    def test_init_fails_loudly_when_collection_missing(self) -> None:
        # A broken/mismatched bundle must raise at startup instead of silently
        # creating an empty collection that returns zero dense hits forever.
        with tempfile.TemporaryDirectory() as temp_dir:
            document_dict_path = self._write_document_dict(temp_dir)

            with self.assertRaises(RuntimeError) as ctx:
                LocalChromaRetriever(
                    db_path=temp_dir,
                    collection_name="missing-collection",
                    document_dict_path=document_dict_path,
                    cohere_api_key="fake",
                )

            message = str(ctx.exception)
            self.assertIn("missing-collection", message)
            self.assertIn(temp_dir, message)

    def test_init_opens_collection_created_beforehand(self) -> None:
        # Mirrors production: create_vector_stores creates the collection; the
        # retriever only opens it.
        with tempfile.TemporaryDirectory() as temp_dir:
            chromadb.PersistentClient(path=temp_dir).create_collection(
                name="test-collection"
            )
            document_dict_path = self._write_document_dict(temp_dir)

            retriever = LocalChromaRetriever(
                db_path=temp_dir,
                collection_name="test-collection",
                document_dict_path=document_dict_path,
                cohere_api_key="fake",
            )

            self.assertEqual(retriever._collection.name, "test-collection")


if __name__ == "__main__":
    unittest.main()