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