diff --git a/polymind/core/agent.py b/polymind/core/agent.py index 21d8c59..e6e8db5 100644 --- a/polymind/core/agent.py +++ b/polymind/core/agent.py @@ -1,64 +1,160 @@ +from abc import ABC, abstractmethod from typing import Dict from pydantic import BaseModel, Field from polymind.core.logger import Logger from polymind.core.message import Message -from polymind.core.tool import BaseTool, LLMTool +from polymind.core.tool import (BaseTool, LLMTool, OptimizableBaseTool, + SyncLLMTool) -class Agent(BaseModel): +class AbstractAgent(BaseModel, ABC): + """ + Abstract base class for all agent types. + + This class defines the common structure and interface for both synchronous + and asynchronous agents. It includes shared attributes and methods, as well + as abstract methods that must be implemented by subclasses. + """ agent_name: str - # Persona of the agent indicates the role of the agent. persona: str - tools: Dict[str, BaseTool] = Field(default=None, description="The tools that the agent can use.") - reasoner: LLMTool = Field(default=None, description="The reasoner that will be used in the thought process.") def __init__(self, **kwargs): super().__init__(**kwargs) self._logger = Logger(__file__) - def __str__(self): + def __str__(self) -> str: return self.agent_name def _input_preprocess(self, input: Message) -> None: - """Preprocess the input message before the agent starts working. - Now now the only thing to do is to add the persona to the input message. + """ + Preprocess the input message before the agent starts working. + + Args: + input (Message): The input message to preprocess. """ input.content["persona"] = self.persona - async def _execute(self, input: Message) -> Message: - """Execute the agent and return the result. - This method defines the behavior of the agent's thought process. + @abstractmethod + def _execute(self, input: Message) -> Message: + """ + Execute the agent and return the result. + + Args: + input (Message): The input message to process. + + Returns: + Message: The result of the agent's execution. + """ + pass + + @abstractmethod + def __call__(self, input: Message) -> Message: + """ + Enable the agent to start working. Args: - input (Message): The input to the thought process carried in a message. + input (Message): The input message to process. Returns: - Message: The result of the thought process carried in a message. + Message: The result of the agent's work. """ - if "requirement" in input.content: - self._logger.thought_process_log( - f"[{self.agent_name}], your requirement is: {input.content['requirement']}" - ) - else: + pass + + +class Agent(AbstractAgent): + """ + Synchronous agent implementation. + + This class represents a synchronous agent that uses OptimizableBaseTool + for its tools and SyncLLMTool for reasoning. + """ + + tools: Dict[str, OptimizableBaseTool] = Field(default=None, description="The tools that the agent can use.") + reasoner: SyncLLMTool = Field(default=None, description="The reasoner that will be used in the thought process.") + + def _execute(self, input: Message) -> Message: + """ + Synchronous execution of the agent. + + Args: + input (Message): The input message to process. + + Returns: + Message: The result of the agent's execution. + + Raises: + ValueError: If the input message doesn't contain the 'requirement' field. + """ + if "requirement" not in input.content: raise ValueError("The input message must contain the 'requirement' field.") + self._logger.thought_process_log(f"[{self.agent_name}], your requirement is: {input.content['requirement']}") + # Add logic for executing the thought process using tools and reasoner. # This is a placeholder implementation. result_content = {"output": f"Processed requirement: {input.content['requirement']}"} return Message(content=result_content) + def __call__(self, input: Message) -> Message: + """ + Synchronous call method. + + Args: + input (Message): The input message to process. + + Returns: + Message: The result of the agent's work. + """ + self._input_preprocess(input=input) + return self._execute(input=input) + + +class AsyncAgent(AbstractAgent): + """ + Asynchronous agent implementation. + + This class represents an asynchronous agent that uses BaseTool + for its tools and LLMTool for reasoning. + """ + + tools: Dict[str, BaseTool] = Field(default=None, description="The tools that the agent can use.") + reasoner: LLMTool = Field(default=None, description="The reasoner that will be used in the thought process.") + + async def _execute(self, input: Message) -> Message: + """ + Asynchronous execution of the agent. + + Args: + input (Message): The input message to process. + + Returns: + Message: The result of the agent's execution. + + Raises: + ValueError: If the input message doesn't contain the 'requirement' field. + """ + if "requirement" not in input.content: + raise ValueError("The input message must contain the 'requirement' field.") + + self._logger.thought_process_log(f"[{self.agent_name}], your requirement is: {input.content['requirement']}") + + # Add async logic for executing the thought process using tools and reasoner. + # This is a placeholder implementation. + result_content = {"output": f"Processed requirement: {input.content['requirement']}"} + return Message(content=result_content) + async def __call__(self, input: Message) -> Message: - """Enable the agent to start working. - The actual processing is driven by the agent itself. + """ + Asynchronous call method. Args: - input (Message): The input message to the agent. + input (Message): The input message to process. Returns: - Message: The output message from the agent. + Message: The result of the agent's work. """ self._input_preprocess(input=input) return await self._execute(input=input) diff --git a/polymind/core/tool.py b/polymind/core/tool.py index e3d67cf..25bc6f6 100644 --- a/polymind/core/tool.py +++ b/polymind/core/tool.py @@ -505,6 +505,129 @@ async def _execute(self, input: Message) -> Message: return response_message +class SyncLLMTool(OptimizableBaseTool): + """Synchronous LLM tool defines the basic properties of the language model tools. + This tool will get the prompt from "input" and return the response to "output". + """ + + llm_name: str = Field(..., description="The name of the model.") + max_tokens: int = Field(..., description="The maximum number of tokens for the chat.") + temperature: float = Field(default=1.0, description="The temperature for the chat.") + top_p: float = Field( + default=0.1, + description="The top p for the chat. Top p is used to prevent the model from generating unlikely words.", + ) + stop: str = Field(default=None, description="The stop sequence for the chat.") + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._logger = Logger(__file__) + self._set_client() + + def _set_client(self): + """Set the client for the language model.""" + # Implement the synchronous client setup here + pass + + def input_spec(self) -> List[Param]: + return [ + Param( + name="input", + type="str", + required=True, + description="The prompt for the chat.", + example="hello, how are you?", + ), + Param( + name="system_prompt", + type="str", + required=False, + example="You are a helpful AI assistant.", + description="The system prompt for the chat.", + ), + Param( + name="max_tokens", + type="int", + required=False, + example="1500", + description="The maximum number of tokens for the chat.", + ), + Param( + name="temperature", + type="float", + required=False, + example="0.7", + description="The temperature for the chat.", + ), + Param( + name="top_p", + type="float", + required=False, + example="0.1", + description="The top p for the chat.", + ), + ] + + def output_spec(self) -> List[Param]: + return [ + Param( + name="output", + type="str", + required=True, + description="The response from the chat.", + ), + ] + + def _invoke(self, input: Message) -> Message: + """Invoke the language model with the input message and return the response message. + + Args: + input (Message): The input message to the language model. The message should contain the below keys: + - prompt: The prompt for the chat. + - system_prompt: The system prompt for the chat. + - max_tokens: The maximum number of tokens for the chat. + - temperature: The temperature for the chat. + - top_p: The top p for the chat. + - stop: The stop sequence for the chat. + + Returns: + Message: The response message from the language model. The actual content is in the "output" field. + """ + # Implement the synchronous invocation of the language model here + # This is a placeholder implementation + return Message(content={"output": "Synchronous LLM response"}) + + def forward(self, **kwargs) -> Message: + """Execute the tool and return the result synchronously.""" + current_datetime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + + # Validate and prepare input + prompt = kwargs.get("input", "") + system_prompt = kwargs.get("system_prompt", "") + if not prompt: + raise ValueError("Prompt in the field 'input' cannot be empty.") + + input_message = Message( + content={ + "input": prompt, + "system_prompt": system_prompt, + "max_tokens": kwargs.get("max_tokens", self.max_tokens), + "temperature": kwargs.get("temperature", self.temperature), + "top_p": kwargs.get("top_p", self.top_p), + "datetime": current_datetime, + } + ) + + if self.stop: + input_message.content["stop"] = self.stop + + response_message = self._invoke(input_message) + if "output" not in response_message.content: + raise ValueError("The response message must contain the 'output' key.") + + return response_message + + class CombinedMeta(type(Module), type(BaseTool)): pass diff --git a/pyproject.toml b/pyproject.toml index adc96cd..9cb246c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "polymind" -version = "0.0.52" # Update this version before publishing to PyPI +version = "0.0.53" # 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" diff --git a/tests/polymind/core/test_agent.py b/tests/polymind/core/test_agent.py index 2ebfceb..b307d3f 100644 --- a/tests/polymind/core/test_agent.py +++ b/tests/polymind/core/test_agent.py @@ -7,10 +7,26 @@ import pytest -from polymind.core.agent import Agent +from polymind.core.agent import Agent, AsyncAgent from polymind.core.memory import LinearMemory from polymind.core.message import Message -from polymind.core.tool import BaseTool, LLMTool, RetrieveTool, ToolManager +from polymind.core.tool import LLMTool, OptimizableBaseTool, Param, SyncLLMTool + + +class MockOptimizableTool(OptimizableBaseTool): + tool_name: str = "mock_tool" + descriptions: List[str] = ["Mock tool for testing."] + + def input_spec(self) -> List[Param]: + return [Param(name="requirement", type="str", description="The requirement to process.")] + + def output_spec(self) -> List[Param]: + return [Param(name="output", type="str", description="The processed requirement.")] + + def _invoke(self, input: Message) -> Message: + # Mock response for testing + response_content = {"output": "Processed requirement: " + input.content.get("requirement")} + return Message(content=response_content) # Mock classes to simulate dependencies @@ -34,25 +50,36 @@ def _set_client(self): pass -class MockToolManager(ToolManager): - pass +class MockSyncLLMTool(SyncLLMTool): + tool_name: str = "mock_sync_tool" + llm_name: str = "mock_sync_llm" + max_tokens: int = 1500 + temperature: float = 0.7 + descriptions: List[str] = ["Mock Sync LLM tool for testing."] + def _invoke(self, input: Message) -> Message: + # Mock response for testing + response_content = { + "output": '{"steps": [{"objective": "mock sync objective", "input": null, "output": {"name": "result", "type": "str"}}]}' + } + return Message(content=response_content) -class MockRetrieveTool(RetrieveTool): - pass + def _set_client(self): + # Mock client setup + pass class MockMemory(LinearMemory): pass -class TestAgent: +class TestAsyncAgent: @pytest.mark.asyncio async def test_process_simple_message(self): # Create a minimal Agent instance for testing reasoner = MockLLMTool(llm_name="mock_llm", max_tokens=1500, temperature=0.7) - agent = Agent( + agent = AsyncAgent( agent_name="TestAgent", persona="Tester", tools={}, @@ -74,7 +101,7 @@ async def test_agent_with_invalid_input(self): # Create a minimal Agent instance for testing reasoner = MockLLMTool(llm_name="mock_llm", max_tokens=1500, temperature=0.7) - agent = Agent( + agent = AsyncAgent( agent_name="TestAgent", persona="Tester", tools={}, @@ -90,3 +117,97 @@ async def test_agent_with_invalid_input(self): await agent(input_message) assert "The input message must contain the 'requirement' field." in str(exc_info.value) + + +class TestAgent: + def test_process_simple_message(self): + # Create a minimal Agent instance for testing + reasoner = MockSyncLLMTool(llm_name="mock_sync_llm", max_tokens=1500, temperature=0.7) + + agent = Agent( + agent_name="TestSyncAgent", + persona="Sync Tester", + tools={}, + reasoner=reasoner, + memory=MockMemory(), + ) + + # Prepare the input message + input_message = Message(content={"requirement": "test sync requirement"}) + + # Now, pass the input_message to the agent + output_message = agent(input_message) + + # Assertions to verify the behavior + assert output_message.content.get("output") == "Processed requirement: test sync requirement" + + def test_agent_with_invalid_input(self): + # Create a minimal Agent instance for testing + reasoner = MockSyncLLMTool(llm_name="mock_sync_llm", max_tokens=1500, temperature=0.7) + + agent = Agent( + agent_name="TestSyncAgent", + persona="Sync Tester", + tools={}, + reasoner=reasoner, + memory=MockMemory(), + ) + + # Prepare the input message without the required 'requirement' field + input_message = Message(content={"hello": "sync world"}) + + # Attempt to process the message and expect a ValueError + with pytest.raises(ValueError) as exc_info: + agent(input_message) + + assert "The input message must contain the 'requirement' field." in str(exc_info.value) + + def test_agent_with_tools(self): + # Create a mock tool + mock_tool = MockOptimizableTool( + tool_name="mock_optimizable_tool", + descriptions=["A mock optimizable tool", "for testing.", "It has no real functionality."], + ) + + # Create a minimal Agent instance for testing with a tool + reasoner = MockSyncLLMTool(llm_name="mock_sync_llm", max_tokens=1500, temperature=0.7) + + agent = Agent( + agent_name="TestSyncAgent", + persona="Sync Tester", + tools={"mock_tool": mock_tool}, + reasoner=reasoner, + memory=MockMemory(), + ) + + # Prepare the input message + input_message = Message(content={"requirement": "test requirement using mock_tool"}) + + # Now, pass the input_message to the agent + output_message = agent(input_message) + + # Assertions to verify the behavior + assert output_message.content.get("output") == "Processed requirement: test requirement using mock_tool" + assert "mock_tool" in agent.tools + + def test_agent_persona(self): + # Create a minimal Agent instance for testing + reasoner = MockSyncLLMTool(llm_name="mock_sync_llm", max_tokens=1500, temperature=0.7) + + agent = Agent( + agent_name="TestSyncAgent", + persona="Helpful Sync Assistant", + tools={}, + reasoner=reasoner, + memory=MockMemory(), + ) + + # Prepare the input message + input_message = Message(content={"requirement": "test sync requirement"}) + + # Now, pass the input_message to the agent + agent(input_message) + + # Assertions to verify the persona is set correctly + assert agent.persona == "Helpful Sync Assistant" + assert input_message.content.get("persona") == "Helpful Sync Assistant"