ai-mirow-projects
/
swarms
mirror of https://github.com/kyegomez/swarms
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[6.5.7]

Browse Source
668
Kye Gomez 1 month ago
parent 40a78c4d06
commit 3b99f3db17
38 changed files with 2189 additions and 1966 deletions
Show Stats Download Patch File Download Diff File

16
api/agent_api.py
Unescape View File

@ -60,7 +60,7 @@ class AgentConfig(BaseModel):
..., description="System prompt for the agent" ..., description="System prompt for the agent"
) )
model_name: str = Field( model_name: str = Field(
default="gpt-4", description="Model name to use" default="gpt-4o-mini", description="Model name to use"
) )
temperature: float = Field( temperature: float = Field(
default=0.1, default=0.1,
@ -102,6 +102,14 @@ class AgentConfig(BaseModel):
default_factory=list, default_factory=list,
description="Tags for categorizing the agent", description="Tags for categorizing the agent",
) )
auto_generate_prompt: bool = Field(
default_factory=bool,
description="Auto generate a prompt based on the input",
)
max_tokens: int = Field(
default_factory=int,
description="The number of max output tokens",
)
class AgentUpdate(BaseModel): class AgentUpdate(BaseModel):
@ -197,9 +205,9 @@ class AgentStore:
user_name=config.user_name, user_name=config.user_name,
retry_attempts=config.retry_attempts, retry_attempts=config.retry_attempts,
context_length=config.context_length, context_length=config.context_length,
return_step_meta=True,
output_type="str", output_type="str",
streaming_on=config.streaming_on, auto_generate_prompt=config.auto_generate_prompt,
max_tokens=config.max_tokens,
) )
agent_id = uuid4() agent_id = uuid4()
@ -441,6 +449,8 @@ class AgentStore:
"agent_name": agent.agent_name, "agent_name": agent.agent_name,
"model_name": agent.llm.model_name, "model_name": agent.llm.model_name,
"temperature": agent.llm.temperature, "temperature": agent.llm.temperature,
"max_loops": agent.max_loops,
"context_window": agent.context_length,
}, },
timestamp=datetime.utcnow(), timestamp=datetime.utcnow(),
processing_time=processing_time, processing_time=processing_time,

112
api/test_api.py
Unescape View File

@ -0,0 +1,112 @@
import requests
import json
from time import sleep
BASE_URL = "http://api.swarms.ai:8000"
def make_request(method, endpoint, data=None):
"""Helper function to make requests with error handling"""
url = f"{BASE_URL}{endpoint}"
try:
if method == "GET":
response = requests.get(url)
elif method == "POST":
response = requests.post(url, json=data)
elif method == "DELETE":
response = requests.delete(url)
response.raise_for_status()
return response.json()
except requests.exceptions.RequestException as e:
print(
f"Error making {method} request to {endpoint}: {str(e)}"
)
if hasattr(e.response, "text"):
print(f"Response text: {e.response.text}")
return None
def create_agent():
"""Create a test agent"""
data = {
"agent_name": "test_agent",
"model_name": "gpt-4",
"system_prompt": "You are a helpful assistant",
"description": "Test agent",
"temperature": 0.7,
"max_loops": 1,
"tags": ["test"],
}
return make_request("POST", "/v1/agent", data)
def list_agents():
"""List all agents"""
return make_request("GET", "/v1/agents")
def test_completion(agent_id):
"""Test a completion with the agent"""
data = {
"prompt": "Say hello!",
"agent_id": agent_id,
"max_tokens": 100,
}
return make_request("POST", "/v1/agent/completions", data)
def get_agent_metrics(agent_id):
"""Get metrics for an agent"""
return make_request("GET", f"/v1/agent/{agent_id}/metrics")
def delete_agent(agent_id):
"""Delete an agent"""
return make_request("DELETE", f"/v1/agent/{agent_id}")
def run_tests():
print("Starting API tests...")
# Create an agent
print("\n1. Creating agent...")
agent_response = create_agent()
if not agent_response:
print("Failed to create agent")
return
agent_id = agent_response.get("agent_id")
print(f"Created agent with ID: {agent_id}")
# Give the server a moment to process
sleep(2)
# List agents
print("\n2. Listing agents...")
agents = list_agents()
print(f"Found {len(agents)} agents")
# Test completion
if agent_id:
print("\n3. Testing completion...")
completion = test_completion(agent_id)
if completion:
print(
f"Completion response: {completion.get('response')}"
)
print("\n4. Getting agent metrics...")
metrics = get_agent_metrics(agent_id)
if metrics:
print(f"Agent metrics: {json.dumps(metrics, indent=2)}")
# Clean up
# print("\n5. Cleaning up - deleting agent...")
# delete_result = delete_agent(agent_id)
# if delete_result:
# print("Successfully deleted agent")
if __name__ == "__main__":
run_tests()

898
byte.py
Unescape View File

