Skip to content

Commit 5a646aa

Browse files
lievanlievan
andauthored
fix(llmobs): fix parsing errors for bedrock converse stream (#13238)
Fixing a noop test discovered a couple bugs for parsing out the response from converse stream 1. We were adding a `[{"role": "assistant"}]` message that didn't contain any content blocks to be parsed out - the fix is to only try to extract a message if it contains associated content block indices 2. `toolName` -> `name` ## Checklist - [x] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [x] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting) --------- Co-authored-by: lievan <[email protected]>
1 parent 346ab1b commit 5a646aa

File tree

4 files changed

+18
-13
lines changed

4 files changed

+18
-13
lines changed

ddtrace/llmobs/_integrations/bedrock.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -186,18 +186,18 @@ def _extract_output_message_for_converse_stream(
186186

187187
if "messageStart" in chunk:
188188
message_data = chunk["messageStart"]
189-
current_message = {"role": message_data.get("role", "assistant"), "context_block_indices": []}
189+
current_message = {"role": message_data.get("role", "assistant"), "content_block_indicies": []}
190190

191191
# always make sure we have a current message
192192
if current_message is None:
193-
current_message = {"role": "assistant", "context_block_indices": []}
193+
current_message = {"role": "assistant", "content_block_indicies": []}
194194

195195
if "contentBlockStart" in chunk:
196196
block_start = chunk["contentBlockStart"]
197197
index = block_start.get("contentBlockIndex")
198198

199199
if index is not None:
200-
current_message["context_block_indices"].append(index)
200+
current_message["content_block_indicies"].append(index)
201201
if "start" in block_start and "toolUse" in block_start["start"]:
202202
tool_content_blocks[index] = block_start["start"]["toolUse"]
203203

@@ -206,8 +206,8 @@ def _extract_output_message_for_converse_stream(
206206
index = content_block_delta.get("contentBlockIndex")
207207

208208
if index is not None and "delta" in content_block_delta:
209-
if index not in current_message.get("context_block_indices", []):
210-
current_message["context_block_indices"].append(index)
209+
if index not in current_message.get("content_block_indicies", []):
210+
current_message["content_block_indicies"].append(index)
211211

212212
delta_content = content_block_delta["delta"]
213213
text_content_blocks[index] = text_content_blocks.get(index, "") + delta_content.get("text", "")
@@ -225,7 +225,7 @@ def _extract_output_message_for_converse_stream(
225225
current_message = None
226226

227227
# Handle the case where we didn't receive an explicit message stop event
228-
if current_message is not None:
228+
if current_message is not None and current_message.get("content_block_indicies"):
229229
messages.append(
230230
get_final_message_converse_stream_message(current_message, text_content_blocks, tool_content_blocks)
231231
)

ddtrace/llmobs/_integrations/utils.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -908,7 +908,7 @@ def get_final_message_converse_stream_message(
908908
Returns:
909909
Dict containing the processed message with content and optional tool calls
910910
"""
911-
indices = sorted(message.get("context_block_indices", []))
911+
indices = sorted(message.get("content_block_indicies", []))
912912
message_output = {"role": message["role"]}
913913

914914
text_contents = [text_blocks[idx] for idx in indices if idx in text_blocks]
@@ -920,7 +920,7 @@ def get_final_message_converse_stream_message(
920920
if not tool_block:
921921
continue
922922
tool_call = {
923-
"name": tool_block.get("toolName", ""),
923+
"name": tool_block.get("name", ""),
924924
"tool_id": tool_block.get("toolUseId", ""),
925925
}
926926
tool_input = tool_block.get("input")
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
fixes:
3+
- |
4+
LLM Observability: This fix resolves an issue where tool call names were not being captured for bedrock `converse_stream` calls.
5+
- |
6+
LLM Observability: This fix resolves an issue where bedrock `converse_stream` calls contained an extra empty output message.

tests/contrib/botocore/test_bedrock_llmobs.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -735,12 +735,12 @@ def test_llmobs_converse(cls, bedrock_client, request_vcr, mock_tracer, llmobs_e
735735
span = mock_tracer.pop_traces()[0][0]
736736
assert len(llmobs_events) == 1
737737

738-
llmobs_events[0] == _expected_llmobs_llm_span_event(
738+
assert llmobs_events[0] == _expected_llmobs_llm_span_event(
739739
span,
740740
model_name="claude-3-sonnet-20240229-v1:0",
741741
model_provider="anthropic",
742742
input_messages=[
743-
{"role": "system", "content": request_params.get("system")},
743+
{"role": "system", "content": request_params.get("system")[0]["text"]},
744744
{"role": "user", "content": request_params.get("messages")[0].get("content")[0].get("text")},
745745
],
746746
output_messages=[
@@ -813,12 +813,12 @@ def test_llmobs_converse_stream(cls, bedrock_client, request_vcr, mock_tracer, l
813813
span = mock_tracer.pop_traces()[0][0]
814814
assert len(llmobs_events) == 1
815815

816-
llmobs_events[0] == _expected_llmobs_llm_span_event(
816+
assert llmobs_events[0] == _expected_llmobs_llm_span_event(
817817
span,
818818
model_name="claude-3-sonnet-20240229-v1:0",
819819
model_provider="anthropic",
820820
input_messages=[
821-
{"role": "system", "content": request_params.get("system")},
821+
{"role": "system", "content": request_params.get("system")[0]["text"]},
822822
{"role": "user", "content": request_params.get("messages")[0].get("content")[0].get("text")},
823823
],
824824
output_messages=[
@@ -835,7 +835,6 @@ def test_llmobs_converse_stream(cls, bedrock_client, request_vcr, mock_tracer, l
835835
}
836836
],
837837
metadata={
838-
"stop_reason": "tool_use",
839838
"temperature": request_params.get("inferenceConfig", {}).get("temperature"),
840839
"max_tokens": request_params.get("inferenceConfig", {}).get("maxTokens"),
841840
},

0 commit comments

Comments
 (0)