core/router.py

import json
import os

from fastapi import APIRouter, Response, Request, Depends, HTTPException
from fastapi.responses import StreamingResponse
from core.auth import verify_app_secret
from core.config import get_settings
from core.logger import setup_logger
from core.models import ChatRequest
from core.utils import process_streaming_response

logger = setup_logger(__name__)
router = APIRouter()
ALLOWED_MODELS = get_settings().ALLOWED_MODELS

@router.post('/test')
async def login(request: Request):
    logger.info(f"body is {json.dumps(await request.json())}")
    logger.info(f"query is {request.query_params}")
    logger.info(f"Request headers: {json.dumps(dict(request.headers), indent=2)}")
    return {
        'status': 'success',
    }
@router.get("/models")
async def list_models():
    return {"object": "list", "data": ALLOWED_MODELS, "success": True}

@router.options("/chat/completions")
async def chat_completions_options():
    return Response(
        status_code=200,
        headers={
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "POST, OPTIONS",
            "Access-Control-Allow-Headers": "Content-Type, Authorization",
        },
    )

@router.post("/chat/completions")
async def chat_completions(
        request: ChatRequest, app_secret: str = Depends(verify_app_secret)
):
    logger.info("Entering chat_completions route")
    # logger.info(f"Received request: {request}")
    logger.info(f"Received request json format: {json.dumps(request.dict(), indent=4)}")
    logger.info(f"App secret: {app_secret}")
    logger.info(f"Received chat completion request for model: {request.model}")

    if request.model not in [model["id"] for model in ALLOWED_MODELS]:
        raise HTTPException(
            status_code=400,
            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
        )

    if request.stream:
        logger.info("Streaming response")
        return StreamingResponse(
            process_streaming_response(request, app_secret),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "Transfer-Encoding": "chunked"
            }
        )
    else:
        logger.info("Non-streaming response")
        # return await process_non_streaming_response(request)

@router.route('/')
@router.route('/healthz')
@router.route('/ready')
@router.route('/alive')
@router.route('/status')
@router.get("/health")
async def health_check(request: Request):
    return Response(content=json.dumps({"status": "ok"}), media_type="application/json")

@router.post("/env")
async def environment(app_secret: str = Depends(verify_app_secret)):
    return Response(content=json.dumps({"token": os.getenv("TOKEN", ""), "refresh_token": os.getenv("REFRESH_TOKEN", ""), "key": os.getenv("FIREBASE_API_KEY", "")}), media_type="application/json")