Merge branch 'main' into fix-assets

src/pieces/app.py

@@ -7,7 +7,7 @@
 from pieces.autocommit import *
 from pieces.copilot import *
 
-from . import __version__
+from pieces import __version__
 ask_stream = AskStream()
 
 class PiecesCLI:

src/pieces/copilot/ask_command.py

@@ -8,6 +8,7 @@
 
 from pieces.gui import show_error
 from pieces.wrapper.basic_identifier.chat import BasicChat
+from pieces.wrapper.websockets.ask_ws import AskStreamWS
 
 if TYPE_CHECKING:
     from pieces_os_client.models.qgpt_stream_output import QGPTStreamOutput
@@ -28,7 +29,7 @@ def on_message(self, response:"QGPTStreamOutput"):
                     text = answer.text
                     self.final_answer += text
                     if text:
-                        self.live.update(Markdown(self.final_answer))
+                        self.live.update(Markdown(self.final_answer), refresh=True)
 
             if response.status == 'COMPLETED':
                 self.live.update(Markdown(self.final_answer), refresh=True)
@@ -42,7 +43,7 @@ def on_message(self, response:"QGPTStreamOutput"):
 
 
     def add_context(self,files,assets_index):
-        
+
         context = Settings.pieces_client.copilot.context
 
         # Files
@@ -54,7 +55,7 @@ def add_context(self,files,assets_index):
                 if os.path.exists(file): # check if file exists
                     show_error(f"{file} is not found","Please enter a valid file path")
                     return
-                
+
                 context.paths.append(os.path.abspath(file)) # Return the abs path
         # snippets
         if assets_index:
@@ -79,6 +80,9 @@ def ask(self,query, **kwargs):
         finishes = self.message_compeleted.wait(Settings.TIMEOUT)
         self.message_compeleted.clear()
 
+        if not Settings.run_in_loop:
+            AskStreamWS.instance.close() # Close the websocket if we are not run in loop
+
         if not finishes and not self.live:
             raise ConnectionError("Failed to get the reponse back")
         return self.final_answer

src/pieces/wrapper/websockets/ask_ws.py

@@ -71,5 +71,5 @@ def send_message(self, message: "QGPTStreamInput"):
 				raise WebSocketConnectionClosedException()
 			self.ws.send(message.to_json())
 		except WebSocketConnectionClosedException:
-			self.on_open = lambda ws: ws.send(message.to_json())  # Send the message on opening
+			self.on_open_callback = lambda ws: ws.send(message.to_json())  # Send the message on opening
 			self.start()  # Start a new WebSocket since we are not connected to any