Merge pull request #31 from small-thinking/update-agent

Update agent

polymind/core/agent.py

@@ -20,25 +20,27 @@ class ThoughtProcess(BaseModel, ABC):
     def __str__(self):
         return self.thought_process_name
 
-    async def __call__(self, input: Message, agent: "Agent") -> Message:
+    async def __call__(self, agent: "Agent", input: Message) -> Message:
         """Makes the instance callable, delegating to the execute method.
         This allows the instance to be used as a callable object, simplifying
         the syntax for executing the thought process.
 
         Args:
+            agent (Agent): The agent who is executing the thought process.
             input (Message): The input message to the thought process.
 
         Returns:
             Message: The output message from the thought process.
         """
-        return await self._execute(input)
+        return await self._execute(agent=agent, input=input)
 
     @abstractmethod
-    async def _execute(self, input: Message) -> Message:
+    async def _execute(self, agent: "Agent", input: Message) -> Message:
         """Execute the thought process and return the result.
         The derived class must implement this method to define the behavior of the thought process.
 
         Args:
+            agent (Agent): The agent who is executing the thought process.
             input (Message): The input to the thought process carried in a message.
 
         Returns:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "polymind"
-version = "0.0.18"  # Update this version before publishing to PyPI
+version = "0.0.19"  # Update this version before publishing to PyPI
 description = "PolyMind is a customizable collaborative multi-agent framework for collective intelligence and distributed problem solving."
 authors = ["TechTao"]
 license = "MIT License"

tests/polymind/core/test_agent.py

@@ -10,7 +10,7 @@
 
 
 class MockThoughtProcess(ThoughtProcess):
-    async def _execute(self, input: Message) -> Message:
+    async def _execute(self, agent: Agent, input: Message) -> Message:
         # Implement a simple test logic, for example, just echo back the input with some modification
         modified_content = {"processed": True, **input.content}
         return Message(content=modified_content)
@@ -30,7 +30,7 @@ async def test_process_simple_message(self):
         input_message = Message(content={"hello": "world"})
 
         # Now, pass both the input_message and agent to the thought_process call
-        output_message = await thought_process(input_message, agent)
+        output_message = await thought_process(agent=agent, input=input_message)
 
         # Assertions to verify the behavior
         assert output_message.content.get("processed") == True