ai-agent-book/tests/test_ch3_hybrid_structured_retriever.py

"""Unit tests for chapter3/structured-index/hybrid_retriever.py (HybridStructuredRetriever)."""

import importlib.util
import os
import sys
from pathlib import Path
import pytest

pytest.importorskip("numpy")
import numpy as np

# Dynamic import for hyphenated module path
_module_path = (
    Path(__file__).resolve().parent.parent
    / "chapter3"
    / "structured-index"
    / "hybrid_retriever.py"
)
_spec = importlib.util.spec_from_file_location("hybrid_retriever", _module_path)
_mod = importlib.util.module_from_spec(_spec)
sys.modules["hybrid_retriever"] = _mod
_spec.loader.exec_module(_mod)

HybridStructuredRetriever = _mod.HybridStructuredRetriever
SearchResult = _mod.SearchResult
EvidenceCitation = _mod.EvidenceCitation


def test_add_nodes_and_basic_retrieval():
    """Verify RAPTOR nodes and GraphRAG entities can be added and retrieved."""
    retriever = HybridStructuredRetriever(rrf_k=60)

    # Add RAPTOR tree summary node
    retriever.add_raptor_node(
        node_id="r1",
        level=2,
        text="Deep learning architectures utilize multi-layer neural networks.",
        summary="Overview of deep learning and multi-layer neural networks.",
        children=["r1_1", "r1_2"],
    )

    # Add GraphRAG entity
    retriever.add_graphrag_entity(
        entity_id="e1",
        name="Neural Network",
        type="ARCHITECTURE",
        description="A machine learning model inspired by biological neural circuits.",
    )

    # Add GraphRAG relationship
    retriever.add_graphrag_relationship(
        relation_id="rel1",
        source="Neural Network",
        target="Deep Learning",
        type="USED_IN",
        description="Neural networks serve as foundational models in deep learning.",
    )

    results = retriever.retrieve("deep learning neural network", top_k=5)

    assert len(results) > 0
    assert isinstance(results[0], SearchResult)
    assert results[0].score > 0.0

    # Verify citation details exist on all results
    for res in results:
        assert isinstance(res.citation, EvidenceCitation)
        assert res.citation.source_type in (
            "raptor_tree",
            "graphrag_entity",
            "graphrag_relation",
            "graphrag_community",
        )
        assert len(res.citation.citation_label) > 0


def test_rrf_scoring_order_and_fusion():
    """Verify Reciprocal Rank Fusion combines RAPTOR and GraphRAG rankings."""
    retriever = HybridStructuredRetriever(rrf_k=60)

    # RAPTOR node relevant to quantum computing
    retriever.add_raptor_node(
        node_id="rap_quantum",
        level=1,
        text="Quantum algorithms exploit superposition and entanglement.",
        summary="Quantum computing algorithms and superposition.",
    )

    # GraphRAG community summary relevant to quantum computing
    retriever.add_graphrag_community(
        community_id="comm_quantum",
        entity_ids=["Qubit", "QuantumGate"],
        summary="Community of quantum hardware components and quantum algorithms.",
        level=0,
    )

    # Irrelevant node
    retriever.add_raptor_node(
        node_id="rap_gardening",
        level=0,
        text="Gardening tips for growing organic tomatoes in summer.",
        summary="Organic tomato gardening guidance.",
    )

    results = retriever.retrieve("quantum algorithms superposition", top_k=2)

    assert len(results) == 2
    retrieved_ids = [r.node_id for r in results]

    assert "rap_quantum" in retrieved_ids or "comm_quantum" in retrieved_ids
    assert "rap_gardening" not in retrieved_ids

    # Check top score calculation aligns with 1 / (60 + rank)
    top_result = results[0]
    assert top_result.score >= 1.0 / 61.0


def test_bulk_ingest_objects_and_dicts():
    """Verify index_raptor_nodes and index_graphrag_data accept lists of dicts or objects."""
    retriever = HybridStructuredRetriever()

    raptor_nodes = [
        {
            "id": "r_node_10",
            "level": 3,
            "text": "Tree root summary of agent memory systems.",
            "summary": "Agent memory hierarchy overview.",
        }
    ]

    graph_entities = [
        {
            "id": "entity_agent",
            "name": "Autonomous Agent",
            "type": "CONCEPT",
            "description": "An entity that perceives its environment and takes actions.",
        }
    ]

    graph_relations = [
        {
            "id": "rel_mem",
            "source": "Autonomous Agent",
            "target": "Memory Store",
            "type": "HAS_COMPONENT",
            "description": "Agents rely on structured memory stores.",
        }
    ]

    retriever.index_raptor_nodes(raptor_nodes)
    retriever.index_graphrag_data(entities=graph_entities, relationships=graph_relations)

    results = retriever.retrieve("agent memory", top_k=3)
    assert len(results) == 3


def test_empty_query_and_edge_cases():
    """Verify empty queries return empty results and custom top_k bounds are respected."""
    retriever = HybridStructuredRetriever()
    retriever.add_raptor_node("1", 0, "Test content", "Test summary")

    assert retriever.retrieve("") == []
    assert retriever.retrieve("   ") == []

    res = retriever.retrieve("Test", top_k=1)
    assert len(res) <= 1