@ -1,898 +0,0 @@
from enum import Enum
from typing import Union, Optional
import io
from PIL import Image
import numpy as np
import torch
import struct
from enum import auto
from typing import List, Dict, Tuple
import wave
from dataclasses import dataclass
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger
from einops import rearrange
from torch import Tensor
@dataclass
class ModelConfig:
"""Configuration for the enhanced BytePredictor model."""
vocab_size: int = 256 # Standard byte range
hidden_size: int = 1024
num_layers: int = 12
num_key_value_heads: int = 8 # For multi-query attention
num_query_heads: int = 32 # More query heads than kv heads
dropout: float = 0.1
max_sequence_length: int = 8192
rope_theta: float = 10000.0
layer_norm_eps: float = 1e-5
vocab_parallel: bool = False
qk_norm: bool = True
qk_norm_scale: float = None
attention_bias: bool = False
class MultiQueryAttention(nn.Module):
"""Fixed Multi-Query Attention implementation."""
def __init__(self, config: ModelConfig):
super().__init__()
self.hidden_size = config.hidden_size
self.num_query_heads = config.num_query_heads
self.num_key_value_heads = config.num_key_value_heads
self.head_dim = config.hidden_size // config.num_query_heads
self.qk_scale = config.qk_norm_scale or (self.head_dim**-0.5)
self.q_proj = nn.Linear(
config.hidden_size, config.num_query_heads * self.head_dim
)
self.k_proj = nn.Linear(
config.hidden_size,
config.num_key_value_heads * self.head_dim,
)
self.v_proj = nn.Linear(
config.hidden_size,
config.num_key_value_heads * self.head_dim,
)
self.o_proj = nn.Linear(
config.num_query_heads * self.head_dim, config.hidden_size
)
self.qk_norm = config.qk_norm
if self.qk_norm:
self.q_norm = nn.LayerNorm(self.head_dim)
self.k_norm = nn.LayerNorm(self.head_dim)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
batch_size, seq_length, _ = hidden_states.shape
# Project and reshape
q = self.q_proj(hidden_states)
k = self.k_proj(hidden_states)
v = self.v_proj(hidden_states)
# Reshape to [seq_len, batch, heads, head_dim]
q = q.view(
batch_size,
seq_length,
self.num_query_heads,
self.head_dim,
).permute(1, 0, 2, 3)
k = k.view(
batch_size,
seq_length,
self.num_key_value_heads,
self.head_dim,
).permute(1, 0, 2, 3)
v = v.view(
batch_size,
seq_length,
self.num_key_value_heads,
self.head_dim,
).permute(1, 0, 2, 3)
# Apply rotary embeddings
# q, k = self.rotary(q, k, seq_length)
# Apply QK normalization if enabled
if self.qk_norm:
q = self.q_norm(q)
k = self.k_norm(k)
# Handle MQA head expansion
if self.num_key_value_heads != self.num_query_heads:
k = k.repeat_interleave(
self.num_query_heads // self.num_key_value_heads,
dim=2,
)
v = v.repeat_interleave(
self.num_query_heads // self.num_key_value_heads,
dim=2,
)
# Compute attention
# Reshape for matmul: [batch, heads, seq_length, head_dim]
q = q.permute(1, 2, 0, 3)
k = k.permute(1, 2, 0, 3)
v = v.permute(1, 2, 0, 3)
attn_weights = (
torch.matmul(q, k.transpose(-2, -1)) * self.qk_scale
)
if attention_mask is not None:
attn_weights = attn_weights + attention_mask
attn_weights = F.softmax(attn_weights, dim=-1)
output = torch.matmul(attn_weights, v)
# Reshape back to [batch, seq_length, hidden_size]
output = (
output.transpose(1, 2)
.contiguous()
.view(batch_size, seq_length, -1)
)
output = self.o_proj(output)
return output
class EnhancedBytePredictor(nn.Module):
"""Enhanced byte prediction model with state-of-the-art features."""
def __init__(self, config: ModelConfig):
super().__init__()
self.config = config
# Token embeddings
self.tok_embeddings = nn.Embedding(
config.vocab_size, config.hidden_size
)
# Transformer layers
self.layers = nn.ModuleList(
[
nn.ModuleDict(
{
"attention": MultiQueryAttention(config),
"attention_norm": nn.LayerNorm(
config.hidden_size,
eps=config.layer_norm_eps,
),
"feed_forward": nn.Sequential(
nn.Linear(
config.hidden_size,
4 * config.hidden_size,
),
nn.GELU(),
nn.Linear(
4 * config.hidden_size,
config.hidden_size,
),
),
"feed_forward_norm": nn.LayerNorm(
config.hidden_size,
eps=config.layer_norm_eps,
),
}
)
for _ in range(config.num_layers)
]
)
self.norm = nn.LayerNorm(
config.hidden_size, eps=config.layer_norm_eps
)
self.output = nn.Linear(
config.hidden_size, config.vocab_size, bias=False
)
# Initialize weights
self.apply(self._init_weights)
def _init_weights(self, module: nn.Module) -> None:
"""Initialize weights with scaled normal distribution."""
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
def forward(
self,
input_ids: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Forward pass of the model.
Args:
input_ids: Tensor of shape (batch_size, sequence_length)
attention_mask: Optional attention mask
Returns:
Tensor of logits with shape (batch_size, sequence_length, vocab_size)
"""
hidden_states = self.tok_embeddings(input_ids)
# Create causal mask if needed
if attention_mask is None:
attention_mask = torch.triu(
torch.ones(
(input_ids.size(1), input_ids.size(1)),
device=input_ids.device,
dtype=torch.bool,
),
diagonal=1,
)
attention_mask = attention_mask.masked_fill(
attention_mask == 1, float("-inf")
)
# Apply transformer layers
for layer in self.layers:
# Attention block
hidden_states = hidden_states + layer["attention"](
layer["attention_norm"](hidden_states), attention_mask
)
# Feed-forward block
hidden_states = hidden_states + layer["feed_forward"](
layer["feed_forward_norm"](hidden_states)
)
hidden_states = self.norm(hidden_states)
logits = self.output(hidden_states)
return logits
def compute_loss(
self,
input_ids: torch.Tensor,
target_ids: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Compute cross entropy loss.
Args:
input_ids: Input token ids
target_ids: Target token ids
attention_mask: Optional attention mask
Returns:
Loss value
"""
logits = self(input_ids, attention_mask)
loss = F.cross_entropy(
rearrange(logits, "b s v -> (b s) v"),
rearrange(target_ids, "b s -> (b s)"),
)
return loss
@torch.no_grad()
def _generate(
self,
input_ids: torch.Tensor,
max_new_tokens: int = 100,
temperature: float = 1.0,
top_k: Optional[int] = None,
top_p: Optional[float] = None,
repetition_penalty: float = 1.0,
) -> torch.Tensor:
"""
Generate new tokens autoregressively.
Args:
input_ids: Starting sequence
max_new_tokens: Number of tokens to generate
temperature: Sampling temperature
top_k: K for top-k sampling
top_p: P for nucleus sampling
repetition_penalty: Penalty for repeating tokens
Returns:
Generated sequence
"""
batch_size, seq_length = input_ids.shape
generated = input_ids.clone()
for _ in range(max_new_tokens):
if generated.size(1) >= self.config.max_sequence_length:
break
# Forward pass
logits = self(generated)[:, -1, :]
# Apply temperature
logits = logits / temperature
# Apply repetition penalty
if repetition_penalty != 1.0:
for i in range(batch_size):
for token_id in set(generated[i].tolist()):
logits[i, token_id] /= repetition_penalty
# Apply top-k sampling
if top_k is not None:
indices_to_remove = (
logits
< torch.topk(logits, top_k)[0][..., -1, None]
)
logits[indices_to_remove] = float("-inf")
# Apply nucleus (top-p) sampling
if top_p is not None:
sorted_logits, sorted_indices = torch.sort(
logits, descending=True
)
cumulative_probs = torch.cumsum(
F.softmax(sorted_logits, dim=-1), dim=-1
)
# Remove tokens with cumulative probability above the threshold
sorted_indices_to_remove = cumulative_probs > top_p
sorted_indices_to_remove[..., 1:] = (
sorted_indices_to_remove[..., :-1].clone()
)
sorted_indices_to_remove[..., 0] = 0
indices_to_remove = torch.zeros_like(
logits, dtype=torch.bool
)
indices_to_remove.scatter_(
1, sorted_indices, sorted_indices_to_remove
)
logits[indices_to_remove] = float("-inf")
# Sample next token
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
# Append to sequence
generated = torch.cat([generated, next_token], dim=1)
return generated
def generate(
self,
input_ids: torch.Tensor,
max_new_tokens: int = 100,
temperature: float = 1.0,
top_k: Optional[int] = None,
top_p: Optional[float] = None,
repetition_penalty: float = 1.0,
):
tensor_data = self._generate(
input_ids=input_ids,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_k=top_k,
top_p=top_p,
repetition_penalty=repetition_penalty,
)
return tensor_to_data(tensor_data)
# import torch
# from typing import Optional
class DataType(Enum):
TEXT = "text"
IMAGE = "image"
AUDIO = "audio"
VIDEO = "video"
BINARY = "binary"
class ByteDetokenizer:
"""Utility class for converting model output bytes back to original data formats."""
@staticmethod
def tensor_to_bytes(tensor: torch.Tensor) -> bytes:
"""Convert model output tensor to bytes."""
# Convert logits/probabilities to byte values
if tensor.dim() > 1:
# If we have logits, convert to byte indices
byte_indices = tensor.argmax(dim=-1)
else:
byte_indices = tensor
# Convert to Python bytes
return bytes(
byte_indices.cpu().numpy().astype(np.uint8).tolist()
)
@staticmethod
def decode_text(byte_sequence: bytes) -> str:
"""Convert bytes to text."""
try:
return byte_sequence.decode("utf-8")
except UnicodeDecodeError:
# Try with error handling
return byte_sequence.decode("utf-8", errors="replace")
@staticmethod
def decode_image(
byte_sequence: bytes,
mode: str = "RGB",
size: Optional[tuple] = None,
) -> Image.Image:
"""Convert bytes to image.
Args:
byte_sequence: Raw image bytes
mode: Image mode (RGB, RGBA, L, etc.)
size: Optional tuple of (width, height)
"""
try:
# Try to load as-is first (for standard image formats)
img = Image.open(io.BytesIO(byte_sequence))
if size:
img = img.resize(size)
return img
except:
# If failed, assume raw pixel data
if not size:
# Try to determine size from byte count
pixel_count = len(byte_sequence) // len(mode)
size = (
int(np.sqrt(pixel_count)),
int(np.sqrt(pixel_count)),
)
# Convert raw bytes to pixel array
pixels = np.frombuffer(byte_sequence, dtype=np.uint8)
pixels = pixels.reshape((*size, len(mode)))
return Image.fromarray(pixels, mode=mode)
@staticmethod
def decode_audio(
byte_sequence: bytes,
sample_rate: int = 44100,
channels: int = 2,
sample_width: int = 2,
) -> np.ndarray:
"""Convert bytes to audio samples.
Args:
byte_sequence: Raw audio bytes
sample_rate: Audio sample rate in Hz
channels: Number of audio channels
sample_width: Bytes per sample (1, 2, or 4)
"""
# Determine format string based on sample width
format_str = {
1: "b", # signed char
2: "h", # short
4: "i", # int
}[sample_width]
# Unpack bytes to samples
sample_count = len(byte_sequence) // (channels * sample_width)
samples = struct.unpack(
f"<{sample_count * channels}{format_str}", byte_sequence
)
# Reshape to [samples, channels]
return np.array(samples).reshape(-1, channels)
def decode_data(
self,
model_output: Union[torch.Tensor, bytes],
data_type: DataType,
**kwargs,
) -> Union[str, Image.Image, np.ndarray, bytes]:
"""Main method to decode model output to desired format.
Args:
model_output: Either tensor from model or raw bytes
data_type: Type of data to decode to
**kwargs: Additional parameters for specific decoders
Returns:
Decoded data in specified format
"""
# Convert tensor to bytes if needed
if isinstance(model_output, torch.Tensor):
byte_sequence = self.tensor_to_bytes(model_output)
else:
byte_sequence = model_output
# Decode based on type
if data_type == DataType.TEXT:
return self.decode_text(byte_sequence)
elif data_type == DataType.IMAGE:
return self.decode_image(byte_sequence, **kwargs)
elif data_type == DataType.AUDIO:
return self.decode_audio(byte_sequence, **kwargs)
elif data_type == DataType.VIDEO:
raise NotImplementedError(
"Video decoding not yet implemented"
)
else: # BINARY
return byte_sequence
# Usage example
class Modality(Enum):
TEXT = auto()
IMAGE = auto()
AUDIO = auto()
VIDEO = auto()
BINARY = auto()
MULTIMODAL = auto()
@dataclass
class ModalityInfo:
"""Information about detected modality."""
modality: Modality
confidence: float
metadata: Dict[str, any]
sub_modalities: Optional[List["ModalityInfo"]] = None
class ModalityDetector:
"""Detects data modalities from byte sequences."""
# Common file signatures (magic numbers)
SIGNATURES = {
# Images
b"\xFF\xD8\xFF": "JPEG",
b"\x89PNG\r\n\x1a\n": "PNG",
b"GIF87a": "GIF",
b"GIF89a": "GIF",
b"RIFF": "WEBP",
# Audio
b"RIFF....WAVE": "WAV",
b"ID3": "MP3",
b"\xFF\xFB": "MP3",
b"OggS": "OGG",
# Video
b"\x00\x00\x00\x18ftypmp42": "MP4",
b"\x00\x00\x00\x1Cftypav01": "MP4",
b"\x1A\x45\xDF\xA3": "WEBM",
}
def __init__(self):
self.magic = magic.Magic(mime=True)
def _check_text_probability(self, data: bytes) -> float:
"""Estimate probability that data is text."""
# Check if data is valid UTF-8
try:
data.decode("utf-8")
# Count printable ASCII characters
printable = sum(1 for b in data if 32 <= b <= 126)
return printable / len(data)
except UnicodeDecodeError:
return 0.0
def _check_image_validity(self, data: bytes) -> Tuple[bool, Dict]:
"""Check if data is a valid image and extract metadata."""
try:
with io.BytesIO(data) as bio:
img = Image.open(bio)
return True, {
"format": img.format,
"size": img.size,
"mode": img.mode,
}
except:
return False, {}
def _check_audio_validity(self, data: bytes) -> Tuple[bool, Dict]:
"""Check if data is valid audio and extract metadata."""
try:
with io.BytesIO(data) as bio:
# Try to parse as WAV
with wave.open(bio) as wav:
return True, {
"channels": wav.getnchannels(),
"sample_width": wav.getsampwidth(),
"framerate": wav.getframerate(),
"frames": wav.getnframes(),
}
except:
# Check for other audio signatures
for sig in [b"ID3", b"\xFF\xFB", b"OggS"]:
if data.startswith(sig):
return True, {"format": "compressed_audio"}
return False, {}
def _detect_boundaries(
self, data: bytes
) -> List[Tuple[int, int, Modality]]:
"""Detect boundaries between different modalities in the data."""
boundaries = []
current_pos = 0
while current_pos < len(data):
# Look for known signatures
for sig, format_type in self.SIGNATURES.items():
if data[current_pos:].startswith(sig):
# Found a signature, determine its length
if format_type in ["JPEG", "PNG", "GIF"]:
# Find image end
try:
with io.BytesIO(
data[current_pos:]
) as bio:
img = Image.open(bio)
img.verify()
size = bio.tell()
boundaries.append(
(
current_pos,
current_pos + size,
Modality.IMAGE,
)
)
current_pos += size
continue
except:
pass
# Check for text sections
text_prob = self._check_text_probability(
data[current_pos : current_pos + 1024]
)
if text_prob > 0.8:
# Look for end of text section
end_pos = current_pos + 1
while end_pos < len(data):
if (
self._check_text_probability(
data[end_pos : end_pos + 32]
)
< 0.5
):
break
end_pos += 1
boundaries.append(
(current_pos, end_pos, Modality.TEXT)
)
current_pos = end_pos
continue
current_pos += 1
return boundaries
def detect_modality(self, data: bytes) -> ModalityInfo:
"""Detect modality of byte sequence."""
# First check for single modality
mime_type = self.magic.from_buffer(data)
# Check text
text_prob = self._check_text_probability(data)
if text_prob > 0.9:
return ModalityInfo(
modality=Modality.TEXT,
confidence=text_prob,
metadata={"mime_type": mime_type},
)
# Check image
is_image, image_meta = self._check_image_validity(data)
if is_image:
return ModalityInfo(
modality=Modality.IMAGE,
confidence=1.0,
metadata={**image_meta, "mime_type": mime_type},
)
# Check audio
is_audio, audio_meta = self._check_audio_validity(data)
if is_audio:
return ModalityInfo(
modality=Modality.AUDIO,
confidence=1.0,
metadata={**audio_meta, "mime_type": mime_type},
)
# Check for multimodal content
boundaries = self._detect_boundaries(data)
if len(boundaries) > 1:
sub_modalities = []
for start, end, modality in boundaries:
chunk_data = data[start:end]
sub_info = self.detect_modality(chunk_data)
if sub_info.modality != Modality.BINARY:
sub_modalities.append(sub_info)
if sub_modalities:
return ModalityInfo(
modality=Modality.MULTIMODAL,
confidence=0.8,
metadata={"mime_type": "multipart/mixed"},
sub_modalities=sub_modalities,
)
# Default to binary
return ModalityInfo(
modality=Modality.BINARY,
confidence=0.5,
metadata={"mime_type": mime_type},
)
def split_modalities(
self, data: bytes
) -> List[Tuple[Modality, bytes, Dict]]:
"""Split multimodal data into separate modalities."""
boundaries = self._detect_boundaries(data)
result = []
for start, end, modality in boundaries:
chunk = data[start:end]
info = self.detect_modality(chunk)
result.append((modality, chunk, info.metadata))
return result
class AutoDetectBytesDecoder:
"""Decoder that automatically detects and decodes different modalities."""
def __init__(self):
self.detector = ModalityDetector()
self.text_decoder = ByteDetokenizer() # From previous example
def decode(
self, data: bytes
) -> Union[str, Image.Image, np.ndarray, List[any]]:
"""Automatically detect and decode byte sequence."""
info = self.detector.detect_modality(data)
if info.modality == Modality.MULTIMODAL:
# Handle multimodal content
parts = self.detector.split_modalities(data)
return [
self.decode(chunk) for modality, chunk, _ in parts
]
if info.modality == Modality.TEXT:
return self.text_decoder.decode_text(data)
elif info.modality == Modality.IMAGE:
return self.text_decoder.decode_image(data)
elif info.modality == Modality.AUDIO:
return self.text_decoder.decode_audio(data)
else:
return data
# # Example usage
# def demo_auto_detection():
# """Demonstrate auto modality detection."""
# # Create mixed content
# text = "Hello, World!".encode('utf-8')
# # Create a small test image
# img = Image.new('RGB', (100, 100), color='red')
# img_bytes = io.BytesIO()
# img.save(img_bytes, format='PNG')
# # Combine into multimodal content
# mixed_content = text + img_bytes.getvalue()
# # Initialize decoder
# decoder = AutoDetectBytesDecoder()
# # Decode
# result = decoder.decode(mixed_content)
# if isinstance(result, list):
# print("Detected multimodal content:")
# for i, part in enumerate(result):
# print(f"Part {i+1}: {type(part)}")
# if __name__ == "__main__":
# demo_auto_detection()
def tensor_to_data(tensor: Tensor):
byte_sequence = ByteDetokenizer.tensor_to_bytes(tensor)
# Initialize auto-detector
decoder = AutoDetectBytesDecoder()
# Decode with automatic detection
result = decoder.decode(byte_sequence)
return result
def demo_byte_predictor():
"""Demo with smaller dimensions to test."""
# Initialize model configuration with adjusted dimensions
config = ModelConfig(
vocab_size=256,
hidden_size=128, # Smaller for testing
num_layers=2, # Fewer layers for testing
num_key_value_heads=2,
num_query_heads=4,
dropout=0.1,
max_sequence_length=1024,
)
# Initialize model
model = EnhancedBytePredictor(config)
logger.info("Model initialized")
# Move to GPU if available
device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
model = model.to(device)
logger.info(f"Using device: {device}")
# Create sample input data
batch_size = 2
seq_length = 16 # Shorter sequence for testing
input_ids = torch.randint(
0, config.vocab_size, (batch_size, seq_length), device=device
)
logger.info(f"Created input tensor of shape: {input_ids.shape}")
# Test forward pass
try:
logits = model(input_ids)
logger.info(
f"Forward pass successful! Output shape: {logits.shape}"
)
# Test loss computation
target_ids = torch.randint(
0,
config.vocab_size,
(batch_size, seq_length),
device=device,
)
loss = model.compute_loss(input_ids, target_ids)
logger.info(
f"Loss computation successful! Loss value: {loss.item():.4f}"
)
# Test generation
prompt = torch.randint(
0,
config.vocab_size,
(1, 4), # Very short prompt for testing
device=device,
)
generated = model.generate(
prompt, max_new_tokens=8, temperature=0.8, top_k=50
)
logger.info(
f"Generation successful! Generated shape: {generated.shape}"
)
except Exception as e:
logger.error(f"Error during execution: {str(e)}")
raise
if __name__ == "__main__":
# Set up logging
# logger.remove() # Remove default handler
# logger.add(sys.stderr, format="<green>{time:HH:mm:ss}</green> | {level} | {message}")
demo_byte_predictor()

10
docs/mkdocs.yml
Unescape View File

@ -222,11 +222,11 @@ nav:
- BaseMultiModalModel: "swarms/models/base_multimodal_model.md" - BaseMultiModalModel: "swarms/models/base_multimodal_model.md"
- Multi Modal Models Available: "swarms/models/multimodal_models.md" - Multi Modal Models Available: "swarms/models/multimodal_models.md"
- GPT4VisionAPI: "swarms/models/gpt4v.md" - GPT4VisionAPI: "swarms/models/gpt4v.md"
- Swarms Cloud API: # - Swarms Cloud API:
# - Overview: "swarms_cloud/main.md" # # - Overview: "swarms_cloud/main.md"
- Overview: "swarms_cloud/vision.md" # - Overview: "swarms_cloud/vision.md"
- Swarms Cloud CLI: "swarms_cloud/cli.md" # - Swarms Cloud CLI: "swarms_cloud/cli.md"
- Add Agents to Marketplace: "swarms_cloud/add_agent.md" # # - Add Agents to Marketplace: "swarms_cloud/add_agent.md"
# - Available Models: "swarms_cloud/available_models.md" # - Available Models: "swarms_cloud/available_models.md"
# - Agent API: "swarms_cloud/agent_api.md" # - Agent API: "swarms_cloud/agent_api.md"
# - Migrate from OpenAI to Swarms in 3 lines of code: "swarms_cloud/migrate_openai.md" # - Migrate from OpenAI to Swarms in 3 lines of code: "swarms_cloud/migrate_openai.md"

25
example.py
Unescape View File

@ -1,28 +1,13 @@
import os
from dotenv import load_dotenv
from swarm_models import OpenAIChat
from swarms import Agent from swarms import Agent
from swarms.prompts.finance_agent_sys_prompt import ( from swarms.prompts.finance_agent_sys_prompt import (
FINANCIAL_AGENT_SYS_PROMPT, FINANCIAL_AGENT_SYS_PROMPT,
) )
load_dotenv()
# Get the OpenAI API key from the environment variable
api_key = os.getenv("OPENAI_API_KEY")
# Create an instance of the OpenAIChat class
model = OpenAIChat(
openai_api_key=api_key, model_name="gpt-4o-mini", temperature=0.1
)
# Initialize the agent # Initialize the agent
agent = Agent( agent = Agent(
agent_name="Financial-Analysis-Agent", agent_name="Financial-Analysis-Agent",
system_prompt=FINANCIAL_AGENT_SYS_PROMPT, system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
llm=model, model_name="gpt-4o-mini",
max_loops=1, max_loops=1,
autosave=True, autosave=True,
dashboard=False, dashboard=False,
@ -33,14 +18,10 @@ agent = Agent(
retry_attempts=1, retry_attempts=1,
streaming_on=True, streaming_on=True,
context_length=200000, context_length=200000,
return_step_meta=True, return_step_meta=False,
output_type="json", # "json", "dict", "csv" OR "string" soon "yaml" and output_type="str", # "json", "dict", "csv" OR "string" soon "yaml" and
auto_generate_prompt=False, # Auto generate prompt for the agent based on name, description, and system prompt, task auto_generate_prompt=False, # Auto generate prompt for the agent based on name, description, and system prompt, task
artifacts_on=True,
artifacts_output_path="roth_ira_report",
artifacts_file_extension=".txt",
max_tokens=8000, max_tokens=8000,
return_history=True,
) )

101
new_features_examples/auto_agent.py
Unescape View File

@ -12,21 +12,31 @@ class DynamicParser:
@staticmethod @staticmethod
def extract_fields(model: Type[BaseModel]) -> Dict[str, Any]: def extract_fields(model: Type[BaseModel]) -> Dict[str, Any]:
return { return {
field_name: (field.annotation, ... if field.is_required() else None) field_name: (
field.annotation,
... if field.is_required() else None,
)
for field_name, field in model.model_fields.items() for field_name, field in model.model_fields.items()
} }
@staticmethod @staticmethod
def create_partial_model(model: Type[BaseModel], data: Dict[str, Any]) -> Type[BaseModel]: def create_partial_model(
model: Type[BaseModel], data: Dict[str, Any]
) -> Type[BaseModel]:
fields = { fields = {
field_name: (field.annotation, ... if field.is_required() else None) field_name: (
field.annotation,
... if field.is_required() else None,
)
for field_name, field in model.model_fields.items() for field_name, field in model.model_fields.items()
if field_name in data if field_name in data
} }
return create_model(f"Partial{model.__name__}", **fields) return create_model(f"Partial{model.__name__}", **fields)
@classmethod @classmethod
def parse(cls, data: Union[str, Dict[str, Any]], model: Type[BaseModel]) -> Optional[BaseModel]: def parse(
cls, data: Union[str, Dict[str, Any]], model: Type[BaseModel]
) -> Optional[BaseModel]:
if isinstance(data, str): if isinstance(data, str):
try: try:
data = json.loads(data) data = json.loads(data)
@ -47,25 +57,52 @@ class DynamicParser:
load_dotenv() load_dotenv()
# Define the Thoughts schema # Define the Thoughts schema
class Thoughts(BaseModel): class Thoughts(BaseModel):
text: str = Field(..., description="Current thoughts or observations regarding the task.") text: str = Field(
reasoning: str = Field(..., description="Logical reasoning behind the thought process.") ...,
plan: str = Field(..., description="A short bulleted list that conveys the immediate and long-term plan.") description="Current thoughts or observations regarding the task.",
criticism: str = Field(..., description="Constructive self-criticism to improve future responses.") )
speak: str = Field(..., description="A concise summary of thoughts intended for the user.") reasoning: str = Field(
...,
description="Logical reasoning behind the thought process.",
)
plan: str = Field(
...,
description="A short bulleted list that conveys the immediate and long-term plan.",
)
criticism: str = Field(
...,
description="Constructive self-criticism to improve future responses.",
)
speak: str = Field(
...,
description="A concise summary of thoughts intended for the user.",
)
# Define the Command schema # Define the Command schema
class Command(BaseModel): class Command(BaseModel):
name: str = Field(..., description="Command name to execute from the provided list of commands.") name: str = Field(
args: Dict[str, Any] = Field(..., description="Arguments required to execute the command.") ...,
description="Command name to execute from the provided list of commands.",
)
args: Dict[str, Any] = Field(
..., description="Arguments required to execute the command."
)
# Define the AgentResponse schema # Define the AgentResponse schema
class AgentResponse(BaseModel): class AgentResponse(BaseModel):
thoughts: Thoughts = Field(..., description="The agent's current thoughts and reasoning.") thoughts: Thoughts = Field(
command: Command = Field(..., description="The command to execute along with its arguments.") ..., description="The agent's current thoughts and reasoning."
)
command: Command = Field(
...,
description="The command to execute along with its arguments.",
)
# Define tool functions # Define tool functions
def fluid_api_command(task: str): def fluid_api_command(task: str):
@ -90,17 +127,26 @@ def do_nothing_command():
def task_complete_command(reason: str): def task_complete_command(reason: str):
"""Mark the task as complete and provide a reason.""" """Mark the task as complete and provide a reason."""
print(f"Task completed: {reason}") print(f"Task completed: {reason}")
return {"status": "success", "message": f"Task completed: {reason}"} return {
"status": "success",
"message": f"Task completed: {reason}",
}
# Dynamic command execution # Dynamic command execution
def execute_command(name: str, args: Dict[str, Any]): def execute_command(name: str, args: Dict[str, Any]):
"""Dynamically execute a command based on its name and arguments.""" """Dynamically execute a command based on its name and arguments."""
command_map: Dict[str, Callable] = { command_map: Dict[str, Callable] = {
"fluid_api": lambda **kwargs: fluid_api_command(task=kwargs.get("task")), "fluid_api": lambda **kwargs: fluid_api_command(
"send_tweet": lambda **kwargs: send_tweet_command(text=kwargs.get("text")), task=kwargs.get("task")
),
"send_tweet": lambda **kwargs: send_tweet_command(
text=kwargs.get("text")
),
"do_nothing": lambda **kwargs: do_nothing_command(), "do_nothing": lambda **kwargs: do_nothing_command(),
"task_complete": lambda **kwargs: task_complete_command(reason=kwargs.get("reason")), "task_complete": lambda **kwargs: task_complete_command(
reason=kwargs.get("reason")
),
} }
if name not in command_map: if name not in command_map:
@ -110,23 +156,26 @@ def execute_command(name: str, args: Dict[str, Any]):
return command_map[name](**args) return command_map[name](**args)
def parse_and_execute_command(response: Union[str, Dict[str, Any]], base_model: Type[BaseModel] = AgentResponse) -> Any: def parse_and_execute_command(
response: Union[str, Dict[str, Any]],
base_model: Type[BaseModel] = AgentResponse,
) -> Any:
"""Enhanced command parser with flexible input handling""" """Enhanced command parser with flexible input handling"""
parsed = DynamicParser.parse(response, base_model) parsed = DynamicParser.parse(response, base_model)
if not parsed: if not parsed:
raise ValueError("Failed to parse response") raise ValueError("Failed to parse response")
if hasattr(parsed, 'command'): if hasattr(parsed, "command"):
command_name = parsed.command.name command_name = parsed.command.name
command_args = parsed.command.args command_args = parsed.command.args
return execute_command(command_name, command_args) return execute_command(command_name, command_args)
return parsed return parsed
ainame = "AutoAgent" ainame = "AutoAgent"
userprovided = "assistant" userprovided = "assistant"
SYSTEM_PROMPT = f""" SYSTEM_PROMPT = f"""
You are {ainame}, an advanced and autonomous {userprovided}. You are {ainame}, an advanced and autonomous {userprovided}.
Your role is to make decisions and complete tasks independently without seeking user assistance. Leverage your strengths as an LLM to solve tasks efficiently, adhering strictly to the commands and resources provided. Your role is to make decisions and complete tasks independently without seeking user assistance. Leverage your strengths as an LLM to solve tasks efficiently, adhering strictly to the commands and resources provided.
@ -174,7 +223,7 @@ model = OpenAIFunctionCaller(
temperature=0.9, temperature=0.9,
base_model=AgentResponse, # Pass the Pydantic schema as the base model base_model=AgentResponse, # Pass the Pydantic schema as the base model
parallel_tool_calls=False, parallel_tool_calls=False,
openai_api_key=os.getenv("OPENAI_API_KEY") openai_api_key=os.getenv("OPENAI_API_KEY"),
) )
# Example usage # Example usage

8
new_features_examples/markdown_agent.py
Unescape View File

@ -0,0 +1,8 @@
from swarms import Agent
Agent(
agent_name="Stock-Analysis-Agent",
model_name="gpt-4o-mini",
max_loops=1,
streaming_on=True,
).run("What are 5 hft algorithms")

39
pyproject.toml
Unescape View File

@ -5,8 +5,8 @@ build-backend = "poetry.core.masonry.api"
[tool.poetry] [tool.poetry]
name = "swarms" name = "swarms"
version = "6.4.7" version = "6.5.7"
description = "Swarms - Pytorch" description = "Swarms - TGSC"
license = "MIT" license = "MIT"
authors = ["Kye Gomez <kye@apac.ai>"] authors = ["Kye Gomez <kye@apac.ai>"]
homepage = "https://github.com/kyegomez/swarms" homepage = "https://github.com/kyegomez/swarms"
@ -57,13 +57,13 @@ classifiers = [
[tool.poetry.dependencies] [tool.poetry.dependencies]
python = ">=3.10,<4.0" python = ">=3.10,<4.0"
torch = ">=2.1.1,<3.0" # torch = ">=2.1.1,<3.0"
transformers = ">= 4.39.0, <5.0.0" # transformers = ">= 4.39.0, <5.0.0"
asyncio = ">=3.4.3,<4.0" asyncio = ">=3.4.3,<4.0"
toml = "*" toml = "*"
pypdf = "4.3.1" pypdf = "4.3.1"
loguru = "*" loguru = "*"
pydantic = "2.8.2" pydantic = ">=2.8.2<3.0"
tenacity = "*" tenacity = "*"
psutil = "*" psutil = "*"
sentry-sdk = {version = "*", extras = ["http"]} # Updated here sentry-sdk = {version = "*", extras = ["http"]} # Updated here
@ -73,12 +73,37 @@ docstring_parser = "0.16"
tiktoken = "*" tiktoken = "*"
networkx = "*" networkx = "*"
aiofiles = "*" aiofiles = "*"
swarm-models = "*"
clusterops = "*" clusterops = "*"
chromadb = "*" # chromadb = "*"
reportlab = "*" reportlab = "*"
doc-master = "*" doc-master = "*"
rich = "*" rich = "*"
# sentence-transformers = "*"
swarm-models = "*"
# [tool.poetry.extras]
# # Extra for NLP-related functionalities
# nlp = [
# "torch>=2.1.1,<3.0",
# "transformers>=4.39.0,<5.0.0",
# "sentence-transformers",
# "swarm-models",
# ]
# # Extra for database-related functionalities
# db = ["chromadb"]
# # All optional dependencies for convenience
# all = [
# "torch>=2.1.1,<3.0",
# "transformers>=4.39.0,<5.0.0",
# "sentence-transformers",
# "chromadb",
# "swarm-models"
# ]
[tool.poetry.scripts] [tool.poetry.scripts]
swarms = "swarms.cli.main:main" swarms = "swarms.cli.main:main"

618
real_time.py
Unescape View File

@ -0,0 +1,618 @@
import torch
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
from loguru import logger
from dataclasses import dataclass
from typing import Optional, Tuple, Dict
import math
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
@dataclass
class TransformerConfig:
"""Configuration class for MoE Transformer model parameters."""
vocab_size: int = 50257
hidden_size: int = 768
num_attention_heads: int = 12
num_expert_layers: int = 4
num_experts: int = 8
expert_capacity: int = 32
max_position_embeddings: int = 1024
dropout_prob: float = 0.1
layer_norm_epsilon: float = 1e-5
initializer_range: float = 0.02
num_query_groups: int = 4 # For multi-query attention
class ExpertLayer(nn.Module):
"""Individual expert neural network."""
def __init__(self, config: TransformerConfig):
super().__init__()
self.fc1 = nn.Linear(
config.hidden_size, 4 * config.hidden_size
)
self.fc2 = nn.Linear(
4 * config.hidden_size, config.hidden_size
)
self.activation = nn.GELU()
self.dropout = nn.Dropout(config.dropout_prob)
def forward(self, x: Tensor) -> Tensor:
x = self.fc1(x)
x = self.activation(x)
x = self.dropout(x)
x = self.fc2(x)
return x
class MixtureOfExperts(nn.Module):
"""Mixture of Experts layer with dynamic routing."""
def __init__(self, config: TransformerConfig):
super().__init__()
self.num_experts = config.num_experts
self.expert_capacity = config.expert_capacity
# Create expert networks
self.experts = nn.ModuleList(
[ExpertLayer(config) for _ in range(config.num_experts)]
)
# Router network
self.router = nn.Linear(
config.hidden_size, config.num_experts
)
def forward(self, x: Tensor) -> Tuple[Tensor, Dict]:
"""Route inputs to experts and combine outputs."""
batch_size, seq_len, hidden_size = x.shape
# Calculate routing probabilities
router_logits = self.router(x)
routing_weights = F.softmax(router_logits, dim=-1)
# Select top-k experts
top_k = 2
gates, indices = torch.topk(routing_weights, top_k, dim=-1)
gates = F.softmax(gates, dim=-1)
# Process inputs through selected experts
final_output = torch.zeros_like(x)
router_load = torch.zeros(self.num_experts, device=x.device)
for i in range(top_k):
expert_index = indices[..., i]
gate = gates[..., i : i + 1]
# Count expert assignments
for j in range(self.num_experts):
router_load[j] += (expert_index == j).float().sum()
# Process through selected experts
for j in range(self.num_experts):
mask = expert_index == j
if not mask.any():
continue
expert_input = x[mask]
expert_output = self.experts[j](expert_input)
final_output[mask] += gate[mask] * expert_output
aux_loss = router_load.float().var() / (
router_load.float().mean() ** 2
)
return final_output, {"load_balancing_loss": aux_loss}
class MultiQueryAttention(nn.Module):
"""Multi-Query Attention mechanism with proper multi-query group handling."""
def __init__(self, config: TransformerConfig):
super().__init__()
self.num_attention_heads = config.num_attention_heads
self.num_query_groups = config.num_query_groups
self.hidden_size = config.hidden_size
self.head_dim = (
config.hidden_size // config.num_attention_heads
)
# Query projection maintains full head dimension
self.q_proj = nn.Linear(
config.hidden_size, config.hidden_size
)
# Key and value projections use reduced number of heads (query groups)
self.k_proj = nn.Linear(
config.hidden_size,
self.head_dim * config.num_query_groups,
)
self.v_proj = nn.Linear(
config.hidden_size,
self.head_dim * config.num_query_groups,
)
self.dropout = nn.Dropout(config.dropout_prob)
# Calculate heads per group for proper reshaping
self.heads_per_group = (
self.num_attention_heads // self.num_query_groups
)
def forward(
self,
hidden_states: Tensor,
attention_mask: Optional[Tensor] = None,
cache: Optional[Dict[str, Tensor]] = None,
) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
batch_size, seq_length, _ = hidden_states.shape
# Project queries, keys, and values
queries = self.q_proj(hidden_states)
keys = self.k_proj(hidden_states)
values = self.v_proj(hidden_states)
# Reshape queries to full number of heads
queries = queries.view(
batch_size,
seq_length,
self.num_attention_heads,
self.head_dim,
)
# Reshape keys and values to number of query groups
keys = keys.view(
batch_size,
seq_length,
self.num_query_groups,
self.head_dim,
)
values = values.view(
batch_size,
seq_length,
self.num_query_groups,
self.head_dim,
)
# Transpose for batch matrix multiplication
queries = queries.transpose(
1, 2
) # (batch, n_heads, seq_len, head_dim)
keys = keys.transpose(
1, 2
) # (batch, n_groups, seq_len, head_dim)
values = values.transpose(
1, 2
) # (batch, n_groups, seq_len, head_dim)
# Repeat keys and values for each head in the group
keys = keys.repeat_interleave(self.heads_per_group, dim=1)
values = values.repeat_interleave(self.heads_per_group, dim=1)
# Compute attention scores
scale = 1.0 / math.sqrt(self.head_dim)
scores = torch.matmul(queries, keys.transpose(-2, -1)) * scale
if attention_mask is not None:
# Expand attention mask to match scores dimensions
expanded_mask = attention_mask.unsqueeze(1).unsqueeze(2)
expanded_mask = expanded_mask.expand(
batch_size,
self.num_attention_heads,
seq_length,
seq_length,
)
mask_value = torch.finfo(scores.dtype).min
attention_mask = expanded_mask.eq(0).float() * mask_value
scores = scores + attention_mask
attention_weights = F.softmax(scores, dim=-1)
attention_weights = self.dropout(attention_weights)
# Compute attention output
attention_output = torch.matmul(attention_weights, values)
attention_output = attention_output.transpose(1, 2)
attention_output = attention_output.reshape(
batch_size, seq_length, -1
)
return attention_output, None
class MoETransformer(nn.Module):
"""
Production-grade Transformer model with Mixture of Experts and Multi-Query Attention.
Features:
- Multi-Query Attention mechanism for efficient inference
- Mixture of Experts for dynamic routing and specialization
- Real-time weight updates based on input similarity
- Built-in logging and monitoring
- Type annotations for better code maintainability
"""
def __init__(self, config: TransformerConfig):
super().__init__()
self.config = config
# Initialize components
self.embedding = nn.Embedding(
config.vocab_size, config.hidden_size
)
self.position_embedding = nn.Embedding(
config.max_position_embeddings, config.hidden_size
)
# Multi-Query Attention layers
self.attention_layers = nn.ModuleList(
[
MultiQueryAttention(config)
for _ in range(config.num_expert_layers)
]
)
# Mixture of Experts layers
self.moe_layers = nn.ModuleList(
[
MixtureOfExperts(config)
for _ in range(config.num_expert_layers)
]
)
# Layer normalization and dropout
self.layer_norm = nn.LayerNorm(
config.hidden_size, eps=config.layer_norm_epsilon
)
self.dropout = nn.Dropout(config.dropout_prob)
# Output projection
self.output_projection = nn.Linear(
config.hidden_size, config.vocab_size
)
# Initialize weights
self.apply(self._init_weights)
logger.info("Initialized MoETransformer model")
def _init_weights(self, module: nn.Module):
"""Initialize model weights."""
if isinstance(module, (nn.Linear, nn.Embedding)):
module.weight.data.normal_(
mean=0.0, std=self.config.initializer_range
)
if (
isinstance(module, nn.Linear)
and module.bias is not None
):
module.bias.data.zero_()
def get_position_embeddings(self, position_ids: Tensor) -> Tensor:
"""Generate position embeddings."""
return self.position_embedding(position_ids)
def forward(
self,
input_ids: Tensor,
attention_mask: Optional[Tensor] = None,
position_ids: Optional[Tensor] = None,
cache: Optional[Dict[str, Tensor]] = None,
) -> Tuple[Tensor, Dict]:
"""
Forward pass through the model.
Args:
input_ids: Input token IDs
attention_mask: Attention mask for padding
position_ids: Position IDs for positioning encoding
cache: Cache for key/value states in generation
Returns:
tuple: (logits, auxiliary_outputs)
"""
batch_size, seq_length = input_ids.shape
if position_ids is None:
position_ids = torch.arange(
seq_length, dtype=torch.long, device=input_ids.device
)
position_ids = position_ids.unsqueeze(0).expand_as(
input_ids
)
# Get embeddings
inputs_embeds = self.embedding(input_ids)
position_embeds = self.get_position_embeddings(position_ids)
hidden_states = inputs_embeds + position_embeds
# Initialize auxiliary outputs
aux_outputs = {"moe_losses": []}
# Process through transformer layers
for attention_layer, moe_layer in zip(
self.attention_layers, self.moe_layers
):
# Multi-Query Attention
attention_output, _ = attention_layer(
hidden_states, attention_mask, cache
)
hidden_states = self.layer_norm(
hidden_states + attention_output
)
# Mixture of Experts
moe_output, moe_aux = moe_layer(hidden_states)
hidden_states = self.layer_norm(
hidden_states + moe_output
)
aux_outputs["moe_losses"].append(
moe_aux["load_balancing_loss"]
)
# Final output projection
logits = self.output_projection(hidden_states)
return logits, aux_outputs
def fetch_loss(
self,
logits: Tensor,
labels: Tensor,
aux_outputs: Dict,
reduction: str = "mean",
) -> Tensor:
"""
Calculate the total loss including MoE balancing losses.
Args:
logits: Model output logits
labels: Ground truth labels
aux_outputs: Auxiliary outputs from forward pass
reduction: Loss reduction method
Returns:
Tensor: Total loss
"""
# Calculate cross entropy loss
ce_loss = F.cross_entropy(
logits.view(-1, self.config.vocab_size),
labels.view(-1),
reduction=reduction,
)
# Calculate MoE loss
moe_loss = torch.stack(aux_outputs["moe_losses"]).mean()
# Combine losses
total_loss = ce_loss + 0.01 * moe_loss
logger.debug(
f"CE Loss: {ce_loss.item():.4f}, "
f"MoE Loss: {moe_loss.item():.4f}"
)
return total_loss
@torch.no_grad()
def generate(
self,
input_ids: Tensor,
max_length: int = 100,
temperature: float = 1.0,
top_k: int = 50,
top_p: float = 0.9,
) -> Tensor:
"""
Generate text using the model.
Args:
input_ids: Initial input tokens
max_length: Maximum sequence length to generate
temperature: Sampling temperature
top_k: Number of highest probability tokens to keep
top_p: Cumulative probability for nucleus sampling
Returns:
Tensor: Generated token IDs
"""
batch_size = input_ids.shape[0]
device = input_ids.device
# Initialize sequence with input_ids
generated = input_ids
# Cache for key-value pairs
cache = {}
for _ in range(max_length):
# Get position IDs for current sequence
position_ids = torch.arange(
generated.shape[1], dtype=torch.long, device=device
)
position_ids = position_ids.unsqueeze(0).expand(
batch_size, -1
)
# Forward pass
logits, _ = self.forward(
generated, position_ids=position_ids, cache=cache
)
# Get next token logits
next_token_logits = logits[:, -1, :] / temperature
# Apply top-k filtering
if top_k > 0:
indices_to_remove = (
next_token_logits
< torch.topk(next_token_logits, top_k)[0][
..., -1, None
]
)
next_token_logits[indices_to_remove] = float("-inf")
# Apply top-p (nucleus) filtering
if top_p < 1.0:
sorted_logits, sorted_indices = torch.sort(
next_token_logits, descending=True
)
cumulative_probs = torch.cumsum(
F.softmax(sorted_logits, dim=-1), dim=-1
)
# Remove tokens with cumulative probability above the threshold
sorted_indices_to_remove = cumulative_probs > top_p
sorted_indices_to_remove[..., 1:] = (
sorted_indices_to_remove[..., :-1].clone()
)
sorted_indices_to_remove[..., 0] = 0
indices_to_remove = sorted_indices[
sorted_indices_to_remove
]
next_token_logits[indices_to_remove] = float("-inf")
# Sample next token
probs = F.softmax(next_token_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
# Append next token to sequence
generated = torch.cat((generated, next_token), dim=1)
# Check for end of sequence token
if (next_token == self.config.vocab_size - 1).all():
break
return generated
# Initialize model configuration
config = TransformerConfig(
vocab_size=50257,
hidden_size=768,
num_attention_heads=12,
num_expert_layers=4,
num_experts=8,
expert_capacity=32,
max_position_embeddings=1024,
num_query_groups=4,
)
def prepare_sample_data(
batch_size: int = 8,
seq_length: int = 512,
vocab_size: int = 50257,
) -> DataLoader:
"""Create sample data for demonstration."""
# Create random input sequences
input_ids = torch.randint(
0, vocab_size, (100, seq_length) # 100 samples
)
# Create target sequences (shifted by 1)
labels = torch.randint(0, vocab_size, (100, seq_length))
# Create attention masks (1 for real tokens, 0 for padding)
attention_mask = torch.ones_like(input_ids)
# Create dataset and dataloader
dataset = TensorDataset(input_ids, attention_mask, labels)
dataloader = DataLoader(
dataset, batch_size=batch_size, shuffle=True
)
return dataloader
def train_step(
model: MoETransformer,
batch: tuple,
optimizer: torch.optim.Optimizer,
device: str = "cuda" if torch.cuda.is_available() else "cpu",
) -> float:
"""Execute single training step."""
model.train()
optimizer.zero_grad()
# Unpack batch
input_ids, attention_mask, labels = [b.to(device) for b in batch]
# Forward pass
logits, aux_outputs = model(
input_ids=input_ids, attention_mask=attention_mask
)
# Calculate loss
loss = model.fetch_loss(logits, labels, aux_outputs)
# Backward pass
loss.backward()
optimizer.step()
return loss.item()
def main():
# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"Using device: {device}")
# Initialize model
model = MoETransformer(config).to(device)
logger.info("Model initialized")
# Setup optimizer
optimizer = torch.optim.AdamW(
model.parameters(), lr=1e-4, weight_decay=0.01
)
# Prepare data
dataloader = prepare_sample_data()
logger.info("Data prepared")
# Training loop
num_epochs = 3
for epoch in range(num_epochs):
epoch_losses = []
for batch_idx, batch in enumerate(dataloader):
loss = train_step(model, batch, optimizer, device)
epoch_losses.append(loss)
if batch_idx % 10 == 0:
logger.info(
f"Epoch {epoch+1}/{num_epochs} "
f"Batch {batch_idx}/{len(dataloader)} "
f"Loss: {loss:.4f}"
)
avg_loss = np.mean(epoch_losses)
logger.info(f"Epoch {epoch+1} average loss: {avg_loss:.4f}")
# Generation example
model.eval()
with torch.no_grad():
# Prepare input prompt
prompt = torch.randint(0, config.vocab_size, (1, 10)).to(
device
)
# Generate sequence
generated = model.generate(
input_ids=prompt,
max_length=50,
temperature=0.7,
top_k=50,
top_p=0.9,
)
logger.info(f"Generated sequence shape: {generated.shape}")
if __name__ == "__main__":
main()

1
requirements.txt
Unescape View File

@ -24,7 +24,6 @@ pytest>=8.1.1
pandas>=2.2.2 pandas>=2.2.2
networkx networkx
aiofiles aiofiles
swarm-models
clusterops clusterops
reportlab reportlab
doc-master doc-master

121
scripts/platform_update/parse_prompts_and_submit_to_marketplace 2.py
Unescape View File

@ -1,121 +0,0 @@
import json
import os
from difflib import SequenceMatcher
import requests
from dotenv import load_dotenv
from loguru import logger
from supabase import Client, create_client
load_dotenv()
# Initialize Supabase client
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
# Swarms API URL and headers
SWARMS_API_URL = "https://swarms.world/api/add-prompt"
SWARMS_API_KEY = os.getenv("SWARMS_API_KEY")
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {SWARMS_API_KEY}",
}
# Configure logger
logger.add(
"fetch_and_publish_prompts.log", rotation="1 MB"
) # Log file with rotation
def fetch_and_publish_prompts():
logger.info("Starting to fetch and publish prompts.")
# Fetch data from Supabase
try:
response = (
supabase.table("swarms_framework_schema")
.select("*")
.execute()
)
rows = response.data
logger.info(f"Fetched {len(rows)} rows from Supabase.")
except Exception as e:
logger.error(f"Failed to fetch data from Supabase: {e}")
return
# Track published prompts to avoid duplicates
published_prompts = set()
for row in rows:
# Extract agent_name and system_prompt
data = row.get("data", {})
agent_name = data.get("agent_name")
system_prompt = data.get("system_prompt")
# Skip if either is missing or duplicate
if not agent_name or not system_prompt:
logger.warning(
f"Skipping row due to missing agent_name or system_prompt: {row}"
)
continue
if is_duplicate(system_prompt, published_prompts):
logger.info(
f"Skipping duplicate prompt for agent: {agent_name}"
)
continue
# Create the data payload for the marketplace
prompt_data = {
"name": f"{agent_name} - System Prompt",
"prompt": system_prompt,
"description": f"System prompt for agent {agent_name}.",
"useCases": extract_use_cases(system_prompt),
"tags": "agent, system-prompt",
}
# Publish to the marketplace
try:
response = requests.post(
SWARMS_API_URL,
headers=headers,
data=json.dumps(prompt_data),
)
if response.status_code == 200:
logger.info(
f"Successfully published prompt for agent: {agent_name}"
)
published_prompts.add(system_prompt)
else:
logger.error(
f"Failed to publish prompt for agent: {agent_name}. Response: {response.text}"
)
except Exception as e:
logger.error(
f"Exception occurred while publishing prompt for agent: {agent_name}. Error: {e}"
)
def is_duplicate(new_prompt, published_prompts):
"""Check if the prompt is a duplicate using semantic similarity."""
for prompt in published_prompts:
similarity = SequenceMatcher(None, new_prompt, prompt).ratio()
if (
similarity > 0.9
): # Threshold for considering prompts as duplicates
return True
return False
def extract_use_cases(prompt):
"""Extract use cases from the prompt by chunking it into meaningful segments."""
# This is a simple placeholder; you can use a more advanced method to extract use cases
chunks = [prompt[i : i + 50] for i in range(0, len(prompt), 50)]
return [
{"title": f"Use case {idx+1}", "description": chunk}
for idx, chunk in enumerate(chunks)
]
# Main execution
fetch_and_publish_prompts()

121
scripts/platform_update/parse_prompts_and_submit_to_marketplace.py
Unescape View File

@ -1,121 +0,0 @@
import json
import os
from difflib import SequenceMatcher
import requests
from dotenv import load_dotenv
from loguru import logger
from supabase import Client, create_client
load_dotenv()
# Initialize Supabase client
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
# Swarms API URL and headers
SWARMS_API_URL = "https://swarms.world/api/add-prompt"
SWARMS_API_KEY = os.getenv("SWARMS_API_KEY")
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {SWARMS_API_KEY}",
}
# Configure logger
logger.add(
"fetch_and_publish_prompts.log", rotation="1 MB"
) # Log file with rotation
def fetch_and_publish_prompts():
logger.info("Starting to fetch and publish prompts.")
# Fetch data from Supabase
try:
response = (
supabase.table("swarms_framework_schema")
.select("*")
.execute()
)
rows = response.data
logger.info(f"Fetched {len(rows)} rows from Supabase.")
except Exception as e:
logger.error(f"Failed to fetch data from Supabase: {e}")
return
# Track published prompts to avoid duplicates
published_prompts = set()
for row in rows:
# Extract agent_name and system_prompt
data = row.get("data", {})
agent_name = data.get("agent_name")
system_prompt = data.get("system_prompt")
# Skip if either is missing or duplicate
if not agent_name or not system_prompt:
logger.warning(
f"Skipping row due to missing agent_name or system_prompt: {row}"
)
continue
if is_duplicate(system_prompt, published_prompts):
logger.info(
f"Skipping duplicate prompt for agent: {agent_name}"
)
continue
# Create the data payload for the marketplace
prompt_data = {
"name": f"{agent_name} - System Prompt",
"prompt": system_prompt,
"description": f"System prompt for agent {agent_name}.",
"useCases": extract_use_cases(system_prompt),
"tags": "agent, system-prompt",
}
# Publish to the marketplace
try:
response = requests.post(
SWARMS_API_URL,
headers=headers,
data=json.dumps(prompt_data),
)
if response.status_code == 200:
logger.info(
f"Successfully published prompt for agent: {agent_name}"
)
published_prompts.add(system_prompt)
else:
logger.error(
f"Failed to publish prompt for agent: {agent_name}. Response: {response.text}"
)
except Exception as e:
logger.error(
f"Exception occurred while publishing prompt for agent: {agent_name}. Error: {e}"
)
def is_duplicate(new_prompt, published_prompts):
"""Check if the prompt is a duplicate using semantic similarity."""
for prompt in published_prompts:
similarity = SequenceMatcher(None, new_prompt, prompt).ratio()
if (
similarity > 0.9
): # Threshold for considering prompts as duplicates
return True
return False
def extract_use_cases(prompt):
"""Extract use cases from the prompt by chunking it into meaningful segments."""
# This is a simple placeholder; you can use a more advanced method to extract use cases
chunks = [prompt[i : i + 50] for i in range(0, len(prompt), 50)]
return [
{"title": f"Use case {idx+1}", "description": chunk}
for idx, chunk in enumerate(chunks)
]
# Main execution
fetch_and_publish_prompts()

2
simple_example.py
Unescape View File

@ -4,4 +4,6 @@ Agent(
agent_name="Stock-Analysis-Agent", agent_name="Stock-Analysis-Agent",
model_name="gpt-4o-mini", model_name="gpt-4o-mini",
max_loops=1, max_loops=1,
interactive=False,
streaming_on=True,
).run("What are 5 hft algorithms") ).run("What are 5 hft algorithms")

