import os
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from loguru import logger
from swarms import Agent, count_tokens
from atp.middleware import ATPSettlementMiddleware
from atp.schemas import PaymentToken
AGENT_MODEL = "gpt-4o-mini"
app = FastAPI(
title="ATP Protocol + Swarms Integration",
description="Example API showing ATP payment processing with Swarms agents",
)
# Read config from env
RECIPIENT_PUBKEY = os.getenv("RECIPIENT_PUBKEY", "YourSolanaWalletHere")
ATP_SETTLEMENT_URL = os.getenv("ATP_SETTLEMENT_URL", "https://facilitator.swarms.world")
agent = Agent(
model_name=AGENT_MODEL,
max_loops=1,
interactive=False,
)
app.add_middleware(
ATPSettlementMiddleware,
allowed_endpoints=["/v1/agent/execute"],
input_cost_per_million_usd=10.0,
output_cost_per_million_usd=30.0,
recipient_pubkey=RECIPIENT_PUBKEY,
payment_token=PaymentToken.SOL,
wallet_private_key_header="x-wallet-private-key",
require_wallet=True,
settlement_service_url=ATP_SETTLEMENT_URL,
)
@app.post("/v1/agent/execute")
async def execute_agent(request: dict):
task = request.get("task", "")
if not task:
raise HTTPException(status_code=400, detail="Task is required")
input_tokens = count_tokens(task, model=AGENT_MODEL)
logger.info(f"Executing agent task: {task[:100]}...")
result = agent.run(task)
output_tokens = count_tokens(str(result), model=AGENT_MODEL)
return JSONResponse(
content={
"output": result,
"task": task,
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
},
}
)