pygpt-net 2.1.37__py3-none-any.whl → 2.4.48__py3-none-any.whl

pygpt_net/controller/assistant/threads.py

@@ -6,35 +6,39 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2024.03.15 10:00:00                  #
+# Updated Date: 2024.12.14 08:00:00                  #
 # ================================================== #
 
 import json
 import time
+from typing import Optional, Any, List, Dict
 
 from PySide6.QtCore import QObject, Signal, Slot, QRunnable
 
+from pygpt_net.core.events import KernelEvent, RenderEvent
+from pygpt_net.core.bridge.context import BridgeContext
 from pygpt_net.item.ctx import CtxItem
-from pygpt_net.utils import trans
+from pygpt_net.utils import trans, get_image_extensions
 
 
-class Threads:
+class Threads(QObject):
     def __init__(self, window=None):
         """
         Assistant threads controller
 
         :param window: Window instance
         """
+        super(Threads, self).__init__()
         self.window = window
         self.started = False
-        self.stop = False
-        self.run_id = None
+        self.stop = False  # force stop run
+        self.run_id = None  # current run ID
         self.tool_calls = []  # list of previous tool calls
-        self.img_ext = ['jpg', 'jpeg', 'png', 'gif', 'bmp', 'tiff', 'webp']
+        self.img_ext = get_image_extensions()  # list, .png, .jpg, etc.
 
     def create_thread(self) -> str:
         """
-        Create assistant thread
+        Create assistant thread and store in context (stream and not stream)
 
         :return: thread id
         """
@@ -43,19 +47,38 @@ class Threads:
         self.window.core.ctx.append_thread(thread_id)
         return thread_id
 
-    def handle_output_message(self, ctx: CtxItem):
+    def handle_output_message(
+            self,
+            ctx: CtxItem,
+            stream: bool = False
+    ):
         """
-        Handle output message
+        Handle output message (not stream ONLY)
 
         :param ctx: CtxItem
+        :param stream: True if stream
         """
         # update ctx
         ctx.from_previous()  # append previous result if exists
         self.window.core.ctx.update_item(ctx)
-        self.window.controller.chat.output.handle(ctx, 'assistant', False)
+        self.window.controller.chat.output.handle(ctx, 'assistant', stream)
+
+        if stream:  # append all output to chat
+            event = RenderEvent(RenderEvent.END, {
+                "meta": ctx.meta,
+                "ctx": ctx,
+                "stream": stream,
+            })  # extra reload for stream markdown needed here
+            self.window.dispatch(event)
 
         ctx.clear_reply()  # reset results
-        self.window.controller.chat.output.handle_cmd(ctx)
+
+        self.log("Handling output message...")
+
+        has_cmd = self.window.core.command.has_cmds(ctx.output)
+        if has_cmd:
+            self.log("Handling message command...")
+            self.window.controller.chat.command.handle(ctx)
 
         # update ctx
         ctx.from_previous()  # append previous result again before save
@@ -67,59 +90,100 @@ class Threads:
         # update ctx list
         self.window.controller.ctx.update()
 
-    def handle_message_data(self, ctx: CtxItem, msg):
+        # if is command execute and not locked yet (not executing)
+        if has_cmd and self.window.controller.kernel.stack.waiting():
+            self.log("Replying for message command...")
+            self.window.controller.kernel.stack.handle()
+
+        self.log("Handled output message.")
+
+    def handle_message_data(
+            self,
+            ctx: CtxItem,
+            msg,
+            stream: bool = False
+    ):
         """
-        Handle message data
+        Handle message data (files, images, text) - stream and not-stream
 
         :param ctx: CtxItem
         :param msg: Message
+        :param stream: True if stream
         """
         paths = []
         file_ids = []
         images_ids = []
         mappings = {}  # file_id: path to sandbox
+        citations = []
 
         for content in msg.content:
             if content.type == "text":
-                ctx.set_output(content.text.value)
+                # text, append only if no stream or if is stream and already streamed output is empty
+                if not stream or (ctx.output is None or ctx.output == ""):
+                    ctx.set_output(content.text.value)
 
                 # annotations
                 if content.text.annotations:
                     self.log("Run: received annotations: {}".format(len(content.text.annotations)))
                     for item in content.text.annotations:
+                        # file
                         if item.type == "file_path":
                             file_id = item.file_path.file_id
                             file_ids.append(file_id)
                             mappings[file_id] = item.text
+                        # citation
+                        elif item.type == "file_citation":
+                            data = {
+                                "type": "assistant",  # store is remote
+                                "text": item.text,
+                                "start_idx": item.start_index,
+                                "end_idx": item.end_index,
+                            }
+                            if item.file_citation:
+                                data["file_id"] = item.file_citation.file_id
+                                data["quote"] = item.file_citation.quote
+                            citations.append(data)
 
             # image file
             elif content.type == "image_file":
                 self.log("Run: received image file")
                 images_ids.append(content.image_file.file_id)
 
