80 lines
3.1 KiB
Python
80 lines
3.1 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
from ea_chatbot.graph.nodes.query_analyzer import query_analyzer_node, QueryAnalysis
|
|
from ea_chatbot.graph.state import AgentState
|
|
|
|
@pytest.fixture
|
|
def mock_state():
|
|
return {
|
|
"messages": [],
|
|
"question": "Show me the 2024 results for Florida",
|
|
"analysis": None,
|
|
"next_action": ""
|
|
}
|
|
|
|
@patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model")
|
|
def test_query_analyzer_data_analysis(mock_get_llm, mock_state):
|
|
"""Test that a clear data analysis query is routed to the planner."""
|
|
# Mock the LLM and the structured output runnable
|
|
mock_llm_instance = MagicMock()
|
|
mock_get_llm.return_value = mock_llm_instance
|
|
mock_structured_llm = MagicMock()
|
|
mock_llm_instance.with_structured_output.return_value = mock_structured_llm
|
|
# Define the expected Pydantic result
|
|
expected_analysis = QueryAnalysis(
|
|
data_required=["2024 results", "Florida"],
|
|
unknowns=[],
|
|
ambiguities=[],
|
|
conditions=[],
|
|
next_action="plan"
|
|
)
|
|
# When structured_llm.invoke is called with messages, return the Pydantic object
|
|
mock_structured_llm.invoke.return_value = expected_analysis
|
|
|
|
new_state_update = query_analyzer_node(mock_state)
|
|
|
|
assert new_state_update["next_action"] == "plan"
|
|
assert "2024 results" in new_state_update["analysis"]["data_required"]
|
|
|
|
@patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model")
|
|
def test_query_analyzer_ambiguous(mock_get_llm, mock_state):
|
|
"""Test that an ambiguous query is routed to clarification."""
|
|
mock_state["question"] = "What happened?"
|
|
|
|
mock_llm_instance = MagicMock()
|
|
mock_get_llm.return_value = mock_llm_instance
|
|
mock_structured_llm = MagicMock()
|
|
mock_llm_instance.with_structured_output.return_value = mock_structured_llm
|
|
expected_analysis = QueryAnalysis(
|
|
data_required=[],
|
|
unknowns=["What event?"],
|
|
ambiguities=[],
|
|
conditions=[],
|
|
next_action="clarify"
|
|
)
|
|
|
|
mock_structured_llm.invoke.return_value = expected_analysis
|
|
|
|
new_state_update = query_analyzer_node(mock_state)
|
|
|
|
assert new_state_update["next_action"] == "clarify"
|
|
assert len(new_state_update["analysis"]["unknowns"]) > 0
|
|
|
|
@patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model")
|
|
def test_query_analyzer_uses_config(mock_get_llm, mock_state, monkeypatch):
|
|
"""Test that the node uses the configured LLM settings."""
|
|
monkeypatch.setenv("QUERY_ANALYZER_LLM__MODEL", "gpt-3.5-turbo")
|
|
|
|
mock_llm_instance = MagicMock()
|
|
mock_get_llm.return_value = mock_llm_instance
|
|
mock_structured_llm = MagicMock()
|
|
mock_llm_instance.with_structured_output.return_value = mock_structured_llm
|
|
mock_structured_llm.invoke.return_value = QueryAnalysis(
|
|
data_required=[], unknowns=[], ambiguities=[], conditions=[], next_action="plan"
|
|
)
|
|
|
|
query_analyzer_node(mock_state)
|
|
|
|
# Verify get_llm_model was called with the overridden config
|
|
called_config = mock_get_llm.call_args[0][0]
|
|
assert called_config.model == "gpt-3.5-turbo" |