Add ID generation methods to TraceProvider (#729)

Author: pakrym-oai

src/agents/__init__.py

@@ -92,6 +92,7 @@
     handoff_span,
     mcp_tools_span,
     set_trace_processors,
+    set_trace_provider,
     set_tracing_disabled,
     set_tracing_export_api_key,
     speech_group_span,
@@ -221,6 +222,7 @@ def enable_verbose_stdout_logging():
     "guardrail_span",
     "handoff_span",
     "set_trace_processors",
+    "set_trace_provider",
     "set_tracing_disabled",
     "speech_group_span",
     "transcription_span",

src/agents/tracing/__init__.py

@@ -18,7 +18,7 @@
 )
 from .processor_interface import TracingProcessor
 from .processors import default_exporter, default_processor
-from .setup import GLOBAL_TRACE_PROVIDER
+from .setup import GLOBAL_TRACE_PROVIDER, set_trace_provider
 from .span_data import (
     AgentSpanData,
     CustomSpanData,
@@ -49,6 +49,7 @@
     "handoff_span",
     "response_span",
     "set_trace_processors",
+    "set_trace_provider",
     "set_tracing_disabled",
     "trace",
     "Trace",

src/agents/tracing/setup.py

@@ -2,6 +2,8 @@
 
 import os
 import threading
+import uuid
+from datetime import datetime, timezone
 from typing import Any
 
 from ..logger import logger
@@ -118,6 +120,22 @@ def set_disabled(self, disabled: bool) -> None:
         """
         self._disabled = disabled
 
+    def time_iso(self) -> str:
+        """Return the current time in ISO 8601 format."""
+        return datetime.now(timezone.utc).isoformat()
+
+    def gen_trace_id(self) -> str:
+        """Generate a new trace ID."""
+        return f"trace_{uuid.uuid4().hex}"
+
+    def gen_span_id(self) -> str:
+        """Generate a new span ID."""
+        return f"span_{uuid.uuid4().hex[:24]}"
+
+    def gen_group_id(self) -> str:
+        """Generate a new group ID."""
+        return f"group_{uuid.uuid4().hex[:24]}"
+
     def create_trace(
         self,
         name: str,
@@ -212,3 +230,9 @@ def shutdown(self) -> None:
 
 
 GLOBAL_TRACE_PROVIDER = TraceProvider()
+
+
+def set_trace_provider(provider: TraceProvider) -> None:
+    """Set the global trace provider used by tracing utilities."""
+    global GLOBAL_TRACE_PROVIDER
+    GLOBAL_TRACE_PROVIDER = provider

src/agents/tracing/util.py

@@ -1,22 +1,21 @@
-import uuid
-from datetime import datetime, timezone
+from .setup import GLOBAL_TRACE_PROVIDER
 
 
 def time_iso() -> str:
-    """Returns the current time in ISO 8601 format."""
-    return datetime.now(timezone.utc).isoformat()
+    """Return the current time in ISO 8601 format."""
+    return GLOBAL_TRACE_PROVIDER.time_iso()
 
 
 def gen_trace_id() -> str:
-    """Generates a new trace ID."""
-    return f"trace_{uuid.uuid4().hex}"
+    """Generate a new trace ID."""
+    return GLOBAL_TRACE_PROVIDER.gen_trace_id()
 
 
 def gen_span_id() -> str:
-    """Generates a new span ID."""
-    return f"span_{uuid.uuid4().hex[:24]}"
+    """Generate a new span ID."""
+    return GLOBAL_TRACE_PROVIDER.gen_span_id()
 
 
 def gen_group_id() -> str:
-    """Generates a new group ID."""
-    return f"group_{uuid.uuid4().hex[:24]}"
+    """Generate a new group ID."""
+    return GLOBAL_TRACE_PROVIDER.gen_group_id()