1 files changed, 70 insertions, 63 deletions
diff --git a/g4f/Provider/H2o.py b/g4f/Provider/H2o.py
index ea4d0825..ac5fcfb7 100644
--- a/g4f/Provider/H2o.py
+++ b/g4f/Provider/H2o.py
@@ -1,78 +1,85 @@
-import json, uuid, requests
+import json
+import uuid
+from aiohttp import ClientSession
 
-from ..typing       import Any, CreateResult
-from .base_provider import BaseProvider
+from ..typing import AsyncGenerator
+from .base_provider import AsyncGeneratorProvider, format_prompt
 
 
-class H2o(BaseProvider):
-    url             = "https://gpt-gm.h2o.ai"
-    working         = True
+class H2o(AsyncGeneratorProvider):
+    url = "https://gpt-gm.h2o.ai"
+    working = True
     supports_stream = True
     model           = "h2oai/h2ogpt-gm-oasst1-en-2048-falcon-40b-v1"
 
-    @staticmethod
-    def create_completion(
+    @classmethod
+    async def create_async_generator(
+        cls,
         model: str,
         messages: list[dict[str, str]],
-        stream: bool, **kwargs: Any) -> CreateResult:
-        
-        conversation = ""
-        for message in messages:
-            conversation += "%s: %s\n" % (message["role"], message["content"])
-        conversation += "assistant: "
-
-        session = requests.Session()
-
-        headers = {"Referer": "https://gpt-gm.h2o.ai/r/jGfKSwU"}
-        data = {
-            "ethicsModalAccepted"               : "true",
-            "shareConversationsWithModelAuthors": "true",
-            "ethicsModalAcceptedAt"             : "",
-            "activeModel"                       : model,
-            "searchEnabled"                     : "true",
-        }
-        
-        session.post("https://gpt-gm.h2o.ai/settings",
-                     headers=headers, data=data)
-
+        proxy: str = None,
+        **kwargs
+    ) -> AsyncGenerator:
+        model = model if model else cls.model
         headers = {"Referer": "https://gpt-gm.h2o.ai/"}
-        data    = {"model": model}
-
-        response = session.post("https://gpt-gm.h2o.ai/conversation",
-                                headers=headers, json=data).json()
-        
-        if "conversationId" not in response:
-            return
 
-        data = {
-            "inputs": conversation,
-            "parameters": {
-                "temperature"   : kwargs.get("temperature", 0.4),
-                "truncate"          : kwargs.get("truncate", 2048),
-                "max_new_tokens"    : kwargs.get("max_new_tokens", 1024),
-                "do_sample"         : kwargs.get("do_sample", True),
-                "repetition_penalty": kwargs.get("repetition_penalty", 1.2),
-                "return_full_text"  : kwargs.get("return_full_text", False),
-            },
-            "stream" : True,
-            "options": {
-                "id"           : kwargs.get("id", str(uuid.uuid4())),
-                "response_id"  : kwargs.get("response_id", str(uuid.uuid4())),
-                "is_retry"     : False,
-                "use_cache"    : False,
-                "web_search_id": "",
-            },
-        }
+        async with ClientSession(
+            headers=headers
+        ) as session:
+            data = {
+                "ethicsModalAccepted": "true",
+                "shareConversationsWithModelAuthors": "true",
+                "ethicsModalAcceptedAt": "",
+                "activeModel": model,
+                "searchEnabled": "true",
+            }
+            async with session.post(
+                "https://gpt-gm.h2o.ai/settings",
+                proxy=proxy,
+                data=data
+            ) as response:
+                response.raise_for_status()
 
-        response = session.post(f"https://gpt-gm.h2o.ai/conversation/{response['conversationId']}",
-            headers=headers, json=data)
-        
-        response.raise_for_status()
-        response.encoding = "utf-8"
-        generated_text    = response.text.replace("\n", "").split("data:")
-        generated_text    = json.loads(generated_text[-1])
+            async with session.post(
+                "https://gpt-gm.h2o.ai/conversation",
+                proxy=proxy,
+                json={"model": model},
+            ) as response:
+                response.raise_for_status()
+                conversationId = (await response.json())["conversationId"]
 
-        yield generated_text["generated_text"]
+            data = {
+                "inputs": format_prompt(messages),
+                "parameters": {
+                    "temperature": 0.4,
+                    "truncate": 2048,
+                    "max_new_tokens": 1024,
+                    "do_sample":  True,
+                    "repetition_penalty": 1.2,
+                    "return_full_text": False,
+                    **kwargs
+                },
+                "stream": True,
+                "options": {
+                    "id": str(uuid.uuid4()),
+                    "response_id": str(uuid.uuid4()),
+                    "is_retry": False,
+                    "use_cache": False,
+                    "web_search_id": "",
+                },
+            }
+            async with session.post(
+                f"https://gpt-gm.h2o.ai/conversation/{conversationId}",
+                proxy=proxy,
+                json=data
+             ) as response:
+                start = "data:"
+                async for line in response.content:
+                    line = line.decode("utf-8")
+                    if line and line.startswith(start):
+                        line = json.loads(line[len(start):-1])
+                        if not line["token"]["special"]:
+                            yield line["token"]["text"]
 
     @classmethod
     @property