433
sol_agent.py
Unescape View File

@ -0,0 +1,433 @@
import asyncio
import json
from dataclasses import asdict, dataclass
from datetime import datetime
from typing import Dict, List, Optional, Set
import aiohttp
import matplotlib.pyplot as plt
import networkx as nx
import websockets
from loguru import logger
from swarms import Agent
TREND_AGENT_PROMPT = """You are a specialized blockchain trend analysis agent. Your role:
1. Analyze transaction patterns in Solana blockchain data
2. Identify volume trends, price movements, and temporal patterns
3. Focus on whale movements and their market impact
4. Format findings in clear, structured JSON
5. Include confidence scores for each insight
6. Flag unusual patterns or anomalies
7. Provide historical context for significant movements
Output format:
{
"trends": [
{"pattern": str, "confidence": float, "impact": str}
],
"whale_activity": {...},
"temporal_analysis": {...}
}"""
RISK_AGENT_PROMPT = """You are a blockchain risk assessment specialist. Your tasks:
1. Identify suspicious transaction patterns
2. Monitor for known exploit signatures
3. Assess wallet clustering and relationship patterns
4. Evaluate transaction velocity and size anomalies
5. Check for bridge-related risks
6. Monitor smart contract interactions
7. Flag potential wash trading
Output format:
{
"risk_score": float,
"flags": [...],
"recommendations": [...]
}"""
SUMMARY_AGENT_PROMPT = """You are a blockchain data synthesis expert. Your responsibilities:
1. Combine insights from trend and risk analyses
2. Prioritize actionable intelligence
3. Highlight critical patterns
4. Generate executive summaries
5. Provide market context
6. Make predictions with confidence intervals
7. Suggest trading strategies based on data
Output format:
{
"key_insights": [...],
"market_impact": str,
"recommendations": {...}
}"""
@dataclass
class Transaction:
signature: str
timestamp: datetime
amount: float
from_address: str
to_address: str
class SolanaRPC:
def __init__(
self, endpoint="https://api.mainnet-beta.solana.com"
):
self.endpoint = endpoint
self.session = None
async def get_signatures(self, address: str) -> List[Dict]:
if not self.session:
self.session = aiohttp.ClientSession()
payload = {
"jsonrpc": "2.0",
"id": 1,
"method": "getSignaturesForAddress",
"params": [address, {"limit": 100}],
}
async with self.session.post(
self.endpoint, json=payload
) as response:
result = await response.json()
return result.get("result", [])
async def get_transaction(self, signature: str) -> Dict:
payload = {
"jsonrpc": "2.0",
"id": 1,
"method": "getTransaction",
"params": [
signature,
{
"encoding": "json",
"maxSupportedTransactionVersion": 0,
},
],
}
async with self.session.post(
self.endpoint, json=payload
) as response:
result = await response.json()
return result.get("result", {})
class AlertSystem:
def __init__(self, email: str, threshold: float = 1000.0):
self.email = email
self.threshold = threshold
self.smtp_server = "smtp.gmail.com"
self.smtp_port = 587
async def check_and_alert(
self, transaction: Transaction, risk_score: float
):
if transaction.amount > self.threshold or risk_score > 0.8:
await self.send_alert(transaction, risk_score)
async def send_alert(
self, transaction: Transaction, risk_score: float
):
# msg = MIMEText(
# f"High-risk transaction detected:\n"
# f"Amount: {transaction.amount} SOL\n"
# f"Risk Score: {risk_score}\n"
# f"Signature: {transaction.signature}"
# )
logger.info(
f"Alert sent for transaction {transaction.signature}"
)
class WalletClusterAnalyzer:
def __init__(self):
self.graph = nx.Graph()
self.known_wallets: Set[str] = set()
def update_graph(self, transaction: Transaction):
self.graph.add_edge(
transaction.from_address,
transaction.to_address,
weight=transaction.amount,
)
self.known_wallets.add(transaction.from_address)
self.known_wallets.add(transaction.to_address)
def identify_clusters(self) -> Dict:
communities = nx.community.greedy_modularity_communities(
self.graph
)
return {
"clusters": [list(c) for c in communities],
"central_wallets": [
wallet
for wallet in self.known_wallets
if self.graph.degree[wallet] > 5
],
}
class TransactionVisualizer:
def __init__(self):
self.transaction_history = []
def add_transaction(self, transaction: Transaction):
self.transaction_history.append(asdict(transaction))
def generate_volume_chart(self) -> str:
volumes = [tx["amount"] for tx in self.transaction_history]
plt.figure(figsize=(12, 6))
plt.plot(volumes)
plt.title("Transaction Volume Over Time")
plt.savefig("volume_chart.png")
return "volume_chart.png"
def generate_network_graph(
self, wallet_analyzer: WalletClusterAnalyzer
) -> str:
plt.figure(figsize=(15, 15))
pos = nx.spring_layout(wallet_analyzer.graph)
nx.draw(
wallet_analyzer.graph,
pos,
node_size=1000,
node_color="lightblue",
with_labels=True,
)
plt.savefig("network_graph.png")
return "network_graph.png"
class SolanaMultiAgentAnalyzer:
def __init__(
self,
min_amount: float = 50.0,
websocket_url: str = "wss://api.mainnet-beta.solana.com",
alert_email: str = None,
):
self.rpc = SolanaRPC()
self.websocket_url = websocket_url
self.min_amount = min_amount
self.transactions = []
self.wallet_analyzer = WalletClusterAnalyzer()
self.visualizer = TransactionVisualizer()
self.alert_system = (
AlertSystem(alert_email) if alert_email else None
)
self.trend_agent = Agent(
agent_name="trend-analyzer",
system_prompt=TREND_AGENT_PROMPT,
model_name="gpt-4o-mini",
max_loops=1,
streaming_on=True,
)
self.risk_agent = Agent(
agent_name="risk-analyzer",
system_prompt=RISK_AGENT_PROMPT,
model_name="gpt-4o-mini",
max_loops=1,
streaming_on=True,
)
self.summary_agent = Agent(
agent_name="summary-agent",
system_prompt=SUMMARY_AGENT_PROMPT,
model_name="gpt-4o-mini",
max_loops=1,
streaming_on=True,
)
logger.add(
"solana_analysis.log", rotation="500 MB", level="INFO"
)
async def start_websocket_stream(self):
async with websockets.connect(
self.websocket_url
) as websocket:
subscribe_message = {
"jsonrpc": "2.0",
"id": 1,
"method": "programSubscribe",
"params": [
"11111111111111111111111111111111",
{"encoding": "json", "commitment": "confirmed"},
],
}
await websocket.send(json.dumps(subscribe_message))
while True:
try:
msg = await websocket.recv()
transaction = await self.parse_websocket_message(
msg
)
if (
transaction
and transaction.amount >= self.min_amount
):
await self.process_transaction(transaction)
except Exception as e:
logger.error(f"Websocket error: {e}")
await asyncio.sleep(5)
async def parse_websocket_message(
self, msg: str
) -> Optional[Transaction]:
try:
data = json.loads(msg)
if "params" in data and "result" in data["params"]:
tx_data = data["params"]["result"]
return Transaction(
signature=tx_data["signature"],
timestamp=datetime.fromtimestamp(
tx_data["blockTime"]
),
amount=float(
tx_data["meta"]["postBalances"][0]
- tx_data["meta"]["preBalances"][0]
)
/ 1e9,
from_address=tx_data["transaction"]["message"][
"accountKeys"
][0],
to_address=tx_data["transaction"]["message"][
"accountKeys"
][1],
)
except Exception as e:
logger.error(f"Error parsing websocket message: {e}")
return None
async def process_transaction(self, transaction: Transaction):
self.wallet_analyzer.update_graph(transaction)
self.visualizer.add_transaction(transaction)
risk_analysis = await self.risk_agent.run(
f"Analyze risk for transaction: {json.dumps(asdict(transaction))}"
)
if self.alert_system:
await self.alert_system.check_and_alert(
transaction, risk_analysis.get("risk_score", 0)
)
async def fetch_transactions(self) -> List[Transaction]:
try:
signatures = await self.rpc.get_signatures(
"11111111111111111111111111111111"
)
transactions = []
for sig_info in signatures:
tx_data = await self.rpc.get_transaction(
sig_info["signature"]
)
if not tx_data or "meta" not in tx_data:
continue
pre_balances = tx_data["meta"]["preBalances"]
post_balances = tx_data["meta"]["postBalances"]
amount = abs(pre_balances[0] - post_balances[0]) / 1e9
if amount >= self.min_amount:
tx = Transaction(
signature=sig_info["signature"],
timestamp=datetime.fromtimestamp(
tx_data["blockTime"]
),
amount=amount,
from_address=tx_data["transaction"][
"message"
]["accountKeys"][0],
to_address=tx_data["transaction"]["message"][
"accountKeys"
][1],
)
transactions.append(tx)
return transactions
except Exception as e:
logger.error(f"Error fetching transactions: {e}")
return []
async def analyze_transactions(
self, transactions: List[Transaction]
) -> Dict:
tx_data = [asdict(tx) for tx in transactions]
cluster_data = self.wallet_analyzer.identify_clusters()
trend_analysis = await self.trend_agent.run(
f"Analyze trends in: {json.dumps(tx_data)}"
)
print(trend_analysis)
risk_analysis = await self.risk_agent.run(
f"Analyze risks in: {json.dumps({'transactions': tx_data, 'clusters': cluster_data})}"
)
print(risk_analysis)
summary = await self.summary_agent.run(
f"Synthesize insights from: {trend_analysis}, {risk_analysis}"
)
print(summary)
volume_chart = self.visualizer.generate_volume_chart()
network_graph = self.visualizer.generate_network_graph(
self.wallet_analyzer
)
return {
"transactions": tx_data,
"trend_analysis": trend_analysis,
"risk_analysis": risk_analysis,
"cluster_analysis": cluster_data,
"summary": summary,
"visualizations": {
"volume_chart": volume_chart,
"network_graph": network_graph,
},
}
async def run_continuous_analysis(self):
logger.info("Starting continuous analysis")
asyncio.create_task(self.start_websocket_stream())
while True:
try:
transactions = await self.fetch_transactions()
if transactions:
analysis = await self.analyze_transactions(
transactions
)
timestamp = datetime.now().strftime(
"%Y%m%d_%H%M%S"
)
with open(f"analysis_{timestamp}.json", "w") as f:
json.dump(analysis, f, indent=2, default=str)
logger.info(
f"Analysis completed: analysis_{timestamp}.json"
)
await asyncio.sleep(60)
except Exception as e:
logger.error(f"Error in analysis loop: {e}")
await asyncio.sleep(60)
# Add to __main__:
if __name__ == "__main__":
logger.info("Starting Solana analyzer...")
analyzer = SolanaMultiAgentAnalyzer(alert_email="your@email.com")
try:
asyncio.run(analyzer.run_continuous_analysis())
except Exception as e:
logger.error(f"Critical error: {e}")

