Kye Gomez
2 weeks ago
parent
b1834a62f8
commit
b53cd15c24
@ -1,276 +0,0 @@
|
||||
import asyncio
|
||||
import pulsar
|
||||
|
||||
from pulsar import ConsumerType
|
||||
from loguru import logger
|
||||
from swarms import Agent
|
||||
from typing import List, Dict, Any
|
||||
import json
|
||||
|
||||
|
||||
class ScalableAsyncAgentSwarm:
|
||||
"""
|
||||
A scalable, asynchronous swarm of agents leveraging Apache Pulsar for inter-agent communication.
|
||||
Provides load balancing, health monitoring, dead letter queues, and centralized logging.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pulsar_url: str,
|
||||
topic: str,
|
||||
dlq_topic: str,
|
||||
agents_config: List[Dict[str, Any]],
|
||||
):
|
||||
"""
|
||||
Initializes the async swarm with agents.
|
||||
|
||||
Args:
|
||||
pulsar_url (str): The URL of the Apache Pulsar broker.
|
||||
topic (str): The main topic for task distribution.
|
||||
dlq_topic (str): The Dead Letter Queue topic for failed messages.
|
||||
agents_config (List[Dict[str, Any]]): List of agent configurations with `name`, `description`, and `model_name`.
|
||||
"""
|
||||
self.pulsar_url = pulsar_url
|
||||
self.topic = topic
|
||||
self.dlq_topic = dlq_topic
|
||||
self.agents_config = agents_config
|
||||
self.client = pulsar.Client(pulsar_url)
|
||||
self.consumer = self.client.subscribe(
|
||||
topic,
|
||||
subscription_name="swarm-task-sub",
|
||||
consumer_type=ConsumerType.Shared,
|
||||
)
|
||||
self.dlq_producer = self.client.create_producer(dlq_topic)
|
||||
self.response_logger = []
|
||||
self.agents = [
|
||||
self.create_agent(config) for config in agents_config
|
||||
]
|
||||
self.agent_index = 0
|
||||
|
||||
logger.info(
|
||||
"Swarm initialized with agents: {}",
|
||||
[agent["name"] for agent in agents_config],
|
||||
)
|
||||
|
||||
def create_agent(
|
||||
self, agent_config: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Creates a new agent configuration with asynchronous capabilities.
|
||||
|
||||
Args:
|
||||
agent_config (Dict[str, Any]): Configuration dictionary with agent details.
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: A dictionary containing agent metadata and functionality.
|
||||
"""
|
||||
agent_name = agent_config["name"]
|
||||
description = agent_config["description"]
|
||||
model_name = agent_config.get("model_name", "gpt-4o-mini")
|
||||
|
||||
class AsyncAgent:
|
||||
"""
|
||||
An asynchronous agent that processes tasks and communicates via Apache Pulsar.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, name: str, description: str, model_name: str
|
||||
):
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.agent = Agent(
|
||||
agent_name=name,
|
||||
model_name=model_name,
|
||||
max_loops="auto",
|
||||
interactive=True,
|
||||
streaming_on=True,
|
||||
)
|
||||
logger.info(
|
||||
f"Initialized agent '{name}' - {description}"
|
||||
)
|
||||
|
||||
async def process_task(
|
||||
self, message: str
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Processes a single task using the agent.
|
||||
|
||||
Args:
|
||||
message (str): The task message.
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: JSON-formatted response.
|
||||
"""
|
||||
try:
|
||||
logger.info(
|
||||
f"Agent {self.name} processing task: {message}"
|
||||
)
|
||||
response = await asyncio.to_thread(
|
||||
self.agent.run, message
|
||||
)
|
||||
logger.info(f"Agent {self.name} completed task.")
|
||||
return {
|
||||
"agent_name": self.name,
|
||||
"response": response,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Agent {self.name} encountered an error: {e}"
|
||||
)
|
||||
return {"agent_name": self.name, "error": str(e)}
|
||||
|
||||
return {
|
||||
"name": agent_name,
|
||||
"instance": AsyncAgent(
|
||||
agent_name, description, model_name
|
||||
),
|
||||
}
|
||||
|
||||
async def distribute_task(self, message: str):
|
||||
"""
|
||||
Distributes a task to the next available agent using round-robin.
|
||||
|
||||
Args:
|
||||
message (str): The task message.
|
||||
"""
|
||||
agent = self.agents[self.agent_index]
|
||||
self.agent_index = (self.agent_index + 1) % len(self.agents)
|
||||
|
||||
try:
|
||||
response = await agent["instance"].process_task(message)
|
||||
self.log_response(response)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error processing task by agent {agent['name']}: {e}"
|
||||
)
|
||||
self.send_to_dlq(message)
|
||||
|
||||
async def monitor_health(self):
|
||||
"""
|
||||
Periodically monitors the health of agents.
|
||||
"""
|
||||
while True:
|
||||
logger.info("Performing health check for all agents.")
|
||||
for agent in self.agents:
|
||||
logger.info(f"Agent {agent['name']} is online.")
|
||||
await asyncio.sleep(10)
|
||||
|
||||
def send_to_dlq(self, message: str):
|
||||
"""
|
||||
Sends a failed message to the Dead Letter Queue (DLQ).
|
||||
|
||||
Args:
|
||||
message (str): The message to send to the DLQ.
|
||||
"""
|
||||
try:
|
||||
self.dlq_producer.send(message.encode("utf-8"))
|
||||
logger.info("Message sent to Dead Letter Queue.")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send message to DLQ: {e}")
|
||||
|
||||
def log_response(self, response: Dict[str, Any]):
|
||||
"""
|
||||
Logs the response to a centralized list for later analysis.
|
||||
|
||||
Args:
|
||||
response (Dict[str, Any]): The agent's response.
|
||||
"""
|
||||
self.response_logger.append(response)
|
||||
logger.info(f"Response logged: {response}")
|
||||
|
||||
async def listen_and_distribute(self):
|
||||
"""
|
||||
Listens to the main Pulsar topic and distributes tasks to agents.
|
||||
"""
|
||||
while True:
|
||||
msg = self.consumer.receive()
|
||||
try:
|
||||
message = msg.data().decode("utf-8")
|
||||
logger.info(f"Received task: {message}")
|
||||
await self.distribute_task(message)
|
||||
self.consumer.acknowledge(msg)
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing message: {e}")
|
||||
self.send_to_dlq(msg.data().decode("utf-8"))
|
||||
self.consumer.negative_acknowledge(msg)
|
||||
|
||||
async def run(self):
|
||||
"""
|
||||
Runs the swarm asynchronously with health monitoring and task distribution.
|
||||
"""
|
||||
logger.info("Starting the async swarm...")
|
||||
task_listener = asyncio.create_task(
|
||||
self.listen_and_distribute()
|
||||
)
|
||||
health_monitor = asyncio.create_task(self.monitor_health())
|
||||
await asyncio.gather(task_listener, health_monitor)
|
||||
|
||||
def shutdown(self):
|
||||
"""
|
||||
Safely shuts down the swarm and logs all responses.
|
||||
"""
|
||||
logger.info("Shutting down the swarm...")
|
||||
self.client.close()
|
||||
with open("responses.json", "w") as f:
|
||||
json.dump(self.response_logger, f, indent=4)
|
||||
logger.info("Responses saved to 'responses.json'.")
|
||||
|
||||
|
||||
# from scalable_agent_swarm import ScalableAsyncAgentSwarm # Assuming your swarm class is saved here
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example Configuration
|
||||
PULSAR_URL = "pulsar://localhost:6650"
|
||||
TOPIC = "stock-analysis"
|
||||
DLQ_TOPIC = "stock-analysis-dlq"
|
||||
|
||||
# Agents configuration
|
||||
AGENTS_CONFIG = [
|
||||
{
|
||||
"name": "Stock-Analysis-Agent-1",
|
||||
"description": "Analyzes stock trends.",
|
||||
"model_name": "gpt-4o-mini",
|
||||
},
|
||||
{
|
||||
"name": "Stock-News-Agent",
|
||||
"description": "Summarizes stock news.",
|
||||
"model_name": "gpt-4o-mini",
|
||||
},
|
||||
{
|
||||
"name": "Tech-Trends-Agent",
|
||||
"description": "Tracks tech sector trends.",
|
||||
"model_name": "gpt-4o-mini",
|
||||
},
|
||||
]
|
||||
|
||||
# Tasks to send
|
||||
TASKS = [
|
||||
"Analyze the trend for tech stocks in Q4 2024",
|
||||
"Summarize the latest news on the S&P 500",
|
||||
"Identify the top-performing sectors in the stock market",
|
||||
"Provide a forecast for AI-related stocks for 2025",
|
||||
]
|
||||
|
||||
# Initialize and run the swarm
|
||||
swarm = ScalableAsyncAgentSwarm(
|
||||
PULSAR_URL, TOPIC, DLQ_TOPIC, AGENTS_CONFIG
|
||||
)
|
||||
try:
|
||||
# Run the swarm in the background
|
||||
swarm_task = asyncio.create_task(swarm.run())
|
||||
|
||||
# Send tasks to the topic
|
||||
client = pulsar.Client(PULSAR_URL)
|
||||
producer = client.create_producer(TOPIC)
|
||||
|
||||
for task in TASKS:
|
||||
producer.send(task.encode("utf-8"))
|
||||
print(f"Sent task: {task}")
|
||||
|
||||
producer.close()
|
||||
client.close()
|
||||
|
||||
# Keep the swarm running
|
||||
asyncio.run(swarm_task)
|
||||
except KeyboardInterrupt:
|
||||
swarm.shutdown()
|
||||
Loading…
Reference in new issue