startrail-org--leann/tests/test_react_dual_source.py

"""
Tests for ReAct agent dual-source routing (issue #283).
Covers: prompt adaptation, web fallback, visit_page errors, routing, and pipeline.
"""

from unittest.mock import MagicMock, patch

from leann.api import SearchResult
from leann.react_agent import ReActAgent
from leann.web_search import WebSearcher


def _make_searcher() -> MagicMock:
    """Create a lightweight mocked searcher for deterministic unit tests."""
    searcher = MagicMock()
    searcher.search.return_value = [
        SearchResult(
            id="1",
            score=0.91,
            text="LEANN achieves 97% storage reduction via graph pruning.",
            metadata={"source": "docs"},
        ),
        SearchResult(
            id="2",
            score=0.83,
            text="The search function uses HNSW for approximate nearest neighbors.",
            metadata={"source": "code"},
        ),
    ]
    return searcher


# ── 1. Dual-source presence ──────────────────────────────────────────


def test_prompt_includes_web_tools_when_key_present():
    """When SERPER_API_KEY is set, prompt should list all three tools."""
    searcher = _make_searcher()
    agent = ReActAgent(searcher=searcher, llm=MagicMock(), serper_api_key="test-key")
    prompt = agent._create_react_prompt("test question", 1, [])
    assert "web_search" in prompt
    assert "visit_page" in prompt
    assert "leann_search" in prompt
    assert agent.web_search_available is True


def test_prompt_excludes_web_tools_when_no_key():
    """When no SERPER_API_KEY, prompt should only show leann_search."""
    searcher = _make_searcher()
    with patch.dict("os.environ", {}, clear=False):
        agent = ReActAgent(searcher=searcher, llm=MagicMock(), serper_api_key=None)
    prompt = agent._create_react_prompt("test question", 1, [])
    assert "leann_search" in prompt
    assert "web_search" not in prompt or "not available" in prompt
    assert agent.web_search_available is False


# ── 2. Routing behavior ─────────────────────────────────────────────


def test_local_only_routing():
    """Agent uses leann_search for local queries."""
    searcher = _make_searcher()
    mock_llm = MagicMock()
    mock_llm.ask.side_effect = [
        'Thought: This is about our codebase.\nAction: leann_search("LEANN storage reduction")',
        "Thought: I have the answer.\nAction: Final Answer: LEANN saves 97% storage.",
    ]
    agent = ReActAgent(searcher=searcher, llm=mock_llm, max_iterations=3)
    agent.run("How does LEANN reduce storage?", top_k=2)

    assert len(agent.search_history) >= 1
    assert agent.search_history[0]["source"] == "local"
    assert "leann_search" in agent.search_history[0]["action"]


def test_web_only_routing():
    """Agent uses web_search for web queries (mocked)."""
    searcher = _make_searcher()

    with patch.object(WebSearcher, "search") as mock_web:
        mock_web.return_value = [
            {
                "title": "Python 3.13 News",
                "link": "https://python.org",
                "snippet": "New features in 3.13",
            }
        ]
        mock_llm = MagicMock()
        mock_llm.ask.side_effect = [
            'Thought: Need current web info.\nAction: web_search("Python 3.13 features")',
            "Thought: Got it.\nAction: Final Answer: Python 3.13 has new features.",
        ]
        agent = ReActAgent(
            searcher=searcher, llm=mock_llm, max_iterations=3, serper_api_key="test-key"
        )
        agent.run("What's new in Python 3.13?", top_k=2)

        assert len(agent.search_history) >= 1
        assert agent.search_history[0]["source"] == "web"
        mock_web.assert_called_once()


def test_mixed_routing_local_then_web():
    """Agent uses leann_search first, then web_search."""
    searcher = _make_searcher()

    with patch.object(WebSearcher, "search") as mock_web:
        mock_web.return_value = [
            {
                "title": "Best Practices",
                "link": "https://example.com",
                "snippet": "Current best practices...",
            }
        ]
        mock_llm = MagicMock()
        mock_llm.ask.side_effect = [
            'Thought: First check our code.\nAction: leann_search("search implementation")',
            'Thought: Now check best practices.\nAction: web_search("vector DB best practices")',
            "Thought: I can compare now.\nAction: Final Answer: Our search is good but could improve.",
        ]
        agent = ReActAgent(
            searcher=searcher, llm=mock_llm, max_iterations=5, serper_api_key="test-key"
        )
        agent.run("Compare our search with best practices", top_k=2)

        assert len(agent.search_history) >= 2
        assert agent.search_history[0]["source"] == "local"
        assert agent.search_history[1]["source"] == "web"


# ── 3. Pipeline correctness ─────────────────────────────────────────


