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

XFEstudio/gpt4free

feat(AiChats): add DALL-E support and improve error handling

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

代码差异

1 个文件 +47 -22
Modified g4f/Provider/AiChats.py +47 -22
@@ -1,19 +1,20 @@
1 1 from __future__ import annotations
2 2
3 from aiohttp import ClientSession
4 3 import json
5
4 import base64
5 from aiohttp import ClientSession
6 6 from ..typing import AsyncResult, Messages
7 from .base_provider import AsyncGeneratorProvider
8 from .helper import format_prompt
9
7 from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
8 from ..image import ImageResponse
10 9
11 class AiChats(AsyncGeneratorProvider):
10 class AiChats(AsyncGeneratorProvider, ProviderModelMixin):
12 11 url = "https://ai-chats.org"
13 12 api_endpoint = "https://ai-chats.org/chat/send2/"
14 13 working = True
15 14 supports_gpt_4 = True
16 15 supports_message_history = True
16 default_model = 'gpt-4'
17 models = ['gpt-4', 'dalle']
17 18
18 19 @classmethod
19 20 async def create_async_generator(
@@ -23,7 +24,6 @@ class AiChats(AsyncGeneratorProvider):
23 24 proxy: str = None,
24 25 **kwargs
25 26 ) -> AsyncResult:
26
27 27 headers = {
28 28 "accept": "application/json, text/event-stream",
29 29 "accept-language": "en-US,en;q=0.9",
@@ -31,7 +31,7 @@ class AiChats(AsyncGeneratorProvider):
31 31 "content-type": "application/json",
32 32 "origin": cls.url,
33 33 "pragma": "no-cache",
34 "referer": f"{cls.url}/chat/",
34 "referer": f"{cls.url}/{'image' if model == 'dalle' else 'chat'}/",
35 35 "sec-ch-ua": '"Chromium";v="127", "Not)A;Brand";v="99"',
36 36 "sec-ch-ua-mobile": "?0",
37 37 "sec-ch-ua-platform": '"Linux"',
@@ -42,9 +42,9 @@ class AiChats(AsyncGeneratorProvider):
42 42 }
43 43
44 44 async with ClientSession(headers=headers) as session:
45 prompt = format_prompt(messages)
45 prompt = cls.format_prompt(messages)
46 46 data = {
47 "type": "chat",
47 "type": "image" if model == 'dalle' else "chat",
48 48 "messagesHistory": [
49 49 {
50 50 "from": "you",
@@ -52,15 +52,40 @@ class AiChats(AsyncGeneratorProvider):
52 52 }
53 53 ]
54 54 }
55 async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
56 response.raise_for_status()
57 full_response = await response.text()
58
59 message = ""
60 for line in full_response.split('\n'):
61 if line.startswith('data: ') and line != 'data: ':
62 message += line[6:]
63
64 message = message.strip()
65
66 yield message
55
56 try:
57 async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
58 response.raise_for_status()
59
60 if model == 'dalle':
61 response_json = await response.json()
62
63 if 'data' in response_json and response_json['data']:
64 image_url = response_json['data'][0].get('url')
65 if image_url:
66 async with session.get(image_url) as img_response:
67 img_response.raise_for_status()
68 image_data = await img_response.read()
69
70 base64_image = base64.b64encode(image_data).decode('utf-8')
71 base64_url = f"data:image/png;base64,{base64_image}"
72 yield ImageResponse(base64_url, prompt)
73 else:
74 yield f"Error: No image URL found in the response. Full response: {response_json}"
75 else:
76 yield f"Error: Unexpected response format. Full response: {response_json}"
77 else:
78 full_response = await response.text()
79 message = ""
80 for line in full_response.split('\n'):
81 if line.startswith('data: ') and line != 'data: ':
82 message += line[6:]
83
84 message = message.strip()
85 yield message
86 except Exception as e:
87 yield f"Error occurred: {str(e)}"
88
89 @classmethod
90 def format_prompt(cls, messages: Messages) -> str:
91 return messages[-1]['content'] if messages else ""