2
swarms/agents/tool_agent.py
Unescape View File

@ -1,10 +1,12 @@
from typing import Any, Optional, Callable from typing import Any, Optional, Callable
from swarms.tools.json_former import Jsonformer from swarms.tools.json_former import Jsonformer
from swarms.utils.loguru_logger import initialize_logger from swarms.utils.loguru_logger import initialize_logger
from swarms.utils.lazy_loader import lazy_import_decorator
logger = initialize_logger(log_folder="tool_agent") logger = initialize_logger(log_folder="tool_agent")
@lazy_import_decorator
class ToolAgent: class ToolAgent:
""" """
Represents a tool agent that performs a specific task using a model and tokenizer. Represents a tool agent that performs a specific task using a model and tokenizer.

19
swarms/cli/create_agent.py
Unescape View File

@ -1,16 +1,6 @@
import os
from swarms.structs.agent import Agent from swarms.structs.agent import Agent
from swarm_models.popular_llms import OpenAIChat
from swarms.structs.agent_registry import AgentRegistry from swarms.structs.agent_registry import AgentRegistry
# Get the OpenAI API key from the environment variable
api_key = os.getenv("OPENAI_API_KEY")
# Create an instance of the OpenAIChat class
model = OpenAIChat(
api_key=api_key, model_name="gpt-4o-mini", temperature=0.1
)
# Registry of agents # Registry of agents
agent_registry = AgentRegistry( agent_registry = AgentRegistry(
@ -19,7 +9,12 @@ agent_registry = AgentRegistry(
) )
def create_agent(name: str, system_prompt: str, max_loops: int = 1): def create_agent(
name: str,
system_prompt: str,
max_loops: int = 1,
model_name: str = "gpt-4o",
):
""" """
Create and initialize an agent with the given parameters. Create and initialize an agent with the given parameters.
@ -36,7 +31,7 @@ def create_agent(name: str, system_prompt: str, max_loops: int = 1):
agent = Agent( agent = Agent(
agent_name=name, agent_name=name,
system_prompt=system_prompt, system_prompt=system_prompt,
llm=model, model_name=model_name,
max_loops=max_loops, max_loops=max_loops,
autosave=True, autosave=True,
dashboard=False, dashboard=False,

4
swarms/structs/__init__.py
Unescape View File

@ -19,7 +19,6 @@ from swarms.structs.majority_voting import (
most_frequent, most_frequent,
parse_code_completion, parse_code_completion,
) )
from swarms.structs.message import Message
from swarms.structs.mixture_of_agents import MixtureOfAgents from swarms.structs.mixture_of_agents import MixtureOfAgents
from swarms.structs.multi_agent_collab import MultiAgentCollaboration from swarms.structs.multi_agent_collab import MultiAgentCollaboration
from swarms.structs.multi_agent_exec import ( from swarms.structs.multi_agent_exec import (
@ -39,7 +38,6 @@ from swarms.structs.round_robin import RoundRobinSwarm
from swarms.structs.sequential_workflow import SequentialWorkflow from swarms.structs.sequential_workflow import SequentialWorkflow
from swarms.structs.spreadsheet_swarm import SpreadSheetSwarm from swarms.structs.spreadsheet_swarm import SpreadSheetSwarm
from swarms.structs.swarm_arange import SwarmRearrange from swarms.structs.swarm_arange import SwarmRearrange
from swarms.structs.swarm_net import SwarmNetwork
from swarms.structs.swarm_router import ( from swarms.structs.swarm_router import (
SwarmRouter, SwarmRouter,
SwarmType, SwarmType,
@ -90,9 +88,7 @@ __all__ = [
"majority_voting", "majority_voting",
"most_frequent", "most_frequent",
"parse_code_completion", "parse_code_completion",
"Message",
"MultiAgentCollaboration", "MultiAgentCollaboration",
"SwarmNetwork",
"AgentRearrange", "AgentRearrange",
"rearrange", "rearrange",
"RoundRobinSwarm", "RoundRobinSwarm",

15
swarms/structs/agent_router.py
Unescape View File

@ -1,14 +1,19 @@
from typing import List, Optional from typing import List, Optional
import chromadb
from tenacity import retry, stop_after_attempt, wait_exponential from tenacity import retry, stop_after_attempt, wait_exponential
from typing import Union, Callable, Any from typing import Union, Callable, Any
from swarms import Agent from swarms import Agent
from swarms.utils.loguru_logger import initialize_logger from swarms.utils.loguru_logger import initialize_logger
from swarms.utils.lazy_loader import lazy_import_decorator
from swarms.utils.auto_download_check_packages import (
auto_check_and_download_package,
)
logger = initialize_logger(log_folder="agent_router") logger = initialize_logger(log_folder="agent_router")
@lazy_import_decorator
class AgentRouter: class AgentRouter:
""" """
Initialize the AgentRouter. Initialize the AgentRouter.
@ -29,6 +34,14 @@ class AgentRouter:
*args, *args,
**kwargs, **kwargs,
): ):
try:
import chromadb
except ImportError:
auto_check_and_download_package(
"chromadb", package_manager="pip", upgrade=True
)
import chromadb
self.collection_name = collection_name self.collection_name = collection_name
self.n_agents = n_agents self.n_agents = n_agents
self.persist_directory = persist_directory self.persist_directory = persist_directory

5
swarms/structs/base_swarm.py
Unescape View File

@ -16,7 +16,6 @@ from typing import (
import yaml import yaml
from swarms_memory import BaseVectorDatabase
from swarms.structs.agent import Agent from swarms.structs.agent import Agent
from swarms.structs.conversation import Conversation from swarms.structs.conversation import Conversation
from swarms.structs.omni_agent_types import AgentType from swarms.structs.omni_agent_types import AgentType
@ -98,9 +97,7 @@ class BaseSwarm(ABC):
agentops_on: Optional[bool] = False, agentops_on: Optional[bool] = False,
speaker_selection_func: Optional[Callable] = None, speaker_selection_func: Optional[Callable] = None,
rules: Optional[str] = None, rules: Optional[str] = None,
collective_memory_system: Optional[ collective_memory_system: Optional[Any] = False,
BaseVectorDatabase
] = False,
agent_ops_on: bool = False, agent_ops_on: bool = False,
output_schema: Optional[BaseModel] = None, output_schema: Optional[BaseModel] = None,
*args, *args,

24
swarms/structs/graph_swarm.py
Unescape View File

@ -1,10 +1,3 @@
"""
GraphSwarm: A production-grade framework for orchestrating swarms of agents
Author: Claude
License: MIT
Version: 2.0.0
"""
import asyncio import asyncio
import json import json
import time import time
@ -12,13 +5,13 @@ from concurrent.futures import ThreadPoolExecutor
from datetime import datetime from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple, Union from typing import Any, Dict, List, Optional, Tuple, Union
import chromadb
import networkx as nx import networkx as nx
from loguru import logger from loguru import logger
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from swarms.utils.auto_download_check_packages import (
from swarms import Agent auto_check_and_download_package,
)
from swarms.structs.agent import Agent
# Configure logging # Configure logging
logger.add( logger.add(
@ -57,6 +50,15 @@ class SwarmMemory:
def __init__(self, collection_name: str = "swarm_memories"): def __init__(self, collection_name: str = "swarm_memories"):
"""Initialize SwarmMemory with ChromaDB.""" """Initialize SwarmMemory with ChromaDB."""
try:
import chromadb
except ImportError:
auto_check_and_download_package(
"chromadb", package_manager="pip", upgrade=True
)
import chromadb
self.client = chromadb.Client() self.client = chromadb.Client()
# Get or create collection # Get or create collection

123
swarms/structs/groupchat_new.py
Unescape View File

@ -3,7 +3,6 @@ import asyncio
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from typing import List, Dict, Any from typing import List, Dict, Any
from swarms import Agent from swarms import Agent
from swarm_models import OpenAIChat
from dotenv import load_dotenv from dotenv import load_dotenv
from swarms.utils.formatter import formatter from swarms.utils.formatter import formatter
@ -181,64 +180,64 @@ class GroupChat:
] ]
# Example Usage # # Example Usage
if __name__ == "__main__": # if __name__ == "__main__":
load_dotenv() # load_dotenv()
# Get the OpenAI API key from the environment variable # # Get the OpenAI API key from the environment variable
api_key = os.getenv("OPENAI_API_KEY") # api_key = os.getenv("OPENAI_API_KEY")
# Create an instance of the OpenAIChat class # # Create an instance of the OpenAIChat class
model = OpenAIChat( # model = OpenAIChat(
openai_api_key=api_key, # openai_api_key=api_key,
model_name="gpt-4o-mini", # model_name="gpt-4o-mini",
temperature=0.1, # temperature=0.1,
) # )
# Example agents # # Example agents
agent1 = Agent( # agent1 = Agent(
agent_name="Financial-Analysis-Agent", # agent_name="Financial-Analysis-Agent",
system_prompt="You are a financial analyst specializing in investment strategies.", # system_prompt="You are a financial analyst specializing in investment strategies.",
llm=model, # llm=model,
max_loops=1, # max_loops=1,
autosave=False, # autosave=False,
dashboard=False, # dashboard=False,
verbose=True, # verbose=True,
dynamic_temperature_enabled=True, # dynamic_temperature_enabled=True,
user_name="swarms_corp", # user_name="swarms_corp",
retry_attempts=1, # retry_attempts=1,
context_length=200000, # context_length=200000,
output_type="string", # output_type="string",
streaming_on=False, # streaming_on=False,
) # )
agent2 = Agent( # agent2 = Agent(
agent_name="Tax-Adviser-Agent", # agent_name="Tax-Adviser-Agent",
system_prompt="You are a tax adviser who provides clear and concise guidance on tax-related queries.", # system_prompt="You are a tax adviser who provides clear and concise guidance on tax-related queries.",
llm=model, # llm=model,
max_loops=1, # max_loops=1,
autosave=False, # autosave=False,
dashboard=False, # dashboard=False,
verbose=True, # verbose=True,
dynamic_temperature_enabled=True, # dynamic_temperature_enabled=True,
user_name="swarms_corp", # user_name="swarms_corp",
retry_attempts=1, # retry_attempts=1,
context_length=200000, # context_length=200000,
output_type="string", # output_type="string",
streaming_on=False, # streaming_on=False,
) # )
# Create group chat # # Create group chat
group_chat = GroupChat( # group_chat = GroupChat(
name="Financial Discussion", # name="Financial Discussion",
description="A group chat for financial analysis and tax advice.", # description="A group chat for financial analysis and tax advice.",
agents=[agent1, agent2], # agents=[agent1, agent2],
) # )
# Run the group chat # # Run the group chat
asyncio.run( # asyncio.run(
group_chat.run( # group_chat.run(
"How can I establish a ROTH IRA to buy stocks and get a tax break? What are the criteria? What do you guys think?" # "How can I establish a ROTH IRA to buy stocks and get a tax break? What are the criteria? What do you guys think?"
) # )
) # )

29
swarms/structs/message.py
Unescape View File

@ -1,29 +0,0 @@
from typing import Dict, Optional
from datetime import datetime
from pydantic import BaseModel, Field
class Message(BaseModel):
"""
Represents a message with timestamp and optional metadata.
Usage
--------------
mes = Message(
sender = "Kye",
content = "message"
)
print(mes)
"""
timestamp: datetime = Field(default_factory=datetime.now)
sender: str
content: str
metadata: Optional[Dict[str, str]] = {}
def __repr__(self) -> str:
"""
__repr__ means...
"""
return f"{self.timestamp} - {self.sender}: {self.content}"

2
swarms/structs/multi_agent_exec.py
Unescape View File

@ -3,7 +3,7 @@ from concurrent.futures import ThreadPoolExecutor
import psutil import psutil
from dataclasses import dataclass from dataclasses import dataclass
import threading import threading
from typing import List, Union, Any, Callable from typing import List, Any
from multiprocessing import cpu_count from multiprocessing import cpu_count
import os import os

15
swarms/structs/output_types.py
Unescape View File

@ -0,0 +1,15 @@
from typing import Literal
# Literal of output types
# Literal of output types
OutputType = Literal[
"all",
"final",
"list",
"dict",
".json",
".md",
".txt",
".yaml",
".toml",
]

19
swarms/structs/rearrange.py
Unescape View File

@ -3,10 +3,9 @@ import traceback
import uuid import uuid
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from datetime import datetime from datetime import datetime
from typing import Callable, Dict, List, Literal, Optional from typing import Any, Callable, Dict, List, Optional
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from swarms_memory import BaseVectorDatabase
from swarms.schemas.agent_step_schemas import ManySteps from swarms.schemas.agent_step_schemas import ManySteps
from swarms.structs.agent import Agent from swarms.structs.agent import Agent
@ -17,22 +16,10 @@ from swarms.utils.loguru_logger import initialize_logger
from swarms.utils.wrapper_clusterop import ( from swarms.utils.wrapper_clusterop import (
exec_callable_with_clusterops, exec_callable_with_clusterops,
) )
from swarms.structs.output_types import OutputType
logger = initialize_logger(log_folder="rearrange") logger = initialize_logger(log_folder="rearrange")
# Literal of output types
OutputType = Literal[
"all",
"final",
"list",
"dict",
".json",
".md",
".txt",
".yaml",
".toml",
]
def swarm_id(): def swarm_id():
return uuid.uuid4().hex return uuid.uuid4().hex
@ -112,7 +99,7 @@ class AgentRearrange(BaseSwarm):
flow: str = None, flow: str = None,
max_loops: int = 1, max_loops: int = 1,
verbose: bool = True, verbose: bool = True,
memory_system: BaseVectorDatabase = None, memory_system: Any = None,
human_in_the_loop: bool = False, human_in_the_loop: bool = False,
custom_human_in_the_loop: Optional[ custom_human_in_the_loop: Optional[
Callable[[str], str] Callable[[str], str]

3
swarms/structs/sequential_workflow.py
Unescape View File

@ -1,6 +1,7 @@
from typing import List, Optional from typing import List, Optional
from swarms.structs.agent import Agent from swarms.structs.agent import Agent
from swarms.structs.rearrange import AgentRearrange, OutputType from swarms.structs.rearrange import AgentRearrange
from swarms.structs.output_types import OutputType
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor, as_completed
from swarms.utils.loguru_logger import initialize_logger from swarms.utils.loguru_logger import initialize_logger

37
swarms/structs/swarm_matcher.py
Unescape View File

@ -1,11 +1,14 @@
from typing import List, Tuple, Optional from typing import List, Tuple, Optional
import numpy as np import numpy as np
import torch from swarms.utils.lazy_loader import lazy_import_decorator
from transformers import AutoTokenizer, AutoModel
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
import json import json
from tenacity import retry, stop_after_attempt, wait_exponential from tenacity import retry, stop_after_attempt, wait_exponential
from swarms.utils.loguru_logger import initialize_logger from swarms.utils.loguru_logger import initialize_logger
from swarms.utils.auto_download_check_packages import (
auto_check_and_download_package,
)
logger = initialize_logger(log_folder="swarm_matcher") logger = initialize_logger(log_folder="swarm_matcher")
@ -25,6 +28,7 @@ class SwarmMatcherConfig(BaseModel):
) )
@lazy_import_decorator
class SwarmMatcher: class SwarmMatcher:
""" """
A class for matching tasks to swarm types based on their descriptions. A class for matching tasks to swarm types based on their descriptions.
@ -41,12 +45,34 @@ class SwarmMatcher:
""" """
logger.add("swarm_matcher_debug.log", level="DEBUG") logger.add("swarm_matcher_debug.log", level="DEBUG")
logger.debug("Initializing SwarmMatcher") logger.debug("Initializing SwarmMatcher")
try:
import torch
except ImportError:
auto_check_and_download_package(
"torch", package_manager="pip", upgrade=True
)
import torch
try:
import transformers
except ImportError:
auto_check_and_download_package(
"transformers", package_manager="pip", upgrade=True
)
import transformers
self.torch = torch
try: try:
self.config = config self.config = config
self.tokenizer = AutoTokenizer.from_pretrained( self.tokenizer = (
transformers.AutoTokenizer.from_pretrained(
config.model_name
)
)
self.model = transformers.AutoModel.from_pretrained(
config.model_name config.model_name
) )
self.model = AutoModel.from_pretrained(config.model_name)
self.swarm_types: List[SwarmType] = [] self.swarm_types: List[SwarmType] = []
logger.debug("SwarmMatcher initialized successfully") logger.debug("SwarmMatcher initialized successfully")
except Exception as e: except Exception as e:
@ -76,7 +102,7 @@ class SwarmMatcher:
truncation=True, truncation=True,
max_length=512, max_length=512,
) )
with torch.no_grad(): with self.torch.no_grad():
outputs = self.model(**inputs) outputs = self.model(**inputs)
embedding = ( embedding = (
outputs.last_hidden_state.mean(dim=1) outputs.last_hidden_state.mean(dim=1)
@ -244,6 +270,7 @@ def initialize_swarm_types(matcher: SwarmMatcher):
logger.debug("Swarm types initialized") logger.debug("Swarm types initialized")
@lazy_import_decorator
def swarm_matcher(task: str, *args, **kwargs): def swarm_matcher(task: str, *args, **kwargs):
""" """
Runs the SwarmMatcher example with predefined tasks and swarm types. Runs the SwarmMatcher example with predefined tasks and swarm types.

511
swarms/structs/swarm_net.py
Unescape View File

@ -1,511 +0,0 @@
"""
Todo
- [ ] Test the new api feature
- [ ] Add the agent schema for every agent -- following OpenAI assistaants schema
- [ ] then add the swarm schema for the swarm url: /v1/swarms/{swarm_name}/agents/{agent_id}
- [ ] Add the agent schema for the agent url: /v1/swarms/{swarm_name}/agents/{agent_id}
"""
import asyncio
import multiprocessing
import queue
import threading
from typing import List, Optional
import tenacity
# from fastapi import FastAPI
from pydantic import BaseModel
from swarms.structs.agent import Agent
from swarms.structs.base_swarm import BaseSwarm
from swarms.utils.loguru_logger import initialize_logger
logger = initialize_logger("swarm-network")
# Pydantic models
class TaskRequest(BaseModel):
task: str
# Pydantic models
class TaskResponse(BaseModel):
result: str
class AgentInfo(BaseModel):
agent_name: str
agent_description: str
class SwarmInfo(BaseModel):
swarm_name: str
swarm_description: str
agents: List[AgentInfo]
# Helper function to get the number of workers
def get_number_of_workers():
return multiprocessing.cpu_count()
# [TODO] Add the agent schema for every agent -- following OpenAI assistaants schema
class SwarmNetwork(BaseSwarm):
"""
SwarmNetwork class
The SwarmNetwork 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.
For example, if the number of pending tasks is greater than the
number of agents in the pool, the SwarmNetwork will scale up the
pool by adding new agents. If the number of pending tasks is less
than the number of agents in the pool, the SwarmNetwork will scale
down the pool by removing agents.
The SwarmNetwork class also provides a simple API for interacting
with the agents pool. The API is implemented using the Flask
framework and is enabled by default. The API can be disabled by
setting the `api_enabled` parameter to False.
Features:
- Agent pool management
- Task queue management
- Agent health monitoring
- Agent pool scaling
- Simple API for interacting with the agent pool
- Simple API for interacting with the task queue
- Simple API for interacting with the agent health monitor
- Simple API for interacting with the agent pool scaler
- Create APIs for each agent in the pool (optional)
- Run each agent on it's own thread
- Run each agent on it's own process
- Run each agent on it's own container
- Run each agent on it's own machine
- Run each agent on it's own cluster
Attributes:
task_queue (queue.Queue): A queue for storing tasks.
idle_threshold (float): The idle threshold for the agents.
busy_threshold (float): The busy threshold for the agents.
agents (List[Agent]): A list of agents in the pool.
api_enabled (bool): A flag to enable/disable the API.
logging_enabled (bool): A flag to enable/disable logging.
Example:
>>> from swarms.structs.agent import Agent
>>> from swarms.structs.swarm_net import SwarmNetwork
>>> agent = Agent()
>>> swarm = SwarmNetwork(agents=[agent])
>>> swarm.add_task("task")
>>> swarm.run()
"""
def __init__(
self,
name: str = None,
description: str = None,
agents: List[Agent] = None,
idle_threshold: float = 0.2,
busy_threshold: float = 0.7,
api_enabled: Optional[bool] = False,
logging_enabled: Optional[bool] = False,
api_on: Optional[bool] = False,
host: str = "0.0.0.0",
port: int = 8000,
swarm_callable: Optional[callable] = None,
*args,
**kwargs,
):
super().__init__(agents=agents, *args, **kwargs)
self.name = name
self.description = description
self.agents = agents
self.task_queue = queue.Queue()
self.idle_threshold = idle_threshold
self.busy_threshold = busy_threshold
self.lock = threading.Lock()
self.api_enabled = api_enabled
self.logging_enabled = logging_enabled
self.host = host
self.port = port
self.swarm_callable = swarm_callable
# Ensure that the agents list is not empty
if not agents:
raise ValueError("The agents list cannot be empty")
# Create a dictionary of agents for easy access
self.agent_dict = {agent.id: agent for agent in agents}
# # Create the FastAPI instance
# if api_on is True:
# logger.info("Creating FastAPI instance")
# self.app = FastAPI(debug=True, *args, **kwargs)
# self.app.add_middleware(
# CORSMiddleware,
# allow_origins=["*"],
# allow_credentials=True,
# allow_methods=["*"],
# allow_headers=["*"],
# )
# logger.info("Routes set for creation")
# self._create_routes()
def add_task(self, task):
"""Add task to the task queue
Args:
task (_type_): _description_
Example:
>>> from swarms.structs.agent import Agent
>>> from swarms.structs.swarm_net import SwarmNetwork
>>> agent = Agent()
>>> swarm = SwarmNetwork(agents=[agent])
>>> swarm.add_task("task")
"""
self.logger.info(f"Adding task {task} to queue")
try:
self.task_queue.put(task)
self.logger.info(f"Task {task} added to queue")
except Exception as error:
print(
f"Error adding task to queue: {error} try again with"
" a new task"
)
raise error
async def async_add_task(self, task):
"""Add task to the task queue
Args:
task (_type_): _description_
Example:
>>> from swarms.structs.agent import Agent
>>> from swarms.structs.swarm_net import SwarmNetwork
>>> agent = Agent()
>>> swarm = SwarmNetwork(agents=[agent])
>>> swarm.add_task("task")
"""
self.logger.info(
f"Adding task {task} to queue asynchronously"
)
try:
# Add task to queue asynchronously with asyncio
loop = asyncio.get_running_loop()
await loop.run_in_executor(
None, self.task_queue.put, task
)
self.logger.info(f"Task {task} added to queue")
except Exception as error:
print(
f"Error adding task to queue: {error} try again with"
" a new task"
)
raise error
# def _create_routes(self) -> None:
# """
# Creates the routes for the API.
# """
# # Extensive logginbg
# logger.info("Creating routes for the API")
# # Routes available
# logger.info(
# "Routes available: /v1/swarms, /v1/health, /v1/swarms/{swarm_name}/agents/{agent_id}, /v1/swarms/{swarm_name}/run"
# )
# @self.app.get("/v1/swarms", response_model=SwarmInfo)
# async def get_swarms() -> SwarmInfo:
# try:
# logger.info("Getting swarm information")
# return SwarmInfo(
# swarm_name=self.swarm_name,
# swarm_description=self.swarm_description,
# agents=[
# AgentInfo(
# agent_name=agent.agent_name,
# agent_description=agent.agent_description,
# )
# for agent in self.agents
# ],
# )
# except Exception as e:
# logger.error(f"Error getting swarm information: {str(e)}")
# raise HTTPException(
# status_code=500, detail="Internal Server Error"
# )
# @self.app.get("/v1/health")
# async def get_health() -> Dict[str, str]:
# try:
# logger.info("Checking health status")
# return {"status": "healthy"}
# except Exception as e:
# logger.error(f"Error checking health status: {str(e)}")
# raise HTTPException(
# status_code=500, detail="Internal Server Error"
# )
# @self.app.get(f"/v1/swarms/{self.swarm_name}/agents/{{agent_id}}")
# async def get_agent_info(agent_id: str) -> AgentInfo:
# try:
# logger.info(f"Getting information for agent {agent_id}")
# agent = self.agent_dict.get(agent_id)
# if not agent:
# raise HTTPException(
# status_code=404, detail="Agent not found"
# )
# return AgentInfo(
# agent_name=agent.agent_name,
# agent_description=agent.agent_description,
# )
# except Exception as e:
# logger.error(f"Error getting agent information: {str(e)}")
# raise HTTPException(
# status_code=500, detail="Internal Server Error"
# )
# @self.app.post(
# f"/v1/swarms/{self.swarm_name}/agents/{{agent_id}}/run",
# response_model=TaskResponse,
# )
# async def run_agent_task(
# task_request: TaskRequest,
# ) -> TaskResponse:
# try:
# logger.info("Running agent task")
# # Assuming only one agent in the swarm for this example
# agent = self.agents[0]
# logger.info(f"Running agent task: {task_request.task}")
# result = agent.run(task_request.task)
# return TaskResponse(result=result)
# except Exception as e:
# logger.error(f"Error running agent task: {str(e)}")
# raise HTTPException(
# status_code=500, detail="Internal Server Error"
# )
# def get_app(self) -> FastAPI:
# """
# Returns the FastAPI instance.
# Returns:
# FastAPI: The FastAPI instance.
# """
# return self.app
def run_single_agent(
self, agent_id, task: Optional[str], *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_
"""
self.logger.info(f"Running task {task} on agent {agent_id}")
try:
for agent in self.agents:
if agent.id == agent_id:
out = agent.run(task, *args, **kwargs)
return out
except Exception as error:
self.logger.error(f"Error running task on agent: {error}")
raise error
def run_many_agents(
self, task: Optional[str] = None, *args, **kwargs
) -> List:
"""Run the task on all agents
Args:
task (str, optional): _description_. Defaults to None.
Returns:
List: _description_
"""
self.logger.info(f"Running task {task} on all agents")
try:
return [
agent.run(task, *args, **kwargs)
for agent in self.agents
]
except Exception as error:
logger.error(f"Error running task on agents: {error}")
raise error
def list_agents(self):
"""List all agents."""
self.logger.info("[Listing all active agents]")
try:
# Assuming self.agents is a list of agent objects
for agent in self.agents:
self.logger.info(
f"[Agent] [ID: {agent.id}] [Name:"
f" {agent.agent_name}] [Description:"
f" {agent.agent_description}] [Status: Running]"
)
except Exception as error:
self.logger.error(f"Error listing agents: {error}")
raise
def get_agent(self, agent_id):
"""Get agent by id
Args:
agent_id (_type_): _description_
Returns:
_type_: _description_
"""
self.logger.info(f"Getting agent {agent_id}")
try:
for agent in self.agents:
if agent.id == agent_id:
return agent
raise ValueError(f"No agent found with ID {agent_id}")
except Exception as error:
self.logger.error(f"Error getting agent: {error}")
raise error
def add_agent(self, agent: Agent):
"""Add agent to the agent pool
Args:
agent (_type_): _description_
"""
self.logger.info(f"Adding agent {agent} to pool")
try:
self.agents.append(agent)
except Exception as error:
print(f"Error adding agent to pool: {error}")
raise error
def remove_agent(self, agent_id):
"""Remove agent from the agent pool
Args:
agent_id (_type_): _description_
"""
self.logger.info(f"Removing agent {agent_id} from pool")
try:
for agent in self.agents:
if agent.id == agent_id:
self.agents.remove(agent)
return
raise ValueError(f"No agent found with ID {agent_id}")
except Exception as error:
print(f"Error removing agent from pool: {error}")
raise error
async def async_remove_agent(self, agent_id):
"""Remove agent from the agent pool
Args:
agent_id (_type_): _description_
"""
self.logger.info(f"Removing agent {agent_id} from pool")
try:
# Remove agent from pool asynchronously with asyncio
loop = asyncio.get_running_loop()
await loop.run_in_executor(
None, self.remove_agent, agent_id
)
except Exception as error:
print(f"Error removing agent from pool: {error}")
raise error
def scale_up(self, num_agents: int = 1):
"""Scale up the agent pool
Args:
num_agents (int, optional): _description_. Defaults to 1.
"""
self.logger.info(f"Scaling up agent pool by {num_agents}")
try:
for _ in range(num_agents):
self.agents.append(Agent())
except Exception as error:
print(f"Error scaling up agent pool: {error}")
raise error
def scale_down(self, num_agents: int = 1):
"""Scale down the agent pool
Args:
num_agents (int, optional): _description_. Defaults to 1.
"""
for _ in range(num_agents):
self.agents.pop()
@tenacity.retry(
wait=tenacity.wait_fixed(1),
stop=tenacity.stop_after_attempt(3),
retry=tenacity.retry_if_exception_type(Exception),
)
def run(self, *args, **kwargs):
"""run the swarm network"""
app = self.get_app()
try:
import uvicorn
logger.info(
f"Running the swarm network with {len(self.agents)} on {self.host}:{self.port}"
)
uvicorn.run(
app,
host=self.host,
port=self.port,
# workers=get_number_of_workers(),
*args,
**kwargs,
)
return app
except Exception as error:
logger.error(f"Error running the swarm network: {error}")
raise error
# # # Example usage
# if __name__ == "__main__":
# agent1 = Agent(
# agent_name="Covid-19-Chat",
# agent_description="This agent provides information about COVID-19 symptoms.",
# llm=OpenAIChat(),
# max_loops="auto",
# autosave=True,
# verbose=True,
# stopping_condition="finish",
# )
# agents = [agent1] # Add more agents as needed
# swarm_name = "HealthSwarm"
# swarm_description = (
# "A swarm of agents providing health-related information."
# )
# agent_api = SwarmNetwork(swarm_name, swarm_description, agents)
# agent_api.run()

46
swarms/structs/tree_swarm.py
Unescape View File

@ -4,17 +4,14 @@ from datetime import datetime
from typing import Any, List, Optional from typing import Any, List, Optional
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from sentence_transformers import SentenceTransformer, util
from swarms.structs.agent import Agent from swarms.structs.agent import Agent
from swarms.utils.loguru_logger import initialize_logger from swarms.utils.loguru_logger import initialize_logger
from swarms.utils.auto_download_check_packages import (
auto_check_and_download_package,
)
logger = initialize_logger(log_folder="tree_swarm")
# Pretrained model for embeddings logger = initialize_logger(log_folder="tree_swarm")
embedding_model = SentenceTransformer(
"all-MiniLM-L6-v2"
) # A small, fast model for embedding
# Pydantic Models for Logging # Pydantic Models for Logging
@ -68,7 +65,7 @@ class TreeAgent(Agent):
name: str = None, name: str = None,
description: str = None, description: str = None,
system_prompt: str = None, system_prompt: str = None,
llm: callable = None, model_name: str = "gpt-4o",
agent_name: Optional[str] = None, agent_name: Optional[str] = None,
*args, *args,
**kwargs, **kwargs,
@ -78,12 +75,29 @@ class TreeAgent(Agent):
name=name, name=name,
description=description, description=description,
system_prompt=system_prompt, system_prompt=system_prompt,
llm=llm, model_name=model_name,
agent_name=agent_name, agent_name=agent_name,
*args, *args,
**kwargs, **kwargs,
) )
self.system_prompt_embedding = embedding_model.encode(
try:
import sentence_transformers
except ImportError:
auto_check_and_download_package(
"sentence-transformers", package_manager="pip"
)
import sentence_transformers
self.sentence_transformers = sentence_transformers
# Pretrained model for embeddings
self.embedding_model = (
sentence_transformers.SentenceTransformer(
"all-MiniLM-L6-v2"
)
)
self.system_prompt_embedding = self.embedding_model.encode(
system_prompt, convert_to_tensor=True system_prompt, convert_to_tensor=True
) )
@ -103,7 +117,7 @@ class TreeAgent(Agent):
Returns: Returns:
float: Distance score between 0 and 1, with 0 being close and 1 being far. float: Distance score between 0 and 1, with 0 being close and 1 being far.
""" """
similarity = util.pytorch_cos_sim( similarity = self.sentence_transformers.util.pytorch_cos_sim(
self.system_prompt_embedding, self.system_prompt_embedding,
other_agent.system_prompt_embedding, other_agent.system_prompt_embedding,
).item() ).item()
@ -154,12 +168,14 @@ class TreeAgent(Agent):
# Perform embedding similarity match if keyword match is not found # Perform embedding similarity match if keyword match is not found
if not keyword_match: if not keyword_match:
task_embedding = embedding_model.encode( task_embedding = self.embedding_model.encode(
task, convert_to_tensor=True task, convert_to_tensor=True
) )
similarity = util.pytorch_cos_sim( similarity = (
self.system_prompt_embedding, task_embedding self.sentence_transformers.util.pytorch_cos_sim(
).item() self.system_prompt_embedding, task_embedding
).item()
)
logger.info( logger.info(
f"Semantic similarity between task and {self.agent_name}: {similarity:.2f}" f"Semantic similarity between task and {self.agent_name}: {similarity:.2f}"
) )

2
swarms/tools/__init__.py
Unescape View File

@ -28,6 +28,7 @@ from swarms.tools.cohere_func_call_schema import (
ParameterDefinition, ParameterDefinition,
) )
from swarms.tools.tool_registry import ToolStorage, tool_registry from swarms.tools.tool_registry import ToolStorage, tool_registry
from swarms.tools.json_utils import base_model_to_json
__all__ = [ __all__ = [
@ -51,4 +52,5 @@ __all__ = [
"ParameterDefinition", "ParameterDefinition",
"ToolStorage", "ToolStorage",
"tool_registry", "tool_registry",
"base_model_to_json",
] ]

19
swarms/tools/json_former.py
Unescape View File

@ -1,7 +1,7 @@
import json import json
from typing import Any, Dict, List, Union from typing import Any, Dict, List, Union
from transformers import PreTrainedModel, PreTrainedTokenizer from swarms.utils.lazy_loader import lazy_import_decorator
from pydantic import BaseModel from pydantic import BaseModel
from swarms.tools.logits_processor import ( from swarms.tools.logits_processor import (
NumberStoppingCriteria, NumberStoppingCriteria,
@ -9,10 +9,23 @@ from swarms.tools.logits_processor import (
StringStoppingCriteria, StringStoppingCriteria,
) )
from swarm_models.base_llm import BaseLLM from swarm_models.base_llm import BaseLLM
from swarms.utils.auto_download_check_packages import (
auto_check_and_download_package,
)
try:
import transformers
except ImportError:
auto_check_and_download_package(
"transformers", package_manager="pip"
)
import transformers
GENERATION_MARKER = "|GENERATION|" GENERATION_MARKER = "|GENERATION|"
@lazy_import_decorator
class Jsonformer: class Jsonformer:
""" """
Initializes the FormatTools class. Initializes the FormatTools class.
@ -35,8 +48,8 @@ class Jsonformer:
def __init__( def __init__(
self, self,
model: PreTrainedModel = None, model: transformers.PreTrainedModel = None, # type: ignore
tokenizer: PreTrainedTokenizer = None, tokenizer: transformers.PreTrainedTokenizer = None, # type: ignore
json_schema: Union[Dict[str, Any], BaseModel] = None, json_schema: Union[Dict[str, Any], BaseModel] = None,
schemas: List[Union[Dict[str, Any], BaseModel]] = [], schemas: List[Union[Dict[str, Any], BaseModel]] = [],
prompt: str = None, prompt: str = None,

42
swarms/tools/logits_processor.py
Unescape View File

@ -1,21 +1,35 @@
import torch from swarms.utils.auto_download_check_packages import (
from transformers import ( auto_check_and_download_package,
LogitsWarper,
PreTrainedTokenizer,
StoppingCriteria,
) )
class StringStoppingCriteria(StoppingCriteria): try:
import torch
except ImportError:
auto_check_and_download_package(
"torch", package_manager="pip", upgrade=True
)
import torch
try:
import transformers
except ImportError:
auto_check_and_download_package(
"transformers", package_manager="pip", upgrade=True
)
import transformers
class StringStoppingCriteria(transformers.StoppingCriteria):
def __init__( def __init__(
self, tokenizer: PreTrainedTokenizer, prompt_length: int self, tokenizer: transformers.PreTrainedTokenizer, prompt_length: int # type: ignore
): ):
self.tokenizer = tokenizer self.tokenizer = tokenizer
self.prompt_length = prompt_length self.prompt_length = prompt_length
def __call__( def __call__(
self, self,
input_ids: torch.LongTensor, input_ids: torch.LongTensor, # type: ignore
_, _,
) -> bool: ) -> bool:
if len(input_ids[0]) <= self.prompt_length: if len(input_ids[0]) <= self.prompt_length:
@ -31,10 +45,10 @@ class StringStoppingCriteria(StoppingCriteria):
return result return result
class NumberStoppingCriteria(StoppingCriteria): class NumberStoppingCriteria(transformers.StoppingCriteria):
def __init__( def __init__(
self, self,
tokenizer: PreTrainedTokenizer, tokenizer: transformers.PreTrainedTokenizer, # type: ignore
prompt_length: int, prompt_length: int,
precision: int = 3, precision: int = 3,
): ):
@ -44,8 +58,8 @@ class NumberStoppingCriteria(StoppingCriteria):
def __call__( def __call__(
self, self,
input_ids: torch.LongTensor, input_ids: torch.LongTensor, # type: ignore
scores: torch.FloatTensor, scores: torch.FloatTensor, # type: ignore
) -> bool: ) -> bool:
decoded = self.tokenizer.decode( decoded = self.tokenizer.decode(
input_ids[0][self.prompt_length :], input_ids[0][self.prompt_length :],
@ -71,8 +85,8 @@ class NumberStoppingCriteria(StoppingCriteria):
return False return False
class OutputNumbersTokens(LogitsWarper): class OutputNumbersTokens(transformers.LogitsWarper):
def __init__(self, tokenizer: PreTrainedTokenizer, prompt: str): def __init__(self, tokenizer: transformers.PreTrainedTokenizer, prompt: str): # type: ignore
self.tokenizer = tokenizer self.tokenizer = tokenizer
self.tokenized_prompt = tokenizer(prompt, return_tensors="pt") self.tokenized_prompt = tokenizer(prompt, return_tensors="pt")
vocab_size = len(tokenizer) vocab_size = len(tokenizer)

146
swarms/utils/auto_download_check_packages.py
Unescape View File

@ -0,0 +1,146 @@
"""
Package installation utility that checks for package existence and installs if needed.
Supports both pip and conda package managers.
"""
import importlib.util
import subprocess
import sys
from typing import Literal, Optional, Union
from swarms.utils.loguru_logger import initialize_logger
import pkg_resources
logger = initialize_logger("autocheckpackages")
def check_and_install_package(
package_name: str,
package_manager: Literal["pip", "conda"] = "pip",
version: Optional[str] = None,
upgrade: bool = False,
) -> bool:
"""
Check if a package is installed and install it if not found.
Args:
package_name: Name of the package to check/install
package_manager: Package manager to use ('pip' or 'conda')
version: Specific version to install (optional)
upgrade: Whether to upgrade the package if it exists
Returns:
bool: True if package is available after check/install, False if installation failed
Raises:
ValueError: If invalid package manager is specified
"""
try:
# Check if package exists
if package_manager == "pip":
try:
pkg_resources.get_distribution(package_name)
if not upgrade:
logger.info(
f"Package {package_name} is already installed"
)
return True
except pkg_resources.DistributionNotFound:
pass
# Construct installation command
cmd = [sys.executable, "-m", "pip", "install"]
if upgrade:
cmd.append("--upgrade")
if version:
cmd.append(f"{package_name}=={version}")
else:
cmd.append(package_name)
elif package_manager == "conda":
# Check if conda is available
try:
subprocess.run(
["conda", "--version"],
check=True,
capture_output=True,
)
except (subprocess.CalledProcessError, FileNotFoundError):
logger.error(
"Conda is not available. Please install conda first."
)
return False
# Construct conda command
cmd = ["conda", "install", "-y"]
if version:
cmd.append(f"{package_name}={version}")
else:
cmd.append(package_name)
else:
raise ValueError(
f"Invalid package manager: {package_manager}"
)
# Run installation
logger.info(f"Installing {package_name}...")
subprocess.run(
cmd, check=True, capture_output=True, text=True
)
# Verify installation
try:
importlib.import_module(package_name)
logger.info(f"Successfully installed {package_name}")
return True
except ImportError:
logger.error(
f"Package {package_name} was installed but cannot be imported"
)
return False
except subprocess.CalledProcessError as e:
logger.error(f"Failed to install {package_name}: {e.stderr}")
return False
except Exception as e:
logger.error(
f"Unexpected error while installing {package_name}: {str(e)}"
)
return False
def auto_check_and_download_package(
packages: Union[str, list[str]],
package_manager: Literal["pip", "conda"] = "pip",
upgrade: bool = False,
) -> bool:
"""
Ensure multiple packages are installed.
Args:
packages: Single package name or list of package names
package_manager: Package manager to use ('pip' or 'conda')
upgrade: Whether to upgrade existing packages
Returns:
bool: True if all packages are available, False if any installation failed
"""
if isinstance(packages, str):
packages = [packages]
success = True
for package in packages:
if ":" in package:
name, version = package.split(":")
if not check_and_install_package(
name, package_manager, version, upgrade
):
success = False
else:
if not check_and_install_package(
package, package_manager, upgrade=upgrade
):
success = False
return success

263
swarms/utils/lazy_loader.py
Unescape View File

@ -0,0 +1,263 @@
"""
Lazy Package Loader
This module provides utilities for lazy loading Python packages to improve startup time
and reduce memory usage by only importing packages when they are actually used.
Features:
- Type-safe lazy loading of packages
- Support for nested module imports
- Auto-completion support in IDEs
- Thread-safe implementation
- Comprehensive test coverage
"""
from types import ModuleType
from typing import (
Optional,
Dict,
Any,
Callable,
Type,
TypeVar,
Union,
cast,
)
import importlib
import functools
import threading
from importlib.util import find_spec
from swarms.utils.auto_download_check_packages import (
auto_check_and_download_package,
)
T = TypeVar("T")
C = TypeVar("C")
class ImportError(Exception):
"""Raised when a lazy import fails."""
pass
class LazyLoader:
"""
A thread-safe lazy loader for Python packages that only imports them when accessed.
Attributes:
_module_name (str): The name of the module to be lazily loaded
_module (Optional[ModuleType]): The cached module instance once loaded
_lock (threading.Lock): Thread lock for safe concurrent access
Examples:
>>> np = LazyLoader('numpy')
>>> # numpy is not imported yet
>>> result = np.array([1, 2, 3])
>>> # numpy is imported only when first used
"""
def __init__(self, module_name: str) -> None:
"""
Initialize the lazy loader with a module name.
Args:
module_name: The fully qualified name of the module to lazily load
Raises:
ImportError: If the module cannot be found in sys.path
"""
self._module_name = module_name
self._module: Optional[ModuleType] = None
self._lock = threading.Lock()
auto_check_and_download_package(
module_name, package_manager="pip"
)
# Verify module exists without importing it
if find_spec(module_name) is None:
raise ImportError(
f"Module '{module_name}' not found in sys.path"
)
def _load_module(self) -> ModuleType:
"""
Thread-safe module loading.
Returns:
ModuleType: The loaded module
Raises:
ImportError: If module import fails
"""
if self._module is None:
with self._lock:
# Double-check pattern
if self._module is None:
try:
self._module = importlib.import_module(
self._module_name
)
except Exception as e:
raise ImportError(
f"Failed to import '{self._module_name}': {str(e)}"
)
return cast(ModuleType, self._module)
def __getattr__(self, name: str) -> Any:
"""
Intercepts attribute access to load the module if needed.
Args:
name: The attribute name being accessed
Returns:
Any: The requested attribute from the loaded module
Raises:
AttributeError: If the attribute doesn't exist in the module
"""
module = self._load_module()
try:
return getattr(module, name)
except AttributeError:
raise AttributeError(
f"Module '{self._module_name}' has no attribute '{name}'"
)
def __dir__(self) -> list[str]:
"""
Returns list of attributes for autocomplete support.
Returns:
List[str]: Available attributes in the module
"""
return dir(self._load_module())
def is_loaded(self) -> bool:
"""
Check if the module has been loaded.
Returns:
bool: True if module is loaded, False otherwise
"""
return self._module is not None
class LazyLoaderMetaclass(type):
"""Metaclass to handle lazy loading behavior"""
def __call__(cls, *args, **kwargs):
if hasattr(cls, "_lazy_loader"):
return super().__call__(*args, **kwargs)
return super().__call__(*args, **kwargs)
class LazyClassLoader:
"""
A descriptor that creates the actual class only when accessed,
with proper inheritance support.
"""
def __init__(
self, class_name: str, bases: tuple, namespace: Dict[str, Any]
):
self.class_name = class_name
self.bases = bases
self.namespace = namespace
self._real_class: Optional[Type] = None
self._lock = threading.Lock()
def _create_class(self) -> Type:
"""Creates the actual class if it hasn't been created yet."""
if self._real_class is None:
with self._lock:
if self._real_class is None:
# Update namespace to include metaclass
namespace = dict(self.namespace)
namespace["__metaclass__"] = LazyLoaderMetaclass
# Create the class with metaclass
new_class = LazyLoaderMetaclass(
self.class_name, self.bases, namespace
)
# Store reference to this loader
new_class._lazy_loader = self
self._real_class = new_class
return cast(Type, self._real_class)
def __call__(self, *args: Any, **kwargs: Any) -> Any:
"""Creates an instance of the lazy loaded class."""
real_class = self._create_class()
# Use the metaclass __call__ method
return real_class(*args, **kwargs)
def __instancecheck__(self, instance: Any) -> bool:
"""Support for isinstance() checks"""
real_class = self._create_class()
return isinstance(instance, real_class)
def __subclasscheck__(self, subclass: Type) -> bool:
"""Support for issubclass() checks"""
real_class = self._create_class()
return issubclass(subclass, real_class)
def lazy_import(*names: str) -> Dict[str, LazyLoader]:
"""
Create multiple lazy loaders at once.
Args:
*names: Module names to create lazy loaders for
Returns:
Dict[str, LazyLoader]: Dictionary mapping module names to their lazy loaders
Examples:
>>> modules = lazy_import('numpy', 'pandas', 'matplotlib.pyplot')
>>> np = modules['numpy']
>>> pd = modules['pandas']
>>> plt = modules['matplotlib.pyplot']
"""
return {name.split(".")[-1]: LazyLoader(name) for name in names}
def lazy_import_decorator(
target: Union[Callable[..., T], Type[C]]
) -> Union[Callable[..., T], Type[C], LazyClassLoader]:
"""
Enhanced decorator that supports both lazy imports and lazy class loading.
"""
if isinstance(target, type):
# Store the original class details
namespace = {
name: value
for name, value in target.__dict__.items()
if not name.startswith("__")
or name in ("__init__", "__new__")
}
# Create lazy loader
loader = LazyClassLoader(
target.__name__, target.__bases__, namespace
)
# Preserve class metadata
loader.__module__ = target.__module__
loader.__doc__ = target.__doc__
# Add reference to original class
loader._original_class = target
return loader
else:
# Handle function decoration
@functools.wraps(target)
def wrapper(*args: Any, **kwargs: Any) -> T:
return target(*args, **kwargs)
return wrapper

3
swarms/utils/litellm.py
Unescape View File

@ -8,6 +8,7 @@ except ImportError:
from litellm import completion from litellm import completion
litellm.set_verbose = True litellm.set_verbose = True
litellm.ssl_verify = False
class LiteLLM: class LiteLLM:
@ -23,6 +24,7 @@ class LiteLLM:
stream: bool = False, stream: bool = False,
temperature: float = 0.5, temperature: float = 0.5,
max_tokens: int = 4000, max_tokens: int = 4000,
ssl_verify: bool = False,
): ):
""" """
Initialize the LiteLLM with the given parameters. Initialize the LiteLLM with the given parameters.
@ -39,6 +41,7 @@ class LiteLLM:
self.stream = stream self.stream = stream
self.temperature = temperature self.temperature = temperature
self.max_tokens = max_tokens self.max_tokens = max_tokens
self.ssl_verify = ssl_verify
def _prepare_messages(self, task: str) -> list: def _prepare_messages(self, task: str) -> list:
""" """

73
swarms/utils/openai_tts.py
Unescape View File

@ -1,73 +0,0 @@
import os
from loguru import logger
import pygame
import requests
import tempfile
from openai import OpenAI
class OpenAITTS:
"""
A class to interact with OpenAI API and play the generated audio with improved streaming capabilities.
"""
def __init__(self, *args, **kwargs):
self.client = OpenAI(
api_key=os.getenv("OPENAI_API_KEY"), *args, **kwargs
)
pygame.init()
def run(
self, task: str, play_sound: bool = True, *args, **kwargs
):
"""
Run a task with the OpenAI API and optionally play the generated audio with improved streaming.
Args:
task (str): The task to be executed.
play_sound (bool): If True, play the generated audio.
Returns:
None
"""
try:
response = self.client.audio.speech.create(
model="tts-1",
voice="nova",
input=task,
*args,
**kwargs,
)
audio_url = response["url"]
logger.info("Task completed successfully.")
if play_sound:
with tempfile.NamedTemporaryFile(
delete=False, suffix=".mp3"
) as tmp_file:
with requests.get(audio_url, stream=True) as r:
r.raise_for_status()
for chunk in r.iter_content(chunk_size=8192):
tmp_file.write(chunk)
pygame.mixer.music.load(tmp_file.name)
pygame.mixer.music.play()
while pygame.mixer.music.get_busy():
pygame.time.Clock().tick(10)
except Exception as e:
logger.error(f"Error during task execution: {str(e)}")
# client = OpenAITTS(api_key=os.getenv("OPENAI_API_KEY"))
# client.run("Hello world! This is a streaming test.", play_sound=True)
def text_to_speech(
task: str, play_sound: bool = True, *args, **kwargs
):
out = OpenAITTS().run(
task, play_sound=play_sound, *args, **kwargs
)
return out
# print(text_to_speech(task="hello"))

42
tree_swarm_test.py
Unescape View File

@ -0,0 +1,42 @@
from swarms.structs.tree_swarm import ForestSwarm, Tree, TreeAgent
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?"
multi_agent_structure.run(task)

206
zpk.py
Unescape View File

@ -0,0 +1,206 @@
from swarms import Agent
from loguru import logger
import random
import re
# Configure loguru
logger.add("zkp_log.log", rotation="500 KB", retention="10 days", level="INFO")
class ProverAgent:
"""
Prover Agent for Zero Knowledge Proof.
Responsibilities:
- Generate commitments based on a secret.
- Respond to challenges from the Verifier.
Attributes:
agent (Agent): Swarms agent instance.
p (int): The prime modulus.
g (int): The generator.
x (int): The Prover's secret.
"""
def __init__(self, p: int, g: int, secret: int):
self.p = p
self.g = g
self.x = secret # Prover's secret
self.agent = Agent(
agent_name="ProverAgent",
model_name="gpt-4o-mini",
max_loop=1,
interactive=False,
streaming_on=True,
system_prompt=(
"You are the Prover in a Zero Knowledge Proof (ZKP) system. "
"Your responsibilities are to generate commitments based on a secret value and "
"respond to challenges from the Verifier without revealing the secret. "
"Follow mathematical rules of modular arithmetic when performing computations."
),
)
logger.info("Initialized ProverAgent with p={}, g={}, secret={}", p, g, secret)
def generate_commitment(self) -> tuple[int, int]:
"""
Generates a random commitment for the proof.
Returns:
tuple[int, int]: The random value (r) and the commitment (t).
"""
r = random.randint(1, self.p - 2)
task = (
f"Compute the commitment t = g^r % p for g={self.g}, r={r}, p={self.p}. "
"Return only the numerical value of t as an integer."
)
t = self.agent.run(task=task)
t_value = self._extract_integer(t, "commitment")
logger.info("Prover generated commitment: r={}, t={}", r, t_value)
return r, t_value
def _extract_integer(self, response: str, label: str) -> int:
"""
Extracts an integer from the LLM response.
Args:
response (str): The response from the agent.
label (str): A label for logging purposes.
Returns:
int: The extracted integer value.
"""
try:
# Use regex to find the first integer in the response
match = re.search(r"\b\d+\b", response)
if match:
value = int(match.group(0))
return value
else:
raise ValueError(f"No integer found in {label} response: {response}")
except Exception as e:
logger.error("Failed to extract integer from {label} response: {response}")
raise ValueError(f"Invalid {label} response: {response}") from e
def respond_to_challenge(self, r: int, c: int) -> int:
"""
Computes the response to a challenge.
Args:
r (int): The random value used in the commitment.
c (int): The challenge issued by the Verifier.
Returns:
int: The response (z).
"""
task = f"Compute the response z = (r + c * x) % (p-1) for r={r}, c={c}, x={self.x}, p={self.p}."
z = self.agent.run(task=task)
logger.info("Prover responded to challenge: z={}", z)
return int(z)
class VerifierAgent:
"""
Verifier Agent for Zero Knowledge Proof.
Responsibilities:
- Issue challenges to the Prover.
- Verify the Prover's response.
Attributes:
agent (Agent): Swarms agent instance.
p (int): The prime modulus.
g (int): The generator.
y (int): The public value from the Prover.
"""
def __init__(self, p: int, g: int, y: int):
self.p = p
self.g = g
self.y = y # Public value
self.agent = Agent(
agent_name="VerifierAgent",
model_name="gpt-4o-mini",
max_loop=1,
interactive=False,
streaming_on=True,
system_prompt=(
"You are the Verifier in a Zero Knowledge Proof (ZKP) system. "
"Your responsibilities are to issue random challenges and verify the Prover's response. "
"Use modular arithmetic to check if the proof satisfies g^z % p == (t * y^c) % p."
),
)
logger.info("Initialized VerifierAgent with p={}, g={}, y={}", p, g, y)
def issue_challenge(self) -> int:
"""
Issues a random challenge to the Prover.
Returns:
int: The challenge value (c).
"""
c = random.randint(1, 10)
logger.info("Verifier issued challenge: c={}", c)
return c
def verify_proof(self, t: int, z: int, c: int) -> bool:
"""
Verifies the Prover's response.
Args:
t (int): The commitment from the Prover.
z (int): The response from the Prover.
c (int): The challenge issued to the Prover.
Returns:
bool: True if the proof is valid, False otherwise.
"""
task = f"Verify if g^z % p == (t * y^c) % p for g={self.g}, z={z}, p={self.p}, t={t}, y={self.y}, c={c}."
verification_result = self.agent.run(task=task)
is_valid = verification_result.strip().lower() == "true"
logger.info("Verifier checked proof: t={}, z={}, c={}, valid={}", t, z, c, is_valid)
return is_valid
class CoordinatorAgent:
"""
Coordinator for orchestrating the Zero Knowledge Proof protocol.
Responsibilities:
- Initialize parameters.
- Facilitate interaction between Prover and Verifier agents.
"""
def __init__(self, p: int, g: int, secret: int):
self.p = p
self.g = g
self.prover = ProverAgent(p, g, secret)
y = pow(g, secret, p) # Public value
self.verifier = VerifierAgent(p, g, y)
logger.info("Coordinator initialized with p={}, g={}, secret={}", p, g, secret)
def orchestrate(self) -> bool:
"""
Orchestrates the Zero Knowledge Proof protocol.
Returns:
bool: True if the proof is valid, False otherwise.
"""
logger.info("Starting ZKP protocol orchestration.")
r, t = self.prover.generate_commitment()
c = self.verifier.issue_challenge()
z = self.prover.respond_to_challenge(r, c)
is_valid = self.verifier.verify_proof(t, z, c)
logger.info("ZKP protocol completed. Valid proof: {}", is_valid)
return is_valid
if __name__ == "__main__":
# Example parameters
p = 23 # Prime number
g = 5 # Generator
secret = 7 # Prover's secret
# Initialize the Coordinator and run the protocol
coordinator = CoordinatorAgent(p, g, secret)
result = coordinator.orchestrate()
print(f"Zero Knowledge Proof Verification Result: {'Valid' if result else 'Invalid'}")
Write Preview
Loading…
Cancel
Save