Update api/routes.py

api/routes.py

@@ -1,9 +1,12 @@
+# api/routes.py
+
 import json
 from fastapi import APIRouter, Depends, HTTPException, Request, Response
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.responses import StreamingResponse
 from api.auth import verify_app_secret
-from api.config import ALLOWED_MODELS
+from api.models import ChatRequest
 from api.utils import process_non_streaming_response, process_streaming_response
+from api.providers import AmigoChat
 from api.logger import setup_logger
 
 logger = setup_logger(__name__)
@@ -25,42 +28,36 @@ async def chat_completions_options():
 @router.get("/v1/models")
 @router.get("/api/v1/models")
 async def list_models():
-    return {"object": "list", "data": ALLOWED_MODELS}
+    allowed_models = [{"id": model, "name": model} for model in AmigoChat.models]
+    return {"object": "list", "data": allowed_models}
 
 @router.post("/v1/chat/completions")
 @router.post("/api/v1/chat/completions")
 async def chat_completions(
-    request: Request, app_secret: str = Depends(verify_app_secret)
+    request: ChatRequest, app_secret: str = Depends(verify_app_secret)
 ):
     logger.info("Entering chat_completions route")
-    request_data = await request.json()
-    model = request_data.get('model')
-    if not model:
-        raise HTTPException(status_code=400, detail="Model is required")
-    logger.info(f"Processing chat completion request for model: {model}")
+    logger.info(f"Processing chat completion request for model: {request.model}")
 
-    if model not in [model_dict["id"] for model_dict in ALLOWED_MODELS]:
+    if request.model not in AmigoChat.models and request.model not in AmigoChat.model_aliases:
         raise HTTPException(
             status_code=400,
-            detail=f"Model {model} is not allowed. Allowed models are: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
+            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(AmigoChat.models)}",
         )
 
-    stream = request_data.get('stream', False)
-
-    if stream:
+    if request.stream:
         logger.info("Streaming response")
-        response = await process_streaming_response(request_data)
-        return StreamingResponse(response, media_type="text/event-stream")
+        return StreamingResponse(process_streaming_response(request), media_type="text/event-stream")
     else:
         logger.info("Non-streaming response")
-        response = await process_non_streaming_response(request_data)
-        return JSONResponse(content=response)
-
-@router.get('/')
-@router.get('/healthz')
-@router.get('/ready')
-@router.get('/alive')
-@router.get('/status')
+        response = await process_non_streaming_response(request)
+        return Response(content=json.dumps(response), media_type="application/json")
+
+@router.route('/')
+@router.route('/healthz')
+@router.route('/ready')
+@router.route('/alive')
+@router.route('/status')
 @router.get("/health")
 def health_check(request: Request):
-    return JSONResponse(content={"status": "ok"})
+    return Response(content=json.dumps({"status": "ok"}), media_type="application/json")