backend/tests/test_rag_tools.py

import pytest

from app.rag import tools
from app.rag.tools import (
    CALCULATOR_TOOL,
    TOOLS,
    WEB_SEARCH_TOOL,
    MathTool,
    PDFSearchTool,
    WebSearchTool,
    calculate_expression,
    execute_tool,
    web_search,
)


def test_calculator_tool_evaluates_basic_expression():
    assert calculate_expression("1000 - 250") == "750"
    assert calculate_expression("10 + 5 * 2") == "20"
    assert calculate_expression("10 / 4") == "2.5"


def test_calculator_tool_rejects_unsafe_expression():
    try:
        calculate_expression("__import__('os').system('echo x')")
    except ValueError as exc:
        assert "Invalid calculator expression" in str(
            exc
        ) or "Unsupported expression" in str(exc)
    else:
        assert False, "Unsafe expressions should not be evaluated"


def test_math_tool_run():
    tool = MathTool()
    result = tool.run({"expression": "12 * 3"})
    assert "Result: 36" in result


def test_math_tool_metadata():
    tool = MathTool()
    assert tool.name == "calculator"
    assert "mathematical calculations" in tool.description
    assert tool.args_schema is not None


def test_execute_tool_dispatches_calculator():
    assert execute_tool("calculator", {"expression": "(1000 - 250) * 0.2"}) == "150"


def test_execute_tool_dispatches_web_search(monkeypatch):
    calls = []

    def fake_web_search(query, max_results=5):
        calls.append((query, max_results))
        return "mock result"

    monkeypatch.setattr(tools, "web_search", fake_web_search)

    assert (
        execute_tool("web_search", {"query": "latest rag papers", "max_results": 2})
        == "mock result"
    )
    assert calls == [("latest rag papers", 2)]


@pytest.mark.parametrize(
    ("name", "arguments", "message"),
    [
        ("calculator", {"expression": ""}, "calculator tool requires"),
        ("web_search", {"query": "   "}, "web_search tool requires"),
        ("unknown", {}, "Unknown tool"),
    ],
)
def test_execute_tool_rejects_invalid_inputs(name, arguments, message):
    with pytest.raises(ValueError, match=message):
        execute_tool(name, arguments)


class FakeDDGS:
    def __init__(self, results=None, error=None):
        self.results = results or []
        self.error = error
        self.calls = []

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, traceback):
        return False

    def text(self, query, max_results=5):
        self.calls.append((query, max_results))
        if self.error:
            raise self.error
        return self.results


def test_web_search_formats_duckduckgo_results(monkeypatch):
    fake_ddgs = FakeDDGS(
        results=[
            {
                "title": "First result",
                "href": "https://example.com/one",
                "body": "Useful snippet",
            },
            {
                "title": "Second result",
                "href": "https://example.com/two",
                "body": "Another snippet",
            },
        ]
    )
    monkeypatch.setattr(tools, "DDGS", lambda: fake_ddgs)

    result = web_search("agentic rag", max_results=2)

    assert fake_ddgs.calls == [("agentic rag", 2)]
    assert "1. **First result**" in result
    assert "URL: https://example.com/one" in result
    assert "Useful snippet" in result
    assert "2. **Second result**" in result


def test_web_search_handles_empty_results(monkeypatch):
    monkeypatch.setattr(tools, "DDGS", lambda: FakeDDGS(results=[]))

    assert web_search("missing topic") == "No web search results found."


def test_web_search_handles_provider_errors(monkeypatch):
    monkeypatch.setattr(
        tools, "DDGS", lambda: FakeDDGS(error=RuntimeError("network unavailable"))
    )

    assert web_search("agentic rag").startswith(
        "Web search failed: network unavailable"
    )


def test_web_search_tool_uses_shared_search_function(monkeypatch):
    monkeypatch.setattr(tools, "web_search", lambda query: f"searched: {query}")

    result = WebSearchTool().run({"query": "retrieval augmented generation"})

    assert result == "searched: retrieval augmented generation"


def test_pdf_search_tool_formats_chunks_and_graph_context(monkeypatch):
    chunks = [
        {"filename": "alpha.pdf", "page": 1, "text": "Alpha revenue grew by 12%."},
        {"filename": "beta.pdf", "page": 4, "text": "Beta margin narrowed."},
    ]
    retrieve_calls = []
    graph_calls = []

    def fake_retrieve(query, user_id, document_id=None, document_ids=None, top_k=None):
        retrieve_calls.append((query, user_id, document_id))
        return chunks

    def fake_get_entity_context(query, user_id, document_id=None):
        graph_calls.append((query, user_id, document_id))
        return "Alpha -> acquired -> Beta"

    monkeypatch.setattr(tools, "retrieve", fake_retrieve)
    monkeypatch.setattr(tools, "get_entity_context", fake_get_entity_context)

    tool = PDFSearchTool(user_id="user-1", document_id="doc-1")
    result = tool.run({"query": "What happened to Alpha?"})

    assert retrieve_calls == [("What happened to Alpha?", "user-1", "doc-1")]
    assert graph_calls == [("What happened to Alpha?", "user-1", "doc-1")]
    assert tool.last_sources == chunks
    assert "Excerpt 1 (alpha.pdf, Page 1):\nAlpha revenue grew by 12%." in result
    assert "Excerpt 2 (beta.pdf, Page 4):\nBeta margin narrowed." in result
    assert "Additional Relationships found:\nAlpha -> acquired -> Beta" in result


def test_pdf_search_tool_skips_graph_lookup_when_no_chunks(monkeypatch):
    graph_called = False

    def fake_get_entity_context(*_args, **_kwargs):
        nonlocal graph_called
        graph_called = True
        return "unused"

    monkeypatch.setattr(tools, "retrieve", lambda **_kwargs: [])
    monkeypatch.setattr(tools, "get_entity_context", fake_get_entity_context)

    tool = PDFSearchTool(user_id="user-1")
    result = tool.run({"query": "missing"})

    assert result == "No relevant information found in the documents."
    assert tool.last_sources == []
    assert graph_called is False


def test_pdf_search_tool_returns_error_message_on_retriever_failure(monkeypatch):
    def failing_retrieve(**_kwargs):
        raise RuntimeError("vector store offline")

    monkeypatch.setattr(tools, "retrieve", failing_retrieve)

    result = PDFSearchTool(user_id="user-1").run({"query": "anything"})

    assert result == "Error searching documents: vector store offline"


def test_huggingface_tool_definitions_expose_expected_schemas():
    calculator_schema = CALCULATOR_TOOL.function.parameters
    web_schema = WEB_SEARCH_TOOL.function.parameters

    assert CALCULATOR_TOOL.function.name == "calculator"
    assert calculator_schema["required"] == ["expression"]
    assert calculator_schema["properties"]["expression"]["type"] == "string"

    assert WEB_SEARCH_TOOL.function.name == "web_search"
    assert web_schema["required"] == ["query"]
    assert web_schema["properties"]["query"]["type"] == "string"
    assert web_schema["properties"]["max_results"]["default"] == 5
    assert [tool.function.name for tool in TOOLS] == ["calculator", "web_search"]