bugfixing; fixed toolbox; isolated [Base|AI|Human]Message crap logic to the agent interface; implemented tests

alfred.py

@@ -1,5 +1,3 @@
-from langchain_core.messages import BaseMessage
-
 from typing import Any, Dict
 
 from args import Args
@@ -35,7 +33,7 @@ class Alfred:
         """
         initial_state: State = {
             "initial_query": query,
-            "messages": [BaseMessage(query)],  # Manager's context
+            "messages": [query],  # Manager's context
             "task_progress": [],  # Solver's context
             "audit_interval": Args.AlfredParams.AUDIT_INTERVAL,
             "manager_queries": 0,

args.py

@@ -4,6 +4,9 @@ from typing import Optional
 from logger import Logger
 
 
+TEST_MODE = False
+
+
 class LLMInterface(Enum):
     OPENAI = "OpenAI"
     HUGGINGFACE = "HuggingFace"
@@ -80,11 +83,11 @@ class Args:
     primary_llm_interface=LLMInterface.OPENAI
     # secondary_llm_interface=LLMInterface.HUGGINGFACE
     vlm_interface=LLMInterface.OPENAI
-    primary_model="qwen2.5-qwq-35b-eureka-cubed-abliterated-uncensored"
-    secondary_model="qwen2.5-7b-instruct-1m"
-    vision_model="qwen/qwen2.5-vl-7b"
+    primary_model="groot" if TEST_MODE else "qwen2.5-qwq-35b-eureka-cubed-abliterated-uncensored"
+    secondary_model="groot" if TEST_MODE else "qwen2.5-7b-instruct-1m"
+    vision_model="groot" if TEST_MODE else "qwen/qwen2.5-vl-7b"
     api_base="http://127.0.0.1:1234/v1"  # LM Studio local endpoint
-    api_key=None
+    api_key="api_key"
     token = ""  # Not needed when using OpenAILike API
     # Agent presets
     PRIMARY_AGENT_PRESET = AgentPreset(

graph.py

@@ -1,13 +1,5 @@
-from langchain_core.messages import AnyMessage, BaseMessage, AIMessage, HumanMessage
-from langgraph.graph import START, END, StateGraph
-from langgraph.graph.message import add_messages
-from langgraph.prebuilt import ToolNode, tools_condition
+from typing import List, Literal, Optional, TypedDict
 
-from typing import Annotated, Any, Dict, List, Literal, Optional, TypedDict
-import logging
-from pathlib import Path
-
-from args import Args
 from agents import *
 from itf_agent import IAgent
 
@@ -15,10 +7,8 @@ from itf_agent import IAgent
 class State(TypedDict):
     """State class for the agent graph."""
     initial_query: str
-    # messages: List[Dict[str, Any]]
-    # messages: Annotated[list[BaseMessage], add_messages]
-    messages: List[BaseMessage]  # Manager's context
-    task_progress: List[BaseMessage]  # Solver's context
+    messages: List[str]  # Manager's context
+    task_progress: List[str]  # Solver's context
     audit_interval: int
     manager_queries: int
     solver_queries: int
@@ -38,7 +28,7 @@ class Agents:
     viewer = Viewer()
 
     @classmethod
-    def guard_output(cls, agent: IAgent, messages: List[BaseMessage]) -> BaseMessage:
+    def guard_output(cls, agent: IAgent, messages: List[str]) -> str:
         response = agent.query(messages)
         guarded_response = cls.guardian.query([response])
         return guarded_response
@@ -65,8 +55,8 @@ class _Helper:
         return first % second == 0
 
     @staticmethod
-    def solver_handler(task_progress: List[BaseMessage]) -> Literal["manager", "researcher", "reasoner", "viewer", "unspecified"]:
-        response = str(task_progress[-1].content)
+    def solver_successor(task_progress: List[str]) -> Literal["manager", "researcher", "reasoner", "viewer", "unspecified"]:
+        response = str(task_progress[-1])
         if "to: researcher" in response.lower():
             return "researcher"
         elif "to: reasoner" in response.lower():
@@ -78,6 +68,18 @@ class _Helper:
         else:
             return "unspecified"
 
+    @staticmethod
+    def manager_successor(state: State) -> Literal["solver", "auditor", "final_answer"]:
+        last_message = state["messages"][-1]
+        answer_ready = "FINAL ANSWER:" in last_message
+        max_interractions_reached = state["manager_queries"] >= state["max_interactions"]
+        if answer_ready or max_interractions_reached:
+            return "final_answer"
+
+        if _Helper._is_divisible(state["manager_queries"], state["audit_interval"]):
+            return "auditor"
+
+        return "solver"
 
 class Nodes:
     """
