You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
swarms/api/app.py

48 lines
1.6 KiB

1 year ago
import logging
import os
2 years ago
from fastapi import FastAPI, HTTPException, Depends
from fastapi_cache.decorator import cache
from fastapi_cache.coder import JsonCoder
1 year ago
2 years ago
from fastapi_cache import FastAPICache
1 year ago
from fastapi_cache.backends.redis import RedisBackend
2 years ago
from aioredis import Redis
1 year ago
2 years ago
from pydantic import BaseModel
2 years ago
from swarms.swarms import swarm
from fastapi_limiter import FastAPILimiter
1 year ago
2 years ago
from fastapi_limiter.depends import RateLimiter
1 year ago
from dotenv import load_dotenv
load_dotenv()
2 years ago
class SwarmInput(BaseModel):
2 years ago
api_key: str
2 years ago
objective: str
2 years ago
app = FastAPI()
@app.on_event("startup")
async def startup():
1 year ago
redis_host = os.getenv("REDIS_HOST", "localhost")
redis_port = int(os.getenv("REDIS_PORT", 6379))
redis = await Redis.create(redis_host, redis_port)
2 years ago
FastAPICache.init(RedisBackend(redis), prefix="fastapi-cache", coder=JsonCoder())
1 year ago
await FastAPILimiter.init(f"redis://{redis_host}:{redis_port}")
2 years ago
2 years ago
@app.post("/chat", dependencies=[Depends(RateLimiter(times=2, minutes=1))])
@cache(expire=60) # Cache results for 1 minute
1 year ago
async def run(swarm_input: SwarmInput):
2 years ago
try:
results = swarm(swarm_input.api_key, swarm_input.objective)
if not results:
2 years ago
raise HTTPException(status_code=500, detail="Failed to run swarms")
2 years ago
return {"results": results}
except ValueError as ve:
1 year ago
logging.error("A ValueError occurred", exc_info=True)
2 years ago
raise HTTPException(status_code=400, detail=str(ve))
1 year ago
except Exception:
1 year ago
logging.error("An error occurred", exc_info=True)
raise HTTPException(status_code=500, detail="An unexpected error occurred")