From 5ca47b44b2b42abb4f48163c17500b5ee67ab28f Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <heiner.lohaus@netformic.com>
Date: Tue, 5 Sep 2023 17:27:24 +0200
Subject: Add to many provider async and stream support, Fix Ails, AItianhu,
 ChatgptAi, ChatgptLogin Provider, Add fallback cookies to Bing, Improve
 OpenaiChat Provider

---
 g4f/Provider/ChatgptAi.py | 84 ++++++++++++++++++++++++++++-------------------
 1 file changed, 50 insertions(+), 34 deletions(-)

(limited to 'g4f/Provider/ChatgptAi.py')

diff --git a/g4f/Provider/ChatgptAi.py b/g4f/Provider/ChatgptAi.py
index 7613ccf1..e6416cc3 100644
--- a/g4f/Provider/ChatgptAi.py
+++ b/g4f/Provider/ChatgptAi.py
@@ -1,32 +1,28 @@
 from __future__ import annotations
 
 import re
+import html
+import json
+from aiohttp import ClientSession
 
-import requests
+from ..typing       import AsyncGenerator
+from .base_provider import AsyncGeneratorProvider
 
-from ..typing import Any, CreateResult
-from .base_provider import BaseProvider
 
+class ChatgptAi(AsyncGeneratorProvider):
+    url: str = "https://chatgpt.ai/"
+    working = True
+    supports_gpt_35_turbo  = True
+    _system_data = None
 
-class ChatgptAi(BaseProvider):
-    url: str        = "https://chatgpt.ai/gpt-4/"
-    working         = True
-    supports_gpt_4  = True
-
-    @staticmethod
-    def create_completion(
+    @classmethod
+    async def create_async_generator(
+        cls,
         model: str,
         messages: list[dict[str, str]],
-        stream: bool, **kwargs: Any) -> CreateResult:
-        
-        chat = "\n".join(f"{message['role']}: {message['content']}" for message in messages)
-        chat += "\nassistant: "
-
-        response = requests.get("https://chatgpt.ai/")
-        nonce, post_id, _, bot_id = re.findall(
-            r'data-nonce="(.*)"\n     data-post-id="(.*)"\n     data-url="(.*)"\n     data-bot-id="(.*)"\n     data-width',
-            response.text)[0]
-
+        proxy: str = None,
+        **kwargs
+    ) -> AsyncGenerator:
         headers = {
             "authority"          : "chatgpt.ai",
             "accept"             : "*/*",
@@ -34,7 +30,7 @@ class ChatgptAi(BaseProvider):
             "cache-control"      : "no-cache",
             "origin"             : "https://chatgpt.ai",
             "pragma"             : "no-cache",
-            "referer"            : "https://chatgpt.ai/gpt-4/",
+            "referer"            : cls.url,
             "sec-ch-ua"          : '"Not.A/Brand";v="8", "Chromium";v="114", "Google Chrome";v="114"',
             "sec-ch-ua-mobile"   : "?0",
             "sec-ch-ua-platform" : '"Windows"',
@@ -43,17 +39,37 @@ class ChatgptAi(BaseProvider):
             "sec-fetch-site"     : "same-origin",
             "user-agent"         : "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36",
         }
-        data = {
-            "_wpnonce" : nonce,
-            "post_id"  : post_id,
-            "url"      : "https://chatgpt.ai/gpt-4",
-            "action"   : "wpaicg_chat_shortcode_message",
-            "message"  : chat,
-            "bot_id"   : bot_id,
-        }
+        async with ClientSession(
+            headers=headers
+        ) as session:
+            if not cls._system_data:
+                async with session.get(cls.url, proxy=proxy) as response:
+                    response.raise_for_status()
+                    match = re.findall(r"data-system='([^']+)'", await response.text())
+                    if not match:
+                        raise RuntimeError("No system data")
+                    cls._system_data = json.loads(html.unescape(match[0]))
 
-        response = requests.post(
-            "https://chatgpt.ai/wp-admin/admin-ajax.php", headers=headers, data=data)
-        
-        response.raise_for_status()
-        yield response.json()["data"]
\ No newline at end of file
+            data = {
+                "botId":      cls._system_data["botId"],
+                "clientId":   "",
+                "contextId":  cls._system_data["contextId"],
+                "id":         cls._system_data["id"],
+                "messages":   messages[:-1],
+                "newMessage": messages[-1]["content"],
+                "session":    cls._system_data["sessionId"],
+                "stream":     True
+            }
+            async with session.post(
+                "https://chatgpt.ai/wp-json/mwai-ui/v1/chats/submit",
+                proxy=proxy,
+                json=data
+            ) as response:
+                response.raise_for_status()
+                start = "data: "
+                async for line in response.content:
+                    line = line.decode('utf-8')
+                    if line.startswith(start):
+                        line = json.loads(line[len(start):-1])
+                        if line["type"] == "live":
+                            yield line["data"]
\ No newline at end of file
-- 
cgit v1.2.3