def test_relationship_target_matching():
    """Verify GraphRAG relationships match queries matching the target entity name."""
    retriever = HybridStructuredRetriever()
    retriever.add_graphrag_relationship(
        relation_id="rel_target",
        source="TransformerModel",
        target="AttentionMechanism",
        type="USES",
        description="Transformer models rely heavily on self-attention.",
    )

    results = retriever.retrieve("AttentionMechanism", top_k=1)
    assert len(results) == 1
    assert results[0].node_id == "rel_target"


def test_integer_ids_and_children_type_safety():
    """Verify integer children and entity_ids do not raise TypeError during citation building."""
    retriever = HybridStructuredRetriever()
    retriever.add_raptor_node(node_id="100", level=1, text="Text", children=[101, 102])
    retriever.add_graphrag_community(community_id="200", entity_ids=[201, 202], summary="Summary")

    results = retriever.retrieve("Text Summary", top_k=2)
    assert len(results) == 2
    for res in results:
        assert isinstance(res.citation.lineage[0], str)


def test_embedding_caching():
    """Verify embedding_fn output is cached on the node dictionary."""
    call_count = 0

    def mock_embed(text: str):
        nonlocal call_count
        call_count += 1
        return np.ones(8, dtype=np.float32)

    retriever = HybridStructuredRetriever(embedding_fn=mock_embed)
    retriever.add_raptor_node(node_id="embed_node", level=0, text="Embedding test text")

    # First retrieval computes embedding
    res1 = retriever.retrieve("Embedding test", top_k=1)
    first_calls = call_count
    assert first_calls > 0

    # Second retrieval reuses cached embedding without re-invoking embedding_fn for the node
    res2 = retriever.retrieve("Embedding test", top_k=1)
    assert call_count == first_calls + 1  # Only +1 for the query embedding


def test_precision_bounds_with_repeated_words():
    """Verify precision score is bounded <= 1.0 even when text contains repeated query terms."""
    retriever = HybridStructuredRetriever()
    retriever.add_raptor_node(
        node_id="rep_node",
        level=0,
        text="apple apple apple apple apple apple apple apple",
        summary="apple apple apple",
    )

    results = retriever.retrieve("apple", top_k=1)
    assert len(results) == 1
    assert results[0].score <= 1.0


def test_deterministic_rrf_ranking():
    """Verify RRF results order is 100% deterministic across multiple invocations."""
    retriever = HybridStructuredRetriever()
    for i in range(10):
        retriever.add_raptor_node(f"r_{i}", 0, f"Common topic text item {i}", f"Summary {i}")
        retriever.add_graphrag_entity(f"e_{i}", f"Entity {i}", "CONCEPT", f"Common topic text item {i}")

    res1 = [r.node_id for r in retriever.retrieve("Common topic text", top_k=5)]
    res2 = [r.node_id for r in retriever.retrieve("Common topic text", top_k=5)]
    assert res1 == res2
def test_index_raptor_nodes_id_zero():
    """Verify node ID 0 is not dropped during index_raptor_nodes."""
    retriever = HybridStructuredRetriever()
    retriever.index_raptor_nodes([{"id": 0, "text": "Zero ID text", "summary": "Zero ID summary"}])
    results = retriever.retrieve("Zero ID", top_k=1)
    assert len(results) == 1
    assert results[0].node_id == "0"


def test_negative_rrf_k_parameter():
    """Verify negative rrf_k override is safely clamped without division by zero."""
    retriever = HybridStructuredRetriever()
    retriever.add_raptor_node("node1", 0, "Quantum physics content", "Quantum physics summary")
    results = retriever.retrieve("Quantum physics", top_k=1, rrf_k=-1)
    assert len(results) == 1
    assert results[0].score > 0


def test_index_graphrag_data_none_id_fallback():
    """Verify items with explicit id=None fall back to entity_id/relation_id/community_id."""
    retriever = HybridStructuredRetriever()
    retriever.index_graphrag_data(
        entities=[{"id": None, "entity_id": "ent_1", "name": "Entity 1", "description": "GraphRAG entity test"}],
        relationships=[{"id": None, "relation_id": "rel_1", "source": "A", "target": "B", "description": "GraphRAG relation test"}],
        communities=[{"id": None, "community_id": "comm_1", "entity_ids": ["ent_1"], "summary": "GraphRAG community test"}],
    )
    res = retriever.retrieve("GraphRAG test", top_k=5)
    assert len(res) == 3


def test_top_k_zero():
    """Verify top_k=0 returns empty results list."""
    retriever = HybridStructuredRetriever()
    retriever.add_raptor_node("node1", 0, "Quantum physics content", "Quantum physics summary")
    results = retriever.retrieve("Quantum physics", top_k=0)
    assert results == []

