feat(trace): add experimental otel trace

Author: ionmincu

src/uipath_langchain/_cli/_runtime/_runtime.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+from contextlib import suppress
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 from langchain_core.callbacks.base import BaseCallbackHandler
@@ -10,15 +11,24 @@
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
 from langgraph.errors import EmptyInputError, GraphRecursionError, InvalidUpdateError
 from langgraph.graph.state import CompiledStateGraph
+from openinference.instrumentation.langchain import LangChainInstrumentor
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from uipath._cli._runtime._contracts import (
     UiPathBaseRuntime,
     UiPathErrorCategory,
     UiPathRuntimeResult,
 )
 
 from ..._utils import _instrument_traceable_attributes
-from ...tracers import AsyncUiPathTracer
-from .._utils._graph import LangGraphConfig
+from ...tracers import (
+    AsyncUiPathTracer,
+    JsonFileExporter,
+    LangchainExporter,
+    SqliteExporter,
+)
+from ...tracers.LangchainSpanProcessor import LangchainSpanProcessor
 from ._context import LangGraphRuntimeContext
 from ._conversation import map_message
 from ._exception import LangGraphRuntimeError
@@ -50,6 +60,27 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         """
         _instrument_traceable_attributes()
 
+        with suppress(Exception):
+            provider = TracerProvider()
+            trace.set_tracer_provider(provider)
+            provider.add_span_processor(BatchSpanProcessor(LangchainExporter()))  # type: ignore
+
+            provider.add_span_processor(
+                BatchSpanProcessor(
+                    JsonFileExporter(".uipath/traces.jsonl", LangchainSpanProcessor())
+                )
+            )
+
+            provider.add_span_processor(
+                BatchSpanProcessor(
+                    SqliteExporter(".uipath/traces.db", LangchainSpanProcessor())
+                )
+            )
+
+            LangChainInstrumentor().instrument(
+                tracer_provider=trace.get_tracer_provider()
+            )
+
         if self.context.state_graph is None:
             return None
 
@@ -74,9 +105,9 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                 # Set up tracing if available
                 callbacks: List[BaseCallbackHandler] = []
 
-                if self.context.job_id and self.context.tracing_enabled:
-                    tracer = AsyncUiPathTracer(context=self.context.trace_context)
-                    callbacks = [tracer]
+                # if self.context.job_id and self.context.tracing_enabled:
+                #     tracer = AsyncUiPathTracer(context=self.context.trace_context)
+                #     callbacks = [tracer]
 
                 graph_config: RunnableConfig = {
                     "configurable": {

src/uipath_langchain/tracers/JsonFileExporter.py

@@ -0,0 +1,63 @@
+import json
+import logging
+import os
+from typing import Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from uipath.tracing._utils import _SpanUtils
+
+from .LangchainSpanProcessor import BaseSpanProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class JsonFileExporter(SpanExporter):
+    """
+    An exporter that writes spans to a file in JSON Lines format.
+
+    This exporter is useful for debugging and local development. It serializes
+    each span to a JSON object and appends it as a new line in the specified
+    file.
+    """
+
+    def __init__(self, file_path: str, processor: BaseSpanProcessor):
+        """
+        Initializes the JsonFileExporter.
+
+        Args:
+            file_path: The path to the JSON file where spans will be written.
+        """
+        self.file_path = file_path
+        self._processor = processor
+        # Ensure the directory exists
+        os.makedirs(os.path.dirname(self.file_path), exist_ok=True)
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """
+        Exports a batch of spans.
+
+        Args:
+            spans: A sequence of ReadableSpan objects.
+
+        Returns:
+            The result of the export operation.
+        """
+        try:
+            uipath_spans = [
+                _SpanUtils.otel_span_to_uipath_span(span).to_dict() for span in spans
+            ]
+            processed_spans = [
+                self._processor.process_span(span) for span in uipath_spans
+            ]
+            with open(self.file_path, "a") as f:
+                for span in processed_spans:
+                    f.write(json.dumps(span) + "\n")
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            logger.error(f"Failed to export spans to {self.file_path}: {e}")
+            return SpanExportResult.FAILURE
+
+    def shutdown(self) -> None:
+        """Shuts down the exporter."""
+        pass

src/uipath_langchain/tracers/LangchainExporter.py

@@ -0,0 +1,27 @@
+import logging
+from typing import Any, Dict
+
+from opentelemetry.sdk.trace.export import (
+    SpanExportResult,
+)
+from uipath.tracing import LlmOpsHttpExporter
+
+from .LangchainSpanProcessor import LangchainSpanProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class LangchainExporter(LlmOpsHttpExporter):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._processor = LangchainSpanProcessor()
+
+    def _send_with_retries(
+        self, url: str, payload: list[Dict[str, Any]], max_retries: int = 4
+    ) -> SpanExportResult:
+        processed_payload = [self._processor.process_span(span) for span in payload]
+        return super()._send_with_retries(
+            url=url,
+            payload=processed_payload,
+            max_retries=max_retries,
+        )