You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
667 lines
22 KiB
667 lines
22 KiB
"""
|
|
GraphSwarm: A production-grade framework for orchestrating swarms of agents
|
|
Author: Claude
|
|
License: MIT
|
|
Version: 2.0.0
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import datetime
|
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
|
|
import chromadb
|
|
import networkx as nx
|
|
from loguru import logger
|
|
from pydantic import BaseModel, Field
|
|
|
|
from swarms import Agent
|
|
|
|
|
|
# Configure logging
|
|
logger.add(
|
|
"graphswarm.log",
|
|
rotation="500 MB",
|
|
retention="10 days",
|
|
level="INFO",
|
|
format="{time:YYYY-MM-DD at HH:mm:ss} | {level} | {message}",
|
|
)
|
|
|
|
|
|
class AgentOutput(BaseModel):
|
|
"""Structured output from an agent."""
|
|
|
|
agent_name: str
|
|
timestamp: float = Field(default_factory=time.time)
|
|
output: Any
|
|
execution_time: float
|
|
error: Optional[str] = None
|
|
metadata: Dict = Field(default_factory=dict)
|
|
|
|
|
|
class SwarmOutput(BaseModel):
|
|
"""Structured output from the entire swarm."""
|
|
|
|
timestamp: float = Field(default_factory=time.time)
|
|
outputs: Dict[str, AgentOutput]
|
|
execution_time: float
|
|
success: bool
|
|
error: Optional[str] = None
|
|
metadata: Dict = Field(default_factory=dict)
|
|
|
|
|
|
class SwarmMemory:
|
|
"""Vector-based memory system for GraphSwarm using ChromaDB."""
|
|
|
|
def __init__(self, collection_name: str = "swarm_memories"):
|
|
"""Initialize SwarmMemory with ChromaDB."""
|
|
self.client = chromadb.Client()
|
|
|
|
# Get or create collection
|
|
self.collection = self.client.get_or_create_collection(
|
|
name=collection_name,
|
|
metadata={"description": "GraphSwarm execution memories"},
|
|
)
|
|
|
|
def store_execution(self, task: str, result: SwarmOutput):
|
|
"""Store execution results in vector memory."""
|
|
try:
|
|
# Create metadata
|
|
metadata = {
|
|
"timestamp": datetime.now().isoformat(),
|
|
"success": result.success,
|
|
"execution_time": result.execution_time,
|
|
"agent_sequence": json.dumps(
|
|
[name for name in result.outputs.keys()]
|
|
),
|
|
"error": result.error if result.error else "",
|
|
}
|
|
|
|
# Create document from outputs
|
|
document = {
|
|
"task": task,
|
|
"outputs": json.dumps(
|
|
{
|
|
name: {
|
|
"output": str(output.output),
|
|
"execution_time": output.execution_time,
|
|
"error": output.error,
|
|
}
|
|
for name, output in result.outputs.items()
|
|
}
|
|
),
|
|
}
|
|
|
|
# Store in ChromaDB
|
|
self.collection.add(
|
|
documents=[json.dumps(document)],
|
|
metadatas=[metadata],
|
|
ids=[f"exec_{datetime.now().timestamp()}"],
|
|
)
|
|
|
|
print("added to database")
|
|
|
|
logger.info(f"Stored execution in memory: {task}")
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Failed to store execution in memory: {str(e)}"
|
|
)
|
|
|
|
def get_similar_executions(self, task: str, limit: int = 5):
|
|
"""Retrieve similar past executions."""
|
|
try:
|
|
# Query ChromaDB for similar executions
|
|
results = self.collection.query(
|
|
query_texts=[task],
|
|
n_results=limit,
|
|
include=["documents", "metadatas"],
|
|
)
|
|
|
|
print(results)
|
|
|
|
if not results["documents"]:
|
|
return []
|
|
|
|
# Process results
|
|
executions = []
|
|
for doc, metadata in zip(
|
|
results["documents"][0], results["metadatas"][0]
|
|
):
|
|
doc_dict = json.loads(doc)
|
|
executions.append(
|
|
{
|
|
"task": doc_dict["task"],
|
|
"outputs": json.loads(doc_dict["outputs"]),
|
|
"success": metadata["success"],
|
|
"execution_time": metadata["execution_time"],
|
|
"agent_sequence": json.loads(
|
|
metadata["agent_sequence"]
|
|
),
|
|
"timestamp": metadata["timestamp"],
|
|
}
|
|
)
|
|
|
|
return executions
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Failed to retrieve similar executions: {str(e)}"
|
|
)
|
|
return []
|
|
|
|
def get_optimal_sequence(self, task: str) -> Optional[List[str]]:
|
|
"""Get the most successful agent sequence for similar tasks."""
|
|
similar_executions = self.get_similar_executions(task)
|
|
print(f"similar_executions {similar_executions}")
|
|
|
|
if not similar_executions:
|
|
return None
|
|
|
|
# Sort by success and execution time
|
|
successful_execs = [
|
|
ex for ex in similar_executions if ex["success"]
|
|
]
|
|
|
|
if not successful_execs:
|
|
return None
|
|
|
|
# Return sequence from most successful execution
|
|
return successful_execs[0]["agent_sequence"]
|
|
|
|
def clear_memory(self):
|
|
"""Clear all memories."""
|
|
self.client.delete_collection(self.collection.name)
|
|
self.collection = self.client.get_or_create_collection(
|
|
name=self.collection.name
|
|
)
|
|
|
|
|
|
class GraphSwarm:
|
|
"""
|
|
Enhanced framework for creating and managing swarms of collaborative agents.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
agents: Union[
|
|
List[Agent], List[Tuple[Agent, List[str]]], None
|
|
] = None,
|
|
max_workers: Optional[int] = None,
|
|
swarm_name: str = "Collaborative Agent Swarm",
|
|
memory_collection: str = "swarm_memory",
|
|
):
|
|
"""Initialize GraphSwarm."""
|
|
self.graph = nx.DiGraph()
|
|
self.agents: Dict[str, Agent] = {}
|
|
self.dependencies: Dict[str, List[str]] = {}
|
|
self.executor = ThreadPoolExecutor(max_workers=max_workers)
|
|
self.swarm_name = swarm_name
|
|
self.memory_collection = memory_collection
|
|
self.memory = SwarmMemory(collection_name=memory_collection)
|
|
|
|
|
|
if agents:
|
|
self.initialize_agents(agents)
|
|
|
|
logger.info(f"Initialized GraphSwarm: {swarm_name}")
|
|
|
|
def initialize_agents(
|
|
self,
|
|
agents: Union[List[Agent], List[Tuple[Agent, List[str]]]],
|
|
):
|
|
"""Initialize agents and their dependencies."""
|
|
try:
|
|
# Handle list of Agents or (Agent, dependencies) tuples
|
|
for item in agents:
|
|
if isinstance(item, tuple):
|
|
agent, dependencies = item
|
|
else:
|
|
agent, dependencies = item, []
|
|
|
|
if not isinstance(agent, Agent):
|
|
raise ValueError(
|
|
f"Expected Agent object, got {type(agent)}"
|
|
)
|
|
|
|
self.agents[agent.agent_name] = agent
|
|
self.dependencies[agent.agent_name] = dependencies
|
|
self.graph.add_node(agent.agent_name, agent=agent)
|
|
|
|
# Add dependencies
|
|
for dep in dependencies:
|
|
if dep not in self.agents:
|
|
raise ValueError(
|
|
f"Dependency {dep} not found for agent {agent.agent_name}"
|
|
)
|
|
self.graph.add_edge(dep, agent.agent_name)
|
|
|
|
self._validate_graph()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize agents: {str(e)}")
|
|
raise
|
|
|
|
def _validate_graph(self):
|
|
"""Validate the agent dependency graph."""
|
|
if not self.graph.nodes():
|
|
raise ValueError("No agents added to swarm")
|
|
|
|
if not nx.is_directed_acyclic_graph(self.graph):
|
|
cycles = list(nx.simple_cycles(self.graph))
|
|
raise ValueError(
|
|
f"Agent dependency graph contains cycles: {cycles}"
|
|
)
|
|
|
|
def _get_agent_role_description(self, agent_name: str) -> str:
|
|
"""Generate a description of the agent's role in the swarm."""
|
|
predecessors = list(self.graph.predecessors(agent_name))
|
|
successors = list(self.graph.successors(agent_name))
|
|
position = (
|
|
"initial"
|
|
if not predecessors
|
|
else ("final" if not successors else "intermediate")
|
|
)
|
|
|
|
role = f"""You are {agent_name}, a specialized agent in the {self.swarm_name}.
|
|
Position: {position} agent in the workflow
|
|
|
|
Your relationships:"""
|
|
|
|
if predecessors:
|
|
role += (
|
|
f"\nYou receive input from: {', '.join(predecessors)}"
|
|
)
|
|
if successors:
|
|
role += f"\nYour output will be used by: {', '.join(successors)}"
|
|
|
|
return role
|
|
|
|
def _generate_workflow_context(self) -> str:
|
|
"""Generate a description of the entire workflow."""
|
|
execution_order = list(nx.topological_sort(self.graph))
|
|
|
|
workflow = f"""Workflow Overview of {self.swarm_name}:
|
|
|
|
Processing Order:
|
|
{' -> '.join(execution_order)}
|
|
|
|
Agent Roles:
|
|
"""
|
|
|
|
for agent_name in execution_order:
|
|
predecessors = list(self.graph.predecessors(agent_name))
|
|
successors = list(self.graph.successors(agent_name))
|
|
|
|
workflow += f"\n\n{agent_name}:"
|
|
if predecessors:
|
|
workflow += (
|
|
f"\n- Receives from: {', '.join(predecessors)}"
|
|
)
|
|
if successors:
|
|
workflow += f"\n- Sends to: {', '.join(successors)}"
|
|
if not predecessors and not successors:
|
|
workflow += "\n- Independent agent"
|
|
|
|
return workflow
|
|
|
|
def _build_agent_prompt(
|
|
self, agent_name: str, task: str, context: Dict = None
|
|
) -> str:
|
|
"""Build a comprehensive prompt for the agent including role and context."""
|
|
prompt_parts = [
|
|
self._get_agent_role_description(agent_name),
|
|
"\nWorkflow Context:",
|
|
self._generate_workflow_context(),
|
|
"\nYour Task:",
|
|
task,
|
|
]
|
|
|
|
if context:
|
|
prompt_parts.extend(
|
|
["\nContext from Previous Agents:", str(context)]
|
|
)
|
|
|
|
prompt_parts.extend(
|
|
[
|
|
"\nInstructions:",
|
|
"1. Process the task according to your role",
|
|
"2. Consider the input from previous agents when available",
|
|
"3. Provide clear, structured output",
|
|
"4. Remember that your output will be used by subsequent agents",
|
|
"\nResponse Guidelines:",
|
|
"- Provide clear, well-organized output",
|
|
"- Include relevant details and insights",
|
|
"- Highlight key findings",
|
|
"- Flag any uncertainties or issues",
|
|
]
|
|
)
|
|
|
|
return "\n".join(prompt_parts)
|
|
|
|
async def _execute_agent(
|
|
self, agent_name: str, task: str, context: Dict = None
|
|
) -> AgentOutput:
|
|
"""Execute a single agent."""
|
|
start_time = time.time()
|
|
agent = self.agents[agent_name]
|
|
|
|
try:
|
|
# Build comprehensive prompt
|
|
full_prompt = self._build_agent_prompt(
|
|
agent_name, task, context
|
|
)
|
|
logger.debug(f"Prompt for {agent_name}:\n{full_prompt}")
|
|
|
|
# Execute agent
|
|
output = await asyncio.to_thread(agent.run, full_prompt)
|
|
|
|
return AgentOutput(
|
|
agent_name=agent_name,
|
|
output=output,
|
|
execution_time=time.time() - start_time,
|
|
metadata={
|
|
"task": task,
|
|
"context": context,
|
|
"position_in_workflow": list(
|
|
nx.topological_sort(self.graph)
|
|
).index(agent_name),
|
|
},
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error executing agent {agent_name}: {str(e)}"
|
|
)
|
|
return AgentOutput(
|
|
agent_name=agent_name,
|
|
output=None,
|
|
execution_time=time.time() - start_time,
|
|
error=str(e),
|
|
metadata={"task": task},
|
|
)
|
|
|
|
async def execute(self, task: str) -> SwarmOutput:
|
|
"""
|
|
Execute the entire swarm of agents with memory integration.
|
|
|
|
Args:
|
|
task: Initial task to execute
|
|
|
|
Returns:
|
|
SwarmOutput: Structured output from all agents
|
|
"""
|
|
start_time = time.time()
|
|
outputs = {}
|
|
success = True
|
|
error = None
|
|
|
|
try:
|
|
# Get similar past executions
|
|
similar_executions = self.memory.get_similar_executions(
|
|
task, limit=3
|
|
)
|
|
optimal_sequence = self.memory.get_optimal_sequence(task)
|
|
|
|
# Get base execution order
|
|
base_execution_order = list(
|
|
nx.topological_sort(self.graph)
|
|
)
|
|
|
|
# Determine final execution order
|
|
if optimal_sequence and all(
|
|
agent in base_execution_order
|
|
for agent in optimal_sequence
|
|
):
|
|
logger.info(
|
|
f"Using optimal sequence from memory: {optimal_sequence}"
|
|
)
|
|
execution_order = optimal_sequence
|
|
else:
|
|
execution_order = base_execution_order
|
|
|
|
# Get historical context if available
|
|
historical_context = {}
|
|
if similar_executions:
|
|
best_execution = similar_executions[0]
|
|
if best_execution["success"]:
|
|
historical_context = {
|
|
"similar_task": best_execution["task"],
|
|
"previous_outputs": best_execution["outputs"],
|
|
"execution_time": best_execution[
|
|
"execution_time"
|
|
],
|
|
"success_patterns": self._extract_success_patterns(
|
|
similar_executions
|
|
),
|
|
}
|
|
|
|
# Execute agents in order
|
|
for agent_name in execution_order:
|
|
try:
|
|
# Get context from dependencies and history
|
|
agent_context = {
|
|
"dependencies": {
|
|
dep: outputs[dep].output
|
|
for dep in self.graph.predecessors(
|
|
agent_name
|
|
)
|
|
if dep in outputs
|
|
},
|
|
"historical": historical_context,
|
|
"position": execution_order.index(agent_name),
|
|
"total_agents": len(execution_order),
|
|
}
|
|
|
|
# Execute agent with enhanced context
|
|
output = await self._execute_agent(
|
|
agent_name, task, agent_context
|
|
)
|
|
outputs[agent_name] = output
|
|
|
|
# Update historical context with current execution
|
|
if output.output:
|
|
historical_context.update(
|
|
{
|
|
f"current_{agent_name}_output": output.output
|
|
}
|
|
)
|
|
|
|
# Check for errors
|
|
if output.error:
|
|
success = False
|
|
error = f"Agent {agent_name} failed: {output.error}"
|
|
|
|
# Try to recover using memory
|
|
if similar_executions:
|
|
recovery_output = self._attempt_recovery(
|
|
agent_name, task, similar_executions
|
|
)
|
|
if recovery_output:
|
|
outputs[agent_name] = recovery_output
|
|
success = True
|
|
error = None
|
|
continue
|
|
break
|
|
|
|
except Exception as agent_error:
|
|
logger.error(
|
|
f"Error executing agent {agent_name}: {str(agent_error)}"
|
|
)
|
|
success = False
|
|
error = f"Agent {agent_name} failed: {str(agent_error)}"
|
|
break
|
|
|
|
# Create result
|
|
result = SwarmOutput(
|
|
outputs=outputs,
|
|
execution_time=time.time() - start_time,
|
|
success=success,
|
|
error=error,
|
|
metadata={
|
|
"task": task,
|
|
"used_optimal_sequence": optimal_sequence
|
|
is not None,
|
|
"similar_executions_found": len(
|
|
similar_executions
|
|
),
|
|
"execution_order": execution_order,
|
|
"historical_context_used": bool(
|
|
historical_context
|
|
),
|
|
},
|
|
)
|
|
|
|
# Store execution in memory
|
|
await self._store_execution_async(task, result)
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
logger.error(f"Swarm execution failed: {str(e)}")
|
|
return SwarmOutput(
|
|
outputs=outputs,
|
|
execution_time=time.time() - start_time,
|
|
success=False,
|
|
error=str(e),
|
|
metadata={"task": task},
|
|
)
|
|
|
|
def run(self, task: str) -> SwarmOutput:
|
|
"""Synchronous interface to execute the swarm."""
|
|
return asyncio.run(self.execute(task))
|
|
|
|
def _extract_success_patterns(
|
|
self, similar_executions: List[Dict]
|
|
) -> Dict:
|
|
"""Extract success patterns from similar executions."""
|
|
patterns = {}
|
|
successful_execs = [
|
|
ex for ex in similar_executions if ex["success"]
|
|
]
|
|
|
|
if successful_execs:
|
|
patterns = {
|
|
"common_sequences": self._find_common_sequences(
|
|
successful_execs
|
|
),
|
|
"avg_execution_time": sum(
|
|
ex["execution_time"] for ex in successful_execs
|
|
)
|
|
/ len(successful_execs),
|
|
"successful_strategies": self._extract_strategies(
|
|
successful_execs
|
|
),
|
|
}
|
|
|
|
return patterns
|
|
|
|
def _attempt_recovery(
|
|
self,
|
|
failed_agent: str,
|
|
task: str,
|
|
similar_executions: List[Dict],
|
|
) -> Optional[AgentOutput]:
|
|
"""Attempt to recover from failure using memory."""
|
|
for execution in similar_executions:
|
|
if (
|
|
execution["success"]
|
|
and failed_agent in execution["outputs"]
|
|
):
|
|
historical_output = execution["outputs"][failed_agent]
|
|
|
|
return AgentOutput(
|
|
agent_name=failed_agent,
|
|
output=historical_output["output"],
|
|
execution_time=historical_output[
|
|
"execution_time"
|
|
],
|
|
metadata={
|
|
"recovered_from_memory": True,
|
|
"original_task": execution["task"],
|
|
},
|
|
)
|
|
return None
|
|
|
|
async def _store_execution_async(
|
|
self, task: str, result: SwarmOutput
|
|
):
|
|
"""Asynchronously store execution in memory."""
|
|
try:
|
|
await asyncio.to_thread(
|
|
self.memory.store_execution, task, result
|
|
)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Failed to store execution in memory: {str(e)}"
|
|
)
|
|
|
|
def add_agent(self, agent: Agent, dependencies: List[str] = None):
|
|
"""Add a new agent to the swarm."""
|
|
dependencies = dependencies or []
|
|
self.agents[agent.agent_name] = agent
|
|
self.dependencies[agent.agent_name] = dependencies
|
|
self.graph.add_node(agent.agent_name, agent=agent)
|
|
|
|
for dep in dependencies:
|
|
if dep not in self.agents:
|
|
raise ValueError(f"Dependency {dep} not found")
|
|
self.graph.add_edge(dep, agent.agent_name)
|
|
|
|
self._validate_graph()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
# Create agents
|
|
data_collector = Agent(
|
|
agent_name="Market-Data-Collector",
|
|
model_name="gpt-4o-mini",
|
|
max_loops=1,
|
|
streaming_on=True,
|
|
)
|
|
|
|
trend_analyzer = Agent(
|
|
agent_name="Market-Trend-Analyzer",
|
|
model_name="gpt-4o-mini",
|
|
max_loops=1,
|
|
streaming_on=True,
|
|
)
|
|
|
|
report_generator = Agent(
|
|
agent_name="Investment-Report-Generator",
|
|
model_name="gpt-4o-mini",
|
|
max_loops=1,
|
|
streaming_on=True,
|
|
)
|
|
|
|
# Create swarm
|
|
swarm = GraphSwarm(
|
|
agents=[
|
|
(data_collector, []),
|
|
(trend_analyzer, ["Market-Data-Collector"]),
|
|
(report_generator, ["Market-Trend-Analyzer"]),
|
|
],
|
|
swarm_name="Market Analysis Intelligence Network",
|
|
)
|
|
|
|
# Run the swarm
|
|
result = swarm.run(
|
|
"Analyze current market trends for tech stocks and provide investment recommendations"
|
|
)
|
|
|
|
# Print results
|
|
print(f"Execution success: {result.success}")
|
|
print(f"Total time: {result.execution_time:.2f} seconds")
|
|
|
|
for agent_name, output in result.outputs.items():
|
|
print(f"\nAgent: {agent_name}")
|
|
print(f"Output: {output.output}")
|
|
if output.error:
|
|
print(f"Error: {output.error}")
|
|
except Exception as error:
|
|
logger.error(error)
|
|
raise error
|