tests/test_chat_agent.py

"""tests/test_chat_agent.py — config threading through the chat Q&A agent."""
from __future__ import annotations

from unittest.mock import MagicMock, patch

from langchain_core.messages import AIMessage

from agent.llm import RunConfig


def _fake_llm_no_tool_calls(answer_text: str = "The answer is 42."):
    """Return a MagicMock standing in for a bound chat model that answers directly."""
    llm = MagicMock()
    llm.bind_tools.return_value = llm
    llm.invoke.return_value = AIMessage(content=answer_text, tool_calls=[])
    return llm


@patch("agent.chat_agent.metrics_db.get_all_metrics")
@patch("agent.chat_agent.make_chat_model")
def test_answer_question_threads_config_to_factory(mock_make_chat_model, mock_get_all_metrics):
    mock_get_all_metrics.return_value = [
        {"period": "Q1 2025", "form_type": "10-Q", "filing_date": "2025-05-01"}
    ]
    mock_make_chat_model.return_value = _fake_llm_no_tool_calls()

    cfg = RunConfig(provider="openai", model="gpt-5-mini", api_key="sk-test")
    from agent.chat_agent import answer_question
    result = answer_question("nvda", "What changed?", [], config=cfg)

    assert result["answer"] == "The answer is 42."
    assert mock_make_chat_model.call_count >= 1
    called_cfg = mock_make_chat_model.call_args_list[0].args[0]
    assert called_cfg is cfg


@patch("agent.chat_agent.metrics_db.get_all_metrics")
@patch("agent.chat_agent.make_chat_model")
def test_answer_question_uses_default_config_when_none_given(mock_make_chat_model, mock_get_all_metrics):
    mock_get_all_metrics.return_value = [
        {"period": "Q1 2025", "form_type": "10-Q", "filing_date": "2025-05-01"}
    ]
    mock_make_chat_model.return_value = _fake_llm_no_tool_calls()

    from agent.chat_agent import answer_question
    answer_question("NVDA", "What changed?", [])

    called_cfg = mock_make_chat_model.call_args_list[0].args[0]
    assert called_cfg.provider == "anthropic"
    assert called_cfg.api_key is None


@patch("agent.chat_agent.metrics_db.get_all_metrics")
@patch("agent.chat_agent.build_system_message")
@patch("agent.chat_agent.make_chat_model")
def test_answer_question_builds_provider_aware_system_message(
    mock_make_chat_model, mock_build_system_message, mock_get_all_metrics
):
    mock_get_all_metrics.return_value = [
        {"period": "Q1 2025", "form_type": "10-Q", "filing_date": "2025-05-01"}
    ]
    mock_make_chat_model.return_value = _fake_llm_no_tool_calls()
    from langchain_core.messages import SystemMessage
    mock_build_system_message.return_value = SystemMessage(content="sys")

    cfg = RunConfig(provider="openai", model="gpt-5-mini", api_key="sk-test")
    from agent.chat_agent import answer_question
    answer_question("NVDA", "hi", [], config=cfg)

    assert mock_build_system_message.call_args.args[0] is cfg