Rename api/provider/gizai.py to api/provider/amigochat.py

api/provider/amigochat.py

@@ -0,0 +1,265 @@
+import json
+import uuid
+import asyncio
+from aiohttp import ClientSession, ClientTimeout, ClientResponseError
+from fastapi import HTTPException
+from api.logger import setup_logger
+from api.config import MODEL_MAPPING
+
+logger = setup_logger(__name__)
+
+# Base URLs
+AMIGOCHAT_URL = "https://amigochat.io/chat/"
+CHAT_API_ENDPOINT = "https://api.amigochat.io/v1/chat/completions"
+IMAGE_API_ENDPOINT = "https://api.amigochat.io/v1/images/generations"
+
+# List of models supported by AmigoChat
+AMIGOCHAT_CHAT_MODELS = [
+    'gpt-4o',
+    'gpt-4o-mini',
+    'o1-preview',
+    'o1-mini',
+    'meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo',
+    'meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo',
+    'claude-3-sonnet-20240229',
+    'gemini-1.5-pro',
+]
+
+AMIGOCHAT_IMAGE_MODELS = [
+    'flux-pro/v1.1',
+    'flux-realism',
+    'flux-pro',
+    'dalle-e-3',
+]
+
+AMIGOCHAT_MODELS = AMIGOCHAT_CHAT_MODELS + AMIGOCHAT_IMAGE_MODELS
+
+AMIGOCHAT_MODEL_ALIASES = {
+    "o1": "o1-preview",
+    "llama-3.1-405b": "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
+    "llama-3.2-90b": "meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo",
+    "claude-3.5-sonnet": "claude-3-sonnet-20240229",
+    "gemini-pro": "gemini-1.5-pro",
+    "dalle-3": "dalle-e-3",
+}
+
+PERSONA_IDS = {
+    'gpt-4o': "gpt",
+    'gpt-4o-mini': "amigo",
+    'o1-preview': "openai-o-one",
+    'o1-mini': "openai-o-one-mini",
+    'meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo': "llama-three-point-one",
+    'meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo': "llama-3-2",
+    'claude-3-sonnet-20240229': "claude",
+    'gemini-1.5-pro': "gemini-1-5-pro",
+    'flux-pro/v1.1': "flux-1-1-pro",
+    'flux-realism': "flux-realism",
+    'flux-pro': "flux-pro",
+    'dalle-e-3': "dalle-three",
+}
+
+def get_amigochat_model(model: str) -> str:
+    model = MODEL_MAPPING.get(model, model)
+    if model in AMIGOCHAT_MODELS:
+        return model
+    elif model in AMIGOCHAT_MODEL_ALIASES:
+        return AMIGOCHAT_MODEL_ALIASES[model]
+    else:
+        # Default model
+        return 'gpt-4o-mini'
+
+def get_persona_id(model: str) -> str:
+    return PERSONA_IDS.get(model, "amigo")
+
+def is_image_model(model: str) -> bool:
+    return model in AMIGOCHAT_IMAGE_MODELS
+
+async def process_streaming_response(request_data):
+    model = get_amigochat_model(request_data.get('model'))
+    messages = request_data.get('messages')
+    stream = request_data.get('stream', False)
+    if not messages:
+        raise HTTPException(status_code=400, detail="Messages are required")
+
+    device_uuid = str(uuid.uuid4())
+
+    headers = {
+        "accept": "*/*",
+        "accept-language": "en-US,en;q=0.9",
+        "authorization": "Bearer",
+        "cache-control": "no-cache",
+        "content-type": "application/json",
+        "origin": AMIGOCHAT_URL,
+        "pragma": "no-cache",
+        "priority": "u=1, i",
+        "referer": f"{AMIGOCHAT_URL}/",
+        "sec-ch-ua": '"Chromium";v="129", "Not=A?Brand";v="8"',
+        "sec-ch-ua-mobile": "?0",
+        "sec-ch-ua-platform": '"Linux"',
+        "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
+        "x-device-language": "en-US",
+        "x-device-platform": "web",
+        "x-device-uuid": device_uuid,
+        "x-device-version": "1.0.32"
+    }
+
+    async with ClientSession(headers=headers) as session:
+        if is_image_model(model):
+            # Image generation does not support streaming
+            response = await process_non_streaming_response(request_data)
+            return iter([json.dumps(response)])
+        else:
+            # Chat completion
+            data = {
+                "messages": [{"role": m["role"], "content": m["content"]} for m in messages],
+                "model": model,
+                "personaId": get_persona_id(model),
+                "frequency_penalty": 0,
+                "max_tokens": 4000,
+                "presence_penalty": 0,
+                "stream": True,
+                "temperature": 0.5,
+                "top_p": 0.95
+            }
+
+            timeout = ClientTimeout(total=300)  # 5 minutes timeout
+
+            async def event_stream():
+                try:
+                    async with session.post(CHAT_API_ENDPOINT, json=data, timeout=timeout) as resp:
+                        if resp.status not in (200, 201):
+                            error_text = await resp.text()
+                            raise HTTPException(status_code=resp.status, detail=error_text)
+
+                        async for line in resp.content:
+                            line = line.decode('utf-8').strip()
+                            if line.startswith('data: '):
+                                if line == 'data: [DONE]':
+                                    break
+                                try:
+                                    chunk = json.loads(line[6:])  # Remove 'data: ' prefix
+                                    if 'choices' in chunk and len(chunk['choices']) > 0:
+                                        choice = chunk['choices'][0]
+                                        delta = choice.get('delta', {})
+                                        content = delta.get('content')
+                                        if content:
+                                            # Yield the chunk in OpenAI format
+                                            response_data = {
+                                                "id": f"chatcmpl-{uuid.uuid4()}",
+                                                "object": "chat.completion.chunk",
+                                                "created": int(datetime.now().timestamp()),
+                                                "model": model,
+                                                "choices": [
+                                                    {
+                                                        "delta": {"content": content},
+                                                        "index": 0,
+                                                        "finish_reason": None,
+                                                    }
+                                                ],
+                                            }
+                                            yield f"data: {json.dumps(response_data)}\n\n"
+                                except json.JSONDecodeError:
+                                    pass
+                    # Signal the end of the stream
+                    yield "data: [DONE]\n\n"
+                except Exception as e:
+                    logger.error(f"Error in streaming response: {e}")
+                    raise HTTPException(status_code=500, detail=str(e))
+
+            return event_stream()
+
+async def process_non_streaming_response(request_data):
+    model = get_amigochat_model(request_data.get('model'))
+    messages = request_data.get('messages')
+    if not messages:
+        raise HTTPException(status_code=400, detail="Messages are required")
+
+    device_uuid = str(uuid.uuid4())
+
+    headers = {
+        "accept": "*/*",
+        "accept-language": "en-US,en;q=0.9",
+        "authorization": "Bearer",
+        "cache-control": "no-cache",
+        "content-type": "application/json",
+        "origin": AMIGOCHAT_URL,
+        "pragma": "no-cache",
+        "priority": "u=1, i",
+        "referer": f"{AMIGOCHAT_URL}/",
+        "sec-ch-ua": '"Chromium";v="129", "Not=A?Brand";v="8"',
+        "sec-ch-ua-mobile": "?0",
+        "sec-ch-ua-platform": '"Linux"',
+        "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
+        "x-device-language": "en-US",
+        "x-device-platform": "web",
+        "x-device-uuid": device_uuid,
+        "x-device-version": "1.0.32"
+    }
+
+    async with ClientSession(headers=headers) as session:
+        if is_image_model(model):
+            # Image generation
+            prompt = messages[-1]['content']
+            data = {
+                "prompt": prompt,
+                "model": model,
+                "personaId": get_persona_id(model)
+            }
+            try:
+                async with session.post(IMAGE_API_ENDPOINT, json=data) as response:
+                    response.raise_for_status()
+                    response_data = await response.json()
+                    if "data" in response_data:
+                        image_urls = []
+                        for item in response_data["data"]:
+                            if "url" in item:
+                                image_url = item["url"]
+                                image_urls.append(image_url)
+                        if image_urls:
+                            return {
+                                "id": f"imggen-{uuid.uuid4()}",
+                                "object": "image_generation",
+                                "created": int(datetime.now().timestamp()),
+                                "model": model,
+                                "data": image_urls,
+                            }
+                    raise HTTPException(status_code=500, detail="Image generation failed")
+            except Exception as e:
+                logger.error(f"Error in image generation: {e}")
+                raise HTTPException(status_code=500, detail=str(e))
+        else:
+            # Chat completion
+            data = {
+                "messages": [{"role": m["role"], "content": m["content"]} for m in messages],
+                "model": model,
+                "personaId": get_persona_id(model),
+                "frequency_penalty": 0,
+                "max_tokens": 4000,
+                "presence_penalty": 0,
+                "stream": False,
+                "temperature": 0.5,
+                "top_p": 0.95
+            }
+
+            try:
+                async with session.post(CHAT_API_ENDPOINT, json=data) as response:
+                    response.raise_for_status()
+                    response_data = await response.json()
+                    output = response_data.get('choices', [{}])[0].get('message', {}).get('content', '')
+                    return {
+                        "id": f"chatcmpl-{uuid.uuid4()}",
+                        "object": "chat.completion",
+                        "created": int(datetime.now().timestamp()),
+                        "model": model,
+                        "choices": [
+                            {
+                                "index": 0,
+                                "message": {"role": "assistant", "content": output},
+                                "finish_reason": "stop",
+                            }
+                        ],
+                        "usage": None,
+                    }
+            except Exception as e:
+                logger.error(f"Error in chat completion: {e}")
+                raise HTTPException(status_code=500, detail=str(e))

