api/utils.py

from fastapi import HTTPException
from api.config import MODEL_PROVIDER_MAPPING
from api.provider import amigochat
from api.logger import setup_logger

logger = setup_logger(__name__)

async def process_streaming_response(request_data):
    provider_name = MODEL_PROVIDER_MAPPING.get(request_data.get('model'))
    if provider_name == 'amigochat':
        return await amigochat.process_streaming_response(request_data)
    else:
        raise HTTPException(status_code=400, detail=f"Model {request_data.get('model')} is not supported for streaming.")

async def process_non_streaming_response(request_data):
    provider_name = MODEL_PROVIDER_MAPPING.get(request_data.get('model'))
    if provider_name == 'amigochat':
        return await amigochat.process_non_streaming_response(request_data)
    else:
        raise HTTPException(status_code=400, detail=f"Model {request_data.get('model')} is not supported.")