diff --git a/py/src/braintrust/auto.py b/py/src/braintrust/auto.py index 48276ea3..689fd435 100644 --- a/py/src/braintrust/auto.py +++ b/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() diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_auto_openai.py b/py/src/braintrust/integrations/auto_test_scripts/test_auto_openai.py index ef5eaf8f..e4aaa48e 100644 --- a/py/src/braintrust/integrations/auto_test_scripts/test_auto_openai.py +++ b/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 diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_auto_openai_agents_sdk.py b/py/src/braintrust/integrations/auto_test_scripts/test_auto_openai_agents_sdk.py new file mode 100644 index 00000000..1722c52a --- /dev/null +++ b/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] diff --git a/py/src/braintrust/wrappers/test_openai.py b/py/src/braintrust/wrappers/test_openai.py index 6ab9b343..a4b52f52 100644 --- a/py/src/braintrust/wrappers/test_openai.py +++ b/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."""