From 433e4decf7b7d5d877df70c55f4917eea8443d52 Mon Sep 17 00:00:00 2001
From: pakrym-oai <pakrym@openai.com>
Date: Wed, 21 May 2025 09:12:19 -0700
Subject: [PATCH] Add ID generation methods to TraceProvider

---
 src/agents/__init__.py         |  2 ++
 src/agents/tracing/__init__.py |  3 ++-
 src/agents/tracing/setup.py    | 24 ++++++++++++++++++++++++
 src/agents/tracing/util.py     | 19 +++++++++----------
 4 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/src/agents/__init__.py b/src/agents/__init__.py
index 3c9c8465..65e44244 100644
--- a/src/agents/__init__.py
+++ b/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",
diff --git a/src/agents/tracing/__init__.py b/src/agents/tracing/__init__.py
index 9df94426..07d8af6d 100644
--- a/src/agents/tracing/__init__.py
+++ b/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",
diff --git a/src/agents/tracing/setup.py b/src/agents/tracing/setup.py
index 9e27d210..daa7b86d 100644
--- a/src/agents/tracing/setup.py
+++ b/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
diff --git a/src/agents/tracing/util.py b/src/agents/tracing/util.py
index f546b4e5..af2f5ff3 100644
--- a/src/agents/tracing/util.py
+++ b/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()