def test_orthogonal_vector_scoring():
    """Verify semantic_score is 0.0 (not replaced by coverage) when orthogonal vector embedding is evaluated."""
    embedding_fn = lambda text: np.array([0.0, 1.0]) if text == "query" else np.array([1.0, 0.0])
    retriever = HybridStructuredRetriever(embedding_fn=embedding_fn)
    retriever.add_raptor_node("node1", 0, "query term content", "query term summary")
    results = retriever.retrieve("query", top_k=1)
    assert len(results) == 1
    # Semantic score should be 0.0 for orthogonal vector
    raw_score, lex_sc, sem_sc = retriever._compute_scores("query", {"query"}, np.array([0.0, 1.0]), retriever.unified_nodes["raptor_node1"])
    assert sem_sc == 0.0


def test_vector_dimension_mismatch_fallback():
    """Verify dimension mismatch during vector comparison safely falls back to coverage score."""
    embedding_fn = lambda text: np.array([1.0, 0.0, 0.0])  # 3D query vector
    retriever = HybridStructuredRetriever(embedding_fn=embedding_fn)
    # Node contains 2D embedding vector
    retriever.add_raptor_node("node1", 0, "query term content", "query term summary", embedding=np.array([1.0, 0.0]))
    results = retriever.retrieve("query", top_k=1)
    assert len(results) == 1
    # Should fall back to lexical / coverage scoring without crashing
    assert results[0].score > 0


def test_results_merged_by_score_not_source():
    """Verify results are ranked by score, not interleaved by source type (Finding 1)."""
    retriever = HybridStructuredRetriever()
    # Two RAPTOR nodes with strong lexical match
    retriever.add_raptor_node("rap_a", 0, "alpha beta gamma", "alpha beta gamma summary")
    retriever.add_raptor_node("rap_b", 0, "alpha beta delta", "alpha beta delta summary")
    # One GraphRAG entity with weaker match
    retriever.add_graphrag_entity("ent_weak", "alpha", "CONCEPT", "alpha description")

    results = retriever.retrieve("alpha beta gamma", top_k=3)
    # Top two results should both be RAPTOR nodes (higher lexical match), not interleaved
    assert results[0].source_type == "raptor_tree"
    assert results[1].source_type == "raptor_tree"
    # Scores must be in descending order
    assert results[0].score >= results[1].score >= results[2].score


def test_negative_vector_similarity_clamped_to_zero():
    """Verify negative cosine similarity is clamped to 0, not ranked above positive text relevance (Finding 10)."""
    # Embedding that produces negative cosine similarity for node text (no "query" in it)
    def mock_embed(text: str) -> np.ndarray:
        if "query" in text.lower():
            return np.array([1.0, 0.0])
        return np.array([-1.0, 0.0])  # Opposite direction → cos_sim = -1.0

    retriever = HybridStructuredRetriever(embedding_fn=mock_embed)
    # Node text must NOT contain "query" so mock_embed returns the opposite vector
    retriever.add_raptor_node("neg_node", 0, "term content", "term summary")

    _, _, sem_sc = retriever._compute_scores(
        "query term", {"query", "term"}, np.array([1.0, 0.0]), retriever.unified_nodes["raptor_neg_node"]
    )
    # Semantic score must be clamped to 0.0, not negative
    assert sem_sc == 0.0
    assert sem_sc >= 0.0


def test_mixed_vector_presence_consistent_scoring():
    """Verify nodes with and without vectors are scored on a consistent scale (Finding 11)."""
    def mock_embed(text: str) -> np.ndarray:
        if "fail" in text.lower():
            raise ValueError("cannot embed")
        return np.array([1.0, 0.0])

    retriever = HybridStructuredRetriever(embedding_fn=mock_embed)
    # Node A: has a pre-computed embedding aligned with query vector
    retriever.add_raptor_node(
        "node_a", 0, "common topic text", "common topic text",
        embedding=np.array([1.0, 0.0]),
    )
    # Node B: no pre-computed embedding; embedding_fn raises → falls back to lexical-only
    retriever.add_raptor_node(
        "node_b", 0, "fail common topic text", "fail common topic text",
    )

    results = retriever.retrieve("common topic", top_k=2)
    assert len(results) == 2
    # Both nodes should have positive scores (lexical match exists for both)
    for res in results:
        assert res.score > 0
    # Node A (has vector, aligned) should rank higher than Node B (no vector, lexical-only fallback)
    assert results[0].node_id == "node_a"


def test_parent_id_zero_preserved_in_citation():
    """Verify parent ID 0 is not dropped from citation lineage (Finding 12)."""
    retriever = HybridStructuredRetriever()
    retriever.add_raptor_node(
        node_id="child_1",
        level=1,
        text="Child node content",
        summary="Child node summary",
        parent=0,
    )

    results = retriever.retrieve("Child node", top_k=1)
    assert len(results) == 1
    citation = results[0].citation
    # Parent ID 0 must appear in lineage, not be dropped by truthiness check
    parent_entries = [lin for lin in citation.lineage if lin.startswith("Parent:")]
    assert len(parent_entries) == 1
    assert "0" in parent_entries[0]