trace: add adapter

Author: ionmincu

pyproject.toml

@@ -7,7 +7,7 @@ requires-python = ">=3.10"
 dependencies = [
     "llama-index>=0.12.36",
     "openinference-instrumentation-llama-index>=4.2.1",
-    "uipath>=2.0.56",
+    "uipath>=2.0.59",
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",

src/uipath_llamaindex/_cli/_runtime/_runtime.py

@@ -5,20 +5,24 @@
 
 from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
 from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from uipath import UiPath
 from uipath._cli._runtime._contracts import (
     UiPathBaseRuntime,
     UiPathErrorCategory,
     UiPathRuntimeResult,
 )
-from uipath.tracing import wait_for_tracers
+from uipath.tracing import LlmOpsHttpExporter, wait_for_tracers
 
+from .._tracing._oteladapter import LlamaIndexAdapter
 from ._context import UiPathLlamaIndexRuntimeContext
 from ._exception import UiPathLlamaIndexRuntimeError
 
 logger = logging.getLogger(__name__)
 
 
+
 class UiPathLlamaIndexRuntime(UiPathBaseRuntime):
     """
     A runtime class for hosting UiPath LlamaIndex agents.
@@ -42,6 +46,9 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         await self.validate()
 
         with suppress(Exception):
+            trace.set_tracer_provider(TracerProvider())
+            trace.get_tracer_provider().add_span_processor(BatchSpanProcessor(LlamaIndexAdapter(LlmOpsHttpExporter())))  # type: ignore
+
             LlamaIndexInstrumentor().instrument(
                 tracer_provider=trace.get_tracer_provider()
             )

src/uipath_llamaindex/_cli/_tracing/__init__.py

@@ -0,0 +1,66 @@
+import json
+import logging
+from typing import Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import (
+    SpanExporter,
+    SpanExportResult,
+)
+
+logger = logging.getLogger(__name__)
+
+class LlamaIndexAdapter(SpanExporter):
+    """A simple wrapper for SpanExporters that allows for customization."""
+
+    # Mapping of old attribute names to new attribute names or (new name, function)
+    ATTRIBUTE_MAPPING = {
+        "input.value": ("input", lambda s: json.loads(s)),
+        "output.value": ("output", lambda s: json.loads(s)),
+        "llm.model_name": "model",
+    }
+
+    def __init__(self, wrapped_exporter: SpanExporter):
+        """Initialize with the exporter to wrap.
+        
+        Args:
+            wrapped_exporter: The underlying SpanExporter to wrap
+        """
+        super().__init__()
+        self.wrapped_exporter = wrapped_exporter
+        
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Export spans, with a hook to transform them first."""
+        transformed_spans = [self.transform_span(span) for span in spans]
+        return self.wrapped_exporter.export(transformed_spans)
+
+    def transform_span(self, span: ReadableSpan) -> ReadableSpan:
+        """Override this method to transform spans before export.
+        
+        Args:
+            span: The original span to transform
+        
+        Returns:
+            The transformed span
+        """
+        for old_key, mapping in self.ATTRIBUTE_MAPPING.items():
+            if old_key in span.attributes:
+                if isinstance(mapping, tuple):
+                    new_key, func = mapping
+                    try:
+                        span.attributes[new_key] = func(span.attributes[old_key])
+                    except Exception:
+                        span.attributes[new_key] = span.attributes[old_key]
+                else:
+                    new_key = mapping
+                    span.attributes[new_key] = span.attributes[old_key]
+                del span.attributes[old_key]
+        return span
+    
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Pass through to the wrapped exporter."""
+        return self.wrapped_exporter.force_flush(timeout_millis)
+    
+    def shutdown(self) -> None:
+        """Pass through to the wrapped exporter."""
+        return self.wrapped_exporter.shutdown()