diff --git a/tests/integrations/bedrock/bedrock_test.py b/tests/integrations/bedrock/bedrock_test.py index a469f3921765..4ec8f5ed798f 100644 --- a/tests/integrations/bedrock/bedrock_test.py +++ b/tests/integrations/bedrock/bedrock_test.py @@ -386,6 +386,19 @@ "sessionId": "test-session", } +MOCK_INVOKE_STREAM_EVENTS = [ + { + "chunk": { + "bytes": b'{"type":"content_block_delta","delta":{"type":"text_delta","text":"Hello"}}' + } + }, + { + "chunk": { + "bytes": b'{"type":"content_block_delta","delta":{"type":"text_delta","text":" World"}}' + } + }, +] + # Original botocore _make_api_call function orig = botocore.client.BaseClient._make_api_call @@ -482,6 +495,30 @@ def mock_invoke_agent_make_api_call( return orig(self, operation_name, api_params) +def mock_invoke_stream_make_api_call( + self, operation_name: str, api_params: dict +) -> dict: + if operation_name == "InvokeModelWithResponseStream": + return { + "ResponseMetadata": { + "RequestId": "b2c3d4e5-f6a7-890b-c1d2-e3f4a5b6c7d8", + "HTTPStatusCode": 200, + "HTTPHeaders": { + "date": "Fri, 20 Dec 2024 16:44:08 GMT", + "content-type": "application/vnd.amazon.eventstream", + "connection": "keep-alive", + "x-amzn-requestid": "b2c3d4e5-f6a7-890b-c1d2-e3f4a5b6c7d8", + "x-amzn-bedrock-input-token-count": "42", + "x-amzn-bedrock-output-token-count": "15", + }, + "RetryAttempts": 0, + }, + "body": iter(MOCK_INVOKE_STREAM_EVENTS), + "contentType": "application/json", + } + return orig(self, operation_name, api_params) + + @mock_aws @pytest.mark.parametrize("model_identifier", [model_id, inference_profile_id]) def test_bedrock_converse( @@ -1049,3 +1086,62 @@ def test_bedrock_agent_invoke_agent( summary = call.summary assert summary is not None, "Summary should not be None" assert "usage" in summary + + +@mock_aws +def test_bedrock_invoke_model_stream( + client: weave.trace.weave_client.WeaveClient, +) -> None: + """invoke_model_with_response_stream is patched and token counts are captured from headers.""" + bedrock_client = boto3.client("bedrock-runtime", region_name="us-east-1") + patch_client(bedrock_client) + + with patch( + "botocore.client.BaseClient._make_api_call", new=mock_invoke_stream_make_api_call + ): + body = json.dumps( + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 30, + "messages": [{"role": "user", "content": invoke_prompt}], + } + ) + response = bedrock_client.invoke_model_with_response_stream( + modelId=model_id, + body=body, + contentType="application/json", + accept="application/json", + ) + + # The caller can consume events from the body stream + received_bytes = b"" + for event in response["body"]: + chunk = event.get("chunk", {}) + if "bytes" in chunk: + received_bytes += chunk["bytes"] + + assert b"Hello" in received_bytes + assert b"World" in received_bytes + + calls = list(client.get_calls()) + assert len(calls) == 1, "Expected exactly one trace call" + call = calls[0] + + assert call.exception is None + assert call.ended_at is not None + assert "invoke_stream" in call.op_name + + # Token counts extracted from HTTP response headers + summary = call.summary + assert summary is not None + model_usage = summary["usage"][model_id] + assert model_usage["requests"] == 1 + assert model_usage["prompt_tokens"] == 42 + assert model_usage["completion_tokens"] == 15 + assert model_usage["total_tokens"] == 57 + + # Logged output contains the buffered stream content + output = call.output + assert "body" in output + assert "Hello" in output["body"]["content"] + assert "World" in output["body"]["content"] diff --git a/weave/integrations/bedrock/bedrock_sdk.py b/weave/integrations/bedrock/bedrock_sdk.py index a1cc83680ce1..84fdb2149dbb 100644 --- a/weave/integrations/bedrock/bedrock_sdk.py +++ b/weave/integrations/bedrock/bedrock_sdk.py @@ -212,6 +212,75 @@ def postprocess_output_invoke( return outputs +def bedrock_on_finish_invoke_stream( + call: Call, output: Any, exception: BaseException | None +) -> None: + model_name = str(call.inputs["modelId"]) + usage = {model_name: {"requests": 1}} + summary_update = {"usage": usage} + + if output and "ResponseMetadata" in output: + headers = output["ResponseMetadata"]["HTTPHeaders"] + prompt_tokens = int(headers.get("x-amzn-bedrock-input-token-count", 0)) + completion_tokens = int(headers.get("x-amzn-bedrock-output-token-count", 0)) + tokens_metrics = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } + usage[model_name].update(tokens_metrics) + if call.summary is not None: + call.summary.update(summary_update) + + +def postprocess_output_invoke_stream( + outputs: dict[str, Any] | None, +) -> dict[str, Any] | None: + """Process invoke_model_with_response_stream outputs for logging. + + Buffers the EventStream body so weave can log content while still providing + callers with an iterable over the original events via a reconstructed generator. + ResponseMetadata is preserved on the original outputs dict so the on_finish + handler can read token counts from HTTP headers. + """ + if outputs is None: + return None + if "body" not in outputs: + return outputs + + outputs_copy = {k: copy.deepcopy(v) for k, v in outputs.items() if k != "body"} + + all_events: list[dict] = [] + all_bytes = b"" + + try: + body_stream = outputs["body"] + for event in body_stream: + all_events.append(event) + chunk = event.get("chunk", {}) + if isinstance(chunk, dict) and "bytes" in chunk: + all_bytes += chunk["bytes"] + + def recreate_stream() -> Any: + yield from all_events + + outputs["body"] = recreate_stream() + + try: + body_text = all_bytes.decode("utf-8") + except UnicodeDecodeError: + body_text = repr(all_bytes) + + outputs_copy["body"] = { + "content": body_text, + "event_count": len(all_events), + } + except Exception as e: + outputs_copy["body"] = {"_stream_error": str(e)} + + return outputs_copy + + def postprocess_output_invoke_agent( outputs: dict[str, Any] | None, ) -> dict[str, Any] | None: @@ -473,6 +542,19 @@ def __getitem__(self, key: str) -> Any: return wrapper +def _patch_invoke_model_stream(bedrock_client: "BaseClient") -> None: + op = weave.op( + bedrock_client.invoke_model_with_response_stream, + name="BedrockRuntime.invoke_stream", + postprocess_inputs=postprocess_inputs_invoke, + postprocess_output=postprocess_output_invoke_stream, + kind="llm", + attributes=BEDROCK_INTEGRATION.as_attributes(), + ) + op._set_on_finish_handler(bedrock_on_finish_invoke_stream) + bedrock_client.invoke_model_with_response_stream = op + + def _patch_converse_stream(bedrock_client: "BaseClient") -> None: """Patches the converse_stream method to handle streaming.""" op = create_stream_wrapper("BedrockRuntime.converse_stream")( @@ -490,6 +572,8 @@ def patch_client(bedrock_client: "BaseClient") -> None: _patch_converse(bedrock_client) _patch_converse_stream(bedrock_client) _patch_invoke(bedrock_client) + if hasattr(bedrock_client, "invoke_model_with_response_stream"): + _patch_invoke_model_stream(bedrock_client) elif hasattr(bedrock_client, "apply_guardrail"): # This is a standard bedrock-runtime client _patch_apply_guardrail(bedrock_client)