diff --git a/src/uipath/core/tracing/_utils.py b/src/uipath/core/tracing/_utils.py index 841aa67..28d25d4 100644 --- a/src/uipath/core/tracing/_utils.py +++ b/src/uipath/core/tracing/_utils.py @@ -86,6 +86,7 @@ def set_span_input_attributes( span_type: str, run_type: Optional[str], input_processor: Optional[Callable[..., Any]], + input_attributes_callable: Optional[Callable[..., dict[str, Any]]] = None, ) -> None: """Set span attributes for metadata and inputs before function execution. @@ -122,11 +123,19 @@ def set_span_input_attributes( span.set_attribute("input.mime_type", "application/json") span.set_attribute("input.value", inputs) + if input_attributes_callable: + try: + for key, value in input_attributes_callable(*args, **kwargs).items(): + span.set_attribute(key, value) + except Exception: + # Suppress exceptions from custom attributes to avoid breaking instrumentation (already-set attributes remain) + pass def set_span_output_attributes( span: Span, result: Any, output_processor: Optional[Callable[..., Any]], + output_attributes_callable: Optional[Callable[..., dict[str, Any]]] = None, ) -> None: """Set span attributes for outputs after function execution. @@ -140,3 +149,10 @@ def set_span_output_attributes( 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") + if output_attributes_callable: + try: + for key, value in output_attributes_callable(result).items(): + span.set_attribute(key, value) + except Exception: + # Suppress exceptions from custom attributes to avoid breaking instrumentation (already-set attributes remain) + pass diff --git a/src/uipath/core/tracing/decorators.py b/src/uipath/core/tracing/decorators.py index d1c751a..3177ed3 100644 --- a/src/uipath/core/tracing/decorators.py +++ b/src/uipath/core/tracing/decorators.py @@ -32,6 +32,8 @@ def _opentelemetry_traced( span_type: Optional[str] = None, input_processor: Optional[Callable[..., Any]] = None, output_processor: Optional[Callable[..., Any]] = None, + input_attributes_callable: Optional[Callable[..., dict[str, Any]]] = None, + output_attributes_callable: Optional[Callable[..., dict[str, Any]]] = None, recording: bool = True, ): """Default tracer implementation using OpenTelemetry. @@ -105,6 +107,7 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: run_type=run_type, span_type=span_type or "function_call_sync", input_processor=input_processor, + input_attributes_callable=input_attributes_callable, ) # Execute the function @@ -115,6 +118,7 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: span, result=result, output_processor=output_processor, + output_attributes_callable=output_attributes_callable, ) return result except Exception as e: @@ -142,6 +146,7 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: run_type=run_type, span_type=span_type or "function_call_async", input_processor=input_processor, + input_attributes_callable=input_attributes_callable, ) # Execute the function @@ -152,6 +157,7 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: span, result=result, output_processor=output_processor, + output_attributes_callable=output_attributes_callable, ) return result except Exception as e: @@ -181,6 +187,7 @@ def generator_wrapper(*args: Any, **kwargs: Any) -> Any: run_type=run_type, span_type=span_type or "function_call_generator_sync", input_processor=input_processor, + input_attributes_callable=input_attributes_callable, ) # Execute the generator and collect outputs @@ -195,6 +202,7 @@ def generator_wrapper(*args: Any, **kwargs: Any) -> Any: span, result=outputs, output_processor=output_processor, + output_attributes_callable=output_attributes_callable, ) except Exception as e: span.record_exception(e) @@ -223,6 +231,7 @@ async def async_generator_wrapper(*args: Any, **kwargs: Any) -> Any: run_type=run_type, span_type=span_type or "function_call_generator_async", input_processor=input_processor, + input_attributes_callable=input_attributes_callable, ) # Execute the generator and collect outputs @@ -237,6 +246,7 @@ async def async_generator_wrapper(*args: Any, **kwargs: Any) -> Any: span, result=outputs, output_processor=output_processor, + output_attributes_callable=output_attributes_callable, ) except Exception as e: span.record_exception(e) @@ -264,6 +274,8 @@ def traced( span_type: Optional[str] = None, input_processor: Optional[Callable[..., Any]] = None, output_processor: Optional[Callable[..., Any]] = None, + input_attributes_callable: Optional[Callable[..., dict[str, Any]]] = None, + output_attributes_callable: Optional[Callable[..., dict[str, Any]]] = None, hide_input: bool = False, hide_output: bool = False, recording: bool = True, @@ -304,6 +316,8 @@ def _default_output_processor(outputs: Any) -> dict[str, str]: "span_type": span_type, "input_processor": input_processor, "output_processor": output_processor, + "input_attributes_callable": input_attributes_callable, + "output_attributes_callable": output_attributes_callable, "recording": recording, }