Update v2.py

v2.py

@@ -1,15 +1,11 @@
 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
+from typing import Any, AsyncGenerator, Dict, Generator
+import os 
+from dotenv import load_dotenv
 load_dotenv()
-
-
 class v2:
     def __init__(
         self,
@@ -17,7 +13,7 @@ class v2:
         proxies: dict = {},
     ):  
         self.session = requests.Session()
-        self.chat_endpoint = os.getenv("v2")
+        self.chat_endpoint = os.get("v2")
         self.stream_chunk_size = 64
         self.timeout = timeout
         self.last_response = {}
@@ -31,14 +27,13 @@ class v2:
         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}"
+    ) -> Generator[Dict[str, Any], None, None]:
+        conversation_prompt = f"Tell me everything about {prompt}"
 
         self.session.headers.update(self.headers)
         payload = {
@@ -53,75 +48,49 @@ class v2:
             "contexts_from": "google"
         }
 
-        def for_stream():
-            response = self.session.post(
-                self.chat_endpoint, json=payload, stream=True, timeout=self.timeout
+        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}"
             )
-            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()
+        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 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()
+    ) -> Generator[str, None, None]:
+        for response in self.ask(prompt, True):
+            yield self.format_text(self.get_message(response))
+        yield "[DONE]"
 
     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)
+    def format_text(self, text: str) -> str:
+        # Convert *text* to <i>text</i> for italic
+        text = re.sub(r'\*(.*?)\*', r'<i>\1</i>', text)
+        return text