Update v2.py

v2.py

@@ -5,7 +5,10 @@ import json
 from typing import Any, AsyncGenerator, Dict, Generator
 import os 
 from dotenv import load_dotenv
+
+# Load environment variables from .env file
 load_dotenv()
+
 class v2:
     def __init__(
         self,
@@ -77,8 +80,16 @@ class v2:
         prompt: str,
         stream: bool = False,
     ) -> Generator[str, None, None]:
+        buffer = ""
         for response in self.ask(prompt, True):
-            yield self.format_text(self.get_message(response))
+            text = self.get_message(response)
+            buffer += text
+            lines = buffer.split('\n')
+            for line in lines[:-1]:
+                yield self.format_text(line) + '\n\n'
+            buffer = lines[-1]
+        if buffer:
+            yield self.format_text(buffer) + '\n\n'
         yield "[DONE]"
 
     def get_message(self, response: dict) -> str: