1 files changed, 69 insertions, 67 deletions
diff --git a/g4f/Provider/H2o.py b/g4f/Provider/H2o.py
index f9b799bb..b62da977 100644
--- a/g4f/Provider/H2o.py
+++ b/g4f/Provider/H2o.py
@@ -1,86 +1,88 @@
+from __future__ import annotations
+
 import json
 import uuid
 
-import requests
+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):
+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,
-        )
-        conversation_id = response.json()["conversationId"]
 
-        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/{conversation_id}",
-            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