-        # handle msg files
+        """
+        # handle message files  -- legacy, deprecated
         for file_id in msg.file_ids:
             if file_id not in images_ids and file_id not in file_ids:
+                self.log("Run: appending file id: {}".format(file_id))
                 file_ids.append(file_id)
+        """
 
         # handle content images
         if images_ids:
             image_paths = self.window.controller.assistant.files.handle_received_ids(images_ids, ".png")
             ctx.images = self.window.core.filesystem.make_local_list(list(image_paths))
 
-        # download msg files
+        # citations
+        if citations:
+            ctx.doc_ids = citations
+
+        # -- files --
+
+        # download message files
         paths += self.window.controller.assistant.files.handle_received_ids(file_ids)
         if paths:
             # convert to local paths
             local_paths = self.window.core.filesystem.make_local_list(list(paths))
-            text_msg = ctx.output
+            text_msg = ctx.output  # use current output
             if text_msg:
                 # map file ids to local paths
                 for file_id in mappings:
                     path_sandbox = mappings[file_id]
                     k = file_ids.index(file_id)
                     if len(local_paths) > k:
-                        text_msg = text_msg.replace(path_sandbox, local_paths[k])  # replace sandbox path
+                        text_msg = text_msg.replace(
+                            path_sandbox,
+                            local_paths[k]
+                        )  # replace sandbox file path with a local path
                 ctx.set_output(text_msg)
             ctx.files = local_paths
 
@@ -132,9 +196,12 @@ class Threads:
                 if img_files:
                     ctx.images = self.window.core.filesystem.make_local_list(list(img_files))
 
+        ctx.from_previous()  # append previous result again before save
+        self.window.core.ctx.update_item(ctx)
+
     def handle_messages(self, ctx: CtxItem):
         """
-        Handle run messages
+        Handle run messages (not stream ONLY)
 
         :param ctx: CtxItem
         """
@@ -142,17 +209,34 @@ class Threads:
         for msg in data:
             if msg.role == "assistant":
                 try:
-                    self.log("Run: handling message...")
+                    self.log("Run: handling non-stream message...")
                     self.handle_message_data(ctx, msg)
                 except Exception as e:
                     self.window.core.debug.log(e)
                     print("Run: handle message error:", e)
                 break
-        self.handle_output_message(ctx)  # send to chat
+        self.handle_output_message(ctx)
+
+    def handle_tool_calls_stream(self, run, ctx: CtxItem):
+        """
+        Handle tool calls (stream ONLY)
+
+        :param run: Run
+        :param ctx: CtxItem
+        """
+        self.stop = False
+        ctx.tool_calls = self.window.core.command.unpack_tool_calls(
+            run.required_action.submit_tool_outputs.tool_calls
+        )
+        event = KernelEvent(KernelEvent.STATE_IDLE, {
+            'msg': trans('assistant.run.func.call'),
+        })
+        self.window.dispatch(event)
+        self.handle_tool_calls(ctx)
 
     def handle_tool_calls(self, ctx: CtxItem):
         """
-        Handle tool calls
+        Handle tool calls (stream and non-stream)
 
         :param ctx: CtxItem
         """
@@ -162,13 +246,18 @@ class Threads:
         self.run_id = ctx.run_id
         self.tool_calls = ctx.tool_calls
 
+        has_calls = len(ctx.tool_calls) > 0
+
         # update ctx
         self.window.core.ctx.update_item(ctx)
 
         ctx.internal = True  # hide in chat input + handle synchronously
 
         self.window.controller.chat.output.handle(ctx, 'assistant', False)
-        self.window.controller.chat.output.handle_cmd(ctx)
+
+        if has_calls:
+            self.log("Handling tool call: {}.".format(ctx.tool_calls))
+            self.window.controller.chat.command.handle(ctx)            
 
         ctx.tool_calls = []  # clear tool calls
 
@@ -178,20 +267,60 @@ class Threads:
         # index ctx (llama-index)
         self.window.controller.idx.on_ctx_end(ctx, mode="assistant")
 
+        if has_calls and self.window.controller.kernel.stack.waiting():
+            self.log("Replying for tool call...")
+            self.window.controller.kernel.stack.handle()
+            return
+
         # check if there are any response, if not send empty response
+        # TODO: if native then there is no cmd response here and response will be sent as tool call result
+        # string response is only for message command response
+        
+        self.log("Sending response reply...")
+
         if not ctx.reply:
             results = {
                 "response": "",
             }
