返回提交历史
Modified
g4f/Provider/AIUncensored.py
+42
-38
XFEstudio/gpt4free
refactor(g4f/Provider/AIUncensored.py): Enhance AIUncensored provider with improved resilience and flexibility
9fe5ac67
代码差异
1 个文件
+42
-38
@@ -2,9 +2,9 @@ from __future__ import annotations
2
2
3
3
import json
4
4
import random
5
import logging
6
5
from aiohttp import ClientSession, ClientError
7
from typing import List
6
import asyncio
7
from itertools import cycle
8
8
9
9
from ..typing import AsyncResult, Messages
10
10
from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
@@ -38,27 +38,9 @@ class AIUncensored(AsyncGeneratorProvider, ProviderModelMixin):
38
38
39
39
@staticmethod
40
40
def generate_cipher() -> str:
41
"""Generate a cipher in format like '3221229284179118'"""
41
42
return ''.join([str(random.randint(0, 9)) for _ in range(16)])
42
43
43
@staticmethod
44
async def try_request(session: ClientSession, endpoints: List[str], data: dict, proxy: str = None):
45
available_endpoints = endpoints.copy()
46
random.shuffle(available_endpoints)
47
48
while available_endpoints:
49
endpoint = available_endpoints.pop()
50
try:
51
async with session.post(endpoint, json=data, proxy=proxy) as response:
52
response.raise_for_status()
53
return response
54
except ClientError as e:
55
logging.warning(f"Failed to connect to {endpoint}: {str(e)}")
56
if not available_endpoints:
57
raise
58
continue
59
60
raise Exception("All endpoints are unavailable")
61
62
44
@classmethod
63
45
def get_model(cls, model: str) -> str:
64
46
if model in cls.models:
@@ -103,26 +85,48 @@ class AIUncensored(AsyncGeneratorProvider, ProviderModelMixin):
103
85
"prompt": prompt,
104
86
"cipher": cls.generate_cipher()
105
87
}
106
response = await cls.try_request(session, cls.api_endpoints_image, data, proxy)
107
response_data = await response.json()
108
image_url = response_data['image_url']
109
image_response = ImageResponse(images=image_url, alt=prompt)
110
yield image_response
111
88
89
endpoints = cycle(cls.api_endpoints_image)
90
91
while True:
92
endpoint = next(endpoints)
93
try:
94
async with session.post(endpoint, json=data, proxy=proxy, timeout=10) as response:
95
response.raise_for_status()
96
response_data = await response.json()
97
image_url = response_data['image_url']
98
image_response = ImageResponse(images=image_url, alt=prompt)
99
yield image_response
100
return
101
except (ClientError, asyncio.TimeoutError):
102
continue
103
112
104
elif model in cls.text_models:
113
105
data = {
114
106
"messages": messages,
115
107
"cipher": cls.generate_cipher()
116
108
}
117
response = await cls.try_request(session, cls.api_endpoints_text, data, proxy)
118
async for line in response.content:
119
line = line.decode('utf-8')
120
if line.startswith("data: "):
121
try:
122
json_str = line[6:]
123
if json_str != "[DONE]":
124
data = json.loads(json_str)
125
if "data" in data:
126
yield data["data"]
127
except json.JSONDecodeError:
128
continue
109
110
endpoints = cycle(cls.api_endpoints_text)
111
112
while True:
113
endpoint = next(endpoints)
114
try:
115
async with session.post(endpoint, json=data, proxy=proxy, timeout=10) as response:
116
response.raise_for_status()
117
full_response = ""
118
async for line in response.content:
119
line = line.decode('utf-8')
120
if line.startswith("data: "):
121
try:
122
json_str = line[6:]
123
if json_str != "[DONE]":
124
data = json.loads(json_str)
125
if "data" in data:
126
full_response += data["data"]
127
yield data["data"]
128
except json.JSONDecodeError:
129
continue
130
return
131
except (ClientError, asyncio.TimeoutError):
132
continue