api/utils.py

# 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))