-            self.window.controller.chat.input.send(
-                text=json.dumps(results),
-                force=True,
-                reply=True,
-                internal=True,
-                prev_ctx=ctx,
-            )
+            context = BridgeContext()
+            context.ctx = ctx
+            context.prompt = json.dumps(results)
+            extra = {
+                "force": True,
+                "reply": True,
+                "internal": True,
+            }
+            event = KernelEvent(KernelEvent.INPUT_SYSTEM, {
+                'context': context,
+                'extra': extra,
+            })
+            self.window.dispatch(event)
+
+    def handle_tool_outputs(
+            self,
+            ctx: CtxItem
+    ) -> List[Dict[str, Any]]:
+        """
+        Handle tool outputs
 
-    def apply_outputs(self, ctx: CtxItem) -> list:
+        :param ctx: CtxItem
+        :return: list of tool calls outputs
+        """
+        tools_outputs = []
+        if self.window.controller.assistant.threads.is_running():
+            tools_outputs = self.window.controller.assistant.threads.apply_outputs(ctx)
+            self.window.controller.assistant.threads.reset()  # reset outputs
+            self.log("Appended tool outputs: {}".format(len(tools_outputs)))
+
+            # clear tool calls to prevent appending cmds to output (otherwise it will call commands again)
+            ctx.tool_calls = []
+        return tools_outputs
+
+
+    def apply_outputs(
+            self,
+            ctx: CtxItem
+    ) -> List[Dict[str, Any]]:
         """
         Apply tool call outputs
 
@@ -199,66 +328,149 @@ class Threads:
         :return: list of tool calls outputs
         """
         self.log("Run: preparing tool calls outputs...")
-
         ctx.run_id = self.run_id
         ctx.tool_calls = self.tool_calls  # set previous tool calls
         return self.window.core.command.get_tool_calls_outputs(ctx)
 
-    def handle_run_created(self, ctx: CtxItem, run):
+    def handle_run_created(
+            self,
+            ctx: CtxItem,
+            run,
+            stream: bool = False
+    ):
         """
-        Handle run created
+        Handle run created (stream and not stream)
 
         :param ctx: context item
         :param run
+        :param stream: True if stream
         """
         ctx.run_id = run.id
         ctx.current = False
         self.window.core.ctx.update_item(ctx)
         self.window.core.ctx.append_run(ctx.run_id)  # get run ID and store in ctx
-        self.handle_run(ctx)  # handle assistant run
+        self.handle_run(ctx, run, stream)  # handle assistant run
 
-    def handle_run_error(self, ctx: CtxItem, err):
+    def handle_run_error(self, ctx: CtxItem, err: Any):
         """
-        Handle run created
+        Handle run created (stream and not stream)
 
         :param ctx: context item
         :param err: error message
         """
+        stream = self.window.core.config.get('stream')
         ctx.current = False  # reset current state
+        ctx.from_previous()
+        event = RenderEvent(RenderEvent.END, {
+            "meta": ctx.meta,
+            "ctx": ctx,
+            "stream": stream,
+        })  # extra reload for stream markdown needed here
+        self.window.dispatch(event)
         self.window.core.ctx.update_item(ctx)
+        self.window.controller.ctx.update()
         self.window.core.debug.log(err)
         self.window.ui.dialogs.alert(err)
-
-    def handle_run(self, ctx: CtxItem):
+        self.window.update_status("An error occurred. Please try again.")
+        self.window.controller.chat.common.unlock_input()  # unlock input
+        # self.handle_stream_end(ctx)
+
+    def handle_run(
+            self,
+            ctx: CtxItem,
+            run,
+            stream: bool = False
+    ):
         """
-        Handle assistant's run
+        Handle assistant's run (not stream ONLY)
 
         :param ctx: CtxItem
+        :param run: Run
+        :param stream: True if stream
         """
+        self.reset()  # clear previous run
+
+        if stream:
+            self.log("Run: finishing stream...")
+            if run and run.usage is not None:
+                    ctx.input_tokens = run.usage.prompt_tokens
+                    ctx.output_tokens = run.usage.completion_tokens
+                    ctx.total_tokens = run.usage.total_tokens
+            self.handle_stream_end(ctx)  # unlock
+            if run:
+                if run.status == "requires_action":  # handle tool calls in stream mode
+                    self.handle_tool_calls_stream(run, ctx)
+            return # stop here, run is finished here in stream mode
+
+        # ---- not stream only ----
+
+        # run status listener
+        self.window.stateChanged.emit(self.window.STATE_BUSY)
+        self.log("Run: starting run worker...")
+
         # worker
         worker = RunWorker()
         worker.window = self.window
         worker.ctx = ctx
 
