Kye Gomez
2 weeks ago
parent
5e4a600d47
commit
d44fbcd700
@ -0,0 +1,43 @@
|
||||
# Swarms Tools Example with HTX + CoinGecko
|
||||
|
||||
- `pip3 install swarms swarms-tools`
|
||||
- Add `OPENAI_API_KEY` to your `.env` file
|
||||
- Run `swarms_tools_htx_gecko.py`
|
||||
- Agent will make a function call to the desired tool
|
||||
- The tool will be executed and the result will be returned to the agent
|
||||
- The agent will then analyze the result and return the final output
|
||||
|
||||
|
||||
```python
|
||||
from swarms import Agent
|
||||
from swarms.prompts.finance_agent_sys_prompt import (
|
||||
FINANCIAL_AGENT_SYS_PROMPT,
|
||||
)
|
||||
from swarms_tools import (
|
||||
fetch_stock_news,
|
||||
coin_gecko_coin_api,
|
||||
fetch_htx_data,
|
||||
)
|
||||
|
||||
# Initialize the agent
|
||||
agent = Agent(
|
||||
agent_name="Financial-Analysis-Agent",
|
||||
agent_description="Personal finance advisor agent",
|
||||
system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
|
||||
max_loops=1,
|
||||
model_name="gpt-4o",
|
||||
dynamic_temperature_enabled=True,
|
||||
user_name="swarms_corp",
|
||||
retry_attempts=3,
|
||||
context_length=8192,
|
||||
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
|
||||
max_tokens=4000, # max output tokens
|
||||
saved_state_path="agent_00.json",
|
||||
interactive=False,
|
||||
tools=[fetch_stock_news, coin_gecko_coin_api, fetch_htx_data],
|
||||
)
|
||||
|
||||
agent.run("Analyze the $swarms token on htx")
|
||||
```
|
||||
@ -0,0 +1,33 @@
|
||||
from swarms.utils.formatter import formatter
|
||||
|
||||
|
||||
def agent_print(
|
||||
agent_name: str,
|
||||
response: str = None,
|
||||
loop_count: int = None,
|
||||
streaming_on: bool = False,
|
||||
):
|
||||
"""
|
||||
Prints the response from an agent based on the streaming mode.
|
||||
|
||||
Args:
|
||||
agent_name (str): The name of the agent.
|
||||
response (str): The response from the agent.
|
||||
loop_count (int): The maximum number of loops.
|
||||
streaming_on (bool): Indicates if streaming is on or off.
|
||||
|
||||
Returns:
|
||||
str: The response from the agent.
|
||||
"""
|
||||
if streaming_on:
|
||||
formatter.print_panel_token_by_token(
|
||||
f"{agent_name}: {response}",
|
||||
title=f"Agent Name: {agent_name} [Max Loops: {loop_count}]",
|
||||
)
|
||||
else:
|
||||
formatter.print_panel(
|
||||
f"{agent_name}: {response}",
|
||||
f"Agent Name {agent_name} [Max Loops: {loop_count} ]",
|
||||
)
|
||||
|
||||
return response
|
||||
@ -0,0 +1,318 @@
|
||||
import base64
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Union, Dict, List
|
||||
|
||||
from cryptography.fernet import Fernet
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import padding, rsa
|
||||
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
||||
|
||||
|
||||
@dataclass
|
||||
class EncryptedMessage:
|
||||
"""Structure for encrypted messages between agents"""
|
||||
|
||||
sender_id: str
|
||||
receiver_id: str
|
||||
encrypted_content: bytes
|
||||
timestamp: float
|
||||
message_id: str
|
||||
session_id: str
|
||||
|
||||
|
||||
class EncryptionSession:
|
||||
"""Represents an encrypted communication session between agents"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
agent_ids: List[str],
|
||||
encrypted_keys: Dict[str, bytes],
|
||||
created_at: datetime,
|
||||
):
|
||||
self.session_id = session_id
|
||||
self.agent_ids = agent_ids
|
||||
self.encrypted_keys = encrypted_keys
|
||||
self.created_at = created_at
|
||||
|
||||
|
||||
class AgentEncryption:
|
||||
"""
|
||||
Handles encryption for agent data both at rest and in transit.
|
||||
Supports both symmetric (for data at rest) and asymmetric (for data in transit) encryption.
|
||||
Also supports secure multi-agent communication.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent_id: Optional[str] = None,
|
||||
encryption_key: Optional[str] = None,
|
||||
enable_transit_encryption: bool = False,
|
||||
enable_rest_encryption: bool = False,
|
||||
enable_multi_agent: bool = False,
|
||||
):
|
||||
self.agent_id = agent_id or str(uuid.uuid4())
|
||||
self.enable_transit_encryption = enable_transit_encryption
|
||||
self.enable_rest_encryption = enable_rest_encryption
|
||||
self.enable_multi_agent = enable_multi_agent
|
||||
|
||||
# Multi-agent communication storage
|
||||
self.sessions: Dict[str, EncryptionSession] = {}
|
||||
self.known_agents: Dict[str, "AgentEncryption"] = {}
|
||||
|
||||
if enable_rest_encryption:
|
||||
# Initialize encryption for data at rest
|
||||
if encryption_key:
|
||||
self.encryption_key = base64.urlsafe_b64encode(
|
||||
PBKDF2HMAC(
|
||||
algorithm=hashes.SHA256(),
|
||||
length=32,
|
||||
salt=f"agent_{self.agent_id}".encode(), # Unique salt per agent
|
||||
iterations=100000,
|
||||
).derive(encryption_key.encode())
|
||||
)
|
||||
else:
|
||||
self.encryption_key = Fernet.generate_key()
|
||||
|
||||
self.cipher_suite = Fernet(self.encryption_key)
|
||||
|
||||
if enable_transit_encryption or enable_multi_agent:
|
||||
# Generate RSA key pair for transit encryption
|
||||
self.private_key = rsa.generate_private_key(
|
||||
public_exponent=65537, key_size=2048
|
||||
)
|
||||
self.public_key = self.private_key.public_key()
|
||||
|
||||
def register_agent(
|
||||
self, agent_id: str, agent_encryption: "AgentEncryption"
|
||||
) -> None:
|
||||
"""Register another agent for secure communication"""
|
||||
if not self.enable_multi_agent:
|
||||
raise ValueError("Multi-agent support is not enabled")
|
||||
self.known_agents[agent_id] = agent_encryption
|
||||
|
||||
def create_session(self, agent_ids: List[str]) -> str:
|
||||
"""Create a new encrypted session between multiple agents"""
|
||||
if not self.enable_multi_agent:
|
||||
raise ValueError("Multi-agent support is not enabled")
|
||||
|
||||
session_id = str(uuid.uuid4())
|
||||
|
||||
# Generate a shared session key
|
||||
session_key = Fernet.generate_key()
|
||||
|
||||
# Create encrypted copies of the session key for each agent
|
||||
encrypted_keys = {}
|
||||
for agent_id in agent_ids:
|
||||
if (
|
||||
agent_id not in self.known_agents
|
||||
and agent_id != self.agent_id
|
||||
):
|
||||
raise ValueError(f"Agent {agent_id} not registered")
|
||||
|
||||
if agent_id == self.agent_id:
|
||||
agent_public_key = self.public_key
|
||||
else:
|
||||
agent_public_key = self.known_agents[
|
||||
agent_id
|
||||
].public_key
|
||||
|
||||
encrypted_key = agent_public_key.encrypt(
|
||||
session_key,
|
||||
padding.OAEP(
|
||||
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
||||
algorithm=hashes.SHA256(),
|
||||
label=None,
|
||||
),
|
||||
)
|
||||
encrypted_keys[agent_id] = encrypted_key
|
||||
|
||||
# Store session information
|
||||
self.sessions[session_id] = EncryptionSession(
|
||||
session_id=session_id,
|
||||
agent_ids=agent_ids,
|
||||
encrypted_keys=encrypted_keys,
|
||||
created_at=datetime.now(),
|
||||
)
|
||||
|
||||
return session_id
|
||||
|
||||
def encrypt_message(
|
||||
self,
|
||||
content: Union[str, dict],
|
||||
receiver_id: str,
|
||||
session_id: str,
|
||||
) -> EncryptedMessage:
|
||||
"""Encrypt a message for another agent within a session"""
|
||||
if not self.enable_multi_agent:
|
||||
raise ValueError("Multi-agent support is not enabled")
|
||||
|
||||
if session_id not in self.sessions:
|
||||
raise ValueError("Invalid session ID")
|
||||
|
||||
session = self.sessions[session_id]
|
||||
if (
|
||||
self.agent_id not in session.agent_ids
|
||||
or receiver_id not in session.agent_ids
|
||||
):
|
||||
raise ValueError("Sender or receiver not in session")
|
||||
|
||||
# Serialize content if it's a dictionary
|
||||
if isinstance(content, dict):
|
||||
content = json.dumps(content)
|
||||
|
||||
# Get the session key
|
||||
encrypted_session_key = session.encrypted_keys[self.agent_id]
|
||||
session_key = self.decrypt_session_key(encrypted_session_key)
|
||||
|
||||
# Create Fernet cipher with session key
|
||||
cipher = Fernet(session_key)
|
||||
|
||||
# Encrypt the message
|
||||
encrypted_content = cipher.encrypt(content.encode())
|
||||
|
||||
return EncryptedMessage(
|
||||
sender_id=self.agent_id,
|
||||
receiver_id=receiver_id,
|
||||
encrypted_content=encrypted_content,
|
||||
timestamp=datetime.now().timestamp(),
|
||||
message_id=str(uuid.uuid4()),
|
||||
session_id=session_id,
|
||||
)
|
||||
|
||||
def decrypt_message(
|
||||
self, message: EncryptedMessage
|
||||
) -> Union[str, dict]:
|
||||
"""Decrypt a message from another agent"""
|
||||
if not self.enable_multi_agent:
|
||||
raise ValueError("Multi-agent support is not enabled")
|
||||
|
||||
if message.session_id not in self.sessions:
|
||||
raise ValueError("Invalid session ID")
|
||||
|
||||
if self.agent_id != message.receiver_id:
|
||||
raise ValueError("Message not intended for this agent")
|
||||
|
||||
session = self.sessions[message.session_id]
|
||||
|
||||
# Get the session key
|
||||
encrypted_session_key = session.encrypted_keys[self.agent_id]
|
||||
session_key = self.decrypt_session_key(encrypted_session_key)
|
||||
|
||||
# Create Fernet cipher with session key
|
||||
cipher = Fernet(session_key)
|
||||
|
||||
# Decrypt the message
|
||||
decrypted_content = cipher.decrypt(
|
||||
message.encrypted_content
|
||||
).decode()
|
||||
|
||||
# Try to parse as JSON
|
||||
try:
|
||||
return json.loads(decrypted_content)
|
||||
except json.JSONDecodeError:
|
||||
return decrypted_content
|
||||
|
||||
def decrypt_session_key(self, encrypted_key: bytes) -> bytes:
|
||||
"""Decrypt a session key using the agent's private key"""
|
||||
return self.private_key.decrypt(
|
||||
encrypted_key,
|
||||
padding.OAEP(
|
||||
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
||||
algorithm=hashes.SHA256(),
|
||||
label=None,
|
||||
),
|
||||
)
|
||||
|
||||
# Original methods preserved below
|
||||
def encrypt_at_rest(self, data: Union[str, dict, bytes]) -> bytes:
|
||||
"""Encrypts data for storage"""
|
||||
if not self.enable_rest_encryption:
|
||||
return (
|
||||
data
|
||||
if isinstance(data, bytes)
|
||||
else str(data).encode()
|
||||
)
|
||||
|
||||
if isinstance(data, dict):
|
||||
data = json.dumps(data)
|
||||
if isinstance(data, str):
|
||||
data = data.encode()
|
||||
|
||||
return self.cipher_suite.encrypt(data)
|
||||
|
||||
def decrypt_at_rest(
|
||||
self, encrypted_data: bytes
|
||||
) -> Union[str, dict]:
|
||||
"""Decrypts stored data"""
|
||||
if not self.enable_rest_encryption:
|
||||
return encrypted_data.decode()
|
||||
|
||||
decrypted_data = self.cipher_suite.decrypt(encrypted_data)
|
||||
|
||||
try:
|
||||
return json.loads(decrypted_data)
|
||||
except json.JSONDecodeError:
|
||||
return decrypted_data.decode()
|
||||
|
||||
def encrypt_for_transit(self, data: Union[str, dict]) -> bytes:
|
||||
"""Encrypts data for transmission"""
|
||||
if not self.enable_transit_encryption:
|
||||
return str(data).encode()
|
||||
|
||||
if isinstance(data, dict):
|
||||
data = json.dumps(data)
|
||||
|
||||
return self.public_key.encrypt(
|
||||
data.encode(),
|
||||
padding.OAEP(
|
||||
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
||||
algorithm=hashes.SHA256(),
|
||||
label=None,
|
||||
),
|
||||
)
|
||||
|
||||
def decrypt_from_transit(
|
||||
self, data: Union[bytes, str]
|
||||
) -> Union[str, dict]:
|
||||
"""Decrypts received data, handling both encrypted and unencrypted inputs"""
|
||||
if not self.enable_transit_encryption:
|
||||
return data.decode() if isinstance(data, bytes) else data
|
||||
|
||||
try:
|
||||
if isinstance(data, bytes) and len(data) == 256:
|
||||
decrypted_data = self.private_key.decrypt(
|
||||
data,
|
||||
padding.OAEP(
|
||||
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
||||
algorithm=hashes.SHA256(),
|
||||
label=None,
|
||||
),
|
||||
).decode()
|
||||
else:
|
||||
return (
|
||||
data.decode() if isinstance(data, bytes) else data
|
||||
)
|
||||
|
||||
try:
|
||||
return json.loads(decrypted_data)
|
||||
except json.JSONDecodeError:
|
||||
return decrypted_data
|
||||
except ValueError:
|
||||
return data.decode() if isinstance(data, bytes) else data
|
||||
|
||||
def get_public_key_pem(self) -> bytes:
|
||||
"""Returns the public key in PEM format for sharing"""
|
||||
if (
|
||||
not self.enable_transit_encryption
|
||||
and not self.enable_multi_agent
|
||||
):
|
||||
return b""
|
||||
|
||||
return self.public_key.public_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PublicFormat.SubjectPublicKeyInfo,
|
||||
)
|
||||
@ -0,0 +1,30 @@
|
||||
from swarms import Agent
|
||||
from swarms.prompts.finance_agent_sys_prompt import (
|
||||
FINANCIAL_AGENT_SYS_PROMPT,
|
||||
)
|
||||
from swarms_tools import (
|
||||
coin_gecko_coin_api,
|
||||
fetch_htx_data,
|
||||
)
|
||||
|
||||
|
||||
# Initialize the agent
|
||||
agent = Agent(
|
||||
agent_name="Financial-Analysis-Agent",
|
||||
agent_description="Personal finance advisor agent",
|
||||
system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
|
||||
max_loops=1,
|
||||
model_name="gpt-4o",
|
||||
dynamic_temperature_enabled=True,
|
||||
user_name="swarms_corp",
|
||||
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
|
||||
max_tokens=4000, # max output tokens
|
||||
saved_state_path="agent_00.json",
|
||||
interactive=False,
|
||||
)
|
||||
|
||||
agent.run(
|
||||
f"Analyze the $swarms token on HTX with data: {fetch_htx_data('swarms')}. Additionally, consider the following CoinGecko data: {coin_gecko_coin_api('swarms')}"
|
||||
)
|
||||
@ -0,0 +1,92 @@
|
||||
# Define a simple testing framework
|
||||
from swarms.tools.tool_parse_exec import parse_and_execute_json
|
||||
|
||||
|
||||
def run_test(test_name, test_func):
|
||||
print(f"Running test: {test_name}")
|
||||
print("------------------------------------------------")
|
||||
try:
|
||||
test_func()
|
||||
print(f"✓ {test_name} passed")
|
||||
print("------------------------------------------------")
|
||||
except Exception as e:
|
||||
print(f"✗ {test_name} failed: {str(e)}")
|
||||
print("------------------------------------------------")
|
||||
|
||||
|
||||
# Mock functions for testing
|
||||
def mock_function_a(param1, param2):
|
||||
return param1 + param2
|
||||
|
||||
|
||||
def mock_function_b(param1):
|
||||
if param1 < 0:
|
||||
raise ValueError("Negative value not allowed")
|
||||
return param1 * 2
|
||||
|
||||
|
||||
# Test cases
|
||||
def test_parse_and_execute_json_success():
|
||||
functions = [mock_function_a, mock_function_b]
|
||||
json_string = '{"functions": [{"name": "mock_function_a", "parameters": {"param1": 1, "param2": 2}}, {"name": "mock_function_b", "parameters": {"param1": 3}}]}'
|
||||
|
||||
result = parse_and_execute_json(functions, json_string)
|
||||
expected_result = {
|
||||
"results": {"mock_function_a": "3", "mock_function_b": "6"},
|
||||
"summary": "mock_function_a: 3\nmock_function_b: 6",
|
||||
}
|
||||
|
||||
assert (
|
||||
result == expected_result
|
||||
), f"Expected {expected_result}, but got {result}"
|
||||
|
||||
|
||||
def test_parse_and_execute_json_function_not_found():
|
||||
functions = [mock_function_a]
|
||||
json_string = '{"functions": [{"name": "non_existent_function", "parameters": {}}]}'
|
||||
|
||||
result = parse_and_execute_json(functions, json_string)
|
||||
expected_result = {
|
||||
"results": {
|
||||
"non_existent_function": "Error: Function non_existent_function not found"
|
||||
},
|
||||
"summary": "non_existent_function: Error: Function non_existent_function not found",
|
||||
}
|
||||
|
||||
assert (
|
||||
result == expected_result
|
||||
), f"Expected {expected_result}, but got {result}"
|
||||
|
||||
|
||||
def test_parse_and_execute_json_error_handling():
|
||||
functions = [mock_function_b]
|
||||
json_string = '{"functions": [{"name": "mock_function_b", "parameters": {"param1": -1}}]}'
|
||||
|
||||
result = parse_and_execute_json(functions, json_string)
|
||||
expected_result = {
|
||||
"results": {
|
||||
"mock_function_b": "Error: Negative value not allowed"
|
||||
},
|
||||
"summary": "mock_function_b: Error: Negative value not allowed",
|
||||
}
|
||||
|
||||
assert (
|
||||
result == expected_result
|
||||
), f"Expected {expected_result}, but got {result}"
|
||||
|
||||
|
||||
# Run tests
|
||||
run_test(
|
||||
"Test parse_and_execute_json success",
|
||||
test_parse_and_execute_json_success,
|
||||
)
|
||||
print("------------------------------------------------")
|
||||
run_test(
|
||||
"Test parse_and_execute_json function not found",
|
||||
test_parse_and_execute_json_function_not_found,
|
||||
)
|
||||
print("------------------------------------------------")
|
||||
run_test(
|
||||
"Test parse_and_execute_json error handling",
|
||||
test_parse_and_execute_json_error_handling,
|
||||
)
|
||||
Loading…
Reference in new issue