Create v2.py

v2.py

@@ -0,0 +1,127 @@
+import re
+from dotenv import load_dotenv
+import requests
+from uuid import uuid4
+import json
+from typing import Any, AsyncGenerator, Dict
+import os
+
+# Load environment variables from .env file
+load_dotenv()
+
+
+class v2:
+    def __init__(
+        self,
+        timeout: int = 100,
+        proxies: dict = {},
+    ):  
+        self.session = requests.Session()
+        self.chat_endpoint = os.getenv("v2")
+        self.stream_chunk_size = 64
+        self.timeout = timeout
+        self.last_response = {}
+        self.headers = {
+            "accept": "*/*",
+            "accept-encoding": "gzip, deflate, br, zstd",
+            "accept-language": "en-US,en;q=0.9,en-IN;q=0.8",
+            "content-type": "application/json",
+        }
+
+        self.session.headers.update(self.headers)
+        self.session.proxies = proxies
+
+
+    def ask(
+        self,
+        prompt: str,
+        stream: bool = False,
+        raw: bool = False,
+    ) -> dict:
+        conversation_prompt =f"Tell me evrything about {prompt}"
+
+        self.session.headers.update(self.headers)
+        payload = {
+            "query": conversation_prompt,
+            "search_uuid": uuid4().hex,
+            "lang": "",
+            "agent_lang": "en",
+            "search_options": {
+                "langcode": "en-US"
+            },
+            "search_video": True,
+            "contexts_from": "google"
+        }
+
+        def for_stream():
+            response = self.session.post(
+                self.chat_endpoint, json=payload, stream=True, timeout=self.timeout
+            )
+            if not response.ok:
+                raise Exception(
+                    f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+                )
+
+            streaming_text = ""
+            for line in response.iter_lines(decode_unicode=True):
+                if line.startswith('data:'):
+                    try:
+                        data = json.loads(line[5:].strip())
+                        if data['type'] == 'answer' and 'text' in data['data']:
+                            new_text = data['data']['text']
+                            if len(new_text) > len(streaming_text):
+                                delta = new_text[len(streaming_text):]
+                                streaming_text = new_text
+                                resp = dict(text=delta)
+                                self.last_response.update(dict(text=streaming_text))
+                                yield line if raw else resp
+                    except json.JSONDecodeError:
+                        pass
+
+        def for_non_stream():
+            full_response = ""
+            for chunk in for_stream():
+                if not raw:
+                    full_response += chunk['text']
+            self.last_response = dict(text=full_response)
+            return self.last_response
+
+        return for_stream() if stream else for_non_stream()
+
+    def chat(
+        self,
+        prompt: str,
+        stream: bool = False,
+    ) -> str:
+        def for_stream():
+            for response in self.ask(
+                prompt, True
+            ):
+                yield self.get_message(response)
+
+        def for_non_stream():
+            return self.get_message(
+                self.ask(
+                    prompt,
+                    False,
+                )
+            )
+
+        return for_stream() if stream else for_non_stream()
+
+    def get_message(self, response: dict) -> str:
+        assert isinstance(response, dict), "Response should be of dict data-type only"
+
+        if "text" in response:
+            text = re.sub(r'\[\[\d+\]\]', '', response["text"])
+            return text
+        else:
+            return ""
+
+
+if __name__ == '__main__':
+    from rich import print
+    ai = v2()
+    response = ai.chat("HelpingAI-9B", stream=True)
+    for chunk in response:
+        print(chunk, end="", flush=True)