Spaces:
Sleeping
Sleeping
| from args import Args | |
| from graph import State, Nodes, Edges | |
| from graph_builder import GraphBuilder | |
| import unittest | |
| class TestAlfredAgent(unittest.TestCase): | |
| """Test suite for the Alfred agent""" | |
| def setUp(self): | |
| """Set up test fixtures""" | |
| self.nodes = Nodes() | |
| self.edges = Edges() | |
| self.builder = GraphBuilder() | |
| self.graph = self.builder.build_agent_graph() | |
| def test_manager_node(self): | |
| """ | |
| Test the manager node functionality. | |
| Orchestrates the workflow by delegating tasks to specialized nodes and integrating their outputs | |
| """ | |
| nodes = Nodes() | |
| # Create a test state | |
| test_state: State = { | |
| "initial_query": "query", | |
| "messages": ["query"], # Manager's context | |
| "task_progress": [], # Solver's context | |
| "audit_interval": 2, | |
| "manager_queries": 0, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| # Test the node function | |
| print(f"Testing 'manager' node...") | |
| nodes.manager_node(test_state) | |
| # Assert that manager_queries has been incremented | |
| self.assertEqual(test_state["manager_queries"], 1, "Manager queries should be incremented from 0 to 1") | |
| # Assert that a new message has been added to the messages list | |
| self.assertEqual(len(test_state["messages"]), 2, "Messages list should contain 2 items: the initial query and a new message from the manager node") | |
| # Test audit interval behaviour | |
| test_state = nodes.manager_node(test_state) | |
| # Assert that manager_queries has been incremented | |
| self.assertEqual(test_state["manager_queries"], 2, "Manager queries should be incremented from 1 to 2") | |
| # Assert that a new message has been added to the messages list | |
| self.assertEqual(len(test_state["messages"]), 2, "Messages list should contain 2 items: the initial 2 messages and no additional message as it is the audit interval") | |
| already_tested_messages = test_state["messages"] | |
| expected_state: State = { | |
| "initial_query": "query", | |
| "messages": already_tested_messages, # Manager's context | |
| "task_progress": [test_state["messages"][-1]], # Solver's context | |
| "audit_interval": 2, | |
| "manager_queries": 2, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| self.assertEqual(test_state, expected_state, "The state after manager node execution should match the expected state with manager_queries=2 and no additional messages added during audit interval") | |
| print(f"State after node execution: {test_state}") | |
| def test_final_answer_node(self): | |
| """ | |
| Test the final_answer node functionality. | |
| Formats and delivers the final response to the user | |
| """ | |
| nodes = Nodes() | |
| # Prepare a state with messages and required fields | |
| test_state: State = { | |
| "initial_query": "What is the capital of France?", | |
| "messages": ["What is the capital of France?", "The capital of France is Paris."], | |
| "task_progress": [], | |
| "audit_interval": 2, | |
| "manager_queries": 2, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| print(f"Testing 'final_answer' node...") | |
| nodes.final_answer_node(test_state) | |
| # The last message should be the instruction | |
| self.assertIn("Formulate a definitive final answer", test_state["messages"][-1]) | |
| # The final_response should be set and not None | |
| self.assertIsNotNone(test_state["final_response"]) | |
| print(f"State after node execution: {test_state}") | |
| def test_auditor_node(self): | |
| """ | |
| Test the auditor node functionality. | |
| Reviews manager's outputs for accuracy, safety, and quality | |
| """ | |
| nodes = Nodes() | |
| test_state: State = { | |
| "initial_query": "What is the capital of France?", | |
| "messages": ["What is the capital of France?", "The capital of France is Paris."], | |
| "task_progress": [], | |
| "audit_interval": 2, | |
| "manager_queries": 2, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| print(f"Testing 'auditor' node...") | |
| nodes.auditor_node(test_state) | |
| # Auditor appends a message | |
| self.assertGreaterEqual(len(test_state["messages"]), 3) | |
| print(f"State after node execution: {test_state}") | |
| def test_solver_node(self): | |
| """ | |
| Test the solver node functionality. | |
| Central problem-solving node that coordinates with specialized experts based on task requirements | |
| """ | |
| nodes = Nodes() | |
| test_state: State = { | |
| "initial_query": "What is the capital of France?", | |
| "messages": ["What is the capital of France?"], | |
| "task_progress": ["Solve: What is the capital of France?"], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| print(f"Testing 'solver' node...") | |
| nodes.solver_node(test_state) | |
| # Solver appends to task_progress | |
| self.assertGreaterEqual(len(test_state["task_progress"]), 2) | |
| print(f"State after node execution: {test_state}") | |
| def test_researcher_node(self): | |
| """ | |
| Test the researcher node functionality. | |
| Retrieves and synthesizes information from various sources to answer knowledge-based questions | |
| """ | |
| nodes = Nodes() | |
| test_state: State = { | |
| "initial_query": "What are the latest news headlines about artificial intelligence published this week?", | |
| "messages": ["What are the latest news headlines about artificial intelligence published this week?"], | |
| "task_progress": ["What are the latest news headlines about artificial intelligence published this week?"], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| print(f"Testing 'researcher' node...") | |
| nodes.researcher_node(test_state) | |
| self.assertGreaterEqual(len(test_state["task_progress"]), 2) | |
| print(f"State after node execution: {test_state}") | |
| def test_reasoner_node(self): | |
| """ | |
| Test the reasoner node functionality. | |
| Performs logical reasoning, inference, and step-by-step problem-solving | |
| """ | |
| nodes = Nodes() | |
| test_state: State = { | |
| "initial_query": "What is the capital of France?", | |
| "messages": ["What is the capital of France?"], | |
| "task_progress": ["Reason: What is the capital of France?"], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| print(f"Testing 'reasoner' node...") | |
| nodes.reasoner_node(test_state) | |
| self.assertGreaterEqual(len(test_state["task_progress"]), 2) | |
| print(f"State after node execution: {test_state}") | |
| def test_viewer_node(self): | |
| """ | |
| Test the viewer node functionality. | |
| Processes, analyzes, and generates vision related information | |
| """ | |
| nodes = Nodes() | |
| test_state: State = { | |
| "initial_query": "Describe the image.", | |
| "messages": ["Describe the image."], | |
| "task_progress": ["View: Describe the image."], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| print(f"Testing 'image_handler' node...") | |
| nodes.viewer_node(test_state) | |
| self.assertGreaterEqual(len(test_state["task_progress"]), 2) | |
| print(f"State after node execution: {test_state}") | |
| def test_manager_edge(self): | |
| """ | |
| Test the conditional edge for manager node. | |
| This edge should return one of: "solver", "auditor", "final_answer" | |
| """ | |
| edges = Edges() | |
| # Test for final_answer by FINAL ANSWER in last message | |
| test_state: State = { | |
| "initial_query": "Q", | |
| "messages": ["Q", "FINAL ANSWER: Paris"], | |
| "task_progress": [], | |
| "audit_interval": 2, | |
| "manager_queries": 2, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| print(f"Testing 'manager' conditional edge...") | |
| result = edges.manager_edge(test_state) | |
| self.assertEqual(result, "final_answer") | |
| # Test for final_answer by max_interactions | |
| test_state2: State = { | |
| "initial_query": "Q", | |
| "messages": ["Q", "Some message"], | |
| "task_progress": [], | |
| "audit_interval": 2, | |
| "manager_queries": 4, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| result2 = edges.manager_edge(test_state2) | |
| self.assertEqual(result2, "final_answer") | |
| # Test for auditor | |
| test_state3: State = { | |
| "initial_query": "Q", | |
| "messages": ["Q", "Some message"], | |
| "task_progress": [], | |
| "audit_interval": 2, | |
| "manager_queries": 2, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| result3 = edges.manager_edge(test_state3) | |
| self.assertEqual(result3, "auditor") | |
| # Test for solver | |
| test_state4: State = { | |
| "initial_query": "Q", | |
| "messages": ["Q", "Some message"], | |
| "task_progress": [], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| result4 = edges.manager_edge(test_state4) | |
| self.assertEqual(result4, "solver") | |
| print(f"Edge decision: {result4}") | |
| def test_solver_edge(self): | |
| """ | |
| Test the conditional edge for solver node. | |
| This edge should return one of: "manager", "researcher", "reasoner", "viewer" | |
| """ | |
| edges = Edges() | |
| # researcher | |
| test_state: State = { | |
| "initial_query": "Q", | |
| "messages": ["Q"], | |
| "task_progress": ["to: researcher"], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| result = edges.solver_edge(test_state) | |
| self.assertEqual(result, "researcher") | |
| # reasoner | |
| test_state2: State = { | |
| "initial_query": "Q", | |
| "messages": ["Q"], | |
| "task_progress": ["to: reasoner"], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| result2 = edges.solver_edge(test_state2) | |
| self.assertEqual(result2, "reasoner") | |
| # viewer | |
| test_state3: State = { | |
| "initial_query": "Q", | |
| "messages": ["Q"], | |
| "task_progress": ["to: viewer"], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| result3 = edges.solver_edge(test_state3) | |
| self.assertEqual(result3, "viewer") | |
| # manager | |
| test_state4: State = { | |
| "initial_query": "Q", | |
| "messages": ["Q"], | |
| "task_progress": ["to: manager"], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| result4 = edges.solver_edge(test_state4) | |
| self.assertEqual(result4, "manager") | |
| # unspecified (should append instruction and return manager) | |
| test_state5: State = { | |
| "initial_query": "Q", | |
| "messages": ["Q"], | |
| "task_progress": ["no receiver"], | |
| "audit_interval": 2, | |
| "manager_queries": 1, | |
| "solver_queries": 0, | |
| "max_interactions": 4, | |
| "max_solving_effort": 4, | |
| "final_response": None | |
| } | |
| result5 = edges.solver_edge(test_state5) | |
| self.assertEqual(result5, "manager") | |
| print(f"Edge decision: {result5}") | |
| def test_full_workflow(self): | |
| """ | |
| Test the Alfred agent full workflow. | |
| """ | |
| # This is a placeholder for a full workflow test. | |
| # For a real test, you would simulate the entire agent graph. | |
| print("Testing Alfred complete workflow...") | |
| # Example test (pseudo, as actual invoke may require more setup) | |
| # result = self.graph.invoke({"input": "Test input"}) | |
| # self.assertIsNotNone(result) | |
| # print(f"Workflow result: {result}") | |
| if __name__ == "__main__": | |
| # test = TestAlfredAgent() | |
| # test.test_researcher_node() | |
| unittest.main() |