From b5897ec51d5417e48efcc0f71ca22e66121b9f78 Mon Sep 17 00:00:00 2001
From: Ion Mincu <ion.mincu@uipath.com>
Date: Thu, 12 Jun 2025 12:21:16 +0300
Subject: [PATCH] fix(trace): fix traces

---
 pyproject.toml                                  | 2 +-
 src/uipath_llamaindex/_cli/_runtime/_runtime.py | 9 +++++++--
 uv.lock                                         | 9 +++++----
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ea8ed35..2e2ea34 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,7 +7,7 @@ requires-python = ">=3.10"
 dependencies = [
     "llama-index>=0.12.38",
     "openinference-instrumentation-llama-index>=4.3.0",
-    "uipath>=2.0.63",
+    "uipath>=2.0.64",
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",
diff --git a/src/uipath_llamaindex/_cli/_runtime/_runtime.py b/src/uipath_llamaindex/_cli/_runtime/_runtime.py
index 3cfe9ee..42a8176 100644
--- a/src/uipath_llamaindex/_cli/_runtime/_runtime.py
+++ b/src/uipath_llamaindex/_cli/_runtime/_runtime.py
@@ -12,7 +12,7 @@
     JsonPickleSerializer,
 )
 from llama_index.core.workflow.handler import WorkflowHandler
-from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
+from openinference.instrumentation.llama_index import LlamaIndexInstrumentor, get_current_span
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
@@ -29,6 +29,8 @@
 from ._exception import UiPathLlamaIndexRuntimeError
 from ._hitl import HitlProcessor, HitlReader
 
+from uipath.tracing import TracingManager
+
 logger = logging.getLogger(__name__)
 
 
@@ -54,15 +56,18 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         await self.validate()
 
         self.trace_provider = TracerProvider()
+        self.tracer = self.trace_provider.get_tracer("uipath.llamaindex.runtime")
 
         with suppress(Exception):
             trace.set_tracer_provider(self.trace_provider)
             self.trace_provider.add_span_processor(
                 BatchSpanProcessor(LlamaIndexExporter())
-            )  # type: ignore
+            )
 
             LlamaIndexInstrumentor().instrument(tracer_provider=self.trace_provider)
 
+            TracingManager.register_current_span_provider(get_current_span)
+
         try:
             if self.context.resume is False and self.context.job_id is None:
                 # Delete the previous graph state file at debug time
diff --git a/uv.lock b/uv.lock
index 451784c..91ca33f 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1,4 +1,5 @@
 version = 1
+revision = 1
 requires-python = ">=3.10"
 resolution-markers = [
     "python_full_version >= '3.12'",
@@ -2818,7 +2819,7 @@ wheels = [
 
 [[package]]
 name = "uipath"
-version = "2.0.63"
+version = "2.0.64"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "azure-monitor-opentelemetry" },
@@ -2832,9 +2833,9 @@ dependencies = [
     { name = "tenacity" },
     { name = "tomli" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/a2/bb/9d921247c7a154699f0b62e5e9acd67b8848441d03fed7345f986371c68c/uipath-2.0.63.tar.gz", hash = "sha256:74a6aef0d0889f93d5b0d2bb3be498e45b589fd71918fa5e44a2fdeab92ab571", size = 1833004 }
+sdist = { url = "https://files.pythonhosted.org/packages/20/8e/1f660202ff18c244012e732bc51f23767536c2e5debc119120ea31b7f19e/uipath-2.0.64.tar.gz", hash = "sha256:f77cd254384a03dd9475f0daa152f5d92e409a1c2b76bc866edb982b071fd988", size = 1840220 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/d7/4f/5afe52f10ecc645acd87a6e9d4b117cb2a7e938d0b6574a933d1703b2c9f/uipath-2.0.63-py3-none-any.whl", hash = "sha256:06f89864ba287638f568938e2654294456369e87fc006e7ca8925b0044a40128", size = 117341 },
+    { url = "https://files.pythonhosted.org/packages/1d/1a/f22e837b8528cf067fb700975a8fdd7403d95c0e7f9ffe9ffae5a276f155/uipath-2.0.64-py3-none-any.whl", hash = "sha256:8ec14ee26a96aa1fad8b19d137013ab5fd599e867f45c11c4643922c0bdcd223", size = 123261 },
 ]
 
 [[package]]
@@ -2862,7 +2863,7 @@ dev = [
 requires-dist = [
     { name = "llama-index", specifier = ">=0.12.38" },
     { name = "openinference-instrumentation-llama-index", specifier = ">=4.3.0" },
-    { name = "uipath", specifier = ">=2.0.63" },
+    { name = "uipath", specifier = ">=2.0.64" },
 ]
 
 [package.metadata.requires-dev]