-        self.reset()  # clear previous run
-
         # signals
         worker.signals.updated.connect(self.handle_status)
         worker.signals.destroyed.connect(self.handle_destroy)
         worker.signals.started.connect(self.handle_started)
 
-        self.window.stateChanged.emit(self.window.STATE_BUSY)
-
-        self.log("Run: starting run worker...")
-
         # start
         self.window.threadpool.start(worker)
         self.started = True
 
+    def handle_stream_begin(self, ctx: CtxItem):
+        """
+        Handle stream end (stream ONLY)
+
+        :param ctx: CtxItem
+        """
+        self.window.stateChanged.emit(self.window.STATE_BUSY)
+
+    def handle_stream_end(self, ctx: CtxItem):
+        """
+        Handle stream end (stream ONLY)
+
+        :param ctx: CtxItem
+        """
+        self.window.controller.chat.common.unlock_input()  # unlock input
+        self.stop = False
+        event = KernelEvent(KernelEvent.STATE_IDLE, {
+            'msg': trans('assistant.run.completed'),
+        })
+        self.window.dispatch(event)
+        self.window.controller.chat.common.show_response_tokens(ctx)  # update tokens
+
+    def handle_status_error(self, ctx: CtxItem):
+        """
+        Handle status error (stream and not stream)
+
+        :param ctx: CtxItem
+        """
+        self.stop = False
+        self.window.controller.chat.common.unlock_input()
+        event = KernelEvent(KernelEvent.STATE_ERROR, {
+            'msg': trans('assistant.run.failed'),
+        })
+        self.window.dispatch(event)
+
+    def handle_run_step_created(self, ctx: CtxItem, stream: bool = False):
+        """
+        Handle run step created (stream and not stream)
+
+        :param ctx: CtxItem
+        :param stream: True if stream
+        """
+        self.window.stateChanged.emit(self.window.STATE_BUSY)
+        self.window.controller.chat.common.lock_input()  # lock input, show stop button
+
     @Slot(object, object)
     def handle_status(self, run, ctx: CtxItem):
         """
-        Handle status
+        Handle status (not stream ONLY)
 
         :param run: Run
         :param ctx: CtxItem
@@ -276,9 +488,11 @@ class Threads:
         if status == "completed":
             self.stop = False
             self.handle_messages(ctx)
-            self.window.statusChanged.emit(trans('assistant.run.completed'))
-            self.window.stateChanged.emit(self.window.STATE_IDLE)
-            self.window.controller.chat.output.show_response_tokens(ctx)  # update tokens
+            event = KernelEvent(KernelEvent.STATE_IDLE, {
+                'msg': trans('assistant.run.completed'),
+            })
+            self.window.dispatch(event)
+            self.window.controller.chat.common.show_response_tokens(ctx)  # update tokens
 
         # function call
         elif status == "requires_action":
@@ -286,16 +500,15 @@ class Threads:
             ctx.tool_calls = self.window.core.command.unpack_tool_calls(
                 run.required_action.submit_tool_outputs.tool_calls
             )
-            self.window.statusChanged.emit(trans('assistant.run.func.call'))
-            self.window.stateChanged.emit(self.window.STATE_IDLE)
+            event = KernelEvent(KernelEvent.STATE_IDLE, {
+                'msg': trans('assistant.run.func.call'),
+            })
+            self.window.dispatch(event)
             self.handle_tool_calls(ctx)
 
         # error
         elif status in ["failed", "cancelled", "expired", "cancelling"]:
-            self.stop = False
-            self.window.controller.chat.common.unlock_input()
-            self.window.statusChanged.emit(trans('assistant.run.failed'))
-            self.window.stateChanged.emit(self.window.STATE_ERROR)
+            self.handle_status_error(ctx)
 
     @Slot()
     def handle_destroy(self):
@@ -307,7 +520,10 @@ class Threads:
     def handle_started(self):
         """Handle listening started"""
         self.log("Run: assistant is listening status...")
-        self.window.statusChanged.emit(trans('assistant.run.listening'))
+        event = KernelEvent(KernelEvent.STATUS, {
+            'status': trans('assistant.run.listening'),
+        })
+        self.window.dispatch(event)
 
     def is_log(self) -> bool:
         """
@@ -338,7 +554,7 @@ class Threads:
 
     def is_running(self) -> bool:
         """
-        Check if tool calls need submit
+        Check if tool calls need submit (tools are running)
 
         :return: True if running
         """
@@ -346,12 +562,14 @@ class Threads:
 
 
 class RunSignals(QObject):
+    """Status check signals"""
     updated = Signal(object, object)
     destroyed = Signal()
     started = Signal()
 
 
 class RunWorker(QRunnable):
+    """Status check async worker"""
     def __init__(self, *args, **kwargs):
         super(RunWorker, self).__init__()
         self.signals = RunSignals()