Merge pull request #7 from UiPath/fix/add_attributes

fix: add span attributes

Author: cristipufu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-core"
-version = "0.0.3"
+version = "0.0.4"
 description = "UiPath Core abstractions"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
@@ -12,7 +12,6 @@ dependencies = [
 classifiers = [
   "Intended Audience :: Developers",
   "Topic :: Software Development :: Build Tools",
-  "Programming Language :: Python :: 3.10",
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
@@ -25,7 +24,7 @@ maintainers = [
 
 [project.urls]
 Homepage = "https://uipath.com"
-Repository = "https://github.com/UiPath/uipath-runtime-python"
+Repository = "https://github.com/UiPath/uipath-core-python"
 Documentation = "https://uipath.github.io/uipath-python/"
 
 [build-system]

src/uipath/core/tracing/__init__.py

@@ -4,14 +4,12 @@
 with OpenTelemetry tracing, including custom processors for UiPath execution tracking.
 """
 
-from uipath.core.tracing.context import UiPathTraceContext
 from uipath.core.tracing.decorators import traced
 from uipath.core.tracing.span_utils import UiPathSpanUtils
 from uipath.core.tracing.trace_manager import UiPathTraceManager
 
 __all__ = [
     "traced",
-    "UiPathTraceContext",
     "UiPathSpanUtils",
     "UiPathTraceManager",
 ]

src/uipath/core/tracing/_utils.py

@@ -6,7 +6,7 @@
 from dataclasses import asdict, is_dataclass
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any, Dict, Mapping
+from typing import Any, Mapping
 from zoneinfo import ZoneInfo
 
 from pydantic import BaseModel
@@ -15,15 +15,15 @@
 def get_supported_params(
     tracer_impl: Callable[..., Any],
     params: Mapping[str, Any],
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """Extract the parameters supported by the tracer implementation."""
     try:
         sig = inspect.signature(tracer_impl)
     except (TypeError, ValueError):
         # If we can't inspect, pass all parameters and let the function handle it
         return dict(params)
 
-    supported: Dict[str, Any] = {}
+    supported: dict[str, Any] = {}
     for name, value in params.items():
         if value is not None and name in sig.parameters:
             supported[name] = value
@@ -90,7 +90,7 @@ def format_object_for_trace_json(
 
 def format_args_for_trace(
     signature: inspect.Signature, *args: Any, **kwargs: Any
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     try:
         """Return a dictionary of inputs from the function signature."""
         # Create a parameter mapping by partially binding the arguments

src/uipath/core/tracing/context.py

@@ -1,7 +1,7 @@
 """Custom OpenTelemetry Span Exporter for UiPath Runtime executions."""
 
 from collections import defaultdict
-from typing import Dict, List, Optional, Sequence
+from typing import Dict, Optional, Sequence
 
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
@@ -12,7 +12,7 @@ class UiPathRuntimeExecutionSpanExporter(SpanExporter):
 
     def __init__(self):
         """Initialize the exporter."""
-        self._spans: Dict[str, List[ReadableSpan]] = defaultdict(list)
+        self._spans: Dict[str, list[ReadableSpan]] = defaultdict(list)
 
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
         """Export spans, grouping them by execution id."""
@@ -24,7 +24,7 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
 
         return SpanExportResult.SUCCESS
 
-    def get_spans(self, execution_id: str) -> List[ReadableSpan]:
+    def get_spans(self, execution_id: str) -> list[ReadableSpan]:
         """Retrieve spans for a given execution id."""
         return self._spans.get(execution_id, [])
 

src/uipath/core/tracing/exporters.py

@@ -1,7 +1,7 @@
 """Utilities for managing UiPath spans."""
 
 import logging
-from typing import Callable, Dict, List, Optional
+from typing import Callable, Optional
 
 from opentelemetry import context, trace
 from opentelemetry.trace import Span, set_span_in_context
@@ -15,8 +15,8 @@ class SpanRegistry:
     MAX_HIERARCHY_DEPTH = 1000  # Hard limit for hierarchy traversal
 
     def __init__(self):
-        self._spans: Dict[int, Span] = {}  # span_id -> span
-        self._parent_map: Dict[int, Optional[int]] = {}  # span_id -> parent_id
+        self._spans: dict[int, Span] = {}  # span_id -> span
+        self._parent_map: dict[int, Optional[int]] = {}  # span_id -> parent_id
 
     def register_span(self, span: Span) -> None:
         """Register a span and its parent relationship."""
@@ -116,7 +116,7 @@ class UiPathSpanUtils:
     """Static utility class to manage tracing implementations and decorated functions."""
 
     _current_span_provider: Optional[Callable[[], Optional[Span]]] = None
-    _current_span_ancestors_provider: Optional[Callable[[], List[Span]]] = None
+    _current_span_ancestors_provider: Optional[Callable[[], list[Span]]] = None
 
     @staticmethod
     def register_current_span_provider(
@@ -279,7 +279,7 @@ def get_external_current_span() -> Optional[Span]:
         return None
 
     @staticmethod
-    def get_ancestor_spans() -> List[Span]:
+    def get_ancestor_spans() -> list[Span]:
         """Get the ancestor spans from the registered provider, if any."""
         if UiPathSpanUtils._current_span_ancestors_provider is not None:
             try:
@@ -290,7 +290,7 @@ def get_ancestor_spans() -> List[Span]:
 
     @staticmethod
     def register_current_span_ancestors_provider(
-        current_span_ancestors_provider: Optional[Callable[[], List[Span]]],
+        current_span_ancestors_provider: Optional[Callable[[], list[Span]]],
     ):
         """Register a custom current span ancestors provider function.
 

src/uipath/core/tracing/span_utils.py

@@ -1,7 +1,7 @@
 """Tracing manager for handling tracer implementations and function registry."""
 
 import contextlib
-from typing import Any, Generator, List
+from typing import Any, Generator, Optional
 
 from opentelemetry import trace
 from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
@@ -30,7 +30,7 @@ def __init__(self):
             "An incompatible Otel TracerProvider was instantiated. Please check runtime configuration."
         )
         self.tracer_provider: TracerProvider = current_provider
-        self.tracer_span_processors: List[SpanProcessor] = []
+        self.tracer_span_processors: list[SpanProcessor] = []
         self.execution_span_exporter = UiPathRuntimeExecutionSpanExporter()
         self.add_span_exporter(self.execution_span_exporter)
 
@@ -52,20 +52,25 @@ def add_span_exporter(
     def get_execution_spans(
         self,
         execution_id: str,
-    ) -> List[ReadableSpan]:
+    ) -> list[ReadableSpan]:
         """Retrieve spans for a given execution id."""
         return self.execution_span_exporter.get_spans(execution_id)
 
     @contextlib.contextmanager
     def start_execution_span(
-        self, root_span: str, execution_id: str
+        self,
+        root_span: str,
+        execution_id: str,
+        attributes: Optional[dict[str, str]] = None,
     ) -> Generator[_AgnosticContextManager[Any] | Any, Any, None]:
         """Start an execution span."""
         try:
             tracer = trace.get_tracer("uipath-runtime")
             span_attributes: dict[str, Any] = {}
             if execution_id:
                 span_attributes["execution.id"] = execution_id
+            if attributes:
+                span_attributes.update(attributes)
             with tracer.start_as_current_span(
                 root_span, attributes=span_attributes
             ) as span:

src/uipath/core/tracing/trace_manager.py

@@ -118,15 +118,15 @@ def process_timestamp(timestamp: datetime):
 
 def test_traced_with_complex_return_value(span_capture: SpanCapture):
     """Test tracing with complex return value."""
-    from typing import Any, Dict
+    from typing import Any
 
     from pydantic import BaseModel
 
     from uipath.core.tracing.decorators import traced
 
     class Result(BaseModel):
         success: bool
-        data: Dict[str, Any]
+        data: dict[str, Any]
 
     @traced(name="get_result")
     def get_result():

tests/tracing/test_serialization.py

@@ -2,7 +2,7 @@
 from asyncio import sleep
 from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Any, Dict, List, Sequence
+from typing import Any, Sequence
 
 import pytest
 from opentelemetry import trace
@@ -30,7 +30,7 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
         self.spans.extend(spans)
         return SpanExportResult.SUCCESS
 
-    def get_exported_spans(self) -> List[ReadableSpan]:
+    def get_exported_spans(self) -> list[ReadableSpan]:
         return self.spans
 
     def clear_exported_spans(self) -> None:
@@ -239,7 +239,7 @@ async def async_operation(message):
     assert output["processed"] is True  # Added by processor
 
 
-def mask_credit_card(inputs: Dict[str, Any]) -> Dict[str, Any]:
+def mask_credit_card(inputs: dict[str, Any]) -> dict[str, Any]:
     """Process inputs to mask credit card information."""
     processed = inputs.copy()
     if "card_number" in processed:
@@ -251,7 +251,7 @@ def mask_credit_card(inputs: Dict[str, Any]) -> Dict[str, Any]:
     return processed
 
 
-def anonymize_single_user_data(output_dict: Dict[str, Any]) -> Dict[str, Any]:
+def anonymize_single_user_data(output_dict: dict[str, Any]) -> dict[str, Any]:
     """Process a single dictionary to anonymize user information."""
     processed = output_dict.copy()
     if "user_info" in processed and isinstance(processed["user_info"], dict):
@@ -669,7 +669,7 @@ class OutputFormat(BaseModel):
         confidence: float = 0.95
 
     @traced()
-    async def llm_chat_completions(messages: List[Any], response_format=None):
+    async def llm_chat_completions(messages: list[Any], response_format=None):
         """Simulate LLM function with BaseModel class as response_format."""
         if response_format:
             mock_content = '{"result": "hi!", "confidence": 0.95}'