mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(bedrock/realtime): content-block session state lifecycle
Nova Sonic streams discrete content blocks. Reset part-local state (delta chunks/type/item) on every contentEnd, mint a new item per assistant contentStart, and only emit response.created for a new response so text/tool/audio blocks cannot leak into each other
This commit is contained in:
parent
7d60bf6126
commit
a098330b51
2 changed files with 194 additions and 38 deletions
|
|
@ -754,6 +754,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
current_response_id: str | None,
|
||||
current_output_item_id: str | None,
|
||||
current_conversation_id: str | None,
|
||||
current_delta_type: ALL_DELTA_TYPES | None = None,
|
||||
) -> tuple[
|
||||
list[OpenAIRealtimeEvents],
|
||||
str | None,
|
||||
|
|
@ -764,14 +765,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
"""
|
||||
Transform Bedrock contentStart event to OpenAI response events.
|
||||
|
||||
Args:
|
||||
event: Bedrock contentStart event
|
||||
current_response_id: Current response ID
|
||||
current_output_item_id: Current output item ID
|
||||
current_conversation_id: Current conversation ID
|
||||
|
||||
Returns:
|
||||
Tuple of (events, response_id, output_item_id, conversation_id, delta_type)
|
||||
Bedrock streams one content block at a time (TEXT, AUDIO, TOOL, …). Only
|
||||
ASSISTANT blocks open an OpenAI response/item lifecycle. Non-assistant
|
||||
blocks must not clobber in-flight assistant part state.
|
||||
"""
|
||||
content_start: Final = event["contentStart"]
|
||||
role: Final = content_start.get("role")
|
||||
|
|
@ -782,40 +778,37 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
current_response_id,
|
||||
current_output_item_id,
|
||||
current_conversation_id,
|
||||
None,
|
||||
current_delta_type,
|
||||
)
|
||||
|
||||
verbose_logger.debug("Handling ASSISTANT contentStart")
|
||||
|
||||
# Initialize IDs if needed
|
||||
is_new_response: Final = not current_response_id
|
||||
if not current_response_id:
|
||||
current_response_id = f"resp_{uuid.uuid4()}"
|
||||
if not current_output_item_id:
|
||||
current_output_item_id = f"item_{uuid.uuid4()}"
|
||||
current_output_item_id = f"item_{uuid.uuid4()}"
|
||||
if not current_conversation_id:
|
||||
current_conversation_id = f"conv_{uuid.uuid4()}"
|
||||
|
||||
# Determine content type
|
||||
content_type: Final = content_start.get("type", "TEXT")
|
||||
current_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio"
|
||||
next_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio"
|
||||
|
||||
returned_messages: Final[list[OpenAIRealtimeEvents]] = []
|
||||
|
||||
# Send response.created
|
||||
response_created: Final = OpenAIRealtimeStreamResponseBaseObject(
|
||||
type="response.created",
|
||||
event_id=f"event_{uuid.uuid4()}",
|
||||
response={
|
||||
"object": "realtime.response",
|
||||
"id": current_response_id,
|
||||
"status": "in_progress",
|
||||
"output": [],
|
||||
"conversation_id": current_conversation_id,
|
||||
},
|
||||
)
|
||||
returned_messages.append(response_created)
|
||||
if is_new_response:
|
||||
response_created: Final = OpenAIRealtimeStreamResponseBaseObject(
|
||||
type="response.created",
|
||||
event_id=f"event_{uuid.uuid4()}",
|
||||
response={
|
||||
"object": "realtime.response",
|
||||
"id": current_response_id,
|
||||
"status": "in_progress",
|
||||
"output": [],
|
||||
"conversation_id": current_conversation_id,
|
||||
},
|
||||
)
|
||||
returned_messages.append(response_created)
|
||||
|
||||
# Send response.output_item.added
|
||||
output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded(
|
||||
type="response.output_item.added",
|
||||
response_id=current_response_id,
|
||||
|
|
@ -831,7 +824,6 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
)
|
||||
returned_messages.append(output_item_added)
|
||||
|
||||
# Send response.content_part.added
|
||||
content_part_added: Final = OpenAIRealtimeResponseContentPartAdded(
|
||||
type="response.content_part.added",
|
||||
content_index=0,
|
||||
|
|
@ -839,7 +831,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
event_id=f"event_{uuid.uuid4()}",
|
||||
item_id=current_output_item_id,
|
||||
part=(
|
||||
{"type": "text", "text": ""} if current_delta_type == "text" else {"type": "audio", "transcript": ""}
|
||||
{"type": "text", "text": ""} if next_delta_type == "text" else {"type": "audio", "transcript": ""}
|
||||
),
|
||||
response_id=current_response_id,
|
||||
)
|
||||
|
|
@ -850,7 +842,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
current_response_id,
|
||||
current_output_item_id,
|
||||
current_conversation_id,
|
||||
current_delta_type,
|
||||
next_delta_type,
|
||||
)
|
||||
|
||||
def transform_text_output_event(
|
||||
|
|
@ -955,10 +947,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
verbose_logger.debug("Handling contentEnd: %s", content_end)
|
||||
|
||||
if not current_output_item_id or not current_response_id:
|
||||
return [], current_delta_chunks
|
||||
return [], None
|
||||
|
||||
if content_end.get("type") == "TOOL" or current_delta_type not in ("text", "audio"):
|
||||
return [], current_delta_chunks
|
||||
return [], None
|
||||
|
||||
returned_messages: Final[list[OpenAIRealtimeEvents]] = []
|
||||
|
||||
|
|
@ -1273,8 +1265,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
current_response_id,
|
||||
current_output_item_id,
|
||||
current_conversation_id,
|
||||
current_delta_type,
|
||||
)
|
||||
returned_messages.extend(events)
|
||||
if events:
|
||||
current_delta_chunks = None
|
||||
|
||||
elif "textOutput" in event:
|
||||
events, current_delta_chunks = self.transform_text_output_event(
|
||||
|
|
@ -1299,10 +1294,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
current_delta_chunks,
|
||||
)
|
||||
returned_messages.extend(events)
|
||||
current_delta_chunks = None
|
||||
current_delta_type = None
|
||||
current_output_item_id = None
|
||||
if content_end.get("type") == "TOOL":
|
||||
current_output_item_id = None # rebind-ok: tool block ends so next ASSISTANT mints fresh ids
|
||||
current_response_id = None # rebind-ok: tool block ends so next ASSISTANT mints fresh ids
|
||||
current_delta_type = None # rebind-ok: tool block ends so next ASSISTANT mints fresh ids
|
||||
current_response_id = None
|
||||
if BedrockContentEnd.model_validate(content_end).stopReason == "END_TURN":
|
||||
(
|
||||
done_events,
|
||||
|
|
@ -1311,6 +1307,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
current_delta_type,
|
||||
) = self._response_done_events(current_response_id, current_conversation_id)
|
||||
returned_messages.extend(done_events)
|
||||
current_delta_chunks = None
|
||||
|
||||
elif "toolUse" in event:
|
||||
(
|
||||
|
|
@ -1321,8 +1318,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
tool_response_id,
|
||||
) = self.transform_tool_use_event(event, current_output_item_id, current_response_id)
|
||||
returned_messages.extend(events)
|
||||
current_output_item_id = tool_output_item_id # rebind-ok: persist minted tool item id into session
|
||||
current_response_id = tool_response_id # rebind-ok: persist minted tool response id into session
|
||||
current_output_item_id = tool_output_item_id
|
||||
current_response_id = tool_response_id
|
||||
current_delta_chunks = None
|
||||
current_delta_type = None
|
||||
verbose_logger.debug("Tool use event: %s (ID: %s)", tool_name, tool_call_id)
|
||||
|
||||
elif "promptEnd" in event or "completionEnd" in event:
|
||||
|
|
@ -1333,6 +1332,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
current_delta_type,
|
||||
) = self.transform_prompt_end_event(event, current_response_id, current_conversation_id)
|
||||
returned_messages.extend(events)
|
||||
current_delta_chunks = None
|
||||
|
||||
return {
|
||||
"response": returned_messages,
|
||||
|
|
|
|||
|
|
@ -1090,6 +1090,162 @@ class TestBedrockRealtimeSessionEvents:
|
|||
assert event["session"]["modalities"] == ["text", "audio"]
|
||||
|
||||
|
||||
class TestBedrockRealtimeContentBlockLifecycle:
|
||||
"""
|
||||
Bedrock streams discrete content blocks. Session state must follow block
|
||||
boundaries so text/audio/tool blocks cannot leak into each other.
|
||||
"""
|
||||
|
||||
def _state(self, **overrides):
|
||||
base = {
|
||||
"session_configuration_request": json.dumps({"configured": True}),
|
||||
"current_output_item_id": None,
|
||||
"current_response_id": None,
|
||||
"current_conversation_id": "conv_1",
|
||||
"current_delta_chunks": None,
|
||||
"current_item_chunks": [],
|
||||
"current_delta_type": None,
|
||||
}
|
||||
base.update(overrides)
|
||||
return base
|
||||
|
||||
def _apply(self, config, logging_obj, state, message):
|
||||
result = config.transform_realtime_response(
|
||||
json.dumps(message),
|
||||
"amazon.nova-2-sonic-v1:0",
|
||||
logging_obj,
|
||||
realtime_response_transform_input=state,
|
||||
)
|
||||
state.update(
|
||||
{
|
||||
"current_output_item_id": result["current_output_item_id"],
|
||||
"current_response_id": result["current_response_id"],
|
||||
"current_conversation_id": result["current_conversation_id"],
|
||||
"current_delta_chunks": result["current_delta_chunks"],
|
||||
"current_item_chunks": result["current_item_chunks"],
|
||||
"current_delta_type": result["current_delta_type"],
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
def test_tool_block_does_not_leak_prior_text_into_next_assistant_turn(self):
|
||||
config = BedrockRealtimeConfig()
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_trace_id = "trace_123"
|
||||
state = self._state()
|
||||
|
||||
self._apply(config, logging_obj, state, {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}})
|
||||
first_response_id = state["current_response_id"]
|
||||
self._apply(
|
||||
config,
|
||||
logging_obj,
|
||||
state,
|
||||
{"event": {"textOutput": {"content": "I will check the weather."}}},
|
||||
)
|
||||
assert state["current_delta_chunks"] is not None
|
||||
assert len(state["current_delta_chunks"]) == 1
|
||||
|
||||
self._apply(
|
||||
config,
|
||||
logging_obj,
|
||||
state,
|
||||
{"event": {"contentEnd": {"stopReason": "PARTIAL_TURN", "type": "TEXT"}}},
|
||||
)
|
||||
assert state["current_delta_chunks"] is None
|
||||
assert state["current_delta_type"] is None
|
||||
assert state["current_output_item_id"] is None
|
||||
assert state["current_response_id"] == first_response_id
|
||||
|
||||
self._apply(config, logging_obj, state, {"event": {"contentStart": {"role": "TOOL", "type": "TOOL"}}})
|
||||
assert state["current_delta_chunks"] is None
|
||||
assert state["current_response_id"] == first_response_id
|
||||
|
||||
tool_result = self._apply(
|
||||
config,
|
||||
logging_obj,
|
||||
state,
|
||||
{
|
||||
"event": {
|
||||
"toolUse": {
|
||||
"toolUseId": "tool_1",
|
||||
"toolName": "get_weather",
|
||||
"content": json.dumps({"location": "Seattle"}),
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
assert tool_result["response"][0]["type"] == "response.function_call_arguments.done"
|
||||
assert state["current_delta_chunks"] is None
|
||||
assert state["current_delta_type"] is None
|
||||
|
||||
self._apply(
|
||||
config,
|
||||
logging_obj,
|
||||
state,
|
||||
{"event": {"contentEnd": {"stopReason": "TOOL_USE", "type": "TOOL"}}},
|
||||
)
|
||||
assert state["current_response_id"] is None
|
||||
assert state["current_output_item_id"] is None
|
||||
assert state["current_delta_chunks"] is None
|
||||
assert state["current_delta_type"] is None
|
||||
|
||||
post_tool = self._apply(
|
||||
config,
|
||||
logging_obj,
|
||||
state,
|
||||
{"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}},
|
||||
)
|
||||
assert state["current_response_id"] != first_response_id
|
||||
assert state["current_delta_chunks"] is None
|
||||
assert [msg["type"] for msg in post_tool["response"]].count("response.created") == 1
|
||||
|
||||
self._apply(
|
||||
config,
|
||||
logging_obj,
|
||||
state,
|
||||
{"event": {"textOutput": {"content": "It is sunny in Seattle."}}},
|
||||
)
|
||||
done = self._apply(
|
||||
config,
|
||||
logging_obj,
|
||||
state,
|
||||
{"event": {"contentEnd": {"stopReason": "END_TURN", "type": "TEXT"}}},
|
||||
)
|
||||
text_done = [msg for msg in done["response"] if msg["type"] == "response.text.done"][0]
|
||||
assert text_done["text"] == "It is sunny in Seattle."
|
||||
assert "I will check the weather." not in text_done["text"]
|
||||
assert any(msg["type"] == "response.done" for msg in done["response"])
|
||||
assert state["current_response_id"] is None
|
||||
assert state["current_delta_chunks"] is None
|
||||
|
||||
def test_second_assistant_content_block_reuses_response_not_item(self):
|
||||
config = BedrockRealtimeConfig()
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_trace_id = "trace_123"
|
||||
state = self._state()
|
||||
|
||||
first = self._apply(
|
||||
config, logging_obj, state, {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}
|
||||
)
|
||||
response_id = state["current_response_id"]
|
||||
first_item = state["current_output_item_id"]
|
||||
assert sum(1 for msg in first["response"] if msg["type"] == "response.created") == 1
|
||||
|
||||
self._apply(
|
||||
config,
|
||||
logging_obj,
|
||||
state,
|
||||
{"event": {"contentEnd": {"stopReason": "PARTIAL_TURN", "type": "TEXT"}}},
|
||||
)
|
||||
second = self._apply(
|
||||
config, logging_obj, state, {"event": {"contentStart": {"role": "ASSISTANT", "type": "AUDIO"}}}
|
||||
)
|
||||
assert state["current_response_id"] == response_id
|
||||
assert state["current_output_item_id"] != first_item
|
||||
assert sum(1 for msg in second["response"] if msg["type"] == "response.created") == 0
|
||||
assert sum(1 for msg in second["response"] if msg["type"] == "response.output_item.added") == 1
|
||||
|
||||
|
||||
class TestBedrockRealtimeUsageAccounting:
|
||||
def _usage_event(
|
||||
self,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue