Kye Gomez
2 months ago
parent
5477635441
commit
007eb5c011
@ -1,910 +0,0 @@
|
|||||||
|
|
||||||
import os
|
|
||||||
from collections import defaultdict
|
|
||||||
from datetime import datetime
|
|
||||||
from decimal import Decimal
|
|
||||||
from functools import lru_cache
|
|
||||||
from threading import Thread
|
|
||||||
from time import sleep, time
|
|
||||||
from typing import Any, Dict, List, Optional, Union
|
|
||||||
|
|
||||||
import pytz
|
|
||||||
import supabase
|
|
||||||
from dotenv import load_dotenv
|
|
||||||
from fastapi import (
|
|
||||||
Depends,
|
|
||||||
FastAPI,
|
|
||||||
Header,
|
|
||||||
HTTPException,
|
|
||||||
Request,
|
|
||||||
status,
|
|
||||||
)
|
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
|
||||||
from loguru import logger
|
|
||||||
from pydantic import BaseModel, Field
|
|
||||||
from swarms import Agent, SwarmRouter, SwarmType
|
|
||||||
from swarms.utils.litellm_tokenizer import count_tokens
|
|
||||||
import asyncio
|
|
||||||
|
|
||||||
load_dotenv()
|
|
||||||
|
|
||||||
# Define rate limit parameters
|
|
||||||
RATE_LIMIT = 100 # Max requests
|
|
||||||
TIME_WINDOW = 60 # Time window in seconds
|
|
||||||
|
|
||||||
# In-memory store for tracking requests
|
|
||||||
request_counts = defaultdict(lambda: {"count": 0, "start_time": time()})
|
|
||||||
|
|
||||||
# In-memory store for scheduled jobs
|
|
||||||
scheduled_jobs: Dict[str, Dict] = {}
|
|
||||||
|
|
||||||
|
|
||||||
def rate_limiter(request: Request):
|
|
||||||
client_ip = request.client.host
|
|
||||||
current_time = time()
|
|
||||||
client_data = request_counts[client_ip]
|
|
||||||
|
|
||||||
# Reset count if time window has passed
|
|
||||||
if current_time - client_data["start_time"] > TIME_WINDOW:
|
|
||||||
client_data["count"] = 0
|
|
||||||
client_data["start_time"] = current_time
|
|
||||||
|
|
||||||
# Increment request count
|
|
||||||
client_data["count"] += 1
|
|
||||||
|
|
||||||
# Check if rate limit is exceeded
|
|
||||||
if client_data["count"] > RATE_LIMIT:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=429, detail="Rate limit exceeded. Please try again later."
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class AgentSpec(BaseModel):
|
|
||||||
agent_name: Optional[str] = Field(None, description="Agent Name", max_length=100)
|
|
||||||
description: Optional[str] = Field(None, description="Description", max_length=500)
|
|
||||||
system_prompt: Optional[str] = Field(
|
|
||||||
None, description="System Prompt", max_length=500
|
|
||||||
)
|
|
||||||
model_name: Optional[str] = Field(
|
|
||||||
"gpt-4o", description="Model Name", max_length=500
|
|
||||||
)
|
|
||||||
auto_generate_prompt: Optional[bool] = Field(
|
|
||||||
False, description="Auto Generate Prompt"
|
|
||||||
)
|
|
||||||
max_tokens: Optional[int] = Field(None, description="Max Tokens")
|
|
||||||
temperature: Optional[float] = Field(0.5, description="Temperature")
|
|
||||||
role: Optional[str] = Field("worker", description="Role")
|
|
||||||
max_loops: Optional[int] = Field(1, description="Max Loops")
|
|
||||||
|
|
||||||
|
|
||||||
# class ExternalAgent(BaseModel):
|
|
||||||
# base_url: str = Field(..., description="Base URL")
|
|
||||||
# parameters: Dict[str, Any] = Field(..., description="Parameters")
|
|
||||||
# headers: Dict[str, Any] = Field(..., description="Headers")
|
|
||||||
|
|
||||||
|
|
||||||
class ScheduleSpec(BaseModel):
|
|
||||||
scheduled_time: datetime = Field(..., description="When to execute the swarm (UTC)")
|
|
||||||
timezone: Optional[str] = Field(
|
|
||||||
"UTC", description="Timezone for the scheduled time"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class SwarmSpec(BaseModel):
|
|
||||||
name: Optional[str] = Field(None, description="Swarm Name", max_length=100)
|
|
||||||
description: Optional[str] = Field(None, description="Description")
|
|
||||||
agents: Optional[Union[List[AgentSpec], Any]] = Field(None, description="Agents")
|
|
||||||
max_loops: Optional[int] = Field(None, description="Max Loops")
|
|
||||||
swarm_type: Optional[SwarmType] = Field(None, description="Swarm Type")
|
|
||||||
rearrange_flow: Optional[str] = Field(None, description="Flow")
|
|
||||||
task: Optional[str] = Field(None, description="Task")
|
|
||||||
img: Optional[str] = Field(None, description="Img")
|
|
||||||
return_history: Optional[bool] = Field(True, description="Return History")
|
|
||||||
rules: Optional[str] = Field(None, description="Rules")
|
|
||||||
schedule: Optional[ScheduleSpec] = Field(None, description="Scheduling information")
|
|
||||||
|
|
||||||
|
|
||||||
class ScheduledJob(Thread):
|
|
||||||
def __init__(
|
|
||||||
self, job_id: str, scheduled_time: datetime, swarm: SwarmSpec, api_key: str
|
|
||||||
):
|
|
||||||
super().__init__()
|
|
||||||
self.job_id = job_id
|
|
||||||
self.scheduled_time = scheduled_time
|
|
||||||
self.swarm = swarm
|
|
||||||
self.api_key = api_key
|
|
||||||
self.daemon = True # Allow the thread to be terminated when main program exits
|
|
||||||
self.cancelled = False
|
|
||||||
|
|
||||||
def run(self):
|
|
||||||
while not self.cancelled:
|
|
||||||
now = datetime.now(pytz.UTC)
|
|
||||||
if now >= self.scheduled_time:
|
|
||||||
try:
|
|
||||||
# Execute the swarm
|
|
||||||
asyncio.run(run_swarm_completion(self.swarm, self.api_key))
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(
|
|
||||||
f"Error executing scheduled swarm {self.job_id}: {str(e)}"
|
|
||||||
)
|
|
||||||
finally:
|
|
||||||
# Remove the job from scheduled_jobs after execution
|
|
||||||
scheduled_jobs.pop(self.job_id, None)
|
|
||||||
break
|
|
||||||
sleep(1) # Check every second
|
|
||||||
|
|
||||||
|
|
||||||
def get_supabase_client():
|
|
||||||
supabase_url = os.getenv("SUPABASE_URL")
|
|
||||||
supabase_key = os.getenv("SUPABASE_KEY")
|
|
||||||
return supabase.create_client(supabase_url, supabase_key)
|
|
||||||
|
|
||||||
|
|
||||||
@lru_cache(maxsize=1000)
|
|
||||||
def check_api_key(api_key: str) -> bool:
|
|
||||||
supabase_client = get_supabase_client()
|
|
||||||
response = (
|
|
||||||
supabase_client.table("swarms_cloud_api_keys")
|
|
||||||
.select("*")
|
|
||||||
.eq("key", api_key)
|
|
||||||
.execute()
|
|
||||||
)
|
|
||||||
return bool(response.data)
|
|
||||||
|
|
||||||
|
|
||||||
# class ExternalAgent:
|
|
||||||
# def __init__(self, base_url: str, parameters: Dict[str, Any], headers: Dict[str, Any]):
|
|
||||||
# self.base_url = base_url
|
|
||||||
# self.parameters = parameters
|
|
||||||
# self.headers = headers
|
|
||||||
|
|
||||||
# def run(self, task: str) -> Dict[str, Any]:
|
|
||||||
# response = requests.post(self.base_url, json=self.parameters, headers=self.headers)
|
|
||||||
# return response.json()
|
|
||||||
|
|
||||||
|
|
||||||
@lru_cache(maxsize=1000)
|
|
||||||
def get_user_id_from_api_key(api_key: str) -> str:
|
|
||||||
"""
|
|
||||||
Maps an API key to its associated user ID.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
api_key (str): The API key to look up
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
str: The user ID associated with the API key
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
ValueError: If the API key is invalid or not found
|
|
||||||
"""
|
|
||||||
supabase_client = get_supabase_client()
|
|
||||||
response = (
|
|
||||||
supabase_client.table("swarms_cloud_api_keys")
|
|
||||||
.select("user_id")
|
|
||||||
.eq("key", api_key)
|
|
||||||
.execute()
|
|
||||||
)
|
|
||||||
if not response.data:
|
|
||||||
raise ValueError("Invalid API key")
|
|
||||||
return response.data[0]["user_id"]
|
|
||||||
|
|
||||||
|
|
||||||
def verify_api_key(x_api_key: str = Header(...)) -> None:
|
|
||||||
"""
|
|
||||||
Dependency to verify the API key.
|
|
||||||
"""
|
|
||||||
if not check_api_key(x_api_key):
|
|
||||||
raise HTTPException(status_code=403, detail="Invalid API Key")
|
|
||||||
|
|
||||||
|
|
||||||
async def get_api_key_logs(api_key: str) -> List[Dict[str, Any]]:
|
|
||||||
"""
|
|
||||||
Retrieve all API request logs for a specific API key.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
api_key: The API key to query logs for
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
List[Dict[str, Any]]: List of log entries for the API key
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
supabase_client = get_supabase_client()
|
|
||||||
|
|
||||||
# Query swarms_api_logs table for entries matching the API key
|
|
||||||
response = (
|
|
||||||
supabase_client.table("swarms_api_logs")
|
|
||||||
.select("*")
|
|
||||||
.eq("api_key", api_key)
|
|
||||||
.execute()
|
|
||||||
)
|
|
||||||
return response.data
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Error retrieving API logs: {str(e)}")
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
||||||
detail=f"Failed to retrieve API logs: {str(e)}",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def create_swarm(swarm_spec: SwarmSpec) -> SwarmRouter:
|
|
||||||
try:
|
|
||||||
# Validate swarm_spec
|
|
||||||
if not swarm_spec.agents:
|
|
||||||
raise ValueError("Swarm specification must include at least one agent.")
|
|
||||||
|
|
||||||
agents = []
|
|
||||||
for agent_spec in swarm_spec.agents:
|
|
||||||
try:
|
|
||||||
# Handle both dict and AgentSpec objects
|
|
||||||
if isinstance(agent_spec, dict):
|
|
||||||
# Convert dict to AgentSpec
|
|
||||||
agent_spec = AgentSpec(**agent_spec)
|
|
||||||
|
|
||||||
# Validate agent_spec fields
|
|
||||||
if not agent_spec.agent_name:
|
|
||||||
raise ValueError("Agent name is required.")
|
|
||||||
if not agent_spec.model_name:
|
|
||||||
raise ValueError("Model name is required.")
|
|
||||||
|
|
||||||
# Create the agent
|
|
||||||
agent = Agent(
|
|
||||||
agent_name=agent_spec.agent_name,
|
|
||||||
description=agent_spec.description,
|
|
||||||
system_prompt=agent_spec.system_prompt,
|
|
||||||
model_name=agent_spec.model_name,
|
|
||||||
auto_generate_prompt=agent_spec.auto_generate_prompt,
|
|
||||||
max_tokens=agent_spec.max_tokens,
|
|
||||||
temperature=agent_spec.temperature,
|
|
||||||
role=agent_spec.role,
|
|
||||||
max_loops=agent_spec.max_loops,
|
|
||||||
)
|
|
||||||
agents.append(agent)
|
|
||||||
logger.info(
|
|
||||||
"Successfully created agent: {}",
|
|
||||||
agent_spec.agent_name,
|
|
||||||
)
|
|
||||||
except ValueError as ve:
|
|
||||||
logger.error(
|
|
||||||
"Validation error for agent {}: {}",
|
|
||||||
getattr(agent_spec, 'agent_name', 'unknown'),
|
|
||||||
str(ve),
|
|
||||||
)
|
|
||||||
raise
|
|
||||||
except Exception as agent_error:
|
|
||||||
logger.error(
|
|
||||||
"Error creating agent {}: {}",
|
|
||||||
getattr(agent_spec, 'agent_name', 'unknown'),
|
|
||||||
str(agent_error),
|
|
||||||
)
|
|
||||||
raise
|
|
||||||
|
|
||||||
if not agents:
|
|
||||||
raise ValueError(
|
|
||||||
"No valid agents could be created from the swarm specification."
|
|
||||||
)
|
|
||||||
|
|
||||||
# Create and configure the swarm
|
|
||||||
swarm = SwarmRouter(
|
|
||||||
name=swarm_spec.name,
|
|
||||||
description=swarm_spec.description,
|
|
||||||
agents=agents,
|
|
||||||
max_loops=swarm_spec.max_loops,
|
|
||||||
swarm_type=swarm_spec.swarm_type,
|
|
||||||
output_type="dict",
|
|
||||||
return_entire_history=False,
|
|
||||||
rules=swarm_spec.rules,
|
|
||||||
rearrange_flow=swarm_spec.rearrange_flow,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Run the swarm task
|
|
||||||
output = swarm.run(task=swarm_spec.task)
|
|
||||||
return output
|
|
||||||
except Exception as e:
|
|
||||||
logger.error("Error creating swarm: {}", str(e))
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
||||||
detail=f"Failed to create swarm: {str(e)}",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# Add this function after your get_supabase_client() function
|
|
||||||
async def log_api_request(api_key: str, data: Dict[str, Any]) -> None:
|
|
||||||
"""
|
|
||||||
Log API request data to Supabase swarms_api_logs table.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
api_key: The API key used for the request
|
|
||||||
data: Dictionary containing request data to log
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
supabase_client = get_supabase_client()
|
|
||||||
|
|
||||||
# Create log entry
|
|
||||||
log_entry = {
|
|
||||||
"api_key": api_key,
|
|
||||||
"data": data,
|
|
||||||
}
|
|
||||||
|
|
||||||
# Insert into swarms_api_logs table
|
|
||||||
response = supabase_client.table("swarms_api_logs").insert(log_entry).execute()
|
|
||||||
|
|
||||||
if not response.data:
|
|
||||||
logger.error("Failed to log API request")
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Error logging API request: {str(e)}")
|
|
||||||
|
|
||||||
|
|
||||||
async def run_swarm_completion(
|
|
||||||
swarm: SwarmSpec, x_api_key: str = None
|
|
||||||
) -> Dict[str, Any]:
|
|
||||||
"""
|
|
||||||
Run a swarm with the specified task.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
swarm_name = swarm.name
|
|
||||||
|
|
||||||
agents = swarm.agents
|
|
||||||
|
|
||||||
await log_api_request(x_api_key, swarm.model_dump())
|
|
||||||
|
|
||||||
# Log start of swarm execution
|
|
||||||
logger.info(f"Starting swarm {swarm_name} with {len(agents)} agents")
|
|
||||||
start_time = time()
|
|
||||||
|
|
||||||
# Create and run the swarm
|
|
||||||
logger.debug(f"Creating swarm object for {swarm_name}")
|
|
||||||
result = create_swarm(swarm)
|
|
||||||
logger.debug(f"Running swarm task: {swarm.task}")
|
|
||||||
|
|
||||||
# Calculate execution time
|
|
||||||
execution_time = time() - start_time
|
|
||||||
logger.info(
|
|
||||||
f"Swarm {swarm_name} executed in {round(execution_time, 2)} seconds"
|
|
||||||
)
|
|
||||||
|
|
||||||
# Calculate costs
|
|
||||||
logger.debug(f"Calculating costs for swarm {swarm_name}")
|
|
||||||
cost_info = calculate_swarm_cost(
|
|
||||||
agents=agents,
|
|
||||||
input_text=swarm.task,
|
|
||||||
agent_outputs=result,
|
|
||||||
execution_time=execution_time,
|
|
||||||
)
|
|
||||||
logger.info(f"Cost calculation completed for swarm {swarm_name}: {cost_info}")
|
|
||||||
|
|
||||||
# Deduct credits based on calculated cost
|
|
||||||
logger.debug(
|
|
||||||
f"Deducting credits for swarm {swarm_name} with cost {cost_info['total_cost']}"
|
|
||||||
)
|
|
||||||
|
|
||||||
deduct_credits(
|
|
||||||
x_api_key,
|
|
||||||
cost_info["total_cost"],
|
|
||||||
f"swarm_execution_{swarm_name}",
|
|
||||||
)
|
|
||||||
|
|
||||||
# Format the response
|
|
||||||
response = {
|
|
||||||
"status": "success",
|
|
||||||
"swarm_name": swarm_name,
|
|
||||||
"description": swarm.description,
|
|
||||||
"swarm_type": swarm.swarm_type,
|
|
||||||
"task": swarm.task,
|
|
||||||
"output": result,
|
|
||||||
"metadata": {
|
|
||||||
"max_loops": swarm.max_loops,
|
|
||||||
"num_agents": len(agents),
|
|
||||||
"execution_time_seconds": round(execution_time, 2),
|
|
||||||
"completion_time": time(),
|
|
||||||
"billing_info": cost_info,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
logger.info(response)
|
|
||||||
await log_api_request(x_api_key, response)
|
|
||||||
|
|
||||||
return response
|
|
||||||
|
|
||||||
except HTTPException as http_exc:
|
|
||||||
logger.error("HTTPException occurred: {}", http_exc.detail)
|
|
||||||
raise
|
|
||||||
except Exception as e:
|
|
||||||
logger.error("Error running swarm {}: {}", swarm_name, str(e))
|
|
||||||
logger.exception(e)
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
||||||
detail=f"Failed to run swarm: {str(e)}",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def deduct_credits(api_key: str, amount: float, product_name: str) -> None:
|
|
||||||
"""
|
|
||||||
Deducts the specified amount of credits for the user identified by api_key,
|
|
||||||
preferring to use free_credit before using regular credit, and logs the transaction.
|
|
||||||
"""
|
|
||||||
supabase_client = get_supabase_client()
|
|
||||||
user_id = get_user_id_from_api_key(api_key)
|
|
||||||
|
|
||||||
# 1. Retrieve the user's credit record
|
|
||||||
response = (
|
|
||||||
supabase_client.table("swarms_cloud_users_credits")
|
|
||||||
.select("*")
|
|
||||||
.eq("user_id", user_id)
|
|
||||||
.execute()
|
|
||||||
)
|
|
||||||
if not response.data:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_404_NOT_FOUND,
|
|
||||||
detail="User credits record not found.",
|
|
||||||
)
|
|
||||||
|
|
||||||
record = response.data[0]
|
|
||||||
# Use Decimal for precise arithmetic
|
|
||||||
available_credit = Decimal(record["credit"])
|
|
||||||
free_credit = Decimal(record.get("free_credit", "0"))
|
|
||||||
deduction = Decimal(str(amount))
|
|
||||||
|
|
||||||
print(
|
|
||||||
f"Available credit: {available_credit}, Free credit: {free_credit}, Deduction: {deduction}"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 2. Verify sufficient total credits are available
|
|
||||||
if (available_credit + free_credit) < deduction:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
|
||||||
detail="Insufficient credits.",
|
|
||||||
)
|
|
||||||
|
|
||||||
# 3. Log the transaction
|
|
||||||
log_response = (
|
|
||||||
supabase_client.table("swarms_cloud_services")
|
|
||||||
.insert(
|
|
||||||
{
|
|
||||||
"user_id": user_id,
|
|
||||||
"api_key": api_key,
|
|
||||||
"charge_credit": int(
|
|
||||||
deduction
|
|
||||||
), # Assuming credits are stored as integers
|
|
||||||
"product_name": product_name,
|
|
||||||
}
|
|
||||||
)
|
|
||||||
.execute()
|
|
||||||
)
|
|
||||||
if not log_response.data:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
||||||
detail="Failed to log the credit transaction.",
|
|
||||||
)
|
|
||||||
|
|
||||||
# 4. Deduct credits: use free_credit first, then deduct the remainder from available_credit
|
|
||||||
if free_credit >= deduction:
|
|
||||||
free_credit -= deduction
|
|
||||||
else:
|
|
||||||
remainder = deduction - free_credit
|
|
||||||
free_credit = Decimal("0")
|
|
||||||
available_credit -= remainder
|
|
||||||
|
|
||||||
update_response = (
|
|
||||||
supabase_client.table("swarms_cloud_users_credits")
|
|
||||||
.update(
|
|
||||||
{
|
|
||||||
"credit": str(available_credit),
|
|
||||||
"free_credit": str(free_credit),
|
|
||||||
}
|
|
||||||
)
|
|
||||||
.eq("user_id", user_id)
|
|
||||||
.execute()
|
|
||||||
)
|
|
||||||
if not update_response.data:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
||||||
detail="Failed to update credits.",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def calculate_swarm_cost(
|
|
||||||
agents: List[Agent],
|
|
||||||
input_text: str,
|
|
||||||
execution_time: float,
|
|
||||||
agent_outputs: Union[List[Dict[str, str]], str] = None, # Update agent_outputs type
|
|
||||||
) -> Dict[str, Any]:
|
|
||||||
"""
|
|
||||||
Calculate the cost of running a swarm based on agents, tokens, and execution time.
|
|
||||||
Includes system prompts, agent memory, and scaled output costs.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
agents: List of agents used in the swarm
|
|
||||||
input_text: The input task/prompt text
|
|
||||||
execution_time: Time taken to execute in seconds
|
|
||||||
agent_outputs: List of output texts from each agent or a list of dictionaries
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Dict containing cost breakdown and total cost
|
|
||||||
"""
|
|
||||||
# Base costs per unit (these could be moved to environment variables)
|
|
||||||
COST_PER_AGENT = 0.01 # Base cost per agent
|
|
||||||
COST_PER_1M_INPUT_TOKENS = 2.00 # Cost per 1M input tokens
|
|
||||||
COST_PER_1M_OUTPUT_TOKENS = 6.00 # Cost per 1M output tokens
|
|
||||||
|
|
||||||
# Get current time in California timezone
|
|
||||||
california_tz = pytz.timezone("America/Los_Angeles")
|
|
||||||
current_time = datetime.now(california_tz)
|
|
||||||
is_night_time = current_time.hour >= 20 or current_time.hour < 6 # 8 PM to 6 AM
|
|
||||||
|
|
||||||
try:
|
|
||||||
# Calculate input tokens for task
|
|
||||||
task_tokens = count_tokens(input_text)
|
|
||||||
|
|
||||||
# Calculate total input tokens including system prompts and memory for each agent
|
|
||||||
total_input_tokens = 0
|
|
||||||
total_output_tokens = 0
|
|
||||||
per_agent_tokens = {}
|
|
||||||
|
|
||||||
for i, agent in enumerate(agents):
|
|
||||||
agent_input_tokens = task_tokens # Base task tokens
|
|
||||||
|
|
||||||
# Add system prompt tokens if present
|
|
||||||
if agent.system_prompt:
|
|
||||||
agent_input_tokens += count_tokens(agent.system_prompt)
|
|
||||||
|
|
||||||
# Add memory tokens if available
|
|
||||||
try:
|
|
||||||
memory = agent.short_memory.return_history_as_string()
|
|
||||||
if memory:
|
|
||||||
memory_tokens = count_tokens(str(memory))
|
|
||||||
agent_input_tokens += memory_tokens
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning(
|
|
||||||
f"Could not get memory for agent {agent.agent_name}: {str(e)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
# Calculate actual output tokens if available, otherwise estimate
|
|
||||||
if agent_outputs:
|
|
||||||
if isinstance(agent_outputs, list):
|
|
||||||
# Sum tokens for each dictionary's content
|
|
||||||
agent_output_tokens = sum(
|
|
||||||
count_tokens(message["content"]) for message in agent_outputs
|
|
||||||
)
|
|
||||||
elif isinstance(agent_outputs, str):
|
|
||||||
agent_output_tokens = count_tokens(agent_outputs)
|
|
||||||
else:
|
|
||||||
agent_output_tokens = int(
|
|
||||||
agent_input_tokens * 2.5
|
|
||||||
) # Estimated output tokens
|
|
||||||
else:
|
|
||||||
agent_output_tokens = int(
|
|
||||||
agent_input_tokens * 2.5
|
|
||||||
) # Estimated output tokens
|
|
||||||
|
|
||||||
# Store per-agent token counts
|
|
||||||
per_agent_tokens[agent.agent_name] = {
|
|
||||||
"input_tokens": agent_input_tokens,
|
|
||||||
"output_tokens": agent_output_tokens,
|
|
||||||
"total_tokens": agent_input_tokens + agent_output_tokens,
|
|
||||||
}
|
|
||||||
|
|
||||||
# Add to totals
|
|
||||||
total_input_tokens += agent_input_tokens
|
|
||||||
total_output_tokens += agent_output_tokens
|
|
||||||
|
|
||||||
# Calculate costs (convert to millions of tokens)
|
|
||||||
agent_cost = len(agents) * COST_PER_AGENT
|
|
||||||
input_token_cost = (
|
|
||||||
(total_input_tokens / 1_000_000) * COST_PER_1M_INPUT_TOKENS * len(agents)
|
|
||||||
)
|
|
||||||
output_token_cost = (
|
|
||||||
(total_output_tokens / 1_000_000) * COST_PER_1M_OUTPUT_TOKENS * len(agents)
|
|
||||||
)
|
|
||||||
|
|
||||||
# Apply discount during California night time hours
|
|
||||||
if is_night_time:
|
|
||||||
input_token_cost *= 0.25 # 75% discount
|
|
||||||
output_token_cost *= 0.25 # 75% discount
|
|
||||||
|
|
||||||
# Calculate total cost
|
|
||||||
total_cost = agent_cost + input_token_cost + output_token_cost
|
|
||||||
|
|
||||||
output = {
|
|
||||||
"cost_breakdown": {
|
|
||||||
"agent_cost": round(agent_cost, 6),
|
|
||||||
"input_token_cost": round(input_token_cost, 6),
|
|
||||||
"output_token_cost": round(output_token_cost, 6),
|
|
||||||
"token_counts": {
|
|
||||||
"total_input_tokens": total_input_tokens,
|
|
||||||
"total_output_tokens": total_output_tokens,
|
|
||||||
"total_tokens": total_input_tokens + total_output_tokens,
|
|
||||||
"per_agent": per_agent_tokens,
|
|
||||||
},
|
|
||||||
"num_agents": len(agents),
|
|
||||||
"execution_time_seconds": round(execution_time, 2),
|
|
||||||
},
|
|
||||||
"total_cost": round(total_cost, 6),
|
|
||||||
}
|
|
||||||
|
|
||||||
return output
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Error calculating swarm cost: {str(e)}")
|
|
||||||
raise ValueError(f"Failed to calculate swarm cost: {str(e)}")
|
|
||||||
|
|
||||||
|
|
||||||
# --- FastAPI Application Setup ---
|
|
||||||
|
|
||||||
app = FastAPI(
|
|
||||||
title="Swarm Agent API",
|
|
||||||
description="API for managing and executing Python agents in the cloud without Docker/Kubernetes.",
|
|
||||||
version="1.0.0",
|
|
||||||
debug=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Enable CORS (adjust origins as needed)
|
|
||||||
app.add_middleware(
|
|
||||||
CORSMiddleware,
|
|
||||||
allow_origins=["*"], # In production, restrict this to specific domains
|
|
||||||
allow_credentials=True,
|
|
||||||
allow_methods=["*"],
|
|
||||||
allow_headers=["*"],
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@app.get("/", dependencies=[Depends(rate_limiter)])
|
|
||||||
def root():
|
|
||||||
return {
|
|
||||||
"status": "Welcome to the Swarm API. Check out the docs at https://docs.swarms.world"
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
@app.get("/health", dependencies=[Depends(rate_limiter)])
|
|
||||||
def health():
|
|
||||||
return {"status": "ok"}
|
|
||||||
|
|
||||||
|
|
||||||
@app.post(
|
|
||||||
"/v1/swarm/completions",
|
|
||||||
dependencies=[
|
|
||||||
Depends(verify_api_key),
|
|
||||||
Depends(rate_limiter),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
async def run_swarm(swarm: SwarmSpec, x_api_key=Header(...)) -> Dict[str, Any]:
|
|
||||||
"""
|
|
||||||
Run a swarm with the specified task.
|
|
||||||
"""
|
|
||||||
return await run_swarm_completion(swarm, x_api_key)
|
|
||||||
|
|
||||||
|
|
||||||
@app.post(
|
|
||||||
"/v1/swarm/batch/completions",
|
|
||||||
dependencies=[
|
|
||||||
Depends(verify_api_key),
|
|
||||||
Depends(rate_limiter),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
async def run_batch_completions(
|
|
||||||
swarms: List[SwarmSpec], x_api_key=Header(...)
|
|
||||||
) -> List[Dict[str, Any]]:
|
|
||||||
"""
|
|
||||||
Run a batch of swarms with the specified tasks.
|
|
||||||
"""
|
|
||||||
results = []
|
|
||||||
for swarm in swarms:
|
|
||||||
try:
|
|
||||||
# Call the existing run_swarm function for each swarm
|
|
||||||
result = await run_swarm_completion(swarm, x_api_key)
|
|
||||||
results.append(result)
|
|
||||||
except HTTPException as http_exc:
|
|
||||||
logger.error("HTTPException occurred: {}", http_exc.detail)
|
|
||||||
results.append(
|
|
||||||
{
|
|
||||||
"status": "error",
|
|
||||||
"swarm_name": swarm.name,
|
|
||||||
"detail": http_exc.detail,
|
|
||||||
}
|
|
||||||
)
|
|
||||||
except Exception as e:
|
|
||||||
logger.error("Error running swarm {}: {}", swarm.name, str(e))
|
|
||||||
logger.exception(e)
|
|
||||||
results.append(
|
|
||||||
{
|
|
||||||
"status": "error",
|
|
||||||
"swarm_name": swarm.name,
|
|
||||||
"detail": f"Failed to run swarm: {str(e)}",
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
return results
|
|
||||||
|
|
||||||
|
|
||||||
# Add this new endpoint
|
|
||||||
@app.get(
|
|
||||||
"/v1/swarm/logs",
|
|
||||||
dependencies=[
|
|
||||||
Depends(verify_api_key),
|
|
||||||
Depends(rate_limiter),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
async def get_logs(x_api_key: str = Header(...)) -> Dict[str, Any]:
|
|
||||||
"""
|
|
||||||
Get all API request logs for the provided API key.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
logs = await get_api_key_logs(x_api_key)
|
|
||||||
return {"status": "success", "count": len(logs), "logs": logs}
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Error in get_logs endpoint: {str(e)}")
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# @app.post("/v1/swarm/cost-prediction")
|
|
||||||
# async def cost_prediction(swarm: SwarmSpec) -> Dict[str, Any]:
|
|
||||||
# """
|
|
||||||
# Predict the cost of running a swarm.
|
|
||||||
# """
|
|
||||||
# return {"status": "success", "cost": calculate_swarm_cost(swarm)})
|
|
||||||
|
|
||||||
|
|
||||||
@app.post(
|
|
||||||
"/v1/swarm/schedule",
|
|
||||||
dependencies=[
|
|
||||||
Depends(verify_api_key),
|
|
||||||
Depends(rate_limiter),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
async def schedule_swarm(
|
|
||||||
swarm: SwarmSpec, x_api_key: str = Header(...)
|
|
||||||
) -> Dict[str, Any]:
|
|
||||||
"""
|
|
||||||
Schedule a swarm to run at a specific time.
|
|
||||||
"""
|
|
||||||
if not swarm.schedule:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_400_BAD_REQUEST,
|
|
||||||
detail="Schedule information is required",
|
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
|
||||||
# Generate a unique job ID
|
|
||||||
job_id = f"swarm_{swarm.name}_{int(time())}"
|
|
||||||
|
|
||||||
# Create and start the scheduled job
|
|
||||||
job = ScheduledJob(
|
|
||||||
job_id=job_id,
|
|
||||||
scheduled_time=swarm.schedule.scheduled_time,
|
|
||||||
swarm=swarm,
|
|
||||||
api_key=x_api_key,
|
|
||||||
)
|
|
||||||
job.start()
|
|
||||||
|
|
||||||
# Store the job information
|
|
||||||
scheduled_jobs[job_id] = {
|
|
||||||
"job": job,
|
|
||||||
"swarm_name": swarm.name,
|
|
||||||
"scheduled_time": swarm.schedule.scheduled_time,
|
|
||||||
"timezone": swarm.schedule.timezone,
|
|
||||||
}
|
|
||||||
|
|
||||||
# Log the scheduling
|
|
||||||
await log_api_request(
|
|
||||||
x_api_key,
|
|
||||||
{
|
|
||||||
"action": "schedule_swarm",
|
|
||||||
"swarm_name": swarm.name,
|
|
||||||
"scheduled_time": swarm.schedule.scheduled_time.isoformat(),
|
|
||||||
"job_id": job_id,
|
|
||||||
},
|
|
||||||
)
|
|
||||||
|
|
||||||
return {
|
|
||||||
"status": "success",
|
|
||||||
"message": "Swarm scheduled successfully",
|
|
||||||
"job_id": job_id,
|
|
||||||
"scheduled_time": swarm.schedule.scheduled_time,
|
|
||||||
"timezone": swarm.schedule.timezone,
|
|
||||||
}
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Error scheduling swarm: {str(e)}")
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
||||||
detail=f"Failed to schedule swarm: {str(e)}",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@app.get(
|
|
||||||
"/v1/swarm/schedule",
|
|
||||||
dependencies=[
|
|
||||||
Depends(verify_api_key),
|
|
||||||
Depends(rate_limiter),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
async def get_scheduled_jobs(x_api_key: str = Header(...)) -> Dict[str, Any]:
|
|
||||||
"""
|
|
||||||
Get all scheduled swarm jobs.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
jobs_list = []
|
|
||||||
current_time = datetime.now(pytz.UTC)
|
|
||||||
|
|
||||||
# Clean up completed jobs
|
|
||||||
completed_jobs = [
|
|
||||||
job_id
|
|
||||||
for job_id, job_info in scheduled_jobs.items()
|
|
||||||
if current_time >= job_info["scheduled_time"]
|
|
||||||
]
|
|
||||||
for job_id in completed_jobs:
|
|
||||||
scheduled_jobs.pop(job_id, None)
|
|
||||||
|
|
||||||
# Get active jobs
|
|
||||||
for job_id, job_info in scheduled_jobs.items():
|
|
||||||
jobs_list.append(
|
|
||||||
{
|
|
||||||
"job_id": job_id,
|
|
||||||
"swarm_name": job_info["swarm_name"],
|
|
||||||
"scheduled_time": job_info["scheduled_time"].isoformat(),
|
|
||||||
"timezone": job_info["timezone"],
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
return {"status": "success", "scheduled_jobs": jobs_list}
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Error retrieving scheduled jobs: {str(e)}")
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
||||||
detail=f"Failed to retrieve scheduled jobs: {str(e)}",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@app.delete(
|
|
||||||
"/v1/swarm/schedule/{job_id}",
|
|
||||||
dependencies=[
|
|
||||||
Depends(verify_api_key),
|
|
||||||
Depends(rate_limiter),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
async def cancel_scheduled_job(
|
|
||||||
job_id: str, x_api_key: str = Header(...)
|
|
||||||
) -> Dict[str, Any]:
|
|
||||||
"""
|
|
||||||
Cancel a scheduled swarm job.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
if job_id not in scheduled_jobs:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_404_NOT_FOUND, detail="Scheduled job not found"
|
|
||||||
)
|
|
||||||
|
|
||||||
# Cancel and remove the job
|
|
||||||
job_info = scheduled_jobs[job_id]
|
|
||||||
job_info["job"].cancelled = True
|
|
||||||
scheduled_jobs.pop(job_id)
|
|
||||||
|
|
||||||
await log_api_request(
|
|
||||||
x_api_key, {"action": "cancel_scheduled_job", "job_id": job_id}
|
|
||||||
)
|
|
||||||
|
|
||||||
return {
|
|
||||||
"status": "success",
|
|
||||||
"message": "Scheduled job cancelled successfully",
|
|
||||||
"job_id": job_id,
|
|
||||||
}
|
|
||||||
|
|
||||||
except HTTPException:
|
|
||||||
raise
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Error cancelling scheduled job: {str(e)}")
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
||||||
detail=f"Failed to cancel scheduled job: {str(e)}",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# --- Main Entrypoint ---
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
import uvicorn
|
|
||||||
|
|
||||||
uvicorn.run(app, host="0.0.0.0", port=8080, workers=os.cpu_count())
|
|
||||||
Loading…
Reference in new issue