Your Name
3 months ago
parent
8cf8761a8f
commit
c51ef4b72d
@ -0,0 +1,141 @@
|
|||||||
|
# Forest Swarm
|
||||||
|
|
||||||
|
This documentation describes the **ForestSwarm** that organizes agents into trees. Each agent specializes in processing specific tasks. Trees are collections of agents, each assigned based on their relevance to a task through keyword extraction and embedding-based similarity.
|
||||||
|
|
||||||
|
The architecture allows for efficient task assignment by selecting the most relevant agent from a set of trees. Tasks are processed asynchronously, with agents selected based on task relevance, calculated by the similarity of system prompts and task keywords.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Class: `TreeAgent`
|
||||||
|
|
||||||
|
`TreeAgent` represents an individual agent responsible for handling a specific task. Agents are initialized with a **system prompt** and are responsible for dynamically determining their relevance to a given task.
|
||||||
|
|
||||||
|
#### Attributes
|
||||||
|
|
||||||
|
| **Attribute** | **Type** | **Description** |
|
||||||
|
|--------------------------|------------------|---------------------------------------------------------------------------------|
|
||||||
|
| `system_prompt` | `str` | A string that defines the agent's area of expertise and task-handling capability.|
|
||||||
|
| `llm` | `callable` | The language model (LLM) used to process tasks (e.g., GPT-4). |
|
||||||
|
| `agent_name` | `str` | The name of the agent. |
|
||||||
|
| `system_prompt_embedding`| `tensor` | Embedding of the system prompt for similarity-based task matching. |
|
||||||
|
| `relevant_keywords` | `List[str]` | Keywords dynamically extracted from the system prompt to assist in task matching.|
|
||||||
|
| `distance` | `Optional[float]`| The computed distance between agents based on embedding similarity. |
|
||||||
|
|
||||||
|
#### Methods
|
||||||
|
|
||||||
|
| **Method** | **Input** | **Output** | **Description** |
|
||||||
|
|--------------------|---------------------------------|--------------------|---------------------------------------------------------------------------------|
|
||||||
|
| `calculate_distance(other_agent: TreeAgent)` | `other_agent: TreeAgent` | `float` | Calculates the cosine similarity between this agent and another agent. |
|
||||||
|
| `run_task(task: str)` | `task: str` | `Any` | Executes the task, logs the input/output, and returns the result. |
|
||||||
|
| `is_relevant_for_task(task: str, threshold: float = 0.7)` | `task: str, threshold: float` | `bool` | Checks if the agent is relevant for the task using keyword matching or embedding similarity.|
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Class: `Tree`
|
||||||
|
|
||||||
|
`Tree` organizes multiple agents into a hierarchical structure, where agents are sorted based on their relevance to tasks.
|
||||||
|
|
||||||
|
#### Attributes
|
||||||
|
|
||||||
|
| **Attribute** | **Type** | **Description** |
|
||||||
|
|--------------------------|------------------|---------------------------------------------------------------------------------|
|
||||||
|
| `tree_name` | `str` | The name of the tree (represents a domain of agents, e.g., "Financial Tree"). |
|
||||||
|
| `agents` | `List[TreeAgent]`| List of agents belonging to this tree. |
|
||||||
|
|
||||||
|
#### Methods
|
||||||
|
|
||||||
|
| **Method** | **Input** | **Output** | **Description** |
|
||||||
|
|--------------------|---------------------------------|--------------------|---------------------------------------------------------------------------------|
|
||||||
|
| `calculate_agent_distances()` | `None` | `None` | Calculates and assigns distances between agents based on similarity of prompts. |
|
||||||
|
| `find_relevant_agent(task: str)` | `task: str` | `Optional[TreeAgent]` | Finds the most relevant agent for a task based on keyword and embedding similarity. |
|
||||||
|
| `log_tree_execution(task: str, selected_agent: TreeAgent, result: Any)` | `task: str, selected_agent: TreeAgent, result: Any` | `None` | Logs details of the task execution by the selected agent. |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Class: `ForestSwarm`
|
||||||
|
|
||||||
|
`ForestSwarm` is the main class responsible for managing multiple trees. It oversees task delegation by finding the most relevant tree and agent for a given task.
|
||||||
|
|
||||||
|
#### Attributes
|
||||||
|
|
||||||
|
| **Attribute** | **Type** | **Description** |
|
||||||
|
|--------------------------|------------------|---------------------------------------------------------------------------------|
|
||||||
|
| `trees` | `List[Tree]` | List of trees containing agents organized by domain. |
|
||||||
|
|
||||||
|
#### Methods
|
||||||
|
|
||||||
|
| **Method** | **Input** | **Output** | **Description** |
|
||||||
|
|--------------------|---------------------------------|--------------------|---------------------------------------------------------------------------------|
|
||||||
|
| `find_relevant_tree(task: str)` | `task: str` | `Optional[Tree]` | Searches across all trees to find the most relevant tree based on task requirements.|
|
||||||
|
| `run(task: str)` | `task: str` | `Any` | Executes the task by finding the most relevant agent from the relevant tree. |
|
||||||
|
|
||||||
|
## Full Code Example
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Example Workflow
|
||||||
|
|
||||||
|
1. **Create Agents**: Agents are initialized with varying system prompts, representing different areas of expertise (e.g., stock analysis, tax filing).
|
||||||
|
2. **Create Trees**: Agents are grouped into trees, with each tree representing a domain (e.g., "Financial Tree", "Investment Tree").
|
||||||
|
3. **Run Task**: When a task is submitted, the system traverses through all trees and finds the most relevant agent to handle the task.
|
||||||
|
4. **Task Execution**: The selected agent processes the task, and the result is returned.
|
||||||
|
|
||||||
|
```plaintext
|
||||||
|
Task: "Our company is incorporated in Delaware, how do we do our taxes for free?"
|
||||||
|
```
|
||||||
|
|
||||||
|
**Process**:
|
||||||
|
- The system searches through the `Financial Tree` and `Investment Tree`.
|
||||||
|
- The most relevant agent (likely the "Tax Filing Agent") is selected based on keyword matching and prompt similarity.
|
||||||
|
- The task is processed, and the result is logged and returned.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Analysis of the Swarm Architecture
|
||||||
|
|
||||||
|
The **Swarm Architecture** leverages a hierarchical structure (forest) composed of individual trees, each containing agents specialized in specific domains. This design allows for:
|
||||||
|
|
||||||
|
- **Modular and Scalable Organization**: By separating agents into trees, it is easy to expand or contract the system by adding or removing trees or agents.
|
||||||
|
- **Task Specialization**: Each agent is specialized, which ensures that tasks are matched with the most appropriate agent based on relevance and expertise.
|
||||||
|
- **Dynamic Matching**: The architecture uses both keyword-based and embedding-based matching to assign tasks, ensuring a high level of accuracy in agent selection.
|
||||||
|
- **Logging and Accountability**: Each task execution is logged in detail, providing transparency and an audit trail of which agent handled which task and the results produced.
|
||||||
|
- **Asynchronous Task Execution**: The architecture can be adapted for asynchronous task processing, making it scalable and suitable for large-scale task handling in real-time systems.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Mermaid Diagram of the Swarm Architecture
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
graph TD
|
||||||
|
A[ForestSwarm] --> B[Financial Tree]
|
||||||
|
A --> C[Investment Tree]
|
||||||
|
|
||||||
|
B --> D[Stock Analysis Agent]
|
||||||
|
B --> E[Financial Planning Agent]
|
||||||
|
B --> F[Retirement Strategy Agent]
|
||||||
|
|
||||||
|
C --> G[Tax Filing Agent]
|
||||||
|
C --> H[Investment Strategy Agent]
|
||||||
|
C --> I[ROTH IRA Agent]
|
||||||
|
|
||||||
|
subgraph Tree Agents
|
||||||
|
D[Stock Analysis Agent]
|
||||||
|
E[Financial Planning Agent]
|
||||||
|
F[Retirement Strategy Agent]
|
||||||
|
G[Tax Filing Agent]
|
||||||
|
H[Investment Strategy Agent]
|
||||||
|
I[ROTH IRA Agent]
|
||||||
|
end
|
||||||
|
```
|
||||||
|
|
||||||
|
### Explanation of the Diagram
|
||||||
|
|
||||||
|
- **ForestSwarm**: Represents the top-level structure managing multiple trees.
|
||||||
|
- **Trees**: In the example, two trees exist—**Financial Tree** and **Investment Tree**—each containing agents related to specific domains.
|
||||||
|
- **Agents**: Each agent within the tree is responsible for handling tasks in its area of expertise. Agents within a tree are organized based on their prompt similarity (distance).
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Summary
|
||||||
|
|
||||||
|
This **Multi-Agent Tree Structure** provides an efficient, scalable, and accurate architecture for delegating and executing tasks based on domain-specific expertise. The combination of hierarchical organization, dynamic task matching, and logging ensures reliability, performance, and transparency in task execution.
|
||||||
@ -0,0 +1,44 @@
|
|||||||
|
from swarms.structs.tree_swarm import TreeAgent, Tree, ForestSwarm
|
||||||
|
# Example Usage:
|
||||||
|
|
||||||
|
# Create agents with varying system prompts and dynamically generated distances/keywords
|
||||||
|
agents_tree1 = [
|
||||||
|
TreeAgent(
|
||||||
|
system_prompt="Stock Analysis Agent",
|
||||||
|
agent_name="Stock Analysis Agent",
|
||||||
|
),
|
||||||
|
TreeAgent(
|
||||||
|
system_prompt="Financial Planning Agent",
|
||||||
|
agent_name="Financial Planning Agent",
|
||||||
|
),
|
||||||
|
TreeAgent(
|
||||||
|
agent_name="Retirement Strategy Agent",
|
||||||
|
system_prompt="Retirement Strategy Agent",
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
agents_tree2 = [
|
||||||
|
TreeAgent(
|
||||||
|
system_prompt="Tax Filing Agent",
|
||||||
|
agent_name="Tax Filing Agent",
|
||||||
|
),
|
||||||
|
TreeAgent(
|
||||||
|
system_prompt="Investment Strategy Agent",
|
||||||
|
agent_name="Investment Strategy Agent",
|
||||||
|
),
|
||||||
|
TreeAgent(
|
||||||
|
system_prompt="ROTH IRA Agent", agent_name="ROTH IRA Agent"
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
# Create trees
|
||||||
|
tree1 = Tree(tree_name="Financial Tree", agents=agents_tree1)
|
||||||
|
tree2 = Tree(tree_name="Investment Tree", agents=agents_tree2)
|
||||||
|
|
||||||
|
# Create the ForestSwarm
|
||||||
|
multi_agent_structure = ForestSwarm(trees=[tree1, tree2])
|
||||||
|
|
||||||
|
# Run a task
|
||||||
|
task = "Our company is incorporated in delaware, how do we do our taxes for free?"
|
||||||
|
output = multi_agent_structure.run(task)
|
||||||
|
print(output)
|
||||||
@ -0,0 +1,360 @@
|
|||||||
|
import asyncio
|
||||||
|
import os
|
||||||
|
import uuid
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
|
import chromadb
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
from loguru import logger
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
from swarm_models import OpenAIChat
|
||||||
|
|
||||||
|
from swarms import Agent
|
||||||
|
from swarms.prompts.finance_agent_sys_prompt import (
|
||||||
|
FINANCIAL_AGENT_SYS_PROMPT,
|
||||||
|
)
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
# Initialize ChromaDB client
|
||||||
|
chroma_client = chromadb.Client()
|
||||||
|
|
||||||
|
# Create a ChromaDB collection to store tasks, responses, and all swarm activity
|
||||||
|
swarm_collection = chroma_client.create_collection(
|
||||||
|
name="swarm_activity"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class InteractionLog(BaseModel):
|
||||||
|
"""
|
||||||
|
Pydantic model to log all interactions between agents, tasks, and responses.
|
||||||
|
"""
|
||||||
|
|
||||||
|
interaction_id: str = Field(
|
||||||
|
default_factory=lambda: str(uuid.uuid4()),
|
||||||
|
description="Unique ID for the interaction.",
|
||||||
|
)
|
||||||
|
agent_name: str
|
||||||
|
task: str
|
||||||
|
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
||||||
|
response: Optional[Dict[str, Any]] = None
|
||||||
|
status: str = Field(
|
||||||
|
description="The status of the interaction, e.g., 'completed', 'failed'."
|
||||||
|
)
|
||||||
|
neighbors: Optional[List[str]] = (
|
||||||
|
None # Names of neighboring agents involved
|
||||||
|
)
|
||||||
|
conversation_id: Optional[str] = Field(
|
||||||
|
default_factory=lambda: str(uuid.uuid4()),
|
||||||
|
description="Unique ID for the conversation history.",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class AgentHealthStatus(BaseModel):
|
||||||
|
"""
|
||||||
|
Pydantic model to log and monitor agent health.
|
||||||
|
"""
|
||||||
|
|
||||||
|
agent_name: str
|
||||||
|
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
||||||
|
status: str = Field(
|
||||||
|
default="available",
|
||||||
|
description="Agent health status, e.g., 'available', 'busy', 'failed'.",
|
||||||
|
)
|
||||||
|
active_tasks: int = Field(
|
||||||
|
0,
|
||||||
|
description="Number of active tasks assigned to this agent.",
|
||||||
|
)
|
||||||
|
load: float = Field(
|
||||||
|
0.0,
|
||||||
|
description="Current load on the agent (CPU or memory usage).",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class Swarm:
|
||||||
|
"""
|
||||||
|
A scalable swarm architecture where agents can communicate by posting and querying all activities to ChromaDB.
|
||||||
|
Every input task, response, and action by the agents is logged to the vector database for persistent tracking.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
agents (List[Agent]): A list of initialized agents.
|
||||||
|
chroma_client (chroma.Client): An instance of the ChromaDB client for agent-to-agent communication.
|
||||||
|
api_key (str): The OpenAI API key.
|
||||||
|
health_statuses (Dict[str, AgentHealthStatus]): A dictionary to monitor agent health statuses.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
agents: List[Agent],
|
||||||
|
chroma_client: chromadb.Client,
|
||||||
|
api_key: str,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Initializes the swarm with agents and a ChromaDB client for vector storage and communication.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
agents (List[Agent]): A list of initialized agents.
|
||||||
|
chroma_client (chroma.Client): The ChromaDB client for handling vector embeddings.
|
||||||
|
api_key (str): The OpenAI API key.
|
||||||
|
"""
|
||||||
|
self.agents = agents
|
||||||
|
self.chroma_client = chroma_client
|
||||||
|
self.api_key = api_key
|
||||||
|
self.health_statuses: Dict[str, AgentHealthStatus] = {
|
||||||
|
agent.agent_name: AgentHealthStatus(
|
||||||
|
agent_name=agent.agent_name
|
||||||
|
)
|
||||||
|
for agent in agents
|
||||||
|
}
|
||||||
|
logger.info(f"Swarm initialized with {len(agents)} agents.")
|
||||||
|
|
||||||
|
def _log_to_db(
|
||||||
|
self, data: Dict[str, Any], description: str
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Logs a dictionary of data into the ChromaDB collection as a new entry.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data (Dict[str, Any]): The data to log in the database (task, response, etc.).
|
||||||
|
description (str): Description of the action (e.g., 'task', 'response').
|
||||||
|
"""
|
||||||
|
logger.info(f"Logging {description} to the database: {data}")
|
||||||
|
swarm_collection.add(
|
||||||
|
documents=[str(data)],
|
||||||
|
ids=[str(uuid.uuid4())], # Unique ID for each entry
|
||||||
|
metadatas=[
|
||||||
|
{
|
||||||
|
"description": description,
|
||||||
|
"timestamp": datetime.utcnow().isoformat(),
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
f"{description.capitalize()} logged successfully."
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _find_most_relevant_agent(
|
||||||
|
self, task: str
|
||||||
|
) -> Optional[Agent]:
|
||||||
|
"""
|
||||||
|
Finds the agent whose system prompt is most relevant to the given task by querying ChromaDB.
|
||||||
|
If no relevant agents are found, return None and log a message.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
task (str): The task for which to find the most relevant agent.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Optional[Agent]: The most relevant agent for the task, or None if no relevant agent is found.
|
||||||
|
"""
|
||||||
|
logger.info(
|
||||||
|
f"Searching for the most relevant agent for the task: {task}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Query ChromaDB collection for nearest neighbor to the task
|
||||||
|
result = swarm_collection.query(
|
||||||
|
query_texts=[task], n_results=4
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check if the query result contains any data
|
||||||
|
if not result["ids"] or not result["ids"][0]:
|
||||||
|
logger.error(
|
||||||
|
"No relevant agents found for the given task."
|
||||||
|
)
|
||||||
|
return None # No agent found, return None
|
||||||
|
|
||||||
|
# Extract the agent ID from the result and find the corresponding agent
|
||||||
|
agent_id = result["ids"][0][0]
|
||||||
|
most_relevant_agent = next(
|
||||||
|
(
|
||||||
|
agent
|
||||||
|
for agent in self.agents
|
||||||
|
if agent.agent_name == agent_id
|
||||||
|
),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
|
||||||
|
if most_relevant_agent:
|
||||||
|
logger.info(
|
||||||
|
f"Most relevant agent for task '{task}' is {most_relevant_agent.agent_name}."
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.error("No matching agent found in the agent list.")
|
||||||
|
|
||||||
|
return most_relevant_agent
|
||||||
|
|
||||||
|
def _monitor_health(self, agent: Agent) -> None:
|
||||||
|
"""
|
||||||
|
Monitors the health status of agents and logs it to the database.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
agent (Agent): The agent whose health is being monitored.
|
||||||
|
"""
|
||||||
|
current_status = self.health_statuses[agent.agent_name]
|
||||||
|
current_status.active_tasks += (
|
||||||
|
1 # Example increment for active tasks
|
||||||
|
)
|
||||||
|
current_status.status = (
|
||||||
|
"busy" if current_status.active_tasks > 0 else "available"
|
||||||
|
)
|
||||||
|
current_status.load = 0.5 # Placeholder for real load data
|
||||||
|
logger.info(
|
||||||
|
f"Agent {agent.agent_name} is currently {current_status.status} with load {current_status.load}."
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log health status to the database
|
||||||
|
self._log_to_db(current_status.dict(), "health status")
|
||||||
|
|
||||||
|
def post_message(self, agent: Agent, message: str) -> None:
|
||||||
|
"""
|
||||||
|
Posts a message from an agent to the shared database.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
agent (Agent): The agent posting the message.
|
||||||
|
message (str): The message to be posted.
|
||||||
|
"""
|
||||||
|
logger.info(
|
||||||
|
f"Agent {agent.agent_name} posting message: {message}"
|
||||||
|
)
|
||||||
|
message_data = {
|
||||||
|
"agent_name": agent.agent_name,
|
||||||
|
"message": message,
|
||||||
|
"timestamp": datetime.utcnow().isoformat(),
|
||||||
|
}
|
||||||
|
self._log_to_db(message_data, "message")
|
||||||
|
|
||||||
|
def query_messages(
|
||||||
|
self, query: str, n_results: int = 5
|
||||||
|
) -> List[Dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Queries the database for relevant messages.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
query (str): The query message or task for which to retrieve related messages.
|
||||||
|
n_results (int, optional): The number of relevant messages to retrieve. Defaults to 5.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List[Dict[str, Any]]: A list of relevant messages and their metadata.
|
||||||
|
"""
|
||||||
|
logger.info(f"Querying the database for query: {query}")
|
||||||
|
results = swarm_collection.query(
|
||||||
|
query_texts=[query], n_results=n_results
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
f"Found {len(results['documents'])} relevant messages."
|
||||||
|
)
|
||||||
|
return results
|
||||||
|
|
||||||
|
async def run_async(self, task: str) -> None:
|
||||||
|
"""
|
||||||
|
Main entry point to find the most relevant agent, submit the task, and allow agents to
|
||||||
|
query the database to understand the task's history. Logs every task and response.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
task (str): The task to be completed.
|
||||||
|
"""
|
||||||
|
# Query past messages to understand task history
|
||||||
|
past_messages = self.query_messages(task)
|
||||||
|
logger.info(
|
||||||
|
f"Past messages related to task '{task}': {past_messages}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Find the most relevant agent
|
||||||
|
agent = await self._find_most_relevant_agent(task)
|
||||||
|
|
||||||
|
if agent is None:
|
||||||
|
logger.error(
|
||||||
|
f"No relevant agent found for task: {task}. Task submission aborted."
|
||||||
|
)
|
||||||
|
return # Exit the function if no relevant agent is found
|
||||||
|
|
||||||
|
# Submit the task to the agent if found
|
||||||
|
await self._submit_task_to_agent(agent, task)
|
||||||
|
|
||||||
|
async def _submit_task_to_agent(
|
||||||
|
self, agent: Agent, task: str
|
||||||
|
) -> Dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Submits a task to the specified agent and logs the result asynchronously.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
agent (Agent): The agent to which the task will be submitted.
|
||||||
|
task (str): The task to be solved.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dict[str, Any]: The result of the task from the agent.
|
||||||
|
"""
|
||||||
|
if agent is None:
|
||||||
|
logger.error("No agent provided for task submission.")
|
||||||
|
return
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
f"Submitting task '{task}' to agent {agent.agent_name}."
|
||||||
|
)
|
||||||
|
|
||||||
|
interaction_log = InteractionLog(
|
||||||
|
agent_name=agent.agent_name, task=task, status="started"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log the task as a message to the shared database
|
||||||
|
self._log_to_db(
|
||||||
|
{"task": task, "agent_name": agent.agent_name}, "task"
|
||||||
|
)
|
||||||
|
|
||||||
|
result = await agent.run(task)
|
||||||
|
|
||||||
|
interaction_log.response = result
|
||||||
|
interaction_log.status = "completed"
|
||||||
|
interaction_log.timestamp = datetime.utcnow()
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
f"Task completed by agent {agent.agent_name}. Logged interaction: {interaction_log.dict()}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log the result as a message to the shared database
|
||||||
|
self._log_to_db(
|
||||||
|
{"response": result, "agent_name": agent.agent_name},
|
||||||
|
"response",
|
||||||
|
)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def run(self, task: str, *args, **kwargs):
|
||||||
|
return asyncio.run(self.run_async(task))
|
||||||
|
|
||||||
|
|
||||||
|
# Initialize the OpenAI model and agents
|
||||||
|
api_key = os.getenv("OPENAI_API_KEY")
|
||||||
|
|
||||||
|
model = OpenAIChat(
|
||||||
|
openai_api_key=api_key, model_name="gpt-4o-mini", temperature=0.1
|
||||||
|
)
|
||||||
|
|
||||||
|
# Example agent creation
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="Financial-Analysis-Agent",
|
||||||
|
system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
|
||||||
|
llm=model,
|
||||||
|
max_loops=1,
|
||||||
|
autosave=True,
|
||||||
|
dashboard=False,
|
||||||
|
verbose=True,
|
||||||
|
dynamic_temperature_enabled=True,
|
||||||
|
saved_state_path="finance_agent.json",
|
||||||
|
user_name="swarms_corp",
|
||||||
|
retry_attempts=1,
|
||||||
|
context_length=200000,
|
||||||
|
return_step_meta=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Example agents list
|
||||||
|
agents_list = [agent]
|
||||||
|
|
||||||
|
# Create the swarm
|
||||||
|
swarm = Swarm(
|
||||||
|
agents=agents_list, chroma_client=chroma_client, api_key=api_key
|
||||||
|
)
|
||||||
|
|
||||||
|
# Execute tasks asynchronously
|
||||||
|
task = "How can I establish a ROTH IRA to buy stocks and get a tax break? What are the criteria?"
|
||||||
|
print(swarm.run(task))
|
||||||
@ -0,0 +1,362 @@
|
|||||||
|
import os
|
||||||
|
import uuid
|
||||||
|
from collections import Counter
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import Any, List, Optional
|
||||||
|
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
from loguru import logger
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
from sentence_transformers import SentenceTransformer, util
|
||||||
|
from swarm_models import OpenAIChat
|
||||||
|
|
||||||
|
from swarms import Agent
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
# Get the OpenAI API key from the environment variable
|
||||||
|
api_key = os.getenv("OPENAI_API_KEY")
|
||||||
|
|
||||||
|
# # Set up loguru to log into a file and console
|
||||||
|
# logger.add(
|
||||||
|
# "multi_agent_log_{time}.log",
|
||||||
|
# format="{time} {level} {message}",
|
||||||
|
# level="DEBUG",
|
||||||
|
# rotation="10 MB",
|
||||||
|
# )
|
||||||
|
|
||||||
|
# Pretrained model for embeddings
|
||||||
|
embedding_model = SentenceTransformer(
|
||||||
|
"all-MiniLM-L6-v2"
|
||||||
|
) # A small, fast model for embedding
|
||||||
|
|
||||||
|
# Create an instance of the OpenAIChat class
|
||||||
|
model = OpenAIChat(
|
||||||
|
openai_api_key=api_key,
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
temperature=0.1,
|
||||||
|
max_tokens=2000,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# Pydantic Models for Logging
|
||||||
|
class AgentLogInput(BaseModel):
|
||||||
|
log_id: str = Field(
|
||||||
|
default_factory=lambda: str(uuid.uuid4()), alias="id"
|
||||||
|
)
|
||||||
|
agent_name: str
|
||||||
|
task: str
|
||||||
|
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
||||||
|
|
||||||
|
|
||||||
|
class AgentLogOutput(BaseModel):
|
||||||
|
log_id: str = Field(
|
||||||
|
default_factory=lambda: str(uuid.uuid4()), alias="id"
|
||||||
|
)
|
||||||
|
agent_name: str
|
||||||
|
result: Any
|
||||||
|
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
||||||
|
|
||||||
|
|
||||||
|
class TreeLog(BaseModel):
|
||||||
|
log_id: str = Field(
|
||||||
|
default_factory=lambda: str(uuid.uuid4()), alias="id"
|
||||||
|
)
|
||||||
|
tree_name: str
|
||||||
|
task: str
|
||||||
|
selected_agent: str
|
||||||
|
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
||||||
|
result: Any
|
||||||
|
|
||||||
|
|
||||||
|
def extract_keywords(prompt: str, top_n: int = 5) -> List[str]:
|
||||||
|
"""
|
||||||
|
A simplified keyword extraction function using basic word splitting instead of NLTK tokenization.
|
||||||
|
"""
|
||||||
|
words = prompt.lower().split()
|
||||||
|
filtered_words = [word for word in words if word.isalnum()]
|
||||||
|
word_counts = Counter(filtered_words)
|
||||||
|
return [word for word, _ in word_counts.most_common(top_n)]
|
||||||
|
|
||||||
|
|
||||||
|
class TreeAgent(Agent):
|
||||||
|
"""
|
||||||
|
A specialized Agent class that contains information about the system prompt's
|
||||||
|
locality and allows for dynamic chaining of agents in trees.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
system_prompt: str = None,
|
||||||
|
llm: callable = model,
|
||||||
|
agent_name: Optional[str] = None,
|
||||||
|
*args,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
agent_name = agent_name
|
||||||
|
super().__init__(
|
||||||
|
system_prompt=system_prompt,
|
||||||
|
llm=llm,
|
||||||
|
agent_name=agent_name,
|
||||||
|
*args,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
self.system_prompt_embedding = embedding_model.encode(
|
||||||
|
system_prompt, convert_to_tensor=True
|
||||||
|
)
|
||||||
|
|
||||||
|
# Automatically extract keywords from system prompt
|
||||||
|
self.relevant_keywords = extract_keywords(system_prompt)
|
||||||
|
|
||||||
|
# Distance is now calculated based on similarity between agents' prompts
|
||||||
|
self.distance = None # Will be dynamically calculated later
|
||||||
|
|
||||||
|
def calculate_distance(self, other_agent: "TreeAgent") -> float:
|
||||||
|
"""
|
||||||
|
Calculate the distance between this agent and another agent using embedding similarity.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
other_agent (TreeAgent): Another agent in the tree.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
float: Distance score between 0 and 1, with 0 being close and 1 being far.
|
||||||
|
"""
|
||||||
|
similarity = util.pytorch_cos_sim(
|
||||||
|
self.system_prompt_embedding,
|
||||||
|
other_agent.system_prompt_embedding,
|
||||||
|
).item()
|
||||||
|
distance = (
|
||||||
|
1 - similarity
|
||||||
|
) # Closer agents have a smaller distance
|
||||||
|
return distance
|
||||||
|
|
||||||
|
def run_task(self, task: str) -> Any:
|
||||||
|
input_log = AgentLogInput(
|
||||||
|
agent_name=self.agent_name,
|
||||||
|
task=task,
|
||||||
|
timestamp=datetime.now(),
|
||||||
|
)
|
||||||
|
logger.info(f"Running task on {self.agent_name}: {task}")
|
||||||
|
logger.debug(f"Input Log: {input_log.json()}")
|
||||||
|
|
||||||
|
result = self.run(task)
|
||||||
|
|
||||||
|
output_log = AgentLogOutput(
|
||||||
|
agent_name=self.agent_name,
|
||||||
|
result=result,
|
||||||
|
timestamp=datetime.now(),
|
||||||
|
)
|
||||||
|
logger.info(f"Task result from {self.agent_name}: {result}")
|
||||||
|
logger.debug(f"Output Log: {output_log.json()}")
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def is_relevant_for_task(
|
||||||
|
self, task: str, threshold: float = 0.7
|
||||||
|
) -> bool:
|
||||||
|
"""
|
||||||
|
Checks if the agent is relevant for the given task using both keyword matching and embedding similarity.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
task (str): The task to be executed.
|
||||||
|
threshold (float): The cosine similarity threshold for embedding-based matching.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if the agent is relevant, False otherwise.
|
||||||
|
"""
|
||||||
|
# Check if any of the relevant keywords are present in the task (case-insensitive)
|
||||||
|
keyword_match = any(
|
||||||
|
keyword.lower() in task.lower()
|
||||||
|
for keyword in self.relevant_keywords
|
||||||
|
)
|
||||||
|
|
||||||
|
# Perform embedding similarity match if keyword match is not found
|
||||||
|
if not keyword_match:
|
||||||
|
task_embedding = embedding_model.encode(
|
||||||
|
task, convert_to_tensor=True
|
||||||
|
)
|
||||||
|
similarity = util.pytorch_cos_sim(
|
||||||
|
self.system_prompt_embedding, task_embedding
|
||||||
|
).item()
|
||||||
|
logger.info(
|
||||||
|
f"Semantic similarity between task and {self.agent_name}: {similarity:.2f}"
|
||||||
|
)
|
||||||
|
return similarity >= threshold
|
||||||
|
|
||||||
|
return True # Return True if keyword match is found
|
||||||
|
|
||||||
|
|
||||||
|
class Tree:
|
||||||
|
def __init__(self, tree_name: str, agents: List[TreeAgent]):
|
||||||
|
"""
|
||||||
|
Initializes a tree of agents.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tree_name (str): The name of the tree.
|
||||||
|
agents (List[TreeAgent]): A list of agents in the tree.
|
||||||
|
"""
|
||||||
|
self.tree_name = tree_name
|
||||||
|
self.agents = agents
|
||||||
|
self.calculate_agent_distances()
|
||||||
|
|
||||||
|
def calculate_agent_distances(self):
|
||||||
|
"""
|
||||||
|
Automatically calculate and assign distances between agents in the tree based on prompt similarity.
|
||||||
|
"""
|
||||||
|
logger.info(
|
||||||
|
f"Calculating distances between agents in tree '{self.tree_name}'"
|
||||||
|
)
|
||||||
|
for i, agent in enumerate(self.agents):
|
||||||
|
if i > 0:
|
||||||
|
agent.distance = agent.calculate_distance(
|
||||||
|
self.agents[i - 1]
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
agent.distance = 0 # First agent is closest
|
||||||
|
|
||||||
|
# Sort agents by distance after calculation
|
||||||
|
self.agents.sort(key=lambda agent: agent.distance)
|
||||||
|
|
||||||
|
def find_relevant_agent(self, task: str) -> Optional[TreeAgent]:
|
||||||
|
"""
|
||||||
|
Finds the most relevant agent in the tree for the given task based on its system prompt.
|
||||||
|
Uses both keyword and semantic similarity matching.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
task (str): The task or query for which we need to find a relevant agent.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Optional[TreeAgent]: The most relevant agent, or None if no match found.
|
||||||
|
"""
|
||||||
|
logger.info(
|
||||||
|
f"Searching relevant agent in tree '{self.tree_name}' for task: {task}"
|
||||||
|
)
|
||||||
|
for agent in self.agents:
|
||||||
|
if agent.is_relevant_for_task(task):
|
||||||
|
return agent
|
||||||
|
logger.warning(
|
||||||
|
f"No relevant agent found in tree '{self.tree_name}' for task: {task}"
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
|
def log_tree_execution(
|
||||||
|
self, task: str, selected_agent: TreeAgent, result: Any
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Logs the execution details of a tree, including selected agent and result.
|
||||||
|
"""
|
||||||
|
tree_log = TreeLog(
|
||||||
|
tree_name=self.tree_name,
|
||||||
|
task=task,
|
||||||
|
selected_agent=selected_agent.agent_name,
|
||||||
|
timestamp=datetime.now(),
|
||||||
|
result=result,
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
f"Tree '{self.tree_name}' executed task with agent '{selected_agent.agent_name}'"
|
||||||
|
)
|
||||||
|
logger.debug(f"Tree Log: {tree_log.json()}")
|
||||||
|
|
||||||
|
|
||||||
|
class ForestSwarm:
|
||||||
|
def __init__(self, trees: List[Tree], *args, **kwargs):
|
||||||
|
"""
|
||||||
|
Initializes the structure with multiple trees of agents.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
trees (List[Tree]): A list of trees in the structure.
|
||||||
|
"""
|
||||||
|
self.trees = trees
|
||||||
|
# Add auto grouping based on trees.
|
||||||
|
# Add auto group agents
|
||||||
|
|
||||||
|
|
||||||
|
def find_relevant_tree(self, task: str) -> Optional[Tree]:
|
||||||
|
"""
|
||||||
|
Finds the most relevant tree based on the given task.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
task (str): The task or query for which we need to find a relevant tree.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Optional[Tree]: The most relevant tree, or None if no match found.
|
||||||
|
"""
|
||||||
|
logger.info(
|
||||||
|
f"Searching for the most relevant tree for task: {task}"
|
||||||
|
)
|
||||||
|
for tree in self.trees:
|
||||||
|
if tree.find_relevant_agent(task):
|
||||||
|
return tree
|
||||||
|
logger.warning(f"No relevant tree found for task: {task}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
def run(self, task: str) -> Any:
|
||||||
|
"""
|
||||||
|
Executes the given task by finding the most relevant tree and agent within that tree.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
task (str): The task or query to be executed.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Any: The result of the task after it has been processed by the agents.
|
||||||
|
"""
|
||||||
|
logger.info(
|
||||||
|
f"Running task across MultiAgentTreeStructure: {task}"
|
||||||
|
)
|
||||||
|
relevant_tree = self.find_relevant_tree(task)
|
||||||
|
if relevant_tree:
|
||||||
|
agent = relevant_tree.find_relevant_agent(task)
|
||||||
|
if agent:
|
||||||
|
result = agent.run_task(task)
|
||||||
|
relevant_tree.log_tree_execution(task, agent, result)
|
||||||
|
return result
|
||||||
|
else:
|
||||||
|
logger.error(
|
||||||
|
"Task could not be completed: No relevant agent or tree found."
|
||||||
|
)
|
||||||
|
return "No relevant agent found to handle this task."
|
||||||
|
|
||||||
|
|
||||||
|
# # Example Usage:
|
||||||
|
|
||||||
|
# # Create agents with varying system prompts and dynamically generated distances/keywords
|
||||||
|
# agents_tree1 = [
|
||||||
|
# TreeAgent(
|
||||||
|
# system_prompt="Stock Analysis Agent",
|
||||||
|
# agent_name="Stock Analysis Agent",
|
||||||
|
# ),
|
||||||
|
# TreeAgent(
|
||||||
|
# system_prompt="Financial Planning Agent",
|
||||||
|
# agent_name="Financial Planning Agent",
|
||||||
|
# ),
|
||||||
|
# TreeAgent(
|
||||||
|
# agent_name="Retirement Strategy Agent",
|
||||||
|
# system_prompt="Retirement Strategy Agent",
|
||||||
|
# ),
|
||||||
|
# ]
|
||||||
|
|
||||||
|
# agents_tree2 = [
|
||||||
|
# TreeAgent(
|
||||||
|
# system_prompt="Tax Filing Agent",
|
||||||
|
# agent_name="Tax Filing Agent",
|
||||||
|
# ),
|
||||||
|
# TreeAgent(
|
||||||
|
# system_prompt="Investment Strategy Agent",
|
||||||
|
# agent_name="Investment Strategy Agent",
|
||||||
|
# ),
|
||||||
|
# TreeAgent(
|
||||||
|
# system_prompt="ROTH IRA Agent", agent_name="ROTH IRA Agent"
|
||||||
|
# ),
|
||||||
|
# ]
|
||||||
|
|
||||||
|
# # Create trees
|
||||||
|
# tree1 = Tree(tree_name="Financial Tree", agents=agents_tree1)
|
||||||
|
# tree2 = Tree(tree_name="Investment Tree", agents=agents_tree2)
|
||||||
|
|
||||||
|
# # Create the ForestSwarm
|
||||||
|
# multi_agent_structure = ForestSwarm(trees=[tree1, tree2])
|
||||||
|
|
||||||
|
# # Run a task
|
||||||
|
# task = "Our company is incorporated in delaware, how do we do our taxes for free?"
|
||||||
|
# output = multi_agent_structure.run(task)
|
||||||
|
# print(output)
|
||||||
Loading…
Reference in new issue