Kye
9 months ago
parent
9cbbdc9a10
commit
5728bff63a
@ -0,0 +1,54 @@
|
|||||||
|
from pydantic import BaseModel, Field
|
||||||
|
from swarms import Anthropic
|
||||||
|
from swarms import Agent
|
||||||
|
|
||||||
|
|
||||||
|
# Initialize the schema for the person's information
|
||||||
|
class Schema(BaseModel):
|
||||||
|
name: str = Field(..., title="Name of the person")
|
||||||
|
agent: int = Field(..., title="Age of the person")
|
||||||
|
is_student: bool = Field(..., title="Whether the person is a student")
|
||||||
|
courses: list[str] = Field(
|
||||||
|
..., title="List of courses the person is taking"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Convert the schema to a JSON string
|
||||||
|
tool_schema = Schema(
|
||||||
|
name="Tool Name",
|
||||||
|
agent=1,
|
||||||
|
is_student=True,
|
||||||
|
courses=["Course1", "Course2"],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Define the task to generate a person's information
|
||||||
|
task = "Generate a person's information based on the following schema:"
|
||||||
|
|
||||||
|
# Initialize the agent
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="Person Information Generator",
|
||||||
|
system_prompt=(
|
||||||
|
"Generate a person's information based on the following schema:"
|
||||||
|
),
|
||||||
|
# Set the tool schema to the JSON string -- this is the key difference
|
||||||
|
tool_schema=tool_schema,
|
||||||
|
llm=Anthropic(),
|
||||||
|
max_loops=3,
|
||||||
|
autosave=True,
|
||||||
|
dashboard=False,
|
||||||
|
streaming_on=True,
|
||||||
|
verbose=True,
|
||||||
|
interactive=True,
|
||||||
|
# Set the output type to the tool schema which is a BaseModel
|
||||||
|
output_type=tool_schema, # or dict, or str
|
||||||
|
metadata_output_type="json",
|
||||||
|
# List of schemas that the agent can handle
|
||||||
|
list_tool_schemas = [tool_schema],
|
||||||
|
function_calling_format_type = "OpenAI",
|
||||||
|
function_calling_type = "json" # or soon yaml
|
||||||
|
)
|
||||||
|
|
||||||
|
# Run the agent to generate the person's information
|
||||||
|
generated_data = agent.run(task)
|
||||||
|
|
||||||
|
# Print the generated data
|
||||||
|
print(f"Generated data: {generated_data}")
|
||||||
@ -0,0 +1,98 @@
|
|||||||
|
swarms
|
||||||
|
|
||||||
|
pip install swarms
|
||||||
|
swarms is the most pythonic way of writing cognitive systems. Leveraging pydantic models as output schemas combined with langchain in the backend allows for a seamless integration of llms into your apps. It utilizes OpenAI Functions or LlamaCpp grammars (json-schema-mode) for efficient structured output. In the backend it compiles the swarms syntax into langchain runnables so you can easily invoke, stream or batch process your pipelines.
|
||||||
|
|
||||||
|
Open in GitHub Codespaces
|
||||||
|
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
from swarms import Anthropic
|
||||||
|
from swarms import Agent
|
||||||
|
|
||||||
|
|
||||||
|
# Initialize the schema for the person's information
|
||||||
|
class Schema(BaseModel):
|
||||||
|
name: str = Field(..., title="Name of the person")
|
||||||
|
agent: int = Field(..., title="Age of the person")
|
||||||
|
is_student: bool = Field(..., title="Whether the person is a student")
|
||||||
|
courses: list[str] = Field(
|
||||||
|
..., title="List of courses the person is taking"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Convert the schema to a JSON string
|
||||||
|
tool_schema = Schema(
|
||||||
|
name="Tool Name",
|
||||||
|
agent=1,
|
||||||
|
is_student=True,
|
||||||
|
courses=["Course1", "Course2"],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Define the task to generate a person's information
|
||||||
|
task = "Generate a person's information based on the following schema:"
|
||||||
|
|
||||||
|
# Initialize the agent
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="Person Information Generator",
|
||||||
|
system_prompt=(
|
||||||
|
"Generate a person's information based on the following schema:"
|
||||||
|
),
|
||||||
|
# Set the tool schema to the JSON string -- this is the key difference
|
||||||
|
tool_schema=tool_schema,
|
||||||
|
llm=Anthropic(),
|
||||||
|
max_loops=3,
|
||||||
|
autosave=True,
|
||||||
|
dashboard=False,
|
||||||
|
streaming_on=True,
|
||||||
|
verbose=True,
|
||||||
|
interactive=True,
|
||||||
|
# Set the output type to the tool schema which is a BaseModel
|
||||||
|
output_type=tool_schema, # or dict, or str
|
||||||
|
metadata_output_type="json",
|
||||||
|
# List of schemas that the agent can handle
|
||||||
|
list_tool_schemas = [tool_schema],
|
||||||
|
function_calling_format_type = "OpenAI",
|
||||||
|
function_calling_type = "json" # or soon yaml
|
||||||
|
)
|
||||||
|
|
||||||
|
# Run the agent to generate the person's information
|
||||||
|
generated_data = agent.run(task)
|
||||||
|
|
||||||
|
# Print the generated data
|
||||||
|
print(f"Generated data: {generated_data}")
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
Features
|
||||||
|
🐍 pythonic
|
||||||
|
🔀 easy swap between openai or local models
|
||||||
|
🔄 dynamic output types (pydantic models, or primitives)
|
||||||
|
👁️ vision llm support
|
||||||
|
🧠 langchain_core as backend
|
||||||
|
📝 jinja templating for prompts
|
||||||
|
🏗️ reliable structured output
|
||||||
|
🔁 auto retry parsing
|
||||||
|
🔧 langsmith support
|
||||||
|
🔄 sync, async, streaming, parallel, fallbacks
|
||||||
|
📦 gguf download from huggingface
|
||||||
|
✅ type hints for all functions and mypy support
|
||||||
|
🗣️ chat router component
|
||||||
|
🧩 composable with langchain LCEL
|
||||||
|
🛠️ easy error handling
|
||||||
|
🚦 enums and literal support
|
||||||
|
📐 custom parsing types
|
||||||
|
Documentation
|
||||||
|
Checkout the docs here 👈
|
||||||
|
|
||||||
|
Also highly recommend to try and run the examples in the ./examples folder.
|
||||||
|
|
||||||
|
Contribution
|
||||||
|
You want to contribute? Thanks, that's great! For more information checkout the Contributing Guide. Please run the dev setup to get started:
|
||||||
|
|
||||||
|
git clone https://github.com/kyegomez/swarms.git && cd swarms
|
||||||
|
|
||||||
|
./dev_setup.sh
|
||||||
|
About
|
||||||
|
⛓️ build cognitive systems, pythonic
|
||||||
@ -0,0 +1,7 @@
|
|||||||
|
from swarms import AutoSwarm
|
||||||
|
|
||||||
|
# Initialize the swarm
|
||||||
|
swarm = AutoSwarm("kyegomez/myswarm")
|
||||||
|
|
||||||
|
# Run the swarm
|
||||||
|
swarm.run("Generate a 10,000 word blog on health and wellness.")
|
||||||
@ -1,329 +0,0 @@
|
|||||||
import logging
|
|
||||||
import queue
|
|
||||||
import threading
|
|
||||||
from time import sleep
|
|
||||||
from typing import Callable, Dict, List, Optional
|
|
||||||
|
|
||||||
from termcolor import colored
|
|
||||||
|
|
||||||
from swarms.structs.agent import Agent
|
|
||||||
from swarms.structs.base_structure import BaseStructure
|
|
||||||
from swarms.utils.decorators import (
|
|
||||||
error_decorator,
|
|
||||||
log_decorator,
|
|
||||||
timing_decorator,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class AutoScaler(BaseStructure):
|
|
||||||
"""
|
|
||||||
AutoScaler class
|
|
||||||
|
|
||||||
The AutoScaler class is responsible for managing the agents pool
|
|
||||||
and the task queue. It also monitors the health of the agents and
|
|
||||||
scales the pool up or down based on the number of pending tasks
|
|
||||||
and the current load of the agents.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
initial_agents (Optional[int], optional): Initial number of
|
|
||||||
agents to start with. Defaults to 10.
|
|
||||||
scale_up_factor (int, optional): Factor by which to scale up
|
|
||||||
the agents pool. Defaults to 1.
|
|
||||||
idle_threshold (float, optional): Threshold for scaling down
|
|
||||||
the agents pool. Defaults to 0.2.
|
|
||||||
busy_threshold (float, optional): Threshold for scaling up
|
|
||||||
the agents pool. Defaults to 0.7.
|
|
||||||
agents (List[Agent], optional): List of agents to use in the
|
|
||||||
pool. Defaults to None.
|
|
||||||
autoscale (Optional[bool], optional): Whether to autoscale
|
|
||||||
the agents pool. Defaults to True.
|
|
||||||
min_agents (Optional[int], optional): Minimum number of
|
|
||||||
agents to keep in the pool. Defaults to 10.
|
|
||||||
max_agents (Optional[int], optional): Maximum number of
|
|
||||||
agents to keep in the pool. Defaults to 100.
|
|
||||||
custom_scale_strategy (Optional[Callable], optional): Custom
|
|
||||||
scaling strategy to use. Defaults to None.
|
|
||||||
|
|
||||||
Methods:
|
|
||||||
add_task: Add tasks to queue
|
|
||||||
scale_up: Add more agents
|
|
||||||
scale_down: scale down
|
|
||||||
run: Run agent the task on the agent id
|
|
||||||
monitor_and_scale: Monitor and scale
|
|
||||||
start: Start scaling
|
|
||||||
check_agent_health: Checks the health of each agent and
|
|
||||||
replaces unhealthy agents.
|
|
||||||
balance_load: Distributes tasks among agents based on their
|
|
||||||
current load.
|
|
||||||
set_scaling_strategy: Set a custom scaling strategy.
|
|
||||||
execute_scaling_strategy: Execute the custom scaling strategy
|
|
||||||
if defined.
|
|
||||||
report_agent_metrics: Collects and reports metrics from each
|
|
||||||
agent.
|
|
||||||
report: Reports the current state of the autoscaler.
|
|
||||||
print_dashboard: Prints a dashboard of the current state of
|
|
||||||
the autoscaler.
|
|
||||||
|
|
||||||
Examples:
|
|
||||||
>>> import os
|
|
||||||
>>> from dotenv import load_dotenv
|
|
||||||
>>> # Import the OpenAIChat model and the Agent struct
|
|
||||||
>>> from swarms.models import OpenAIChat
|
|
||||||
>>> from swarms.structs import Agent
|
|
||||||
>>> from swarms.structs.autoscaler import AutoScaler
|
|
||||||
>>> # Load the environment variables
|
|
||||||
>>> load_dotenv()
|
|
||||||
>>> # Get the API key from the environment
|
|
||||||
>>> api_key = os.environ.get("OPENAI_API_KEY")
|
|
||||||
>>> # Initialize the language model
|
|
||||||
>>> llm = OpenAIChat(
|
|
||||||
... temperature=0.5,
|
|
||||||
... openai_api_key=api_key,
|
|
||||||
... )
|
|
||||||
>>> ## Initialize the workflow
|
|
||||||
>>> agent = Agent(llm=llm, max_loops=1, dashboard=True)
|
|
||||||
>>> # Load the autoscaler
|
|
||||||
>>> autoscaler = AutoScaler(
|
|
||||||
... initial_agents=2,
|
|
||||||
... scale_up_factor=1,
|
|
||||||
... idle_threshold=0.2,
|
|
||||||
... busy_threshold=0.7,
|
|
||||||
... agents=[agent],
|
|
||||||
... autoscale=True,
|
|
||||||
... min_agents=1,
|
|
||||||
... max_agents=5,
|
|
||||||
... custom_scale_strategy=None,
|
|
||||||
... )
|
|
||||||
>>> print(autoscaler)
|
|
||||||
>>> # Run the workflow on a task
|
|
||||||
>>> out = autoscaler.run(agent.id, "Generate a 10,000 word blog on health and wellness.")
|
|
||||||
>>> print(out)
|
|
||||||
|
|
||||||
"""
|
|
||||||
|
|
||||||
@log_decorator
|
|
||||||
@error_decorator
|
|
||||||
@timing_decorator
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
initial_agents: Optional[int] = 10,
|
|
||||||
scale_up_factor: int = 1,
|
|
||||||
idle_threshold: float = 0.2,
|
|
||||||
busy_threshold: float = 0.7,
|
|
||||||
agents: List[Agent] = None,
|
|
||||||
autoscale: Optional[bool] = True,
|
|
||||||
min_agents: Optional[int] = 10,
|
|
||||||
max_agents: Optional[int] = 100,
|
|
||||||
custom_scale_strategy: Optional[Callable] = None,
|
|
||||||
*args,
|
|
||||||
**kwargs,
|
|
||||||
):
|
|
||||||
self.agents_pool = agents or [
|
|
||||||
agents[0]() for _ in range(initial_agents)
|
|
||||||
]
|
|
||||||
self.task_queue = queue.Queue()
|
|
||||||
self.scale_up_factor = scale_up_factor
|
|
||||||
self.idle_threshold = idle_threshold
|
|
||||||
self.busy_threshold = busy_threshold
|
|
||||||
self.lock = threading.Lock()
|
|
||||||
self.agents = agents
|
|
||||||
self.autoscale = autoscale
|
|
||||||
self.min_agents = min_agents
|
|
||||||
self.max_agents = max_agents
|
|
||||||
self.custom_scale_strategy = custom_scale_strategy
|
|
||||||
|
|
||||||
def add_task(self, task):
|
|
||||||
"""Add tasks to queue"""
|
|
||||||
try:
|
|
||||||
self.task_queue.put(task)
|
|
||||||
except Exception as error:
|
|
||||||
print(
|
|
||||||
f"Error adding task to queue: {error} try again with"
|
|
||||||
" a new task"
|
|
||||||
)
|
|
||||||
|
|
||||||
@log_decorator
|
|
||||||
@error_decorator
|
|
||||||
@timing_decorator
|
|
||||||
def scale_up(self):
|
|
||||||
"""Add more agents"""
|
|
||||||
try:
|
|
||||||
with self.lock:
|
|
||||||
new_agents_counts = (
|
|
||||||
len(self.agents_pool) * self.scale_up_factor
|
|
||||||
)
|
|
||||||
for _ in range(new_agents_counts):
|
|
||||||
self.agents_pool.append(self.agents[0]())
|
|
||||||
except Exception as error:
|
|
||||||
print(f"Error scaling up: {error} try again with a new task")
|
|
||||||
|
|
||||||
def scale_down(self):
|
|
||||||
"""scale down"""
|
|
||||||
try:
|
|
||||||
with self.lock:
|
|
||||||
if (
|
|
||||||
len(self.agents_pool) > 10
|
|
||||||
): # ensure minmum of 10 agents
|
|
||||||
del self.agents_pool[-1] # remove last agent
|
|
||||||
except Exception as error:
|
|
||||||
print(
|
|
||||||
f"Error scaling down: {error} try again with a new" " task"
|
|
||||||
)
|
|
||||||
|
|
||||||
def run(self, agent_id, task: Optional[str] = None, *args, **kwargs):
|
|
||||||
"""Run agent the task on the agent id
|
|
||||||
|
|
||||||
Args:
|
|
||||||
agent_id (_type_): _description_
|
|
||||||
task (str, optional): _description_. Defaults to None.
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
ValueError: _description_
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
_type_: _description_
|
|
||||||
"""
|
|
||||||
for agent in self.agents_pool:
|
|
||||||
if agent.id == agent_id:
|
|
||||||
return agent.run(task, *args, **kwargs)
|
|
||||||
raise ValueError(f"No agent found with ID {agent_id}")
|
|
||||||
|
|
||||||
@log_decorator
|
|
||||||
@error_decorator
|
|
||||||
@timing_decorator
|
|
||||||
def monitor_and_scale(self):
|
|
||||||
"""Monitor and scale"""
|
|
||||||
try:
|
|
||||||
while True:
|
|
||||||
sleep(60) # check minute
|
|
||||||
pending_tasks = self.task_queue.qsize()
|
|
||||||
active_agents = sum(
|
|
||||||
[1 for agent in self.agents_pool if agent.is_busy()]
|
|
||||||
)
|
|
||||||
|
|
||||||
if (
|
|
||||||
pending_tasks / len(self.agents_pool)
|
|
||||||
> self.busy_threshold
|
|
||||||
):
|
|
||||||
self.scale_up()
|
|
||||||
elif (
|
|
||||||
active_agents / len(self.agents_pool)
|
|
||||||
< self.idle_threshold
|
|
||||||
):
|
|
||||||
self.scale_down()
|
|
||||||
except Exception as error:
|
|
||||||
print(
|
|
||||||
f"Error monitoring and scaling: {error} try again"
|
|
||||||
" with a new task"
|
|
||||||
)
|
|
||||||
|
|
||||||
@log_decorator
|
|
||||||
@error_decorator
|
|
||||||
@timing_decorator
|
|
||||||
def start(self):
|
|
||||||
"""Start scaling"""
|
|
||||||
try:
|
|
||||||
monitor_thread = threading.Thread(
|
|
||||||
target=self.monitor_and_scale
|
|
||||||
)
|
|
||||||
monitor_thread.start()
|
|
||||||
|
|
||||||
while True:
|
|
||||||
task = self.task_queue.get()
|
|
||||||
if task:
|
|
||||||
available_agent = next(
|
|
||||||
agent for agent in self.agents_pool
|
|
||||||
)
|
|
||||||
if available_agent:
|
|
||||||
available_agent.run(task)
|
|
||||||
except Exception as error:
|
|
||||||
print(f"Error starting: {error} try again with a new task")
|
|
||||||
|
|
||||||
def check_agent_health(self):
|
|
||||||
"""Checks the health of each agent and replaces unhealthy agents."""
|
|
||||||
for i, agent in enumerate(self.agents_pool):
|
|
||||||
if not agent.is_healthy():
|
|
||||||
logging.warning(f"Replacing unhealthy agent at index {i}")
|
|
||||||
self.agents_pool[i] = self.agent()
|
|
||||||
|
|
||||||
def balance_load(self):
|
|
||||||
"""Distributes tasks among agents based on their current load."""
|
|
||||||
try:
|
|
||||||
while not self.task_queue.empty():
|
|
||||||
for agent in self.agents_pool:
|
|
||||||
if agent.can_accept_task():
|
|
||||||
task = self.task_queue.get()
|
|
||||||
agent.run(task)
|
|
||||||
except Exception as error:
|
|
||||||
print(
|
|
||||||
f"Error balancing load: {error} try again with a new"
|
|
||||||
" task"
|
|
||||||
)
|
|
||||||
|
|
||||||
def set_scaling_strategy(self, strategy: Callable[[int, int], int]):
|
|
||||||
"""Set a custom scaling strategy."""
|
|
||||||
self.custom_scale_strategy = strategy
|
|
||||||
|
|
||||||
def execute_scaling_strategy(self):
|
|
||||||
"""Execute the custom scaling strategy if defined."""
|
|
||||||
try:
|
|
||||||
if hasattr(self, "custom_scale_strategy"):
|
|
||||||
scale_amount = self.custom_scale_strategy(
|
|
||||||
self.task_queue.qsize(), len(self.agents_pool)
|
|
||||||
)
|
|
||||||
if scale_amount > 0:
|
|
||||||
for _ in range(scale_amount):
|
|
||||||
self.agents_pool.append(self.agent())
|
|
||||||
elif scale_amount < 0:
|
|
||||||
for _ in range(abs(scale_amount)):
|
|
||||||
if len(self.agents_pool) > 10:
|
|
||||||
del self.agents_pool[-1]
|
|
||||||
except Exception as error:
|
|
||||||
print(
|
|
||||||
f"Error executing scaling strategy: {error} try again"
|
|
||||||
" with a new task"
|
|
||||||
)
|
|
||||||
|
|
||||||
def report_agent_metrics(self) -> Dict[str, List[float]]:
|
|
||||||
"""Collects and reports metrics from each agent."""
|
|
||||||
metrics = {
|
|
||||||
"completion_time": [],
|
|
||||||
"success_rate": [],
|
|
||||||
"error_rate": [],
|
|
||||||
}
|
|
||||||
for agent in self.agents_pool:
|
|
||||||
agent_metrics = agent.get_metrics()
|
|
||||||
for key in metrics.keys():
|
|
||||||
metrics[key].append(agent_metrics.get(key, 0.0))
|
|
||||||
return metrics
|
|
||||||
|
|
||||||
def report(self):
|
|
||||||
"""Reports the current state of the autoscaler."""
|
|
||||||
self.check_agent_health()
|
|
||||||
self.balance_load()
|
|
||||||
self.execute_scaling_strategy()
|
|
||||||
metrics = self.report_agent_metrics()
|
|
||||||
print(metrics)
|
|
||||||
|
|
||||||
def print_dashboard(self):
|
|
||||||
"""Prints a dashboard of the current state of the autoscaler."""
|
|
||||||
print(
|
|
||||||
colored(
|
|
||||||
f"""
|
|
||||||
|
|
||||||
Autoscaler Dashboard
|
|
||||||
--------------------
|
|
||||||
Agents: {len(self.agents_pool)}
|
|
||||||
Initial Agents: {self.initial_agents}
|
|
||||||
self.scale_up_factor: {self.scale_up_factor}
|
|
||||||
self.idle_threshold: {self.idle_threshold}
|
|
||||||
self.busy_threshold: {self.busy_threshold}
|
|
||||||
self.task_queue.qsize(): {self.task_queue.qsize()}
|
|
||||||
self.task_queue.empty(): {self.task_queue.empty()}
|
|
||||||
self.task_queue.full(): {self.task_queue.full()}
|
|
||||||
self.task_queue.maxsize: {self.task_queue.maxsize}
|
|
||||||
|
|
||||||
""",
|
|
||||||
"cyan",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
@ -1,25 +0,0 @@
|
|||||||
from dataclasses import dataclass
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
|
||||||
class BlockDevice:
|
|
||||||
"""
|
|
||||||
Represents a block device.
|
|
||||||
|
|
||||||
Attributes:
|
|
||||||
device (str): The device name.
|
|
||||||
cluster (str): The cluster name.
|
|
||||||
description (str): A description of the block device.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, device: str, cluster: str, description: str):
|
|
||||||
self.device = device
|
|
||||||
self.cluster = cluster
|
|
||||||
self.description = description
|
|
||||||
|
|
||||||
def __str__(self):
|
|
||||||
return (
|
|
||||||
f"BlockDevice(device={self.device},"
|
|
||||||
f" cluster={self.cluster},"
|
|
||||||
f" description={self.description})"
|
|
||||||
)
|
|
||||||
@ -1,70 +0,0 @@
|
|||||||
import multiprocessing as mp
|
|
||||||
from typing import List, Optional
|
|
||||||
|
|
||||||
from swarms.structs.base_structure import BaseStructure
|
|
||||||
|
|
||||||
|
|
||||||
class LoadBalancer(BaseStructure):
|
|
||||||
"""
|
|
||||||
A load balancer class that distributes tasks among multiple workers.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
num_workers (int): The number of worker processes to create.
|
|
||||||
agents (Optional[List]): A list of agents to assign to the load balancer.
|
|
||||||
*args: Variable length argument list.
|
|
||||||
**kwargs: Arbitrary keyword arguments.
|
|
||||||
|
|
||||||
Attributes:
|
|
||||||
num_workers (int): The number of worker processes.
|
|
||||||
agents (Optional[List]): A list of agents assigned to the load balancer.
|
|
||||||
tasks (List): A list of tasks to be executed.
|
|
||||||
results (List): A list of results from the executed tasks.
|
|
||||||
workers (List): A list of worker processes.
|
|
||||||
|
|
||||||
Methods:
|
|
||||||
add_task: Add a task to the load balancer.
|
|
||||||
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
num_workers: int = 1,
|
|
||||||
agents: Optional[List] = None,
|
|
||||||
*args,
|
|
||||||
**kwargs,
|
|
||||||
):
|
|
||||||
super().__init__(*args, **kwargs)
|
|
||||||
self.num_workers = num_workers
|
|
||||||
self.agents = agents
|
|
||||||
self.tasks = []
|
|
||||||
self.results = []
|
|
||||||
self.workers = []
|
|
||||||
self._init_workers()
|
|
||||||
|
|
||||||
def _init_workers(self):
|
|
||||||
for i in range(self.num_workers):
|
|
||||||
worker = mp.Process(target=self._worker)
|
|
||||||
worker.start()
|
|
||||||
self.workers.append(worker)
|
|
||||||
|
|
||||||
def _worker(self):
|
|
||||||
while True:
|
|
||||||
task = self._get_task()
|
|
||||||
if task is None:
|
|
||||||
break
|
|
||||||
result = self._run_task(task)
|
|
||||||
self._add_result(result)
|
|
||||||
|
|
||||||
def _get_task(self):
|
|
||||||
if len(self.tasks) == 0:
|
|
||||||
return None
|
|
||||||
return self.tasks.pop(0)
|
|
||||||
|
|
||||||
def _run_task(self, task):
|
|
||||||
return task()
|
|
||||||
|
|
||||||
def _add_result(self, result):
|
|
||||||
self.results.append(result)
|
|
||||||
|
|
||||||
def add_task(self, task):
|
|
||||||
self.tasks.append(task)
|
|
||||||
@ -1,176 +0,0 @@
|
|||||||
from dataclasses import asdict
|
|
||||||
from typing import List
|
|
||||||
|
|
||||||
import networkx as nx
|
|
||||||
import redis
|
|
||||||
from redis.commands.graph import Graph, Node
|
|
||||||
|
|
||||||
from swarms.structs.agent import Agent
|
|
||||||
from swarms.structs.base_swarm import BaseSwarm
|
|
||||||
|
|
||||||
|
|
||||||
class SwarmRelationship:
|
|
||||||
JOINED = "joined"
|
|
||||||
|
|
||||||
|
|
||||||
class RedisSwarmRegistry(BaseSwarm):
|
|
||||||
"""
|
|
||||||
Initialize the SwarmRedisRegistry object.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
host (str): The hostname or IP address of the Redis server. Default is "localhost".
|
|
||||||
port (int): The port number of the Redis server. Default is 6379.
|
|
||||||
db: The Redis database number. Default is 0.
|
|
||||||
graph_name (str): The name of the RedisGraph graph. Default is "swarm_registry".
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
host: str = "localhost",
|
|
||||||
port: int = 6379,
|
|
||||||
db=0,
|
|
||||||
graph_name: str = "swarm_registry",
|
|
||||||
):
|
|
||||||
self.redis = redis.StrictRedis(
|
|
||||||
host=host, port=port, db=db, decode_responses=True
|
|
||||||
)
|
|
||||||
self.redis_graph = Graph(self.redis, graph_name)
|
|
||||||
self.graph = nx.DiGraph()
|
|
||||||
|
|
||||||
def _entity_to_node(self, entity: Agent | Agent) -> Node:
|
|
||||||
"""
|
|
||||||
Converts an Agent or Swarm object to a Node object.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
entity (Agent | Agent): The Agent or Swarm object to convert.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Node: The converted Node object.
|
|
||||||
"""
|
|
||||||
return Node(
|
|
||||||
node_id=entity.id,
|
|
||||||
alias=entity.agent_name,
|
|
||||||
label=entity.agent_description,
|
|
||||||
properties=asdict(entity),
|
|
||||||
)
|
|
||||||
|
|
||||||
def _add_node(self, node: Agent | Agent):
|
|
||||||
"""
|
|
||||||
Adds a node to the graph.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
node (Agent | Agent): The Agent or Swarm node to add.
|
|
||||||
"""
|
|
||||||
self.graph.add_node(node.id)
|
|
||||||
if isinstance(node, Agent):
|
|
||||||
self.add_swarm_entry(node)
|
|
||||||
elif isinstance(node, Agent):
|
|
||||||
self.add_agent_entry(node)
|
|
||||||
|
|
||||||
def _add_edge(self, from_node: Node, to_node: Node, relationship):
|
|
||||||
"""
|
|
||||||
Adds an edge between two nodes in the graph.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
from_node (Node): The source node of the edge.
|
|
||||||
to_node (Node): The target node of the edge.
|
|
||||||
relationship: The relationship type between the nodes.
|
|
||||||
"""
|
|
||||||
match_query = (
|
|
||||||
f"MATCH (a:{from_node.label}),(b:{to_node.label}) WHERE"
|
|
||||||
f" a.id = {from_node.id} AND b.id = {to_node.id}"
|
|
||||||
)
|
|
||||||
|
|
||||||
query = f"""
|
|
||||||
{match_query}
|
|
||||||
CREATE (a)-[r:joined]->(b) RETURN r
|
|
||||||
""".replace(
|
|
||||||
"\n", ""
|
|
||||||
)
|
|
||||||
|
|
||||||
self.redis_graph.query(query)
|
|
||||||
|
|
||||||
def add_swarm_entry(self, swarm: Agent):
|
|
||||||
"""
|
|
||||||
Adds a swarm entry to the graph.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
swarm (Agent): The swarm object to add.
|
|
||||||
"""
|
|
||||||
node = self._entity_to_node(swarm)
|
|
||||||
self._persist_node(node)
|
|
||||||
|
|
||||||
def add_agent_entry(self, agent: Agent):
|
|
||||||
"""
|
|
||||||
Adds an agent entry to the graph.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
agent (Agent): The agent object to add.
|
|
||||||
"""
|
|
||||||
node = self._entity_to_node(agent)
|
|
||||||
self._persist_node(node)
|
|
||||||
|
|
||||||
def join_swarm(
|
|
||||||
self,
|
|
||||||
from_entity: Agent | Agent,
|
|
||||||
to_entity: Agent,
|
|
||||||
):
|
|
||||||
"""
|
|
||||||
Adds an edge between two nodes in the graph.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
from_entity (Agent | Agent): The source entity of the edge.
|
|
||||||
to_entity (Agent): The target entity of the edge.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Any: The result of adding the edge.
|
|
||||||
"""
|
|
||||||
from_node = self._entity_to_node(from_entity)
|
|
||||||
to_node = self._entity_to_node(to_entity)
|
|
||||||
|
|
||||||
return self._add_edge(from_node, to_node, SwarmRelationship.JOINED)
|
|
||||||
|
|
||||||
def _persist_node(self, node: Node):
|
|
||||||
"""
|
|
||||||
Persists a node in the graph.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
node (Node): The node to persist.
|
|
||||||
"""
|
|
||||||
query = f"CREATE {node}"
|
|
||||||
self.redis_graph.query(query)
|
|
||||||
|
|
||||||
def retrieve_swarm_information(self, swarm_id: int) -> Agent:
|
|
||||||
"""
|
|
||||||
Retrieves swarm information from the registry.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
swarm_id (int): The ID of the swarm to retrieve.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Agent: The retrieved swarm information as an Agent object.
|
|
||||||
"""
|
|
||||||
swarm_key = f"swarm:{swarm_id}"
|
|
||||||
swarm_data = self.redis.hgetall(swarm_key)
|
|
||||||
if swarm_data:
|
|
||||||
# Parse the swarm_data and return an instance of AgentBase
|
|
||||||
# You can use the retrieved data to populate the AgentBase attributes
|
|
||||||
|
|
||||||
return Agent(**swarm_data)
|
|
||||||
return None
|
|
||||||
|
|
||||||
def retrieve_joined_agents(self) -> List[Agent]:
|
|
||||||
"""
|
|
||||||
Retrieves a list of joined agents from the registry.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
List[Agent]: The retrieved joined agents as a list of Agent objects.
|
|
||||||
"""
|
|
||||||
agent_data = self.redis_graph.query(
|
|
||||||
"MATCH (a:agent)-[:joined]->(b:manager) RETURN a"
|
|
||||||
)
|
|
||||||
if agent_data:
|
|
||||||
# Parse the agent_data and return an instance of AgentBase
|
|
||||||
# You can use the retrieved data to populate the AgentBase attributes
|
|
||||||
return [Agent(**agent_data) for agent_data in agent_data]
|
|
||||||
return None
|
|
||||||
@ -0,0 +1,134 @@
|
|||||||
|
from typing import Any, Optional, List
|
||||||
|
|
||||||
|
from docstring_parser import parse
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
|
||||||
|
def _remove_a_key(d: dict, remove_key: str) -> None:
|
||||||
|
"""Remove a key from a dictionary recursively"""
|
||||||
|
if isinstance(d, dict):
|
||||||
|
for key in list(d.keys()):
|
||||||
|
if key == remove_key and "type" in d.keys():
|
||||||
|
del d[key]
|
||||||
|
else:
|
||||||
|
_remove_a_key(d[key], remove_key)
|
||||||
|
|
||||||
|
|
||||||
|
def pydantic_to_functions(
|
||||||
|
pydantic_type: type[BaseModel],
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Convert a Pydantic model to a dictionary representation of functions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pydantic_type (type[BaseModel]): The Pydantic model type to convert.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[str, Any]: A dictionary representation of the functions.
|
||||||
|
|
||||||
|
"""
|
||||||
|
schema = pydantic_type.model_json_schema()
|
||||||
|
|
||||||
|
docstring = parse(pydantic_type.__doc__ or "")
|
||||||
|
parameters = {
|
||||||
|
k: v
|
||||||
|
for k, v in schema.items()
|
||||||
|
if k not in ("title", "description")
|
||||||
|
}
|
||||||
|
|
||||||
|
for param in docstring.params:
|
||||||
|
if (name := param.arg_name) in parameters["properties"] and (
|
||||||
|
description := param.description
|
||||||
|
):
|
||||||
|
if "description" not in parameters["properties"][name]:
|
||||||
|
parameters["properties"][name]["description"] = description
|
||||||
|
|
||||||
|
parameters["type"] = "object"
|
||||||
|
|
||||||
|
if "description" not in schema:
|
||||||
|
if docstring.short_description:
|
||||||
|
schema["description"] = docstring.short_description
|
||||||
|
else:
|
||||||
|
schema["description"] = (
|
||||||
|
f"Correctly extracted `{pydantic_type.__class__.__name__.lower()}` with all "
|
||||||
|
f"the required parameters with correct types"
|
||||||
|
)
|
||||||
|
|
||||||
|
_remove_a_key(parameters, "title")
|
||||||
|
_remove_a_key(parameters, "additionalProperties")
|
||||||
|
|
||||||
|
return {
|
||||||
|
"function_call": {
|
||||||
|
"name": pydantic_type.__class__.__name__.lower(),
|
||||||
|
},
|
||||||
|
"functions": [
|
||||||
|
{
|
||||||
|
"name": pydantic_type.__class__.__name__.lower(),
|
||||||
|
"description": schema["description"],
|
||||||
|
"parameters": parameters,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def multi_pydantic_to_functions(
|
||||||
|
pydantic_types: List[BaseModel] = None
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Converts multiple Pydantic types to a dictionary of functions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pydantic_types (List[BaseModel]]): A list of Pydantic types to convert.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[str, Any]: A dictionary containing the converted functions.
|
||||||
|
|
||||||
|
"""
|
||||||
|
functions: list[dict[str, Any]] = [
|
||||||
|
pydantic_to_functions(pydantic_type)["functions"][0]
|
||||||
|
for pydantic_type in pydantic_types
|
||||||
|
]
|
||||||
|
|
||||||
|
return {
|
||||||
|
"function_call": "auto",
|
||||||
|
"functions": functions,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def function_to_str(function: dict[str, Any]) -> str:
|
||||||
|
"""
|
||||||
|
Convert a function dictionary to a string representation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
function (dict[str, Any]): The function dictionary to convert.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: The string representation of the function.
|
||||||
|
|
||||||
|
"""
|
||||||
|
function_str = f"Function: {function['name']}\n"
|
||||||
|
function_str += f"Description: {function['description']}\n"
|
||||||
|
function_str += "Parameters:\n"
|
||||||
|
|
||||||
|
for param, details in function["parameters"]["properties"].items():
|
||||||
|
function_str += f" {param} ({details['type']}): {details.get('description', '')}\n"
|
||||||
|
|
||||||
|
return function_str
|
||||||
|
|
||||||
|
|
||||||
|
def functions_to_str(functions: list[dict[str, Any]]) -> str:
|
||||||
|
"""
|
||||||
|
Convert a list of function dictionaries to a string representation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
functions (list[dict[str, Any]]): The list of function dictionaries to convert.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: The string representation of the functions.
|
||||||
|
|
||||||
|
"""
|
||||||
|
functions_str = ""
|
||||||
|
for function in functions:
|
||||||
|
functions_str += function_to_str(function) + "\n"
|
||||||
|
|
||||||
|
return functions_str
|
||||||
Loading…
Reference in new issue