fix: refactor span attribute setting to rm code duplication

Author: andrei-rusu

src/uipath/core/tracing/_utils.py

@@ -6,9 +6,10 @@
 from dataclasses import asdict, is_dataclass
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any, Mapping
+from typing import Any, Mapping, Optional
 from zoneinfo import ZoneInfo
 
+from opentelemetry.trace import Span
 from pydantic import BaseModel
 
 
@@ -30,7 +31,7 @@ def get_supported_params(
     return supported
 
 
-def _simple_serialize_defaults(obj):
+def _simple_serialize_defaults(obj: Any) -> Any:
     # Handle Pydantic BaseModel instances
     if hasattr(obj, "model_dump") and not isinstance(obj, type):
         return obj.model_dump(exclude_none=True, mode="json")
@@ -56,8 +57,8 @@ def _simple_serialize_defaults(obj):
         return _simple_serialize_defaults(obj.value)
 
     if isinstance(obj, (set, tuple)):
-        if hasattr(obj, "_asdict") and callable(obj._asdict):
-            return obj._asdict()
+        if hasattr(obj, "_asdict") and callable(obj._asdict):  # pyright: ignore[reportAttributeAccessIssue]
+            return obj._asdict()  # pyright: ignore[reportAttributeAccessIssue]
         return list(obj)
 
     if isinstance(obj, datetime):
@@ -120,3 +121,68 @@ def format_args_for_trace(
         return result
     except Exception:
         return {"args": args, "kwargs": kwargs}
+
+
+def set_span_input_attributes(
+    span: Span,
+    trace_name: str,
+    wrapped_func: Callable,
+    args: Any,
+    kwargs: Any,
+    span_type: str,
+    run_type: Optional[str],
+    input_processor: Optional[Callable[..., Any]],
+) -> None:
+    """Set span attributes for metadata and inputs before function execution.
+
+    This should be called BEFORE the wrapped function executes to ensure
+    input context is captured even if the function raises an exception.
+
+    Args:
+        span: The OpenTelemetry span to set attributes on
+        trace_name: Name of the trace/span
+        wrapped_func: The function being traced
+        args: Positional arguments passed to the function
+        kwargs: Keyword arguments passed to the function
+        span_type: Span type categorization (set to "TOOL" for OpenInference tool calls)
+        run_type: Optional run type categorization
+        input_processor: Optional function to process inputs before recording
+    """
+    is_tool = span_type and span_type.upper() == "TOOL"
+    if is_tool:
+        span.set_attribute("openinference.span.kind", "TOOL")
+        span.set_attribute("tool.name", trace_name)
+        span.set_attribute("span_type", "TOOL")
+    else:
+        span.set_attribute("span_type", span_type)
+
+    if run_type is not None:
+        span.set_attribute("run_type", run_type)
+
+    inputs = format_args_for_trace_json(
+        inspect.signature(wrapped_func), *args, **kwargs
+    )
+    if input_processor:
+        processed_inputs = input_processor(json.loads(inputs))
+        inputs = json.dumps(processed_inputs, default=str)
+    span.set_attribute("input.mime_type", "application/json")
+    span.set_attribute("input.value", inputs)
+
+
+def set_span_output_attributes(
+    span: Span,
+    result: Any,
+    output_processor: Optional[Callable[..., Any]],
+) -> None:
+    """Set span attributes for outputs after function execution.
+
+    This should be called AFTER the wrapped function executes successfully.
+
+    Args:
+        span: The OpenTelemetry span to set attributes on
+        result: The result from the function execution
+        output_processor: Optional function to process outputs before recording
+    """
+    output = output_processor(result) if output_processor else result
+    span.set_attribute("output.value", format_object_for_trace_json(output))
+    span.set_attribute("output.mime_type", "application/json")

src/uipath/core/tracing/decorators.py

@@ -1,7 +1,6 @@
 """Tracing decorators for function instrumentation."""
 
 import inspect
-import json
 import logging
 import random
 from functools import wraps
@@ -12,9 +11,9 @@
 from opentelemetry.trace.status import StatusCode
 
 from uipath.core.tracing._utils import (
-    format_args_for_trace_json,
-    format_object_for_trace_json,
     get_supported_params,
+    set_span_input_attributes,
+    set_span_output_attributes,
 )
 from uipath.core.tracing.span_utils import UiPathSpanUtils
 
@@ -45,7 +44,7 @@ def _opentelemetry_traced(
         recording: If False, span is not recorded
     """
 
-    def decorator(func):
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         trace_name = name or func.__name__
 
         def get_span():
@@ -78,38 +77,30 @@ def get_span():
 
         # --------- Sync wrapper ---------
         @wraps(func)
-        def sync_wrapper(*args, **kwargs):
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
             span_cm, span = get_span()
             try:
-                # Check if this should be treated as a tool call
-                is_tool = span_type and span_type.upper() == "TOOL"
-
-                if is_tool:
-                    # Set OpenInference tool call attributes
-                    span.set_attribute("openinference.span.kind", "TOOL")
-                    span.set_attribute("tool.name", trace_name)
-                    span.set_attribute("span_type", "TOOL")
-                else:
-                    span.set_attribute("span_type", span_type or "function_call_sync")
-
-                if run_type is not None:
-                    span.set_attribute("run_type", run_type)
-
-                inputs = format_args_for_trace_json(
-                    inspect.signature(func), *args, **kwargs
+                # Set input attributes BEFORE execution
+                set_span_input_attributes(
+                    span,
+                    trace_name=trace_name,
+                    wrapped_func=func,
+                    args=args,
+                    kwargs=kwargs,
+                    run_type=run_type,
+                    span_type=span_type or "function_call_sync",
+                    input_processor=input_processor,
                 )
-                if input_processor:
-                    processed_inputs = input_processor(json.loads(inputs))
-                    inputs = json.dumps(processed_inputs, default=str)
-
-                span.set_attribute("input.mime_type", "application/json")
-                span.set_attribute("input.value", inputs)
 
+                # Execute the function
                 result = func(*args, **kwargs)
-                output = output_processor(result) if output_processor else result
 
-                span.set_attribute("output.value", format_object_for_trace_json(output))
-                span.set_attribute("output.mime_type", "application/json")
+                # Set output attributes AFTER execution
+                set_span_output_attributes(
+                    span,
+                    result=result,
+                    output_processor=output_processor,
+                )
                 return result
             except Exception as e:
                 span.record_exception(e)
@@ -121,38 +112,30 @@ def sync_wrapper(*args, **kwargs):
 
         # --------- Async wrapper ---------
         @wraps(func)
-        async def async_wrapper(*args, **kwargs):
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
             span_cm, span = get_span()
             try:
-                # Check if this should be treated as a tool call
-                is_tool = span_type and span_type.upper() == "TOOL"
-
-                if is_tool:
-                    # Set OpenInference tool call attributes
-                    span.set_attribute("openinference.span.kind", "TOOL")
-                    span.set_attribute("tool.name", trace_name)
-                    span.set_attribute("span_type", "TOOL")
-                else:
-                    span.set_attribute("span_type", span_type or "function_call_async")
-
-                if run_type is not None:
-                    span.set_attribute("run_type", run_type)
-
-                inputs = format_args_for_trace_json(
-                    inspect.signature(func), *args, **kwargs
+                # Set input attributes BEFORE execution
+                set_span_input_attributes(
+                    span,
+                    trace_name=trace_name,
+                    wrapped_func=func,
+                    args=args,
+                    kwargs=kwargs,
+                    run_type=run_type,
+                    span_type=span_type or "function_call_async",
+                    input_processor=input_processor,
                 )
-                if input_processor:
-                    processed_inputs = input_processor(json.loads(inputs))
-                    inputs = json.dumps(processed_inputs, default=str)
-
-                span.set_attribute("input.mime_type", "application/json")
-                span.set_attribute("input.value", inputs)
 
+                # Execute the function
                 result = await func(*args, **kwargs)
-                output = output_processor(result) if output_processor else result
 
-                span.set_attribute("output.value", format_object_for_trace_json(output))
-                span.set_attribute("output.mime_type", "application/json")
+                # Set output attributes AFTER execution
+                set_span_output_attributes(
+                    span,
+                    result=result,
+                    output_processor=output_processor,
+                )
                 return result
             except Exception as e:
                 span.record_exception(e)
@@ -164,42 +147,34 @@ async def async_wrapper(*args, **kwargs):
 
         # --------- Generator wrapper ---------
         @wraps(func)
-        def generator_wrapper(*args, **kwargs):
+        def generator_wrapper(*args: Any, **kwargs: Any) -> Any:
             span_cm, span = get_span()
             try:
-                # Check if this should be treated as a tool call
-                is_tool = span_type and span_type.upper() == "TOOL"
-
-                if is_tool:
-                    # Set OpenInference tool call attributes
-                    span.set_attribute("openinference.span.kind", "TOOL")
-                    span.set_attribute("tool.name", trace_name)
-                    span.set_attribute("span_type", "TOOL")
-                else:
-                    span.set_attribute(
-                        "span_type", span_type or "function_call_generator_sync"
-                    )
-
-                if run_type is not None:
-                    span.set_attribute("run_type", run_type)
-
-                inputs = format_args_for_trace_json(
-                    inspect.signature(func), *args, **kwargs
+                # Set input attributes BEFORE execution
+                set_span_input_attributes(
+                    span,
+                    trace_name=trace_name,
+                    wrapped_func=func,
+                    args=args,
+                    kwargs=kwargs,
+                    run_type=run_type,
+                    span_type=span_type or "function_call_generator_sync",
+                    input_processor=input_processor,
                 )
-                if input_processor:
-                    processed_inputs = input_processor(json.loads(inputs))
-                    inputs = json.dumps(processed_inputs, default=str)
-                span.set_attribute("input.mime_type", "application/json")
-                span.set_attribute("input.value", inputs)
 
+                # Execute the generator and collect outputs
                 outputs = []
                 for item in func(*args, **kwargs):
                     outputs.append(item)
                     span.add_event(f"Yielded: {item}")
                     yield item
-                output = output_processor(outputs) if output_processor else outputs
-                span.set_attribute("output.value", format_object_for_trace_json(output))
-                span.set_attribute("output.mime_type", "application/json")
+
+                # Set output attributes AFTER execution
+                set_span_output_attributes(
+                    span,
+                    result=outputs,
+                    output_processor=output_processor,
+                )
             except Exception as e:
                 span.record_exception(e)
                 span.set_status(StatusCode.ERROR, str(e))
@@ -210,44 +185,34 @@ def generator_wrapper(*args, **kwargs):
 
         # --------- Async generator wrapper ---------
         @wraps(func)
-        async def async_generator_wrapper(*args, **kwargs):
+        async def async_generator_wrapper(*args: Any, **kwargs: Any) -> Any:
             span_cm, span = get_span()
             try:
-                # Check if this should be treated as a tool call
-                is_tool = span_type and span_type.upper() == "TOOL"
-
-                if is_tool:
-                    # Set OpenInference tool call attributes
-                    span.set_attribute("openinference.span.kind", "TOOL")
-                    span.set_attribute("tool.name", trace_name)
-                    span.set_attribute("span_type", "TOOL")
-                else:
-                    span.set_attribute(
-                        "span_type", span_type or "function_call_generator_async"
-                    )
-
-                if run_type is not None:
-                    span.set_attribute("run_type", run_type)
-
-                inputs = format_args_for_trace_json(
-                    inspect.signature(func), *args, **kwargs
+                # Set input attributes BEFORE execution
+                set_span_input_attributes(
+                    span,
+                    trace_name=trace_name,
+                    wrapped_func=func,
+                    args=args,
+                    kwargs=kwargs,
+                    run_type=run_type,
+                    span_type=span_type or "function_call_generator_async",
+                    input_processor=input_processor,
                 )
-                if input_processor:
-                    processed_inputs = input_processor(json.loads(inputs))
-                    inputs = json.dumps(processed_inputs, default=str)
-
-                span.set_attribute("input.mime_type", "application/json")
-                span.set_attribute("input.value", inputs)
 
+                # Execute the generator and collect outputs
                 outputs = []
                 async for item in func(*args, **kwargs):
                     outputs.append(item)
                     span.add_event(f"Yielded: {item}")
                     yield item
 
-                output = output_processor(outputs) if output_processor else outputs
-                span.set_attribute("output.value", format_object_for_trace_json(output))
-                span.set_attribute("output.mime_type", "application/json")
+                # Set output attributes AFTER execution
+                set_span_output_attributes(
+                    span,
+                    result=outputs,
+                    output_processor=output_processor,
+                )
             except Exception as e:
                 span.record_exception(e)
                 span.set_status(StatusCode.ERROR, str(e))
@@ -294,11 +259,11 @@ def traced(
     """
 
     # Apply default processors selectively based on hide flags
-    def _default_input_processor(inputs):
+    def _default_input_processor(inputs: Any) -> Any:
         """Default input processor that doesn't log any actual input data."""
         return {"redacted": "Input data not logged for privacy/security"}
 
-    def _default_output_processor(outputs):
+    def _default_output_processor(outputs: Any) -> Any:
         """Default output processor that doesn't log any actual output data."""
         return {"redacted": "Output data not logged for privacy/security"}
 
@@ -319,7 +284,7 @@ def _default_output_processor(outputs):
 
     tracer_impl = _opentelemetry_traced
 
-    def decorator(func):
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         # Check which parameters are supported by the tracer_impl
         supported_params = get_supported_params(tracer_impl, params)