94 lines
2.6 KiB
Python
94 lines
2.6 KiB
Python
|
# api_server.py
|
||
|
|
import os
|
||
|
|
from fastapi import FastAPI, HTTPException
|
||
|
|
from fastapi.middleware.cors import CORSMiddleware
|
||
|
|
from pydantic import BaseModel
|
||
|
|
import uvicorn
|
||
|
|
import uuid
|
||
|
|
from typing import Dict, Optional
|
||
|
|
import gc
|
||
|
|
import psutil
|
||
|
|
|
||
|
|
from agents.compass import compass
|
||
|
|
|
||
|
|
app = FastAPI()
|
||
|
|
|
||
|
|
# Add CORS middleware
|
||
|
|
app.add_middleware(
|
||
|
|
CORSMiddleware,
|
||
|
|
allow_origins=["*"], # This allows all origins
|
||
|
|
allow_credentials=True,
|
||
|
|
allow_methods=["*"],
|
||
|
|
allow_headers=["*"],
|
||
|
|
)
|
||
|
|
|
||
|
|
@app.middleware("http")
|
||
|
|
async def check_memory_usage(request, call_next):
|
||
|
|
# Force garbage collection before processing request
|
||
|
|
gc.collect()
|
||
|
|
|
||
|
|
# Get current memory usage
|
||
|
|
process = psutil.Process(os.getpid())
|
||
|
|
mem_before = process.memory_info().rss / 1024 / 1024 # Memory in MB
|
||
|
|
|
||
|
|
# If memory usage is too high, refuse new requests
|
||
|
|
if mem_before > 900: # 900MB threshold
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=503,
|
||
|
|
detail="Server is currently overloaded. Please try again later."
|
||
|
|
)
|
||
|
|
|
||
|
|
response = await call_next(request)
|
||
|
|
return response
|
||
|
|
|
||
|
|
class ChatRequest(BaseModel):
|
||
|
|
message: str
|
||
|
|
session_id: Optional[str] = None
|
||
|
|
|
||
|
|
class ChatResponse(BaseModel):
|
||
|
|
response: str
|
||
|
|
session_id: str
|
||
|
|
|
||
|
|
# Store active sessions
|
||
|
|
active_sessions: Dict[str, dict] = {}
|
||
|
|
|
||
|
|
@app.post("/api/chat", response_model=ChatResponse)
|
||
|
|
async def chat_endpoint(request: ChatRequest):
|
||
|
|
try:
|
||
|
|
# Generate session_id if not provided
|
||
|
|
session_id = request.session_id or str(uuid.uuid4())
|
||
|
|
|
||
|
|
# Initialize session if it doesn't exist
|
||
|
|
if session_id not in active_sessions:
|
||
|
|
active_sessions[session_id] = {
|
||
|
|
"history": []
|
||
|
|
}
|
||
|
|
|
||
|
|
# Process message through Compass agent
|
||
|
|
response = compass(session_id, request.message)
|
||
|
|
# print(response)
|
||
|
|
|
||
|
|
# Update session history
|
||
|
|
active_sessions[session_id]["history"].append({
|
||
|
|
"role": "user",
|
||
|
|
"content": request.message
|
||
|
|
})
|
||
|
|
active_sessions[session_id]["history"].append({
|
||
|
|
"role": "assistant",
|
||
|
|
"content": response
|
||
|
|
})
|
||
|
|
|
||
|
|
return ChatResponse(
|
||
|
|
response=response,
|
||
|
|
session_id=session_id
|
||
|
|
)
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
raise HTTPException(status_code=500, detail=str(e))
|
||
|
|
|
||
|
|
@app.get("/api/health")
|
||
|
|
async def health_check():
|
||
|
|
return {"status": "healthy"}
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
uvicorn.run(app, port=8000)
|