diff --git a/py/noxfile.py b/py/noxfile.py index a7bffce0..3916f62d 100644 --- a/py/noxfile.py +++ b/py/noxfile.py @@ -230,6 +230,18 @@ def test_cohere(session, version): _run_tests(session, f"{INTEGRATION_DIR}/cohere/test_cohere.py", version=version) +BOTO3_VERSIONS = _get_matrix_versions("boto3") + + +@nox.session() +@nox.parametrize("version", BOTO3_VERSIONS, ids=BOTO3_VERSIONS) +def test_boto3(session, version): + """Test the native boto3 SDK integration.""" + _install_test_deps(session) + _install_matrix_dep(session, "boto3", version) + _run_tests(session, f"{INTEGRATION_DIR}/boto3/test_boto3.py", version=version) + + INSTRUCTOR_VERSIONS = _get_matrix_versions("instructor") diff --git a/py/pyproject.toml b/py/pyproject.toml index 92c56879..c9577039 100644 --- a/py/pyproject.toml +++ b/py/pyproject.toml @@ -438,6 +438,9 @@ latest = "pytest==9.1.0" [tool.braintrust.matrix.braintrust-core] latest = "braintrust-core==0.0.59" +[tool.braintrust.matrix.boto3] +latest = "boto3==1.43.26" + # --------------------------------------------------------------------------- # Vendor packages — optional third-party packages the SDK can work without. # Keys are matrix keys; values are Python import names. The noxfile uses this @@ -461,6 +464,7 @@ agentscope = ["agentscope"] agno = ["agno"] autogen = ["autogen-agentchat"] anthropic = ["anthropic"] +boto3 = ["boto3"] cohere = ["cohere"] claude_agent_sdk = ["claude-agent-sdk"] crewai = ["crewai"] @@ -485,6 +489,7 @@ agentscope = "agentscope" autogen-agentchat = "autogen_agentchat" autogen-ext = "autogen_ext" anthropic = "anthropic" +boto3 = "boto3" cohere = "cohere" autoevals = "autoevals" braintrust-core = "braintrust_core" diff --git a/py/src/braintrust/auto.py b/py/src/braintrust/auto.py index 4d61cc2c..51dd0b61 100644 --- a/py/src/braintrust/auto.py +++ b/py/src/braintrust/auto.py @@ -13,6 +13,7 @@ AgnoIntegration, AnthropicIntegration, AutoGenIntegration, + Boto3Integration, ClaudeAgentSDKIntegration, CohereIntegration, CrewAIIntegration, @@ -76,6 +77,7 @@ def auto_instrument( strands: bool = True, temporal: bool = True, livekit_agents: bool = True, + boto3: bool = True, ) -> dict[str, bool]: """ Auto-instrument supported AI/ML libraries for Braintrust tracing. @@ -203,7 +205,8 @@ def auto_instrument( results["temporal"] = _instrument_integration(TemporalIntegration) if livekit_agents: results["livekit_agents"] = _instrument_integration(LiveKitAgentsIntegration) - + if boto3: + results["boto3"] = _instrument_integration(Boto3Integration) return results diff --git a/py/src/braintrust/integrations/__init__.py b/py/src/braintrust/integrations/__init__.py index 52ecca8e..e039826e 100644 --- a/py/src/braintrust/integrations/__init__.py +++ b/py/src/braintrust/integrations/__init__.py @@ -3,6 +3,7 @@ from .agno import AgnoIntegration from .anthropic import AnthropicIntegration from .autogen import AutoGenIntegration +from .boto3 import Boto3Integration from .claude_agent_sdk import ClaudeAgentSDKIntegration from .cohere import CohereIntegration from .crewai import CrewAIIntegration @@ -29,6 +30,7 @@ "AgnoIntegration", "AnthropicIntegration", "AutoGenIntegration", + "Boto3Integration", "ClaudeAgentSDKIntegration", "CohereIntegration", "CrewAIIntegration", diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_auto_boto3_sdk.py b/py/src/braintrust/integrations/auto_test_scripts/test_auto_boto3_sdk.py new file mode 100644 index 00000000..5dc2931e --- /dev/null +++ b/py/src/braintrust/integrations/auto_test_scripts/test_auto_boto3_sdk.py @@ -0,0 +1,65 @@ +import os + +import boto3 +from braintrust.auto import auto_instrument +from braintrust.integrations.boto3.patchers import Boto3ConversePatcher +from braintrust.integrations.test_utils import autoinstrument_test_context + + +def is_patched(target, patcher): + return bool(getattr(target, patcher.nested_patch_marker_attr(), False)) + + +# verify not patched initially +original_client = boto3.client( + "bedrock-runtime", + aws_access_key_id="test-key", + aws_secret_access_key="test-access-key", + aws_session_token="test-session-key", + region_name="us-east-1", +) + +assert is_patched(original_client, Boto3ConversePatcher) is False + +# instrument +results = auto_instrument() +assert results.get("boto3") is True + +# patch +patched_client = boto3.client( + "bedrock-runtime", + aws_access_key_id="", + aws_secret_access_key="", + aws_session_token="", + region_name="us-east-1", +) + +# idempotent +results = auto_instrument() +assert results.get("boto3") is True + + +assert is_patched(patched_client, Boto3ConversePatcher) is True + +# make api call +with autoinstrument_test_context("test_auto_boto3", integration="boto3") as memory_logger: + client = boto3.client( + "bedrock-runtime", + region_name="us-east-1", + aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"), + aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"), + aws_session_token=os.getenv("AWS_SESSION_TOKEN"), + ) + + response = response = client.converse( + modelId="anthropic.claude-3-haiku-20240307-v1:0", + messages=[{"role": "user", "content": [{"text": "what's 2+2"}]}], + system=[{"text": "answer only in single integer"}], + ) + + spans = memory_logger.pop() + assert len(spans) == 1, f"Expected 1 span, got {len(spans)}" + + span = spans[0] + assert span["metadata"]["model_provider"] == "anthropic" + assert "claude" in span["metadata"]["model"] diff --git a/py/src/braintrust/integrations/base.py b/py/src/braintrust/integrations/base.py index b9664e15..39ab341d 100644 --- a/py/src/braintrust/integrations/base.py +++ b/py/src/braintrust/integrations/base.py @@ -437,6 +437,30 @@ def _wrapper(cls, wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any: return cls.wrapper(wrapped, instance, args, kwargs) +class NestedFunctionWrapperPatcher(FunctionWrapperPatcher): + """Add Wrapping for target attribute when the object is formed; useful for dynamic sdk builders such as boto3""" + + target_attribute: ClassVar[str] + nested_wrapper: ClassVar[Any] + wrapper: ClassVar[Any] = staticmethod(_call_wrapped) + + @classmethod + def nested_patch_marker_attr(cls) -> str: + """Return the sentinel attribute used to mark an instance as set up.""" + suffix = re.sub(r"\W+", "_", cls.name).strip("_") + return f"__braintrust_nested_patched_{suffix}__" + + @classmethod + def _wrapper(cls, wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any: + obj = cls.wrapper(wrapped, instance, args, kwargs) + if not getattr(obj, cls.target_attribute, False): + return + marker = cls.nested_patch_marker_attr() + wrap_function_wrapper(obj, cls.target_attribute, cls.nested_wrapper) + setattr(obj, marker, True) + return obj + + class ClassReplacementPatcher(BasePatcher): """Base patcher for replacing an exported class with a tracing wrapper class. diff --git a/py/src/braintrust/integrations/boto3/__init__.py b/py/src/braintrust/integrations/boto3/__init__.py new file mode 100644 index 00000000..0420e45e --- /dev/null +++ b/py/src/braintrust/integrations/boto3/__init__.py @@ -0,0 +1,23 @@ +import logging + +from braintrust.logger import NOOP_SPAN, current_span, init_logger + +from .integration import Boto3Integration + + +logger = logging.getLogger(__name__) + +__all__ = ["Boto3Integration", "setup_boto3"] + + +def setup_boto3( + api_key: str | None = None, + project_id: str | None = None, + project_name: str | None = None, +) -> bool: + span = current_span() + + if span == NOOP_SPAN: + init_logger(project=project_name, api_key=api_key, project_id=project_id) + + return Boto3Integration.setup() diff --git a/py/src/braintrust/integrations/boto3/cassettes/latest/test_auto_boto3.yaml b/py/src/braintrust/integrations/boto3/cassettes/latest/test_auto_boto3.yaml new file mode 100644 index 00000000..1b521cd0 --- /dev/null +++ b/py/src/braintrust/integrations/boto3/cassettes/latest/test_auto_boto3.yaml @@ -0,0 +1,62 @@ +interactions: +- request: + body: '{"messages": [{"role": "user", "content": [{"text": "what''s 2+2"}]}], + "system": [{"text": "answer only in single integer"}]}' + headers: + Content-Length: + - '124' + Content-Type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + User-Agent: + - !!binary | + Qm90bzMvMS40My4yNiBtZC9Cb3RvY29yZSMxLjQzLjI4IHVhLzIuMSBvcy9saW51eCM2LjguMC0x + MDUyLWF6dXJlIG1kL2FyY2gjeDg2XzY0IGxhbmcvcHl0aG9uIzMuMTIuMSBtZC9weWltcGwjQ1B5 + dGhvbiBtL1osYixELGUgY2ZnL3JldHJ5LW1vZGUjbGVnYWN5IEJvdG9jb3JlLzEuNDMuMjg= + X-Amz-Date: + - !!binary | + MjAyNjA2MTJUMDkwMjE5Wg== + X-Amz-Security-Token: + - !!binary | + SVFvSmIzSnBaMmx1WDJWakVFWWFDWFZ6TFhkbGMzUXRNaUpITUVVQ0lFSlVDZDFxQlZncWhMd3BM + cFluTW92aWxVbnRZOVFTNm1HWitkT2lGdWxBQWlFQXFzYzJ0bXloWjk0N2J2ZXJsVW9yV3BkU0xK + UURmMEdlK0pIUHo4U0dhNWdxZ3dNSUR4QUFHZ3czTWpjMk5EWTBOelE0TkRVaURGWWh1dm52blpJ + bjJtMjYvaXJnQXZiNVIvV2FzOFFSeHBEQXozQ0ppM1hNcUFjR2JlU1FCS0psZklMUDBJWjQ2cW90 + d1VNeGRRR0hkbnVvQk9XcnJGTzNGdVlrVkxqYXZHd0tDeWRhSkFlcC9iODR6QzRlbGhWZVNQNUNx + L1hNeGl5SHZaSHROZFdvMEVsN0Fma2dQNVhSUnNYM3ZRM0VER1Y4VW9JL0tSa3JHcGhwbWxtWWQv + dG9JdzJzeTRTYWJoNXpBL0NDRUJWRC9ETlJUaDkxR0MrQ1kvR0YrTk9JUk9lVFY1VmtzWndyOVpz + Tmh2NExqcURmeVdHVXoxbERwZmRzRUhPNHpRbjVBVy9oRGNuc04rYlFuR1NtSThaTjk1VTJsdVlr + elBPOGJUUzJobFJNeFN0KzZvbE1NZ09CUkFkbW1OamlYVGNaK3FkNnBCM2paWDcwVWpOZURIdCtJ + SDRoNUozVGxmVXR5ZG01TlBXdzZsb1l2WitPV1M2dUwyTVg1ZlhtWTlJMVBlZlV1M21KQTRFNlhN + aHNncmlYaWQ5THY1UFJoMU1ybTUvN2pOZTVkLzhuSTdJYTZlLytwTVJXYXEreHVvYzIxQWdob2JF + Q3lWUU1XcWZKSWZrclcxMGpMU3FPYlMwd2tMbXUwUVk2bVFHdmZCaitadTN6TnNpL0hRRUIzRjB6 + QWtmRENHYmhYbXJrTzVjRkJPNG0zdFdjNmJoWGlST2o0QVRmVVRzSGVvOFlzOExhTEdsSGtSbmdt + WWc4Y3dPVG5FSTBUNldUQjVEYm41L2IybUN6dHdJNHdCTFBUUUhkcVRxVDJ2M05JbC8zTitpd0dQ + UytGRGRnYldPZXVjZEV6TWluWnNNRXJ4RUExK1dZMVNtMkp6alN6VmFOc0NZTDZXQ3A5dlNneUE3 + U001aDFQVjlLRCtnPQ== + amz-sdk-invocation-id: + - !!binary | + MzkwYzg5YjktMjM4Ni00NWI1LTkzZTYtYjgyMjc1YWM2OWY2 + amz-sdk-request: + - !!binary | + YXR0ZW1wdD0x + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse + response: + body: + string: '{"metrics":{"latencyMs":269},"output":{"message":{"content":[{"text":"4"}],"role":"assistant"}},"stopReason":"end_turn","usage":{"inputTokens":18,"outputTokens":5,"serverToolUsage":{},"totalTokens":23}}' + headers: + Connection: + - keep-alive + Content-Length: + - '202' + Content-Type: + - application/json + Date: + - Fri, 12 Jun 2026 09:02:20 GMT + x-amzn-RequestId: + - 20c1e6b5-f545-444a-8de6-a47e4c4c295a + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/boto3/cassettes/latest/test_boto3_converse.yaml b/py/src/braintrust/integrations/boto3/cassettes/latest/test_boto3_converse.yaml new file mode 100644 index 00000000..49085757 --- /dev/null +++ b/py/src/braintrust/integrations/boto3/cassettes/latest/test_boto3_converse.yaml @@ -0,0 +1,62 @@ +interactions: +- request: + body: '{"messages": [{"role": "user", "content": [{"text": "what''s 2+2"}]}], + "system": [{"text": "answer only in single integer"}]}' + headers: + Content-Length: + - '124' + Content-Type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + User-Agent: + - !!binary | + Qm90bzMvMS40My4yNiBtZC9Cb3RvY29yZSMxLjQzLjI4IHVhLzIuMSBvcy9saW51eCM2LjguMC0x + MDUyLWF6dXJlIG1kL2FyY2gjeDg2XzY0IGxhbmcvcHl0aG9uIzMuMTIuMSBtZC9weWltcGwjQ1B5 + dGhvbiBtL2UsYixaLEQgY2ZnL3JldHJ5LW1vZGUjbGVnYWN5IEJvdG9jb3JlLzEuNDMuMjg= + X-Amz-Date: + - !!binary | + MjAyNjA2MTJUMDYxMzAxWg== + X-Amz-Security-Token: + - !!binary | + SVFvSmIzSnBaMmx1WDJWakVFWWFDWFZ6TFhkbGMzUXRNaUpITUVVQ0lFSlVDZDFxQlZncWhMd3BM + cFluTW92aWxVbnRZOVFTNm1HWitkT2lGdWxBQWlFQXFzYzJ0bXloWjk0N2J2ZXJsVW9yV3BkU0xK + UURmMEdlK0pIUHo4U0dhNWdxZ3dNSUR4QUFHZ3czTWpjMk5EWTBOelE0TkRVaURGWWh1dm52blpJ + bjJtMjYvaXJnQXZiNVIvV2FzOFFSeHBEQXozQ0ppM1hNcUFjR2JlU1FCS0psZklMUDBJWjQ2cW90 + d1VNeGRRR0hkbnVvQk9XcnJGTzNGdVlrVkxqYXZHd0tDeWRhSkFlcC9iODR6QzRlbGhWZVNQNUNx + L1hNeGl5SHZaSHROZFdvMEVsN0Fma2dQNVhSUnNYM3ZRM0VER1Y4VW9JL0tSa3JHcGhwbWxtWWQv + dG9JdzJzeTRTYWJoNXpBL0NDRUJWRC9ETlJUaDkxR0MrQ1kvR0YrTk9JUk9lVFY1VmtzWndyOVpz + Tmh2NExqcURmeVdHVXoxbERwZmRzRUhPNHpRbjVBVy9oRGNuc04rYlFuR1NtSThaTjk1VTJsdVlr + elBPOGJUUzJobFJNeFN0KzZvbE1NZ09CUkFkbW1OamlYVGNaK3FkNnBCM2paWDcwVWpOZURIdCtJ + SDRoNUozVGxmVXR5ZG01TlBXdzZsb1l2WitPV1M2dUwyTVg1ZlhtWTlJMVBlZlV1M21KQTRFNlhN + aHNncmlYaWQ5THY1UFJoMU1ybTUvN2pOZTVkLzhuSTdJYTZlLytwTVJXYXEreHVvYzIxQWdob2JF + Q3lWUU1XcWZKSWZrclcxMGpMU3FPYlMwd2tMbXUwUVk2bVFHdmZCaitadTN6TnNpL0hRRUIzRjB6 + QWtmRENHYmhYbXJrTzVjRkJPNG0zdFdjNmJoWGlST2o0QVRmVVRzSGVvOFlzOExhTEdsSGtSbmdt + WWc4Y3dPVG5FSTBUNldUQjVEYm41L2IybUN6dHdJNHdCTFBUUUhkcVRxVDJ2M05JbC8zTitpd0dQ + UytGRGRnYldPZXVjZEV6TWluWnNNRXJ4RUExK1dZMVNtMkp6alN6VmFOc0NZTDZXQ3A5dlNneUE3 + U001aDFQVjlLRCtnPQ== + amz-sdk-invocation-id: + - !!binary | + MzBlZjQ1MzktMmFlNS00M2IwLTkwMjUtMWQwZDU3MDkzMjQ5 + amz-sdk-request: + - !!binary | + YXR0ZW1wdD0x + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse + response: + body: + string: '{"metrics":{"latencyMs":257},"output":{"message":{"content":[{"text":"4"}],"role":"assistant"}},"stopReason":"end_turn","usage":{"inputTokens":18,"outputTokens":5,"serverToolUsage":{},"totalTokens":23}}' + headers: + Connection: + - keep-alive + Content-Length: + - '202' + Content-Type: + - application/json + Date: + - Fri, 12 Jun 2026 06:13:02 GMT + x-amzn-RequestId: + - f9de30ea-4863-4786-9de1-a01f0232e54c + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/boto3/integration.py b/py/src/braintrust/integrations/boto3/integration.py new file mode 100644 index 00000000..b49af4bd --- /dev/null +++ b/py/src/braintrust/integrations/boto3/integration.py @@ -0,0 +1,9 @@ +from braintrust.integrations.base import BaseIntegration + +from .patchers import Boto3ConversePatcher + + +class Boto3Integration(BaseIntegration): + name = "boto3_integration" + patchers = (Boto3ConversePatcher,) + import_names = ("botocore",) diff --git a/py/src/braintrust/integrations/boto3/patchers.py b/py/src/braintrust/integrations/boto3/patchers.py new file mode 100644 index 00000000..fbf88d1e --- /dev/null +++ b/py/src/braintrust/integrations/boto3/patchers.py @@ -0,0 +1,13 @@ +from braintrust.integrations.base import NestedFunctionWrapperPatcher + +from .tracing import converse_tracer + + +class Boto3ConversePatcher(NestedFunctionWrapperPatcher): + """Patcher class for Boto3.BedrockRuntime.Client.Converse API""" + + name = "boto3" + target_module = "botocore.client" + target_path = "ClientCreator._create_client_class" + target_attribute = "converse" + nested_wrapper = converse_tracer diff --git a/py/src/braintrust/integrations/boto3/test_boto3.py b/py/src/braintrust/integrations/boto3/test_boto3.py new file mode 100644 index 00000000..ffa0948c --- /dev/null +++ b/py/src/braintrust/integrations/boto3/test_boto3.py @@ -0,0 +1,141 @@ +import boto3 +import pytest +from braintrust import logger +from braintrust.integrations.test_utils import verify_autoinstrument_script +from braintrust.logger import Attachment +from braintrust.test_logger import init_test_logger + +from .integration import Boto3Integration +from .tracing import _normalize_converse_message_content + + +PROJECT_NAME = "test-boto3" + +Boto3Integration.setup() + + +@pytest.fixture +def bedrock_client(): + bedrock_client = boto3.client( + "bedrock-runtime", + aws_access_key_id="", + aws_secret_access_key="", + aws_session_token="", + region_name="us-east-1", + ) + + return bedrock_client + + +@pytest.fixture +def memory_logger(): + init_test_logger(PROJECT_NAME) + with logger._internal_with_memory_background_logger() as bgl: + yield bgl + + +def test_normalize_converse_message_content_converts_inline_media_bytes(): + normalized = _normalize_converse_message_content( + [ + {"text": "summarize these"}, + {"image": {"format": "png", "source": {"bytes": b"image-data"}}}, + { + "document": { + "format": "pdf", + "name": "report", + "source": {"bytes": b"pdf-data"}, + "context": "annual report", + "citations": {"enabled": True}, + } + }, + {"video": {"format": "mp4", "source": {"bytes": b"video-data"}}}, + {"audio": {"format": "mp3", "source": {"bytes": b"audio-data"}}}, + ] + ) + + assert normalized[0] == {"text": "summarize these"} + image_attachment = normalized[1]["image"]["source"]["image_url"]["url"] + assert isinstance(image_attachment, Attachment) + assert image_attachment.reference["content_type"] == "image/png" + + document = normalized[2]["document"] + assert document["context"] == "annual report" + assert document["citations"] == {"enabled": True} + document_attachment = document["source"]["file"]["file_data"] + assert isinstance(document_attachment, Attachment) + assert document_attachment.reference["content_type"] == "application/pdf" + + video_attachment = normalized[3]["video"]["source"]["file"]["file_data"] + assert isinstance(video_attachment, Attachment) + assert video_attachment.reference["content_type"] == "video/mp4" + + audio_attachment = normalized[4]["audio"]["source"]["file"]["file_data"] + assert isinstance(audio_attachment, Attachment) + assert audio_attachment.reference["content_type"] == "audio/mp3" + + +def test_normalize_converse_message_content_preserves_non_binary_sources_and_errors(): + normalized = _normalize_converse_message_content( + [ + { + "document": { + "format": "txt", + "name": "notes", + "source": {"text": "plain text", "content": [{"text": "chunk"}]}, + } + }, + { + "video": { + "format": "mp4", + "source": {"s3Location": {"uri": "s3://bucket/video.mp4", "bucketOwner": "123"}}, + } + }, + {"audio": {"format": "mp3", "error": {"message": "could not decode"}}}, + ] + ) + + assert normalized[0]["document"]["source"] == {"text": "plain text", "content": [{"text": "chunk"}]} + assert normalized[1]["video"]["source"] == {"s3Location": {"uri": "s3://bucket/video.mp4", "bucketOwner": "123"}} + assert normalized[2]["audio"]["error"] == {"message": "could not decode"} + + +@pytest.mark.vcr() +def test_boto3_converse(bedrock_client, memory_logger): + assert not memory_logger.pop() + + messages = [{"role": "user", "content": [{"text": "what's 2+2"}]}] + model_provider = "anthropic" + model_name = "claude-3-haiku-20240307-v1" + + response = bedrock_client.converse( + modelId="anthropic.claude-3-haiku-20240307-v1:0", + messages=[{"role": "user", "content": [{"text": "what's 2+2"}]}], + system=[{"text": "answer only in single integer"}], + ) + + assert response["output"]["message"]["role"] == "assistant" + assert "4" in response["output"]["message"]["content"][0]["text"] + + spans = memory_logger.pop() + + assert len(spans) == 1 + traced_span = spans[0] + + assert isinstance(traced_span["metrics"], dict) + assert traced_span["span_attributes"]["name"] == "bedrock_runtime.converse" + assert traced_span["span_attributes"]["type"] == "llm" + + assert traced_span["input"]["messages"] == messages + + assert traced_span["metadata"]["model"] == model_name + assert traced_span["metadata"]["model_provider"] == model_provider + + output = response.get("output", None) + print(traced_span) + assert output is not None + assert traced_span["output"] == output + + +class TestAutoInstrumentBoto3: + def test_auto_instrument_anthropic(self): + verify_autoinstrument_script("test_auto_boto3_sdk.py") diff --git a/py/src/braintrust/integrations/boto3/tracing.py b/py/src/braintrust/integrations/boto3/tracing.py new file mode 100644 index 00000000..d1742ac5 --- /dev/null +++ b/py/src/braintrust/integrations/boto3/tracing.py @@ -0,0 +1,215 @@ +import time +from typing import Any + +from braintrust.integrations.utils import _camel_to_snake, _resolved_attachment_from_bytes +from braintrust.logger import start_span +from braintrust.span_types import SpanTypeAttribute +from braintrust.util import merge_dicts + + +_RUNTIME_PROVIDER = "aws_bedrock_runtime" + +_MIME_TYPES = { + "document": { + "pdf": "application/pdf", + "csv": "text/csv", + "doc": "application/msword", + "html": "text/html", + "txt": "text/plain", + "md": "text/markdown", + }, + "image": { + "gif": "image/gif", + "jpeg": "image/jpeg", + "jpg": "image/jpeg", + "png": "image/png", + "webp": "image/webp", + }, + "video": { + "flv": "video/flv", + "mkv": "video/x-matroska", + "mov": "video/quicktime", + "mp4": "video/mp4", + "mpeg": "video/mpeg", + "mpg": "video/mpeg", + "three_gp": "video/3gpp", + "webm": "video/webm", + "wmv": "video/x-ms-wmv", + }, + "audio": { + "aac": "audio/aac", + "flac": "audio/flac", + "m4a": "audio/mp4", + "mka": "audio/x-matroska", + "mkv": "audio/x-matroska", + "mp3": "audio/mp3", + "mp4": "audio/mp4", + "mpeg": "audio/mpeg", + "mpga": "audio/mpeg", + "ogg": "audio/ogg", + "opus": "audio/opus", + "pcm": "audio/pcm", + "wav": "audio/wav", + "webm": "audio/webm", + "x-aac": "audio/aac", + }, +} + + +def _normalize_converse_media_block(block_type: str, block: dict[str, Any]) -> dict[str, Any]: + normalized_block = {key: value for key, value in block.items() if key != "source"} + + source = block.get("source") + if not isinstance(source, dict): + return normalized_block + + source_bytes = source.get("bytes") + if source_bytes is not None: + file_format = block.get("format", "unknown") + mime_type = _MIME_TYPES.get(block_type, {}).get(file_format, "application/octet-stream") + resolved_attachment = _resolved_attachment_from_bytes( + source_bytes, + mime_type=mime_type, + prefix=block_type, + ) + normalized_block["source"] = resolved_attachment.multimodal_part_payload + else: + normalized_block["source"] = source + + return normalized_block + + +def _normalize_converse_message_content(message_contents: list[dict[str, Any]]) -> list[dict[str, Any]]: + """helper function for `_get_converse_request`""" + normalized_content: list[dict[str, Any]] = [] + for part_content in message_contents: + if part_content.get("text") is not None: + normalized_content.append({"text": part_content.get("text")}) + continue + + for block_type in ("image", "document", "video", "audio"): + block = part_content.get(block_type) + if isinstance(block, dict): + normalized_content.append({block_type: _normalize_converse_media_block(block_type, block)}) + break + else: + normalized_content.append(part_content) + + return normalized_content + + +def _get_converse_input(**kwargs: dict[str, Any]) -> dict[str, Any]: + """normalize raw input into braintrust supporeted Attachments for multimodel inputs""" + input_data: dict[str, Any] = {} + + messages = kwargs.get("messages", None) + if messages is not None: + normalized_messages = [] + for message in messages: + content = message.get("content") + normalized_messages.append( + { + **message, + "content": _normalize_converse_message_content(content), + } + ) + input_data["messages"] = normalized_messages + + return input_data + + +def _get_converse_request_metadata(**kwargs: dict[str, Any]): + """Log metadata for boto3; track model id, request-id, system instructions, inference config""" + model_id = kwargs.get("modelId", None) + request_id = kwargs.get("amz-sdk-invocation-id") + instructions = kwargs.get("system") + inference_config = kwargs.get("inferenceConfig") + + metadata: dict[str, Any] = {} + metadata["runtime_provider"] = _RUNTIME_PROVIDER + + if model_id is not None: + model_details = model_id.split(".") + if len(model_details) < 2: + metadata["model"] = model_id + else: + model_provider, model_name = model_details + metadata["model"] = model_name.split(":")[0] + metadata["model_provider"] = model_provider + + if request_id is not None: + metadata["request_id"] = request_id + + if instructions is not None: + final_instruction = "" + for instruction in instructions: + insturction_chunk = instruction.get("text", None) + if insturction_chunk is not None: + final_instruction += insturction_chunk + + if final_instruction != "": + metadata["instructions"] = final_instruction + + if inference_config is not None: + inference_dict = {} + for key, value in inference_config.items(): + inference_dict[_camel_to_snake(key)] = value + merge_dicts(metadata, inference_dict) + return metadata + + +def parse_converse_result(result: dict[str, Any]): + """""" + message = result.get("output", {}).get("message", {}) + content = message.get("content") + + _normalized_content = _normalize_converse_message_content(content) + + return {"message": {**message, "content": _normalized_content}} + + +def _extract_converse_metrics(response: dict[str, Any], start_time: float, end_time: float) -> dict[str, float]: + """Extract metrics such as input tokens/output tokens and total tokens""" + metrics: dict[str, float] = {} + + metrics["start"] = start_time + metrics["end"] = end_time + metrics["duration"] = end_time - start_time + + usage: dict[str, float] | None = response.get("usage", None) + + if usage is not None: + if usage.get("inputTokens") is not None: + metrics["prompt_tokens"] = float(usage.get("inputTokens", 0)) + if usage.get("outputTokens") is not None: + metrics["completion_tokens"] = float(usage.get("outputTokens", 0)) + if usage.get("totalTokens") is not None: + metrics["tokens"] = float(usage.get("totalTokens", 0)) + if usage.get("cacheReadInputTokens") is not None: + metrics["prompt_cached_tokens"] = float(usage.get("cacheReadInputTokens", 0)) + if usage.get("cacheWriteInputTokens") is not None: + metrics["completion_cached_tokens"] = float(usage.get("cacheWriteInputTokens", 0)) + + return metrics + + +def converse_tracer(wrapped: Any, instance: Any, args: Any, kwargs: Any): + """trace helper for `boto3.BedrockRuntime.Client.Converse` API""" + + input_data = _get_converse_input(**kwargs) + metadata = _get_converse_request_metadata(**kwargs) + + with start_span( + name="bedrock_runtime.converse", + type=SpanTypeAttribute.LLM, + input=input_data, + metadata=metadata, + ) as span: + start_time = time.time() + result: dict[str, Any] = wrapped(*args, **kwargs) + end_time = time.time() + + output = parse_converse_result(result) + metrics = _extract_converse_metrics(result, start_time, end_time) + span.log(output=output, metrics=metrics) + return result