XFE Git
XFE Studio Git
Git 首页 全局搜索
XFE 主站 文档 NuGet
公开
关注 0 Fork 0 Star 1
返回提交历史

XFEstudio/gpt4free

refactor(g4f/Provider/AIUncensored.py): Enhance robustness and add features

d41f599a
kqlio67 <kqlio67@users.noreply.github.com>
提交于

代码差异

1 个文件 +49 -37
Modified g4f/Provider/AIUncensored.py +49 -37
@@ -1,17 +1,17 @@
1 1 from __future__ import annotations
2 2
3 3 import json
4 from aiohttp import ClientSession
5 from itertools import cycle
4 import random
5 import logging
6 from aiohttp import ClientSession, ClientError
7 from typing import List
6 8
7 9 from ..typing import AsyncResult, Messages
8 10 from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
9 from .helper import format_prompt
10 11 from ..image import ImageResponse
11 12
12
13 13 class AIUncensored(AsyncGeneratorProvider, ProviderModelMixin):
14 url = "https://www.aiuncensored.info"
14 url = "https://www.aiuncensored.info/ai_uncensored"
15 15 api_endpoints_text = [
16 16 "https://twitterclone-i0wr.onrender.com/api/chat",
17 17 "https://twitterclone-4e8t.onrender.com/api/chat",
@@ -22,8 +22,6 @@ class AIUncensored(AsyncGeneratorProvider, ProviderModelMixin):
22 22 "https://twitterclone-i0wr.onrender.com/api/image",
23 23 "https://twitterclone-8wd1.onrender.com/api/image",
24 24 ]
25 api_endpoints_cycle_text = cycle(api_endpoints_text)
26 api_endpoints_cycle_image = cycle(api_endpoints_image)
27 25 working = True
28 26 supports_stream = True
29 27 supports_system_message = True
@@ -35,10 +33,32 @@ class AIUncensored(AsyncGeneratorProvider, ProviderModelMixin):
35 33 models = [*text_models, *image_models]
36 34
37 35 model_aliases = {
38 #"": "TextGenerations",
39 36 "flux": "ImageGenerations",
40 37 }
41 38
39 @staticmethod
40 def generate_cipher() -> str:
41 return ''.join([str(random.randint(0, 9)) for _ in range(16)])
42
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
42 62 @classmethod
43 63 def get_model(cls, model: str) -> str:
44 64 if model in cls.models:
@@ -81,36 +101,28 @@ class AIUncensored(AsyncGeneratorProvider, ProviderModelMixin):
81 101 prompt = messages[-1]['content']
82 102 data = {
83 103 "prompt": prompt,
104 "cipher": cls.generate_cipher()
84 105 }
85 api_endpoint = next(cls.api_endpoints_cycle_image)
86 async with session.post(api_endpoint, json=data, proxy=proxy) as response:
87 response.raise_for_status()
88 response_data = await response.json()
89 image_url = response_data['image_url']
90 image_response = ImageResponse(images=image_url, alt=prompt)
91 yield image_response
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
92 112 elif model in cls.text_models:
93 113 data = {
94 "messages": [
95 {
96 "role": "user",
97 "content": format_prompt(messages)
98 }
99 ]
114 "messages": messages,
115 "cipher": cls.generate_cipher()
100 116 }
101 api_endpoint = next(cls.api_endpoints_cycle_text)
102 async with session.post(api_endpoint, json=data, proxy=proxy) as response:
103 response.raise_for_status()
104 full_response = ""
105 async for line in response.content:
106 line = line.decode('utf-8')
107 if line.startswith("data: "):
108 try:
109 json_str = line[6:]
110 if json_str != "[DONE]":
111 data = json.loads(json_str)
112 if "data" in data:
113 full_response += data["data"]
114 yield data["data"]
115 except json.JSONDecodeError:
116 continue
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