From 2ebfa0f36a947bc80fc0930a689a148a37a2a1f0 Mon Sep 17 00:00:00 2001 From: Kye Date: Tue, 18 Jul 2023 16:49:10 -0400 Subject: [PATCH] clean up fastapi --- api/app.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/api/app.py b/api/app.py index 0158dc88..c4afcf24 100644 --- a/api/app.py +++ b/api/app.py @@ -1,23 +1,34 @@ -from fastapi import FastAPI, HTTPException -from swarms.swarms import swarm +from fastapi import FastAPI, HTTPException, Depends +from fastapi_cache.decorator import cache +from fastapi_cache.coder import JsonCoder +from fastapi_cache import FastAPICache +from aioredis import Redis from pydantic import BaseModel - -app = FastAPI() +from swarms.swarms import swarm +from fastapi_limiter import FastAPILimiter +from fastapi_limiter.depends import RateLimiter class SwarmInput(BaseModel): - api_key: set + api_key: str objective: str +app = FastAPI() + +@app.on_event("startup") +async def startup(): + redis = Redis(host="localhost", port=6379) + FastAPICache.init(RedisBackend(redis), prefix="fastapi-cache", coder=JsonCoder()) + await FastAPILimiter.init("redis://localhost:6379") -@app.post("/chat") +@app.post("/chat", dependencies=[Depends(RateLimiter(times=2, minutes=1))]) +@cache(expire=60) # Cache results for 1 minute async def run_swarms(swarm_input: SwarmInput): try: results = swarm(swarm_input.api_key, swarm_input.objective) if not results: - raise HTTPException(status_code=500, detaile="Failed to run swarms") + raise HTTPException(status_code=500, detail="Failed to run swarms") return {"results": results} except ValueError as ve: raise HTTPException(status_code=400, detail=str(ve)) except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) - + raise HTTPException(status_code=500, detail=str(e)) \ No newline at end of file