@@ -88,10 +90,13 @@ class Nodes:
         Orchestrates the workflow by delegating tasks to specialized nodes and integrating their outputs
         """
         state["manager_queries"] += 1
-        if not _Helper._is_divisible(state["manager_queries"], state["audit_interval"]):
+        successor = _Helper.manager_successor(state)
+        if successor == "solver":
             response = Agents.guard_output(Agents.manager, state["messages"])
             state["messages"].append(response)
-        # else: wait for auditor's feedback !
+            # Prepare task for Solver
+            state["task_progress"] = [response]
+        # else: [wait for auditor's feedback] or [is final answer]
 
         return state
 
@@ -99,10 +104,10 @@ class Nodes:
         """
         Formats and delivers the final response to the user
         """
-        instruction = BaseMessage("Formulate a definitive final answer in english. Be very concise and use no redundant words !")
+        instruction = "Formulate a definitive final answer in english. Be very concise and use no redundant words !"
         state["messages"].append(instruction)
         response = Agents.manager.query(state["messages"])
-        state["final_response"] = str(response.content)
+        state["final_response"] = response
         return state
 
     def auditor_node(self, state: State) -> State:
@@ -119,6 +124,16 @@ class Nodes:
         """
         response = Agents.guard_output(Agents.solver, state["task_progress"])
         state["task_progress"].append(response)
+
+        successor = _Helper.solver_successor(state["task_progress"])
+        if successor == "unspecified":
+            instruction = "Formulate an answer for the manager with your findings so far !"
+            state["task_progress"].append(instruction)
+            response = Agents.solver.query(state["task_progress"])
+            state["messages"].append(response)
+        elif successor == "manager":
+            state["messages"].append(response)
+
         return state
 
     def researcher_node(self, state: State) -> State:
