diff options
Diffstat (limited to '')
| -rw-r--r-- | g4f/Provider/AmigoChat.py | 174 |
1 files changed, 93 insertions, 81 deletions
diff --git a/g4f/Provider/AmigoChat.py b/g4f/Provider/AmigoChat.py index 6e22094c..5e896dc8 100644 --- a/g4f/Provider/AmigoChat.py +++ b/g4f/Provider/AmigoChat.py @@ -2,7 +2,7 @@ from __future__ import annotations import json import uuid -from aiohttp import ClientSession, ClientTimeout +from aiohttp import ClientSession, ClientTimeout, ClientResponseError from ..typing import AsyncResult, Messages from .base_provider import AsyncGeneratorProvider, ProviderModelMixin @@ -92,87 +92,99 @@ class AmigoChat(AsyncGeneratorProvider, ProviderModelMixin): model = cls.get_model(model) device_uuid = str(uuid.uuid4()) + max_retries = 3 + retry_count = 0 - headers = { - "accept": "*/*", - "accept-language": "en-US,en;q=0.9", - "authorization": "Bearer", # You need to implement proper authorization - "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: - # Chat completion - data = { - "messages": [{"role": m["role"], "content": m["content"]} for m in messages], - "model": model, - "personaId": cls.get_personaId(model), - "frequency_penalty": 0, - "max_tokens": 4000, - "presence_penalty": 0, - "stream": stream, - "temperature": 0.5, - "top_p": 0.95 + while retry_count < max_retries: + try: + 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" } - timeout = ClientTimeout(total=300) # 5 minutes timeout - 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}") - - 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:]) # Remove 'data: ' prefix - 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 - else: - # Image generation - prompt = messages[0]['content'] - data = { - "prompt": prompt, - "model": model, - "personaId": cls.get_personaId(model) - } - async with session.post(cls.image_api_endpoint, json=data, proxy=proxy) 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: - yield ImageResponse(image_urls, prompt) + async with ClientSession(headers=headers) as session: + if model in cls.chat_models: + # Chat completion + data = { + "messages": [{"role": m["role"], "content": m["content"]} for m in messages], + "model": model, + "personaId": cls.get_personaId(model), + "frequency_penalty": 0, + "max_tokens": 4000, + "presence_penalty": 0, + "stream": stream, + "temperature": 0.5, + "top_p": 0.95 + } + + timeout = ClientTimeout(total=300) # 5 minutes timeout + 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}") + + 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:]) # Remove 'data: ' prefix + 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 else: - yield None + # Image generation + prompt = messages[0]['content'] + data = { + "prompt": prompt, + "model": model, + "personaId": cls.get_personaId(model) + } + async with session.post(cls.image_api_endpoint, json=data, proxy=proxy) 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: + yield ImageResponse(image_urls, prompt) + else: + yield None + + break + + except (ClientResponseError, Exception) as e: + retry_count += 1 + if retry_count >= max_retries: + raise e + device_uuid = str(uuid.uuid4()) |