| import unittest |
| from unittest.mock import MagicMock, patch |
| import os |
|
|
| |
| if "GEMINI_API_KEY" not in os.environ: |
| os.environ["GEMINI_API_KEY"] = "dummy_key" |
| if "TAVILY_API_KEY" not in os.environ: |
| os.environ["TAVILY_API_KEY"] = "dummy_key" |
|
|
| from agentic_rag_v2_graph import build_agentic_rag_v2_graph |
|
|
| class TestRagGraph(unittest.TestCase): |
| @patch('agentic_rag_v2_graph.genai.GenerativeModel') |
| @patch('agentic_rag_v2_graph.TavilyClient') |
| def test_web_search_flow(self, mock_tavily, mock_genai): |
| print("\n\n=== π§ͺ STARTING DRY RUN GRAPH TEST ===") |
| print("Goal: Verify 'research_web' -> 'verifier' -> 'responder' flow without API calls.\n") |
|
|
| |
| |
| DEMO_QUERY = "Who is the father of the computer?" |
| EXPECTED_WEB_CONTENT = "Charles Babbage is considered by many as the father of the computer." |
| VERIFIER_NOTE = "β
VERIFIED: Search results confirm Charles Babbage invented the Analytical Engine." |
| FINAL_ANSWER = "Charles Babbage is the father of the computer." |
| |
|
|
| |
| mock_model = MagicMock() |
| mock_genai.return_value = mock_model |
| |
| |
| def create_response(text): |
| r = MagicMock() |
| r.text = text |
| return r |
|
|
| |
| |
| |
| |
| |
| |
| mock_model.generate_content.side_effect = [ |
| create_response("research_web"), |
| create_response(VERIFIER_NOTE), |
| create_response("responder"), |
| create_response(FINAL_ANSWER) |
| ] |
| |
| |
| mock_tavily_instance = MagicMock() |
| mock_tavily.return_value = mock_tavily_instance |
| mock_tavily_instance.get_search_context.return_value = EXPECTED_WEB_CONTENT |
|
|
| |
| print("π οΈ Building Graph...") |
| try: |
| graph = build_agentic_rag_v2_graph() |
| print("β
Graph built successfully.") |
| except Exception as e: |
| self.fail(f"β Graph build failed: {e}") |
| |
| |
| initial_state = { |
| "messages": [], |
| "query": DEMO_QUERY, |
| "final_answer": "", |
| "next_node": "", |
| "current_tool": "", |
| "tool_outputs": [], |
| "verification_notes": "", |
| "retries": 0 |
| } |
| |
| print("\nπ Invoking Graph (Mocked LLM)...") |
| result = graph.invoke(initial_state, config={"configurable": {"thread_id": "test_dry_run"}}) |
| |
| |
| print("\n\n=== π TEST RESULT ANALYSIS ===") |
| print(f"Final Answer: {result['final_answer']}") |
| print(f"Verification Notes: {result['verification_notes']}") |
| |
| self.assertIn("VERIFIED", result['verification_notes'], "β verifier_node did not populate verification_notes!") |
| self.assertIn(FINAL_ANSWER, result['final_answer'], "β Responder did not fail gracefully.") |
| |
| print("\nβ
SUCCESS: The Graph followed the correct path: Supervisor -> Web -> Verifier -> Supervisor -> Responder") |
| print("β
SUCCESS: 'verifier_node' executed and produced notes.") |
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|