diff --git a/py/src/braintrust/integrations/bedrock_runtime/cassettes/1.34.116/test_wrap_bedrock_converse_stream_parses_tool_input.yaml b/py/src/braintrust/integrations/bedrock_runtime/cassettes/1.34.116/test_wrap_bedrock_converse_stream_parses_tool_input.yaml new file mode 100644 index 000000000..ee6bbd481 --- /dev/null +++ b/py/src/braintrust/integrations/bedrock_runtime/cassettes/1.34.116/test_wrap_bedrock_converse_stream_parses_tool_input.yaml @@ -0,0 +1,61 @@ +interactions: +- request: + body: '{"messages": [{"role": "user", "content": [{"text": "What''s the weather + in Paris? Use the tool."}]}], "inferenceConfig": {"maxTokens": 200, "temperature": + 0}, "toolConfig": {"tools": [{"toolSpec": {"name": "get_weather", "description": + "Get the current weather for a city.", "inputSchema": {"json": {"type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, "required": + ["city"]}}}}], "toolChoice": {"tool": {"name": "get_weather"}}}}' + headers: + Content-Length: + - '466' + Content-Type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + User-Agent: + - !!binary | + Qm90bzMvMS40My4xMTEgbWQvQm90b2NvcmUjMS40My4xMTEgdWEvMi4xIG9zL2xpbnV4IzYuNi44 + Ny4yLW1pY3Jvc29mdC1zdGFuZGFyZC1XU0wyIG1kL2FyY2gjeDg2XzY0IGxhbmcvcHl0aG9uIzMu + MTQuOCBtZC9weWltcGwjQ1B5dGhvbiBtL0QsWixiLG4sMyBjZmcvcmV0cnktbW9kZSNsZWdhY3kg + Qm90b2NvcmUvMS40My4xMTE= + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.amazon.nova-lite-v1%3A0/converse-stream + response: + body: + string: !!binary | + AAAAsQAAAFLN4M4VCzpldmVudC10eXBlBwAMbWVzc2FnZVN0YXJ0DTpjb250ZW50LXR5cGUHABBh + cHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVldmVudHsicCI6ImFiY2RlZmdoaWprbG1u + b3BxcnN0dXZ3eHl6QUJDREVGR0hJSktMTU5PUFFSU1RVVldYWVoiLCJyb2xlIjoiYXNzaXN0YW50 + In3DJvYkAAAA5gAAAFc3uSlBCzpldmVudC10eXBlBwARY29udGVudEJsb2NrU3RhcnQNOmNvbnRl + bnQtdHlwZQcAEGFwcGxpY2F0aW9uL2pzb24NOm1lc3NhZ2UtdHlwZQcABWV2ZW50eyJjb250ZW50 + QmxvY2tJbmRleCI6MCwicCI6ImFiY2RlZmdoaSIsInN0YXJ0Ijp7InRvb2xVc2UiOnsibmFtZSI6 + ImdldF93ZWF0aGVyIiwidG9vbFVzZUlkIjoidG9vbHVzZV9od0tPMkljSkdLSzRSTWxlSndZblRt + In19fWPJlGIAAADEAAAAV4y4VSULOmV2ZW50LXR5cGUHABFjb250ZW50QmxvY2tEZWx0YQ06Y29u + dGVudC10eXBlBwAQYXBwbGljYXRpb24vanNvbg06bWVzc2FnZS10eXBlBwAFZXZlbnR7ImNvbnRl + bnRCbG9ja0luZGV4IjowLCJkZWx0YSI6eyJ0b29sVXNlIjp7ImlucHV0Ijoie1wiY2l0eVwiOlwi + UGFyaXNcIn0ifX0sInAiOiJhYmNkZWZnaGlqIn3jnRuKAAAArAAAAFZS/Vk/CzpldmVudC10eXBl + BwAQY29udGVudEJsb2NrU3RvcA06Y29udGVudC10eXBlBwAQYXBwbGljYXRpb24vanNvbg06bWVz + c2FnZS10eXBlBwAFZXZlbnR7ImNvbnRlbnRCbG9ja0luZGV4IjowLCJwIjoiYWJjZGVmZ2hpamts + bW5vcHFyc3R1dnd4eXpBQkNERUZHSElKS0xNTiJ9MZDPXQAAAJMAAABR7+jjyws6ZXZlbnQtdHlw + ZQcAC21lc3NhZ2VTdG9wDTpjb250ZW50LXR5cGUHABBhcHBsaWNhdGlvbi9qc29uDTptZXNzYWdl + LXR5cGUHAAVldmVudHsicCI6ImFiY2RlZmdoaWprbG1ub3BxciIsInN0b3BSZWFzb24iOiJ0b29s + X3VzZSJ9FcGaWgAAAQkAAABOMgoCKgs6ZXZlbnQtdHlwZQcACG1ldGFkYXRhDTpjb250ZW50LXR5 + cGUHABBhcHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVldmVudHsibWV0cmljcyI6eyJs + YXRlbmN5TXMiOjY3N30sInAiOiJhYmNkZWZnaGlqa2xtbm9wcXJzdHV2d3h5ekFCQ0RFRkdISUpL + TE1OT1BRUlNUVVZXWCIsInVzYWdlIjp7ImlucHV0VG9rZW5zIjo0MTEsIm91dHB1dFRva2VucyI6 + MTYsInNlcnZlclRvb2xVc2FnZSI6e30sInRvdGFsVG9rZW5zIjo0Mjd9fRkMP4U= + headers: + Connection: + - keep-alive + Content-Type: + - application/vnd.amazon.eventstream + Date: + - Sat, 17 Oct 2026 22:49:01 GMT + Transfer-Encoding: + - chunked + x-amzn-RequestId: + - 5ceccdfb-1053-424c-b3d8-050a251c63bd + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/bedrock_runtime/cassettes/latest/test_wrap_bedrock_converse_stream_parses_tool_input.yaml b/py/src/braintrust/integrations/bedrock_runtime/cassettes/latest/test_wrap_bedrock_converse_stream_parses_tool_input.yaml new file mode 100644 index 000000000..ee6bbd481 --- /dev/null +++ b/py/src/braintrust/integrations/bedrock_runtime/cassettes/latest/test_wrap_bedrock_converse_stream_parses_tool_input.yaml @@ -0,0 +1,61 @@ +interactions: +- request: + body: '{"messages": [{"role": "user", "content": [{"text": "What''s the weather + in Paris? Use the tool."}]}], "inferenceConfig": {"maxTokens": 200, "temperature": + 0}, "toolConfig": {"tools": [{"toolSpec": {"name": "get_weather", "description": + "Get the current weather for a city.", "inputSchema": {"json": {"type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, "required": + ["city"]}}}}], "toolChoice": {"tool": {"name": "get_weather"}}}}' + headers: + Content-Length: + - '466' + Content-Type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + User-Agent: + - !!binary | + Qm90bzMvMS40My4xMTEgbWQvQm90b2NvcmUjMS40My4xMTEgdWEvMi4xIG9zL2xpbnV4IzYuNi44 + Ny4yLW1pY3Jvc29mdC1zdGFuZGFyZC1XU0wyIG1kL2FyY2gjeDg2XzY0IGxhbmcvcHl0aG9uIzMu + MTQuOCBtZC9weWltcGwjQ1B5dGhvbiBtL0QsWixiLG4sMyBjZmcvcmV0cnktbW9kZSNsZWdhY3kg + Qm90b2NvcmUvMS40My4xMTE= + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.amazon.nova-lite-v1%3A0/converse-stream + response: + body: + string: !!binary | + AAAAsQAAAFLN4M4VCzpldmVudC10eXBlBwAMbWVzc2FnZVN0YXJ0DTpjb250ZW50LXR5cGUHABBh + cHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVldmVudHsicCI6ImFiY2RlZmdoaWprbG1u + b3BxcnN0dXZ3eHl6QUJDREVGR0hJSktMTU5PUFFSU1RVVldYWVoiLCJyb2xlIjoiYXNzaXN0YW50 + In3DJvYkAAAA5gAAAFc3uSlBCzpldmVudC10eXBlBwARY29udGVudEJsb2NrU3RhcnQNOmNvbnRl + bnQtdHlwZQcAEGFwcGxpY2F0aW9uL2pzb24NOm1lc3NhZ2UtdHlwZQcABWV2ZW50eyJjb250ZW50 + QmxvY2tJbmRleCI6MCwicCI6ImFiY2RlZmdoaSIsInN0YXJ0Ijp7InRvb2xVc2UiOnsibmFtZSI6 + ImdldF93ZWF0aGVyIiwidG9vbFVzZUlkIjoidG9vbHVzZV9od0tPMkljSkdLSzRSTWxlSndZblRt + In19fWPJlGIAAADEAAAAV4y4VSULOmV2ZW50LXR5cGUHABFjb250ZW50QmxvY2tEZWx0YQ06Y29u + dGVudC10eXBlBwAQYXBwbGljYXRpb24vanNvbg06bWVzc2FnZS10eXBlBwAFZXZlbnR7ImNvbnRl + bnRCbG9ja0luZGV4IjowLCJkZWx0YSI6eyJ0b29sVXNlIjp7ImlucHV0Ijoie1wiY2l0eVwiOlwi + UGFyaXNcIn0ifX0sInAiOiJhYmNkZWZnaGlqIn3jnRuKAAAArAAAAFZS/Vk/CzpldmVudC10eXBl + BwAQY29udGVudEJsb2NrU3RvcA06Y29udGVudC10eXBlBwAQYXBwbGljYXRpb24vanNvbg06bWVz + c2FnZS10eXBlBwAFZXZlbnR7ImNvbnRlbnRCbG9ja0luZGV4IjowLCJwIjoiYWJjZGVmZ2hpamts + bW5vcHFyc3R1dnd4eXpBQkNERUZHSElKS0xNTiJ9MZDPXQAAAJMAAABR7+jjyws6ZXZlbnQtdHlw + ZQcAC21lc3NhZ2VTdG9wDTpjb250ZW50LXR5cGUHABBhcHBsaWNhdGlvbi9qc29uDTptZXNzYWdl + LXR5cGUHAAVldmVudHsicCI6ImFiY2RlZmdoaWprbG1ub3BxciIsInN0b3BSZWFzb24iOiJ0b29s + X3VzZSJ9FcGaWgAAAQkAAABOMgoCKgs6ZXZlbnQtdHlwZQcACG1ldGFkYXRhDTpjb250ZW50LXR5 + cGUHABBhcHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVldmVudHsibWV0cmljcyI6eyJs + YXRlbmN5TXMiOjY3N30sInAiOiJhYmNkZWZnaGlqa2xtbm9wcXJzdHV2d3h5ekFCQ0RFRkdISUpL + TE1OT1BRUlNUVVZXWCIsInVzYWdlIjp7ImlucHV0VG9rZW5zIjo0MTEsIm91dHB1dFRva2VucyI6 + MTYsInNlcnZlclRvb2xVc2FnZSI6e30sInRvdGFsVG9rZW5zIjo0Mjd9fRkMP4U= + headers: + Connection: + - keep-alive + Content-Type: + - application/vnd.amazon.eventstream + Date: + - Sat, 17 Oct 2026 22:49:01 GMT + Transfer-Encoding: + - chunked + x-amzn-RequestId: + - 5ceccdfb-1053-424c-b3d8-050a251c63bd + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/bedrock_runtime/test_bedrock_runtime.py b/py/src/braintrust/integrations/bedrock_runtime/test_bedrock_runtime.py index cfa227f2c..b979c4543 100644 --- a/py/src/braintrust/integrations/bedrock_runtime/test_bedrock_runtime.py +++ b/py/src/braintrust/integrations/bedrock_runtime/test_bedrock_runtime.py @@ -289,6 +289,14 @@ def _converse_tool_kwargs(model_id=CONVERSE_MODEL): } +def test_parse_streamed_tool_input_keeps_input_the_decoder_rejects(): + from braintrust.integrations.bedrock_runtime.tracing import _parse_streamed_tool_input + + assert _parse_streamed_tool_input('{"city": "Par') == '{"city": "Par' + deeply_nested = "[" * 100_000 + "]" * 100_000 + assert _parse_streamed_tool_input(deeply_nested) == deeply_nested + + @pytest.mark.vcr def test_wrap_bedrock_converse_normalizes_tool_config_and_output(memory_logger): assert not memory_logger.pop() @@ -331,6 +339,28 @@ def test_wrap_bedrock_converse_normalizes_tool_config_and_output(memory_logger): assert isinstance(tool_use.get("input"), dict) +@pytest.mark.vcr +def test_wrap_bedrock_converse_stream_parses_tool_input(memory_logger): + assert not memory_logger.pop() + client = wrap_bedrock(_bedrock_client()) + + response = client.converse_stream(**_converse_tool_kwargs(CONVERSE_STREAM_MODEL)) + tool_input = "" + for event in response["stream"]: + delta = event.get("contentBlockDelta", {}).get("delta", {}) + if "toolUse" in delta: + tool_input += delta["toolUse"]["input"] + + spans = memory_logger.pop() + assert len(spans) == 1 + content = spans[0]["output"][0]["content"] + tool_uses = [block for block in content if block.get("type") == "tool_use"] + assert tool_uses, f"expected a tool_use block in {content!r}" + # Same shape as the non-streamed converse output: the streamed JSON string is parsed. + assert tool_uses[0]["input"] == json.loads(tool_input) + assert isinstance(tool_uses[0]["input"], dict) + + @pytest.mark.vcr def test_wrap_bedrock_converse_with_image_input_materializes_attachment(memory_logger): assert not memory_logger.pop() diff --git a/py/src/braintrust/integrations/bedrock_runtime/tracing.py b/py/src/braintrust/integrations/bedrock_runtime/tracing.py index c778e359e..d029db192 100644 --- a/py/src/braintrust/integrations/bedrock_runtime/tracing.py +++ b/py/src/braintrust/integrations/bedrock_runtime/tracing.py @@ -560,6 +560,18 @@ def _anthropic_invoke_usage_metrics(model_id: Any, output: Any) -> dict[str, Any return metrics +def _parse_streamed_tool_input(text: str) -> Any: + """Parse streamed tool input into the JSON value the non-streamed converse output has.""" + if not text: + return {} + try: + return json.loads(text) + except (ValueError, RecursionError): + # e.g. cut off by maxTokens, or nested deeper than the decoder allows: + # keep what was streamed + return text + + class _TracedConverseStream: """Wrap a botocore EventStream and log the final span when it drains.""" @@ -656,7 +668,7 @@ def _finish(self, error: BaseException | None) -> None: block = self._content_blocks.setdefault(idx, {}) text = "".join(pieces) if block.get("type") == "tool_use": - block["input"] = text + block["input"] = _parse_streamed_tool_input(text) else: block["text"] = text content = [self._content_blocks[idx] for idx in sorted(self._content_blocks)]