Skip to content

Commit 427b3a2

Browse files
AbhiPrasadAbhiPrasad
authored
feat(autogen): add automatic instrumentation (#345)
Instrument AutoGen AgentChat agents, teams, and function tools through the integrations API and wire the integration into auto_instrument(). Add nox matrix coverage for latest and 0.7.0 with VCR-backed cassettes and subprocess auto-instrument checks. resolves #336
1 parent 8b5aae5 commit 427b3a2

19 files changed

Lines changed: 2213 additions & 46 deletions

‎py/noxfile.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,6 +259,18 @@ def test_agentscope(session, version):
259259
_run_tests(session, f"{INTEGRATION_DIR}/agentscope/test_agentscope.py", version=version)
260260

261261

262+
AUTOGEN_VERSIONS = _get_matrix_versions("autogen-agentchat")
263+
264+
265+
@nox.session()
266+
@nox.parametrize("version", AUTOGEN_VERSIONS, ids=AUTOGEN_VERSIONS)
267+
def test_autogen(session, version):
268+
_install_test_deps(session)
269+
_install_matrix_dep(session, "autogen-agentchat", version)
270+
_install_matrix_dep(session, "autogen-ext", version)
271+
_run_tests(session, f"{INTEGRATION_DIR}/autogen/test_autogen.py", version=version)
272+
273+
262274
# Two test suites with different version requirements:
263275
# 1. wrap_openai approach: works with older versions (0.1.9+)
264276
# 2. Direct wrapper (setup_pydantic_ai): requires 1.10.0+ for all features

‎py/pyproject.toml‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,8 @@ lint = [
167167
"agentscope",
168168
"agno",
169169
"anthropic",
170+
"autogen-agentchat",
171+
"autogen-ext[openai]",
170172
"cohere",
171173
"autoevals",
172174
"braintrust-core",
@@ -278,6 +280,14 @@ latest = "agno==2.6.0"
278280
latest = "agentscope==1.0.19"
279281
"1.0.0" = "agentscope==1.0.0"
280282

283+
[tool.braintrust.matrix.autogen-agentchat]
284+
latest = "autogen-agentchat==0.7.5"
285+
"0.7.0" = "autogen-agentchat==0.7.0"
286+
287+
[tool.braintrust.matrix.autogen-ext]
288+
latest = "autogen-ext[openai]==0.7.5"
289+
"0.7.0" = "autogen-ext[openai]==0.7.0"
290+
281291
[tool.braintrust.matrix.pydantic-ai-integration]
282292
latest = "pydantic-ai==1.86.1"
283293
"1.10.0" = "pydantic-ai==1.10.0"
@@ -348,6 +358,7 @@ latest = "braintrust-core==0.0.59"
348358
adk = ["google-adk"]
349359
agentscope = ["agentscope"]
350360
agno = ["agno"]
361+
autogen = ["autogen-agentchat"]
351362
anthropic = ["anthropic"]
352363
cohere = ["cohere"]
353364
claude_agent_sdk = ["claude-agent-sdk"]
@@ -364,6 +375,8 @@ pydantic_ai = ["pydantic-ai-integration", "pydantic-ai-wrap-openai"]
364375
[tool.braintrust.vendor-packages]
365376
agno = "agno"
366377
agentscope = "agentscope"
378+
autogen-agentchat = "autogen_agentchat"
379+
autogen-ext = "autogen_ext"
367380
anthropic = "anthropic"
368381
cohere = "cohere"
369382
autoevals = "autoevals"

‎py/src/braintrust/auto.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
AgentScopeIntegration,
1313
AgnoIntegration,
1414
AnthropicIntegration,
15+
AutoGenIntegration,
1516
ClaudeAgentSDKIntegration,
1617
CohereIntegration,
1718
DSPyIntegration,
@@ -60,6 +61,7 @@ def auto_instrument(
6061
langchain: bool = True,
6162
openai_agents: bool = True,
6263
cohere: bool = True,
64+
autogen: bool = True,
6365
) -> dict[str, bool]:
6466
"""
6567
Auto-instrument supported AI/ML libraries for Braintrust tracing.
@@ -86,6 +88,7 @@ def auto_instrument(
8688
langchain: Enable LangChain instrumentation (default: True)
8789
openai_agents: Enable OpenAI Agents SDK instrumentation (default: True)
8890
cohere: Enable Cohere instrumentation (default: True)
91+
autogen: Enable AutoGen instrumentation (default: True)
8992
9093
Returns:
9194
Dict mapping integration name to whether it was successfully instrumented.
@@ -163,6 +166,8 @@ def auto_instrument(
163166
results["openai_agents"] = _instrument_integration(OpenAIAgentsIntegration)
164167
if cohere:
165168
results["cohere"] = _instrument_integration(CohereIntegration)
169+
if autogen:
170+
results["autogen"] = _instrument_integration(AutoGenIntegration)
166171

167172
return results
168173

‎py/src/braintrust/integrations/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from .agentscope import AgentScopeIntegration
33
from .agno import AgnoIntegration
44
from .anthropic import AnthropicIntegration
5+
from .autogen import AutoGenIntegration
56
from .claude_agent_sdk import ClaudeAgentSDKIntegration
67
from .cohere import CohereIntegration
78
from .dspy import DSPyIntegration
@@ -20,6 +21,7 @@
2021
"AgentScopeIntegration",
2122
"AgnoIntegration",
2223
"AnthropicIntegration",
24+
"AutoGenIntegration",
2325
"ClaudeAgentSDKIntegration",
2426
"CohereIntegration",
2527
"DSPyIntegration",

‎py/src/braintrust/integrations/auto_test_scripts/test_auto_autogen.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
from autogen_agentchat.agents import BaseChatAgent
2+
from autogen_agentchat.teams import BaseGroupChat
3+
from autogen_core.tools import FunctionTool
4+
from braintrust.auto import auto_instrument
5+
6+
7+
assert not getattr(BaseChatAgent.run, "__braintrust_patched_autogen_chat_agent_run__", False)
8+
assert not getattr(BaseGroupChat.run, "__braintrust_patched_autogen_team_run__", False)
9+
assert not getattr(FunctionTool.run, "__braintrust_patched_autogen_function_tool_run__", False)
10+
11+
results = auto_instrument()
12+
assert results.get("autogen") == True
13+
assert auto_instrument().get("autogen") == True
14+
15+
assert getattr(BaseChatAgent.run, "__braintrust_patched_autogen_chat_agent_run__", False)
16+
assert getattr(BaseChatAgent.run_stream, "__braintrust_patched_autogen_chat_agent_run_stream__", False)
17+
assert getattr(BaseGroupChat.run, "__braintrust_patched_autogen_team_run__", False)
18+
assert getattr(BaseGroupChat.run_stream, "__braintrust_patched_autogen_team_run_stream__", False)
19+
assert getattr(FunctionTool.run, "__braintrust_patched_autogen_function_tool_run__", False)
20+
21+
print("SUCCESS")

‎py/src/braintrust/integrations/autogen/__init__.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
"""Braintrust AutoGen integration."""
2+
3+
from braintrust.logger import NOOP_SPAN, current_span, init_logger
4+
5+
from .integration import AutoGenIntegration
6+
7+
8+
def setup_autogen(
9+
api_key: str | None = None,
10+
project_id: str | None = None,
11+
project_name: str | None = None,
12+
) -> bool:
13+
"""Setup Braintrust integration with AutoGen."""
14+
if current_span() == NOOP_SPAN:
15+
init_logger(project=project_name, api_key=api_key, project_id=project_id)
16+
17+
return AutoGenIntegration.setup()
18+
19+
20+
__all__ = ["AutoGenIntegration", "setup_autogen"]

‎py/src/braintrust/integrations/autogen/cassettes/0.7.0/test_autogen_agent_run_creates_braintrust_spans.yaml‎

Lines changed: 216 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
1+
interactions:
2+
- request:
3+
body: '{"messages":[{"content":"You are concise. Answer directly.","role":"system"},{"role":"user","name":"user","content":"Say
4+
hello in exactly two words."}],"model":"gpt-4o-mini","stream":false,"temperature":0}'
5+
headers:
6+
Accept:
7+
- application/json
8+
Accept-Encoding:
9+
- gzip, deflate
10+
Connection:
11+
- keep-alive
12+
Content-Length:
13+
- '205'
14+
Content-Type:
15+
- application/json
16+
Host:
17+
- api.openai.com
18+
User-Agent:
19+
- AsyncOpenAI/Python 2.32.0
20+
X-Stainless-Arch:
21+
- arm64
22+
X-Stainless-Async:
23+
- async:asyncio
24+
X-Stainless-Lang:
25+
- python
26+
X-Stainless-OS:
27+
- MacOS
28+
X-Stainless-Package-Version:
29+
- 2.32.0
30+
X-Stainless-Runtime:
31+
- CPython
32+
X-Stainless-Runtime-Version:
33+
- 3.12.12
34+
x-stainless-read-timeout:
35+
- '600'
36+
x-stainless-retry-count:
37+
- '0'
38+
method: POST
39+
uri: https://api.openai.com/v1/chat/completions
40+
response:
41+
body:
42+
string: "{\n \"id\": \"chatcmpl-DYGF4xklfNjHQ3dNZDbuSrGAFsYDu\",\n \"object\":
43+
\"chat.completion\",\n \"created\": 1777057958,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n
44+
\ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\":
45+
\"assistant\",\n \"content\": \"Hello there.\",\n \"refusal\":
46+
null,\n \"annotations\": []\n },\n \"logprobs\": null,\n
47+
\ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\":
48+
26,\n \"completion_tokens\": 3,\n \"total_tokens\": 29,\n \"prompt_tokens_details\":
49+
{\n \"cached_tokens\": 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\":
50+
{\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\":
51+
0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\":
52+
\"default\",\n \"system_fingerprint\": \"fp_a7190374f3\"\n}\n"
53+
headers:
54+
CF-Cache-Status:
55+
- DYNAMIC
56+
CF-Ray:
57+
- 9f176baf4e10ac31-YYZ
58+
Connection:
59+
- keep-alive
60+
Content-Type:
61+
- application/json
62+
Date:
63+
- Fri, 24 Apr 2026 19:12:39 GMT
64+
Server:
65+
- cloudflare
66+
Strict-Transport-Security:
67+
- max-age=31536000; includeSubDomains; preload
68+
Transfer-Encoding:
69+
- chunked
70+
X-Content-Type-Options:
71+
- nosniff
72+
access-control-expose-headers:
73+
- X-Request-ID
74+
alt-svc:
75+
- h3=":443"; ma=86400
76+
content-length:
77+
- '818'
78+
openai-organization:
79+
- braintrust-data
80+
openai-processing-ms:
81+
- '242'
82+
openai-project:
83+
- proj_vsCSXafhhByzWOThMrJcZiw9
84+
openai-version:
85+
- '2020-10-01'
86+
set-cookie:
87+
- __cf_bm=Nc7FmCaFJfxm0G8e_ebd8Eba0ln5XCCF5rMUIYKGhm8-1777057958.2822428-1.0.1.1-dWCiuBiWLzZ_uGu88vNOw4PYjIEEuhEMzc3AoxClyqguU9ItaGD3mo.kpjSf9TKQwRzNOtwJJ9jsn6lNJiHYUMOn.P9rBhdiokBoqTvIe6iVhLEAv9XyQnBiEjn73kyE;
88+
HttpOnly; Secure; Path=/; Domain=api.openai.com; Expires=Fri, 24 Apr 2026
89+
19:42:39 GMT
90+
x-openai-proxy-wasm:
91+
- v0.1
92+
x-ratelimit-limit-requests:
93+
- '30000'
94+
x-ratelimit-limit-tokens:
95+
- '150000000'
96+
x-ratelimit-remaining-requests:
97+
- '29999'
98+
x-ratelimit-remaining-tokens:
99+
- '149999980'
100+
x-ratelimit-reset-requests:
101+
- 2ms
102+
x-ratelimit-reset-tokens:
103+
- 0s
104+
x-request-id:
105+
- req_16ebc0bb94f54d5d92879f3f653c4473
106+
status:
107+
code: 200
108+
message: OK
109+
- request:
110+
body: '{"messages":[{"content":"You are concise. Answer directly.","role":"system"},{"role":"user","name":"user","content":"Say
111+
hello in exactly two words."}],"model":"gpt-4o-mini","stream":false,"temperature":0}'
112+
headers:
113+
Accept:
114+
- application/json
115+
Accept-Encoding:
116+
- gzip, deflate
117+
Connection:
118+
- keep-alive
119+
Content-Length:
120+
- '205'
121+
Content-Type:
122+
- application/json
123+
Host:
124+
- api.openai.com
125+
User-Agent:
126+
- AsyncOpenAI/Python 2.32.0
127+
X-Stainless-Arch:
128+
- arm64
129+
X-Stainless-Async:
130+
- async:asyncio
131+
X-Stainless-Lang:
132+
- python
133+
X-Stainless-OS:
134+
- MacOS
135+
X-Stainless-Package-Version:
136+
- 2.32.0
137+
X-Stainless-Runtime:
138+
- CPython
139+
X-Stainless-Runtime-Version:
140+
- 3.12.12
141+
x-stainless-read-timeout:
142+
- '600'
143+
x-stainless-retry-count:
144+
- '0'
145+
method: POST
146+
uri: https://api.openai.com/v1/chat/completions
147+
response:
148+
body:
149+
string: "{\n \"id\": \"chatcmpl-DYGHRGb8gFqrGHLu01pO9PifTTNAT\",\n \"object\":
150+
\"chat.completion\",\n \"created\": 1777058105,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n
151+
\ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\":
152+
\"assistant\",\n \"content\": \"Hello there.\",\n \"refusal\":
153+
null,\n \"annotations\": []\n },\n \"logprobs\": null,\n
154+
\ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\":
155+
26,\n \"completion_tokens\": 3,\n \"total_tokens\": 29,\n \"prompt_tokens_details\":
156+
{\n \"cached_tokens\": 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\":
157+
{\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\":
158+
0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\":
159+
\"default\",\n \"system_fingerprint\": \"fp_fa616991a3\"\n}\n"
160+
headers:
161+
CF-Cache-Status:
162+
- DYNAMIC
163+
CF-Ray:
164+
- 9f176f43992c5d15-YYZ
165+
Connection:
166+
- keep-alive
167+
Content-Type:
168+
- application/json
169+
Date:
170+
- Fri, 24 Apr 2026 19:15:05 GMT
171+
Server:
172+
- cloudflare
173+
Strict-Transport-Security:
174+
- max-age=31536000; includeSubDomains; preload
175+
Transfer-Encoding:
176+
- chunked
177+
X-Content-Type-Options:
178+
- nosniff
179+
access-control-expose-headers:
180+
- X-Request-ID
181+
alt-svc:
182+
- h3=":443"; ma=86400
183+
content-length:
184+
- '818'
185+
openai-organization:
186+
- braintrust-data
187+
openai-processing-ms:
188+
- '298'
189+
openai-project:
190+
- proj_vsCSXafhhByzWOThMrJcZiw9
191+
openai-version:
192+
- '2020-10-01'
193+
set-cookie:
194+
- __cf_bm=8tkx7xc7WnAjyJWBvTnV.2hAVD3Xd6z1iKcrqd22KKk-1777058104.8957334-1.0.1.1-htZq5VoU00hy9oOwPpLlgjXRvNiv06KVGIalt1HU8QrRn0TOHRU5NJvDT0KbgX8BVE_ZzME0Cx_jEhn4r1li3rtjYfRO2tzr0CsEuzBtFV9cdvvpeyGISz5l2Ht6uHIm;
195+
HttpOnly; Secure; Path=/; Domain=api.openai.com; Expires=Fri, 24 Apr 2026
196+
19:45:05 GMT
197+
x-openai-proxy-wasm:
198+
- v0.1
199+
x-ratelimit-limit-requests:
200+
- '30000'
201+
x-ratelimit-limit-tokens:
202+
- '150000000'
203+
x-ratelimit-remaining-requests:
204+
- '29999'
205+
x-ratelimit-remaining-tokens:
206+
- '149999980'
207+
x-ratelimit-reset-requests:
208+
- 2ms
209+
x-ratelimit-reset-tokens:
210+
- 0s
211+
x-request-id:
212+
- req_c6c492f11afb43fb97f328b69ad54803
213+
status:
214+
code: 200
215+
message: OK
216+
version: 1

0 commit comments

Comments
 (0)