@@ -155,35 +170,16 @@ class Edges:
         Conditional edge for manager node.
         Returns one of: "solver", "auditor", "final_answer"
         """
-        last_message = state["messages"][-1]
-        answer_ready = "FINAL ANSWER:" in str(last_message.content)
-        max_interractions_reached = state["manager_queries"] >= state["max_interactions"]
-        if answer_ready or max_interractions_reached:
-            return "final_answer"
-
-        if _Helper._is_divisible(state["manager_queries"], state["audit_interval"]):
-            return "auditor"
-
-        # Prepare task for Solver
-        state["task_progress"] = [last_message]
-        return "solver"
+        return _Helper.manager_successor(state)
 
     def solver_edge(self, state: State) -> Literal["manager", "researcher", "reasoner", "viewer"]:
         """
         Conditional edge for solver node.
         Returns one of: "manager", "researcher", "reasoner", "viewer"
         """
-        receiver = _Helper.solver_handler(state["task_progress"])
+        receiver = _Helper.solver_successor(state["task_progress"])
 
         if receiver == "unspecified":
-            instruction = BaseMessage("Formulate an answer for the manager with your findings so far !")
-            state["task_progress"].append(instruction)
-            response = Agents.solver.query(state["task_progress"])
-            state["messages"].append(response)
             return "manager"
 
-        if receiver == "manager":
-            response = state["task_progress"][-1]
-            state["messages"].append(response)
-
         return receiver

graph_builder.py

@@ -23,16 +23,14 @@ class GraphBuilder:
         graph.add_node("solver", self.nodes.solver_node)
         graph.add_node("researcher", self.nodes.researcher_node)
         graph.add_node("reasoner", self.nodes.reasoner_node)
-        graph.add_node("image_handler", self.nodes.image_handler_node)
-        graph.add_node("video_handler", self.nodes.video_handler_node)
+        graph.add_node("viewer", self.nodes.viewer_node)
         
         graph.add_edge(START, "manager")
         graph.add_edge("final_answer", END)
         graph.add_edge("auditor", "manager")
         graph.add_edge("researcher", "solver")
         graph.add_edge("reasoner", "solver")
-        graph.add_edge("image_handler", "solver")
-        graph.add_edge("video_handler", "solver")
+        graph.add_edge("viewer", "solver")
 
         graph.add_conditional_edges(
             "manager",
@@ -45,7 +43,7 @@ class GraphBuilder:
             "solver",
             self.edges.solver_edge,
             {
-                "manager": "manager", "researcher": "researcher", "reasoner": "reasoner", "image_handler": "image_handler", "video_handler": "video_handler"
+                "manager": "manager", "researcher": "researcher", "reasoner": "reasoner", "viewer": "viewer"
             }
         )
         

itf_agent.py

@@ -1,4 +1,4 @@
-from langchain_core.messages import BaseMessage, SystemMessage
+from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage
 
 import logging
 import os
@@ -13,6 +13,7 @@ class IAgent():
     def __init__(self, sys_prompt_filename, agent_preset: AgentPreset, tools: List = [], parallel_tool_calls=False):
         self.name = self._format_name(sys_prompt_filename)
         self.interface = agent_preset.get_interface()
+        self.mock = (agent_preset.get_model_name() == "groot")
 
         # Load the system prompt from a file
         system_prompt_path = os.path.join(os.getcwd(), "system_prompts", sys_prompt_filename)
@@ -37,6 +38,34 @@ class IAgent():
         cleaned_name = re.sub(r'^[^a-zA-Z]+', '', name_without_ext)
         return cleaned_name
     
+    @staticmethod
+    def _bake_roles(messages: List[str]) -> List[AnyMessage]:
+        """
+        Assigns roles to messages in reverse order: last message is HumanMessage,
+        previous is AIMessage, and so on, alternating backwards.
+
+        Args:
+            messages (List[str]): List of message strings.
+
+        Returns:
+            List[AnyMessage]: List of messages wrapped with appropriate role classes.
+
+        Raises:
+            ValueError: If messages is empty.
+        """
+        if not messages:
+            raise ValueError("The list of messages cannot be empty !")
+        messages_with_roles = []
+        total_messages = len(messages)
+        for idx, msg in enumerate(messages):
+            # Assign roles in reverse: last is Human, previous is AI, etc.
+            reverse_idx = total_messages - idx - 1
+            if reverse_idx % 2 == 0:
+                messages_with_roles.append(HumanMessage(content=msg))
+            else:
+                messages_with_roles.append(AIMessage(content=msg))
+        return messages_with_roles
+
     def get_system_prompt(self) -> str:
         """
         Retrieves the system prompt.
@@ -46,7 +75,7 @@ class IAgent():
         """
         return self.system_prompt
 
-    def query(self, messages: List[BaseMessage]) -> BaseMessage:
+    def query(self, messages: List[str]) -> str:
         """
         Asynchronously queries the agent with a given question and returns the response.
 
@@ -56,15 +85,22 @@ class IAgent():
         Returns:
             str: The response from the agent as a string.
         """
+
         if Args.LOGGER is None:
             raise RuntimeError("LOGGER must be defined before querying the agent.")
 
         separator = "=============================="
         Args.LOGGER.log(logging.INFO, f"\n{separator}\nAgent '{self.name}' has been queried !\nINPUT:\n{messages}\n")
 
+        if self.mock:
+            response = str("I am GROOT !")
+            Args.LOGGER.log(logging.INFO, f"\nAgent '{self.name}' produced OUTPUT:\n{response}\n{separator}\n")
+            return response
+
         system_prompt = self.get_system_prompt()
-        conversation = [SystemMessage(content=system_prompt)] + messages
-        response = self.model.invoke(conversation)
+        messages_with_roles = self._bake_roles(messages)
+        conversation = [SystemMessage(content=system_prompt)] + messages_with_roles
+        response = str(self.model.invoke(conversation).content)
 
         Args.LOGGER.log(logging.INFO, f"\nAgent '{self.name}' produced OUTPUT:\n{response}\n{separator}\n")
         return response

llm_factory.py

@@ -27,12 +27,12 @@ class LLMFactory():
         repeat_penalty = agent_preset.get_repeat_penalty()
 
         kwargs = {
+            "name": model_name,
             "model": model_name,
             "base_url": Args.api_base,
             "api_key": Args.api_key,
             "temperature": temperature,
             "max_completion_tokens": max_tokens,
-            # "presence_penalty": repeat_penalty,
             "frequency_penalty": repeat_penalty
         }
 
@@ -48,6 +48,7 @@ class LLMFactory():
         repeat_penalty = agent_preset.get_repeat_penalty()
 
         kwargs = {
+            "name": model_name,
             "model": model_name,
             "temperature": temperature,
             "max_new_tokens": max_tokens,

test.py

@@ -1,3 +1,4 @@
+from args import Args
 from graph import State, Nodes, Edges
 from graph_builder import GraphBuilder
 
@@ -20,17 +21,52 @@ class TestAlfredAgent(unittest.TestCase):
         
         Orchestrates the workflow by delegating tasks to specialized nodes and integrating their outputs
         """
