evelynmitchell
3 weeks ago
committed by
GitHub
commit
7f15fd23b1
@ -1,107 +1,320 @@
|
||||
import requests
|
||||
from loguru import logger
|
||||
import time
|
||||
|
||||
# Configure loguru
|
||||
logger.add(
|
||||
"api_tests_{time}.log",
|
||||
rotation="100 MB",
|
||||
level="DEBUG",
|
||||
format="{time} {level} {message}",
|
||||
)
|
||||
from typing import Dict, Optional, Tuple
|
||||
from uuid import UUID
|
||||
import sys
|
||||
|
||||
BASE_URL = "http://localhost:8000/v1"
|
||||
|
||||
|
||||
def test_create_agent():
|
||||
def check_api_server() -> bool:
|
||||
"""Check if the API server is running and accessible."""
|
||||
try:
|
||||
response = requests.get(f"{BASE_URL}/docs")
|
||||
return response.status_code == 200
|
||||
except requests.exceptions.ConnectionError:
|
||||
logger.error("API server is not running at {BASE_URL}")
|
||||
logger.error("Please start the API server first with:")
|
||||
logger.error(" python main.py")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking API server: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
class TestSession:
|
||||
"""Manages test session state and authentication."""
|
||||
|
||||
def __init__(self):
|
||||
self.user_id: Optional[UUID] = None
|
||||
self.api_key: Optional[str] = None
|
||||
self.test_agents: list[UUID] = []
|
||||
|
||||
@property
|
||||
def headers(self) -> Dict[str, str]:
|
||||
"""Get headers with authentication."""
|
||||
return {"api-key": self.api_key} if self.api_key else {}
|
||||
|
||||
|
||||
def create_test_user(session: TestSession) -> Tuple[bool, str]:
|
||||
"""Create a test user and store credentials in session."""
|
||||
logger.info("Creating test user")
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{BASE_URL}/users",
|
||||
json={"username": f"test_user_{int(time.time())}"},
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
session.user_id = data["user_id"]
|
||||
session.api_key = data["api_key"]
|
||||
logger.success(f"Created user with ID: {session.user_id}")
|
||||
return True, "Success"
|
||||
else:
|
||||
logger.error(f"Failed to create user: {response.text}")
|
||||
return False, response.text
|
||||
except Exception as e:
|
||||
logger.exception("Exception during user creation")
|
||||
return False, str(e)
|
||||
|
||||
|
||||
def create_additional_api_key(
|
||||
session: TestSession,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Test creating an additional API key."""
|
||||
logger.info("Creating additional API key")
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{BASE_URL}/users/{session.user_id}/api-keys",
|
||||
headers=session.headers,
|
||||
json={"name": "Test Key"},
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
logger.success("Created additional API key")
|
||||
return True, response.json()["key"]
|
||||
else:
|
||||
logger.error(f"Failed to create API key: {response.text}")
|
||||
return False, response.text
|
||||
except Exception as e:
|
||||
logger.exception("Exception during API key creation")
|
||||
return False, str(e)
|
||||
|
||||
|
||||
def test_create_agent(
|
||||
session: TestSession,
|
||||
) -> Tuple[bool, Optional[UUID]]:
|
||||
"""Test creating a new agent."""
|
||||
logger.info("Testing agent creation")
|
||||
|
||||
payload = {
|
||||
"agent_name": "Test Agent",
|
||||
"agent_name": f"Test Agent {int(time.time())}",
|
||||
"system_prompt": "You are a helpful assistant",
|
||||
"model_name": "gpt-4",
|
||||
"description": "Test agent",
|
||||
"tags": ["test"],
|
||||
"tags": ["test", "automated"],
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/agent", json=payload)
|
||||
logger.debug(f"Create response: {response.json()}")
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{BASE_URL}/agent", headers=session.headers, json=payload
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
agent_id = response.json()["agent_id"]
|
||||
session.test_agents.append(agent_id)
|
||||
logger.success(f"Created agent with ID: {agent_id}")
|
||||
return True, agent_id
|
||||
else:
|
||||
logger.error(f"Failed to create agent: {response.text}")
|
||||
return False, None
|
||||
except Exception:
|
||||
logger.exception("Exception during agent creation")
|
||||
return False, None
|
||||
|
||||
|
||||
def test_list_user_agents(session: TestSession) -> bool:
|
||||
"""Test listing user's agents."""
|
||||
logger.info("Testing user agent listing")
|
||||
|
||||
try:
|
||||
response = requests.get(
|
||||
f"{BASE_URL}/users/me/agents", headers=session.headers
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
agents = response.json()
|
||||
logger.success(f"Found {len(agents)} user agents")
|
||||
return True
|
||||
else:
|
||||
logger.error(
|
||||
f"Failed to list user agents: {response.text}"
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.exception("Exception during agent listing")
|
||||
return False
|
||||
|
||||
if response.status_code == 200:
|
||||
logger.success("Successfully created agent")
|
||||
return response.json()["agent_id"]
|
||||
else:
|
||||
logger.error(f"Failed to create agent: {response.text}")
|
||||
return None
|
||||
|
||||
def test_agent_operations(
|
||||
session: TestSession, agent_id: UUID
|
||||
) -> bool:
|
||||
"""Test various operations on an agent."""
|
||||
logger.info(f"Testing operations for agent {agent_id}")
|
||||
|
||||
def test_list_agents():
|
||||
"""Test listing all agents."""
|
||||
logger.info("Testing agent listing")
|
||||
# Test update
|
||||
try:
|
||||
update_response = requests.patch(
|
||||
f"{BASE_URL}/agent/{agent_id}",
|
||||
headers=session.headers,
|
||||
json={
|
||||
"description": "Updated description",
|
||||
"tags": ["test", "updated"],
|
||||
},
|
||||
)
|
||||
if update_response.status_code != 200:
|
||||
logger.error(
|
||||
f"Failed to update agent: {update_response.text}"
|
||||
)
|
||||
return False
|
||||
|
||||
response = requests.get(f"{BASE_URL}/agents")
|
||||
logger.debug(f"List response: {response.json()}")
|
||||
# Test metrics
|
||||
metrics_response = requests.get(
|
||||
f"{BASE_URL}/agent/{agent_id}/metrics",
|
||||
headers=session.headers,
|
||||
)
|
||||
if metrics_response.status_code != 200:
|
||||
logger.error(
|
||||
f"Failed to get agent metrics: {metrics_response.text}"
|
||||
)
|
||||
return False
|
||||
|
||||
if response.status_code == 200:
|
||||
logger.success(f"Found {len(response.json())} agents")
|
||||
else:
|
||||
logger.error(f"Failed to list agents: {response.text}")
|
||||
logger.success("Successfully performed agent operations")
|
||||
return True
|
||||
except Exception:
|
||||
logger.exception("Exception during agent operations")
|
||||
return False
|
||||
|
||||
|
||||
def test_completion(agent_id):
|
||||
def test_completion(session: TestSession, agent_id: UUID) -> bool:
|
||||
"""Test running a completion."""
|
||||
logger.info("Testing completion")
|
||||
|
||||
payload = {
|
||||
"prompt": "What is the weather like today?",
|
||||
"agent_id": agent_id,
|
||||
"max_tokens": 100,
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
f"{BASE_URL}/agent/completions", json=payload
|
||||
)
|
||||
logger.debug(f"Completion response: {response.json()}")
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{BASE_URL}/agent/completions",
|
||||
headers=session.headers,
|
||||
json=payload,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
logger.success("Successfully got completion")
|
||||
else:
|
||||
logger.error(f"Failed to get completion: {response.text}")
|
||||
if response.status_code == 200:
|
||||
completion_data = response.json()
|
||||
logger.success(
|
||||
f"Got completion, used {completion_data['token_usage']['total_tokens']} tokens"
|
||||
)
|
||||
return True
|
||||
else:
|
||||
logger.error(f"Failed to get completion: {response.text}")
|
||||
return False
|
||||
except Exception:
|
||||
logger.exception("Exception during completion")
|
||||
return False
|
||||
|
||||
|
||||
def test_delete_agent(agent_id):
|
||||
"""Test deleting an agent."""
|
||||
logger.info("Testing agent deletion")
|
||||
def cleanup_test_resources(session: TestSession):
|
||||
"""Clean up all test resources."""
|
||||
logger.info("Cleaning up test resources")
|
||||
|
||||
response = requests.delete(f"{BASE_URL}/agent/{agent_id}")
|
||||
logger.debug(f"Delete response: {response.json()}")
|
||||
# Delete test agents
|
||||
for agent_id in session.test_agents:
|
||||
try:
|
||||
response = requests.delete(
|
||||
f"{BASE_URL}/agent/{agent_id}",
|
||||
headers=session.headers,
|
||||
)
|
||||
if response.status_code == 200:
|
||||
logger.debug(f"Deleted agent {agent_id}")
|
||||
else:
|
||||
logger.warning(
|
||||
f"Failed to delete agent {agent_id}: {response.text}"
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(f"Exception deleting agent {agent_id}")
|
||||
|
||||
if response.status_code == 200:
|
||||
logger.success("Successfully deleted agent")
|
||||
else:
|
||||
logger.error(f"Failed to delete agent: {response.text}")
|
||||
# Revoke API keys
|
||||
if session.user_id:
|
||||
try:
|
||||
response = requests.get(
|
||||
f"{BASE_URL}/users/{session.user_id}/api-keys",
|
||||
headers=session.headers,
|
||||
)
|
||||
if response.status_code == 200:
|
||||
for key in response.json():
|
||||
try:
|
||||
revoke_response = requests.delete(
|
||||
f"{BASE_URL}/users/{session.user_id}/api-keys/{key['key']}",
|
||||
headers=session.headers,
|
||||
)
|
||||
if revoke_response.status_code == 200:
|
||||
logger.debug(
|
||||
f"Revoked API key {key['name']}"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Failed to revoke API key {key['name']}"
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
f"Exception revoking API key {key['name']}"
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Exception getting API keys for cleanup")
|
||||
|
||||
|
||||
def run_tests():
|
||||
"""Run all tests in sequence."""
|
||||
def run_test_workflow():
|
||||
"""Run complete test workflow."""
|
||||
logger.info("Starting API tests")
|
||||
|
||||
# Create agent and get ID
|
||||
agent_id = test_create_agent()
|
||||
if not agent_id:
|
||||
logger.error("Cannot continue tests without agent ID")
|
||||
return
|
||||
# Check if API server is running first
|
||||
if not check_api_server():
|
||||
return False
|
||||
|
||||
session = TestSession()
|
||||
|
||||
try:
|
||||
# Create user
|
||||
user_success, message = create_test_user(session)
|
||||
if not user_success:
|
||||
logger.error(f"User creation failed: {message}")
|
||||
return False
|
||||
|
||||
# Create additional API key
|
||||
key_success, key = create_additional_api_key(session)
|
||||
if not key_success:
|
||||
logger.error(f"API key creation failed: {key}")
|
||||
return False
|
||||
|
||||
# Create agent
|
||||
agent_success, agent_id = test_create_agent(session)
|
||||
if not agent_success or not agent_id:
|
||||
logger.error("Agent creation failed")
|
||||
return False
|
||||
|
||||
# Test user agent listing
|
||||
if not test_list_user_agents(session):
|
||||
logger.error("Agent listing failed")
|
||||
return False
|
||||
|
||||
# Test agent operations
|
||||
if not test_agent_operations(session, agent_id):
|
||||
logger.error("Agent operations failed")
|
||||
return False
|
||||
|
||||
# Wait a bit for agent to be ready
|
||||
time.sleep(1)
|
||||
# Test completion
|
||||
if not test_completion(session, agent_id):
|
||||
logger.error("Completion test failed")
|
||||
return False
|
||||
|
||||
# Run other tests
|
||||
test_list_agents()
|
||||
test_completion(agent_id)
|
||||
test_delete_agent(agent_id)
|
||||
logger.success("All tests completed successfully")
|
||||
return True
|
||||
|
||||
logger.info("Tests completed")
|
||||
except Exception:
|
||||
logger.exception("Exception during test workflow")
|
||||
return False
|
||||
finally:
|
||||
cleanup_test_resources(session)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_tests()
|
||||
success = run_test_workflow()
|
||||
sys.exit(0 if success else 1)
|
||||
|
||||
@ -0,0 +1,6 @@
|
||||
fastapi
|
||||
uvicorn
|
||||
pydantic
|
||||
loguru
|
||||
python-dotenv
|
||||
swarms # Specify the version or source if it's not on PyPI
|
||||
@ -0,0 +1,37 @@
|
||||
service:
|
||||
readiness_probe:
|
||||
path: /docs
|
||||
initial_delay_seconds: 300
|
||||
timeout_seconds: 30
|
||||
|
||||
replica_policy:
|
||||
min_replicas: 1
|
||||
max_replicas: 50
|
||||
target_qps_per_replica: 5
|
||||
upscale_delay_seconds: 180
|
||||
downscale_delay_seconds: 600
|
||||
|
||||
resources:
|
||||
ports: 8000 # FastAPI default port
|
||||
cpus: 16
|
||||
memory: 64
|
||||
disk_size: 100
|
||||
use_spot: true
|
||||
|
||||
workdir: /app
|
||||
|
||||
setup: |
|
||||
git clone https://github.com/kyegomez/swarms.git
|
||||
cd swarms/api
|
||||
pip install -r requirements.txt
|
||||
pip install swarms
|
||||
|
||||
run: |
|
||||
cd swarms/api
|
||||
uvicorn main:app --host 0.0.0.0 --port 8000 --workers 4
|
||||
|
||||
# env:
|
||||
# PYTHONPATH: /app/swarms
|
||||
# LOG_LEVEL: "INFO"
|
||||
# # MAX_WORKERS: "4"
|
||||
|
||||
@ -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()
|
||||
@ -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()
|
||||
@ -1,50 +1,31 @@
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from swarm_models import OpenAIChat
|
||||
|
||||
from swarms import Agent
|
||||
from swarms.prompts.finance_agent_sys_prompt import (
|
||||
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
|
||||
agent = Agent(
|
||||
agent_name="Financial-Analysis-Agent",
|
||||
system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
dashboard=False,
|
||||
verbose=True,
|
||||
agent_description="Personal finance advisor agent",
|
||||
system_prompt=FINANCIAL_AGENT_SYS_PROMPT
|
||||
+ "Output the <DONE> token when you're done creating a portfolio of etfs, index, funds, and more for AI",
|
||||
model_name="gpt-4o", # Use any model from litellm
|
||||
max_loops="auto",
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="finance_agent.json",
|
||||
user_name="swarms_corp",
|
||||
retry_attempts=1,
|
||||
user_name="Kye",
|
||||
retry_attempts=3,
|
||||
streaming_on=True,
|
||||
context_length=200000,
|
||||
return_step_meta=True,
|
||||
output_type="json", # "json", "dict", "csv" OR "string" soon "yaml" and
|
||||
context_length=16000,
|
||||
return_step_meta=False,
|
||||
output_type="str", # "json", "dict", "csv" OR "string" "yaml" and
|
||||
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,
|
||||
return_history=True,
|
||||
max_tokens=16000, # max output tokens
|
||||
interactive=True,
|
||||
stopping_token="<DONE>",
|
||||
execute_tool=True,
|
||||
)
|
||||
|
||||
|
||||
agent.run(
|
||||
"How can I establish a ROTH IRA to buy stocks and get a tax break? What are the criteria. Create a report on this question.",
|
||||
"Create a table of super high growth opportunities for AI. I have $40k to invest in ETFs, index funds, and more. Please create a table in markdown.",
|
||||
all_cores=True,
|
||||
)
|
||||
|
||||
@ -0,0 +1,147 @@
|
||||
from swarms.structs.tree_swarm import ForestSwarm, Tree, TreeAgent
|
||||
|
||||
# Fund Analysis Tree
|
||||
fund_agents = [
|
||||
TreeAgent(
|
||||
system_prompt="""Mutual Fund Analysis Agent:
|
||||
- Analyze mutual fund performance metrics and ratios
|
||||
- Evaluate fund manager track records and strategy consistency
|
||||
- Compare expense ratios and fee structures
|
||||
- Assess fund holdings and sector allocations
|
||||
- Monitor fund inflows/outflows and size implications
|
||||
- Analyze risk-adjusted returns (Sharpe, Sortino ratios)
|
||||
- Consider tax efficiency and distribution history
|
||||
- Track style drift and benchmark adherence
|
||||
Knowledge base: Mutual fund operations, portfolio management, fee structures
|
||||
Output format: Fund analysis report with recommendations""",
|
||||
agent_name="Mutual Fund Analyst",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Index Fund Specialist Agent:
|
||||
- Evaluate index tracking accuracy and tracking error
|
||||
- Compare different index methodologies
|
||||
- Analyze index fund costs and tax efficiency
|
||||
- Monitor index rebalancing impacts
|
||||
- Assess market capitalization weightings
|
||||
- Compare similar indices and their differences
|
||||
- Evaluate smart beta and factor strategies
|
||||
Knowledge base: Index construction, passive investing, market efficiency
|
||||
Output format: Index fund comparison and selection recommendations""",
|
||||
agent_name="Index Fund Specialist",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""ETF Strategy Agent:
|
||||
- Analyze ETF liquidity and trading volumes
|
||||
- Evaluate creation/redemption mechanisms
|
||||
- Compare ETF spreads and premium/discount patterns
|
||||
- Assess underlying asset liquidity
|
||||
- Monitor authorized participant activity
|
||||
- Analyze securities lending revenue
|
||||
- Compare similar ETFs and their structures
|
||||
Knowledge base: ETF mechanics, trading strategies, market making
|
||||
Output format: ETF analysis with trading recommendations""",
|
||||
agent_name="ETF Strategist",
|
||||
),
|
||||
]
|
||||
|
||||
# Sector Specialist Tree
|
||||
sector_agents = [
|
||||
TreeAgent(
|
||||
system_prompt="""Energy Sector Analysis Agent:
|
||||
- Track global energy market trends
|
||||
- Analyze traditional and renewable energy companies
|
||||
- Monitor regulatory changes and policy impacts
|
||||
- Evaluate commodity price influences
|
||||
- Assess geopolitical risk factors
|
||||
- Track technological disruption in energy
|
||||
- Analyze energy infrastructure investments
|
||||
Knowledge base: Energy markets, commodities, regulatory environment
|
||||
Output format: Energy sector analysis with investment opportunities""",
|
||||
agent_name="Energy Sector Analyst",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""AI and Technology Specialist Agent:
|
||||
- Research AI company fundamentals and growth metrics
|
||||
- Evaluate AI technology adoption trends
|
||||
- Analyze AI chip manufacturers and supply chains
|
||||
- Monitor AI software and service providers
|
||||
- Track AI patent filings and R&D investments
|
||||
- Assess competitive positioning in AI market
|
||||
- Consider regulatory risks and ethical factors
|
||||
Knowledge base: AI technology, semiconductor industry, tech sector dynamics
|
||||
Output format: AI sector analysis with investment recommendations""",
|
||||
agent_name="AI Technology Analyst",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Market Infrastructure Agent:
|
||||
- Monitor trading platform stability
|
||||
- Analyze market maker activity
|
||||
- Track exchange system updates
|
||||
- Evaluate clearing house operations
|
||||
- Monitor settlement processes
|
||||
- Assess cybersecurity measures
|
||||
- Track regulatory compliance updates
|
||||
Knowledge base: Market structure, trading systems, regulatory requirements
|
||||
Output format: Market infrastructure assessment and risk analysis""",
|
||||
agent_name="Infrastructure Monitor",
|
||||
),
|
||||
]
|
||||
|
||||
# Trading Strategy Tree
|
||||
strategy_agents = [
|
||||
TreeAgent(
|
||||
system_prompt="""Portfolio Strategy Agent:
|
||||
- Develop asset allocation strategies
|
||||
- Implement portfolio rebalancing rules
|
||||
- Monitor portfolio risk metrics
|
||||
- Optimize position sizing
|
||||
- Calculate portfolio correlation matrices
|
||||
- Implement tax-loss harvesting strategies
|
||||
- Track portfolio performance attribution
|
||||
Knowledge base: Portfolio theory, risk management, asset allocation
|
||||
Output format: Portfolio strategy recommendations with implementation plan""",
|
||||
agent_name="Portfolio Strategist",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Technical Analysis Agent:
|
||||
- Analyze price patterns and trends
|
||||
- Calculate technical indicators
|
||||
- Identify support/resistance levels
|
||||
- Monitor volume and momentum indicators
|
||||
- Track market breadth metrics
|
||||
- Analyze intermarket relationships
|
||||
- Generate trading signals
|
||||
Knowledge base: Technical analysis, chart patterns, market indicators
|
||||
Output format: Technical analysis report with trade signals""",
|
||||
agent_name="Technical Analyst",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Risk Management Agent:
|
||||
- Calculate position-level risk metrics
|
||||
- Monitor portfolio VaR and stress tests
|
||||
- Track correlation changes
|
||||
- Implement stop-loss strategies
|
||||
- Monitor margin requirements
|
||||
- Assess liquidity risk factors
|
||||
- Generate risk alerts and warnings
|
||||
Knowledge base: Risk metrics, position sizing, risk modeling
|
||||
Output format: Risk assessment report with mitigation recommendations""",
|
||||
agent_name="Risk Manager",
|
||||
),
|
||||
]
|
||||
|
||||
# Create trees
|
||||
fund_tree = Tree(tree_name="Fund Analysis", agents=fund_agents)
|
||||
sector_tree = Tree(tree_name="Sector Analysis", agents=sector_agents)
|
||||
strategy_tree = Tree(
|
||||
tree_name="Trading Strategy", agents=strategy_agents
|
||||
)
|
||||
|
||||
# Create the ForestSwarm
|
||||
trading_forest = ForestSwarm(
|
||||
trees=[fund_tree, sector_tree, strategy_tree]
|
||||
)
|
||||
|
||||
# Example usage
|
||||
task = "Analyze current opportunities in AI sector ETFs considering market conditions and provide a risk-adjusted portfolio allocation strategy. Add in the names of the best AI etfs that are reliable and align with this strategy and also include where to purchase the etfs"
|
||||
result = trading_forest.run(task)
|
||||
@ -0,0 +1,150 @@
|
||||
from swarms.structs.tree_swarm import ForestSwarm, Tree, TreeAgent
|
||||
|
||||
# Diagnostic Specialists Tree
|
||||
diagnostic_agents = [
|
||||
TreeAgent(
|
||||
system_prompt="""Primary Care Diagnostic Agent:
|
||||
- Conduct initial patient assessment and triage
|
||||
- Analyze patient symptoms, vital signs, and medical history
|
||||
- Identify red flags and emergency conditions
|
||||
- Coordinate with specialist agents for complex cases
|
||||
- Provide preliminary diagnosis recommendations
|
||||
- Consider common conditions and their presentations
|
||||
- Factor in patient demographics and risk factors
|
||||
Medical knowledge base: General medicine, common conditions, preventive care
|
||||
Output format: Structured assessment with recommended next steps""",
|
||||
agent_name="Primary Diagnostician",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Laboratory Analysis Agent:
|
||||
- Interpret complex laboratory results
|
||||
- Recommend appropriate test panels based on symptoms
|
||||
- Analyze blood work, urinalysis, and other diagnostic tests
|
||||
- Identify abnormal results and their clinical significance
|
||||
- Suggest follow-up tests when needed
|
||||
- Consider test accuracy and false positive/negative rates
|
||||
- Integrate lab results with clinical presentation
|
||||
Medical knowledge base: Clinical pathology, laboratory medicine, test interpretation
|
||||
Output format: Detailed lab analysis with clinical correlations""",
|
||||
agent_name="Lab Analyst",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Medical Imaging Specialist Agent:
|
||||
- Analyze radiological images (X-rays, CT, MRI, ultrasound)
|
||||
- Identify anatomical abnormalities and pathological changes
|
||||
- Recommend appropriate imaging studies
|
||||
- Correlate imaging findings with clinical symptoms
|
||||
- Provide differential diagnoses based on imaging
|
||||
- Consider radiation exposure and cost-effectiveness
|
||||
- Suggest follow-up imaging when needed
|
||||
Medical knowledge base: Radiology, anatomy, pathological imaging patterns
|
||||
Output format: Structured imaging report with findings and recommendations""",
|
||||
agent_name="Imaging Specialist",
|
||||
),
|
||||
]
|
||||
|
||||
# Treatment Specialists Tree
|
||||
treatment_agents = [
|
||||
TreeAgent(
|
||||
system_prompt="""Treatment Planning Agent:
|
||||
- Develop comprehensive treatment plans based on diagnosis
|
||||
- Consider evidence-based treatment guidelines
|
||||
- Account for patient factors (age, comorbidities, preferences)
|
||||
- Evaluate treatment risks and benefits
|
||||
- Consider cost-effectiveness and accessibility
|
||||
- Plan for treatment monitoring and adjustment
|
||||
- Coordinate multi-modal treatment approaches
|
||||
Medical knowledge base: Clinical guidelines, treatment protocols, medical management
|
||||
Output format: Detailed treatment plan with rationale and monitoring strategy""",
|
||||
agent_name="Treatment Planner",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Medication Management Agent:
|
||||
- Recommend appropriate medications and dosing
|
||||
- Check for drug interactions and contraindications
|
||||
- Consider patient-specific factors affecting medication choice
|
||||
- Provide medication administration guidelines
|
||||
- Monitor for adverse effects and therapeutic response
|
||||
- Suggest alternatives for contraindicated medications
|
||||
- Plan medication tapering or adjustments
|
||||
Medical knowledge base: Pharmacology, drug interactions, clinical pharmacotherapy
|
||||
Output format: Medication plan with monitoring parameters""",
|
||||
agent_name="Medication Manager",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Specialist Intervention Agent:
|
||||
- Recommend specialized procedures and interventions
|
||||
- Evaluate need for surgical vs. non-surgical approaches
|
||||
- Consider procedural risks and benefits
|
||||
- Provide pre- and post-procedure care guidelines
|
||||
- Coordinate with other specialists
|
||||
- Plan follow-up care and monitoring
|
||||
- Handle complex cases requiring multiple interventions
|
||||
Medical knowledge base: Surgical procedures, specialized interventions, perioperative care
|
||||
Output format: Intervention plan with risk assessment and care protocol""",
|
||||
agent_name="Intervention Specialist",
|
||||
),
|
||||
]
|
||||
|
||||
# Follow-up and Monitoring Tree
|
||||
followup_agents = [
|
||||
TreeAgent(
|
||||
system_prompt="""Recovery Monitoring Agent:
|
||||
- Track patient progress and treatment response
|
||||
- Identify complications or adverse effects early
|
||||
- Adjust treatment plans based on response
|
||||
- Coordinate follow-up appointments and tests
|
||||
- Monitor vital signs and symptoms
|
||||
- Evaluate treatment adherence and barriers
|
||||
- Recommend lifestyle modifications
|
||||
Medical knowledge base: Recovery patterns, complications, monitoring protocols
|
||||
Output format: Progress report with recommendations""",
|
||||
agent_name="Recovery Monitor",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Preventive Care Agent:
|
||||
- Develop preventive care strategies
|
||||
- Recommend appropriate screening tests
|
||||
- Provide lifestyle and dietary guidance
|
||||
- Monitor risk factors for disease progression
|
||||
- Coordinate vaccination schedules
|
||||
- Suggest health maintenance activities
|
||||
- Plan long-term health monitoring
|
||||
Medical knowledge base: Preventive medicine, health maintenance, risk reduction
|
||||
Output format: Preventive care plan with timeline""",
|
||||
agent_name="Prevention Specialist",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""Patient Education Agent:
|
||||
- Provide comprehensive patient education
|
||||
- Explain conditions and treatments in accessible language
|
||||
- Develop self-management strategies
|
||||
- Create educational materials and resources
|
||||
- Address common questions and concerns
|
||||
- Provide lifestyle modification guidance
|
||||
- Support treatment adherence
|
||||
Medical knowledge base: Patient education, health literacy, behavior change
|
||||
Output format: Educational plan with resources and materials""",
|
||||
agent_name="Patient Educator",
|
||||
),
|
||||
]
|
||||
|
||||
# Create trees
|
||||
diagnostic_tree = Tree(
|
||||
tree_name="Diagnostic Specialists", agents=diagnostic_agents
|
||||
)
|
||||
treatment_tree = Tree(
|
||||
tree_name="Treatment Specialists", agents=treatment_agents
|
||||
)
|
||||
followup_tree = Tree(
|
||||
tree_name="Follow-up and Monitoring", agents=followup_agents
|
||||
)
|
||||
|
||||
# Create the ForestSwarm
|
||||
medical_forest = ForestSwarm(
|
||||
trees=[diagnostic_tree, treatment_tree, followup_tree]
|
||||
)
|
||||
|
||||
# Example usage
|
||||
task = "Patient presents with persistent headache for 2 weeks, accompanied by visual disturbances and neck stiffness. Need comprehensive evaluation and treatment plan."
|
||||
result = medical_forest.run(task)
|
||||
@ -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)
|
||||
@ -0,0 +1,169 @@
|
||||
from swarms import Agent
|
||||
|
||||
|
||||
# Claims Processing Agent system prompt
|
||||
CLAIMS_PROCESSING_AGENT_SYS_PROMPT = """
|
||||
Here's an extended and detailed system prompt for the **Claims Processing Agent**, incorporating reasoning steps, output format, and examples for structured responses:
|
||||
You are a Claims Processing Agent specializing in automating and accelerating claims processing workflows. Your primary goal is to ensure Accuracy, reduce processing time, and flag potential fraud while providing clear and actionable insights. You must follow the detailed steps below to process claims efficiently and provide consistent, structured output.
|
||||
|
||||
### Primary Objectives:
|
||||
1. **Extract Information**:
|
||||
- Identify and extract key details from claim documents such as:
|
||||
- Claimant name, date of incident, and location.
|
||||
- Relevant policy numbers and coverage details.
|
||||
- Information from supporting documents like police reports, medical bills, or repair estimates.
|
||||
- For images (e.g., accident photos), extract descriptive metadata and identify key features (e.g., vehicle damage, environmental factors).
|
||||
|
||||
2. **Cross-Reference**:
|
||||
- Compare details across documents and media:
|
||||
- Validate consistency between police reports, medical bills, and other supporting documents.
|
||||
- Cross-check dates, times, and locations for coherence.
|
||||
- Analyze image evidence and correlate it with textual claims for verification.
|
||||
|
||||
3. **Fraud Detection**:
|
||||
- Apply analytical reasoning to identify inconsistencies or suspicious patterns, such as:
|
||||
- Discrepancies in timelines, damages, or descriptions.
|
||||
- Repetitive or unusually frequent claims involving the same parties or locations.
|
||||
- Signs of manipulated or altered evidence.
|
||||
|
||||
4. **Provide a Risk Assessment**:
|
||||
- Assign a preliminary risk level to the claim based on your analysis (e.g., Low, Medium, High).
|
||||
- Justify the risk level with a clear explanation.
|
||||
|
||||
5. **Flag and Recommend**:
|
||||
- Highlight any flagged concerns for human review and provide actionable recommendations.
|
||||
- Indicate documents, areas, or sections requiring further investigation.
|
||||
|
||||
---
|
||||
|
||||
### Reasoning Steps:
|
||||
Follow these steps to ensure comprehensive and accurate claim processing:
|
||||
1. **Document Analysis**:
|
||||
- Analyze each document individually to extract critical details.
|
||||
- Identify any missing or incomplete information.
|
||||
2. **Data Cross-Referencing**:
|
||||
- Check for consistency between documents.
|
||||
- Use contextual reasoning to spot potential discrepancies.
|
||||
3. **Fraud Pattern Analysis**:
|
||||
- Apply pattern recognition to flag anomalies or unusual claims.
|
||||
4. **Risk Assessment**:
|
||||
- Summarize your findings and categorize the risk.
|
||||
5. **Final Recommendations**:
|
||||
- Provide clear next steps for resolution or escalation.
|
||||
|
||||
---
|
||||
|
||||
### Output Format:
|
||||
Your output must be structured as follows:
|
||||
|
||||
#### 1. Extracted Information:
|
||||
```
|
||||
Claimant Name: [Name]
|
||||
Date of Incident: [Date]
|
||||
Location: [Location]
|
||||
Policy Number: [Policy Number]
|
||||
Summary of Incident: [Brief Summary]
|
||||
Supporting Documents:
|
||||
- Police Report: [Key Details]
|
||||
- Medical Bills: [Key Details]
|
||||
- Repair Estimate: [Key Details]
|
||||
- Photos: [Key Observations]
|
||||
```
|
||||
|
||||
#### 2. Consistency Analysis:
|
||||
```
|
||||
[Provide a detailed comparison of documents, highlighting any inconsistencies or gaps in data.]
|
||||
```
|
||||
|
||||
#### 3. Risk Assessment:
|
||||
```
|
||||
Risk Level: [Low / Medium / High]
|
||||
Reasoning: [Provide justification for the assigned risk level, supported by evidence from the analysis.]
|
||||
```
|
||||
|
||||
#### 4. Flagged Concerns and Recommendations:
|
||||
```
|
||||
Flagged Concerns:
|
||||
- [Detail specific issues or inconsistencies, e.g., timeline mismatch, suspicious patterns, etc.]
|
||||
|
||||
Recommendations:
|
||||
- [Provide actionable next steps for resolving the claim or escalating for human review.]
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Example Task:
|
||||
**Input**:
|
||||
"Process the attached car accident claim. Extract details from the police report, analyze the attached images, and provide an initial risk assessment. Highlight any inconsistencies for human review."
|
||||
|
||||
**Output**:
|
||||
#### 1. Extracted Information:
|
||||
```
|
||||
Claimant Name: John Doe
|
||||
Date of Incident: 2024-01-15
|
||||
Location: Miami, FL
|
||||
Policy Number: ABC-12345
|
||||
Summary of Incident: The claimant reports being rear-ended at a traffic light.
|
||||
|
||||
Supporting Documents:
|
||||
- Police Report: Incident verified by Officer Jane Smith; driver's statement matches claimant's report.
|
||||
- Medical Bills: $1,500 for physical therapy; injury type aligns with collision severity.
|
||||
- Repair Estimate: $4,000 for rear bumper and trunk damage.
|
||||
- Photos: Damage visible to rear bumper; no damage visible to other vehicle.
|
||||
```
|
||||
|
||||
#### 2. Consistency Analysis:
|
||||
```
|
||||
- Police Report and Claimant Statement: Consistent.
|
||||
- Medical Bills and Injury Details: Consistent with collision type.
|
||||
- Repair Estimate and Photos: Consistent; no indications of additional hidden damage.
|
||||
- No discrepancies in timeline or location details.
|
||||
```
|
||||
|
||||
#### 3. Risk Assessment:
|
||||
```
|
||||
Risk Level: Low
|
||||
Reasoning: All supporting documents align with the claimant's statement, and no unusual patterns or inconsistencies were identified.
|
||||
```
|
||||
|
||||
#### 4. Flagged Concerns and Recommendations:
|
||||
```
|
||||
Flagged Concerns:
|
||||
- None identified.
|
||||
|
||||
Recommendations:
|
||||
- Proceed with claim approval and settlement.
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Additional Notes:
|
||||
- Always ensure outputs are clear, professional, and comprehensive.
|
||||
- Use concise, evidence-backed reasoning to justify all conclusions.
|
||||
- Where relevant, prioritize human review for flagged concerns or high-risk cases.
|
||||
"""
|
||||
|
||||
# Initialize the Claims Processing Agent with RAG capabilities
|
||||
agent = Agent(
|
||||
agent_name="Claims-Processing-Agent",
|
||||
system_prompt=CLAIMS_PROCESSING_AGENT_SYS_PROMPT,
|
||||
agent_description="Agent automates claims processing and fraud detection.",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops="auto", # Auto-adjusts loops based on task complexity
|
||||
autosave=True, # Automatically saves agent state
|
||||
dashboard=False, # Disables dashboard for this example
|
||||
verbose=True, # Enables verbose mode for detailed output
|
||||
streaming_on=True, # Enables streaming for real-time processing
|
||||
dynamic_temperature_enabled=True, # Dynamically adjusts temperature for optimal performance
|
||||
saved_state_path="claims_processing_agent.json", # Path to save agent state
|
||||
user_name="swarms_corp", # User name for the agent
|
||||
retry_attempts=3, # Number of retry attempts for failed tasks
|
||||
context_length=200000, # Maximum length of the context to consider
|
||||
return_step_meta=False,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Sample task for the Claims Processing Agent
|
||||
agent.run(
|
||||
"Process the attached car accident claim. Extract details from the police report, analyze the attached images, and provide an initial risk assessment. Highlight any inconsistencies for human review."
|
||||
)
|
||||
@ -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")
|
||||
@ -0,0 +1,209 @@
|
||||
Agent Name: Chief Medical Officer
|
||||
Output: **Initial Assessment:**
|
||||
|
||||
The patient is a 45-year-old female presenting with fever, dry cough, fatigue, and mild shortness of breath. She has a medical history of controlled hypertension, is fully vaccinated for COVID-19, and reports no recent travel or known sick contacts. These symptoms are nonspecific but could be indicative of a viral respiratory infection.
|
||||
|
||||
**Differential Diagnoses:**
|
||||
|
||||
1. **Influenza:** Given the time of year (December), influenza is a possibility, especially with symptoms like fever, cough, and fatigue. Vaccination status for influenza should be checked.
|
||||
|
||||
2. **COVID-19:** Despite being fully vaccinated, breakthrough infections can occur. The symptoms align with COVID-19, and testing should be considered.
|
||||
|
||||
3. **Respiratory Syncytial Virus (RSV):** RSV can present with similar symptoms in adults, especially those with underlying health conditions like hypertension.
|
||||
|
||||
4. **Common Cold (Rhinovirus):** Although less likely given the fever, it is still a consideration.
|
||||
|
||||
5. **Other Viral Infections:** Adenovirus, parainfluenza, and human metapneumovirus could also present with these symptoms.
|
||||
|
||||
**Specialist Consultations Needed:**
|
||||
|
||||
- **Infectious Disease Specialist:** To help narrow down the viral causes and suggest specific tests.
|
||||
- **Pulmonologist:** Given the mild shortness of breath, a pulmonologist could provide insights into any underlying respiratory issues or complications.
|
||||
|
||||
**Recommended Next Steps:**
|
||||
|
||||
1. **Diagnostic Testing:**
|
||||
- Perform a rapid influenza test.
|
||||
- Conduct a COVID-19 PCR test to rule out a breakthrough infection.
|
||||
- Consider a respiratory viral panel to detect other viruses like RSV or adenovirus.
|
||||
|
||||
2. **Symptom Management:**
|
||||
- Recommend supportive care including hydration, rest, and antipyretics (e.g., acetaminophen) for fever.
|
||||
|
||||
3. **Monitoring:**
|
||||
- Advise the patient to monitor symptoms closely, especially the shortness of breath, and seek immediate care if symptoms worsen.
|
||||
|
||||
4. **Review Vaccination History:**
|
||||
- Confirm influenza vaccination status for this season.
|
||||
|
||||
5. **Follow-Up:**
|
||||
- Schedule a follow-up appointment to review test results and adjust the treatment plan as necessary.
|
||||
|
||||
**Limitations/Uncertainties:**
|
||||
|
||||
- The absence of known sick contacts and travel history makes exposure assessment challenging.
|
||||
- The possibility of co-infection with multiple viruses or secondary bacterial infection should be considered if symptoms worsen or do not improve with initial management. Agent Name: Virologist
|
||||
Output: **Detailed Analysis:**
|
||||
|
||||
**Characteristic Viral Symptoms:**
|
||||
|
||||
1. **Influenza:** Typically presents with sudden onset of high fever, cough, sore throat, muscle aches, fatigue, and headache. The dry cough and fatigue in this patient are consistent with influenza, but muscle aches and sore throat are not mentioned.
|
||||
|
||||
2. **COVID-19:** Symptoms can vary widely but often include fever, cough, fatigue, shortness of breath, and loss of taste or smell. The patient's symptoms align well with COVID-19, though the absence of anosmia (loss of smell) is noted.
|
||||
|
||||
3. **RSV:** In adults, RSV often presents with mild cold-like symptoms such as cough, fever, and fatigue. Shortness of breath can occur, especially in those with pre-existing conditions.
|
||||
|
||||
4. **Common Cold (Rhinovirus):** Typically causes milder symptoms like runny nose, cough, and sore throat. Fever is less common, making it a less likely primary cause in this case.
|
||||
|
||||
5. **Other Viral Infections:** Adenovirus and human metapneumovirus can present with respiratory symptoms similar to those of influenza and COVID-19, including fever and cough.
|
||||
|
||||
**Disease Progression Timeline:**
|
||||
|
||||
- **Influenza:** Symptoms usually appear 1-4 days after exposure and can last for about a week, with cough and fatigue potentially persisting longer.
|
||||
- **COVID-19:** Incubation period ranges from 2-14 days, with symptoms lasting from a few days to weeks depending on severity.
|
||||
- **RSV:** Incubation is 4-6 days, and symptoms typically last 1-2 weeks.
|
||||
- **Common Cold:** Symptoms usually appear 1-3 days after exposure and last about 7-10 days.
|
||||
|
||||
**Risk Factors for Severe Disease:**
|
||||
|
||||
- Controlled hypertension may increase the risk of complications from respiratory viruses like influenza and COVID-19.
|
||||
- Age (45 years) is not a significant risk factor for severe disease, but vigilance is needed.
|
||||
- The absence of other chronic conditions or immunosuppression reduces the risk of severe outcomes.
|
||||
|
||||
**Potential Complications:**
|
||||
|
||||
- **Influenza:** Can lead to pneumonia, bronchitis, or exacerbation of chronic conditions.
|
||||
- **COVID-19:** Risk of pneumonia, acute respiratory distress syndrome (ARDS), and long COVID symptoms.
|
||||
- **RSV:** Can cause bronchitis or pneumonia, particularly in older adults with underlying health issues.
|
||||
- **Common Cold:** Rarely leads to complications, but secondary bacterial infections are possible.
|
||||
|
||||
**Recommendations:**
|
||||
|
||||
1. **Diagnostic Testing:**
|
||||
- Rapid influenza and COVID-19 PCR tests are appropriate initial steps.
|
||||
- A respiratory viral panel can provide a comprehensive assessment for other viral pathogens like RSV and adenovirus.
|
||||
|
||||
2. **Symptom Management:**
|
||||
- Supportive care remains crucial. Hydration, rest, and antipyretics are recommended.
|
||||
- Consider cough suppressants or expectorants if cough is bothersome.
|
||||
|
||||
3. **Monitoring and Follow-Up:**
|
||||
- Close monitoring of symptoms, particularly shortness of breath, is essential.
|
||||
- Follow-up should be scheduled to review test results and adjust treatment.
|
||||
|
||||
4. **Vaccination Review:**
|
||||
- Confirm influenza vaccination status and encourage vaccination if not already received.
|
||||
|
||||
5. **Consideration of Co-Infections:**
|
||||
- Be vigilant for signs of bacterial superinfection, particularly if symptoms worsen or do not improve with initial management.
|
||||
|
||||
**Epidemiological Considerations:**
|
||||
|
||||
- Seasonal factors (December) increase the likelihood of influenza and RSV.
|
||||
- Current COVID-19 variants should be considered, even in vaccinated individuals.
|
||||
- Geographic prevalence and local outbreak data can provide additional context for risk assessment. Agent Name: Internist
|
||||
Output: **Internal Medicine Analysis:**
|
||||
|
||||
**1. Vital Signs and Their Implications:**
|
||||
- **Temperature:** Elevated temperature would suggest an active infection or inflammatory process.
|
||||
- **Blood Pressure:** Controlled hypertension is noted, which could predispose the patient to complications from respiratory infections.
|
||||
- **Heart Rate:** Tachycardia can be a response to fever or infection.
|
||||
- **Respiratory Rate and Oxygen Saturation:** Increased respiratory rate or decreased oxygen saturation may indicate respiratory distress or hypoxemia, particularly in the context of viral infections like COVID-19 or influenza.
|
||||
|
||||
**2. System-by-System Review:**
|
||||
|
||||
- **Cardiovascular:**
|
||||
- Monitor for signs of myocarditis or pericarditis, which can be complications of viral infections.
|
||||
- Controlled hypertension should be managed to minimize cardiovascular stress.
|
||||
|
||||
- **Respiratory:**
|
||||
- Assess for signs of pneumonia or bronchitis, common complications of viral infections.
|
||||
- Shortness of breath is a critical symptom that may indicate lower respiratory tract involvement.
|
||||
|
||||
- **Neurological:**
|
||||
- Fatigue and headache are common in viral illnesses but monitor for any signs of neurological involvement.
|
||||
|
||||
- **Musculoskeletal:**
|
||||
- Absence of muscle aches reduces the likelihood of influenza but does not rule it out.
|
||||
|
||||
**3. Impact of Existing Medical Conditions:**
|
||||
- Controlled hypertension may increase the risk of complications from respiratory infections.
|
||||
- No other chronic conditions or immunosuppression are noted, which reduces the risk of severe outcomes.
|
||||
|
||||
**4. Medication Interactions and Contraindications:**
|
||||
- Review any current medications for potential interactions with antiviral or symptomatic treatments.
|
||||
- Ensure medications for hypertension do not exacerbate respiratory symptoms or interact with treatments for the viral infection.
|
||||
|
||||
**5. Risk Stratification:**
|
||||
- Age (45 years) is not a significant risk factor for severe disease but requires vigilance.
|
||||
- Controlled hypertension is a relevant risk factor for complications.
|
||||
- Absence of other chronic conditions suggests a lower risk for severe outcomes.
|
||||
|
||||
**Physical Examination Findings:**
|
||||
- Focus on respiratory examination for signs of distress, consolidation, or wheezing.
|
||||
- Cardiovascular examination should assess for any signs of increased workload or complications.
|
||||
- General examination should assess for signs of systemic involvement or secondary bacterial infection.
|
||||
|
||||
**System-Specific Symptoms:**
|
||||
- Respiratory: Cough, shortness of breath.
|
||||
- General: Fatigue, fever.
|
||||
- Neurological: Headache.
|
||||
|
||||
**Relevant Lab Abnormalities:**
|
||||
- Elevated inflammatory markers (CRP, ESR) may indicate an active infection.
|
||||
- CBC may show leukocytosis or lymphopenia, common in viral infections.
|
||||
- Abnormal liver function tests could indicate systemic involvement.
|
||||
|
||||
**Risk Factors for Complications:**
|
||||
- Controlled hypertension.
|
||||
- Potential for bacterial superinfection if symptoms worsen or do not improve.
|
||||
|
||||
**Recommendations:**
|
||||
- **Diagnostic Testing:** Rapid influenza and COVID-19 tests, respiratory viral panel if needed.
|
||||
- **Symptom Management:** Supportive care, hydration, rest, antipyretics, cough suppressants if needed.
|
||||
- **Monitoring and Follow-Up:** Monitor respiratory symptoms closely, schedule follow-up for test results and treatment adjustment.
|
||||
- **Vaccination Review:** Confirm influenza vaccination status.
|
||||
- **Consideration of Co-Infections:** Be vigilant for bacterial superinfection signs.
|
||||
|
||||
**Epidemiological Considerations:**
|
||||
- Seasonal factors and current viral prevalence should guide diagnostic suspicion and management. Agent Name: Diagnostic Synthesizer
|
||||
Output: **Final Diagnostic Assessment**
|
||||
|
||||
**1. Primary Diagnosis with Confidence Level:**
|
||||
- **Viral Respiratory Infection (likely COVID-19 or Influenza):** Confidence Level: Moderate to High
|
||||
- Based on the presence of fever, cough, shortness of breath, fatigue, and elevated inflammatory markers, a viral respiratory infection is the most probable diagnosis. The seasonality and current prevalence of COVID-19 and influenza further support this diagnosis.
|
||||
|
||||
**2. Supporting Evidence Summary:**
|
||||
- **Clinical Presentation:** Fever, cough, shortness of breath, fatigue, and headache are indicative of a viral infection.
|
||||
- **Vital Signs:** Tachycardia and potential respiratory distress align with an active infection.
|
||||
- **Lab Abnormalities:** Elevated CRP/ESR and possible leukocytosis or lymphopenia are common in viral infections.
|
||||
- **Epidemiological Factors:** Current high prevalence of COVID-19 and influenza.
|
||||
|
||||
**3. Alternative Diagnoses to Consider:**
|
||||
- **Bacterial Pneumonia:** Consider if symptoms persist or worsen, particularly if there is consolidation on examination or imaging.
|
||||
- **Myocarditis or Pericarditis:** These are potential complications of viral infections, especially if there are cardiovascular symptoms or abnormalities.
|
||||
- **Non-Infectious Causes:** Less likely given the acute presentation but consider if symptoms do not resolve with typical viral course.
|
||||
|
||||
**4. Recommended Confirmatory Tests:**
|
||||
- **Rapid Influenza Test**
|
||||
- **COVID-19 PCR or Antigen Test**
|
||||
- **Respiratory Viral Panel:** If initial tests are negative and symptoms persist.
|
||||
- **Chest X-ray or CT Scan:** If there is suspicion of pneumonia or other complications.
|
||||
|
||||
**5. Red Flags or Warning Signs:**
|
||||
- Worsening shortness of breath or persistent hypoxemia.
|
||||
- Chest pain or signs of cardiovascular involvement.
|
||||
- Persistent high fever or new onset of symptoms indicating a secondary bacterial infection.
|
||||
|
||||
**6. Follow-up Recommendations:**
|
||||
- **Symptom Monitoring:** Close monitoring of respiratory symptoms and general condition.
|
||||
- **Follow-up Appointment:** Schedule follow-up to review test results and adjust treatment as necessary.
|
||||
- **Vaccination Status:** Ensure influenza vaccination is up to date and consider COVID-19 booster if eligible.
|
||||
- **Patient Education:** Inform about signs of worsening condition and when to seek immediate care.
|
||||
|
||||
**Documentation Requirements:**
|
||||
- **Reasoning Chain:** The diagnosis is based on clinical presentation, lab findings, and current epidemiological data.
|
||||
- **Evidence Quality Assessment:** Moderate to high confidence based on reliable clinical and laboratory evidence.
|
||||
- **Confidence Levels for Each Diagnosis:** Primary diagnosis is given moderate to high confidence, while alternatives are considered with lower probability unless symptoms evolve.
|
||||
- **Knowledge Gaps Identified:** Awaiting confirmatory testing results to solidify the diagnosis.
|
||||
- **Risk Assessment:** Controlled hypertension presents a moderate risk for complications, necessitating vigilant monitoring.
|
||||
@ -0,0 +1,248 @@
|
||||
"""
|
||||
- For each diagnosis, pull lab results,
|
||||
- egfr
|
||||
- for each diagnosis, pull lab ranges,
|
||||
- pull ranges for diagnosis
|
||||
|
||||
- if the diagnosis is x, then the lab ranges should be a to b
|
||||
- train the agents, increase the load of input
|
||||
- medical history sent to the agent
|
||||
- setup rag for the agents
|
||||
- run the first agent -> kidney disease -> don't know the stage -> stage 2 -> lab results -> indicative of stage 3 -> the case got elavated ->
|
||||
- how to manage diseases and by looking at correlating lab, docs, diagnoses
|
||||
- put docs in rag ->
|
||||
- monitoring, evaluation, and treatment
|
||||
- can we confirm for every diagnosis -> monitoring, evaluation, and treatment, specialized for these things
|
||||
- find diagnosis -> or have diagnosis, -> for each diagnosis are there evidence of those 3 things
|
||||
- swarm of those 4 agents, ->
|
||||
- fda api for healthcare for commerically available papers
|
||||
-
|
||||
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from swarms import Agent, AgentRearrange, create_file_in_folder
|
||||
|
||||
chief_medical_officer = Agent(
|
||||
agent_name="Chief Medical Officer",
|
||||
system_prompt="""You are the Chief Medical Officer coordinating a team of medical specialists for viral disease diagnosis.
|
||||
Your responsibilities include:
|
||||
- Gathering initial patient symptoms and medical history
|
||||
- Coordinating with specialists to form differential diagnoses
|
||||
- Synthesizing different specialist opinions into a cohesive diagnosis
|
||||
- Ensuring all relevant symptoms and test results are considered
|
||||
- Making final diagnostic recommendations
|
||||
- Suggesting treatment plans based on team input
|
||||
- Identifying when additional specialists need to be consulted
|
||||
- For each diferrential diagnosis provide minimum lab ranges to meet that diagnosis or be indicative of that diagnosis minimum and maximum
|
||||
|
||||
Format all responses with clear sections for:
|
||||
- Initial Assessment (include preliminary ICD-10 codes for symptoms)
|
||||
- Differential Diagnoses (with corresponding ICD-10 codes)
|
||||
- Specialist Consultations Needed
|
||||
- Recommended Next Steps
|
||||
|
||||
|
||||
""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
virologist = Agent(
|
||||
agent_name="Virologist",
|
||||
system_prompt="""You are a specialist in viral diseases. For each case, provide:
|
||||
|
||||
Clinical Analysis:
|
||||
- Detailed viral symptom analysis
|
||||
- Disease progression timeline
|
||||
- Risk factors and complications
|
||||
|
||||
Coding Requirements:
|
||||
- List relevant ICD-10 codes for:
|
||||
* Confirmed viral conditions
|
||||
* Suspected viral conditions
|
||||
* Associated symptoms
|
||||
* Complications
|
||||
- Include both:
|
||||
* Primary diagnostic codes
|
||||
* Secondary condition codes
|
||||
|
||||
Document all findings using proper medical coding standards and include rationale for code selection.""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
internist = Agent(
|
||||
agent_name="Internist",
|
||||
system_prompt="""You are an Internal Medicine specialist responsible for comprehensive evaluation.
|
||||
|
||||
For each case, provide:
|
||||
|
||||
Clinical Assessment:
|
||||
- System-by-system review
|
||||
- Vital signs analysis
|
||||
- Comorbidity evaluation
|
||||
|
||||
Medical Coding:
|
||||
- ICD-10 codes for:
|
||||
* Primary conditions
|
||||
* Secondary diagnoses
|
||||
* Complications
|
||||
* Chronic conditions
|
||||
* Signs and symptoms
|
||||
- Include hierarchical condition category (HCC) codes where applicable
|
||||
|
||||
Document supporting evidence for each code selected.""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
medical_coder = Agent(
|
||||
agent_name="Medical Coder",
|
||||
system_prompt="""You are a certified medical coder responsible for:
|
||||
|
||||
Primary Tasks:
|
||||
1. Reviewing all clinical documentation
|
||||
2. Assigning accurate ICD-10 codes
|
||||
3. Ensuring coding compliance
|
||||
4. Documenting code justification
|
||||
|
||||
Coding Process:
|
||||
- Review all specialist inputs
|
||||
- Identify primary and secondary diagnoses
|
||||
- Assign appropriate ICD-10 codes
|
||||
- Document supporting evidence
|
||||
- Note any coding queries
|
||||
|
||||
Output Format:
|
||||
1. Primary Diagnosis Codes
|
||||
- ICD-10 code
|
||||
- Description
|
||||
- Supporting documentation
|
||||
2. Secondary Diagnosis Codes
|
||||
- Listed in order of clinical significance
|
||||
3. Symptom Codes
|
||||
4. Complication Codes
|
||||
5. Coding Notes""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
synthesizer = Agent(
|
||||
agent_name="Diagnostic Synthesizer",
|
||||
system_prompt="""You are responsible for creating the final diagnostic and coding assessment.
|
||||
|
||||
Synthesis Requirements:
|
||||
1. Integrate all specialist findings
|
||||
2. Reconcile any conflicting diagnoses
|
||||
3. Verify coding accuracy and completeness
|
||||
|
||||
Final Report Sections:
|
||||
1. Clinical Summary
|
||||
- Primary diagnosis with ICD-10
|
||||
- Secondary diagnoses with ICD-10
|
||||
- Supporting evidence
|
||||
2. Coding Summary
|
||||
- Complete code list with descriptions
|
||||
- Code hierarchy and relationships
|
||||
- Supporting documentation
|
||||
3. Recommendations
|
||||
- Additional testing needed
|
||||
- Follow-up care
|
||||
- Documentation improvements needed
|
||||
|
||||
Include confidence levels and evidence quality for all diagnoses and codes.""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create agent list
|
||||
agents = [
|
||||
chief_medical_officer,
|
||||
virologist,
|
||||
internist,
|
||||
medical_coder,
|
||||
synthesizer,
|
||||
]
|
||||
|
||||
# Define diagnostic flow
|
||||
flow = f"""{chief_medical_officer.agent_name} -> {virologist.agent_name} -> {internist.agent_name} -> {medical_coder.agent_name} -> {synthesizer.agent_name}"""
|
||||
|
||||
# Create the swarm system
|
||||
diagnosis_system = AgentRearrange(
|
||||
name="Medical-coding-diagnosis-swarm",
|
||||
description="Comprehensive medical diagnosis and coding system",
|
||||
agents=agents,
|
||||
flow=flow,
|
||||
max_loops=1,
|
||||
output_type="all",
|
||||
)
|
||||
|
||||
|
||||
def generate_coding_report(diagnosis_output: str) -> str:
|
||||
"""
|
||||
Generate a structured medical coding report from the diagnosis output.
|
||||
"""
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
report = f"""# Medical Diagnosis and Coding Report
|
||||
Generated: {timestamp}
|
||||
|
||||
## Clinical Summary
|
||||
{diagnosis_output}
|
||||
|
||||
## Coding Summary
|
||||
### Primary Diagnosis Codes
|
||||
[Extracted from synthesis]
|
||||
|
||||
### Secondary Diagnosis Codes
|
||||
[Extracted from synthesis]
|
||||
|
||||
### Symptom Codes
|
||||
[Extracted from synthesis]
|
||||
|
||||
### Procedure Codes (if applicable)
|
||||
[Extracted from synthesis]
|
||||
|
||||
## Documentation and Compliance Notes
|
||||
- Code justification
|
||||
- Supporting documentation references
|
||||
- Any coding queries or clarifications needed
|
||||
|
||||
## Recommendations
|
||||
- Additional documentation needed
|
||||
- Suggested follow-up
|
||||
- Coding optimization opportunities
|
||||
"""
|
||||
return report
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example patient case
|
||||
patient_case = """
|
||||
Patient: 45-year-old White Male
|
||||
|
||||
Lab Results:
|
||||
- egfr
|
||||
- 59 ml / min / 1.73
|
||||
- non african-american
|
||||
|
||||
"""
|
||||
|
||||
# Add timestamp to the patient case
|
||||
case_info = f"Timestamp: {datetime.now()}\nPatient Information: {patient_case}"
|
||||
|
||||
# Run the diagnostic process
|
||||
diagnosis = diagnosis_system.run(case_info)
|
||||
|
||||
# Generate coding report
|
||||
coding_report = generate_coding_report(diagnosis)
|
||||
|
||||
# Create reports
|
||||
create_file_in_folder(
|
||||
"reports", "medical_diagnosis_report.md", diagnosis
|
||||
)
|
||||
create_file_in_folder(
|
||||
"reports", "medical_coding_report.md", coding_report
|
||||
)
|
||||
@ -0,0 +1,342 @@
|
||||
# Medical Diagnosis and Coding Report
|
||||
Generated: 2024-12-09 11:17:38
|
||||
|
||||
## Clinical Summary
|
||||
Agent Name: Chief Medical Officer
|
||||
Output: **Initial Assessment**
|
||||
|
||||
- **Patient Information**: 45-year-old White Male
|
||||
- **Key Lab Result**:
|
||||
- eGFR (Estimated Glomerular Filtration Rate): 59 ml/min/1.73 m²
|
||||
- **Preliminary ICD-10 Codes for Symptoms**:
|
||||
- N18.3: Chronic kidney disease, stage 3 (moderate)
|
||||
|
||||
**Differential Diagnoses**
|
||||
|
||||
1. **Chronic Kidney Disease (CKD)**
|
||||
- **ICD-10 Code**: N18.3
|
||||
- **Minimum Lab Range**: eGFR 30-59 ml/min/1.73 m² (indicative of stage 3 CKD)
|
||||
- **Maximum Lab Range**: eGFR 59 ml/min/1.73 m²
|
||||
|
||||
2. **Possible Acute Kidney Injury (AKI) on Chronic Kidney Disease**
|
||||
- **ICD-10 Code**: N17.9 (Acute kidney failure, unspecified) superimposed on N18.3
|
||||
- **Minimum Lab Range**: Rapid decline in eGFR or increase in serum creatinine
|
||||
- **Maximum Lab Range**: Dependent on baseline kidney function and rapidity of change
|
||||
|
||||
3. **Hypertensive Nephropathy**
|
||||
- **ICD-10 Code**: I12.9 (Hypertensive chronic kidney disease with stage 1 through stage 4 chronic kidney disease, or unspecified chronic kidney disease)
|
||||
- **Minimum Lab Range**: eGFR 30-59 ml/min/1.73 m² with evidence of hypertension
|
||||
- **Maximum Lab Range**: eGFR 59 ml/min/1.73 m²
|
||||
|
||||
**Specialist Consultations Needed**
|
||||
|
||||
- **Nephrologist**: To assess the kidney function and evaluate for CKD or other renal pathologies.
|
||||
- **Cardiologist**: If hypertensive nephropathy is suspected, to manage associated cardiovascular risks and blood pressure.
|
||||
- **Endocrinologist**: If there are any signs of diabetes or metabolic syndrome contributing to renal impairment.
|
||||
|
||||
**Recommended Next Steps**
|
||||
|
||||
1. **Detailed Medical History and Physical Examination**:
|
||||
- Assess for symptoms such as fatigue, swelling, changes in urination, or hypertension.
|
||||
- Review any history of diabetes, hypertension, or cardiovascular disease.
|
||||
|
||||
2. **Additional Laboratory Tests**:
|
||||
- Serum creatinine and blood urea nitrogen (BUN) to further evaluate kidney function.
|
||||
- Urinalysis to check for proteinuria or hematuria.
|
||||
- Lipid profile and fasting glucose to assess for metabolic syndrome.
|
||||
|
||||
3. **Imaging Studies**:
|
||||
- Renal ultrasound to evaluate kidney size and rule out obstructive causes.
|
||||
|
||||
4. **Blood Pressure Monitoring**:
|
||||
- Regular monitoring to assess for hypertension which could contribute to kidney damage.
|
||||
|
||||
5. **Referral to Nephrology**:
|
||||
- For comprehensive evaluation and management of kidney disease.
|
||||
|
||||
6. **Patient Education**:
|
||||
- Discuss lifestyle modifications such as diet, exercise, and smoking cessation to slow the progression of kidney disease.
|
||||
|
||||
By following these steps, we can ensure a thorough evaluation of the patient's condition and formulate an appropriate management plan. Agent Name: Virologist
|
||||
Output: **Clinical Analysis for Viral Diseases**
|
||||
|
||||
Given the current patient information, there is no direct indication of a viral disease from the provided data. However, if a viral etiology is suspected or confirmed, the following analysis can be applied:
|
||||
|
||||
### Clinical Analysis:
|
||||
- **Detailed Viral Symptom Analysis**:
|
||||
- Symptoms of viral infections can be diverse but often include fever, fatigue, muscle aches, and respiratory symptoms such as cough or sore throat. In the context of renal impairment, certain viral infections can lead to or exacerbate kidney issues, such as Hepatitis B or C, HIV, or cytomegalovirus (CMV).
|
||||
|
||||
- **Disease Progression Timeline**:
|
||||
- Viral infections typically have an incubation period ranging from a few days to weeks. The acute phase can last from several days to weeks, with symptoms peaking and then gradually resolving. Chronic viral infections, such as Hepatitis B or C, can lead to long-term complications, including kidney damage.
|
||||
|
||||
- **Risk Factors and Complications**:
|
||||
- Risk factors for viral infections include immunosuppression, exposure to infected individuals, travel history, and underlying health conditions. Complications can include acute kidney injury, chronic kidney disease progression, and systemic involvement leading to multi-organ dysfunction.
|
||||
|
||||
### Coding Requirements:
|
||||
|
||||
#### Relevant ICD-10 Codes:
|
||||
|
||||
- **Confirmed Viral Conditions**:
|
||||
- **B18.1**: Chronic viral hepatitis B
|
||||
- **B18.2**: Chronic viral hepatitis C
|
||||
- **B20**: HIV disease resulting in infectious and parasitic diseases
|
||||
|
||||
- **Suspected Viral Conditions**:
|
||||
- **B34.9**: Viral infection, unspecified
|
||||
|
||||
- **Associated Symptoms**:
|
||||
- **R50.9**: Fever, unspecified
|
||||
- **R53.83**: Other fatigue
|
||||
- **R05**: Cough
|
||||
|
||||
- **Complications**:
|
||||
- **N17.9**: Acute kidney failure, unspecified (if viral infection leads to AKI)
|
||||
- **N18.9**: Chronic kidney disease, unspecified (if progression due to viral infection)
|
||||
|
||||
#### Primary and Secondary Diagnostic Codes:
|
||||
|
||||
- **Primary Diagnostic Codes**:
|
||||
- Use the specific viral infection code as primary if confirmed (e.g., B18.2 for Hepatitis C).
|
||||
|
||||
- **Secondary Condition Codes**:
|
||||
- Use codes for symptoms or complications as secondary (e.g., N17.9 for AKI if due to viral infection).
|
||||
|
||||
### Rationale for Code Selection:
|
||||
|
||||
- **B18.1 and B18.2**: Selected for confirmed chronic hepatitis B or C, which can have renal complications.
|
||||
- **B20**: Used if HIV is confirmed, given its potential impact on renal function.
|
||||
- **B34.9**: Utilized when a viral infection is suspected but not yet identified.
|
||||
- **R50.9, R53.83, R05**: Common symptoms associated with viral infections.
|
||||
- **N17.9, N18.9**: Codes for renal complications potentially exacerbated by viral infections.
|
||||
|
||||
### Documentation:
|
||||
|
||||
- Ensure thorough documentation of clinical findings, suspected or confirmed viral infections, and associated symptoms or complications to justify the selected ICD-10 codes.
|
||||
- Follow-up with additional testing or specialist referrals as needed to confirm or rule out viral etiologies and manage complications effectively. Agent Name: Internist
|
||||
Output: To provide a comprehensive evaluation as an Internal Medicine specialist, let's conduct a detailed clinical assessment and medical coding for the presented case. This will involve a system-by-system review, analysis of vital signs, and evaluation of comorbidities, followed by appropriate ICD-10 coding.
|
||||
|
||||
### Clinical Assessment:
|
||||
|
||||
#### System-by-System Review:
|
||||
1. **Respiratory System:**
|
||||
- Evaluate for symptoms such as cough, shortness of breath, or wheezing.
|
||||
- Consider potential viral or bacterial infections affecting the respiratory tract.
|
||||
|
||||
2. **Cardiovascular System:**
|
||||
- Assess for any signs of heart failure or hypertension.
|
||||
- Look for symptoms like chest pain, palpitations, or edema.
|
||||
|
||||
3. **Gastrointestinal System:**
|
||||
- Check for symptoms such as nausea, vomiting, diarrhea, or abdominal pain.
|
||||
- Consider liver function if hepatitis is suspected.
|
||||
|
||||
4. **Renal System:**
|
||||
- Monitor for signs of acute kidney injury or chronic kidney disease.
|
||||
- Evaluate urine output and creatinine levels.
|
||||
|
||||
5. **Neurological System:**
|
||||
- Assess for headaches, dizziness, or any focal neurological deficits.
|
||||
- Consider viral encephalitis if neurological symptoms are present.
|
||||
|
||||
6. **Musculoskeletal System:**
|
||||
- Look for muscle aches or joint pain, common in viral infections.
|
||||
|
||||
7. **Integumentary System:**
|
||||
- Check for rashes or skin lesions, which may indicate viral infections like herpes or CMV.
|
||||
|
||||
8. **Immune System:**
|
||||
- Consider immunosuppression status, especially in the context of HIV or other chronic infections.
|
||||
|
||||
#### Vital Signs Analysis:
|
||||
- **Temperature:** Evaluate for fever, which may indicate an infection.
|
||||
- **Blood Pressure:** Check for hypertension or hypotension.
|
||||
- **Heart Rate:** Assess for tachycardia or bradycardia.
|
||||
- **Respiratory Rate:** Monitor for tachypnea.
|
||||
- **Oxygen Saturation:** Ensure adequate oxygenation, especially in respiratory infections.
|
||||
|
||||
#### Comorbidity Evaluation:
|
||||
- Assess for chronic conditions such as diabetes, hypertension, or chronic kidney disease.
|
||||
- Consider the impact of these conditions on the current clinical presentation and potential complications.
|
||||
|
||||
### Medical Coding:
|
||||
|
||||
#### ICD-10 Codes:
|
||||
|
||||
1. **Primary Conditions:**
|
||||
- If a specific viral infection is confirmed, use the appropriate code (e.g., B18.2 for chronic hepatitis C).
|
||||
|
||||
2. **Secondary Diagnoses:**
|
||||
- **B34.9:** Viral infection, unspecified (if viral etiology is suspected but not confirmed).
|
||||
- **R50.9:** Fever, unspecified (common symptom in infections).
|
||||
- **R53.83:** Other fatigue (common in viral infections).
|
||||
|
||||
3. **Complications:**
|
||||
- **N17.9:** Acute kidney failure, unspecified (if there is renal involvement).
|
||||
- **N18.9:** Chronic kidney disease, unspecified (if there is progression due to infection).
|
||||
|
||||
4. **Chronic Conditions:**
|
||||
- **I10:** Essential (primary) hypertension (if present).
|
||||
- **E11.9:** Type 2 diabetes mellitus without complications (if present).
|
||||
|
||||
5. **Signs and Symptoms:**
|
||||
- **R05:** Cough (common respiratory symptom).
|
||||
- **M79.1:** Myalgia (muscle pain).
|
||||
|
||||
#### Hierarchical Condition Category (HCC) Codes:
|
||||
- **HCC 18:** Diabetes with chronic complications (if applicable).
|
||||
- **HCC 85:** Congestive heart failure (if applicable).
|
||||
|
||||
### Documentation Supporting Evidence:
|
||||
- Ensure documentation includes detailed clinical findings, symptoms, and any laboratory or imaging results that support the diagnosis.
|
||||
- Include any history of chronic conditions or recent changes in health status.
|
||||
- Document any suspected or confirmed viral infections, along with their impact on the patient's health.
|
||||
|
||||
### Conclusion:
|
||||
This comprehensive evaluation and coding approach allows for accurate diagnosis and management of the patient's condition, considering both acute and chronic aspects of their health. Proper documentation and coding facilitate effective communication and continuity of care. Agent Name: Medical Coder
|
||||
Output: ### Medical Coding Summary
|
||||
|
||||
#### 1. Primary Diagnosis Codes
|
||||
- **ICD-10 Code:** B18.2
|
||||
- **Description:** Chronic viral hepatitis C
|
||||
- **Supporting Documentation:** The diagnosis of chronic hepatitis C is confirmed through serological testing and liver function tests indicating chronic viral infection.
|
||||
|
||||
#### 2. Secondary Diagnosis Codes
|
||||
- **B34.9:** Viral infection, unspecified
|
||||
- **Supporting Documentation:** Suspected viral etiology without specific identification.
|
||||
- **R50.9:** Fever, unspecified
|
||||
- **Supporting Documentation:** Documented fever without a definitive cause.
|
||||
- **R53.83:** Other fatigue
|
||||
- **Supporting Documentation:** Patient reports persistent fatigue, common in viral infections.
|
||||
- **I10:** Essential (primary) hypertension
|
||||
- **Supporting Documentation:** History of hypertension with current blood pressure readings.
|
||||
- **E11.9:** Type 2 diabetes mellitus without complications
|
||||
- **Supporting Documentation:** Documented history of type 2 diabetes, managed with oral hypoglycemics.
|
||||
|
||||
#### 3. Symptom Codes
|
||||
- **R05:** Cough
|
||||
- **Supporting Documentation:** Patient presents with a persistent cough, noted in the respiratory evaluation.
|
||||
- **M79.1:** Myalgia
|
||||
- **Supporting Documentation:** Patient reports muscle pain, consistent with viral infections.
|
||||
|
||||
#### 4. Complication Codes
|
||||
- **N17.9:** Acute kidney failure, unspecified
|
||||
- **Supporting Documentation:** Elevated creatinine levels and reduced urine output indicative of renal involvement.
|
||||
- **N18.9:** Chronic kidney disease, unspecified
|
||||
- **Supporting Documentation:** Documented chronic kidney disease stage, with baseline creatinine levels.
|
||||
|
||||
#### 5. Coding Notes
|
||||
- Ensure all clinical findings and laboratory results supporting the diagnoses are documented in the patient's medical record.
|
||||
- Confirm the presence of chronic conditions and their management strategies.
|
||||
- Monitor for any changes in the patient's condition that may require code updates or additions.
|
||||
- Address any coding queries related to unspecified viral infections by seeking further diagnostic clarification if possible.
|
||||
|
||||
This coding summary provides a structured approach to documenting the patient's current health status, ensuring accurate and compliant ICD-10 coding. Agent Name: Diagnostic Synthesizer
|
||||
Output: ### Final Diagnostic and Coding Assessment
|
||||
|
||||
#### Clinical Summary
|
||||
|
||||
**Primary Diagnosis:**
|
||||
- **ICD-10 Code:** B18.2
|
||||
- **Description:** Chronic viral hepatitis C
|
||||
- **Supporting Evidence:** This diagnosis is substantiated by serological testing and liver function tests indicating a chronic viral infection. The confidence level for this diagnosis is high, with high-quality evidence from laboratory results.
|
||||
|
||||
**Secondary Diagnoses:**
|
||||
1. **ICD-10 Code:** B34.9
|
||||
- **Description:** Viral infection, unspecified
|
||||
- **Supporting Evidence:** The suspected viral etiology lacks specific identification. Confidence level is moderate due to limited specificity in viral identification.
|
||||
|
||||
2. **ICD-10 Code:** R50.9
|
||||
- **Description:** Fever, unspecified
|
||||
- **Supporting Evidence:** Documented fever without a definitive cause. Confidence level is moderate, supported by clinical observation.
|
||||
|
||||
3. **ICD-10 Code:** R53.83
|
||||
- **Description:** Other fatigue
|
||||
- **Supporting Evidence:** Patient reports persistent fatigue, often associated with viral infections. Confidence level is moderate, based on patient-reported symptoms.
|
||||
|
||||
4. **ICD-10 Code:** I10
|
||||
- **Description:** Essential (primary) hypertension
|
||||
- **Supporting Evidence:** History of hypertension corroborated by current blood pressure readings. Confidence level is high, with consistent clinical evidence.
|
||||
|
||||
5. **ICD-10 Code:** E11.9
|
||||
- **Description:** Type 2 diabetes mellitus without complications
|
||||
- **Supporting Evidence:** Managed with oral hypoglycemics, with a documented history. Confidence level is high, with strong management records.
|
||||
|
||||
**Symptom Codes:**
|
||||
- **ICD-10 Code:** R05
|
||||
- **Description:** Cough
|
||||
- **Supporting Evidence:** Persistent cough noted in respiratory evaluation. Confidence level is moderate, based on clinical observation.
|
||||
|
||||
- **ICD-10 Code:** M79.1
|
||||
- **Description:** Myalgia
|
||||
- **Supporting Evidence:** Muscle pain reported by the patient, consistent with viral infections. Confidence level is moderate, based on patient-reported symptoms.
|
||||
|
||||
**Complication Codes:**
|
||||
1. **ICD-10 Code:** N17.9
|
||||
- **Description:** Acute kidney failure, unspecified
|
||||
- **Supporting Evidence:** Elevated creatinine levels and reduced urine output suggest renal involvement. Confidence level is high, supported by laboratory data.
|
||||
|
||||
2. **ICD-10 Code:** N18.9
|
||||
- **Description:** Chronic kidney disease, unspecified
|
||||
- **Supporting Evidence:** Documented chronic kidney disease stage with baseline creatinine levels. Confidence level is high, with consistent clinical data.
|
||||
|
||||
#### Coding Summary
|
||||
|
||||
**Complete Code List with Descriptions:**
|
||||
- B18.2: Chronic viral hepatitis C
|
||||
- B34.9: Viral infection, unspecified
|
||||
- R50.9: Fever, unspecified
|
||||
- R53.83: Other fatigue
|
||||
- I10: Essential (primary) hypertension
|
||||
- E11.9: Type 2 diabetes mellitus without complications
|
||||
- R05: Cough
|
||||
- M79.1: Myalgia
|
||||
- N17.9: Acute kidney failure, unspecified
|
||||
- N18.9: Chronic kidney disease, unspecified
|
||||
|
||||
**Code Hierarchy and Relationships:**
|
||||
- The primary diagnosis (B18.2) is the central focus, with secondary diagnoses and symptoms potentially related to or exacerbated by the chronic hepatitis C infection.
|
||||
- Complications (N17.9 and N18.9) may be linked to the primary diagnosis and other chronic conditions like diabetes and hypertension.
|
||||
|
||||
**Supporting Documentation:**
|
||||
- Ensure that all clinical findings and laboratory results supporting the diagnoses are documented in the patient's medical record.
|
||||
- Confirm the presence of chronic conditions and their management strategies.
|
||||
- Monitor for any changes in the patient's condition that may require code updates or additions.
|
||||
|
||||
#### Recommendations
|
||||
|
||||
1. **Additional Testing Needed:**
|
||||
- Further diagnostic testing is recommended to clarify the unspecified viral infection (B34.9) and to monitor kidney function.
|
||||
|
||||
2. **Follow-up Care:**
|
||||
- Regular follow-up appointments to manage chronic conditions such as hypertension and diabetes.
|
||||
- Monitor renal function and adjust treatment plans as necessary.
|
||||
|
||||
3. **Documentation Improvements Needed:**
|
||||
- Enhance documentation specificity for the unspecified viral infection.
|
||||
- Ensure comprehensive records of all chronic conditions and their management strategies.
|
||||
|
||||
These recommendations aim to improve diagnostic accuracy and patient care continuity.
|
||||
|
||||
## Coding Summary
|
||||
### Primary Diagnosis Codes
|
||||
[Extracted from synthesis]
|
||||
|
||||
### Secondary Diagnosis Codes
|
||||
[Extracted from synthesis]
|
||||
|
||||
### Symptom Codes
|
||||
[Extracted from synthesis]
|
||||
|
||||
### Procedure Codes (if applicable)
|
||||
[Extracted from synthesis]
|
||||
|
||||
## Documentation and Compliance Notes
|
||||
- Code justification
|
||||
- Supporting documentation references
|
||||
- Any coding queries or clarifications needed
|
||||
|
||||
## Recommendations
|
||||
- Additional documentation needed
|
||||
- Suggested follow-up
|
||||
- Coding optimization opportunities
|
||||
|
||||
@ -0,0 +1,314 @@
|
||||
Agent Name: Chief Medical Officer
|
||||
Output: **Initial Assessment**
|
||||
|
||||
- **Patient Information**: 45-year-old White Male
|
||||
- **Key Lab Result**:
|
||||
- eGFR (Estimated Glomerular Filtration Rate): 59 ml/min/1.73 m²
|
||||
- **Preliminary ICD-10 Codes for Symptoms**:
|
||||
- N18.3: Chronic kidney disease, stage 3 (moderate)
|
||||
|
||||
**Differential Diagnoses**
|
||||
|
||||
1. **Chronic Kidney Disease (CKD)**
|
||||
- **ICD-10 Code**: N18.3
|
||||
- **Minimum Lab Range**: eGFR 30-59 ml/min/1.73 m² (indicative of stage 3 CKD)
|
||||
- **Maximum Lab Range**: eGFR 59 ml/min/1.73 m²
|
||||
|
||||
2. **Possible Acute Kidney Injury (AKI) on Chronic Kidney Disease**
|
||||
- **ICD-10 Code**: N17.9 (Acute kidney failure, unspecified) superimposed on N18.3
|
||||
- **Minimum Lab Range**: Rapid decline in eGFR or increase in serum creatinine
|
||||
- **Maximum Lab Range**: Dependent on baseline kidney function and rapidity of change
|
||||
|
||||
3. **Hypertensive Nephropathy**
|
||||
- **ICD-10 Code**: I12.9 (Hypertensive chronic kidney disease with stage 1 through stage 4 chronic kidney disease, or unspecified chronic kidney disease)
|
||||
- **Minimum Lab Range**: eGFR 30-59 ml/min/1.73 m² with evidence of hypertension
|
||||
- **Maximum Lab Range**: eGFR 59 ml/min/1.73 m²
|
||||
|
||||
**Specialist Consultations Needed**
|
||||
|
||||
- **Nephrologist**: To assess the kidney function and evaluate for CKD or other renal pathologies.
|
||||
- **Cardiologist**: If hypertensive nephropathy is suspected, to manage associated cardiovascular risks and blood pressure.
|
||||
- **Endocrinologist**: If there are any signs of diabetes or metabolic syndrome contributing to renal impairment.
|
||||
|
||||
**Recommended Next Steps**
|
||||
|
||||
1. **Detailed Medical History and Physical Examination**:
|
||||
- Assess for symptoms such as fatigue, swelling, changes in urination, or hypertension.
|
||||
- Review any history of diabetes, hypertension, or cardiovascular disease.
|
||||
|
||||
2. **Additional Laboratory Tests**:
|
||||
- Serum creatinine and blood urea nitrogen (BUN) to further evaluate kidney function.
|
||||
- Urinalysis to check for proteinuria or hematuria.
|
||||
- Lipid profile and fasting glucose to assess for metabolic syndrome.
|
||||
|
||||
3. **Imaging Studies**:
|
||||
- Renal ultrasound to evaluate kidney size and rule out obstructive causes.
|
||||
|
||||
4. **Blood Pressure Monitoring**:
|
||||
- Regular monitoring to assess for hypertension which could contribute to kidney damage.
|
||||
|
||||
5. **Referral to Nephrology**:
|
||||
- For comprehensive evaluation and management of kidney disease.
|
||||
|
||||
6. **Patient Education**:
|
||||
- Discuss lifestyle modifications such as diet, exercise, and smoking cessation to slow the progression of kidney disease.
|
||||
|
||||
By following these steps, we can ensure a thorough evaluation of the patient's condition and formulate an appropriate management plan. Agent Name: Virologist
|
||||
Output: **Clinical Analysis for Viral Diseases**
|
||||
|
||||
Given the current patient information, there is no direct indication of a viral disease from the provided data. However, if a viral etiology is suspected or confirmed, the following analysis can be applied:
|
||||
|
||||
### Clinical Analysis:
|
||||
- **Detailed Viral Symptom Analysis**:
|
||||
- Symptoms of viral infections can be diverse but often include fever, fatigue, muscle aches, and respiratory symptoms such as cough or sore throat. In the context of renal impairment, certain viral infections can lead to or exacerbate kidney issues, such as Hepatitis B or C, HIV, or cytomegalovirus (CMV).
|
||||
|
||||
- **Disease Progression Timeline**:
|
||||
- Viral infections typically have an incubation period ranging from a few days to weeks. The acute phase can last from several days to weeks, with symptoms peaking and then gradually resolving. Chronic viral infections, such as Hepatitis B or C, can lead to long-term complications, including kidney damage.
|
||||
|
||||
- **Risk Factors and Complications**:
|
||||
- Risk factors for viral infections include immunosuppression, exposure to infected individuals, travel history, and underlying health conditions. Complications can include acute kidney injury, chronic kidney disease progression, and systemic involvement leading to multi-organ dysfunction.
|
||||
|
||||
### Coding Requirements:
|
||||
|
||||
#### Relevant ICD-10 Codes:
|
||||
|
||||
- **Confirmed Viral Conditions**:
|
||||
- **B18.1**: Chronic viral hepatitis B
|
||||
- **B18.2**: Chronic viral hepatitis C
|
||||
- **B20**: HIV disease resulting in infectious and parasitic diseases
|
||||
|
||||
- **Suspected Viral Conditions**:
|
||||
- **B34.9**: Viral infection, unspecified
|
||||
|
||||
- **Associated Symptoms**:
|
||||
- **R50.9**: Fever, unspecified
|
||||
- **R53.83**: Other fatigue
|
||||
- **R05**: Cough
|
||||
|
||||
- **Complications**:
|
||||
- **N17.9**: Acute kidney failure, unspecified (if viral infection leads to AKI)
|
||||
- **N18.9**: Chronic kidney disease, unspecified (if progression due to viral infection)
|
||||
|
||||
#### Primary and Secondary Diagnostic Codes:
|
||||
|
||||
- **Primary Diagnostic Codes**:
|
||||
- Use the specific viral infection code as primary if confirmed (e.g., B18.2 for Hepatitis C).
|
||||
|
||||
- **Secondary Condition Codes**:
|
||||
- Use codes for symptoms or complications as secondary (e.g., N17.9 for AKI if due to viral infection).
|
||||
|
||||
### Rationale for Code Selection:
|
||||
|
||||
- **B18.1 and B18.2**: Selected for confirmed chronic hepatitis B or C, which can have renal complications.
|
||||
- **B20**: Used if HIV is confirmed, given its potential impact on renal function.
|
||||
- **B34.9**: Utilized when a viral infection is suspected but not yet identified.
|
||||
- **R50.9, R53.83, R05**: Common symptoms associated with viral infections.
|
||||
- **N17.9, N18.9**: Codes for renal complications potentially exacerbated by viral infections.
|
||||
|
||||
### Documentation:
|
||||
|
||||
- Ensure thorough documentation of clinical findings, suspected or confirmed viral infections, and associated symptoms or complications to justify the selected ICD-10 codes.
|
||||
- Follow-up with additional testing or specialist referrals as needed to confirm or rule out viral etiologies and manage complications effectively. Agent Name: Internist
|
||||
Output: To provide a comprehensive evaluation as an Internal Medicine specialist, let's conduct a detailed clinical assessment and medical coding for the presented case. This will involve a system-by-system review, analysis of vital signs, and evaluation of comorbidities, followed by appropriate ICD-10 coding.
|
||||
|
||||
### Clinical Assessment:
|
||||
|
||||
#### System-by-System Review:
|
||||
1. **Respiratory System:**
|
||||
- Evaluate for symptoms such as cough, shortness of breath, or wheezing.
|
||||
- Consider potential viral or bacterial infections affecting the respiratory tract.
|
||||
|
||||
2. **Cardiovascular System:**
|
||||
- Assess for any signs of heart failure or hypertension.
|
||||
- Look for symptoms like chest pain, palpitations, or edema.
|
||||
|
||||
3. **Gastrointestinal System:**
|
||||
- Check for symptoms such as nausea, vomiting, diarrhea, or abdominal pain.
|
||||
- Consider liver function if hepatitis is suspected.
|
||||
|
||||
4. **Renal System:**
|
||||
- Monitor for signs of acute kidney injury or chronic kidney disease.
|
||||
- Evaluate urine output and creatinine levels.
|
||||
|
||||
5. **Neurological System:**
|
||||
- Assess for headaches, dizziness, or any focal neurological deficits.
|
||||
- Consider viral encephalitis if neurological symptoms are present.
|
||||
|
||||
6. **Musculoskeletal System:**
|
||||
- Look for muscle aches or joint pain, common in viral infections.
|
||||
|
||||
7. **Integumentary System:**
|
||||
- Check for rashes or skin lesions, which may indicate viral infections like herpes or CMV.
|
||||
|
||||
8. **Immune System:**
|
||||
- Consider immunosuppression status, especially in the context of HIV or other chronic infections.
|
||||
|
||||
#### Vital Signs Analysis:
|
||||
- **Temperature:** Evaluate for fever, which may indicate an infection.
|
||||
- **Blood Pressure:** Check for hypertension or hypotension.
|
||||
- **Heart Rate:** Assess for tachycardia or bradycardia.
|
||||
- **Respiratory Rate:** Monitor for tachypnea.
|
||||
- **Oxygen Saturation:** Ensure adequate oxygenation, especially in respiratory infections.
|
||||
|
||||
#### Comorbidity Evaluation:
|
||||
- Assess for chronic conditions such as diabetes, hypertension, or chronic kidney disease.
|
||||
- Consider the impact of these conditions on the current clinical presentation and potential complications.
|
||||
|
||||
### Medical Coding:
|
||||
|
||||
#### ICD-10 Codes:
|
||||
|
||||
1. **Primary Conditions:**
|
||||
- If a specific viral infection is confirmed, use the appropriate code (e.g., B18.2 for chronic hepatitis C).
|
||||
|
||||
2. **Secondary Diagnoses:**
|
||||
- **B34.9:** Viral infection, unspecified (if viral etiology is suspected but not confirmed).
|
||||
- **R50.9:** Fever, unspecified (common symptom in infections).
|
||||
- **R53.83:** Other fatigue (common in viral infections).
|
||||
|
||||
3. **Complications:**
|
||||
- **N17.9:** Acute kidney failure, unspecified (if there is renal involvement).
|
||||
- **N18.9:** Chronic kidney disease, unspecified (if there is progression due to infection).
|
||||
|
||||
4. **Chronic Conditions:**
|
||||
- **I10:** Essential (primary) hypertension (if present).
|
||||
- **E11.9:** Type 2 diabetes mellitus without complications (if present).
|
||||
|
||||
5. **Signs and Symptoms:**
|
||||
- **R05:** Cough (common respiratory symptom).
|
||||
- **M79.1:** Myalgia (muscle pain).
|
||||
|
||||
#### Hierarchical Condition Category (HCC) Codes:
|
||||
- **HCC 18:** Diabetes with chronic complications (if applicable).
|
||||
- **HCC 85:** Congestive heart failure (if applicable).
|
||||
|
||||
### Documentation Supporting Evidence:
|
||||
- Ensure documentation includes detailed clinical findings, symptoms, and any laboratory or imaging results that support the diagnosis.
|
||||
- Include any history of chronic conditions or recent changes in health status.
|
||||
- Document any suspected or confirmed viral infections, along with their impact on the patient's health.
|
||||
|
||||
### Conclusion:
|
||||
This comprehensive evaluation and coding approach allows for accurate diagnosis and management of the patient's condition, considering both acute and chronic aspects of their health. Proper documentation and coding facilitate effective communication and continuity of care. Agent Name: Medical Coder
|
||||
Output: ### Medical Coding Summary
|
||||
|
||||
#### 1. Primary Diagnosis Codes
|
||||
- **ICD-10 Code:** B18.2
|
||||
- **Description:** Chronic viral hepatitis C
|
||||
- **Supporting Documentation:** The diagnosis of chronic hepatitis C is confirmed through serological testing and liver function tests indicating chronic viral infection.
|
||||
|
||||
#### 2. Secondary Diagnosis Codes
|
||||
- **B34.9:** Viral infection, unspecified
|
||||
- **Supporting Documentation:** Suspected viral etiology without specific identification.
|
||||
- **R50.9:** Fever, unspecified
|
||||
- **Supporting Documentation:** Documented fever without a definitive cause.
|
||||
- **R53.83:** Other fatigue
|
||||
- **Supporting Documentation:** Patient reports persistent fatigue, common in viral infections.
|
||||
- **I10:** Essential (primary) hypertension
|
||||
- **Supporting Documentation:** History of hypertension with current blood pressure readings.
|
||||
- **E11.9:** Type 2 diabetes mellitus without complications
|
||||
- **Supporting Documentation:** Documented history of type 2 diabetes, managed with oral hypoglycemics.
|
||||
|
||||
#### 3. Symptom Codes
|
||||
- **R05:** Cough
|
||||
- **Supporting Documentation:** Patient presents with a persistent cough, noted in the respiratory evaluation.
|
||||
- **M79.1:** Myalgia
|
||||
- **Supporting Documentation:** Patient reports muscle pain, consistent with viral infections.
|
||||
|
||||
#### 4. Complication Codes
|
||||
- **N17.9:** Acute kidney failure, unspecified
|
||||
- **Supporting Documentation:** Elevated creatinine levels and reduced urine output indicative of renal involvement.
|
||||
- **N18.9:** Chronic kidney disease, unspecified
|
||||
- **Supporting Documentation:** Documented chronic kidney disease stage, with baseline creatinine levels.
|
||||
|
||||
#### 5. Coding Notes
|
||||
- Ensure all clinical findings and laboratory results supporting the diagnoses are documented in the patient's medical record.
|
||||
- Confirm the presence of chronic conditions and their management strategies.
|
||||
- Monitor for any changes in the patient's condition that may require code updates or additions.
|
||||
- Address any coding queries related to unspecified viral infections by seeking further diagnostic clarification if possible.
|
||||
|
||||
This coding summary provides a structured approach to documenting the patient's current health status, ensuring accurate and compliant ICD-10 coding. Agent Name: Diagnostic Synthesizer
|
||||
Output: ### Final Diagnostic and Coding Assessment
|
||||
|
||||
#### Clinical Summary
|
||||
|
||||
**Primary Diagnosis:**
|
||||
- **ICD-10 Code:** B18.2
|
||||
- **Description:** Chronic viral hepatitis C
|
||||
- **Supporting Evidence:** This diagnosis is substantiated by serological testing and liver function tests indicating a chronic viral infection. The confidence level for this diagnosis is high, with high-quality evidence from laboratory results.
|
||||
|
||||
**Secondary Diagnoses:**
|
||||
1. **ICD-10 Code:** B34.9
|
||||
- **Description:** Viral infection, unspecified
|
||||
- **Supporting Evidence:** The suspected viral etiology lacks specific identification. Confidence level is moderate due to limited specificity in viral identification.
|
||||
|
||||
2. **ICD-10 Code:** R50.9
|
||||
- **Description:** Fever, unspecified
|
||||
- **Supporting Evidence:** Documented fever without a definitive cause. Confidence level is moderate, supported by clinical observation.
|
||||
|
||||
3. **ICD-10 Code:** R53.83
|
||||
- **Description:** Other fatigue
|
||||
- **Supporting Evidence:** Patient reports persistent fatigue, often associated with viral infections. Confidence level is moderate, based on patient-reported symptoms.
|
||||
|
||||
4. **ICD-10 Code:** I10
|
||||
- **Description:** Essential (primary) hypertension
|
||||
- **Supporting Evidence:** History of hypertension corroborated by current blood pressure readings. Confidence level is high, with consistent clinical evidence.
|
||||
|
||||
5. **ICD-10 Code:** E11.9
|
||||
- **Description:** Type 2 diabetes mellitus without complications
|
||||
- **Supporting Evidence:** Managed with oral hypoglycemics, with a documented history. Confidence level is high, with strong management records.
|
||||
|
||||
**Symptom Codes:**
|
||||
- **ICD-10 Code:** R05
|
||||
- **Description:** Cough
|
||||
- **Supporting Evidence:** Persistent cough noted in respiratory evaluation. Confidence level is moderate, based on clinical observation.
|
||||
|
||||
- **ICD-10 Code:** M79.1
|
||||
- **Description:** Myalgia
|
||||
- **Supporting Evidence:** Muscle pain reported by the patient, consistent with viral infections. Confidence level is moderate, based on patient-reported symptoms.
|
||||
|
||||
**Complication Codes:**
|
||||
1. **ICD-10 Code:** N17.9
|
||||
- **Description:** Acute kidney failure, unspecified
|
||||
- **Supporting Evidence:** Elevated creatinine levels and reduced urine output suggest renal involvement. Confidence level is high, supported by laboratory data.
|
||||
|
||||
2. **ICD-10 Code:** N18.9
|
||||
- **Description:** Chronic kidney disease, unspecified
|
||||
- **Supporting Evidence:** Documented chronic kidney disease stage with baseline creatinine levels. Confidence level is high, with consistent clinical data.
|
||||
|
||||
#### Coding Summary
|
||||
|
||||
**Complete Code List with Descriptions:**
|
||||
- B18.2: Chronic viral hepatitis C
|
||||
- B34.9: Viral infection, unspecified
|
||||
- R50.9: Fever, unspecified
|
||||
- R53.83: Other fatigue
|
||||
- I10: Essential (primary) hypertension
|
||||
- E11.9: Type 2 diabetes mellitus without complications
|
||||
- R05: Cough
|
||||
- M79.1: Myalgia
|
||||
- N17.9: Acute kidney failure, unspecified
|
||||
- N18.9: Chronic kidney disease, unspecified
|
||||
|
||||
**Code Hierarchy and Relationships:**
|
||||
- The primary diagnosis (B18.2) is the central focus, with secondary diagnoses and symptoms potentially related to or exacerbated by the chronic hepatitis C infection.
|
||||
- Complications (N17.9 and N18.9) may be linked to the primary diagnosis and other chronic conditions like diabetes and hypertension.
|
||||
|
||||
**Supporting Documentation:**
|
||||
- Ensure that all clinical findings and laboratory results supporting the diagnoses are documented in the patient's medical record.
|
||||
- Confirm the presence of chronic conditions and their management strategies.
|
||||
- Monitor for any changes in the patient's condition that may require code updates or additions.
|
||||
|
||||
#### Recommendations
|
||||
|
||||
1. **Additional Testing Needed:**
|
||||
- Further diagnostic testing is recommended to clarify the unspecified viral infection (B34.9) and to monitor kidney function.
|
||||
|
||||
2. **Follow-up Care:**
|
||||
- Regular follow-up appointments to manage chronic conditions such as hypertension and diabetes.
|
||||
- Monitor renal function and adjust treatment plans as necessary.
|
||||
|
||||
3. **Documentation Improvements Needed:**
|
||||
- Enhance documentation specificity for the unspecified viral infection.
|
||||
- Ensure comprehensive records of all chronic conditions and their management strategies.
|
||||
|
||||
These recommendations aim to improve diagnostic accuracy and patient care continuity.
|
||||
@ -0,0 +1,177 @@
|
||||
from datetime import datetime
|
||||
|
||||
from swarms import Agent, AgentRearrange, create_file_in_folder
|
||||
|
||||
chief_medical_officer = Agent(
|
||||
agent_name="Chief Medical Officer",
|
||||
system_prompt="""You are the Chief Medical Officer coordinating a team of medical specialists for viral disease diagnosis.
|
||||
Your responsibilities include:
|
||||
- Gathering initial patient symptoms and medical history
|
||||
- Coordinating with specialists to form differential diagnoses
|
||||
- Synthesizing different specialist opinions into a cohesive diagnosis
|
||||
- Ensuring all relevant symptoms and test results are considered
|
||||
- Making final diagnostic recommendations
|
||||
- Suggesting treatment plans based on team input
|
||||
- Identifying when additional specialists need to be consulted
|
||||
|
||||
Guidelines:
|
||||
1. Always start with a comprehensive patient history
|
||||
2. Consider both common and rare viral conditions
|
||||
3. Factor in patient demographics and risk factors
|
||||
4. Document your reasoning process clearly
|
||||
5. Highlight any critical or emergency symptoms
|
||||
6. Note any limitations or uncertainties in the diagnosis
|
||||
|
||||
Format all responses with clear sections for:
|
||||
- Initial Assessment
|
||||
- Differential Diagnoses
|
||||
- Specialist Consultations Needed
|
||||
- Recommended Next Steps""",
|
||||
model_name="gpt-4o", # Models from litellm -> claude-2
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Viral Disease Specialist
|
||||
virologist = Agent(
|
||||
agent_name="Virologist",
|
||||
system_prompt="""You are a specialist in viral diseases with expertise in:
|
||||
- Respiratory viruses (Influenza, Coronavirus, RSV)
|
||||
- Systemic viral infections (EBV, CMV, HIV)
|
||||
- Childhood viral diseases (Measles, Mumps, Rubella)
|
||||
- Emerging viral threats
|
||||
|
||||
Your role involves:
|
||||
1. Analyzing symptoms specific to viral infections
|
||||
2. Distinguishing between different viral pathogens
|
||||
3. Assessing viral infection patterns and progression
|
||||
4. Recommending specific viral tests
|
||||
5. Evaluating epidemiological factors
|
||||
|
||||
For each case, consider:
|
||||
- Incubation periods
|
||||
- Transmission patterns
|
||||
- Seasonal factors
|
||||
- Geographic prevalence
|
||||
- Patient immune status
|
||||
- Current viral outbreaks
|
||||
|
||||
Provide detailed analysis of:
|
||||
- Characteristic viral symptoms
|
||||
- Disease progression timeline
|
||||
- Risk factors for severe disease
|
||||
- Potential complications""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Internal Medicine Specialist
|
||||
internist = Agent(
|
||||
agent_name="Internist",
|
||||
system_prompt="""You are an Internal Medicine specialist responsible for:
|
||||
- Comprehensive system-based evaluation
|
||||
- Integration of symptoms across organ systems
|
||||
- Identification of systemic manifestations
|
||||
- Assessment of comorbidities
|
||||
|
||||
For each case, analyze:
|
||||
1. Vital signs and their implications
|
||||
2. System-by-system review (cardiovascular, respiratory, etc.)
|
||||
3. Impact of existing medical conditions
|
||||
4. Medication interactions and contraindications
|
||||
5. Risk stratification
|
||||
|
||||
Consider these aspects:
|
||||
- Age-related factors
|
||||
- Chronic disease impact
|
||||
- Medication history
|
||||
- Social and environmental factors
|
||||
|
||||
Document:
|
||||
- Physical examination findings
|
||||
- System-specific symptoms
|
||||
- Relevant lab abnormalities
|
||||
- Risk factors for complications""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Diagnostic Synthesizer
|
||||
synthesizer = Agent(
|
||||
agent_name="Diagnostic Synthesizer",
|
||||
system_prompt="""You are responsible for synthesizing all specialist inputs to create a final diagnostic assessment:
|
||||
|
||||
Core responsibilities:
|
||||
1. Integrate findings from all specialists
|
||||
2. Identify patterns and correlations
|
||||
3. Resolve conflicting opinions
|
||||
4. Generate probability-ranked differential diagnoses
|
||||
5. Recommend additional testing if needed
|
||||
|
||||
Analysis framework:
|
||||
- Weight evidence based on reliability and specificity
|
||||
- Consider epidemiological factors
|
||||
- Evaluate diagnostic certainty
|
||||
- Account for test limitations
|
||||
|
||||
Provide structured output including:
|
||||
1. Primary diagnosis with confidence level
|
||||
2. Supporting evidence summary
|
||||
3. Alternative diagnoses to consider
|
||||
4. Recommended confirmatory tests
|
||||
5. Red flags or warning signs
|
||||
6. Follow-up recommendations
|
||||
|
||||
Documentation requirements:
|
||||
- Clear reasoning chain
|
||||
- Evidence quality assessment
|
||||
- Confidence levels for each diagnosis
|
||||
- Knowledge gaps identified
|
||||
- Risk assessment""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create agent list
|
||||
agents = [chief_medical_officer, virologist, internist, synthesizer]
|
||||
|
||||
# Define diagnostic flow
|
||||
flow = f"""{chief_medical_officer.agent_name} -> {virologist.agent_name} -> {internist.agent_name} -> {synthesizer.agent_name}"""
|
||||
|
||||
# Create the swarm system
|
||||
diagnosis_system = AgentRearrange(
|
||||
name="Medical-nlp-diagnosis-swarm",
|
||||
description="natural language symptions to diagnosis report",
|
||||
agents=agents,
|
||||
flow=flow,
|
||||
max_loops=1,
|
||||
output_type="all",
|
||||
)
|
||||
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
# Example patient case
|
||||
patient_case = """
|
||||
Patient: 45-year-old female
|
||||
Presenting symptoms:
|
||||
- Fever (101.5°F) for 3 days
|
||||
- Dry cough
|
||||
- Fatigue
|
||||
- Mild shortness of breath
|
||||
Medical history:
|
||||
- Controlled hypertension
|
||||
- No recent travel
|
||||
- Fully vaccinated for COVID-19
|
||||
- No known sick contacts
|
||||
"""
|
||||
|
||||
# Add timestamp to the patient case
|
||||
case_info = f"Timestamp: {datetime.now()}\nPatient Information: {patient_case}"
|
||||
|
||||
# Run the diagnostic process
|
||||
diagnosis = diagnosis_system.run(case_info)
|
||||
|
||||
# Create a folder and file called reports
|
||||
create_file_in_folder(
|
||||
"reports", "medical_analysis_agent_rearrange.md", diagnosis
|
||||
)
|
||||
@ -0,0 +1,243 @@
|
||||
from datetime import datetime
|
||||
from swarms import Agent, AgentRearrange, create_file_in_folder
|
||||
|
||||
# Lead Investment Analyst
|
||||
lead_analyst = Agent(
|
||||
agent_name="Lead Investment Analyst",
|
||||
system_prompt="""You are the Lead Investment Analyst coordinating document analysis for venture capital investments.
|
||||
|
||||
Core responsibilities:
|
||||
- Coordinating overall document review process
|
||||
- Identifying key terms and conditions
|
||||
- Flagging potential risks and concerns
|
||||
- Synthesizing specialist inputs into actionable insights
|
||||
- Recommending negotiation points
|
||||
|
||||
Document Analysis Framework:
|
||||
1. Initial document classification and overview
|
||||
2. Key terms identification
|
||||
3. Risk assessment
|
||||
4. Market context evaluation
|
||||
5. Recommendation formulation
|
||||
|
||||
Output Format Requirements:
|
||||
- Executive Summary
|
||||
- Key Terms Analysis
|
||||
- Risk Factors
|
||||
- Negotiation Points
|
||||
- Recommended Actions
|
||||
- Areas Requiring Specialist Review""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# SAFE Agreement Specialist
|
||||
safe_specialist = Agent(
|
||||
agent_name="SAFE Specialist",
|
||||
system_prompt="""You are a specialist in SAFE (Simple Agreement for Future Equity) agreements with expertise in:
|
||||
|
||||
Technical Analysis Areas:
|
||||
- Valuation caps and discount rates
|
||||
- Conversion mechanisms and triggers
|
||||
- Pro-rata rights
|
||||
- Most Favored Nation (MFN) provisions
|
||||
- Dilution and anti-dilution provisions
|
||||
|
||||
Required Assessments:
|
||||
1. Cap table impact analysis
|
||||
2. Conversion scenarios modeling
|
||||
3. Rights and preferences evaluation
|
||||
4. Standard vs. non-standard terms identification
|
||||
5. Post-money vs. pre-money SAFE analysis
|
||||
|
||||
Consider and Document:
|
||||
- Valuation implications
|
||||
- Future round impacts
|
||||
- Investor rights and limitations
|
||||
- Comparative market terms
|
||||
- Potential conflicts with other securities
|
||||
|
||||
Output Requirements:
|
||||
- Term-by-term analysis
|
||||
- Conversion mechanics explanation
|
||||
- Risk assessment for non-standard terms
|
||||
- Recommendations for negotiations""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Term Sheet Analyst
|
||||
term_sheet_analyst = Agent(
|
||||
agent_name="Term Sheet Analyst",
|
||||
system_prompt="""You are a Term Sheet Analyst specialized in venture capital financing documents.
|
||||
|
||||
Core Analysis Areas:
|
||||
- Economic terms (valuation, option pool, etc.)
|
||||
- Control provisions
|
||||
- Investor rights and protections
|
||||
- Governance structures
|
||||
- Exit and liquidity provisions
|
||||
|
||||
Detailed Review Requirements:
|
||||
1. Economic Terms Analysis:
|
||||
- Pre/post-money valuation
|
||||
- Share price calculation
|
||||
- Capitalization analysis
|
||||
- Option pool sizing
|
||||
|
||||
2. Control Provisions Review:
|
||||
- Board composition
|
||||
- Voting rights
|
||||
- Protective provisions
|
||||
- Information rights
|
||||
|
||||
3. Investor Rights Assessment:
|
||||
- Pro-rata rights
|
||||
- Anti-dilution protection
|
||||
- Registration rights
|
||||
- Right of first refusal
|
||||
|
||||
Output Format:
|
||||
- Term-by-term breakdown
|
||||
- Market standard comparison
|
||||
- Founder impact analysis
|
||||
- Investor rights summary
|
||||
- Governance implications""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Legal Compliance Analyst
|
||||
legal_analyst = Agent(
|
||||
agent_name="Legal Compliance Analyst",
|
||||
system_prompt="""You are a Legal Compliance Analyst for venture capital documentation.
|
||||
|
||||
Primary Focus Areas:
|
||||
- Securities law compliance
|
||||
- Corporate governance requirements
|
||||
- Regulatory restrictions
|
||||
- Standard market practices
|
||||
- Legal risk assessment
|
||||
|
||||
Analysis Framework:
|
||||
1. Regulatory Compliance:
|
||||
- Securities regulations
|
||||
- Disclosure requirements
|
||||
- Investment company considerations
|
||||
- Blue sky laws
|
||||
|
||||
2. Documentation Review:
|
||||
- Legal definitions accuracy
|
||||
- Enforceability concerns
|
||||
- Jurisdiction issues
|
||||
- Amendment provisions
|
||||
|
||||
3. Risk Assessment:
|
||||
- Legal precedent analysis
|
||||
- Regulatory exposure
|
||||
- Enforcement mechanisms
|
||||
- Dispute resolution provisions
|
||||
|
||||
Output Requirements:
|
||||
- Compliance checklist
|
||||
- Risk assessment summary
|
||||
- Required disclosures list
|
||||
- Recommended legal modifications
|
||||
- Jurisdiction-specific concerns""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Market Comparison Analyst
|
||||
market_analyst = Agent(
|
||||
agent_name="Market Comparison Analyst",
|
||||
system_prompt="""You are a Market Comparison Analyst for venture capital terms and conditions.
|
||||
|
||||
Core Responsibilities:
|
||||
- Benchmark terms against market standards
|
||||
- Identify industry-specific variations
|
||||
- Track emerging market trends
|
||||
- Assess term competitiveness
|
||||
|
||||
Analysis Framework:
|
||||
1. Market Comparison:
|
||||
- Stage-appropriate terms
|
||||
- Industry-standard provisions
|
||||
- Geographic variations
|
||||
- Recent trend analysis
|
||||
|
||||
2. Competitive Assessment:
|
||||
- Investor-friendliness rating
|
||||
- Founder-friendliness rating
|
||||
- Term flexibility analysis
|
||||
- Market positioning
|
||||
|
||||
3. Trend Analysis:
|
||||
- Emerging terms and conditions
|
||||
- Shifting market standards
|
||||
- Industry-specific adaptations
|
||||
- Regional variations
|
||||
|
||||
Output Format:
|
||||
- Market positioning summary
|
||||
- Comparative analysis
|
||||
- Trend implications
|
||||
- Negotiation leverage points
|
||||
- Recommended modifications""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create agent list
|
||||
agents = [
|
||||
lead_analyst,
|
||||
safe_specialist,
|
||||
term_sheet_analyst,
|
||||
legal_analyst,
|
||||
market_analyst,
|
||||
]
|
||||
|
||||
# Define analysis flow
|
||||
flow = f"""{lead_analyst.agent_name} -> {safe_specialist.agent_name} -> {term_sheet_analyst.agent_name} -> {legal_analyst.agent_name} -> {market_analyst.agent_name}"""
|
||||
|
||||
# Create the swarm system
|
||||
vc_analysis_system = AgentRearrange(
|
||||
name="VC-Document-Analysis-Swarm",
|
||||
description="SAFE and Term Sheet document analysis and Q&A system",
|
||||
agents=agents,
|
||||
flow=flow,
|
||||
max_loops=1,
|
||||
output_type="all",
|
||||
)
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
# Example document for analysis
|
||||
document_text = """
|
||||
SAFE AGREEMENT
|
||||
|
||||
Valuation Cap: $10,000,000
|
||||
Discount Rate: 20%
|
||||
|
||||
Investment Amount: $500,000
|
||||
|
||||
Conversion Provisions:
|
||||
- Automatic conversion upon Equity Financing of at least $1,000,000
|
||||
- Optional conversion upon Liquidity Event
|
||||
- Most Favored Nation provision included
|
||||
|
||||
Pro-rata Rights: Included for future rounds
|
||||
"""
|
||||
|
||||
# Add timestamp to the analysis
|
||||
analysis_request = f"Timestamp: {datetime.now()}\nDocument for Analysis: {document_text}"
|
||||
|
||||
# Run the analysis
|
||||
analysis = vc_analysis_system.run(analysis_request)
|
||||
|
||||
# Create analysis report
|
||||
create_file_in_folder(
|
||||
"reports", "vc_document_analysis.md", analysis
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
@ -0,0 +1,252 @@
|
||||
"""
|
||||
- For each diagnosis, pull lab results,
|
||||
- egfr
|
||||
- for each diagnosis, pull lab ranges,
|
||||
- pull ranges for diagnosis
|
||||
|
||||
- if the diagnosis is x, then the lab ranges should be a to b
|
||||
- train the agents, increase the load of input
|
||||
- medical history sent to the agent
|
||||
- setup rag for the agents
|
||||
- run the first agent -> kidney disease -> don't know the stage -> stage 2 -> lab results -> indicative of stage 3 -> the case got elavated ->
|
||||
- how to manage diseases and by looking at correlating lab, docs, diagnoses
|
||||
- put docs in rag ->
|
||||
- monitoring, evaluation, and treatment
|
||||
- can we confirm for every diagnosis -> monitoring, evaluation, and treatment, specialized for these things
|
||||
- find diagnosis -> or have diagnosis, -> for each diagnosis are there evidence of those 3 things
|
||||
- swarm of those 4 agents, ->
|
||||
- fda api for healthcare for commerically available papers
|
||||
-
|
||||
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from swarms import Agent, AgentRearrange, create_file_in_folder
|
||||
|
||||
from swarm_models import OllamaModel
|
||||
|
||||
model = OllamaModel(model_name="llama3.2")
|
||||
|
||||
chief_medical_officer = Agent(
|
||||
agent_name="Chief Medical Officer",
|
||||
system_prompt="""You are the Chief Medical Officer coordinating a team of medical specialists for viral disease diagnosis.
|
||||
Your responsibilities include:
|
||||
- Gathering initial patient symptoms and medical history
|
||||
- Coordinating with specialists to form differential diagnoses
|
||||
- Synthesizing different specialist opinions into a cohesive diagnosis
|
||||
- Ensuring all relevant symptoms and test results are considered
|
||||
- Making final diagnostic recommendations
|
||||
- Suggesting treatment plans based on team input
|
||||
- Identifying when additional specialists need to be consulted
|
||||
- For each diferrential diagnosis provide minimum lab ranges to meet that diagnosis or be indicative of that diagnosis minimum and maximum
|
||||
|
||||
Format all responses with clear sections for:
|
||||
- Initial Assessment (include preliminary ICD-10 codes for symptoms)
|
||||
- Differential Diagnoses (with corresponding ICD-10 codes)
|
||||
- Specialist Consultations Needed
|
||||
- Recommended Next Steps
|
||||
|
||||
|
||||
""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
virologist = Agent(
|
||||
agent_name="Virologist",
|
||||
system_prompt="""You are a specialist in viral diseases. For each case, provide:
|
||||
|
||||
Clinical Analysis:
|
||||
- Detailed viral symptom analysis
|
||||
- Disease progression timeline
|
||||
- Risk factors and complications
|
||||
|
||||
Coding Requirements:
|
||||
- List relevant ICD-10 codes for:
|
||||
* Confirmed viral conditions
|
||||
* Suspected viral conditions
|
||||
* Associated symptoms
|
||||
* Complications
|
||||
- Include both:
|
||||
* Primary diagnostic codes
|
||||
* Secondary condition codes
|
||||
|
||||
Document all findings using proper medical coding standards and include rationale for code selection.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
internist = Agent(
|
||||
agent_name="Internist",
|
||||
system_prompt="""You are an Internal Medicine specialist responsible for comprehensive evaluation.
|
||||
|
||||
For each case, provide:
|
||||
|
||||
Clinical Assessment:
|
||||
- System-by-system review
|
||||
- Vital signs analysis
|
||||
- Comorbidity evaluation
|
||||
|
||||
Medical Coding:
|
||||
- ICD-10 codes for:
|
||||
* Primary conditions
|
||||
* Secondary diagnoses
|
||||
* Complications
|
||||
* Chronic conditions
|
||||
* Signs and symptoms
|
||||
- Include hierarchical condition category (HCC) codes where applicable
|
||||
|
||||
Document supporting evidence for each code selected.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
medical_coder = Agent(
|
||||
agent_name="Medical Coder",
|
||||
system_prompt="""You are a certified medical coder responsible for:
|
||||
|
||||
Primary Tasks:
|
||||
1. Reviewing all clinical documentation
|
||||
2. Assigning accurate ICD-10 codes
|
||||
3. Ensuring coding compliance
|
||||
4. Documenting code justification
|
||||
|
||||
Coding Process:
|
||||
- Review all specialist inputs
|
||||
- Identify primary and secondary diagnoses
|
||||
- Assign appropriate ICD-10 codes
|
||||
- Document supporting evidence
|
||||
- Note any coding queries
|
||||
|
||||
Output Format:
|
||||
1. Primary Diagnosis Codes
|
||||
- ICD-10 code
|
||||
- Description
|
||||
- Supporting documentation
|
||||
2. Secondary Diagnosis Codes
|
||||
- Listed in order of clinical significance
|
||||
3. Symptom Codes
|
||||
4. Complication Codes
|
||||
5. Coding Notes""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
synthesizer = Agent(
|
||||
agent_name="Diagnostic Synthesizer",
|
||||
system_prompt="""You are responsible for creating the final diagnostic and coding assessment.
|
||||
|
||||
Synthesis Requirements:
|
||||
1. Integrate all specialist findings
|
||||
2. Reconcile any conflicting diagnoses
|
||||
3. Verify coding accuracy and completeness
|
||||
|
||||
Final Report Sections:
|
||||
1. Clinical Summary
|
||||
- Primary diagnosis with ICD-10
|
||||
- Secondary diagnoses with ICD-10
|
||||
- Supporting evidence
|
||||
2. Coding Summary
|
||||
- Complete code list with descriptions
|
||||
- Code hierarchy and relationships
|
||||
- Supporting documentation
|
||||
3. Recommendations
|
||||
- Additional testing needed
|
||||
- Follow-up care
|
||||
- Documentation improvements needed
|
||||
|
||||
Include confidence levels and evidence quality for all diagnoses and codes.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create agent list
|
||||
agents = [
|
||||
chief_medical_officer,
|
||||
virologist,
|
||||
internist,
|
||||
medical_coder,
|
||||
synthesizer,
|
||||
]
|
||||
|
||||
# Define diagnostic flow
|
||||
flow = f"""{chief_medical_officer.agent_name} -> {virologist.agent_name} -> {internist.agent_name} -> {medical_coder.agent_name} -> {synthesizer.agent_name}"""
|
||||
|
||||
# Create the swarm system
|
||||
diagnosis_system = AgentRearrange(
|
||||
name="Medical-coding-diagnosis-swarm",
|
||||
description="Comprehensive medical diagnosis and coding system",
|
||||
agents=agents,
|
||||
flow=flow,
|
||||
max_loops=1,
|
||||
output_type="all",
|
||||
)
|
||||
|
||||
|
||||
def generate_coding_report(diagnosis_output: str) -> str:
|
||||
"""
|
||||
Generate a structured medical coding report from the diagnosis output.
|
||||
"""
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
report = f"""# Medical Diagnosis and Coding Report
|
||||
Generated: {timestamp}
|
||||
|
||||
## Clinical Summary
|
||||
{diagnosis_output}
|
||||
|
||||
## Coding Summary
|
||||
### Primary Diagnosis Codes
|
||||
[Extracted from synthesis]
|
||||
|
||||
### Secondary Diagnosis Codes
|
||||
[Extracted from synthesis]
|
||||
|
||||
### Symptom Codes
|
||||
[Extracted from synthesis]
|
||||
|
||||
### Procedure Codes (if applicable)
|
||||
[Extracted from synthesis]
|
||||
|
||||
## Documentation and Compliance Notes
|
||||
- Code justification
|
||||
- Supporting documentation references
|
||||
- Any coding queries or clarifications needed
|
||||
|
||||
## Recommendations
|
||||
- Additional documentation needed
|
||||
- Suggested follow-up
|
||||
- Coding optimization opportunities
|
||||
"""
|
||||
return report
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example patient case
|
||||
patient_case = """
|
||||
Patient: 45-year-old White Male
|
||||
|
||||
Lab Results:
|
||||
- egfr
|
||||
- 59 ml / min / 1.73
|
||||
- non african-american
|
||||
|
||||
"""
|
||||
|
||||
# Add timestamp to the patient case
|
||||
case_info = f"Timestamp: {datetime.now()}\nPatient Information: {patient_case}"
|
||||
|
||||
# Run the diagnostic process
|
||||
diagnosis = diagnosis_system.run(case_info)
|
||||
|
||||
# Generate coding report
|
||||
coding_report = generate_coding_report(diagnosis)
|
||||
|
||||
# Create reports
|
||||
create_file_in_folder(
|
||||
"reports", "medical_diagnosis_report.md", diagnosis
|
||||
)
|
||||
create_file_in_folder(
|
||||
"reports", "medical_coding_report.md", coding_report
|
||||
)
|
||||
@ -0,0 +1,354 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
import asyncio
|
||||
from loguru import logger
|
||||
import json
|
||||
import base58
|
||||
from decimal import Decimal
|
||||
|
||||
# Swarms imports
|
||||
from swarms import Agent
|
||||
|
||||
# Solana imports
|
||||
from solders.rpc.responses import GetTransactionResp
|
||||
from solders.transaction import Transaction
|
||||
from anchorpy import Provider, Wallet
|
||||
from solders.keypair import Keypair
|
||||
import aiohttp
|
||||
|
||||
# Specialized Solana Analysis System Prompt
|
||||
SOLANA_ANALYSIS_PROMPT = """You are a specialized Solana blockchain analyst agent. Your role is to:
|
||||
|
||||
1. Analyze real-time Solana transactions for patterns and anomalies
|
||||
2. Identify potential market-moving transactions and whale movements
|
||||
3. Detect important DeFi interactions across major protocols
|
||||
4. Monitor program interactions for suspicious or notable activity
|
||||
5. Track token movements across significant protocols like:
|
||||
- Serum DEX
|
||||
- Raydium
|
||||
- Orca
|
||||
- Marinade
|
||||
- Jupiter
|
||||
- Other major Solana protocols
|
||||
|
||||
When analyzing transactions, consider:
|
||||
- Transaction size relative to protocol norms
|
||||
- Historical patterns for involved addresses
|
||||
- Impact on protocol liquidity
|
||||
- Relationship to known market events
|
||||
- Potential wash trading or suspicious patterns
|
||||
- MEV opportunities and arbitrage patterns
|
||||
- Program interaction sequences
|
||||
|
||||
Provide analysis in the following format:
|
||||
{
|
||||
"analysis_type": "[whale_movement|program_interaction|defi_trade|suspicious_activity]",
|
||||
"severity": "[high|medium|low]",
|
||||
"details": {
|
||||
"transaction_context": "...",
|
||||
"market_impact": "...",
|
||||
"recommended_actions": "...",
|
||||
"related_patterns": "..."
|
||||
}
|
||||
}
|
||||
|
||||
Focus on actionable insights that could affect:
|
||||
1. Market movements
|
||||
2. Protocol stability
|
||||
3. Trading opportunities
|
||||
4. Risk management
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransactionData:
|
||||
"""Data structure for parsed Solana transaction information"""
|
||||
|
||||
signature: str
|
||||
block_time: datetime
|
||||
slot: int
|
||||
fee: int
|
||||
lamports: int
|
||||
from_address: str
|
||||
to_address: str
|
||||
program_id: str
|
||||
instruction_data: Optional[str] = None
|
||||
program_logs: List[str] = None
|
||||
|
||||
@property
|
||||
def sol_amount(self) -> Decimal:
|
||||
"""Convert lamports to SOL"""
|
||||
return Decimal(self.lamports) / Decimal(1e9)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert transaction data to dictionary for agent analysis"""
|
||||
return {
|
||||
"signature": self.signature,
|
||||
"timestamp": self.block_time.isoformat(),
|
||||
"slot": self.slot,
|
||||
"fee": self.fee,
|
||||
"amount_sol": str(self.sol_amount),
|
||||
"from_address": self.from_address,
|
||||
"to_address": self.to_address,
|
||||
"program_id": self.program_id,
|
||||
"instruction_data": self.instruction_data,
|
||||
"program_logs": self.program_logs,
|
||||
}
|
||||
|
||||
|
||||
class SolanaSwarmAgent:
|
||||
"""Intelligent agent for analyzing Solana transactions using swarms"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent_name: str = "Solana-Analysis-Agent",
|
||||
model_name: str = "gpt-4",
|
||||
):
|
||||
self.agent = Agent(
|
||||
agent_name=agent_name,
|
||||
system_prompt=SOLANA_ANALYSIS_PROMPT,
|
||||
model_name=model_name,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
dashboard=False,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="solana_agent.json",
|
||||
user_name="solana_analyzer",
|
||||
retry_attempts=3,
|
||||
context_length=4000,
|
||||
)
|
||||
|
||||
# Initialize known patterns database
|
||||
self.known_patterns = {
|
||||
"whale_addresses": set(),
|
||||
"program_interactions": {},
|
||||
"recent_transactions": [],
|
||||
}
|
||||
logger.info(
|
||||
f"Initialized {agent_name} with specialized Solana analysis capabilities"
|
||||
)
|
||||
|
||||
async def analyze_transaction(
|
||||
self, tx_data: TransactionData
|
||||
) -> Dict[str, Any]:
|
||||
"""Analyze a transaction using the specialized agent"""
|
||||
try:
|
||||
# Update recent transactions for pattern analysis
|
||||
self.known_patterns["recent_transactions"].append(
|
||||
tx_data.signature
|
||||
)
|
||||
if len(self.known_patterns["recent_transactions"]) > 1000:
|
||||
self.known_patterns["recent_transactions"].pop(0)
|
||||
|
||||
# Prepare context for agent
|
||||
context = {
|
||||
"transaction": tx_data.to_dict(),
|
||||
"known_patterns": {
|
||||
"recent_similar_transactions": [
|
||||
tx
|
||||
for tx in self.known_patterns[
|
||||
"recent_transactions"
|
||||
][-5:]
|
||||
if abs(
|
||||
TransactionData(tx).sol_amount
|
||||
- tx_data.sol_amount
|
||||
)
|
||||
< 1
|
||||
],
|
||||
"program_statistics": self.known_patterns[
|
||||
"program_interactions"
|
||||
].get(tx_data.program_id, {}),
|
||||
},
|
||||
}
|
||||
|
||||
# Get analysis from agent
|
||||
analysis = await self.agent.run_async(
|
||||
f"Analyze the following Solana transaction and provide insights: {json.dumps(context, indent=2)}"
|
||||
)
|
||||
|
||||
# Update pattern database
|
||||
if tx_data.sol_amount > 1000: # Track whale addresses
|
||||
self.known_patterns["whale_addresses"].add(
|
||||
tx_data.from_address
|
||||
)
|
||||
|
||||
# Update program interaction statistics
|
||||
if (
|
||||
tx_data.program_id
|
||||
not in self.known_patterns["program_interactions"]
|
||||
):
|
||||
self.known_patterns["program_interactions"][
|
||||
tx_data.program_id
|
||||
] = {"total_interactions": 0, "total_volume": 0}
|
||||
self.known_patterns["program_interactions"][
|
||||
tx_data.program_id
|
||||
]["total_interactions"] += 1
|
||||
self.known_patterns["program_interactions"][
|
||||
tx_data.program_id
|
||||
]["total_volume"] += float(tx_data.sol_amount)
|
||||
|
||||
return json.loads(analysis)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in agent analysis: {str(e)}")
|
||||
return {
|
||||
"analysis_type": "error",
|
||||
"severity": "low",
|
||||
"details": {
|
||||
"error": str(e),
|
||||
"transaction": tx_data.signature,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class SolanaTransactionMonitor:
|
||||
"""Main class for monitoring and analyzing Solana transactions"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rpc_url: str,
|
||||
swarm_agent: SolanaSwarmAgent,
|
||||
min_sol_threshold: Decimal = Decimal("100"),
|
||||
):
|
||||
self.rpc_url = rpc_url
|
||||
self.swarm_agent = swarm_agent
|
||||
self.min_sol_threshold = min_sol_threshold
|
||||
self.wallet = Wallet(Keypair())
|
||||
self.provider = Provider(rpc_url, self.wallet)
|
||||
logger.info("Initialized Solana transaction monitor")
|
||||
|
||||
async def parse_transaction(
|
||||
self, tx_resp: GetTransactionResp
|
||||
) -> Optional[TransactionData]:
|
||||
"""Parse transaction response into TransactionData object"""
|
||||
try:
|
||||
if not tx_resp.value:
|
||||
return None
|
||||
|
||||
tx_value = tx_resp.value
|
||||
meta = tx_value.transaction.meta
|
||||
if not meta:
|
||||
return None
|
||||
|
||||
tx: Transaction = tx_value.transaction.transaction
|
||||
|
||||
# Extract transaction details
|
||||
from_pubkey = str(tx.message.account_keys[0])
|
||||
to_pubkey = str(tx.message.account_keys[1])
|
||||
program_id = str(tx.message.account_keys[-1])
|
||||
|
||||
# Calculate amount from balance changes
|
||||
amount = abs(meta.post_balances[0] - meta.pre_balances[0])
|
||||
|
||||
return TransactionData(
|
||||
signature=str(tx_value.transaction.signatures[0]),
|
||||
block_time=datetime.fromtimestamp(
|
||||
tx_value.block_time or 0
|
||||
),
|
||||
slot=tx_value.slot,
|
||||
fee=meta.fee,
|
||||
lamports=amount,
|
||||
from_address=from_pubkey,
|
||||
to_address=to_pubkey,
|
||||
program_id=program_id,
|
||||
program_logs=(
|
||||
meta.log_messages if meta.log_messages else []
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse transaction: {str(e)}")
|
||||
return None
|
||||
|
||||
async def start_monitoring(self):
|
||||
"""Start monitoring for new transactions"""
|
||||
logger.info(
|
||||
"Starting transaction monitoring with swarm agent analysis"
|
||||
)
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.ws_connect(self.rpc_url) as ws:
|
||||
await ws.send_json(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "transactionSubscribe",
|
||||
"params": [
|
||||
{"commitment": "finalized"},
|
||||
{
|
||||
"encoding": "jsonParsed",
|
||||
"commitment": "finalized",
|
||||
},
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
data = json.loads(msg.data)
|
||||
if "params" in data:
|
||||
signature = data["params"]["result"][
|
||||
"value"
|
||||
]["signature"]
|
||||
|
||||
# Fetch full transaction data
|
||||
tx_response = await self.provider.connection.get_transaction(
|
||||
base58.b58decode(signature)
|
||||
)
|
||||
|
||||
if tx_response:
|
||||
tx_data = (
|
||||
await self.parse_transaction(
|
||||
tx_response
|
||||
)
|
||||
)
|
||||
if (
|
||||
tx_data
|
||||
and tx_data.sol_amount
|
||||
>= self.min_sol_threshold
|
||||
):
|
||||
# Get agent analysis
|
||||
analysis = await self.swarm_agent.analyze_transaction(
|
||||
tx_data
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Transaction Analysis:\n"
|
||||
f"Signature: {tx_data.signature}\n"
|
||||
f"Amount: {tx_data.sol_amount} SOL\n"
|
||||
f"Analysis: {json.dumps(analysis, indent=2)}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error processing message: {str(e)}"
|
||||
)
|
||||
continue
|
||||
|
||||
|
||||
async def main():
|
||||
"""Example usage"""
|
||||
|
||||
# Start monitoring
|
||||
try:
|
||||
# Initialize swarm agent
|
||||
swarm_agent = SolanaSwarmAgent(
|
||||
agent_name="Solana-Whale-Detector", model_name="gpt-4"
|
||||
)
|
||||
|
||||
# Initialize monitor
|
||||
monitor = SolanaTransactionMonitor(
|
||||
rpc_url="wss://api.mainnet-beta.solana.com",
|
||||
swarm_agent=swarm_agent,
|
||||
min_sol_threshold=Decimal("100"),
|
||||
)
|
||||
|
||||
await monitor.start_monitoring()
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Shutting down gracefully...")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@ -0,0 +1,164 @@
|
||||
from swarms import Agent, SwarmRouter
|
||||
|
||||
# Portfolio Analysis Specialist
|
||||
portfolio_analyzer = Agent(
|
||||
agent_name="Portfolio-Analysis-Specialist",
|
||||
system_prompt="""You are an expert portfolio analyst specializing in fund analysis and selection. Your core competencies include:
|
||||
- Comprehensive analysis of mutual funds, ETFs, and index funds
|
||||
- Evaluation of fund performance metrics (expense ratios, tracking error, Sharpe ratio)
|
||||
- Assessment of fund composition and strategy alignment
|
||||
- Risk-adjusted return analysis
|
||||
- Tax efficiency considerations
|
||||
|
||||
For each portfolio analysis:
|
||||
1. Evaluate fund characteristics and performance metrics
|
||||
2. Analyze expense ratios and fee structures
|
||||
3. Assess historical performance and volatility
|
||||
4. Compare funds within same category
|
||||
5. Consider tax implications
|
||||
6. Review fund manager track record and strategy consistency
|
||||
|
||||
Maintain focus on cost-efficiency and alignment with investment objectives.""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
saved_state_path="portfolio_analyzer.json",
|
||||
user_name="investment_team",
|
||||
retry_attempts=2,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Asset Allocation Strategist
|
||||
allocation_strategist = Agent(
|
||||
agent_name="Asset-Allocation-Strategist",
|
||||
system_prompt="""You are a specialized asset allocation strategist focused on portfolio construction and optimization. Your expertise includes:
|
||||
- Strategic and tactical asset allocation
|
||||
- Risk tolerance assessment and portfolio matching
|
||||
- Geographic and sector diversification
|
||||
- Rebalancing strategy development
|
||||
- Portfolio optimization using modern portfolio theory
|
||||
|
||||
For each allocation:
|
||||
1. Analyze investor risk tolerance and objectives
|
||||
2. Develop appropriate asset class weights
|
||||
3. Select optimal fund combinations
|
||||
4. Design rebalancing triggers and schedules
|
||||
5. Consider tax-efficient fund placement
|
||||
6. Account for correlation between assets
|
||||
|
||||
Focus on creating well-diversified portfolios aligned with client goals and risk tolerance.""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
saved_state_path="allocation_strategist.json",
|
||||
user_name="investment_team",
|
||||
retry_attempts=2,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Risk Management Specialist
|
||||
risk_manager = Agent(
|
||||
agent_name="Risk-Management-Specialist",
|
||||
system_prompt="""You are a risk management specialist focused on portfolio risk assessment and mitigation. Your expertise covers:
|
||||
- Portfolio risk metrics analysis
|
||||
- Downside protection strategies
|
||||
- Correlation analysis between funds
|
||||
- Stress testing and scenario analysis
|
||||
- Market condition impact assessment
|
||||
|
||||
For each portfolio:
|
||||
1. Calculate key risk metrics (Beta, Standard Deviation, etc.)
|
||||
2. Analyze correlation matrices
|
||||
3. Perform stress tests under various scenarios
|
||||
4. Evaluate liquidity risks
|
||||
5. Assess concentration risks
|
||||
6. Monitor factor exposures
|
||||
|
||||
Focus on maintaining appropriate risk levels while maximizing risk-adjusted returns.""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
saved_state_path="risk_manager.json",
|
||||
user_name="investment_team",
|
||||
retry_attempts=2,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Portfolio Implementation Specialist
|
||||
implementation_specialist = Agent(
|
||||
agent_name="Portfolio-Implementation-Specialist",
|
||||
system_prompt="""You are a portfolio implementation specialist focused on efficient execution and maintenance. Your responsibilities include:
|
||||
- Fund selection for specific asset class exposure
|
||||
- Tax-efficient implementation strategies
|
||||
- Portfolio rebalancing execution
|
||||
- Trading cost analysis
|
||||
- Cash flow management
|
||||
|
||||
For each implementation:
|
||||
1. Select most efficient funds for desired exposure
|
||||
2. Plan tax-efficient transitions
|
||||
3. Design rebalancing schedule
|
||||
4. Optimize trade execution
|
||||
5. Manage cash positions
|
||||
6. Monitor tracking error
|
||||
|
||||
Maintain focus on minimizing costs and maximizing tax efficiency during implementation.""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
saved_state_path="implementation_specialist.json",
|
||||
user_name="investment_team",
|
||||
retry_attempts=2,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Portfolio Monitoring Specialist
|
||||
monitoring_specialist = Agent(
|
||||
agent_name="Portfolio-Monitoring-Specialist",
|
||||
system_prompt="""You are a portfolio monitoring specialist focused on ongoing portfolio oversight and optimization. Your expertise includes:
|
||||
- Regular portfolio performance review
|
||||
- Drift monitoring and rebalancing triggers
|
||||
- Fund changes and replacements
|
||||
- Tax loss harvesting opportunities
|
||||
- Performance attribution analysis
|
||||
|
||||
For each review:
|
||||
1. Track portfolio drift from targets
|
||||
2. Monitor fund performance and changes
|
||||
3. Identify tax loss harvesting opportunities
|
||||
4. Analyze tracking error and expenses
|
||||
5. Review risk metrics evolution
|
||||
6. Generate performance attribution reports
|
||||
|
||||
Ensure continuous alignment with investment objectives while maintaining optimal portfolio efficiency.""",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
saved_state_path="monitoring_specialist.json",
|
||||
user_name="investment_team",
|
||||
retry_attempts=2,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# List of all agents for portfolio management
|
||||
portfolio_agents = [
|
||||
portfolio_analyzer,
|
||||
allocation_strategist,
|
||||
risk_manager,
|
||||
implementation_specialist,
|
||||
monitoring_specialist,
|
||||
]
|
||||
|
||||
|
||||
# Router
|
||||
router = SwarmRouter(
|
||||
name="etf-portfolio-management-swarm",
|
||||
description="Creates and suggests an optimal portfolio",
|
||||
agents=portfolio_agents,
|
||||
swarm_type="SequentialWorkflow", # ConcurrentWorkflow
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
router.run(
|
||||
task="I have 10,000$ and I want to create a porfolio based on energy, ai, and datacenter companies. high growth."
|
||||
)
|
||||
@ -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()
|
||||
@ -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()
|
||||
@ -0,0 +1,333 @@
|
||||
import os
|
||||
from typing import List, Optional
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.v1 import validator
|
||||
from loguru import logger
|
||||
from tenacity import (
|
||||
retry,
|
||||
stop_after_attempt,
|
||||
wait_exponential,
|
||||
)
|
||||
|
||||
from swarm_models import OpenAIFunctionCaller, OpenAIChat
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.swarm_router import SwarmRouter
|
||||
from swarms.structs.agents_available import showcase_available_agents
|
||||
|
||||
|
||||
BOSS_SYSTEM_PROMPT = """
|
||||
Manage a swarm of worker agents to efficiently serve the user by deciding whether to create new agents or delegate tasks. Ensure operations are efficient and effective.
|
||||
|
||||
### Instructions:
|
||||
|
||||
1. **Task Assignment**:
|
||||
- Analyze available worker agents when a task is presented.
|
||||
- Delegate tasks to existing agents with clear, direct, and actionable instructions if an appropriate agent is available.
|
||||
- If no suitable agent exists, create a new agent with a fitting system prompt to handle the task.
|
||||
|
||||
2. **Agent Creation**:
|
||||
- Name agents according to the task they are intended to perform (e.g., "Twitter Marketing Agent").
|
||||
- Provide each new agent with a concise and clear system prompt that includes its role, objectives, and any tools it can utilize.
|
||||
|
||||
3. **Efficiency**:
|
||||
- Minimize redundancy and maximize task completion speed.
|
||||
- Avoid unnecessary agent creation if an existing agent can fulfill the task.
|
||||
|
||||
4. **Communication**:
|
||||
- Be explicit in task delegation instructions to avoid ambiguity and ensure effective task execution.
|
||||
- Require agents to report back on task completion or encountered issues.
|
||||
|
||||
5. **Reasoning and Decisions**:
|
||||
- Offer brief reasoning when selecting or creating agents to maintain transparency.
|
||||
- Avoid using an agent if unnecessary, with a clear explanation if no agents are suitable for a task.
|
||||
|
||||
# Output Format
|
||||
|
||||
Present your plan in clear, bullet-point format or short concise paragraphs, outlining task assignment, agent creation, efficiency strategies, and communication protocols.
|
||||
|
||||
# Notes
|
||||
|
||||
- Preserve transparency by always providing reasoning for task-agent assignments and creation.
|
||||
- Ensure instructions to agents are unambiguous to minimize error.
|
||||
|
||||
"""
|
||||
|
||||
|
||||
class AgentConfig(BaseModel):
|
||||
"""Configuration for an individual agent in a swarm"""
|
||||
|
||||
name: str = Field(
|
||||
description="The name of the agent", example="Research-Agent"
|
||||
)
|
||||
description: str = Field(
|
||||
description="A description of the agent's purpose and capabilities",
|
||||
example="Agent responsible for researching and gathering information",
|
||||
)
|
||||
system_prompt: str = Field(
|
||||
description="The system prompt that defines the agent's behavior",
|
||||
example="You are a research agent. Your role is to gather and analyze information...",
|
||||
)
|
||||
|
||||
@validator("name")
|
||||
def validate_name(cls, v):
|
||||
if not v.strip():
|
||||
raise ValueError("Agent name cannot be empty")
|
||||
return v.strip()
|
||||
|
||||
@validator("system_prompt")
|
||||
def validate_system_prompt(cls, v):
|
||||
if not v.strip():
|
||||
raise ValueError("System prompt cannot be empty")
|
||||
return v.strip()
|
||||
|
||||
|
||||
class SwarmConfig(BaseModel):
|
||||
"""Configuration for a swarm of cooperative agents"""
|
||||
|
||||
name: str = Field(
|
||||
description="The name of the swarm",
|
||||
example="Research-Writing-Swarm",
|
||||
)
|
||||
description: str = Field(
|
||||
description="The description of the swarm's purpose and capabilities",
|
||||
example="A swarm of agents that work together to research topics and write articles",
|
||||
)
|
||||
agents: List[AgentConfig] = Field(
|
||||
description="The list of agents that make up the swarm",
|
||||
min_items=1,
|
||||
)
|
||||
|
||||
@validator("agents")
|
||||
def validate_agents(cls, v):
|
||||
if not v:
|
||||
raise ValueError("Swarm must have at least one agent")
|
||||
return v
|
||||
|
||||
|
||||
class AutoSwarmBuilder:
|
||||
"""A class that automatically builds and manages swarms of AI agents with enhanced error handling."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
verbose: bool = True,
|
||||
api_key: Optional[str] = None,
|
||||
model_name: str = "gpt-4",
|
||||
):
|
||||
self.name = name or "DefaultSwarm"
|
||||
self.description = description or "Generic AI Agent Swarm"
|
||||
self.verbose = verbose
|
||||
self.agents_pool = []
|
||||
self.api_key = api_key or os.getenv("OPENAI_API_KEY")
|
||||
self.model_name = model_name
|
||||
|
||||
if not self.api_key:
|
||||
raise ValueError(
|
||||
"OpenAI API key must be provided either through initialization or environment variable"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Initialized AutoSwarmBuilder",
|
||||
extra={
|
||||
"swarm_name": self.name,
|
||||
"description": self.description,
|
||||
"model": self.model_name,
|
||||
},
|
||||
)
|
||||
|
||||
# Initialize OpenAI chat model
|
||||
try:
|
||||
self.chat_model = OpenAIChat(
|
||||
openai_api_key=self.api_key,
|
||||
model_name=self.model_name,
|
||||
temperature=0.1,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to initialize OpenAI chat model: {str(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=10),
|
||||
)
|
||||
def run(self, task: str, image_url: Optional[str] = None) -> str:
|
||||
"""Run the swarm on a given task with error handling and retries."""
|
||||
if not task or not task.strip():
|
||||
raise ValueError("Task cannot be empty")
|
||||
|
||||
logger.info("Starting swarm execution", extra={"task": task})
|
||||
|
||||
try:
|
||||
# Create agents for the task
|
||||
agents = self._create_agents(task, image_url)
|
||||
if not agents:
|
||||
raise ValueError(
|
||||
"No agents were created for the task"
|
||||
)
|
||||
|
||||
# Execute the task through the swarm router
|
||||
logger.info(
|
||||
"Routing task through swarm",
|
||||
extra={"num_agents": len(agents)},
|
||||
)
|
||||
output = self.swarm_router(agents, task, image_url)
|
||||
|
||||
logger.info("Swarm execution completed successfully")
|
||||
return output
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error during swarm execution: {str(e)}",
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
def _create_agents(
|
||||
self, task: str, image_url: Optional[str] = None
|
||||
) -> List[Agent]:
|
||||
"""Create the necessary agents for a task with enhanced error handling."""
|
||||
logger.info("Creating agents for task", extra={"task": task})
|
||||
|
||||
try:
|
||||
model = OpenAIFunctionCaller(
|
||||
system_prompt=BOSS_SYSTEM_PROMPT,
|
||||
api_key=self.api_key,
|
||||
temperature=0.1,
|
||||
base_model=SwarmConfig,
|
||||
)
|
||||
|
||||
agents_config = model.run(task)
|
||||
print(f"{agents_config}")
|
||||
|
||||
if isinstance(agents_config, dict):
|
||||
agents_config = SwarmConfig(**agents_config)
|
||||
|
||||
# Update swarm configuration
|
||||
self.name = agents_config.name
|
||||
self.description = agents_config.description
|
||||
|
||||
# Create agents from configuration
|
||||
agents = []
|
||||
for agent_config in agents_config.agents:
|
||||
if isinstance(agent_config, dict):
|
||||
agent_config = AgentConfig(**agent_config)
|
||||
|
||||
agent = self.build_agent(
|
||||
agent_name=agent_config.name,
|
||||
agent_description=agent_config.description,
|
||||
agent_system_prompt=agent_config.system_prompt,
|
||||
)
|
||||
agents.append(agent)
|
||||
|
||||
# Add available agents showcase to system prompts
|
||||
agents_available = showcase_available_agents(
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
agents=agents,
|
||||
)
|
||||
|
||||
for agent in agents:
|
||||
agent.system_prompt += "\n" + agents_available
|
||||
|
||||
logger.info(
|
||||
"Successfully created agents",
|
||||
extra={"num_agents": len(agents)},
|
||||
)
|
||||
return agents
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error creating agents: {str(e)}", exc_info=True
|
||||
)
|
||||
raise
|
||||
|
||||
def build_agent(
|
||||
self,
|
||||
agent_name: str,
|
||||
agent_description: str,
|
||||
agent_system_prompt: str,
|
||||
) -> Agent:
|
||||
"""Build a single agent with enhanced error handling."""
|
||||
logger.info(
|
||||
"Building agent", extra={"agent_name": agent_name}
|
||||
)
|
||||
|
||||
try:
|
||||
agent = Agent(
|
||||
agent_name=agent_name,
|
||||
description=agent_description,
|
||||
system_prompt=agent_system_prompt,
|
||||
llm=self.chat_model,
|
||||
autosave=True,
|
||||
dashboard=False,
|
||||
verbose=self.verbose,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path=f"states/{agent_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
|
||||
user_name="swarms_corp",
|
||||
retry_attempts=3,
|
||||
context_length=200000,
|
||||
return_step_meta=False,
|
||||
output_type="str",
|
||||
streaming_on=False,
|
||||
auto_generate_prompt=True,
|
||||
)
|
||||
return agent
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error building agent: {str(e)}", exc_info=True
|
||||
)
|
||||
raise
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=10),
|
||||
)
|
||||
def swarm_router(
|
||||
self,
|
||||
agents: List[Agent],
|
||||
task: str,
|
||||
image_url: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Route tasks between agents in the swarm with error handling and retries."""
|
||||
logger.info(
|
||||
"Initializing swarm router",
|
||||
extra={"num_agents": len(agents)},
|
||||
)
|
||||
|
||||
try:
|
||||
swarm_router_instance = SwarmRouter(
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
agents=agents,
|
||||
swarm_type="auto",
|
||||
)
|
||||
|
||||
formatted_task = f"{self.name} {self.description} {task}"
|
||||
result = swarm_router_instance.run(formatted_task)
|
||||
|
||||
logger.info("Successfully completed swarm routing")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error in swarm router: {str(e)}", exc_info=True
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
swarm = AutoSwarmBuilder(
|
||||
name="ChipDesign-Swarm",
|
||||
description="A swarm of specialized AI agents for chip design",
|
||||
api_key="your-api-key", # Optional if set in environment
|
||||
model_name="gpt-4", # Optional, defaults to gpt-4
|
||||
)
|
||||
|
||||
result = swarm.run(
|
||||
"Design a new AI accelerator chip optimized for transformer model inference..."
|
||||
)
|
||||
@ -1,84 +1,43 @@
|
||||
import os
|
||||
from swarms.structs.agent import Agent
|
||||
from swarm_models.popular_llms import OpenAIChat
|
||||
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
|
||||
agent_registry = AgentRegistry(
|
||||
name="Swarms CLI",
|
||||
description="A registry of agents for the Swarms CLI",
|
||||
)
|
||||
|
||||
|
||||
def create_agent(name: str, system_prompt: str, max_loops: int = 1):
|
||||
# Run the agents in the registry
|
||||
def run_agent_by_name(
|
||||
name: str,
|
||||
system_prompt: str,
|
||||
model_name: str,
|
||||
max_loops: int,
|
||||
task: str,
|
||||
img: str,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Create and initialize an agent with the given parameters.
|
||||
This function creates an Agent instance and runs a task on it.
|
||||
|
||||
Args:
|
||||
name (str): The name of the agent.
|
||||
system_prompt (str): The system prompt for the agent.
|
||||
max_loops (int, optional): The maximum number of loops the agent can perform. Defaults to 1.
|
||||
model_name (str): The name of the model used by the agent.
|
||||
max_loops (int): The maximum number of loops the agent can run.
|
||||
task (str): The task to be run by the agent.
|
||||
*args: Variable length arguments.
|
||||
**kwargs: Keyword arguments.
|
||||
|
||||
Returns:
|
||||
Agent: The initialized agent.
|
||||
|
||||
The output of the task run by the agent.
|
||||
"""
|
||||
# Initialize the agent
|
||||
agent = Agent(
|
||||
agent_name=name,
|
||||
system_prompt=system_prompt,
|
||||
llm=model,
|
||||
max_loops=max_loops,
|
||||
autosave=True,
|
||||
dashboard=False,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path=f"{name}.json",
|
||||
user_name="swarms_corp",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
# return_step_meta=True,
|
||||
# disable_print_every_step=True,
|
||||
# output_type="json",
|
||||
interactive=True,
|
||||
)
|
||||
|
||||
agent_registry.add(agent)
|
||||
|
||||
return agent
|
||||
|
||||
|
||||
# Run the agents in the registry
|
||||
def run_agent_by_name(name: str, task: str, *args, **kwargs):
|
||||
"""
|
||||
Run an agent by its name and perform a specified task.
|
||||
|
||||
Parameters:
|
||||
- name (str): The name of the agent.
|
||||
- task (str): The task to be performed by the agent.
|
||||
- *args: Variable length argument list.
|
||||
- **kwargs: Arbitrary keyword arguments.
|
||||
|
||||
Returns:
|
||||
- output: The output of the agent's task.
|
||||
|
||||
"""
|
||||
agent = agent_registry.get_agent_by_name(name)
|
||||
|
||||
output = agent.run(task, *args, **kwargs)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# # Test
|
||||
# out = create_agent("Accountant1", "Prepares financial statements")
|
||||
# print(out)
|
||||
try:
|
||||
agent = Agent(
|
||||
agent_name=name,
|
||||
system_prompt=system_prompt,
|
||||
model_name=model_name,
|
||||
max_loops=max_loops,
|
||||
)
|
||||
|
||||
output = agent.run(task=task, img=img, *args, **kwargs)
|
||||
|
||||
return output
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {str(e)}")
|
||||
return None
|
||||
|
||||
@ -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}"
|
||||
@ -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",
|
||||
]
|
||||
@ -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()
|
||||
@ -1,40 +0,0 @@
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
from swarms.utils.loguru_logger import initialize_logger
|
||||
from swarms.telemetry.check_update import check_for_update
|
||||
|
||||
logger = initialize_logger(log_folder="auto_upgrade_swarms")
|
||||
|
||||
|
||||
def auto_update():
|
||||
"""auto update swarms"""
|
||||
try:
|
||||
# Check if auto-update is disabled
|
||||
auto_update_enabled = os.getenv(
|
||||
"SWARMS_AUTOUPDATE_ON", "false"
|
||||
).lower()
|
||||
if auto_update_enabled == "false":
|
||||
logger.info(
|
||||
"Auto-update is disabled via SWARMS_AUTOUPDATE_ON"
|
||||
)
|
||||
return
|
||||
|
||||
outcome = check_for_update()
|
||||
if outcome is True:
|
||||
logger.info(
|
||||
"There is a new version of swarms available! Downloading..."
|
||||
)
|
||||
try:
|
||||
subprocess.run(
|
||||
["pip", "install", "-U", "swarms"], check=True
|
||||
)
|
||||
except subprocess.CalledProcessError:
|
||||
logger.info("Attempting to install with pip3...")
|
||||
subprocess.run(
|
||||
["pip3", "install", "-U", "swarms"], check=True
|
||||
)
|
||||
else:
|
||||
logger.info("swarms is up to date!")
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
@ -1,73 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
|
||||
import pkg_resources
|
||||
import requests
|
||||
from packaging import version
|
||||
from swarms.utils.loguru_logger import initialize_logger
|
||||
|
||||
logger = initialize_logger("check-update")
|
||||
|
||||
|
||||
# borrowed from: https://stackoverflow.com/a/1051266/656011
|
||||
def check_for_package(package: str) -> bool:
|
||||
"""
|
||||
Checks if a package is installed and available for import.
|
||||
|
||||
Args:
|
||||
package (str): The name of the package to check.
|
||||
|
||||
Returns:
|
||||
bool: True if the package is installed and can be imported, False otherwise.
|
||||
"""
|
||||
if package in sys.modules:
|
||||
return True
|
||||
elif (spec := importlib.util.find_spec(package)) is not None:
|
||||
try:
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
|
||||
sys.modules[package] = module
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
logger.error(f"Failed to import {package}")
|
||||
return False
|
||||
else:
|
||||
logger.info(f"{package} not found")
|
||||
return False
|
||||
|
||||
|
||||
def check_for_update() -> bool:
|
||||
"""
|
||||
Checks if there is an update available for the swarms package.
|
||||
|
||||
Returns:
|
||||
bool: True if an update is available, False otherwise.
|
||||
"""
|
||||
try:
|
||||
# Fetch the latest version from the PyPI API
|
||||
response = requests.get("https://pypi.org/pypi/swarms/json")
|
||||
response.raise_for_status() # Raises an HTTPError if the response status code is 4XX/5XX
|
||||
latest_version = response.json()["info"]["version"]
|
||||
|
||||
# Get the current version using pkg_resources
|
||||
current_version = pkg_resources.get_distribution(
|
||||
"swarms"
|
||||
).version
|
||||
|
||||
if version.parse(latest_version) > version.parse(
|
||||
current_version
|
||||
):
|
||||
logger.info(
|
||||
f"Update available: {latest_version} > {current_version}"
|
||||
)
|
||||
return True
|
||||
else:
|
||||
logger.info(
|
||||
f"No update available: {latest_version} <= {current_version}"
|
||||
)
|
||||
return False
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Failed to check for update: {e}")
|
||||
return False
|
||||
@ -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
|
||||
@ -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
|
||||
@ -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"))
|
||||
Loading…
Reference in new issue