WIP: emit all events on debug

Author: radu-mocanu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-llamaindex"
-version = "0.0.28"
+version = "0.0.29"
 description = "UiPath LlamaIndex SDK"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
@@ -9,7 +9,7 @@ dependencies = [
     "llama-index-embeddings-azure-openai>=0.3.8",
     "llama-index-llms-azure-openai>=0.3.2",
     "openinference-instrumentation-llama-index>=4.3.0",
-    "uipath>=2.0.71, <2.1.0",
+    "uipath>=3.0.0.dev1004330509,<3.0.0.dev1004340000"
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",
@@ -91,3 +91,5 @@ url = "https://test.pypi.org/simple/"
 publish-url = "https://test.pypi.org/legacy/"
 explicit = true
 
+[tool.uv.sources]
+uipath = { index = "testpypi" }

src/uipath_llamaindex/_cli/_runtime/_runtime.py

@@ -2,14 +2,15 @@
 import logging
 import os
 import pickle
+from cgitb import handler
 from contextlib import suppress
 from typing import Optional, cast
 
 from llama_index.core.workflow import (
     Context,
     HumanResponseEvent,
     InputRequiredEvent,
-    JsonPickleSerializer,
+    JsonPickleSerializer, StartEvent,
 )
 from llama_index.core.workflow.errors import WorkflowTimeoutError
 from llama_index.core.workflow.handler import WorkflowHandler
@@ -87,30 +88,12 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
             if self.context.workflow_context is None:
                 return None
 
-            handler: WorkflowHandler = self.context.workflow.run(
-                start_event=ev if self.context.resume else None,
-                ctx=self.context.workflow_context,
-                **(self.context.input_json or {}),
-            )
-
             resume_trigger: Optional[UiPathResumeTrigger] = None
 
-            response_applied = False
-            async for event in handler.stream_events():
-                # log the received event on trace level
-                if isinstance(event, InputRequiredEvent):
-                    # for api trigger hitl scenarios only pass the str input for processing
-                    hitl_processor = HitlProcessor(value=event.prefix)
-                    if self.context.resume and not response_applied:
-                        # If we are resuming, we need to apply the response to the event stream.
-                        response_applied = True
-                        response_event = await self.get_response_event()
-                        if response_event:
-                            # If we have a response event, send it to the workflow context.
-                            self.context.workflow_context.send_event(response_event)
-                    else:
-                        resume_trigger = await hitl_processor.create_resume_trigger()
-                        break
+            if self.context.debug and os.getenv("VSCODE_EXPLORER_EXECUTION") == "True":
+                handler = await self.debug_workflow()
+            else:
+                handler, resume_trigger = await self.run_workflow(ev)
 
             if resume_trigger is None:
                 try:
@@ -179,6 +162,63 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         finally:
             self.trace_provider.shutdown()
 
+    async def debug_workflow(self):
+        try:
+            import debugpy
+        except ImportError:
+            raise Exception("DebugPy not available")
+
+        handler: WorkflowHandler = self.context.workflow.run(
+            stepwise=True,
+            **(self.context.input_json or {}),
+        )
+        # if we are in debug mode, emit all events
+        while produced_events := await handler.run_step():
+            for ev in produced_events:
+                debug_file_path = os.environ.get('VSCODE_EXTENSION_TEMP_DEBUG_FILE')
+                if debug_file_path:
+                    # Write the debug data
+                    debug_data = {
+                        'llama_event': self._serialize_object(ev)
+                    }
+                    with open(debug_file_path, 'w') as f:
+                        json.dump(debug_data, f)
+
+                    # Trigger breakpoint
+                    debugpy.breakpoint()
+
+                # Send the event
+                handler.ctx.send_event(ev)
+        return handler
+
+    async def run_workflow(self, start_event :type[StartEvent]) -> (WorkflowHandler, Optional[UiPathResumeTrigger]):
+        resume_trigger: Optional[UiPathResumeTrigger] = None
+
+        handler: WorkflowHandler = self.context.workflow.run(
+            start_event=start_event if self.context.resume else None,
+            ctx=self.context.workflow_context,
+            **(self.context.input_json or {}),
+        )
+
+        response_applied = False
+        async for event in handler.stream_events():
+            # log the received event on trace level
+            if isinstance(event, InputRequiredEvent):
+                # for api trigger hitl scenarios only pass the str input for processing
+                hitl_processor = HitlProcessor(value=event.prefix)
+                if self.context.resume and not response_applied:
+                    # If we are resuming, we need to apply the response to the event stream.
+                    response_applied = True
+                    response_event = await self.get_response_event()
+                    if response_event:
+                        # If we have a response event, send it to the workflow context.
+                        self.context.workflow_context.send_event(response_event)
+                else:
+                    resume_trigger = await hitl_processor.create_resume_trigger()
+                    break
+
+        return handler, resume_trigger
+
     async def validate(self) -> None:
         """Validate runtime inputs and load Llama agent configuration."""
         try:

src/uipath_llamaindex/_cli/cli_run.py

@@ -6,18 +6,18 @@
 from dotenv import load_dotenv
 from uipath._cli._runtime._contracts import UiPathTraceContext
 from uipath._cli.middlewares import MiddlewareResult
-
 from ._runtime._context import UiPathLlamaIndexRuntimeContext
 from ._runtime._exception import UiPathLlamaIndexRuntimeError
 from ._runtime._runtime import UiPathLlamaIndexRuntime
 from ._utils._config import LlamaIndexConfig
+from uipath._cli._utils._common import serialize_object
 
 logger = logging.getLogger(__name__)
 load_dotenv()
 
 
 def llamaindex_run_middleware(
-    entrypoint: Optional[str], input: Optional[str], resume: bool
+    entrypoint: Optional[str], input: Optional[str], resume: bool, **kwargs
 ) -> MiddlewareResult:
     """Middleware to handle LlamaIndex agent execution"""
 
@@ -37,6 +37,7 @@ async def execute():
             context.entrypoint = entrypoint
             context.input = input
             context.resume = resume
+            context.debug = kwargs.get("debug", False)
             context.logs_min_level = env.get("LOG_LEVEL", "INFO")
             context.job_id = env.get("UIPATH_JOB_KEY")
             context.trace_id = env.get("UIPATH_TRACE_ID")
@@ -59,9 +60,9 @@ async def execute():
             async with UiPathLlamaIndexRuntime.from_context(context) as runtime:
                 return await runtime.execute()
 
-        asyncio.run(execute())
+        result = asyncio.run(execute())
 
-        return MiddlewareResult(should_continue=False, error_message=None)
+        return MiddlewareResult(should_continue=False, error_message=None, output=serialize_object(result.output))
 
     except UiPathLlamaIndexRuntimeError as e:
         return MiddlewareResult(