From 6401084fd0d23b75a883888f70cd705e4cacd202 Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <heiner@lohaus.eu>
Date: Sun, 8 Oct 2023 13:59:56 +0200
Subject: Add Messages and AsyncResult typing Add system_message in Yqcloud

---
 g4f/Provider/Yqcloud.py | 32 ++++++++++++++++++++------------
 1 file changed, 20 insertions(+), 12 deletions(-)

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

diff --git a/g4f/Provider/Yqcloud.py b/g4f/Provider/Yqcloud.py
index 74c998e2..0e96e20f 100644
--- a/g4f/Provider/Yqcloud.py
+++ b/g4f/Provider/Yqcloud.py
@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+import random
 from aiohttp import ClientSession
 
-from ..typing import AsyncGenerator
+from ..typing import AsyncResult, Messages
 from .base_provider import AsyncGeneratorProvider, format_prompt
 
 
@@ -14,22 +15,22 @@ class Yqcloud(AsyncGeneratorProvider):
     @staticmethod
     async def create_async_generator(
         model: str,
-        messages: list[dict[str, str]],
+        messages: Messages,
         proxy: str = None,
         **kwargs,
-    ) -> AsyncGenerator:
+    ) -> AsyncResult:
         async with ClientSession(
             headers=_create_header()
         ) as session:
-            payload = _create_payload(messages)
+            payload = _create_payload(messages, **kwargs)
             async with session.post("https://api.aichatos.cloud/api/generateStream", proxy=proxy, json=payload) as response:
                 response.raise_for_status()
-                async for stream in response.content.iter_any():
-                    if stream:
-                        stream = stream.decode()
-                        if "sorry, 您的ip已由于触发防滥用检测而被封禁" in stream:
+                async for chunk in response.content.iter_any():
+                    if chunk:
+                        chunk = chunk.decode()
+                        if "sorry, 您的ip已由于触发防滥用检测而被封禁" in chunk:
                             raise RuntimeError("IP address is blocked by abuse detection.")
-                        yield stream.decode()
+                        yield chunk
 
 
 def _create_header():
@@ -40,12 +41,19 @@ def _create_header():
     }
 
 
-def _create_payload(messages: list[dict[str, str]]):
+def _create_payload(
+    messages: Messages,
+    system_message: str = "",
+    user_id: int = None,
+    **kwargs
+):
+    if not user_id:
+        user_id = random.randint(1690000544336, 2093025544336)
     return {
         "prompt": format_prompt(messages),
         "network": True,
-        "system": "",
+        "system": system_message,
         "withoutContext": False,
         "stream": True,
-        "userId": "#/chat/1693025544336"
+        "userId": f"#/chat/{user_id}"
     }
-- 
cgit v1.2.3