lambdagent/tests/test_rag.py

"""
Direct unit tests for `lambdagent.rag` — Paper III's headline RAG layer.

audit/AUDIT_2026-06-05.md #35 baseline: 438 LOC of retrieval +
RAG-tool + agentic-RAG code with ZERO direct test. The integration
test paths in tests/test_extensions.py exercise Memory but not
retrieval — a regression in TF-IDF scoring, in score-then-format
glue, or in RAGTool's `Term` apply contract could ship to production
with green CI.

This file covers the dependency-free SimpleVectorStore + RAGTool +
create_rag paths. ChromaDB-backed paths are excluded — they need an
external service and that lifts test infra cost out of proportion;
add an opt-in `tests/test_rag_chroma.py` if/when needed.
"""
from __future__ import annotations

import pytest

from lambdagent.rag import (
    Document,
    SearchResult,
    SimpleVectorStore,
    RAGTool,
    create_rag,
)


# ─────────────────────────────────────────────────────────────────────────────
# SimpleVectorStore — pure data layer (no LLM)
# ─────────────────────────────────────────────────────────────────────────────

def test_simple_vector_store_add_returns_doc_id():
    store = SimpleVectorStore()
    doc_id = store.add("Python is a programming language")
    assert isinstance(doc_id, str) and len(doc_id) > 0
    assert len(store.documents) == 1
    assert store.documents[0].content == "Python is a programming language"


def test_simple_vector_store_search_ranks_relevant_first():
    """TF-IDF cosine should rank a topically-on-target document above
    off-target ones. We use distinctive words so the result is
    deterministic regardless of tokenizer details."""
    store = SimpleVectorStore()
    store.add("Lambda calculus is the foundation of computation theory")
    store.add("Python is a popular programming language for data science")
    store.add("The Eiffel Tower is in Paris and was built in 1889")

    results = store.search("what is lambda calculus", top_k=3)

    assert len(results) > 0, "no results returned at all"
    # The top-ranked result must be the lambda-calculus doc.
    assert "lambda" in results[0].document.content.lower(), (
        f"top result was not the lambda doc; got: "
        f"{results[0].document.content[:80]}"
    )
    # And it must outrank the unrelated Eiffel-Tower doc.
    scores_by_topic = {
        ("lambda" in r.document.content.lower()): r.score
        for r in results
    }
    if True in scores_by_topic and False in scores_by_topic:
        assert scores_by_topic[True] > scores_by_topic[False]


def test_simple_vector_store_empty_query_returns_empty():
    """Defense: an empty query shouldn't blow up — it should just
    return nothing useful. This guards a fragile Counter() path that
    used to ZeroDivisionError when the query tokenized to []."""
    store = SimpleVectorStore()
    store.add("anything")
    results = store.search("", top_k=3)
    # Either an empty list or all-zero-score results are acceptable;
    # the contract is "no crash + nothing claimed relevant".
    if results:
        for r in results:
            assert r.score <= 0.0 + 1e-9


def test_simple_vector_store_no_documents_returns_empty():
    store = SimpleVectorStore()
    results = store.search("anything", top_k=3)
    assert results == []


# ─────────────────────────────────────────────────────────────────────────────
# RAGTool — wraps store as a lambdagent Term
# ─────────────────────────────────────────────────────────────────────────────

def test_ragtool_apply_returns_formatted_string():
    """RAGTool must satisfy the Term apply contract:
      - input = query string
      - output = string suitable for injection into a downstream prompt
    The 'numbered' format produces `[Source N]` markers."""
    store = SimpleVectorStore()
    store.add("Lambda calculus underlies functional programming")
    store.add("Lambda functions in Python are anonymous functions")

    rag = RAGTool(store, top_k=2, format="numbered")
    out = rag.apply("lambda")
    assert isinstance(out, str)
    # `numbered` format emits `[Source 1]` or `[Source 1, score=...]`
    # depending on rag.py's format flag. Both forms start with the same
    # prefix.
    assert "[Source 1" in out, f"output missing [Source 1 marker: {out[:120]}"


def test_ragtool_min_score_filters_low_matches():
    """A query that matches NOTHING in the store should yield an empty
    or near-empty output, not a hallucinated reference."""
    store = SimpleVectorStore()
    store.add("Cats are friendly mammals")
    store.add("Dogs need daily walks")

    # min_score forces the filter even on the top-scored irrelevant doc.
    rag = RAGTool(store, top_k=3, min_score=0.5)
    out = rag.apply("quantum field theory renormalization")
    # Accept either an empty string or "no results found" message —
    # the contract is "don't lie about relevance".
    if out.strip():
        assert "[Source 1]" not in out or "no results" in out.lower()


# ─────────────────────────────────────────────────────────────────────────────
# create_rag — public one-liner
# ─────────────────────────────────────────────────────────────────────────────

def test_create_rag_simple_backend_smoke():
    """The README's headline 4-line RAG usage must keep working."""
    rag = create_rag([
        "Python is a programming language",
        "Lambda calculus is the foundation of computation",
        "AI agents are autonomous task executors",
    ])
    assert isinstance(rag, RAGTool)
    out = rag.apply("what is lambda calculus")
    assert "[Source" in out, f"create_rag output missing source markers: {out[:120]}"


def test_create_rag_unknown_backend_raises():
    """Defensive: typo in backend name should not silently fall back."""
    from lambdagent.core import LambdagentError
    with pytest.raises(LambdagentError):
        create_rag(["doc"], backend="not-a-real-backend")