|
|
|
|
|
|
|
|
from __future__ import annotations |
|
|
|
|
|
import json |
|
|
import uuid |
|
|
from aiohttp import ClientSession, ClientTimeout |
|
|
from typing import AsyncGenerator, List, Dict, Any, Union |
|
|
|
|
|
from api.logger import setup_logger |
|
|
|
|
|
logger = setup_logger(__name__) |
|
|
|
|
|
class AmigoChat: |
|
|
url = "https://amigochat.io" |
|
|
chat_api_endpoint = "https://api.amigochat.io/v1/chat/completions" |
|
|
image_api_endpoint = "https://api.amigochat.io/v1/images/generations" |
|
|
default_model = 'gpt-4o-mini' |
|
|
|
|
|
chat_models = [ |
|
|
'gpt-4o', |
|
|
default_model, |
|
|
'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', |
|
|
] |
|
|
|
|
|
image_models = [ |
|
|
'flux-pro/v1.1', |
|
|
'flux-realism', |
|
|
'flux-pro', |
|
|
'dalle-e-3', |
|
|
] |
|
|
|
|
|
models = chat_models + image_models |
|
|
|
|
|
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", |
|
|
"flux-pro": "flux-pro/v1.1", |
|
|
"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", |
|
|
} |
|
|
|
|
|
@classmethod |
|
|
def get_model(cls, model: str) -> str: |
|
|
return cls.model_aliases.get(model, model if model in cls.models else cls.default_model) |
|
|
|
|
|
@classmethod |
|
|
def get_persona_id(cls, model: str) -> str: |
|
|
return cls.persona_ids.get(model, "amigo") |
|
|
|
|
|
@classmethod |
|
|
async def generate_response( |
|
|
cls, |
|
|
model: str, |
|
|
messages: List[Dict[str, Any]], |
|
|
stream: bool = False, |
|
|
proxy: str = None, |
|
|
) -> Union[AsyncGenerator[str, None], str]: |
|
|
model = cls.get_model(model) |
|
|
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": cls.url, |
|
|
"pragma": "no-cache", |
|
|
"priority": "u=1, i", |
|
|
"referer": f"{cls.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 model in cls.chat_models: |
|
|
data = { |
|
|
"messages": messages, |
|
|
"model": model, |
|
|
"personaId": cls.get_persona_id(model), |
|
|
"frequency_penalty": 0, |
|
|
"max_tokens": 4000, |
|
|
"presence_penalty": 0, |
|
|
"stream": stream, |
|
|
"temperature": 0.5, |
|
|
"top_p": 0.95 |
|
|
} |
|
|
|
|
|
timeout = ClientTimeout(total=300) |
|
|
try: |
|
|
async with session.post(cls.chat_api_endpoint, json=data, proxy=proxy, timeout=timeout) as response: |
|
|
if response.status not in (200, 201): |
|
|
error_text = await response.text() |
|
|
raise Exception(f"Error {response.status}: {error_text}") |
|
|
|
|
|
if stream: |
|
|
async def stream_content(): |
|
|
async for line in response.content: |
|
|
line = line.decode('utf-8').strip() |
|
|
if line.startswith('data: '): |
|
|
if line == 'data: [DONE]': |
|
|
break |
|
|
try: |
|
|
chunk = json.loads(line[6:]) |
|
|
if 'choices' in chunk and len(chunk['choices']) > 0: |
|
|
choice = chunk['choices'][0] |
|
|
if 'delta' in choice: |
|
|
content = choice['delta'].get('content') |
|
|
elif 'text' in choice: |
|
|
content = choice['text'] |
|
|
else: |
|
|
content = None |
|
|
if content: |
|
|
yield content |
|
|
except json.JSONDecodeError: |
|
|
pass |
|
|
return stream_content() |
|
|
else: |
|
|
response_data = await response.json() |
|
|
if 'choices' in response_data and len(response_data['choices']) > 0: |
|
|
content = response_data['choices'][0]['message']['content'] |
|
|
return content |
|
|
else: |
|
|
return "" |
|
|
except Exception as e: |
|
|
logger.error(f"Error during request: {e}") |
|
|
raise |
|
|
else: |
|
|
|
|
|
pass |
|
|
|
