Update api/utils.py

api/utils.py

@@ -1,229 +1,22 @@
-# api/utils.py
-
-from datetime import datetime
-import json
-from typing import AsyncGenerator, Union
-
-import aiohttp
 from fastapi import HTTPException
-from api.config import GIZAI_API_ENDPOINT, GIZAI_BASE_URL
-from api.models import ChatRequest, ImageResponseModel, ChatCompletionResponse
-from api.logger import setup_logger
-
-logger = setup_logger(__name__)
-
-class GizAI:
-    # Chat models
-    default_model = 'chat-gemini-flash'
-    chat_models = [
-        default_model,
-        'chat-gemini-pro',
-        'chat-gpt4m',
-        'chat-gpt4',
-        'claude-sonnet',
-        'claude-haiku',
-        'llama-3-70b',
-        'llama-3-8b',
-        'mistral-large',
-        'chat-o1-mini'
-    ]
-
-    # Image models
-    image_models = [
-        'flux1',
-        'sdxl',
-        'sd',
-        'sd35',
-    ]
-
-    models = [*chat_models, *image_models]
-    
-    model_aliases = {
-        # Chat model aliases
-        "gemini-flash": "chat-gemini-flash",
-        "gemini-pro": "chat-gemini-pro",
-        "gpt-4o-mini": "chat-gpt4m",
-        "gpt-4o": "chat-gpt4",
-        "claude-3.5-sonnet": "claude-sonnet",
-        "claude-3-haiku": "claude-haiku",
-        "llama-3.1-70b": "llama-3-70b",
-        "llama-3.1-8b": "llama-3-8b",
-        "o1-mini": "chat-o1-mini",
-        # Image model aliases
-        "sd-1.5": "sd",
-        "sd-3.5": "sd35",
-        "flux-schnell": "flux1",
-    }
-
-    @classmethod
-    def get_model(cls, model: str) -> str:
-        if model in cls.models:
-            return model
-        elif model in cls.model_aliases:
-            return cls.model_aliases[model]
-        else:
-            return cls.default_model
-
-    @classmethod
-    def is_image_model(cls, model: str) -> bool:
-        return model in cls.image_models
-
-async def process_gizai_stream_response(request: ChatRequest, model: str) -> AsyncGenerator[str, None]:
-    async with aiohttp.ClientSession() as session:
-        # Set up headers
-        headers = {
-            'Accept': 'application/json, text/plain, */*',
-            'Accept-Language': 'en-US,en;q=0.9',
-            'Cache-Control': 'no-cache',
-            'Connection': 'keep-alive',
-            'Content-Type': 'application/json',
-            'Origin': 'https://app.giz.ai',
-            'Pragma': 'no-cache',
-            'Sec-Fetch-Dest': 'empty',
-            'Sec-Fetch-Mode': 'cors',
-            'Sec-Fetch-Site': 'same-origin',
-            'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36',
-            'sec-ch-ua': '"Not?A_Brand";v="99", "Chromium";v="130"',
-            'sec-ch-ua-mobile': '?0',
-            'sec-ch-ua-platform': '"Linux"'
-        }
-
-        if GizAI.is_image_model(model):
-            # Image generation logic (streaming might not make sense here)
-            raise HTTPException(status_code=400, detail="Image generation does not support streaming.")
-        else:
-            # Chat completion logic
-            messages_formatted = [
-                {
-                    "type": "human",
-                    "content": msg.content if isinstance(msg.content, str) else msg.content[0].get("text", "")
-                } for msg in request.messages
-            ]
-            data = {
-                "model": model,
-                "input": {
-                    "messages": messages_formatted,
-                    "mode": "plan"
-                },
-                "noStream": False  # Enable streaming
-            }
-            try:
-                async with session.post(
-                    GIZAI_API_ENDPOINT,
-                    headers=headers,
-                    json=data
-                ) as response:
-                    response.raise_for_status()
-                    async for line in response.content:
-                        if line:
-                            decoded_line = line.decode('utf-8').strip()
-                            if decoded_line.startswith("data:"):
-                                content = decoded_line.replace("data: ", "")
-                                yield f"data: {content}\n\n"
-                # Indicate the end of the stream
-                yield "data: [DONE]\n\n"
-            except aiohttp.ClientResponseError as e:
-                logger.error(f"HTTP error occurred: {e.status} - {e.message}")
-                raise HTTPException(status_code=e.status, detail=str(e))
-            except Exception as e:
-                logger.error(f"Unexpected error: {str(e)}")
-                raise HTTPException(status_code=500, detail=str(e))
-
-async def process_gizai_non_stream_response(request: ChatRequest, model: str) -> Union[ImageResponseModel, ChatCompletionResponse]:
-    async with aiohttp.ClientSession() as session:
-        # Set up headers
-        headers = {
-            'Accept': 'application/json, text/plain, */*',
-            'Accept-Language': 'en-US,en;q=0.9',
-            'Cache-Control': 'no-cache',
-            'Connection': 'keep-alive',
-            'Content-Type': 'application/json',
-            'Origin': 'https://app.giz.ai',
-            'Pragma': 'no-cache',
-            'Sec-Fetch-Dest': 'empty',
-            'Sec-Fetch-Mode': 'cors',
-            'Sec-Fetch-Site': 'same-origin',
-            'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36',
-            'sec-ch-ua': '"Not?A_Brand";v="99", "Chromium";v="130"',
-            'sec-ch-ua-mobile': '?0',
-            'sec-ch-ua-platform': '"Linux"'
-        }
-
-        if GizAI.is_image_model(model):
-            # Image generation logic
-            prompt = request.messages[-1].content if isinstance(request.messages[-1].content, str) else request.messages[-1].content[0].get("text", "")
-            data = {
-                "model": model,
-                "input": {
-                    "width": "1024",
-                    "height": "1024",
-                    "steps": 4,
-                    "output_format": "webp",
-                    "batch_size": 1,
-                    "mode": "plan",
-                    "prompt": prompt
-                }
-            }
-            try:
-                async with session.post(
-                    GIZAI_API_ENDPOINT,
-                    headers=headers,
-                    json=data
-                ) as response:
-                    response.raise_for_status()
-                    response_data = await response.json()
-                    if response_data.get('status') == 'completed' and response_data.get('output'):
-                        images = response_data['output']
-                        return ImageResponseModel(images=images, alt="Generated Image")
-                    else:
-                        raise HTTPException(status_code=500, detail="Image generation failed.")
-            except aiohttp.ClientResponseError as e:
-                logger.error(f"HTTP error occurred: {e.status} - {e.message}")
-                raise HTTPException(status_code=e.status, detail=str(e))
-            except Exception as e:
-                logger.error(f"Unexpected error: {str(e)}")
-                raise HTTPException(status_code=500, detail=str(e))
-        else:
-            # Chat completion logic
-            messages_formatted = [
-                {
-                    "type": "human",
-                    "content": msg.content if isinstance(msg.content, str) else msg.content[0].get("text", "")
-                } for msg in request.messages
-            ]
-            data = {
-                "model": model,
-                "input": {
-                    "messages": messages_formatted,
-                    "mode": "plan"
-                },
-                "noStream": True  # Disable streaming
-            }
-            try:
-                async with session.post(
-                    GIZAI_API_ENDPOINT,
-                    headers=headers,
-                    json=data
-                ) as response:
-                    response.raise_for_status()
-                    result = await response.json()
-                    return ChatCompletionResponse(
-                        id=f"chatcmpl-{uuid.uuid4()}",
-                        object="chat.completion",
-                        created=int(datetime.now().timestamp()),
-                        model=model,
-                        choices=[
-                            {
-                                "index": 0,
-                                "message": {"role": "assistant", "content": result.get('output', '')},
-                                "finish_reason": "stop",
-                            }
-                        ],
-                        usage=None,
-                    )
-            except aiohttp.ClientResponseError as e:
-                logger.error(f"HTTP error occurred: {e.status} - {e.message}")
-                raise HTTPException(status_code=e.status, detail=str(e))
-            except Exception as e:
-                logger.error(f"Unexpected error: {str(e)}")
-                raise HTTPException(status_code=500, detail=str(e))
+from api.config import MODEL_PROVIDER_MAPPING
+from api.models import ChatRequest
+from api.provider import blackboxai, gizai
+
+async def process_streaming_response(request: ChatRequest):
+    provider_name = MODEL_PROVIDER_MAPPING.get(request.model)
+    if provider_name == 'blackboxai':
+        return await blackboxai.process_streaming_response(request)
+    elif provider_name == 'gizai':
+        return await gizai.process_streaming_response(request)
+    else:
+        raise HTTPException(status_code=400, detail=f"Model {request.model} is not supported for streaming.")
+
+async def process_non_streaming_response(request: ChatRequest):
+    provider_name = MODEL_PROVIDER_MAPPING.get(request.model)
+    if provider_name == 'blackboxai':
+        return await blackboxai.process_non_streaming_response(request)
+    elif provider_name == 'gizai':
+        return await gizai.process_non_streaming_response(request)
+    else:
+        raise HTTPException(status_code=400, detail=f"Model {request.model} is not supported.")