feat(traced): instrument traceable with otel

Author: ionmincu

pyproject.toml

@@ -5,7 +5,7 @@ description = "UiPath Langchain"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
 dependencies = [
-    "uipath>=2.1.10, <2.2.0",
+    "uipath>=2.1.28.dev1005170000,<2.1.28.dev1005180000",
     "langgraph>=0.5.0, <0.7.0",
     "langchain-core>=0.3.34",
     "langgraph-checkpoint-sqlite>=2.0.3",
@@ -111,3 +111,7 @@ name = "testpypi"
 url = "https://test.pypi.org/simple/"
 publish-url = "https://test.pypi.org/legacy/"
 explicit = true
+
+
+[tool.uv.sources]
+uipath = { index = "testpypi" }

src/uipath_langchain/_cli/_runtime/_runtime.py

@@ -11,7 +11,10 @@
 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 openinference.instrumentation.langchain import (
+    LangChainInstrumentor,
+    get_current_span,
+)
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
@@ -20,6 +23,7 @@
     UiPathErrorCategory,
     UiPathRuntimeResult,
 )
+from uipath.tracing import TracingManager
 
 from ..._utils import _instrument_traceable_attributes
 from ...tracers import (
@@ -57,30 +61,43 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         Raises:
             LangGraphRuntimeError: If execution fails
         """
-        _instrument_traceable_attributes()
+
+        use_otel_ff = os.getenv("UIPATH_USE_OTEL_TRACING", "false").lower() == "true"
+
+        # Instrument LangSmith @traceable to use appropriate adapter
+        # When OTEL is enabled: uses UiPath @traced() 
+        # When OTEL is disabled: uses existing dispatch_event system
+        _instrument_traceable_attributes(useOtel=use_otel_ff)
 
         await self.validate()
 
-        with suppress(Exception):
-            provider = TracerProvider()
-            trace.set_tracer_provider(provider)
-            provider.add_span_processor(BatchSpanProcessor(LangchainExporter()))  # type: ignore
+        if use_otel_ff:
+            try:
+                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(
-                    JsonFileExporter(".uipath/traces.jsonl", LangchainSpanProcessor())
+                provider.add_span_processor(
+                    BatchSpanProcessor(
+                        SqliteExporter(".uipath/traces.db", LangchainSpanProcessor())
+                    )
                 )
-            )
 
-            provider.add_span_processor(
-                BatchSpanProcessor(
-                    SqliteExporter(".uipath/traces.db", LangchainSpanProcessor())
+                LangChainInstrumentor().instrument(
+                    tracer_provider=trace.get_tracer_provider()
                 )
-            )
 
-            LangChainInstrumentor().instrument(
-                tracer_provider=trace.get_tracer_provider()
-            )
+                TracingManager.register_current_span_provider(get_current_span)
+            except Exception as e:
+                logger.error(f"Failed to initialize OpenTelemetry tracing: {e}")
 
         if self.context.state_graph is None:
             return None
@@ -111,9 +128,10 @@ 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 not use_otel_ff:
+                    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/LangchainSpanProcessor.py

@@ -1,7 +1,9 @@
 import json
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, Dict, MutableMapping, override
+from typing import Any, Dict, MutableMapping
+
+from typing_extensions import override
 
 logger = logging.getLogger(__name__)
 
@@ -209,14 +211,25 @@ def extract_attributes(self, span_data: MutableMapping[str, Any]) -> Dict[str, A
         for key in ["Attributes", "attributes"]:
             if key in span_data:
                 value = span_data.pop(key)
-                return safe_parse_json(value) if isinstance(value, str) else value
-        return None
+                if isinstance(value, str):
+                    try:
+                        parsed_value = json.loads(value)
+                        return parsed_value if isinstance(parsed_value, dict) else {}
+                    except json.JSONDecodeError:
+                        logger.warning(f"Failed to parse attributes JSON: {value}")
+                        return {}
+                elif isinstance(value, dict):
+                    return value
+                else:
+                    return {}
+        return {}
 
     @override
     def process_span(self, span_data: MutableMapping[str, Any]) -> Dict[str, Any]:
+        logger.info(f"Processing span: {span_data}")
         attributes = self.extract_attributes(span_data)
 
-        if attributes:
+        if attributes and isinstance(attributes, dict):
             if "openinference.span.kind" in attributes:
                 # Remove the span kind attribute
                 span_type = attributes["openinference.span.kind"]

src/uipath_langchain/tracers/_instrument_traceable.py

@@ -7,7 +7,7 @@
 from typing import Any, Callable, Dict, List, Literal, Optional
 
 from langchain_core.callbacks import dispatch_custom_event
-from uipath.tracing import TracingManager
+from uipath.tracing import TracingManager, traced
 
 from ._events import CustomTraceEvents, FunctionCallEventData
 
@@ -387,12 +387,112 @@ def register_uipath_tracing():
 
 
 # Apply the patch
-def _instrument_traceable_attributes():
+def _map_traceable_to_traced_args(
+    run_type: Optional[str] = None,
+    name: Optional[str] = None,
+    tags: Optional[List[str]] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs: Any,
+) -> Dict[str, Any]:
+    """
+    Map LangSmith @traceable arguments to UiPath @traced() arguments.
+    
+    Args:
+        run_type: Function type (tool, chain, llm, retriever, etc.)
+        name: Custom name for the traced function
+        tags: List of tags for categorization
+        metadata: Additional metadata dictionary
+        **kwargs: Additional arguments (ignored)
+    
+    Returns:
+        Dict containing mapped arguments for @traced()
+    """
+    traced_args = {}
+    
+    # Direct mappings
+    if name is not None:
+        traced_args["name"] = name
+    
+    # Pass through run_type directly to UiPath @traced()
+    if run_type:
+        traced_args["run_type"] = run_type
+        
+    # For span_type, we can derive from run_type or use a default
+    if run_type:
+        # Map run_type to appropriate span_type for OpenTelemetry
+        span_type_mapping = {
+            "tool": "tool_call",
+            "chain": "chain_execution", 
+            "llm": "llm_call",
+            "retriever": "retrieval",
+            "embedding": "embedding",
+            "prompt": "prompt_template",
+            "parser": "output_parser"
+        }
+        traced_args["span_type"] = span_type_mapping.get(run_type, run_type)
+    
+    # Note: UiPath @traced() doesn't support custom attributes directly
+    # Tags and metadata information is lost in the current mapping
+    # This could be enhanced in future versions
+    
+    return traced_args
+
+
+def otel_traceable_adapter(
+    func: Optional[Callable] = None,
+    *,
+    run_type: Optional[str] = None,
+    name: Optional[str] = None,
+    tags: Optional[List[str]] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs: Any,
+):
+    """
+    OTEL-based adapter that converts LangSmith @traceable decorator calls to UiPath @traced().
+    
+    This function maintains the same interface as LangSmith's @traceable but uses
+    UiPath's OpenTelemetry-based tracing system underneath.
+    
+    Args:
+        func: Function to be decorated (when used without parentheses)
+        run_type: Type of function (tool, chain, llm, etc.)
+        name: Custom name for tracing
+        tags: List of tags for categorization
+        metadata: Additional metadata dictionary
+        **kwargs: Additional arguments (for future compatibility)
+        
+    Returns:
+        Decorated function or decorator function
+    """
+    def decorator(f: Callable) -> Callable:
+        # Map arguments to @traced() format
+        traced_args = _map_traceable_to_traced_args(
+            run_type=run_type,
+            name=name,
+            tags=tags,
+            metadata=metadata,
+            **kwargs
+        )
+        
+        # Apply UiPath @traced() decorator
+        return traced(**traced_args)(f)
+    
+    # Handle both @traceable and @traceable(...) usage patterns
+    if func is None:
+        # Called as @traceable(...) - return decorator
+        return decorator
+    else:
+        # Called as @traceable - apply decorator directly
+        return decorator(func)
+
+
+def _instrument_traceable_attributes(useOtel: bool = False):
     """Apply the patch to langsmith module at import time."""
     global original_langsmith, original_traceable
 
-    # Register our custom tracing decorator
-    register_uipath_tracing()
+    if not useOtel:
+        # Register our custom tracing decorator when not using opentelemetry
+        register_uipath_tracing()
 
     # Import the original module if not already done
     if original_langsmith is None:
@@ -408,7 +508,12 @@ def _instrument_traceable_attributes():
         original_traceable = original_langsmith.traceable
 
         # Replace the traceable function with our patched version
-        original_langsmith.traceable = patched_traceable
+        if useOtel:
+            # Use OTEL-based adapter when OTEL is enabled
+            original_langsmith.traceable = otel_traceable_adapter
+        else:
+            # Use existing dispatch_event-based adapter
+            original_langsmith.traceable = patched_traceable
 
         # Put our modified module back
         sys.modules["langsmith"] = original_langsmith