Merge pull request #180 from UiPath/fix/otel_tracing_run

fix: use otel for traces

Author: cristipufu

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.54, <2.2.0",
+    "uipath>=2.1.56, <2.2.0",
     "langgraph>=0.5.0, <0.7.0",
     "langchain-core>=0.3.34",
     "langgraph-checkpoint-sqlite>=2.0.3",

src/uipath_langchain/_cli/_runtime/_runtime.py

@@ -6,7 +6,6 @@
 from langchain_core.callbacks.base import BaseCallbackHandler
 from langchain_core.messages import BaseMessage
 from langchain_core.runnables.config import RunnableConfig
-from langchain_core.tracers.langchain import wait_for_all_tracers
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
 from langgraph.errors import EmptyInputError, GraphRecursionError, InvalidUpdateError
 from langgraph.graph.state import CompiledStateGraph
@@ -16,8 +15,6 @@
     UiPathRuntimeResult,
 )
 
-from ..._utils import _instrument_traceable_attributes
-from ...tracers import AsyncUiPathTracer
 from .._utils._graph import LangGraphConfig
 from ._context import LangGraphRuntimeContext
 from ._conversation import map_message
@@ -48,13 +45,10 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         Raises:
             LangGraphRuntimeError: If execution fails
         """
-        _instrument_traceable_attributes()
 
         if self.context.state_graph is None:
             return None
 
-        tracer = None
-
         try:
             async with AsyncSqliteSaver.from_conn_string(
                 self.state_file_path
@@ -71,13 +65,8 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
 
                 processed_input = await input_processor.process()
 
-                # 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]
-
                 graph_config: RunnableConfig = {
                     "configurable": {
                         "thread_id": (
@@ -185,11 +174,7 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                 UiPathErrorCategory.USER,
             ) from e
         finally:
-            if tracer is not None:
-                await tracer.wait_for_all_tracers()
-
-            if self.context.langsmith_tracing_enabled:
-                wait_for_all_tracers()
+            pass
 
     async def validate(self) -> None:
         """Validate runtime inputs."""

src/uipath_langchain/_cli/cli_dev.py

@@ -10,6 +10,7 @@
 from uipath._cli._utils._console import ConsoleLogger
 from uipath._cli.middlewares import MiddlewareResult
 
+from .._tracing import _instrument_traceable_attributes
 from ._runtime._context import LangGraphRuntimeContext
 from ._runtime._runtime import LangGraphRuntime
 
@@ -24,6 +25,8 @@ def langgraph_dev_middleware(interface: Optional[str]) -> MiddlewareResult:
             runtime_factory = UiPathRuntimeFactory(
                 LangGraphRuntime, LangGraphRuntimeContext
             )
+
+            _instrument_traceable_attributes()
             runtime_factory.add_instrumentor(LangChainInstrumentor, get_current_span)
             app = UiPathDevTerminal(runtime_factory)
             asyncio.run(app.run_async())

src/uipath_langchain/_cli/cli_eval.py

@@ -1,7 +1,10 @@
 import asyncio
-from os import environ as env
 from typing import List, Optional
 
+from openinference.instrumentation.langchain import (
+    LangChainInstrumentor,
+    get_current_span,
+)
 from uipath._cli._evals._runtime import UiPathEvalContext, UiPathEvalRuntime
 from uipath._cli._runtime._contracts import (
     UiPathRuntimeFactory,
@@ -13,15 +16,15 @@
 from uipath_langchain._cli._runtime._context import LangGraphRuntimeContext
 from uipath_langchain._cli._runtime._runtime import LangGraphRuntime
 from uipath_langchain._cli._utils._graph import LangGraphConfig
+from uipath_langchain._tracing import (
+    LangChainExporter,
+    _instrument_traceable_attributes,
+)
 
 
 def langgraph_eval_middleware(
     entrypoint: Optional[str], eval_set: Optional[str], eval_ids: List[str], **kwargs
 ) -> MiddlewareResult:
-    # Add default env variables
-    env["UIPATH_REQUESTING_PRODUCT"] = "uipath-python-sdk"
-    env["UIPATH_REQUESTING_FEATURE"] = "langgraph-agent"
-
     config = LangGraphConfig()
     if not config.exists:
         return MiddlewareResult(
@@ -33,6 +36,8 @@ def langgraph_eval_middleware(
     eval_context.eval_ids = eval_ids
 
     try:
+        _instrument_traceable_attributes()
+
         runtime_entrypoint = entrypoint or auto_discover_entrypoint()
 
         def generate_runtime_context(
@@ -53,6 +58,11 @@ def generate_runtime_context(
             ),
         )
 
+        if eval_context.job_id:
+            runtime_factory.add_span_exporter(LangChainExporter())
+
+        runtime_factory.add_instrumentor(LangChainInstrumentor, get_current_span)
+
         async def execute():
             async with UiPathEvalRuntime.from_eval_context(
                 factory=runtime_factory, context=eval_context

src/uipath_langchain/_cli/cli_run.py

@@ -1,10 +1,17 @@
 import asyncio
 import os
-from os import environ as env
 from typing import Optional
 
+from openinference.instrumentation.langchain import (
+    LangChainInstrumentor,
+    get_current_span,
+)
+from uipath._cli._runtime._contracts import (
+    UiPathRuntimeFactory,
+)
 from uipath._cli.middlewares import MiddlewareResult
 
+from .._tracing import LangChainExporter, _instrument_traceable_attributes
 from ._runtime._exception import LangGraphRuntimeError
 from ._runtime._runtime import (  # type: ignore[attr-defined]
     LangGraphRuntime,
@@ -24,23 +31,35 @@ def langgraph_run_middleware(
         )  # Continue with normal flow if no langgraph.json
 
     try:
-        # Add default env variables
-        env["UIPATH_REQUESTING_PRODUCT"] = "uipath-python-sdk"
-        env["UIPATH_REQUESTING_FEATURE"] = "langgraph-agent"
-
         context = LangGraphRuntimeContext.with_defaults(**kwargs)
         context.langgraph_config = config
         context.entrypoint = entrypoint
         context.input = input
         context.resume = resume
 
+        _instrument_traceable_attributes()
+
+        def generate_runtime(ctx: LangGraphRuntimeContext) -> LangGraphRuntime:
+            runtime = LangGraphRuntime(ctx)
+            # If not resuming and no job id, delete the previous state file
+            if not ctx.resume and ctx.job_id is None:
+                if os.path.exists(runtime.state_file_path):
+                    os.remove(runtime.state_file_path)
+            return runtime
+
         async def execute():
-            async with LangGraphRuntime.from_context(context) as runtime:
-                if context.resume is False and context.job_id is None:
-                    # Delete the previous graph state file at debug time
-                    if os.path.exists(runtime.state_file_path):
-                        os.remove(runtime.state_file_path)
-                await runtime.execute()
+            runtime_factory = UiPathRuntimeFactory(
+                LangGraphRuntime,
+                LangGraphRuntimeContext,
+                runtime_generator=generate_runtime,
+            )
+
+            if context.job_id:
+                runtime_factory.add_span_exporter(LangChainExporter())
+
+            runtime_factory.add_instrumentor(LangChainInstrumentor, get_current_span)
+
+            await runtime_factory.execute(context)
 
         asyncio.run(execute())
 

src/uipath_langchain/_tracing/__init__.py

@@ -0,0 +1,4 @@
+from ._instrument_traceable import _instrument_traceable_attributes
+from ._oteladapter import LangChainExporter
+
+__all__ = ["LangChainExporter", "_instrument_traceable_attributes"]

src/uipath_langchain/_tracing/_instrument_traceable.py

@@ -0,0 +1,135 @@
+import importlib
+import logging
+import sys
+from typing import Any, Callable, Dict, List, Optional
+
+from uipath.tracing import traced
+
+# Original module and traceable function references
+original_langsmith: Any = None
+original_traceable: Any = None
+
+logger = logging.getLogger(__name__)
+
+
+# Apply the patch
+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[..., Any]] = 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[..., Any]) -> Callable[..., Any]:
+        # 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():
+    """Apply the patch to langsmith module at import time."""
+    global original_langsmith, original_traceable
+
+    # Import the original module if not already done
+    if original_langsmith is None:
+        # Temporarily remove our custom module from sys.modules
+        if "langsmith" in sys.modules:
+            original_langsmith = sys.modules["langsmith"]
+            del sys.modules["langsmith"]
+
+        # Import the original module
+        original_langsmith = importlib.import_module("langsmith")
+
+        # Store the original traceable
+        original_traceable = original_langsmith.traceable
+
+        # Replace the traceable function with our patched version
+        original_langsmith.traceable = otel_traceable_adapter
+
+        # Put our modified module back
+        sys.modules["langsmith"] = original_langsmith
+
+    return original_langsmith