Spaces:
Runtime error
Runtime error
| import pytest | |
| from unittest.mock import patch, MagicMock | |
| from langchain_core.messages import AIMessage, HumanMessage | |
| def _run_input_agent_with_llm_response(llm_content, initial_message="I have an idea about bike lanes"): | |
| """Helper: run input_agent with a mocked LLM returning llm_content.""" | |
| mock_llm = MagicMock() | |
| mock_llm.invoke.return_value = AIMessage(content=llm_content) | |
| state = { | |
| "messages": [HumanMessage(content=initial_message)], | |
| "location": "New York City, NY", | |
| "org": "nyc", | |
| "category": None, | |
| "collected_details": {}, | |
| "email_cc": [], | |
| "conversation_complete": False, | |
| "draft_approved": False, | |
| } | |
| with patch("src.agents.input_agent.get_llm", return_value=mock_llm): | |
| from src.agents.input_agent import input_agent | |
| return input_agent(state) | |
| def test_extracts_category_idea(): | |
| result = _run_input_agent_with_llm_response( | |
| "Great idea!\nCATEGORY: idea\nLOCATION: New York City" | |
| ) | |
| assert result["category"] == "idea" | |
| def test_extracts_category_complaint(): | |
| result = _run_input_agent_with_llm_response( | |
| "I hear you.\nCATEGORY: complaint\nLOCATION: New York City" | |
| ) | |
| assert result["category"] == "complaint" | |
| def test_extracts_location_from_tags(): | |
| result = _run_input_agent_with_llm_response( | |
| "Understood.\nCATEGORY: idea\nLOCATION: Tamil Nadu" | |
| ) | |
| assert result["location"] == "Tamil Nadu" | |
| def test_category_is_lowercased(): | |
| result = _run_input_agent_with_llm_response( | |
| "CATEGORY: IDEA\nLOCATION: London" | |
| ) | |
| assert result["category"] == "idea" | |
| def test_display_text_strips_tag_lines(): | |
| result = _run_input_agent_with_llm_response( | |
| "Here is my response.\nCATEGORY: idea\nLOCATION: NYC" | |
| ) | |
| last_ai = result["messages"][-1] | |
| assert "CATEGORY:" not in last_ai.content | |
| assert "LOCATION:" not in last_ai.content | |
| assert "Here is my response." in last_ai.content | |
| def test_category_found_without_location_still_routes(): | |
| result = _run_input_agent_with_llm_response( | |
| "Your request is clear.\nCATEGORY: complaint" | |
| ) | |
| assert result["category"] == "complaint" | |
| assert "location" not in result | |
| def test_sets_conversation_complete_false_on_category(): | |
| result = _run_input_agent_with_llm_response( | |
| "CATEGORY: idea\nLOCATION: NYC" | |
| ) | |
| assert result["conversation_complete"] is False | |
| def test_sets_draft_approved_false_on_category(): | |
| result = _run_input_agent_with_llm_response( | |
| "CATEGORY: idea\nLOCATION: NYC" | |
| ) | |
| assert result["draft_approved"] is False | |
| def test_no_category_calls_interrupt(mocker): | |
| mock_llm = MagicMock() | |
| mock_llm.invoke.return_value = AIMessage(content="Which city is your email about?") | |
| mock_interrupt = mocker.patch("src.agents.input_agent.interrupt", return_value="London") | |
| state = { | |
| "messages": [HumanMessage(content="I have a complaint")], | |
| "location": "New York City, NY", | |
| "org": "nyc", | |
| "category": None, | |
| "collected_details": {}, | |
| "email_cc": [], | |
| "conversation_complete": False, | |
| "draft_approved": False, | |
| } | |
| with patch("src.agents.input_agent.get_llm", return_value=mock_llm): | |
| from src.agents.input_agent import input_agent | |
| result = input_agent(state) | |
| mock_interrupt.assert_called_once() | |
| # User reply should be added as HumanMessage | |
| human_messages = [m for m in result["messages"] if isinstance(m, HumanMessage)] | |
| assert any(m.content == "London" for m in human_messages) | |
| def test_location_saved_even_when_no_category(mocker): | |
| mock_llm = MagicMock() | |
| mock_llm.invoke.return_value = AIMessage( | |
| content="Which city?\nLOCATION: Chennai" | |
| ) | |
| mocker.patch("src.agents.input_agent.interrupt", return_value="complaint about roads") | |
| state = { | |
| "messages": [HumanMessage(content="CM Vijay should fix this")], | |
| "location": "New York City, NY", | |
| "org": "nyc", | |
| "category": None, | |
| "collected_details": {}, | |
| "email_cc": [], | |
| "conversation_complete": False, | |
| "draft_approved": False, | |
| } | |
| with patch("src.agents.input_agent.get_llm", return_value=mock_llm): | |
| from src.agents.input_agent import input_agent | |
| result = input_agent(state) | |
| assert result.get("location") == "Chennai" | |