Fix for new OpenAI Agents SDK (#1152)

Author: alexmojaki

logfire/_internal/integrations/openai_agents.py

@@ -40,7 +40,7 @@
 from logfire._internal.utils import handle_internal_errors, log_internal_error, truncate_string
 
 if TYPE_CHECKING:  # pragma: no cover
-    from agents.tracing.setup import TraceProvider
+    from agents.tracing import TraceProvider
     from openai.types.responses import Response
 
     from logfire import Logfire, LogfireSpan
@@ -130,18 +130,28 @@ def __getattr__(self, item: Any) -> Any:
 
     @classmethod
     def install(cls, logfire_instance: Logfire) -> None:
-        name = 'GLOBAL_TRACE_PROVIDER'
-        original = getattr(agents.tracing, name)
-        if isinstance(original, cls):
-            return
-        wrapper = cls(original, logfire_instance)
-        for module_name, mod in sys.modules.items():
-            if module_name.startswith('agents'):
-                try:
-                    if getattr(mod, name, None) is original:
-                        setattr(mod, name, wrapper)
-                except Exception:  # pragma: no cover
-                    pass
+        try:
+            from agents.tracing import get_trace_provider, set_trace_provider
+        except ImportError:  # pragma: no cover
+            # Handle older versions of agents where these functions are not available
+            name = 'GLOBAL_TRACE_PROVIDER'
+            original = getattr(agents.tracing, name)
+            if isinstance(original, cls):
+                return
+            wrapper = cls(original, logfire_instance)
+            for module_name, mod in sys.modules.items():
+                if module_name.startswith('agents'):
+                    try:
+                        if getattr(mod, name, None) is original:
+                            setattr(mod, name, wrapper)
+                    except Exception:  # pragma: no cover
+                        pass
+        else:
+            original = get_trace_provider()
+            if isinstance(original, cls):
+                return
+            wrapper = cls(original, logfire_instance)
+            set_trace_provider(wrapper)  # type: ignore
 
 
 @dataclass
@@ -185,13 +195,13 @@ class LogfireWrapperBase(Generic[T]):
     span_helper: LogfireSpanHelper
     token: contextvars.Token[T | None] | None = None
 
-    def start(self, mark_as_current: bool = False):
+    def start(self, mark_as_current: bool = False) -> None:
         self.span_helper.start(mark_as_current)
         if mark_as_current:
             self.attach()
         return self.wrapped.start()
 
-    def finish(self, reset_current: bool = False):
+    def finish(self, reset_current: bool = False) -> None:
         self.on_ending()
         self.span_helper.end(reset_current)
         if reset_current:
@@ -204,7 +214,7 @@ def __enter__(self) -> Self:
         self.attach()
         return self
 
-    def __exit__(self, exc_type: type[BaseException], exc_val: BaseException, exc_tb: TracebackType):
+    def __exit__(self, exc_type: type[BaseException], exc_val: BaseException, exc_tb: TracebackType) -> None:
         self.on_ending()
         self.span_helper.__exit__(exc_type, exc_val, exc_tb)
         self.wrapped.finish()

tests/conftest.py

@@ -7,7 +7,7 @@
 
 import anyio._backends._asyncio  # noqa  # type: ignore
 import pytest
-from agents.tracing.setup import GLOBAL_TRACE_PROVIDER
+from agents.tracing import get_trace_provider
 from opentelemetry import trace
 from opentelemetry.sdk._logs.export import SimpleLogRecordProcessor
 from opentelemetry.sdk.metrics.export import InMemoryMetricReader
@@ -28,8 +28,8 @@
 os.environ['LOGFIRE_TOKEN'] = ''
 
 
-GLOBAL_TRACE_PROVIDER.shutdown()
-GLOBAL_TRACE_PROVIDER.set_processors([])
+get_trace_provider().shutdown()
+get_trace_provider().set_processors([])
 
 
 @pytest.fixture(scope='session', autouse=True)

tests/otel_integrations/test_openai_agents.py

@@ -30,6 +30,7 @@
     input_guardrail,
     trace,
 )
+from agents.tracing import get_trace_provider
 from agents.tracing.span_data import MCPListToolsSpanData, ResponseSpanData
 from agents.tracing.spans import NoOpSpan
 from agents.tracing.traces import NoOpTrace
@@ -1073,8 +1074,6 @@ def all_subclasses(cls: type) -> list[type]:
 def test_unknown_span(exporter: TestExporter):
     logfire.instrument_openai_agents()
 
-    from agents.tracing.setup import GLOBAL_TRACE_PROVIDER
-
     class MySpanData(SpanData):
         def export(self):
             return {'foo': 'bar', 'type': self.type}
@@ -1085,7 +1084,7 @@ def type(self) -> str:
 
     with trace('my_trace', trace_id='trace_123', group_id='456') as t:
         assert t.name == 'my_trace'
-        with GLOBAL_TRACE_PROVIDER.create_span(span_data=MySpanData(), span_id='span_789') as s:
+        with get_trace_provider().create_span(span_data=MySpanData(), span_id='span_789') as s:
             assert s.trace_id == 'trace_123'
             assert s.span_id == 'span_789'
             assert s.parent_id is None