-        # Create an instance of Nodes class
         nodes = Nodes()
         
         # Create a test state
-        test_state = {}  # TODO: Initialize with appropriate test data
+        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)
-        
-        # TODO: Add assertions to verify the state changes
+
+        # 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):
@@ -39,17 +75,25 @@ class TestAlfredAgent(unittest.TestCase):
         
         Formats and delivers the final response to the user
         """
-        # Create an instance of Nodes class
         nodes = Nodes()
-        
-        # Create a test state
-        test_state = {}  # TODO: Initialize with appropriate test data
-        
-        # Test the node function
+        # 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)
-        
-        # TODO: Add assertions to verify the state changes
+        # 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):
@@ -58,17 +102,22 @@ class TestAlfredAgent(unittest.TestCase):
         
         Reviews manager's outputs for accuracy, safety, and quality
         """
-        # Create an instance of Nodes class
         nodes = Nodes()
-        
-        # Create a test state
-        test_state = {}  # TODO: Initialize with appropriate test data
-        
-        # Test the node function
+        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)
-        
-        # TODO: Add assertions to verify the state changes
+        # Auditor appends a message
+        self.assertGreaterEqual(len(test_state["messages"]), 3)
         print(f"State after node execution: {test_state}")
 
     def test_solver_node(self):
@@ -77,17 +126,22 @@ class TestAlfredAgent(unittest.TestCase):
         
         Central problem-solving node that coordinates with specialized experts based on task requirements
         """
-        # Create an instance of Nodes class
         nodes = Nodes()
-        
-        # Create a test state
-        test_state = {}  # TODO: Initialize with appropriate test data
-        
-        # Test the node function
+        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)
-        
-        # TODO: Add assertions to verify the state changes
+        # 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):
@@ -96,17 +150,21 @@ class TestAlfredAgent(unittest.TestCase):
         
         Retrieves and synthesizes information from various sources to answer knowledge-based questions
         """
-        # Create an instance of Nodes class
         nodes = Nodes()
-        
-        # Create a test state
-        test_state = {}  # TODO: Initialize with appropriate test data
-        
-        # Test the node function
+        test_state: State = {
+            "initial_query": "What is the capital of France?",
+            "messages": ["What is the capital of France?"],
+            "task_progress": ["Research: 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 'researcher' node...")
         nodes.researcher_node(test_state)
-        
-        # TODO: Add assertions to verify the state changes
+        self.assertGreaterEqual(len(test_state["task_progress"]), 2)
         print(f"State after node execution: {test_state}")
 
     def test_reasoner_node(self):
@@ -115,17 +173,21 @@ class TestAlfredAgent(unittest.TestCase):
         
         Performs logical reasoning, inference, and step-by-step problem-solving
         """
-        # Create an instance of Nodes class
         nodes = Nodes()
-        
-        # Create a test state
-        test_state = {}  # TODO: Initialize with appropriate test data
-        
-        # Test the node function
+        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)
-        
-        # TODO: Add assertions to verify the state changes
+        self.assertGreaterEqual(len(test_state["task_progress"]), 2)
         print(f"State after node execution: {test_state}")
 
     def test_viewer_node(self):
@@ -134,17 +196,21 @@ class TestAlfredAgent(unittest.TestCase):
         
         Processes, analyzes, and generates vision related information
         """
-        # Create an instance of Nodes class
         nodes = Nodes()
-        
-        # Create a test state
-        test_state = {}  # TODO: Initialize with appropriate test data
-        
-        # Test the node function
+        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)
-        
-        # TODO: Add assertions to verify the state changes
+        self.assertGreaterEqual(len(test_state["task_progress"]), 2)
         print(f"State after node execution: {test_state}")
 
     def test_manager_edge(self):
