1 files changed, 38 insertions, 11 deletions

diff --git a/g4f/Provider/AItianhu.py b/g4f/Provider/AItianhu.py
index 0f01e536..c6e4dbad 100644
--- a/g4f/Provider/AItianhu.py
+++ b/g4f/Provider/AItianhu.py
@@ -1,24 +1,29 @@
 from __future__ import annotations
 
 import json
-from curl_cffi.requests import AsyncSession
 
-from .base_provider import AsyncProvider, format_prompt
+from ..typing import AsyncGenerator
+from ..requests import StreamSession
+from .base_provider import AsyncGeneratorProvider, format_prompt, get_cookies
 
 
-class AItianhu(AsyncProvider):
+class AItianhu(AsyncGeneratorProvider):
     url = "https://www.aitianhu.com"
     working = True
     supports_gpt_35_turbo = True
 
     @classmethod
-    async def create_async(
+    async def create_async_generator(
         cls,
         model: str,
         messages: list[dict[str, str]],
         proxy: str = None,
+        cookies: dict = None,
+        timeout: int = 30,
         **kwargs
-    ) -> str:
+    ) -> AsyncGenerator:
+        if not cookies:
+            cookies = get_cookies("www.aitianhu.com")
         data = {
             "prompt": format_prompt(messages),
             "options": {},
@@ -27,12 +32,34 @@ class AItianhu(AsyncProvider):
             "top_p": 1,
             **kwargs
         }
-        async with AsyncSession(proxies={"https": proxy}, impersonate="chrome107", verify=False) as session:
-            response = await session.post(cls.url + "/api/chat-process", json=data)
-            response.raise_for_status()
-            line = response.text.splitlines()[-1]
-            line = json.loads(line)
-            return line["text"]
+        headers = {
+            "Authority": cls.url,
+            "Accept": "application/json, text/plain, */*",
+            "Origin": cls.url,
+            "Referer": f"{cls.url}/"
+        }
+        async with StreamSession(
+            headers=headers,
+            cookies=cookies,
+            timeout=timeout,
+            proxies={"https": proxy},
+            impersonate="chrome107",
+            verify=False
+        ) as session:
+            async with session.post(f"{cls.url}/api/chat-process", json=data) as response:
+                response.raise_for_status()
+                async for line in response.iter_lines():
+                    if line == b"<script>":
+                        raise RuntimeError("Solve challenge and pass cookies")
+                    if b"platform's risk control" in line:
+                        raise RuntimeError("Platform's Risk Control")
+                    line = json.loads(line)
+                    if "detail" in line:
+                        content = line["detail"]["choices"][0]["delta"].get("content")
+                        if content:
+                            yield content
+                    else:
+                        raise RuntimeError(f"Response: {line}")
 
 
     @classmethod