auto_instrument openai agents sdk

Author: Mahhheshh

py/src/braintrust/auto.py

@@ -41,6 +41,7 @@ def auto_instrument(
     google_genai: bool = True,
     agno: bool = True,
     claude_agent_sdk: bool = True,
+    openai_agent_sdk: bool = True,
     dspy: bool = True,
     adk: bool = True,
 ) -> dict[str, bool]:
@@ -61,6 +62,7 @@ def auto_instrument(
         google_genai: Enable Google GenAI instrumentation (default: True)
         agno: Enable Agno instrumentation (default: True)
         claude_agent_sdk: Enable Claude Agent SDK instrumentation (default: True)
+        openai_agent_sdk: Enable OpenAI Agent SDK instrumentation (default: True)
         dspy: Enable DSPy instrumentation (default: True)
         adk: Enable Google ADK instrumentation (default: True)
 
@@ -124,6 +126,8 @@ def auto_instrument(
         results["agno"] = _instrument_integration(AgnoIntegration)
     if claude_agent_sdk:
         results["claude_agent_sdk"] = _instrument_integration(ClaudeAgentSDKIntegration)
+    if openai_agent_sdk:
+        results["openai_agent_sdk"] = _instrument_openai_agent_sdk()
     if dspy:
         results["dspy"] = _instrument_dspy()
     if adk:
@@ -140,6 +144,16 @@ def _instrument_openai() -> bool:
     return False
 
 
+def _instrument_openai_agent_sdk() -> bool:
+    with _try_patch():
+        from agents import set_trace_processors
+        from braintrust.wrappers.openai import BraintrustTracingProcessor
+
+        set_trace_processors([BraintrustTracingProcessor()])
+        return True
+    return False
+
+
 def _instrument_integration(integration) -> bool:
     with _try_patch():
         return integration.setup()

py/src/braintrust/integrations/auto_test_scripts/test_auto_openai.py

@@ -8,11 +8,13 @@
 # 1. Verify not patched initially
 assert not getattr(openai, "__braintrust_wrapped__", False)
 
+
 # 2. Instrument
 results = auto_instrument()
 assert results.get("openai") == True
 assert getattr(openai, "__braintrust_wrapped__", False)
 
+
 # 3. Idempotent
 results2 = auto_instrument()
 assert results2.get("openai") == True

py/src/braintrust/integrations/auto_test_scripts/test_auto_openai_agents_sdk.py

@@ -0,0 +1,26 @@
+"""Test auto_instrument for OpenAI Agents SDK."""
+
+from agents.tracing import get_trace_provider
+from braintrust.auto import auto_instrument
+from braintrust.wrappers.openai import BraintrustTracingProcessor
+
+
+results = auto_instrument()
+assert results.get("openai_agent_sdk") == True
+bt_processor = get_trace_provider()._multi_processor._processors[0]
+assert isinstance(bt_processor, BraintrustTracingProcessor)
+
+results2 = auto_instrument()
+assert results2.get("openai_agent_sdk") == True
+
+# with autoinstrument_test_context("test_auto_openai_agents_sdk") as memory_logger:
+#   agent = Agent(
+#         name="Assistant",
+#         instructions="You only respond in haikus.",
+#     )
+
+#   result = Runner.run_sync(agent, "Tell me about recursion in programming.")
+
+#   span = memory_logger.pop()
+#   assert len(spans) == 1, f"Expected 1 span, got {len(spans)}"
+#   span = spans[0]

py/src/braintrust/wrappers/test_openai.py

@@ -2172,6 +2172,10 @@ def test_auto_instrument_openai(self):
         """Test auto_instrument patches OpenAI, creates spans, and uninstrument works."""
         verify_autoinstrument_script("test_auto_openai.py")
 
+    def test_auto_instrument_openai_agents_sdk(self):
+        """Test auto_instrument patches OpenAI, creates spans, and uninstrument works."""
+        verify_autoinstrument_script("test_auto_openai_agents_sdk.py")
+
 
 class TestZAICompatibleOpenAI:
     """Tests for validating some ZAI compatibility with OpenAI wrapper."""