fix: Change instruction insertion to respect tool call/response pairs

Make sure _add_instructions_to_user_content skips over user messages that carry function_response parts so tool_use/tool_result blocks stay together

Close #3229

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 826076141
This commit is contained in:
George Weale
2025-10-30 09:54:49 -07:00
committed by Copybara-Service
parent d3796f9b33
commit 1e6a9daa63
2 changed files with 50 additions and 4 deletions
+15 -4
View File
@@ -668,6 +668,16 @@ def _is_live_model_audio_event(event: Event) -> bool:
return False
def _content_contains_function_response(content: types.Content) -> bool:
"""Checks whether the content includes any function response parts."""
if not content.parts:
return False
for part in content.parts:
if part.function_response:
return True
return False
async def _add_instructions_to_user_content(
invocation_context: InvocationContext,
llm_request: LlmRequest,
@@ -695,13 +705,14 @@ async def _add_instructions_to_user_content(
if llm_request.contents:
for i in range(len(llm_request.contents) - 1, -1, -1):
if llm_request.contents[i].role != 'user':
content = llm_request.contents[i]
if content.role != 'user':
insert_index = i + 1
break
elif i == 0:
# All content from start is user content
insert_index = 0
if _content_contains_function_response(content):
insert_index = i + 1
break
insert_index = i
else:
# No contents remaining, just append at the end
insert_index = 0
@@ -923,6 +923,41 @@ async def test_no_dynamic_instructions_when_no_static(llm_backend):
assert llm_request.contents[0].parts[0].text == "Hello world"
@pytest.mark.asyncio
async def test_instructions_insert_after_function_response():
"""Ensure instruction insertion does not split tool_use/tool_result pairs."""
agent = LlmAgent(name="test_agent")
invocation_context = await _create_invocation_context(agent)
tool_call = types.Part.from_function_call(
name="echo_tool", args={"echo": "value"}
)
tool_response = types.Part.from_function_response(
name="echo_tool", response={"result": "value"}
)
llm_request = LlmRequest(
contents=[
types.Content(role="assistant", parts=[tool_call]),
types.Content(role="user", parts=[tool_response]),
]
)
instruction_contents = [
types.Content(
role="user", parts=[types.Part.from_text(text="Dynamic instruction")]
)
]
await _add_instructions_to_user_content(
invocation_context, llm_request, instruction_contents
)
assert len(llm_request.contents) == 3
assert llm_request.contents[0].parts[0].function_call
assert llm_request.contents[1].parts[0].function_response
assert llm_request.contents[2].parts[0].text == "Dynamic instruction"
@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI", "VERTEX"])
@pytest.mark.asyncio
async def test_static_instruction_with_files_and_text(llm_backend):