api/provider/gizai.py

@@ -1,186 +0,0 @@
-import uuid
-from datetime import datetime
-import json
-from typing import Any, Dict
-
-import httpx
-from fastapi import HTTPException
-from api.logger import setup_logger
-from api.config import MODEL_MAPPING
-
-logger = setup_logger(__name__)
-
-# Base URL and API Endpoint for GizAI
-GIZAI_BASE_URL = "https://app.giz.ai"
-GIZAI_API_ENDPOINT = f"{GIZAI_BASE_URL}/api/data/users/inferenceServer.infer"
-
-# Headers for GizAI
-GIZAI_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"'
-}
-
-# List of models supported by GizAI
-GIZAI_CHAT_MODELS = [
-    'chat-gemini-flash',
-    'chat-gemini-pro',
-    'chat-gpt4m',
-    'chat-gpt4',
-    'claude-sonnet',
-    'claude-haiku',
-    'llama-3-70b',
-    'llama-3-8b',
-    'mistral-large',
-    'chat-o1-mini'
-]
-
-GIZAI_IMAGE_MODELS = [
-    'flux1',
-    'sdxl',
-    'sd',
-    'sd35',
-]
-
-GIZAI_MODELS = GIZAI_CHAT_MODELS + GIZAI_IMAGE_MODELS
-
-GIZAI_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",
-}
-
-def get_gizai_model(model: str) -> str:
-    model = MODEL_MAPPING.get(model, model)
-    if model in GIZAI_MODELS:
-        return model
-    elif model in GIZAI_MODEL_ALIASES:
-        return GIZAI_MODEL_ALIASES[model]
-    else:
-        # Default model
-        return 'chat-gemini-flash'
-
-def is_image_model(model: str) -> bool:
-    return model in GIZAI_IMAGE_MODELS
-
-async def process_streaming_response(request_data):
-    # GizAI does not support streaming; handle as non-streaming
-    response = await process_non_streaming_response(request_data)
-    # Return the response wrapped in an iterator
-    return iter([json.dumps(response)])
-
-async def process_non_streaming_response(request_data):
-    model = get_gizai_model(request_data.get('model'))
-
-    async with httpx.AsyncClient() as client:
-        if is_image_model(model):
-            # Image generation
-            prompt = request_data['messages'][-1]['content']
-            data = {
-                "model": model,
-                "input": {
-                    "width": "1024",
-                    "height": "1024",
-                    "steps": 4,
-                    "output_format": "webp",
-                    "batch_size": 1,
-                    "mode": "plan",
-                    "prompt": prompt
-                }
-            }
-            try:
-                response = await client.post(
-                    GIZAI_API_ENDPOINT,
-                    headers=GIZAI_HEADERS,
-                    json=data,
-                    timeout=100,
-                )
-                response.raise_for_status()
-                response_data = response.json()
-                if response_data.get('status') == 'completed' and response_data.get('output'):
-                    images = response_data['output']
-                    # Return image response (e.g., URLs)
-                    return {
-                        "id": f"imggen-{uuid.uuid4()}",
-                        "object": "image_generation",
-                        "created": int(datetime.now().timestamp()),
-                        "model": request_data['model'],
-                        "data": images,
-                    }
-                else:
-                    raise HTTPException(status_code=500, detail="Image generation failed")
-            except httpx.HTTPStatusError as e:
-                logger.error(f"HTTP error occurred: {e}")
-                raise HTTPException(status_code=e.response.status_code, detail=str(e))
-            except httpx.RequestError as e:
-                logger.error(f"Error occurred during request: {e}")
-                raise HTTPException(status_code=500, detail=str(e))
-        else:
-            # Chat completion
-            messages = request_data['messages']
-            messages_content = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
-            data = {
-                "model": model,
-                "input": {
-                    "messages": [
-                        {
-                            "type": "human",
-                            "content": messages_content
-                        }
-                    ],
-                    "mode": "plan"
-                },
-                "noStream": True
-            }
-            try:
-                response = await client.post(
-                    GIZAI_API_ENDPOINT,
-                    headers=GIZAI_HEADERS,
-                    json=data,
-                    timeout=100,
-                )
-                response.raise_for_status()
-                response_data = response.json()
-                output = response_data.get('output', '')
-                return {
-                    "id": f"chatcmpl-{uuid.uuid4()}",
-                    "object": "chat.completion",
-                    "created": int(datetime.now().timestamp()),
-                    "model": request_data['model'],
-                    "choices": [
-                        {
-                            "index": 0,
-                            "message": {"role": "assistant", "content": output},
-                            "finish_reason": "stop",
-                        }
-                    ],
-                    "usage": None,
-                }
-            except httpx.HTTPStatusError as e:
-                logger.error(f"HTTP error occurred: {e}")
-                raise HTTPException(status_code=e.response.status_code, detail=str(e))
-            except httpx.RequestError as e:
-                logger.error(f"Error occurred during request: {e}")
-                raise HTTPException(status_code=500, detail=str(e))