@@ -153,48 +219,160 @@ class TestAlfredAgent(unittest.TestCase):
         
         This edge should return one of: "solver", "auditor", "final_answer"
         """
-        # Create an instance of Edges class
         edges = Edges()
-        
-        # Create a test state
-        test_state = {}  # TODO: Initialize with appropriate test data
-        
-        # Test the edge function
+        # 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)
-        
-        # TODO: Add assertions to verify the result
-        print(f"Edge decision: {result}")
-        assert result in ["solver", "auditor", "final_answer"], f"Edge result '{result}' not in expected values"
+        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", "image_handler", "video_handler"
+        This edge should return one of: "manager", "researcher", "reasoner", "viewer"
         """
-        # Create an instance of Edges class
         edges = Edges()
-        
-        # Create a test state
-        test_state = {}  # TODO: Initialize with appropriate test data
-        
-        # Test the edge function
-        print(f"Testing 'solver' conditional edge...")
+        # 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)
-        
-        # TODO: Add assertions to verify the result
-        print(f"Edge decision: {result}")
-        assert result in ["manager", "researcher", "reasoner", "image_handler", "video_handler"], f"Edge result '{result}' not in expected values"
+        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.
         """
-        # TODO: Add test code here
+        # 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
+        # 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}")

toolbox.py

@@ -1,3 +1,5 @@
+from langchain_core.tools.simple import Tool
+
 from duckduckgo_search import DDGS
 import pint
 import sympy as sp
@@ -9,10 +11,10 @@ class _Math:
     def symbolic_calc(expression: str) -> str:
         """
         Evaluates complex mathematical expressions using SymPy.
-        
+
         Args:
-            expression: Mathematical expression as string
-        
+            expression (str): Mathematical expression as string
+
         Returns:
             Result of the calculation
         """
@@ -28,12 +30,12 @@ class _Math:
     def unit_converter(value: float, from_unit: str, to_unit: str) -> str:
         """
         Converts values between different units of measurement.
-        
+
         Args:
-            value: The numerical value to convert
-            from_unit: The source unit (e.g., 'meter', 'kg', 'celsius')
-            to_unit: The target unit (e.g., 'feet', 'pound', 'fahrenheit')
-        
+            value (float): The numerical value to convert
+            from_unit (str): The source unit (e.g., 'meter', 'kg', 'celsius')
+            to_unit (str): The target unit (e.g., 'feet', 'pound', 'fahrenheit')
+
         Returns:
             The converted value with appropriate units
         """
@@ -55,14 +57,27 @@ class _Math:
             return f"Error in unit conversion: {str(e)}"
 
 
+class MathToolbox:
+    symbolic_calc = Tool(
+        name="symbolic_calc",
+        func=_Math.symbolic_calc,
+        description=_Math.symbolic_calc. __doc__ or ""
+    )
+    unit_converter = Tool(
+        name="unit_converter",
+        func=_Math.unit_converter,
+        description=_Math.unit_converter. __doc__ or ""
+    )
+
+
 class _WebSearch:
     @staticmethod
     def duckduckgo_text_search(keywords, max_results=5) -> list[dict[str, str]]:
         """DuckDuckGo text search.
 
         Args:
-            keywords: keywords for query.
-            max_results: max number of results. If None, returns results only from the first response. Defaults to 5.
+            keywords (str): keywords for query.
+            max_results (int): max number of results. If None, returns results only from the first response. Defaults to 5.
 
         Returns:
             List of dictionaries with search results, or None if there was an error.