def test_web_results_formatted_as_observations():
    """Web search results are properly formatted and passed to the next iteration."""
    searcher = _make_searcher()

    with patch.object(WebSearcher, "search") as mock_web:
        mock_web.return_value = [
            {"title": "Result A", "link": "https://a.com", "snippet": "Snippet A"},
            {"title": "Result B", "link": "https://b.com", "snippet": "Snippet B"},
        ]
        mock_llm = MagicMock()
        mock_llm.ask.side_effect = [
            'Thought: Search the web.\nAction: web_search("test query")',
            "Thought: Done.\nAction: Final Answer: Found it.",
        ]
        agent = ReActAgent(
            searcher=searcher, llm=mock_llm, max_iterations=3, serper_api_key="test-key"
        )
        agent.run("test", top_k=2)

        # The second LLM call should contain the formatted web results
        second_prompt = mock_llm.ask.call_args_list[1][0][0]
        assert "Result A" in second_prompt or "Snippet A" in second_prompt


def test_visit_page_content_truncated():
    """visit_page content is truncated to 15k chars."""
    searcher = _make_searcher()

    with patch.object(WebSearcher, "get_page_content") as mock_fetch:
        mock_fetch.return_value = "x" * 20000
        mock_llm = MagicMock()
        mock_llm.ask.side_effect = [
            'Thought: Read this page.\nAction: visit_page("https://docs.python.org/3")',
            "Thought: Done.\nAction: Final Answer: Got the content.",
        ]
        agent = ReActAgent(
            searcher=searcher, llm=mock_llm, max_iterations=3, serper_api_key="test-key"
        )
        agent.run("read docs", top_k=2)

        second_prompt = mock_llm.ask.call_args_list[1][0][0]
        # Content should be truncated — observation contains at most 15000 chars of content
        assert "x" * 15000 in second_prompt
        assert "x" * 15001 not in second_prompt


def test_leann_search_results_include_scores():
    """Local search results include scores and text snippets."""
    searcher = _make_searcher()
    mock_llm = MagicMock()
    mock_llm.ask.side_effect = [
        'Thought: Search locally.\nAction: leann_search("LEANN")',
        "Thought: Done.\nAction: Final Answer: Found results.",
    ]
    agent = ReActAgent(searcher=searcher, llm=mock_llm, max_iterations=3)
    agent.run("tell me about LEANN", top_k=2)

    second_prompt = mock_llm.ask.call_args_list[1][0][0]
    assert "Score:" in second_prompt
    assert "Result" in second_prompt


# ── 4. Edge cases ────────────────────────────────────────────────────


def test_web_search_no_api_key_graceful():
    """web_search without API key returns clear fallback message, not a crash."""
    searcher = _make_searcher()

    mock_llm = MagicMock()
    mock_llm.ask.side_effect = [
        'Thought: Try web.\nAction: web_search("Python 3.13")',
        'Thought: Web failed, try local.\nAction: leann_search("search")',
        "Thought: Done.\nAction: Final Answer: Here's what I found locally.",
    ]
    with patch.dict("os.environ", {}, clear=False):
        agent = ReActAgent(searcher=searcher, llm=mock_llm, max_iterations=5, serper_api_key=None)
    answer = agent.run("test", top_k=2)

    assert agent.search_history[0]["results_count"] == 0
    assert agent.search_history[0]["source"] == "web"
    # Agent should not crash and should produce an answer
    assert answer is not None and len(answer) > 0


def test_web_search_invalid_key_graceful():
    """Invalid Serper key returns error, agent continues."""
    searcher = _make_searcher()

    with patch.object(WebSearcher, "search") as mock_web:
        mock_web.return_value = [
            {"title": "Error", "link": "", "snippet": "Web Search failed:401 Unauthorized"}
        ]
        mock_llm = MagicMock()
        mock_llm.ask.side_effect = [
            'Thought: Try web.\nAction: web_search("test")',
            "Thought: Web failed.\nAction: Final Answer: Could not search web.",
        ]
        agent = ReActAgent(
            searcher=searcher, llm=mock_llm, max_iterations=3, serper_api_key="bad-key"
        )
        answer = agent.run("test", top_k=2)

        assert agent.search_history[0]["results_count"] == 0
        assert answer is not None


def test_visit_page_404_graceful():
    """visit_page on a 404 URL returns error string, agent continues."""
    searcher = _make_searcher()

    with patch.object(WebSearcher, "get_page_content") as mock_fetch:
        mock_fetch.return_value = "Error fetching content: 404 Not Found"
        mock_llm = MagicMock()
        mock_llm.ask.side_effect = [
            'Thought: Read page.\nAction: visit_page("https://example.com/404")',
            "Thought: Page not found.\nAction: Final Answer: Page was not accessible.",
        ]
        agent = ReActAgent(
            searcher=searcher, llm=mock_llm, max_iterations=3, serper_api_key="test-key"
        )
        answer = agent.run("read page", top_k=2)

        assert agent.search_history[0]["results_count"] == 0
        assert "not accessible" in answer.lower() or len(answer) > 0


