pull/25/head
Kye 1 year ago
parent 4b50568bec
commit 4a5d3546e9

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

Loading…
Cancel
Save