@@ -74,18 +89,17 @@ class _WebSearch:
         """
         return DDGS().text(keywords, max_results=max_results)
 
-
     @staticmethod
     def duckduckgo_images_search(keywords, license = None, max_results=5) -> list[dict[str, str]]:
         """DuckDuckGo images search.
 
         Args:
-            keywords: keywords for query.
-            license: any (All Creative Commons), Public (PublicDomain),
+            keywords (str): keywords for query.
+            license (str|None): any (All Creative Commons), Public (PublicDomain),
                 Share (Free to Share and Use), ShareCommercially (Free to Share and Use Commercially),
                 Modify (Free to Modify, Share, and Use), ModifyCommercially (Free to Modify, Share, and
                 Use Commercially). Defaults to None.
-            max_results: max number of results. If None, returns results only from the first response. Defaults to 5.
+            max_results (int): max number of results. If None, returns results only from the first response. Defaults to 5.
 
         Returns:
             List of dictionaries with images search results.
@@ -97,15 +111,14 @@ class _WebSearch:
         """
         return DDGS().images(keywords, license_image=license, max_results=max_results)
 
-
     @staticmethod
     def duckduckgo_videos_search(keywords, license = None, max_results=5) -> list[dict[str, str]]:
         """DuckDuckGo videos search.
 
         Args:
-            keywords: keywords for query.
-            license: creativeCommon, youtube. Defaults to None.
-            max_results: max number of results. If None, returns results only from the first response. Defaults to 5.
+            keywords (str): keywords for query.
+            license (str|None): creativeCommon, youtube. Defaults to None.
+            max_results (int): max number of results. If None, returns results only from the first response. Defaults to 5.
 
         Returns:
             List of dictionaries with videos search results.
@@ -118,16 +131,34 @@ class _WebSearch:
         return DDGS().videos(keywords, license_videos=license, max_results=max_results)
 
 
+class WebSearchToolbox:
+    duckduckgo_text_search = Tool(
+        name="duckduckgo_text_search",
+        func=_WebSearch.duckduckgo_text_search,
+        description=_WebSearch.duckduckgo_text_search. __doc__ or ""
+    )
+    duckduckgo_images_search = Tool(
+        name="duckduckgo_images_search",
+        func=_WebSearch.duckduckgo_images_search,
+        description=_WebSearch.duckduckgo_images_search. __doc__ or ""
+    )
+    duckduckgo_videos_search = Tool(
+        name="duckduckgo_videos_search",
+        func=_WebSearch.duckduckgo_videos_search,
+        description=_WebSearch.duckduckgo_videos_search. __doc__ or ""
+    )
+
+
 class _Encryption:
-    
+
     @staticmethod
     def ascii_encode(text: str) -> str:
         """
         Converts each character in a string to its ASCII value.
-        
+
         Args:
-            text: The text to encode
-            
+            text (str): The text to encode
+
         Returns:
             Space-separated ASCII values
         """
@@ -138,15 +169,15 @@ class _Encryption:
             return result
         except Exception as e:
             return f"Error in ASCII encoding: {str(e)}"
-    
+
     @staticmethod
     def ascii_decode(text: str) -> str:
         """
         Converts space-separated ASCII values back to characters.
-        
+
         Args:
-            text: Space-separated ASCII values
-            
+            text (str): Space-separated ASCII values
+
         Returns:
             Decoded string
         """
@@ -162,10 +193,10 @@ class _Encryption:
     def base64_encode(text: str) -> str:
         """
         Encodes a string to base64.
-        
+
         Args:
-            text: The text to encode
-            
+            text (str): The text to encode
+
         Returns:
             Base64 encoded string
         """
@@ -177,15 +208,15 @@ class _Encryption:
             return encoded_text
         except Exception as e:
             return f"Error in base64 encoding: {str(e)}"
-    
+
     @staticmethod
     def base64_decode(encoded_text: str) -> str:
         """
         Decodes a base64 string to plain text.
-        
+
         Args:
-            encoded_text: The base64 encoded text
-            
+            encoded_text (str): The base64 encoded text
+
         Returns:
             Decoded string
         """
@@ -197,16 +228,16 @@ class _Encryption:
             return decoded_text
         except Exception as e:
             return f"Error in base64 decoding: {str(e)}"
