From b0b9c4e012462112499021dc344fc29fb84b38fc Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 15:28:18 +0000 Subject: [PATCH 01/11] fix(adk): align integration with sdk-integrations SKILL - Add metadata.provider="google" and metadata.tools on llm spans; lift tools/tool_config out of input.config via _capture_config(exclude=...) - Pass Exception instance to span.log(error=...) instead of str(e) in the tool and mcp_tool wrappers - Drop raw ADK objects (invocation_context, model_response_event, parent_context, kwargs spread) from span metadata; keep only the small named fields that are actually useful - Assert span_origin coverage across task/llm/tool types and validate metadata.provider / metadata.tools in the existing VCR-backed tests - Prune mock-heavy MCP tests; real behavior is covered end-to-end by the VCR suite in test_adk.py Co-Authored-By: Claude Opus 4.7 --- .../braintrust/integrations/adk/test_adk.py | 28 ++ .../integrations/adk/test_adk_mcp_tool.py | 338 +----------------- py/src/braintrust/integrations/adk/tracing.py | 75 ++-- 3 files changed, 83 insertions(+), 358 deletions(-) diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index 462f2eba..d234344a 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -482,6 +482,33 @@ async def test_adk_braintrust_integration(memory_logger): assert function_call["name"] == "get_weather" assert function_call["args"]["location"] == "San Francisco" + # Every integration-owned span carries the adk-auto instrumentation name + # (SKILL: "Every span an integration creates MUST carry ... instrumentation.name") + adk_spans = [ + row + for row in spans + if row["context"]["span_origin"]["instrumentation"]["name"] == "adk-auto" + ] + span_types_by_origin = {row["span_attributes"]["type"] for row in adk_spans} + assert {"task", "llm", "tool"} <= span_types_by_origin, ( + f"adk-auto origin missing on task/llm/tool spans: {span_types_by_origin}" + ) + + # Every llm span carries metadata.model, metadata.provider=google, and metadata.tools + # (with tools NOT leaking into input.config). + for span in llm_spans: + meta = span["metadata"] + assert meta.get("provider") == "google", f"Missing metadata.provider=google on {span['span_attributes']['name']}" + assert meta.get("model"), "Missing metadata.model on llm span" + assert meta.get("tools"), "metadata.tools should be non-empty for a tool-using agent" + tool_names = [ + fn.get("name") + for tool_entry in meta["tools"] + for fn in (tool_entry.get("function_declarations") or []) + ] + assert "get_weather" in tool_names, f"get_weather missing from metadata.tools: {tool_names}" + assert "tools" not in span["input"].get("config", {}), "tools should not be in input.config" + # Check response generation LLM call response_gen_spans = [span for span in llm_spans if "response_generation" in span["span_attributes"]["name"]] assert len(response_gen_spans) > 0, "Missing response generation LLM call" @@ -896,6 +923,7 @@ async def test_adk_captures_metrics(memory_logger): metadata = llm_span_with_metrics.get("metadata", {}) assert "model" in metadata, "Metadata should include model name" assert metadata["model"] == ADK_MODEL, "Model name should match the agent's model" + assert metadata.get("provider") == "google", "Metadata should include provider=google" def test_determine_llm_call_type_direct_response(): diff --git a/py/src/braintrust/integrations/adk/test_adk_mcp_tool.py b/py/src/braintrust/integrations/adk/test_adk_mcp_tool.py index 3474c215..b7d7c827 100644 --- a/py/src/braintrust/integrations/adk/test_adk_mcp_tool.py +++ b/py/src/braintrust/integrations/adk/test_adk_mcp_tool.py @@ -1,348 +1,30 @@ -"""Tests for MCP tool tracing integration.""" +"""Tests for MCP tool tracing integration. -from unittest.mock import AsyncMock, MagicMock, patch +Behavioral coverage of the MCP tool span is provided by the VCR-backed suite in +``test_adk.py`` — when an ADK agent invokes a real tool through the flow, the +``tool [...]`` / ``mcp_tool [...]`` span is exercised end-to-end. The tests here +only cover the patcher wiring, which does not require a live MCP server. +""" import pytest from braintrust.integrations.adk import setup_adk, wrap_mcp_tool +from braintrust.integrations.adk.patchers import McpToolPatcher @pytest.mark.asyncio async def test_wrap_mcp_tool_marks_as_patched(): - """Test that wrap_mcp_tool marks the class as patched.""" + """wrap_mcp_tool marks the class via the patcher marker (idempotence signal).""" - # Create a real class to wrap class MockMcpTool: async def run_async(self, *, args, tool_context): return {"result": "success"} - # Wrap the class wrapped_class = wrap_mcp_tool(MockMcpTool) - - # Verify it's marked as patched via the patcher marker - from braintrust.integrations.adk.patchers import McpToolPatcher - assert getattr(wrapped_class, McpToolPatcher.patch_marker_attr(), False) -@pytest.mark.asyncio -async def test_mcp_tool_execution_creates_span(): - """Test that MCP tool execution creates proper trace spans.""" - - with patch("braintrust.integrations.adk.tracing.start_span") as mock_start_span: - # Setup mock span - mock_span = MagicMock() - mock_span.__enter__ = MagicMock(return_value=mock_span) - mock_span.__exit__ = MagicMock(return_value=False) - mock_start_span.return_value = mock_span - - # Mock McpTool class and instance - MockMcpTool = MagicMock() - mock_instance = MagicMock() - mock_instance.name = "read_file" - mock_instance.run_async = AsyncMock(return_value={"content": [{"type": "text", "text": "file contents"}]}) - - # Wrap the class - wrapped_class = wrap_mcp_tool(MockMcpTool) - - # Simulate tool execution - tool_args = {"path": "/tmp/test.txt"} - tool_context = None - - # Call the wrapped method directly on the mock instance - # We need to manually trigger the wrapper - - # Get the original method - original_run_async = mock_instance.run_async - - # Create wrapped version by calling wrap_mcp_tool's wrapper - # This simulates what wrapt does - async def call_wrapped(): - return await original_run_async(args=tool_args, tool_context=tool_context) - - result = await call_wrapped() - - # For now, just verify the mock was called - mock_instance.run_async.assert_called_once_with(args=tool_args, tool_context=tool_context) - - -@pytest.mark.asyncio -async def test_mcp_tool_span_captures_tool_info(): - """Test that MCP tool spans capture tool name, args, and results.""" - from braintrust.span_types import SpanTypeAttribute - - with patch("braintrust.integrations.adk.tracing.start_span") as mock_start_span: - mock_span = MagicMock() - mock_span.__enter__ = MagicMock(return_value=mock_span) - mock_span.__exit__ = MagicMock(return_value=False) - mock_start_span.return_value = mock_span - - # Create a real-ish McpTool mock - class MockMcpTool: - def __init__(self): - self.name = "list_directory" - self._original_run_async = AsyncMock( - return_value={"content": [{"type": "text", "text": "file1.txt\nfile2.txt"}]} - ) - - async def run_async(self, *, args, tool_context): - return await self._original_run_async(args=args, tool_context=tool_context) - - # Wrap the class - wrap_mcp_tool(MockMcpTool) - - # Create instance and call - tool = MockMcpTool() - tool_args = {"path": "/tmp"} - tool_context = None - - result = await tool.run_async(args=tool_args, tool_context=tool_context) - - # Verify span was created - assert mock_start_span.called - call_kwargs = mock_start_span.call_args[1] - - # Check span name includes tool name - assert "list_directory" in call_kwargs["name"] - - # Check span type is TOOL - assert call_kwargs["type"] == SpanTypeAttribute.TOOL - - # Check input contains tool name and arguments - assert "tool_name" in call_kwargs["input"] - assert call_kwargs["input"]["tool_name"] == "list_directory" - assert call_kwargs["input"]["arguments"] == tool_args - - # Verify output was logged - mock_span.log.assert_called_once() - log_call = mock_span.log.call_args[1] - assert "output" in log_call - - -@pytest.mark.asyncio -async def test_mcp_tool_error_handling(): - """Test that MCP tool errors are captured in spans.""" - with patch("braintrust.integrations.adk.tracing.start_span") as mock_start_span: - mock_span = MagicMock() - mock_span.__enter__ = MagicMock(return_value=mock_span) - mock_span.__exit__ = MagicMock(return_value=False) - mock_start_span.return_value = mock_span - - # Create mock tool that raises error - class MockMcpTool: - def __init__(self): - self.name = "failing_tool" - - async def run_async(self, *, args, tool_context): - raise ValueError("Tool execution failed") - - # Wrap the class - wrap_mcp_tool(MockMcpTool) - - # Create instance and call (should raise) - tool = MockMcpTool() - - with pytest.raises(ValueError, match="Tool execution failed"): - await tool.run_async(args={}, tool_context=None) - - # Verify error was logged to span - assert mock_span.log.called - # Check if error was logged - log_calls = [call for call in mock_span.log.call_args_list] - # Should have logged the error - - @pytest.mark.asyncio async def test_setup_adk_patches_mcp_tool(): - """Test that setup_adk automatically patches McpTool via ADKIntegration.""" - result = setup_adk(project_name="test") - assert result is True - - # Verify McpTool got patched. The google-adk nox matrix installs the - # optional MCP extra so this integration surface stays covered. - from braintrust.integrations.adk.patchers import McpToolPatcher - + """setup_adk patches the real McpTool class via ADKIntegration.""" + assert setup_adk(project_name="test") is True assert McpToolPatcher.is_patched(None, None), "McpTool should be patched" - - -@pytest.mark.asyncio -async def test_setup_adk_graceful_fallback_when_mcp_unavailable(): - """Test that setup_adk gracefully handles MCP not being installed.""" - # setup_adk delegates to ADKIntegration.setup() which handles ImportError - # in the McpToolPatcher gracefully - result = setup_adk(project_name="test") - - # Should succeed - MCP is optional - assert result is True - - -@pytest.mark.asyncio -async def test_mcp_tool_async_context_preservation(): - """Test that MCP tool spans handle async context switching correctly. - - This test reproduces the "was created in a different Context" error that occurs - when async generators yield control and resume in a different async context. - This is the issue we're seeing in the trace screenshot where mcp_tool spans - lose their parent context. - """ - import contextvars - - from braintrust.integrations.adk import wrap_mcp_tool - - # Track context switches - context_var = contextvars.ContextVar("test_context", default=None) - - class MockMcpTool: - def __init__(self): - self.name = "test_tool" - - async def run_async(self, *, args, tool_context): - # Simulate async work that might switch contexts - import asyncio - - await asyncio.sleep(0.001) - return {"result": "success"} - - # Wrap the tool - wrap_mcp_tool(MockMcpTool) - - # Create tool instance - tool = MockMcpTool() - - # Set initial context - context_var.set("initial") - - # Create an async generator that yields and switches contexts - async def context_switching_generator(): - # Call the tool (creates span) - context_var.set("during_call") - result = await tool.run_async(args={"test": "value"}, tool_context=None) - yield result - - # Switch context after yield - context_var.set("after_yield") - - # Execute the generator - this should trigger the context switch issue - results = [] - async for result in context_switching_generator(): - results.append(result) - - # Verify the tool executed successfully despite context switches - assert len(results) == 1 - assert results[0]["result"] == "success" - - # The test passes if no ValueError about "different Context" is raised - # The aclosing wrapper in __init__.py should suppress this error - - -@pytest.mark.asyncio -async def test_mcp_tool_nested_async_generators(): - """Test MCP tool execution within nested async generators. - - This simulates the real-world scenario where: - 1. Runner.run_async creates an async generator with a span - 2. Agent.run_async creates another async generator with a span - 3. MCP tool execution happens deep in the stack - 4. All generators yield and resume, potentially in different contexts - """ - from braintrust.integrations.adk import wrap_mcp_tool - - class MockMcpTool: - def __init__(self): - self.name = "nested_tool" - - async def run_async(self, *, args, tool_context): - import asyncio - - await asyncio.sleep(0.001) - return {"nested": "result"} - - wrap_mcp_tool(MockMcpTool) - tool = MockMcpTool() - - # Simulate nested async generators like Runner -> Agent -> Tool - async def outer_generator(): - """Simulates Runner.run_async""" - async for event in middle_generator(): - yield event - - async def middle_generator(): - """Simulates Agent.run_async""" - # Execute tool in the middle of generator execution - result = await tool.run_async(args={"nested": "test"}, tool_context=None) - yield {"type": "tool_result", "data": result} - - # Yield more events after tool execution - yield {"type": "final", "done": True} - - # Collect all events - events = [] - async for event in outer_generator(): - events.append(event) - - # Verify execution completed successfully - assert len(events) == 2 - assert events[0]["type"] == "tool_result" - assert events[0]["data"]["nested"] == "result" - assert events[1]["type"] == "final" - - # If we get here without ValueError, the context handling is working - - -@pytest.mark.asyncio -async def test_real_context_loss_with_braintrust_spans(): - """Test that demonstrates the actual context loss issue with real Braintrust spans. - - This test creates a scenario that matches the real-world issue: - 1. Create a span in an async generator - 2. Yield from that generator - 3. Try to clean up the span after context has switched - - This should trigger the "was created in a different Context" error that we're - suppressing in the aclosing.__aexit__ method. - """ - import asyncio - from contextlib import aclosing - - from braintrust import init_logger - - # Initialize a test logger - logger = init_logger(project="test-context-loss") - - # Track if we hit the context error - context_error_occurred = False - - async def problematic_generator(): - """Generator that creates a span and yields, simulating the Flow behavior.""" - from braintrust import start_span - - with start_span(name="test_span", type="task") as span: - # Yield some events - yield {"event": 1} - await asyncio.sleep(0.001) - yield {"event": 2} - # Span cleanup happens in __exit__, which may be in different context - - # Create a new async context and run the generator - async def outer_context(): - """Simulates the outer runner context.""" - events = [] - - # Use aclosing which has the error suppression - async with aclosing(problematic_generator()) as gen: - async for event in gen: - events.append(event) - # Force context switch - await asyncio.sleep(0.001) - - return events - - # Run in a fresh event loop context - events = await outer_context() - - # Verify we got the events - assert len(events) == 2 - assert events[0]["event"] == 1 - assert events[1]["event"] == 2 - - # If we get here without an unhandled ValueError, the suppression is working - # The aclosing.__aexit__ should have caught and suppressed any context errors diff --git a/py/src/braintrust/integrations/adk/tracing.py b/py/src/braintrust/integrations/adk/tracing.py index c6367e8e..72e14d08 100644 --- a/py/src/braintrust/integrations/adk/tracing.py +++ b/py/src/braintrust/integrations/adk/tracing.py @@ -4,7 +4,7 @@ import inspect import logging import time -from collections.abc import Iterable, Mapping +from collections.abc import Mapping from contextlib import aclosing from functools import lru_cache from itertools import chain @@ -118,13 +118,15 @@ def _serialize_pydantic_schema(schema_class: Any) -> dict[str, Any]: return {"__class__": schema_class.__name__ if inspect.isclass(schema_class) else str(type(schema_class).__name__)} -def _capture_config(config: Any) -> dict[str, Any] | Any: +def _capture_config(config: Any, exclude: tuple[str, ...] = ()) -> dict[str, Any] | Any: """ Capture the ADK config fields that make LLM spans readable. Google ADK uses these fields for schemas: - response_schema, response_json_schema (in GenerateContentConfig for LLM requests) - input_schema, output_schema (in agent config) + + ``exclude`` drops named fields (e.g. tools/tool_config that belong in metadata). """ if config is None or not config: return config @@ -143,9 +145,13 @@ def _capture_config(config: Any) -> dict[str, Any] | Any: "stop_sequences", "candidate_count", ] - captured: dict[str, Any] = dict(config) if isinstance(config, dict) else {} + captured: dict[str, Any] = ( + {k: v for k, v in config.items() if k not in exclude} if isinstance(config, dict) else {} + ) for field in config_fields: + if field in exclude: + continue value = _get_field(config, field) if value is None: continue @@ -163,10 +169,6 @@ def _capture_config(config: Any) -> dict[str, Any] | Any: return captured or config -def _omit(obj: Any, keys: Iterable[str]): - return {k: v for k, v in obj.items() if k not in keys} - - def _extract_metrics(response: Any) -> dict[str, float] | None: """Extract token usage metrics from Google GenAI response.""" if not response: @@ -252,13 +254,34 @@ def _capture_llm_request_input(llm_request: Any) -> Any: [_serialize_content(c) for c in contents] if isinstance(contents, list) else _serialize_content(contents) ) if config: - captured["config"] = _capture_config(config) + captured["config"] = _capture_config(config, exclude=("tools", "tool_config")) if live_connect_config is not None or hasattr(llm_request, "live_connect_config") or isinstance(llm_request, dict): captured["live_connect_config"] = live_connect_config return captured or llm_request +def _extract_tool_metadata(llm_request: Any) -> dict[str, Any]: + """Extract tool definitions and tool_config from an ADK LLM request for metadata.tools. + + Google-native shape is preserved; ADK is Google-backed and metadata.provider="google" + lets the UI apply the Google normalizer. + """ + if llm_request is None: + return {} + config = _get_field(llm_request, "config") + if config is None: + return {} + result: dict[str, Any] = {} + tools = _get_field(config, "tools") + if tools: + result["tools"] = bt_safe_deep_copy(tools) + tool_config = _get_field(config, "tool_config") + if tool_config: + result["tool_config"] = bt_safe_deep_copy(tool_config) + return result + + def _event_output_with_content(last_event: Any, event_with_content: Any | None) -> Any: if event_with_content is None or _get_field(last_event, "content") is not None: return last_event @@ -349,13 +372,11 @@ def _run_in_context(*target_args: Any, **target_kwargs: Any) -> Any: async def _agent_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): - parent_context = args[0] if len(args) > 0 else kwargs.get("parent_context") - async def _trace(): with start_span( name=f"agent_run [{instance.name}]", type=SpanTypeAttribute.TASK, - metadata={"parent_context": parent_context, **_omit(kwargs, ["parent_context"])}, + metadata={"agent_name": instance.name}, ) as agent_span: last_event = None async with aclosing(wrapped(*args, **kwargs)) as agen: @@ -372,16 +393,11 @@ async def _trace(): async def _flow_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): - invocation_context = args[0] if len(args) > 0 else kwargs.get("invocation_context") - async def _trace(): with start_span( name="call_llm", type=SpanTypeAttribute.TASK, - metadata={ - "invocation_context": invocation_context, - **_omit(kwargs, ["invocation_context"]), - }, + metadata={"flow_class": instance.__class__.__name__}, ) as llm_span: last_event = None async with aclosing(wrapped(*args, **kwargs)) as agen: @@ -397,9 +413,7 @@ async def _trace(): async def _flow_call_llm_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): - invocation_context = args[0] if len(args) > 0 else kwargs.get("invocation_context") llm_request = args[1] if len(args) > 1 else kwargs.get("llm_request") - model_response_event = args[2] if len(args) > 2 else kwargs.get("model_response_event") async def _trace(): # Capture only the fields we need to alter: contents may contain binary @@ -410,19 +424,20 @@ async def _trace(): # Extract model name from request or instance model_name = _extract_model_name(None, llm_request, instance) + metadata: dict[str, Any] = { + "flow_class": instance.__class__.__name__, + "model": model_name, + "provider": "google", + } + metadata.update(_extract_tool_metadata(llm_request)) + # Create span BEFORE execution so child spans (like mcp_tool) have proper parent # Start with generic name - we'll update it after we see the response with start_span( name="llm_call", type=SpanTypeAttribute.LLM, input=captured_request, - metadata={ - "invocation_context": invocation_context, - "model_response_event": model_response_event, - "flow_class": instance.__class__.__name__, - "model": model_name, - **_omit(kwargs, ["invocation_context", "model_response_event", "flow_class", "llm_call_type"]), - }, + metadata=metadata, ) as llm_span: # Execute the LLM call and yield events while span is active last_event = None @@ -487,10 +502,10 @@ async def _trace(): type=SpanTypeAttribute.TASK, input={"new_message": serialized_message}, metadata={ + "app_name": instance.app_name, "user_id": user_id, "session_id": session_id, "state_delta": state_delta, - **_omit(kwargs, ["user_id", "session_id", "new_message", "state_delta"]), }, ) as runner_span: last_event = None @@ -528,7 +543,7 @@ async def _tool_call_async_wrapper(wrapped: Any, instance: Any, args: Any, kwarg tool_span.log(output=result) return result except Exception as e: - tool_span.log(error=str(e)) + tool_span.log(error=e) raise @@ -541,7 +556,7 @@ async def _mcp_tool_run_async_wrapper_async(wrapped: Any, instance: Any, args: A name=f"mcp_tool [{tool_name}]", type=SpanTypeAttribute.TOOL, input={"tool_name": tool_name, "arguments": tool_args}, - metadata=_omit(kwargs, ["args"]), + metadata={"tool_class": instance.__class__.__name__}, ) as tool_span: try: result = await wrapped(*args, **kwargs) @@ -549,5 +564,5 @@ async def _mcp_tool_run_async_wrapper_async(wrapped: Any, instance: Any, args: A return result except Exception as e: # Log error to span but re-raise for ADK to handle - tool_span.log(error=str(e)) + tool_span.log(error=e) raise From 00013d6363fc1adf0a60869b354e9cc339bcf2b3 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 15:30:14 +0000 Subject: [PATCH 02/11] chore(adk): drop test_adk_mcp_tool.py The two remaining tests only asserted patcher wiring (marker set, setup_adk returns True). Patcher idempotence is enforced by the base patcher and covered indirectly through the auto-instrument subprocess test; real MCP tool span behavior is covered end-to-end by the VCR-backed suite in test_adk.py. Co-Authored-By: Claude Opus 4.7 --- py/noxfile.py | 1 - .../integrations/adk/test_adk_mcp_tool.py | 30 ------------------- 2 files changed, 31 deletions(-) delete mode 100644 py/src/braintrust/integrations/adk/test_adk_mcp_tool.py diff --git a/py/noxfile.py b/py/noxfile.py index 32a437a2..22c8f64a 100644 --- a/py/noxfile.py +++ b/py/noxfile.py @@ -513,7 +513,6 @@ def test_google_adk(session, version): _install_test_deps(session) _install_matrix_dep(session, "google-adk", version) _run_tests(session, f"{INTEGRATION_DIR}/adk/test_adk.py", version=version) - _run_tests(session, f"{INTEGRATION_DIR}/adk/test_adk_mcp_tool.py", version=version) LANGCHAIN_VERSIONS = _get_matrix_versions("langchain-core") diff --git a/py/src/braintrust/integrations/adk/test_adk_mcp_tool.py b/py/src/braintrust/integrations/adk/test_adk_mcp_tool.py deleted file mode 100644 index b7d7c827..00000000 --- a/py/src/braintrust/integrations/adk/test_adk_mcp_tool.py +++ /dev/null @@ -1,30 +0,0 @@ -"""Tests for MCP tool tracing integration. - -Behavioral coverage of the MCP tool span is provided by the VCR-backed suite in -``test_adk.py`` — when an ADK agent invokes a real tool through the flow, the -``tool [...]`` / ``mcp_tool [...]`` span is exercised end-to-end. The tests here -only cover the patcher wiring, which does not require a live MCP server. -""" - -import pytest -from braintrust.integrations.adk import setup_adk, wrap_mcp_tool -from braintrust.integrations.adk.patchers import McpToolPatcher - - -@pytest.mark.asyncio -async def test_wrap_mcp_tool_marks_as_patched(): - """wrap_mcp_tool marks the class via the patcher marker (idempotence signal).""" - - class MockMcpTool: - async def run_async(self, *, args, tool_context): - return {"result": "success"} - - wrapped_class = wrap_mcp_tool(MockMcpTool) - assert getattr(wrapped_class, McpToolPatcher.patch_marker_attr(), False) - - -@pytest.mark.asyncio -async def test_setup_adk_patches_mcp_tool(): - """setup_adk patches the real McpTool class via ADKIntegration.""" - assert setup_adk(project_name="test") is True - assert McpToolPatcher.is_patched(None, None), "McpTool should be patched" From f197cf923a2947f7a92792dffa4c8ee9bdaffbc1 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 15:34:39 +0000 Subject: [PATCH 03/11] refactor(adk): make _capture_config a strict allowlist; drop deep-copy on tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _capture_config no longer seeds from dict(config); only fields in _CAPTURED_CONFIG_FIELDS are returned. Removes the pass-through behavior that let tools/tool_config leak into input.config (previously patched with an `exclude=` bandaid). - _extract_tool_metadata no longer wraps tools/tool_config in bt_safe_deep_copy — the log pipeline serializes Pydantic objects at send time (SKILL: "do not over-serialize"). - Update the two _capture_config unit tests to reflect the strict allowlist behavior. Co-Authored-By: Claude Opus 4.7 --- .../braintrust/integrations/adk/test_adk.py | 12 ++--- py/src/braintrust/integrations/adk/tracing.py | 54 +++++++++---------- 2 files changed, 33 insertions(+), 33 deletions(-) diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index d234344a..054ba6a2 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -1501,7 +1501,7 @@ class TestSchema(BaseModel): "response_json_schema": TestSchema, "input_schema": TestSchema, "output_schema": TestSchema, - "other_field": "keep me", + "other_field": "drop me", } serialized = _capture_config(config) @@ -1518,8 +1518,8 @@ class TestSchema(BaseModel): assert "value" in properties assert properties["value"]["description"] == "Test value" - # Other fields should be preserved - assert "other_field" in serialized + # _capture_config is a strict allowlist — non-listed fields are dropped + assert "other_field" not in serialized @pytest.mark.asyncio @@ -1527,15 +1527,15 @@ async def test_capture_config_handles_non_pydantic(): """Test that _capture_config handles non-Pydantic values gracefully.""" from braintrust.integrations.adk.tracing import _capture_config - # Test with non-Pydantic values + # Test with non-Pydantic values — allowlisted field passes through as-is, + # non-listed field is dropped. config = {"response_schema": "not a pydantic model", "other_field": {"key": "value"}} serialized = _capture_config(config) assert isinstance(serialized, dict) - # Non-Pydantic schema should remain as-is - assert "response_schema" in serialized assert serialized["response_schema"] == "not a pydantic model" + assert "other_field" not in serialized @pytest.mark.asyncio diff --git a/py/src/braintrust/integrations/adk/tracing.py b/py/src/braintrust/integrations/adk/tracing.py index 72e14d08..fc1fa27f 100644 --- a/py/src/braintrust/integrations/adk/tracing.py +++ b/py/src/braintrust/integrations/adk/tracing.py @@ -118,40 +118,40 @@ def _serialize_pydantic_schema(schema_class: Any) -> dict[str, Any]: return {"__class__": schema_class.__name__ if inspect.isclass(schema_class) else str(type(schema_class).__name__)} -def _capture_config(config: Any, exclude: tuple[str, ...] = ()) -> dict[str, Any] | Any: +_CAPTURED_CONFIG_FIELDS = ( + "system_instruction", + "response_mime_type", + "response_schema", + "response_json_schema", + "input_schema", + "output_schema", + "max_output_tokens", + "temperature", + "top_p", + "top_k", + "stop_sequences", + "candidate_count", +) + + +def _capture_config(config: Any) -> dict[str, Any] | Any: """ Capture the ADK config fields that make LLM spans readable. + Strict allowlist — only ``_CAPTURED_CONFIG_FIELDS`` are returned. Callers + that need additional fields (e.g. ``tools`` for ``metadata.tools``) should + extract them separately. Returns the original ``config`` untouched when it + is falsy or when no allowlisted field is set. + Google ADK uses these fields for schemas: - response_schema, response_json_schema (in GenerateContentConfig for LLM requests) - input_schema, output_schema (in agent config) - - ``exclude`` drops named fields (e.g. tools/tool_config that belong in metadata). """ if config is None or not config: return config - config_fields = [ - "system_instruction", - "response_mime_type", - "response_schema", - "response_json_schema", - "input_schema", - "output_schema", - "max_output_tokens", - "temperature", - "top_p", - "top_k", - "stop_sequences", - "candidate_count", - ] - captured: dict[str, Any] = ( - {k: v for k, v in config.items() if k not in exclude} if isinstance(config, dict) else {} - ) - - for field in config_fields: - if field in exclude: - continue + captured: dict[str, Any] = {} + for field in _CAPTURED_CONFIG_FIELDS: value = _get_field(config, field) if value is None: continue @@ -254,7 +254,7 @@ def _capture_llm_request_input(llm_request: Any) -> Any: [_serialize_content(c) for c in contents] if isinstance(contents, list) else _serialize_content(contents) ) if config: - captured["config"] = _capture_config(config, exclude=("tools", "tool_config")) + captured["config"] = _capture_config(config) if live_connect_config is not None or hasattr(llm_request, "live_connect_config") or isinstance(llm_request, dict): captured["live_connect_config"] = live_connect_config @@ -275,10 +275,10 @@ def _extract_tool_metadata(llm_request: Any) -> dict[str, Any]: result: dict[str, Any] = {} tools = _get_field(config, "tools") if tools: - result["tools"] = bt_safe_deep_copy(tools) + result["tools"] = tools tool_config = _get_field(config, "tool_config") if tool_config: - result["tool_config"] = bt_safe_deep_copy(tool_config) + result["tool_config"] = tool_config return result From 5378dff819195da8ab0a1fefbe157123a4d425ed Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 15:36:45 +0000 Subject: [PATCH 04/11] chore(adk): drop explanatory comments in test assertions Co-Authored-By: Claude Opus 4.7 --- py/src/braintrust/integrations/adk/test_adk.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index 054ba6a2..31c77bb8 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -482,8 +482,6 @@ async def test_adk_braintrust_integration(memory_logger): assert function_call["name"] == "get_weather" assert function_call["args"]["location"] == "San Francisco" - # Every integration-owned span carries the adk-auto instrumentation name - # (SKILL: "Every span an integration creates MUST carry ... instrumentation.name") adk_spans = [ row for row in spans @@ -494,8 +492,6 @@ async def test_adk_braintrust_integration(memory_logger): f"adk-auto origin missing on task/llm/tool spans: {span_types_by_origin}" ) - # Every llm span carries metadata.model, metadata.provider=google, and metadata.tools - # (with tools NOT leaking into input.config). for span in llm_spans: meta = span["metadata"] assert meta.get("provider") == "google", f"Missing metadata.provider=google on {span['span_attributes']['name']}" From 224a476ee2040796b7a971bc3544fec9d6b12720 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 15:37:04 +0000 Subject: [PATCH 05/11] chore(adk): trim _capture_config docstring Co-Authored-By: Claude Opus 4.7 --- py/src/braintrust/integrations/adk/tracing.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/py/src/braintrust/integrations/adk/tracing.py b/py/src/braintrust/integrations/adk/tracing.py index fc1fa27f..0672e123 100644 --- a/py/src/braintrust/integrations/adk/tracing.py +++ b/py/src/braintrust/integrations/adk/tracing.py @@ -135,18 +135,7 @@ def _serialize_pydantic_schema(schema_class: Any) -> dict[str, Any]: def _capture_config(config: Any) -> dict[str, Any] | Any: - """ - Capture the ADK config fields that make LLM spans readable. - - Strict allowlist — only ``_CAPTURED_CONFIG_FIELDS`` are returned. Callers - that need additional fields (e.g. ``tools`` for ``metadata.tools``) should - extract them separately. Returns the original ``config`` untouched when it - is falsy or when no allowlisted field is set. - - Google ADK uses these fields for schemas: - - response_schema, response_json_schema (in GenerateContentConfig for LLM requests) - - input_schema, output_schema (in agent config) - """ + """Capture the ADK config fields that make LLM spans readable.""" if config is None or not config: return config From 9e752f61b23d4125e4a8fbb30c5647cb30460090 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 15:43:21 +0000 Subject: [PATCH 06/11] refactor(adk): drop _get_field, use plain getattr MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Production ADK always passes Pydantic objects (LlmRequest, Content, Event, Part, GenerateContentConfig) — the dict-vs-attribute duality existed only for unit tests. Rewrite the _determine_llm_call_type and _capture_config unit tests to build fixtures with SimpleNamespace via a small _ns() helper so getattr works uniformly. Also drop the now- unused Mapping import. Co-Authored-By: Claude Opus 4.7 --- .../braintrust/integrations/adk/test_adk.py | 107 ++++++++---------- py/src/braintrust/integrations/adk/tracing.py | 41 +++---- 2 files changed, 63 insertions(+), 85 deletions(-) diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index 31c77bb8..2ef2d7df 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -922,12 +922,26 @@ async def test_adk_captures_metrics(memory_logger): assert metadata.get("provider") == "google", "Metadata should include provider=google" +def _ns(obj): + """Recursively convert nested dicts/lists into SimpleNamespace for attribute access. + + _determine_llm_call_type / _capture_config read production ADK Pydantic + objects via getattr; these test fixtures mirror that access pattern. + """ + from types import SimpleNamespace + + if isinstance(obj, dict): + return SimpleNamespace(**{k: _ns(v) for k, v in obj.items()}) + if isinstance(obj, list): + return [_ns(item) for item in obj] + return obj + + def test_determine_llm_call_type_direct_response(): """Test that _determine_llm_call_type returns 'direct_response' when tools are available but not used.""" from braintrust.integrations.adk.tracing import _determine_llm_call_type - # Request with tools available - llm_request = { + llm_request = _ns({ "config": { "tools": [ { @@ -939,75 +953,63 @@ def test_determine_llm_call_type_direct_response(): ] }, "contents": [{"parts": [{"text": "What is 2+2?"}], "role": "user"}], - } + }) - # Response without function calls - model_response = { + model_response = _ns({ "content": {"parts": [{"text": "4\n"}], "role": "model"}, "finish_reason": "STOP", - } + }) - call_type = _determine_llm_call_type(llm_request, model_response) - assert call_type == "direct_response", "Should be direct_response when tools available but not used" + assert _determine_llm_call_type(llm_request, model_response) == "direct_response" def test_determine_llm_call_type_tool_selection(): """Test that _determine_llm_call_type returns 'tool_selection' when LLM calls a tool.""" from braintrust.integrations.adk.tracing import _determine_llm_call_type - # Request with tools available - llm_request = { + llm_request = _ns({ "config": { - "tools": [ - { - "function_declarations": [ - {"name": "get_weather", "description": "Get weather"}, - ] - } - ] + "tools": [{"function_declarations": [{"name": "get_weather", "description": "Get weather"}]}] }, "contents": [{"parts": [{"text": "What's the weather?"}], "role": "user"}], - } + }) - # Response with function call (camelCase) - model_response = { + # camelCase function call + model_response = _ns({ "content": { "parts": [{"functionCall": {"name": "get_weather", "args": {"location": "SF"}}}], "role": "model", }, - } + }) - call_type = _determine_llm_call_type(llm_request, model_response) - assert call_type == "tool_selection", "Should be tool_selection when LLM calls a tool" + assert _determine_llm_call_type(llm_request, model_response) == "tool_selection" def test_determine_llm_call_type_tool_selection_snake_case(): """Test that _determine_llm_call_type handles snake_case function_call.""" from braintrust.integrations.adk.tracing import _determine_llm_call_type - llm_request = { + llm_request = _ns({ "config": {"tools": [{"function_declarations": [{"name": "search"}]}]}, "contents": [{"parts": [{"text": "Search for pizza"}], "role": "user"}], - } + }) - # Response with function call (snake_case) - model_response = { + # snake_case function_call + model_response = _ns({ "content": { "parts": [{"function_call": {"name": "search", "args": {"query": "pizza"}}}], "role": "model", }, - } + }) - call_type = _determine_llm_call_type(llm_request, model_response) - assert call_type == "tool_selection", "Should be tool_selection for snake_case function_call" + assert _determine_llm_call_type(llm_request, model_response) == "tool_selection" def test_determine_llm_call_type_response_generation(): """Test that _determine_llm_call_type returns 'response_generation' after tool execution.""" from braintrust.integrations.adk.tracing import _determine_llm_call_type - # Request with function_response in history - llm_request = { + llm_request = _ns({ "config": {"tools": [{"function_declarations": [{"name": "get_weather"}]}]}, "contents": [ {"parts": [{"text": "What's the weather?"}], "role": "user"}, @@ -1017,46 +1019,33 @@ def test_determine_llm_call_type_response_generation(): "role": "user", }, ], - } + }) - # Response after tool execution - model_response = { - "content": {"parts": [{"text": "It's 72 degrees"}], "role": "model"}, - } + model_response = _ns({"content": {"parts": [{"text": "It's 72 degrees"}], "role": "model"}}) - call_type = _determine_llm_call_type(llm_request, model_response) - assert call_type == "response_generation", "Should be response_generation after tool execution" + assert _determine_llm_call_type(llm_request, model_response) == "response_generation" def test_determine_llm_call_type_no_tools(): """Test that _determine_llm_call_type returns 'direct_response' when no tools configured.""" from braintrust.integrations.adk.tracing import _determine_llm_call_type - llm_request = { - "config": {}, - "contents": [{"parts": [{"text": "Hello"}], "role": "user"}], - } - - model_response = { - "content": {"parts": [{"text": "Hi there"}], "role": "model"}, - } + llm_request = _ns({"config": {}, "contents": [{"parts": [{"text": "Hello"}], "role": "user"}]}) + model_response = _ns({"content": {"parts": [{"text": "Hi there"}], "role": "model"}}) - call_type = _determine_llm_call_type(llm_request, model_response) - assert call_type == "direct_response", "Should be direct_response when no tools configured" + assert _determine_llm_call_type(llm_request, model_response) == "direct_response" def test_determine_llm_call_type_no_response(): """Test that _determine_llm_call_type handles missing model_response gracefully.""" from braintrust.integrations.adk.tracing import _determine_llm_call_type - llm_request = { + llm_request = _ns({ "config": {"tools": [{"function_declarations": [{"name": "tool1"}]}]}, "contents": [{"parts": [{"text": "Test"}], "role": "user"}], - } + }) - # No model_response provided - call_type = _determine_llm_call_type(llm_request, None) - assert call_type == "direct_response", "Should default to direct_response when no response available" + assert _determine_llm_call_type(llm_request, None) == "direct_response" @pytest.mark.asyncio @@ -1491,14 +1480,13 @@ async def test_capture_config_handles_all_schema_fields(): class TestSchema(BaseModel): value: str = Field(description="Test value") - # Test with a dict config that has all schema fields - config = { + config = _ns({ "response_schema": TestSchema, "response_json_schema": TestSchema, "input_schema": TestSchema, "output_schema": TestSchema, "other_field": "drop me", - } + }) serialized = _capture_config(config) @@ -1523,9 +1511,8 @@ async def test_capture_config_handles_non_pydantic(): """Test that _capture_config handles non-Pydantic values gracefully.""" from braintrust.integrations.adk.tracing import _capture_config - # Test with non-Pydantic values — allowlisted field passes through as-is, - # non-listed field is dropped. - config = {"response_schema": "not a pydantic model", "other_field": {"key": "value"}} + # Allowlisted field passes through as-is; non-listed field is dropped. + config = _ns({"response_schema": "not a pydantic model", "other_field": {"key": "value"}}) serialized = _capture_config(config) diff --git a/py/src/braintrust/integrations/adk/tracing.py b/py/src/braintrust/integrations/adk/tracing.py index 0672e123..fa47fc97 100644 --- a/py/src/braintrust/integrations/adk/tracing.py +++ b/py/src/braintrust/integrations/adk/tracing.py @@ -4,7 +4,6 @@ import inspect import logging import time -from collections.abc import Mapping from contextlib import aclosing from functools import lru_cache from itertools import chain @@ -141,7 +140,7 @@ def _capture_config(config: Any) -> dict[str, Any] | Any: captured: dict[str, Any] = {} for field in _CAPTURED_CONFIG_FIELDS: - value = _get_field(config, field) + value = getattr(config, field, None) if value is None: continue if inspect.isclass(value): @@ -217,12 +216,8 @@ def _extract_model_name(response: Any, llm_request: Any, instance: Any) -> str | return None -def _get_field(value: Any, field: str, default: Any = None) -> Any: - return value.get(field, default) if isinstance(value, Mapping) else getattr(value, field, default) - - def _part_has_field(part: Any, *field_names: str) -> bool: - return any(_get_field(part, field_name) is not None for field_name in field_names) + return any(getattr(part, field_name, None) is not None for field_name in field_names) def _capture_llm_request_input(llm_request: Any) -> Any: @@ -230,10 +225,9 @@ def _capture_llm_request_input(llm_request: Any) -> Any: if llm_request is None: return None - contents = _get_field(llm_request, "contents") - config = _get_field(llm_request, "config") - model = _get_field(llm_request, "model") - live_connect_config = _get_field(llm_request, "live_connect_config") + contents = getattr(llm_request, "contents", None) + config = getattr(llm_request, "config", None) + model = getattr(llm_request, "model", None) captured: dict[str, Any] = {} if model: @@ -244,8 +238,8 @@ def _capture_llm_request_input(llm_request: Any) -> Any: ) if config: captured["config"] = _capture_config(config) - if live_connect_config is not None or hasattr(llm_request, "live_connect_config") or isinstance(llm_request, dict): - captured["live_connect_config"] = live_connect_config + if hasattr(llm_request, "live_connect_config"): + captured["live_connect_config"] = getattr(llm_request, "live_connect_config", None) return captured or llm_request @@ -258,30 +252,27 @@ def _extract_tool_metadata(llm_request: Any) -> dict[str, Any]: """ if llm_request is None: return {} - config = _get_field(llm_request, "config") + config = getattr(llm_request, "config", None) if config is None: return {} result: dict[str, Any] = {} - tools = _get_field(config, "tools") + tools = getattr(config, "tools", None) if tools: result["tools"] = tools - tool_config = _get_field(config, "tool_config") + tool_config = getattr(config, "tool_config", None) if tool_config: result["tool_config"] = tool_config return result def _event_output_with_content(last_event: Any, event_with_content: Any | None) -> Any: - if event_with_content is None or _get_field(last_event, "content") is not None: + if event_with_content is None or getattr(last_event, "content", None) is not None: return last_event - content = _get_field(event_with_content, "content") + content = getattr(event_with_content, "content", None) if content is None: return last_event - if isinstance(last_event, dict): - return {**last_event, "content": content} - # Keep the original event instead of recursively serializing it; add the # captured content alongside it so Braintrust can serialize both values. return {"event": last_event, "content": content} @@ -299,8 +290,8 @@ def _determine_llm_call_type(llm_request: Any, model_response: Any = None) -> st try: has_function_response = any( _part_has_field(part, "function_response", "functionResponse") - for content in (_get_field(llm_request, "contents", []) or []) - for part in (_get_field(content, "parts", []) or []) + for content in (getattr(llm_request, "contents", None) or []) + for part in (getattr(content, "parts", None) or []) ) response_has_function_call = False @@ -313,11 +304,11 @@ def _determine_llm_call_type(llm_request: Any, model_response: Any = None) -> st pass if not response_has_function_call: - content = _get_field(model_response, "content") + content = getattr(model_response, "content", None) response_has_function_call = any( _part_has_field(part, "function_call", "functionCall") for part in chain( - _get_field(content, "parts", []) or [], _get_field(model_response, "parts", []) or [] + getattr(content, "parts", None) or [], getattr(model_response, "parts", None) or [] ) ) From 3bc65a44ec603c66f1f8b638867ba4d0a2d0f417 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 15:47:48 +0000 Subject: [PATCH 07/11] test(adk): drop fake-fixture unit tests, lean on VCR coverage The 6 test_determine_llm_call_type_* and 2 test_capture_config_* tests built SimpleNamespace fixtures to exercise internal helpers. Every path they covered is already asserted through real ADK calls in the VCR-backed suite: - tool_selection / response_generation: test_adk_braintrust_integration - direct_response: test_adk_captures_metrics (now asserts the name) - response_schema Pydantic: test_adk_structured_output_pydantic - response_schema nested: test_adk_complex_nested_schema - input_schema: test_adk_input_schema_serialization - response_json_schema dict: test_adk_response_json_schema_dict Also drops the _ns() helper. Co-Authored-By: Claude Opus 4.7 --- .../braintrust/integrations/adk/test_adk.py | 183 ++---------------- 1 file changed, 13 insertions(+), 170 deletions(-) diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index 2ef2d7df..236645d6 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -921,131 +921,16 @@ async def test_adk_captures_metrics(memory_logger): assert metadata["model"] == ADK_MODEL, "Model name should match the agent's model" assert metadata.get("provider") == "google", "Metadata should include provider=google" + # No tools configured — _determine_llm_call_type should mark this as direct_response + assert "direct_response" in llm_span_with_metrics["span_attributes"]["name"], ( + f"Expected direct_response call type, got {llm_span_with_metrics['span_attributes']['name']}" + ) -def _ns(obj): - """Recursively convert nested dicts/lists into SimpleNamespace for attribute access. - - _determine_llm_call_type / _capture_config read production ADK Pydantic - objects via getattr; these test fixtures mirror that access pattern. - """ - from types import SimpleNamespace - - if isinstance(obj, dict): - return SimpleNamespace(**{k: _ns(v) for k, v in obj.items()}) - if isinstance(obj, list): - return [_ns(item) for item in obj] - return obj - - -def test_determine_llm_call_type_direct_response(): - """Test that _determine_llm_call_type returns 'direct_response' when tools are available but not used.""" - from braintrust.integrations.adk.tracing import _determine_llm_call_type - - llm_request = _ns({ - "config": { - "tools": [ - { - "function_declarations": [ - {"name": "read_file", "description": "Read a file"}, - {"name": "list_directory", "description": "List directory"}, - ] - } - ] - }, - "contents": [{"parts": [{"text": "What is 2+2?"}], "role": "user"}], - }) - - model_response = _ns({ - "content": {"parts": [{"text": "4\n"}], "role": "model"}, - "finish_reason": "STOP", - }) - - assert _determine_llm_call_type(llm_request, model_response) == "direct_response" - - -def test_determine_llm_call_type_tool_selection(): - """Test that _determine_llm_call_type returns 'tool_selection' when LLM calls a tool.""" - from braintrust.integrations.adk.tracing import _determine_llm_call_type - - llm_request = _ns({ - "config": { - "tools": [{"function_declarations": [{"name": "get_weather", "description": "Get weather"}]}] - }, - "contents": [{"parts": [{"text": "What's the weather?"}], "role": "user"}], - }) - - # camelCase function call - model_response = _ns({ - "content": { - "parts": [{"functionCall": {"name": "get_weather", "args": {"location": "SF"}}}], - "role": "model", - }, - }) - - assert _determine_llm_call_type(llm_request, model_response) == "tool_selection" - - -def test_determine_llm_call_type_tool_selection_snake_case(): - """Test that _determine_llm_call_type handles snake_case function_call.""" - from braintrust.integrations.adk.tracing import _determine_llm_call_type - - llm_request = _ns({ - "config": {"tools": [{"function_declarations": [{"name": "search"}]}]}, - "contents": [{"parts": [{"text": "Search for pizza"}], "role": "user"}], - }) - - # snake_case function_call - model_response = _ns({ - "content": { - "parts": [{"function_call": {"name": "search", "args": {"query": "pizza"}}}], - "role": "model", - }, - }) - - assert _determine_llm_call_type(llm_request, model_response) == "tool_selection" - - -def test_determine_llm_call_type_response_generation(): - """Test that _determine_llm_call_type returns 'response_generation' after tool execution.""" - from braintrust.integrations.adk.tracing import _determine_llm_call_type - - llm_request = _ns({ - "config": {"tools": [{"function_declarations": [{"name": "get_weather"}]}]}, - "contents": [ - {"parts": [{"text": "What's the weather?"}], "role": "user"}, - {"parts": [{"functionCall": {"name": "get_weather", "args": {}}}], "role": "model"}, - { - "parts": [{"function_response": {"name": "get_weather", "response": {"temp": "72F"}}}], - "role": "user", - }, - ], - }) - - model_response = _ns({"content": {"parts": [{"text": "It's 72 degrees"}], "role": "model"}}) - - assert _determine_llm_call_type(llm_request, model_response) == "response_generation" - - -def test_determine_llm_call_type_no_tools(): - """Test that _determine_llm_call_type returns 'direct_response' when no tools configured.""" - from braintrust.integrations.adk.tracing import _determine_llm_call_type - - llm_request = _ns({"config": {}, "contents": [{"parts": [{"text": "Hello"}], "role": "user"}]}) - model_response = _ns({"content": {"parts": [{"text": "Hi there"}], "role": "model"}}) - - assert _determine_llm_call_type(llm_request, model_response) == "direct_response" - - -def test_determine_llm_call_type_no_response(): - """Test that _determine_llm_call_type handles missing model_response gracefully.""" - from braintrust.integrations.adk.tracing import _determine_llm_call_type - - llm_request = _ns({ - "config": {"tools": [{"function_declarations": [{"name": "tool1"}]}]}, - "contents": [{"parts": [{"text": "Test"}], "role": "user"}], - }) - assert _determine_llm_call_type(llm_request, None) == "direct_response" +# _determine_llm_call_type paths are exercised through the VCR-backed +# integration tests: `test_adk_braintrust_integration` asserts the +# tool_selection / response_generation span names, and the direct_response +# path is asserted in `test_adk_captures_metrics`. @pytest.mark.asyncio @@ -1472,53 +1357,11 @@ class Person(BaseModel): assert output["avg_logprobs"] is not None -@pytest.mark.asyncio -async def test_capture_config_handles_all_schema_fields(): - """Test that _capture_config handles all 4 schema fields.""" - from braintrust.integrations.adk.tracing import _capture_config - - class TestSchema(BaseModel): - value: str = Field(description="Test value") - - config = _ns({ - "response_schema": TestSchema, - "response_json_schema": TestSchema, - "input_schema": TestSchema, - "output_schema": TestSchema, - "other_field": "drop me", - }) - - serialized = _capture_config(config) - - assert isinstance(serialized, dict) - - # All schema fields should be serialized to JSON Schema format - for field in ["response_schema", "response_json_schema", "input_schema", "output_schema"]: - assert field in serialized, f"Missing {field}" - schema = serialized[field] - assert isinstance(schema, dict) - properties = schema.get("properties") - assert isinstance(properties, dict) - assert "value" in properties - assert properties["value"]["description"] == "Test value" - - # _capture_config is a strict allowlist — non-listed fields are dropped - assert "other_field" not in serialized - - -@pytest.mark.asyncio -async def test_capture_config_handles_non_pydantic(): - """Test that _capture_config handles non-Pydantic values gracefully.""" - from braintrust.integrations.adk.tracing import _capture_config - - # Allowlisted field passes through as-is; non-listed field is dropped. - config = _ns({"response_schema": "not a pydantic model", "other_field": {"key": "value"}}) - - serialized = _capture_config(config) - - assert isinstance(serialized, dict) - assert serialized["response_schema"] == "not a pydantic model" - assert "other_field" not in serialized +# _capture_config's allowlisted fields are exercised through the VCR-backed +# integration tests: response_schema (`test_adk_structured_output_pydantic`, +# `test_adk_complex_nested_schema`), input_schema (`test_adk_input_schema_serialization`), +# response_json_schema (`test_adk_response_json_schema_dict`), and the sampling +# params (`test_adk_generation_config_is_logged`). @pytest.mark.asyncio From 5cda5fb360931903f42ff70c9241ff087cabee87 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 15:58:13 +0000 Subject: [PATCH 08/11] test(adk): drop remaining fake-LLM and non-ADK unit tests Removes tests that used fake LlmRegistry entries, MagicMock, or exercised code that isn't ADK-specific: - test_adk_thread_context_propagation: folded into test_adk_sync_runner_run_does_not_duplicate_invocation_spans by wrapping runner.run() in an outer parent span and asserting every adk-auto span shares its root_span_id. - test_adk_generation_config_is_logged: folded into test_adk_max_tokens_captures_content (asserts max_output_tokens and temperature against the real logged span input). - test_adk_document_inline_data_attachment_conversion: same code path as test_adk_binary_data_attachment_conversion (VCR-backed image). - test_llm_call_span_wraps_child_spans (MagicMock Flow), and the test_serialize_*, test_bt_safe_deep_copy_*, test_serialize_pydantic_ schema_direct, test_capture_config_preserves_none, test_adk_agent_metadata_with_attachment (patched Attachment.upload), test_adk_bytes_and_attachment_in_structure, and test_async_context_preservation_across_yields tests: covered by the VCR suite or not ADK-specific. The one non-VCR test kept is test_create_thread_wrapper_exception_does_not_double_invoke_target - a plain-Python regression test for a wrapper bug (no mocks, no HTTP). Co-Authored-By: Claude Opus 4.7 --- ...st_adk_agent_metadata_with_attachment.yaml | 254 ------- ...st_adk_agent_metadata_with_attachment.yaml | 142 ---- .../braintrust/integrations/adk/test_adk.py | 674 +----------------- 3 files changed, 28 insertions(+), 1042 deletions(-) delete mode 100644 py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_agent_metadata_with_attachment.yaml delete mode 100644 py/src/braintrust/integrations/adk/cassettes/latest/test_adk_agent_metadata_with_attachment.yaml diff --git a/py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_agent_metadata_with_attachment.yaml b/py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_agent_metadata_with_attachment.yaml deleted file mode 100644 index 7d0fe0bc..00000000 --- a/py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_agent_metadata_with_attachment.yaml +++ /dev/null @@ -1,254 +0,0 @@ -interactions: -- request: - body: null - headers: - Accept: - - '*/*' - Accept-Encoding: - - gzip, deflate - Connection: - - keep-alive - User-Agent: - - python-requests/2.32.5 - method: GET - uri: https://staging-api.braintrust.dev/version - response: - body: - string: '{"version":"1.1.31","date_version":"20260303","ff_version":21,"commit":"ef190e7cc21d4a7447c7d5714d94519d3e11abe6","deployment_mode":"lambda","deployment_type":"custom","brainstore_default":"force","brainstore_can_contain_row_refs":true,"skip_pg_config":"all","has_realtime_wal_bucket":true,"has_logs2":true,"js":true,"universal":true,"code_execution":true,"logs3_payload_max_bytes":5242880,"control_plane_telemetry":["status","metrics","logs","traces","memprof","usage"]}' - headers: - Connection: - - keep-alive - Content-Length: - - '471' - Content-Type: - - application/json; charset=utf-8 - Date: - - Wed, 04 Mar 2026 18:53:38 GMT - Via: - - 1.1 24365d50ec90c9fb2b814e9d6c2f8b8c.cloudfront.net (CloudFront), 1.1 df34ce5bf73c140dc63a22fa17a4dcda.cloudfront.net - (CloudFront) - X-Amz-Cf-Id: - - iZSRCMFDldoyMgPWaFWbIytLGFZDbXXVq3vT1kx3mD2s0z5JPeYL5g== - X-Amz-Cf-Pop: - - YTO53-P2 - - YTO50-P1 - X-Amzn-Trace-Id: - - Root=1-69a87fb2-0563d3ca08011a600e7c0c30;Parent=29a9a49d07eb3aa6;Sampled=0;Lineage=1:fc3b4ff1:0 - X-Cache: - - Miss from cloudfront - access-control-allow-credentials: - - 'true' - access-control-expose-headers: - - x-bt-cursor,x-bt-found-existing,x-bt-query-plan,x-bt-api-duration-ms,x-bt-brainstore-duration-ms - etag: - - W/"1d7-AZtEQwxyAKme2P8gABnHCeVZRA4" - vary: - - Origin - x-amz-apigw-id: - - Ztjj5Hn9oAMES6Q= - x-amzn-Remapped-content-length: - - '471' - x-amzn-RequestId: - - 2e376c2d-4dc9-4eb3-ae5c-befa424bc3f8 - x-bt-internal-trace-id: - - 69a87fb2000000004178719e0e9e5aa1 - status: - code: 200 - message: OK -- request: - body: '{"contents": [{"parts": [{"text": "Use the tool with query: test"}], "role": - "user"}], "systemInstruction": {"parts": [{"text": "You are a helpful assistant - with tools.\n\nYou are an agent. Your internal name is \"tool_agent\"."}], "role": - "user"}, "tools": [{"functionDeclarations": [{"description": "A simple tool.", - "name": "simple_tool", "parameters": {"properties": {"query": {"type": "STRING"}}, - "required": ["query"], "type": "OBJECT"}}]}], "generationConfig": {}}' - headers: - Accept: - - '*/*' - Accept-Encoding: - - gzip, deflate - Connection: - - keep-alive - Content-Length: - - '471' - Content-Type: - - application/json - Host: - - generativelanguage.googleapis.com - user-agent: - - google-genai-sdk/1.65.0 gl-python/3.13.3 google-adk/1.26.0 gl-python/3.13.3 - x-goog-api-client: - - google-genai-sdk/1.65.0 gl-python/3.13.3 google-adk/1.26.0 gl-python/3.13.3 - method: POST - uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent - response: - body: - string: !!binary | - H4sIAAAAAAAC/61S0U6DMBR971eQPo9lONimb2b6YDJ1ccTMGGOu4451lhbbYmKW/bulwAZT3+SB - lHvuPef0HnbE8+gKRMISMKjphfdsK563c+8Sk8KgMBZoSraYgzLH3urZtc62ZV2IlWFSTIHzznCN - C8jQ1qlmWc7x1UjJae+0CVSqfxm2yEeB6quct6YNPcH35K+v4/nlKEaV5M5LJhPkDdm+aaBrJpje - PCBoKcq2RXw/P3il8JnOZJor+VZa9c/74+E4mIwHoyiYhEEwCtEfRKQRd7K00JDiLRqwS4fDBakl - yXITy3cUU1m4pYdBJdTKqINHNWykAd6dHPV+sOorq8l4O7pWqvb6wJlxa42vl3ErD8vfMdXsiLRW - eWrxn8Sirhapk6nCekSlWZVKipnNyT/rD/w1B71xhFShzqXQeJO4n23pS4Dt9k5th/7kac7MbKEv - Q0r25BtzH/bJCQMAAA== - headers: - Alt-Svc: - - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 - Content-Encoding: - - gzip - Content-Type: - - application/json; charset=UTF-8 - Date: - - Wed, 04 Mar 2026 18:53:38 GMT - Server: - - scaffolding on HTTPServer2 - Server-Timing: - - gfet4t7; dur=533 - Transfer-Encoding: - - chunked - Vary: - - Origin - - X-Origin - - Referer - X-Content-Type-Options: - - nosniff - X-Frame-Options: - - SAMEORIGIN - X-XSS-Protection: - - '0' - status: - code: 200 - message: OK -- request: - body: '{"rows": [{"_is_merge": false, "context": {"caller_filename": "/Users/abhijeetprasad/.local/share/mise/installs/python/3.13.3/lib/python3.13/asyncio/events.py", - "caller_functionname": "_run", "caller_lineno": 89}, "created": "2026-03-04T18:53:37.537733+00:00", - "id": "67c0d81c-587d-434f-a173-8274bb5247ba", "log_id": "g", "metrics": {"start": - 1772650417.537733}, "project_id": "test-context", "root_span_id": "a4266ce9-8861-4723-ac6e-688ba177b395", - "span_attributes": {"exec_counter": 21, "name": "outer_span", "type": "task"}, - "span_id": "a4266ce9-8861-4723-ac6e-688ba177b395", "span_parents": null}], "api_version": - 2}' - headers: - Accept: - - '*/*' - Accept-Encoding: - - gzip, deflate - Connection: - - keep-alive - Content-Length: - - '620' - User-Agent: - - python-requests/2.32.5 - method: POST - uri: https://staging-api.braintrust.dev/logs3 - response: - body: - string: '{"Code":"ForbiddenError","Message":"Missing read access to project_log - id test-context, or the project_log does not exist [user_email=abhijeet@braintrustdata.com] - [user_org=braintrustdata.com] [timestamp=1772650418.546]","InternalTraceId":"69a87fb2000000000cf539bdeabb9690","Path":"/logs3","Service":"api"}' - headers: - Connection: - - keep-alive - Content-Length: - - '237' - Content-Type: - - application/json; charset=utf-8 - Date: - - Wed, 04 Mar 2026 18:53:38 GMT - Via: - - 1.1 5785adb181f17ab60069f54a29dd7b7a.cloudfront.net (CloudFront), 1.1 6477e7b623b71ec66bc28ed8e271db7e.cloudfront.net - (CloudFront) - X-Amz-Cf-Id: - - XAH7plInwceQZyVAeRsFyCbHK3wPol5Q3FzOnLOSOLcXlWAEIJxpQw== - X-Amz-Cf-Pop: - - YTO53-P2 - - YTO50-P1 - X-Amzn-Trace-Id: - - Root=1-69a87fb2-040ba02c60171b0a65f5688d;Parent=40c51c3702a2cb1f;Sampled=0;Lineage=1:fc3b4ff1:0 - X-Cache: - - Error from cloudfront - access-control-allow-credentials: - - 'true' - access-control-expose-headers: - - x-bt-cursor,x-bt-found-existing,x-bt-query-plan,x-bt-api-duration-ms,x-bt-brainstore-duration-ms - content-encoding: - - gzip - etag: - - W/"132-zIQvupRcq5HDv40+Qcx4u9krY4Q" - vary: - - Origin, Accept-Encoding - x-amz-apigw-id: - - Ztjj7G-roAMEKIw= - x-amzn-RequestId: - - 4586197f-a71c-4ba7-89e8-21ad190ee055 - x-bt-internal-trace-id: - - 69a87fb2000000000cf539bdeabb9690 - status: - code: 403 - message: Forbidden -- request: - body: '{"contents": [{"parts": [{"text": "Use the tool with query: test"}], "role": - "user"}, {"parts": [{"functionCall": {"args": {"query": "test"}, "name": "simple_tool"}}], - "role": "model"}, {"parts": [{"functionResponse": {"name": "simple_tool", "response": - {"result": "Processed: test"}}}], "role": "user"}], "systemInstruction": {"parts": - [{"text": "You are a helpful assistant with tools.\n\nYou are an agent. Your - internal name is \"tool_agent\"."}], "role": "user"}, "tools": [{"functionDeclarations": - [{"description": "A simple tool.", "name": "simple_tool", "parameters": {"properties": - {"query": {"type": "STRING"}}, "required": ["query"], "type": "OBJECT"}}]}], - "generationConfig": {}}' - headers: - Accept: - - '*/*' - Accept-Encoding: - - gzip, deflate - Connection: - - keep-alive - Content-Length: - - '690' - Content-Type: - - application/json - Host: - - generativelanguage.googleapis.com - user-agent: - - google-genai-sdk/1.65.0 gl-python/3.13.3 google-adk/1.26.0 gl-python/3.13.3 - x-goog-api-client: - - google-genai-sdk/1.65.0 gl-python/3.13.3 google-adk/1.26.0 gl-python/3.13.3 - method: POST - uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent - response: - body: - string: !!binary | - H4sIAAAAAAAC/61RTU+DQBC98yvInkvDxzbSXtVDE42oxNhUD9vulBJhl+wupkr47+5CoYte5UA2 - 897Mm/emcVwX7QmjOSUKJFq5W11x3ab7G4wzBUxpYCjpYkWEunD7r7HemqLgZJpQIvgepAS6cvV8 - 9caQxWvH9/vsMl3wAkxrySkUA70dCOiQs1wen4BIzgztOX1I0IiSz+yOZ5XgO7Og5899jBdBHOMo - XEbB0hlkO0FUS5LBPSii3ZPRI9LtZaVS/gHsmted+0XUS1hhTXB8hhVXpJh2Xs3+TJU3WjMv7Ayt - eLVxUuTqy7hLb19TZIWjpksN6ThWiL9X/CcxPNVyzjfpz/QCQub9PTIo9YW8cO57h4LIYzcQCZAV - ZxLW1HAk8zjZfQfrDcNevElCGp3qR4yc1vkBvZPsN5ICAAA= - headers: - Alt-Svc: - - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 - Content-Encoding: - - gzip - Content-Type: - - application/json; charset=UTF-8 - Date: - - Wed, 04 Mar 2026 18:53:38 GMT - Server: - - scaffolding on HTTPServer2 - Server-Timing: - - gfet4t7; dur=510 - Transfer-Encoding: - - chunked - Vary: - - Origin - - X-Origin - - Referer - X-Content-Type-Options: - - nosniff - X-Frame-Options: - - SAMEORIGIN - X-XSS-Protection: - - '0' - status: - code: 200 - message: OK -version: 1 diff --git a/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_agent_metadata_with_attachment.yaml b/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_agent_metadata_with_attachment.yaml deleted file mode 100644 index 4ab36f7b..00000000 --- a/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_agent_metadata_with_attachment.yaml +++ /dev/null @@ -1,142 +0,0 @@ -interactions: -- request: - body: '{"contents": [{"parts": [{"text": "Use the tool with query: test"}], "role": - "user"}], "systemInstruction": {"parts": [{"text": "You are a helpful assistant - with tools.\n\nYou are an agent. Your internal name is \"tool_agent\"."}], "role": - "user"}, "tools": [{"functionDeclarations": [{"description": "A simple tool.", - "name": "simple_tool", "parameters_json_schema": {"properties": {"query": {"title": - "Query", "type": "string"}}, "required": ["query"], "title": "simple_toolParams", - "type": "object"}}]}], "generationConfig": {}}' - headers: - accept: - - '*/*' - accept-encoding: - - gzip, deflate - connection: - - keep-alive - content-length: - - '531' - content-type: - - application/json - host: - - generativelanguage.googleapis.com - user-agent: - - google-genai-sdk/2.9.0 gl-python/3.14.3 google-adk/2.3.0 gl-python/3.14.3 - x-goog-api-client: - - google-genai-sdk/2.9.0 gl-python/3.14.3 google-adk/2.3.0 gl-python/3.14.3 - method: POST - uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent - response: - body: - string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": - [\n {\n \"functionCall\": {\n \"name\": \"simple_tool\",\n - \ \"args\": {\n \"query\": \"test\"\n }\n - \ }\n }\n ],\n \"role\": \"model\"\n },\n - \ \"finishReason\": \"STOP\",\n \"index\": 0,\n \"finishMessage\": - \"Model generated function call(s).\"\n }\n ],\n \"usageMetadata\": {\n - \ \"promptTokenCount\": 68,\n \"candidatesTokenCount\": 15,\n \"totalTokenCount\": - 83,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n - \ \"tokenCount\": 68\n }\n ],\n \"serviceTier\": \"standard\"\n - \ },\n \"modelVersion\": \"gemini-2.5-flash-lite\",\n \"responseId\": \"SnQ5auuVE5nQ_uMPxZSWqAU\"\n}\n" - headers: - alt-svc: - - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 - content-length: - - '738' - content-type: - - application/json; charset=UTF-8 - date: - - Mon, 22 Jun 2026 17:43:40 GMT - server: - - scaffolding on HTTPServer2 - server-timing: - - gfet4t7; dur=2673 - transfer-encoding: - - chunked - vary: - - Origin - - X-Origin - - Referer - x-content-type-options: - - nosniff - x-frame-options: - - SAMEORIGIN - x-gemini-service-tier: - - standard - x-xss-protection: - - '0' - status: - code: 200 - message: OK -- request: - body: '{"contents": [{"parts": [{"text": "Use the tool with query: test"}], "role": - "user"}, {"parts": [{"functionCall": {"args": {"query": "test"}, "name": "simple_tool"}}], - "role": "model"}, {"parts": [{"functionResponse": {"name": "simple_tool", "response": - {"result": "Processed: test"}}}], "role": "user"}], "systemInstruction": {"parts": - [{"text": "You are a helpful assistant with tools.\n\nYou are an agent. Your - internal name is \"tool_agent\"."}], "role": "user"}, "tools": [{"functionDeclarations": - [{"description": "A simple tool.", "name": "simple_tool", "parameters_json_schema": - {"properties": {"query": {"title": "Query", "type": "string"}}, "required": - ["query"], "title": "simple_toolParams", "type": "object"}}]}], "generationConfig": - {}}' - headers: - accept: - - '*/*' - accept-encoding: - - gzip, deflate - connection: - - keep-alive - content-length: - - '750' - content-type: - - application/json - host: - - generativelanguage.googleapis.com - user-agent: - - google-genai-sdk/2.9.0 gl-python/3.14.3 google-adk/2.3.0 gl-python/3.14.3 - x-goog-api-client: - - google-genai-sdk/2.9.0 gl-python/3.14.3 google-adk/2.3.0 gl-python/3.14.3 - method: POST - uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent - response: - body: - string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": - [\n {\n \"text\": \"The tool was able to process the query - \\\"test\\\" and the response was \\\"Processed: test\\\".\"\n }\n - \ ],\n \"role\": \"model\"\n },\n \"finishReason\": - \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": - 100,\n \"candidatesTokenCount\": 20,\n \"totalTokenCount\": 120,\n \"promptTokensDetails\": - [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 100\n - \ }\n ],\n \"serviceTier\": \"standard\"\n },\n \"modelVersion\": - \"gemini-2.5-flash-lite\",\n \"responseId\": \"TXQ5atCBBtiR-8YP0YW90Aw\"\n}\n" - headers: - alt-svc: - - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 - content-length: - - '640' - content-type: - - application/json; charset=UTF-8 - date: - - Mon, 22 Jun 2026 17:43:41 GMT - server: - - scaffolding on HTTPServer2 - server-timing: - - gfet4t7; dur=519 - transfer-encoding: - - chunked - vary: - - Origin - - X-Origin - - Referer - x-content-type-options: - - nosniff - x-frame-options: - - SAMEORIGIN - x-gemini-service-tier: - - standard - x-xss-protection: - - '0' - status: - code: 200 - message: OK -version: 1 diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index 236645d6..ff71f776 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -4,7 +4,6 @@ import pytest from braintrust import logger -from braintrust.bt_json import bt_safe_deep_copy from braintrust.integrations.adk import setup_adk from braintrust.integrations.adk.tracing import _create_thread_wrapper from braintrust.logger import Attachment @@ -85,64 +84,6 @@ def _extract_text_parts(contents): return texts -def test_adk_thread_context_propagation(memory_logger): - """Runner.run should preserve Braintrust context across its thread bridge.""" - import asyncio - - from braintrust import current_span, start_span - from google.adk.agents import LlmAgent - from google.adk.models.base_llm import BaseLlm - from google.adk.models.llm_request import LlmRequest - from google.adk.models.llm_response import LlmResponse - from google.adk.models.registry import LLMRegistry - from google.adk.runners import Runner - from google.adk.sessions import InMemorySessionService - from google.genai import types - - assert not memory_logger.pop() - - parent_seen = [] - - class TestLlm(BaseLlm): - @classmethod - def supported_models(cls) -> list[str]: - return [r"test-llm-context-prop"] - - async def generate_content_async(self, llm_request: LlmRequest, stream: bool = False): - parent_seen.append(current_span()) - yield LlmResponse(content=types.Content(role="model", parts=[types.Part(text="ok")])) - - LLMRegistry.register(TestLlm) - - agent = LlmAgent( - name="echo_agent", - model="test-llm-context-prop", - instruction="Respond with ok.", - ) - session_service = InMemorySessionService() - app_name = "thread_bridge_app" - user_id = "test-user" - session_id = "test-session-thread" - asyncio.run( - session_service.create_session( - app_name=app_name, - user_id=user_id, - session_id=session_id, - ) - ) - runner = Runner(agent=agent, app_name=app_name, session_service=session_service) - user_msg = types.Content(role="user", parts=[types.Part(text="hello")]) - - with start_span(name="adk_thread_parent") as parent_span: - events = list(runner.run(user_id=user_id, session_id=session_id, new_message=user_msg)) - - assert events - assert parent_seen - thread_root = getattr(parent_seen[0], "root_span_id", None) - assert thread_root is not None - assert thread_root == parent_span.root_span_id - - def test_create_thread_wrapper_exception_does_not_double_invoke_target(): """Regression test: target exceptions must not cause a second invocation.""" call_count = 0 @@ -220,141 +161,13 @@ async def run_message(text: str) -> str: assert "alice" in second_response_text.lower() -@pytest.mark.asyncio -async def test_adk_generation_config_is_logged(memory_logger): - """Sampling and stop-sequence config should be captured in the LLM span input.""" - from google.adk.models.base_llm import BaseLlm - from google.adk.models.llm_request import LlmRequest - from google.adk.models.llm_response import LlmResponse - from google.adk.models.registry import LLMRegistry - - assert not memory_logger.pop() - - class ConfigCaptureLlm(BaseLlm): - @classmethod - def supported_models(cls) -> list[str]: - return [r"test-llm-config-capture"] - - async def generate_content_async(self, llm_request: LlmRequest, stream: bool = False): - yield LlmResponse(content=types.Content(role="model", parts=[types.Part(text="configured")])) - - LLMRegistry.register(ConfigCaptureLlm) - - app_name = "config_app" - user_id = "test-user" - session_id = "test-session-config" - agent = LlmAgent( - name="config_agent", - model="test-llm-config-capture", - instruction="Reply with the word configured.", - generate_content_config=types.GenerateContentConfig( - max_output_tokens=23, - temperature=0.7, - top_p=0.9, - stop_sequences=["END", "\n\n"], - ), - ) - runner = await _create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id) - - user_msg = types.Content(role="user", parts=[types.Part(text="Please answer.")]) - responses = [] - async for event in runner.run_async(user_id=user_id, session_id=session_id, new_message=user_msg): - if event.is_final_response(): - responses.append(event) - - assert responses - - spans = memory_logger.pop() - llm_spans = [row for row in spans if row["span_attributes"]["type"] == "llm"] - assert llm_spans - - config = llm_spans[0]["input"]["config"] - assert config["max_output_tokens"] == 23 - assert config["temperature"] == 0.7 - assert config["top_p"] == 0.9 - assert config["stop_sequences"] == ["END", "\n\n"] - - -@pytest.mark.asyncio -async def test_adk_document_inline_data_attachment_conversion(memory_logger): - """Document bytes should be logged as attachment references, not raw payloads.""" - from google.adk.models.base_llm import BaseLlm - from google.adk.models.llm_request import LlmRequest - from google.adk.models.llm_response import LlmResponse - from google.adk.models.registry import LLMRegistry - - assert not memory_logger.pop() - - class DocumentCaptureLlm(BaseLlm): - @classmethod - def supported_models(cls) -> list[str]: - return [r"test-llm-document-capture"] - - async def generate_content_async(self, llm_request: LlmRequest, stream: bool = False): - yield LlmResponse(content=types.Content(role="model", parts=[types.Part(text="document received")])) - - LLMRegistry.register(DocumentCaptureLlm) - - app_name = "document_app" - user_id = "test-user" - session_id = "test-session-document" - agent = LlmAgent( - name="document_agent", - model="test-llm-document-capture", - instruction="Acknowledge the uploaded document.", - ) - runner = await _create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id) - - pdf_path = FIXTURES_DIR / "test-document.pdf" - with open(pdf_path, "rb") as f: - pdf_data = f.read() - - user_msg = types.Content( - role="user", - parts=[ - types.Part(inline_data=types.Blob(mime_type="application/pdf", data=pdf_data)), - types.Part(text="Summarize this document."), - ], - ) - - responses = [] - async for event in runner.run_async(user_id=user_id, session_id=session_id, new_message=user_msg): - if event.is_final_response(): - responses.append(event) - - assert responses - - spans = memory_logger.pop() - invocation_span = next(row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]") - new_message = invocation_span["input"]["new_message"] - assert len(new_message["parts"]) == 2 - - document_part = new_message["parts"][0] - assert "file" in document_part - assert document_part["file"]["filename"] == "file.pdf" - attachment = document_part["file"]["file_data"] - assert isinstance(attachment, Attachment) - assert attachment.reference["content_type"] == "application/pdf" - assert attachment.reference["filename"] == "file.pdf" - - text_part = new_message["parts"][1] - assert text_part == {"text": "Summarize this document."} - - logged_payload = str(invocation_span).lower() - assert pdf_data[:8].hex() not in logged_payload - - llm_span = next(row for row in spans if row["span_attributes"]["type"] == "llm") - llm_contents = llm_span["input"]["contents"] - llm_document_part = llm_contents[0]["parts"][0] - assert isinstance(llm_document_part["file"]["file_data"], Attachment) - assert llm_document_part["file"]["file_data"].reference["content_type"] == "application/pdf" - - @pytest.mark.vcr def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger): - """Runner.run() should emit a single invocation span even though it delegates to run_async().""" + """Runner.run() emits one invocation span AND preserves Braintrust context through + ADK's thread bridge (Runner.run dispatches to a background thread).""" import asyncio + from braintrust import start_span from braintrust.util import LazyValue assert not memory_logger.pop() @@ -380,11 +193,12 @@ def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger): original_global_bg_logger = logger._state._global_bg_logger logger._state._global_bg_logger = LazyValue(lambda: memory_logger, use_mutex=False) try: - responses = [ - event - for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg) - if event.is_final_response() - ] + with start_span(name="adk_thread_parent") as parent_span: + responses = [ + event + for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg) + if event.is_final_response() + ] finally: logger._state._global_bg_logger = original_global_bg_logger @@ -405,6 +219,20 @@ def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger): f"got parents {agent_spans[0].get('span_parents')}" ) + # Thread-bridge context propagation: every ADK span emitted on the worker + # thread should share the outer parent's root_span_id. + adk_spans = [ + row + for row in spans + if row["context"]["span_origin"]["instrumentation"]["name"] == "adk-auto" + ] + assert adk_spans + for row in adk_spans: + assert row["root_span_id"] == parent_span.root_span_id, ( + f"{row['span_attributes']['name']} lost thread context: " + f"{row['root_span_id']} != {parent_span.root_span_id}" + ) + @pytest.mark.vcr @pytest.mark.asyncio @@ -621,6 +449,11 @@ async def test_adk_max_tokens_captures_content(memory_logger): llm_span = llm_spans[0] assert "output" in llm_span, "Missing output in LLM span" + # Sampling config from generate_content_config is captured in input.config + config = llm_span["input"]["config"] + assert config["max_output_tokens"] == 50 + assert config["temperature"] == 0.7 + output = llm_span["output"] # When MAX_TOKENS is hit, we should still have content captured @@ -640,121 +473,6 @@ async def test_adk_max_tokens_captures_content(memory_logger): assert "usage_metadata" in output, "Should have usage metadata" -def test_serialize_content_with_binary_data(): - """Test that _serialize_content converts binary data to Attachment references.""" - from braintrust.integrations.adk.tracing import _serialize_content, _serialize_part - from braintrust.logger import Attachment - - # Create a minimal PNG image (1x1 red pixel) - minimal_png = ( - b"\x89PNG\r\n\x1a\n" # PNG signature - b"\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" - b"\x08\x02\x00\x00\x00\x90wS\xde" # IHDR - b"\x00\x00\x00\x0cIDATx\x9cc\xf8\xcf\xc0\x00\x00\x00\x03\x00\x01\x00\x18\xdd\x8d\xb4" # IDAT - b"\x00\x00\x00\x00IEND\xaeB`\x82" # IEND - ) - - # Create a mock Part with inline_data - class MockBlob: - def __init__(self, data, mime_type): - self.data = data - self.mime_type = mime_type - - class MockPart: - def __init__(self, inline_data=None, text=None): - self.inline_data = inline_data - self.text = text - - # Test serializing a Part with binary data - part_with_image = MockPart(inline_data=MockBlob(minimal_png, "image/png")) - serialized_part = _serialize_part(part_with_image) - - # Verify structure - assert "image_url" in serialized_part, "Should have image_url field" - assert "url" in serialized_part["image_url"], "Should have url field" - - attachment = serialized_part["image_url"]["url"] - # The Attachment object should be in the serialized output - assert isinstance(attachment, Attachment), "Should be an Attachment object" - assert attachment.reference["type"] == "braintrust_attachment" - assert attachment.reference["content_type"] == "image/png" - assert attachment.reference["filename"] == "image.png" - assert "key" in attachment.reference - - # Test serializing a Part with text - part_with_text = MockPart(text="Hello, world!") - serialized_text_part = _serialize_part(part_with_text) - assert serialized_text_part == {"text": "Hello, world!"}, "Text part should serialize correctly" - - # Test serializing Content with multiple parts - class MockContent: - def __init__(self, parts, role): - self.parts = parts - self.role = role - - content = MockContent( - parts=[ - MockPart(inline_data=MockBlob(minimal_png, "image/png")), - MockPart(text="What's in this image?"), - ], - role="user", - ) - - serialized_content = _serialize_content(content) - assert "parts" in serialized_content - assert "role" in serialized_content - assert serialized_content["role"] == "user" - assert len(serialized_content["parts"]) == 2 - - # First part should be the image as Attachment - assert "image_url" in serialized_content["parts"][0] - assert isinstance(serialized_content["parts"][0]["image_url"]["url"], Attachment) - - # Second part should be text - assert serialized_content["parts"][1] == {"text": "What's in this image?"} - - -def test_serialize_part_with_file_data(): - """Test that _serialize_part handles file_data (file references) correctly.""" - from braintrust.integrations.adk.tracing import _serialize_part - - class MockFileData: - def __init__(self, file_uri, mime_type): - self.file_uri = file_uri - self.mime_type = mime_type - - class MockPart: - def __init__(self, file_data=None, text=None): - self.file_data = file_data - self.text = text - - # Test serializing a Part with file_data - part_with_file = MockPart(file_data=MockFileData("gs://bucket/file.pdf", "application/pdf")) - serialized_part = _serialize_part(part_with_file) - - assert "file_data" in serialized_part - assert serialized_part["file_data"]["file_uri"] == "gs://bucket/file.pdf" - assert serialized_part["file_data"]["mime_type"] == "application/pdf" - - -def test_serialize_part_with_dict(): - """Test that _serialize_part handles dict input correctly.""" - from braintrust.integrations.adk.tracing import _serialize_part - - # Test that dicts pass through unchanged - dict_part = {"text": "Hello", "custom": "field"} - serialized = _serialize_part(dict_part) - assert serialized == dict_part, "Dict should pass through unchanged" - - -def test_serialize_content_with_none(): - """Test that _serialize_content handles None correctly.""" - from braintrust.integrations.adk.tracing import _serialize_content - - result = _serialize_content(None) - assert result is None, "None should serialize to None" - - @pytest.mark.vcr @pytest.mark.asyncio async def test_adk_binary_data_attachment_conversion(memory_logger): @@ -933,150 +651,6 @@ async def test_adk_captures_metrics(memory_logger): # path is asserted in `test_adk_captures_metrics`. -@pytest.mark.asyncio -async def test_llm_call_span_wraps_child_spans(memory_logger): - """Test that llm_call span is created BEFORE yielding events, so child spans have proper parent. - - This test validates the fix for the issue where mcp_tool and other child spans - were losing their parent context because the llm_call span was created AFTER - all events were yielded. - - The fix ensures: - 1. llm_call span is created BEFORE wrapped() is called - 2. Child spans (like mcp_tool) created during execution have proper parent - 3. Span is updated with correct call_type after response is received - """ - from unittest.mock import MagicMock - - from braintrust import current_span, start_span - from braintrust.integrations.adk import wrap_flow - - # Clear any existing logs - memory_logger.pop() - - # Mock Flow class - class MockFlow: - def __init__(self): - self.llm = MagicMock() - self.llm.model = "test-model" - - async def run_async(self, invocation_context, llm_request=None, model_response_event=None): - """Method that wrap_flow will wrap.""" - async for event in self._call_llm_async(invocation_context, llm_request, model_response_event): - yield event - - async def _call_llm_async(self, invocation_context, llm_request, model_response_event): - """Simulates the flow making LLM calls and potentially calling tools.""" - # Simulate an event stream - yield {"type": "start"} - - # During execution, child spans might be created (like mcp_tool calls) - # This simulates an MCP tool being called during LLM execution - with start_span(name="mcp_tool [test_tool]", type="tool") as tool_span: - tool_span.log(output={"result": "success"}) - - yield {"type": "complete", "content": {"parts": [{"text": "Done"}], "role": "model"}} - - # Wrap the flow - wrap_flow(MockFlow) - - # Create flow instance - flow = MockFlow() - - # Track parent span during execution - parent_spans_during_execution = [] - - async def wrapped_execution(): - """Wrapper that tracks parent span during execution.""" - async for event in flow.run_async( - invocation_context={"test": "context"}, - llm_request={"contents": [{"parts": [{"text": "test"}], "role": "user"}]}, - model_response_event=None, - ): - # Check what the current parent span is during execution - parent = current_span() - if parent and hasattr(parent, "id"): - parent_spans_during_execution.append(parent.id) - - # Execute - await wrapped_execution() - - # Give background logger time to flush - memory_logger.flush() - - # Get all logged spans - logs = memory_logger.pop() - - # Find the spans by name - llm_call_spans = [log for log in logs if "llm_call" in log.get("span_attributes", {}).get("name", "")] - mcp_tool_spans = [log for log in logs if "mcp_tool" in log.get("span_attributes", {}).get("name", "")] - - # Verify llm_call span exists - assert len(llm_call_spans) > 0, "Should have created llm_call span" - - # Verify mcp_tool span exists - assert len(mcp_tool_spans) > 0, "Should have created mcp_tool span" - - # Verify mcp_tool span has the llm_call span as parent - llm_call_span_id = llm_call_spans[0]["span_id"] - mcp_tool_span = mcp_tool_spans[0] - - # The mcp_tool span should have the llm_call span in its parent chain - assert "span_parents" in mcp_tool_span, "mcp_tool span should have span_parents" - assert llm_call_span_id in mcp_tool_span["span_parents"], ( - f"mcp_tool span should have llm_call span as parent. " - f"Expected {llm_call_span_id} in {mcp_tool_span['span_parents']}" - ) - - # Verify llm_call span name was updated with call_type - llm_call_name = llm_call_spans[0]["span_attributes"]["name"] - assert "[" in llm_call_name, f"llm_call span name should include call_type in brackets: {llm_call_name}" - - -@pytest.mark.asyncio -async def test_async_context_preservation_across_yields(): - """Test that async context is preserved across generator yields. - - This validates that the aclosing wrapper properly handles ContextVar errors - that occur when async generators yield control and resume in different contexts. - """ - import asyncio - from contextlib import aclosing - - from braintrust import start_span - - # Initialize logger - init_test_logger("test-context") - - async def context_switching_generator(): - """Generator that creates spans and yields, potentially switching contexts.""" - with start_span(name="outer_span", type="task") as outer: - yield {"event": 1} - await asyncio.sleep(0.001) # Force context switch - - with start_span(name="inner_span", type="task") as inner: - inner.log(output={"data": "test"}) - yield {"event": 2} - await asyncio.sleep(0.001) # Another context switch - - yield {"event": 3} - - # Collect events using aclosing - events = [] - async with aclosing(context_switching_generator()) as gen: - async for event in gen: - events.append(event) - await asyncio.sleep(0.001) # Force context switches during iteration - - # Verify all events were collected successfully - assert len(events) == 3 - assert events[0]["event"] == 1 - assert events[1]["event"] == 2 - assert events[2]["event"] == 3 - - # If we get here, the context error suppression in aclosing.__aexit__ worked correctly - - class CapitalOutput(BaseModel): capital: str = Field(description="The capital of the country.") @@ -1364,76 +938,6 @@ class Person(BaseModel): # params (`test_adk_generation_config_is_logged`). -@pytest.mark.asyncio -async def test_serialize_pydantic_schema_direct(): - """Test _serialize_pydantic_schema directly with various inputs.""" - from braintrust.integrations.adk.tracing import _serialize_pydantic_schema - - class SimpleSchema(BaseModel): - name: str = Field(description="A name") - count: int = Field(description="A count", ge=0) - - # Test with Pydantic class - result = _serialize_pydantic_schema(SimpleSchema) - assert isinstance(result, dict) - assert result["type"] == "object" - assert "properties" in result - assert "name" in result["properties"] - assert result["properties"]["name"]["description"] == "A name" - assert "count" in result["properties"] - - # Test with non-Pydantic class - class NotPydantic: - pass - - result = _serialize_pydantic_schema(NotPydantic) - assert isinstance(result, dict) - assert "__class__" in result - assert result["__class__"] == "NotPydantic" - - # Test with non-class object - result = _serialize_pydantic_schema("not a class") - assert isinstance(result, dict) - assert "__class__" in result - - -@pytest.mark.asyncio -async def test_bt_safe_deep_copy_never_raises(): - """Test that bt_safe_deep_copy never raises exceptions.""" - from braintrust.bt_json import bt_safe_deep_copy - - class BrokenModel: - def model_dump(self): - raise ValueError("I'm broken!") - - # Should not raise - result = bt_safe_deep_copy(BrokenModel()) - assert result is not None - - # Test with various types - assert bt_safe_deep_copy({"key": "value"}) == {"key": "value"} - assert bt_safe_deep_copy([1, 2, 3]) == [1, 2, 3] - assert bt_safe_deep_copy("string") == "string" - assert bt_safe_deep_copy(123) == 123 - assert bt_safe_deep_copy(None) is None - - # Test with Pydantic model instance - class WorkingModel(BaseModel): - value: str = "test" - - instance = WorkingModel() - result = bt_safe_deep_copy(instance) - assert isinstance(result, dict) - assert result["value"] == "test" - - # Test with Pydantic model class (not instance) - # bt_safe_deep_copy now returns the JSON schema for Pydantic model classes - result = bt_safe_deep_copy(WorkingModel) - assert isinstance(result, dict) - assert "properties" in result - assert "value" in result["properties"] - - @pytest.mark.vcr @pytest.mark.asyncio async def test_adk_response_json_schema_dict(memory_logger): @@ -1555,128 +1059,6 @@ async def test_adk_response_json_schema_dict(memory_logger): assert output["avg_logprobs"] is not None -@pytest.mark.asyncio -async def test_capture_config_preserves_none(): - """Test that _capture_config returns None when config is None (not empty dict).""" - from braintrust.integrations.adk.tracing import _capture_config - - # None should be preserved as None, not converted to {} - result = _capture_config(None) - assert result is None, f"Expected None, got {result}" - - # Empty dict should remain empty dict - result = _capture_config({}) - assert result == {} - - # False should be preserved as False - result = _capture_config(False) - assert result is False - - # 0 should be preserved as 0 - result = _capture_config(0) - assert result == 0 - - # Empty string should be preserved - result = _capture_config("") - assert result == "" - - -@pytest.mark.asyncio -async def test_bt_safe_deep_copy_with_attachments(memory_logger): - """Test that bt_safe_deep_copy preserves Attachment objects in ADK context.""" - from braintrust.bt_json import bt_safe_deep_copy - - attachment = Attachment(data=b"test data", filename="test.txt", content_type="text/plain") - - # Test preserving attachment in nested structure (simulating ADK metadata) - metadata = {"file": attachment, "nested": {"also_file": attachment}} - - result = bt_safe_deep_copy(metadata) - - # Attachment identity should be preserved - assert result["file"] is attachment - assert result["nested"]["also_file"] is attachment - - -@pytest.mark.vcr -@pytest.mark.asyncio -async def test_adk_agent_metadata_with_attachment(memory_logger): - """Test that attachments in ADK agent metadata are preserved and uploaded.""" - from unittest.mock import patch - - assert not memory_logger.pop() - - attachment = Attachment(data=b"context data", filename="context.txt", content_type="text/plain") - - def simple_tool(query: str): - """A simple tool.""" - return {"result": f"Processed: {query}"} - - agent = Agent( - name="tool_agent", - model=ADK_MODEL, - instruction="You are a helpful assistant with tools.", - tools=[simple_tool], - ) - - APP_NAME = "attachment_app" - USER_ID = "test-user" - SESSION_ID = "test-session-attachment" - - session_service = InMemorySessionService() - await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID) - - runner = Runner(agent=agent, app_name=APP_NAME, session_service=session_service) - - # Create message with attachment in metadata context - user_msg = types.Content(role="user", parts=[types.Part(text="Use the tool with query: test")]) - - with patch.object(Attachment, "upload", return_value={"upload_status": "done"}) as mock_upload: - responses = [] - # We can't directly inject attachment into ADK's internal flow, - # but we can test that if an attachment appears in logged metadata, - # bt_safe_deep_copy preserves it - async for event in runner.run_async(user_id=USER_ID, session_id=SESSION_ID, new_message=user_msg): - if event.is_final_response(): - responses.append(event) - - memory_logger.flush() - - spans = memory_logger.pop() - assert len(spans) > 0, "Should have logged spans" - - # Verify bt_safe_deep_copy behavior with attachment - test_data = {"metadata": {"context_file": attachment}} - copied = bt_safe_deep_copy(test_data) - assert copied["metadata"]["context_file"] is attachment - - -@pytest.mark.asyncio -async def test_adk_bytes_and_attachment_in_structure(): - """Test that dataclass/dict with both bytes and attachment fields are handled correctly.""" - from braintrust.bt_json import bt_safe_deep_copy - - attachment = Attachment(data=b"attachment data", filename="file.txt", content_type="text/plain") - - # Simulate ADK structure with both bytes and attachments - structure = { - "binary_data": b"some bytes", - "attachment": attachment, - "nested": {"more_bytes": bytearray(b"more data"), "another_attachment": attachment}, - } - - result = bt_safe_deep_copy(structure) - - # Attachment should be preserved - assert result["attachment"] is attachment - assert result["nested"]["another_attachment"] is attachment - - # Bytes should be handled (converted via bt_dumps/bt_loads) - assert "binary_data" in result - assert "nested" in result - assert "more_bytes" in result["nested"] - - class TestAutoInstrumentADK: """Tests for auto_instrument() with Google ADK.""" From 7c2c14c99ae8e56f3f589b1373c816053bc436c1 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 16:01:35 +0000 Subject: [PATCH 09/11] style(adk): apply ruff format Co-Authored-By: Claude Opus 4.7 --- .../braintrust/integrations/adk/test_adk.py | 20 ++++++------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index ff71f776..bf4f7709 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -221,11 +221,7 @@ def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger): # Thread-bridge context propagation: every ADK span emitted on the worker # thread should share the outer parent's root_span_id. - adk_spans = [ - row - for row in spans - if row["context"]["span_origin"]["instrumentation"]["name"] == "adk-auto" - ] + adk_spans = [row for row in spans if row["context"]["span_origin"]["instrumentation"]["name"] == "adk-auto"] assert adk_spans for row in adk_spans: assert row["root_span_id"] == parent_span.root_span_id, ( @@ -310,11 +306,7 @@ async def test_adk_braintrust_integration(memory_logger): assert function_call["name"] == "get_weather" assert function_call["args"]["location"] == "San Francisco" - adk_spans = [ - row - for row in spans - if row["context"]["span_origin"]["instrumentation"]["name"] == "adk-auto" - ] + adk_spans = [row for row in spans if row["context"]["span_origin"]["instrumentation"]["name"] == "adk-auto"] span_types_by_origin = {row["span_attributes"]["type"] for row in adk_spans} assert {"task", "llm", "tool"} <= span_types_by_origin, ( f"adk-auto origin missing on task/llm/tool spans: {span_types_by_origin}" @@ -322,13 +314,13 @@ async def test_adk_braintrust_integration(memory_logger): for span in llm_spans: meta = span["metadata"] - assert meta.get("provider") == "google", f"Missing metadata.provider=google on {span['span_attributes']['name']}" + assert meta.get("provider") == "google", ( + f"Missing metadata.provider=google on {span['span_attributes']['name']}" + ) assert meta.get("model"), "Missing metadata.model on llm span" assert meta.get("tools"), "metadata.tools should be non-empty for a tool-using agent" tool_names = [ - fn.get("name") - for tool_entry in meta["tools"] - for fn in (tool_entry.get("function_declarations") or []) + fn.get("name") for tool_entry in meta["tools"] for fn in (tool_entry.get("function_declarations") or []) ] assert "get_weather" in tool_names, f"get_weather missing from metadata.tools: {tool_names}" assert "tools" not in span["input"].get("config", {}), "tools should not be in input.config" From b9ff136def2cb22a83ba6709ab369650c79f97ac Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 16:10:06 +0000 Subject: [PATCH 10/11] ci: retrigger checks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pydantic_ai_wrap_openai(0.1.9) shard 3 failure looks unrelated to this PR — pydantic_ai 0.1.9 imports opentelemetry._events (added in opentelemetry-api 1.30.0) but the session resolved an older api. Same session passed on main at c0567629 (this branch's base) an hour before. Co-Authored-By: Claude Opus 4.7 From 6ace474ad6de87e0879b842a071f61fdfe87a91a Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 16:23:59 +0000 Subject: [PATCH 11/11] fix(ci): pin opentelemetry-api<1.40 for pydantic-ai 0.1.9 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pydantic-ai 0.1.9 unconditionally imports ``from opentelemetry._events import Event``. The ``_events`` module was removed from ``opentelemetry-api`` in 1.40+, so once opentelemetry-api 1.44.0 landed on PyPI the fresh resolution in the test_pydantic_ai_wrap_openai(0.1.9) session picked a version that fails at import time. Verified locally: opentelemetry-api 1.39.1 exposes _events, 1.44.0 does not. Constraint only applies to the 0.1.9 install; newer matrix versions (1.0.1, 2.9.0) remain unpinned. Unrelated to this PR's ADK work — same session passed on main at c0567629 (this branch's base) before the opentelemetry release. Co-Authored-By: Claude Opus 4.7 --- py/noxfile.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/py/noxfile.py b/py/noxfile.py index 22c8f64a..1eb9b8b8 100644 --- a/py/noxfile.py +++ b/py/noxfile.py @@ -442,6 +442,11 @@ def test_pydantic_ai_wrap_openai(session, version): """Test pydantic_ai with wrap_openai() approach - supports older versions.""" _install_test_deps(session) _install_matrix_dep(session, "pydantic-ai-wrap-openai", version) + # pydantic-ai 0.1.9 unconditionally imports ``from opentelemetry._events import Event``. + # The ``_events`` module was removed from ``opentelemetry-api`` in 1.40+, so a fresh + # resolution on newer opentelemetry releases picks a version that breaks import. + if version == "0.1.9": + session.install("opentelemetry-api<1.40", silent=SILENT_INSTALLS) _run_tests(session, f"{INTEGRATION_DIR}/pydantic_ai/test_pydantic_ai_wrap_openai.py", version=version)