| | 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() |
| |
|