-    
+
     @staticmethod
     def caesar_cipher_encode(text: str, shift: int) -> str:
         """
         Encodes text using Caesar cipher with specified shift.
-        
+
         Args:
-            text: The text to encode
-            shift: Number of positions to shift each character
-            
+            text (str): The text to encode
+            shift (int): Number of positions to shift each character
+
         Returns:
             Caesar cipher encoded string
         """
@@ -223,42 +254,42 @@ class _Encryption:
             return result
         except Exception as e:
             return f"Error in Caesar cipher encoding: {str(e)}"
-    
+
     @classmethod
     def caesar_cipher_decode(cls, encoded_text: str, shift: int) -> str:
         """
         Decodes Caesar cipher text with specified shift.
-        
+
         Args:
-            encoded_text: The encoded text
-            shift: Number of positions the text was shifted
-            
+            encoded_text (str): The encoded text
+            shift (int): Number of positions the text was shifted
+
         Returns:
             Decoded string
         """
         print(f"-> caesar_cipher_decode tool used (input: {encoded_text[:30]}..., shift: {shift}) !")
         # To decode, we shift in the opposite direction
         return cls.caesar_cipher_encode(encoded_text, -shift)
-    
+
     @classmethod
     def caesar_cipher_brute_force(cls, text: str) -> str:
         """
         Performs a brute force attack on a Caesar cipher by trying all 26 shifts.
-        
+
         Args:
-            text: The Caesar cipher encoded text
-            
+            text (str): The Caesar cipher encoded text
+
         Returns:
             All possible decoding results with their respective shifts
         """
         print(f"-> caesar_cipher_brute_force tool used (input: {text[:30]}...) !")
         results = []
-        
+
         # Try all 26 possible shifts for English alphabet
         for shift in range(26):
             decoded = cls.caesar_cipher_decode(text, shift)
             results.append(f"Shift {shift}: {decoded}")
-        
+
         output = "\n".join(results)
         return output
 
@@ -266,10 +297,10 @@ class _Encryption:
     def reverse_string(text: str) -> str:
         """
         Reverses a string.
-        
+
         Args:
-            text: The text to reverse
-            
+            text (str): The text to reverse
+
         Returns:
             Reversed string
         """
@@ -278,7 +309,50 @@ class _Encryption:
         return reversed_text
 
 
+class EncryptionToolbox:
+    ascii_encode = Tool(
+        name="ascii_encode",
+        func=_Encryption.ascii_encode,
+        description=_Encryption.ascii_encode. __doc__ or ""
+    )
+    ascii_decode = Tool(
+        name="ascii_decode",
+        func=_Encryption.ascii_decode,
+        description=_Encryption.ascii_decode. __doc__ or ""
+    )
+    base64_encode = Tool(
+        name="base64_encode",
+        func=_Encryption.base64_encode,
+        description=_Encryption.base64_encode. __doc__ or ""
+    )
+    base64_decode = Tool(
+        name="base64_decode",
+        func=_Encryption.base64_decode,
+        description=_Encryption.base64_decode. __doc__ or ""
+    )
+    caesar_cipher_encode = Tool(
+        name="caesar_cipher_encode",
+        func=_Encryption.caesar_cipher_encode,
+        description=_Encryption.caesar_cipher_encode. __doc__ or ""
+    )
+    caesar_cipher_decode = Tool(
+        name="caesar_cipher_decode",
+        func=_Encryption.caesar_cipher_decode,
+        description=_Encryption.caesar_cipher_decode. __doc__ or ""
+    )
+    caesar_cipher_brute_force = Tool(
+        name="caesar_cipher_brute_force",
+        func=_Encryption.caesar_cipher_brute_force,
+        description=_Encryption.caesar_cipher_brute_force. __doc__ or ""
+    )
+    reverse_string = Tool(
+        name="reverse_string",
+        func=_Encryption.reverse_string,
+        description=_Encryption.reverse_string. __doc__ or ""
+    )
+
+
 class Toolbox:
-    math = _Math()
-    web_search = _WebSearch()
-    encryption = _Encryption()
+    math = MathToolbox()
+    web_search = WebSearchToolbox()
+    encryption = EncryptionToolbox()