def test_max_iterations_with_mixed_sources():
    """Agent hitting max iterations with mixed sources still produces an answer."""
    searcher = _make_searcher()

    with patch.object(WebSearcher, "search") as mock_web:
        mock_web.return_value = [
            {"title": "Web Result", "link": "https://example.com", "snippet": "Some web info"}
        ]
        mock_llm = MagicMock()
        # Never gives a Final Answer — forces max iterations
        mock_llm.ask.side_effect = [
            'Thought: Search locally.\nAction: leann_search("storage")',
            'Thought: Now web.\nAction: web_search("vector databases")',
            # Max iterations reached — the agent will ask for a final answer
            "Based on all searches, LEANN is a storage-efficient vector DB.",
        ]
        agent = ReActAgent(
            searcher=searcher, llm=mock_llm, max_iterations=2, serper_api_key="test-key"
        )
        answer = agent.run("Compare approaches", top_k=2)

        assert len(agent.search_history) == 2
        sources = {h["source"] for h in agent.search_history}
        assert "local" in sources
        assert "web" in sources
        assert answer is not None and len(answer) > 0


def test_search_history_has_source_field():
    """search_history entries include 'source' field for easy inspection."""
    searcher = _make_searcher()
    mock_llm = MagicMock()
    mock_llm.ask.side_effect = [
        'Thought: Local search.\nAction: leann_search("test")',
        "Thought: Done.\nAction: Final Answer: Done.",
    ]
    agent = ReActAgent(searcher=searcher, llm=mock_llm, max_iterations=3)
    agent.run("test", top_k=2)

    assert "source" in agent.search_history[0]
    assert agent.search_history[0]["source"] in ("local", "web")


def test_zero_results_on_iteration_two_does_not_abort_early():
    """A single empty tool result must not force an ungrounded final answer (#381)."""
    searcher = _make_searcher()
    searcher.search.side_effect = [
        [
            SearchResult(
                id="1",
                score=0.9,
                text="Partial LEANN benchmark mention.",
                metadata={"source": "docs"},
            )
        ],
        [],
        [
            SearchResult(
                id="2",
                score=0.95,
                text="LEANN recall@10 = 0.95 vs FAISS 0.97 per the paper.",
                metadata={"source": "web"},
            )
        ],
    ]

    mock_llm = MagicMock()
    mock_llm.ask.side_effect = [
        'Thought: Start local.\nAction: leann_search("LEANN benchmark")',
        'Thought: Narrow local query.\nAction: leann_search("LEANN recall numbers vs faiss")',
        'Thought: Try web.\nAction: web_search("LEANN paper recall comparison FAISS")',
        "Thought: Found it.\nAction: Final Answer: LEANN achieves recall@10 = 0.95 vs FAISS 0.97.",
    ]

    with patch.object(WebSearcher, "search") as mock_web:
        mock_web.return_value = [
            {
                "title": "LEANN paper",
                "link": "https://example.com/paper",
                "snippet": "recall@10 comparison",
            }
        ]
        agent = ReActAgent(
            searcher=searcher,
            llm=mock_llm,
            max_iterations=5,
            serper_api_key="test-key",
        )
        answer = agent.run("What is LEANN's recall vs FAISS at recall@10?", top_k=2)

    assert len(agent.search_history) == 3
    assert agent.search_history[1]["results_count"] == 0
    assert "0.95" in answer
    assert mock_llm.ask.call_count == 4


def test_transient_web_failure_allows_retry_on_next_iteration():
    """Transient web_search failure should be retried, not terminate the loop (#381)."""
    searcher = _make_searcher()
    searcher.search.return_value = [
        SearchResult(
            id="1",
            score=0.8,
            text="Local setup notes.",
            metadata={"source": "docs"},
        )
    ]

    mock_llm = MagicMock()
    mock_llm.ask.side_effect = [
        'Thought: Check local docs.\nAction: leann_search("project setup")',
        'Thought: Check web.\nAction: web_search("latest setup guide")',
        'Thought: Retry web.\nAction: web_search("setup guide tutorial")',
        "Thought: Found it.\nAction: Final Answer: per the latest docs, run `make install`.",
    ]

    with patch.object(WebSearcher, "search") as mock_web:
        mock_web.side_effect = [
            [{"title": "Error", "link": "", "snippet": "Web Search failed:502 Bad Gateway"}],
            [
                {
                    "title": "Setup guide",
                    "link": "https://example.com/setup",
                    "snippet": "Run `make install` to set up the project.",
                }
            ],
        ]
        agent = ReActAgent(
            searcher=searcher,
            llm=mock_llm,
            max_iterations=5,
            serper_api_key="test-key",
        )
        answer = agent.run("How do I set up the project?", top_k=2)

    assert len(agent.search_history) == 3
    assert agent.search_history[1]["results_count"] == 0
    assert agent.search_history[2]["results_count"] == 1
    assert "make install" in answer
    assert mock_llm.ask.call_count == 4