Kye Gomez
7 months ago
parent
20a994b62f
commit
5a3fb81b77
@ -1,182 +0,0 @@
|
||||
from typing import List
|
||||
|
||||
from langchain_experimental.autonomous_agents import AutoGPT
|
||||
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.tools.base_tool import BaseTool
|
||||
from swarms.utils.decorators import error_decorator, timing_decorator
|
||||
|
||||
|
||||
class Worker(Agent):
|
||||
"""
|
||||
The Worker class represents an autonomous agent that can perform tassks through
|
||||
function calls or by running a chat.
|
||||
|
||||
Args:
|
||||
name (str, optional): Name of the agent. Defaults to "Autobot Swarm Worker".
|
||||
role (str, optional): Role of the agent. Defaults to "Worker in a swarm".
|
||||
external_tools (list, optional): List of external tools. Defaults to None.
|
||||
human_in_the_loop (bool, optional): Whether to include human in the loop. Defaults to False.
|
||||
temperature (float, optional): Temperature for the agent. Defaults to 0.5.
|
||||
llm ([type], optional): Language model. Defaults to None.
|
||||
openai_api_key (str, optional): OpenAI API key. Defaults to None.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If there is an error while setting up the agent.
|
||||
|
||||
Example:
|
||||
>>> worker = Worker(
|
||||
... name="My Worker",
|
||||
... role="Worker",
|
||||
... external_tools=[MyTool1(), MyTool2()],
|
||||
... human_in_the_loop=False,
|
||||
... temperature=0.5,
|
||||
... )
|
||||
>>> worker.run("What's the weather in Miami?")
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "WorkerAgent",
|
||||
role: str = "Worker in a swarm",
|
||||
external_tools=None,
|
||||
human_in_the_loop: bool = False,
|
||||
temperature: float = 0.5,
|
||||
llm=None,
|
||||
openai_api_key: str = None,
|
||||
tools: List[BaseTool] = None,
|
||||
embedding_size: int = 1536,
|
||||
search_kwargs: dict = {"k": 8},
|
||||
verbose: bool = False,
|
||||
memory: callable = None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
self.name = name
|
||||
self.role = role
|
||||
self.external_tools = external_tools
|
||||
self.human_in_the_loop = human_in_the_loop
|
||||
self.temperature = temperature
|
||||
self.llm = llm
|
||||
self.openai_api_key = openai_api_key
|
||||
self.tools = tools
|
||||
self.embedding_size = embedding_size
|
||||
self.search_kwargs = search_kwargs
|
||||
self.verbose = verbose
|
||||
self.memory = memory
|
||||
|
||||
self.setup_tools(external_tools)
|
||||
self.setup_memory()
|
||||
self.setup_agent()
|
||||
|
||||
def reset(self):
|
||||
"""
|
||||
Reset the message history.
|
||||
"""
|
||||
self.message_history = []
|
||||
|
||||
def receieve(self, name: str, message: str) -> None:
|
||||
"""
|
||||
Receive a message and update the message history.
|
||||
|
||||
Parameters:
|
||||
- `name` (str): The name of the sender.
|
||||
- `message` (str): The received message.
|
||||
"""
|
||||
self.message_history.append(f"{name}: {message}")
|
||||
|
||||
def send(self) -> str:
|
||||
"""Send message history."""
|
||||
self.agent.run(task=self.message_history)
|
||||
|
||||
def setup_tools(self, external_tools):
|
||||
"""
|
||||
Set up tools for the worker.
|
||||
|
||||
Parameters:
|
||||
- `external_tools` (list): List of external tools (optional).
|
||||
|
||||
Example:
|
||||
```
|
||||
external_tools = [MyTool1(), MyTool2()]
|
||||
worker = Worker(model_name="gpt-4",
|
||||
openai_api_key="my_key",
|
||||
name="My Worker",
|
||||
role="Worker",
|
||||
external_tools=external_tools,
|
||||
human_in_the_loop=False,
|
||||
temperature=0.5)
|
||||
```
|
||||
"""
|
||||
if self.tools is None:
|
||||
self.tools = []
|
||||
|
||||
if external_tools is not None:
|
||||
self.tools.extend(external_tools)
|
||||
|
||||
def setup_memory(self):
|
||||
"""
|
||||
Set up memory for the worker.
|
||||
"""
|
||||
try:
|
||||
self.vectorstore = self.memory
|
||||
|
||||
except Exception as error:
|
||||
raise RuntimeError(
|
||||
"Error setting up memory perhaps try try tuning the"
|
||||
f" embedding size: {error}"
|
||||
)
|
||||
|
||||
def setup_agent(self):
|
||||
"""
|
||||
Set up the autonomous agent.
|
||||
"""
|
||||
try:
|
||||
self.agent = AutoGPT.from_llm_and_tools(
|
||||
ai_name=self.name,
|
||||
ai_role=self.role,
|
||||
tools=self.tools,
|
||||
llm=self.llm,
|
||||
# memory = None,
|
||||
human_in_the_loop=self.human_in_the_loop,
|
||||
)
|
||||
|
||||
except Exception as error:
|
||||
raise RuntimeError(f"Error setting up agent: {error}")
|
||||
|
||||
@error_decorator
|
||||
@timing_decorator
|
||||
def run(self, task: str = None, *args, **kwargs):
|
||||
"""
|
||||
Run the autonomous agent on a given task.
|
||||
|
||||
Parameters:
|
||||
- `task`: The task to be processed.
|
||||
|
||||
Returns:
|
||||
- `result`: The result of the agent's processing.
|
||||
"""
|
||||
try:
|
||||
result = self.agent.run([task], *args, **kwargs)
|
||||
return result
|
||||
except Exception as error:
|
||||
raise RuntimeError(f"Error while running agent: {error}")
|
||||
|
||||
@error_decorator
|
||||
@timing_decorator
|
||||
def __call__(self, task: str = None, *args, **kwargs):
|
||||
"""
|
||||
Make the worker callable to run the agent on a given task.
|
||||
|
||||
Parameters:
|
||||
- `task`: The task to be processed.
|
||||
|
||||
Returns:
|
||||
- `results`: The results of the agent's processing.
|
||||
"""
|
||||
try:
|
||||
result = self.agent.run([task], *args, **kwargs)
|
||||
return result
|
||||
except Exception as error:
|
||||
raise RuntimeError(f"Error while running agent: {error}")
|
||||
@ -0,0 +1,74 @@
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
from swarms.structs.mixture_of_agents import MixtureOfAgents
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.memory.base_vectordb import BaseVectorDatabase
|
||||
|
||||
|
||||
def test_init():
|
||||
with patch.object(
|
||||
MixtureOfAgents, "agent_check"
|
||||
) as mock_agent_check, patch.object(
|
||||
MixtureOfAgents, "final_agent_check"
|
||||
) as mock_final_agent_check, patch.object(
|
||||
MixtureOfAgents, "swarm_initialization"
|
||||
) as mock_swarm_initialization, patch.object(
|
||||
MixtureOfAgents, "communication_protocol"
|
||||
) as mock_communication_protocol:
|
||||
agents = [Mock(spec=Agent)]
|
||||
final_agent = Mock(spec=Agent)
|
||||
scp = Mock(spec=BaseVectorDatabase)
|
||||
MixtureOfAgents(agents=agents, final_agent=final_agent, scp=scp)
|
||||
mock_agent_check.assert_called_once()
|
||||
mock_final_agent_check.assert_called_once()
|
||||
mock_swarm_initialization.assert_called_once()
|
||||
mock_communication_protocol.assert_called_once()
|
||||
|
||||
|
||||
def test_communication_protocol():
|
||||
agents = [Mock(spec=Agent)]
|
||||
final_agent = Mock(spec=Agent)
|
||||
scp = Mock(spec=BaseVectorDatabase)
|
||||
swarm = MixtureOfAgents(
|
||||
agents=agents, final_agent=final_agent, scp=scp
|
||||
)
|
||||
swarm.communication_protocol()
|
||||
for agent in agents:
|
||||
agent.long_term_memory.assert_called_once_with(scp)
|
||||
|
||||
|
||||
def test_agent_check():
|
||||
final_agent = Mock(spec=Agent)
|
||||
with pytest.raises(TypeError):
|
||||
MixtureOfAgents(agents="not a list", final_agent=final_agent)
|
||||
with pytest.raises(TypeError):
|
||||
MixtureOfAgents(agents=["not an agent"], final_agent=final_agent)
|
||||
|
||||
|
||||
def test_final_agent_check():
|
||||
agents = [Mock(spec=Agent)]
|
||||
with pytest.raises(TypeError):
|
||||
MixtureOfAgents(agents=agents, final_agent="not an agent")
|
||||
|
||||
|
||||
def test_swarm_initialization():
|
||||
with patch("swarms.structs.mixture_of_agents.logger") as mock_logger:
|
||||
agents = [Mock(spec=Agent)]
|
||||
final_agent = Mock(spec=Agent)
|
||||
swarm = MixtureOfAgents(agents=agents, final_agent=final_agent)
|
||||
swarm.swarm_initialization()
|
||||
assert mock_logger.info.call_count == 3
|
||||
|
||||
|
||||
def test_run():
|
||||
with patch("swarms.structs.mixture_of_agents.logger"), patch(
|
||||
"builtins.open", new_callable=Mock
|
||||
) as mock_open:
|
||||
agents = [Mock(spec=Agent)]
|
||||
final_agent = Mock(spec=Agent)
|
||||
swarm = MixtureOfAgents(agents=agents, final_agent=final_agent)
|
||||
swarm.run("task")
|
||||
for agent in agents:
|
||||
agent.run.assert_called_once()
|
||||
final_agent.run.assert_called_once()
|
||||
mock_open.assert_called_once_with(swarm.saved_file_name, "w")
|
||||
Loading…
Reference in new issue