feat(weave): instrument invoke_model_with_response_stream for token tracking

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"]

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)