fix: Conditionally preserve function call IDs in LLM requests

Function call and response IDs generated by ADK are now preserved in the LLM request contents when the agent is using a Gemini model with `use_interactions_api` enabled

Close #4381

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 867675945
This commit is contained in:
George Weale
2026-02-09 10:46:55 -08:00
committed by Copybara-Service
parent 64a44c2897
commit 663cb75b32
2 changed files with 167 additions and 4 deletions
+29 -3
View File
@@ -42,8 +42,16 @@ class _ContentLlmRequestProcessor(BaseLlmRequestProcessor):
self, invocation_context: InvocationContext, llm_request: LlmRequest
) -> AsyncGenerator[Event, None]:
from ...agents.llm_agent import LlmAgent
from ...models.google_llm import Gemini
agent = invocation_context.agent
preserve_function_call_ids = False
if isinstance(agent, LlmAgent):
canonical_model = agent.canonical_model
preserve_function_call_ids = (
isinstance(canonical_model, Gemini)
and canonical_model.use_interactions_api
)
# Preserve all contents that were added by instruction processor
# (since llm_request.contents will be completely reassigned below)
@@ -55,6 +63,7 @@ class _ContentLlmRequestProcessor(BaseLlmRequestProcessor):
invocation_context.branch,
invocation_context.session.events,
agent.name,
preserve_function_call_ids=preserve_function_call_ids,
)
else:
# Include current turn context only (no conversation history)
@@ -62,6 +71,7 @@ class _ContentLlmRequestProcessor(BaseLlmRequestProcessor):
invocation_context.branch,
invocation_context.session.events,
agent.name,
preserve_function_call_ids=preserve_function_call_ids,
)
# Add instruction-related contents to proper position in conversation
@@ -360,7 +370,11 @@ def _process_compaction_events(events: list[Event]) -> list[Event]:
def _get_contents(
current_branch: Optional[str], events: list[Event], agent_name: str = ''
current_branch: Optional[str],
events: list[Event],
agent_name: str = '',
*,
preserve_function_call_ids: bool = False,
) -> list[types.Content]:
"""Get the contents for the LLM request.
@@ -370,6 +384,7 @@ def _get_contents(
current_branch: The current branch of the agent.
events: Events to process.
agent_name: The name of the agent.
preserve_function_call_ids: Whether to preserve function call ids.
Returns:
A list of processed contents.
@@ -469,13 +484,18 @@ def _get_contents(
for event in result_events:
content = copy.deepcopy(event.content)
if content:
if not preserve_function_call_ids:
remove_client_function_call_id(content)
contents.append(content)
return contents
def _get_current_turn_contents(
current_branch: Optional[str], events: list[Event], agent_name: str = ''
current_branch: Optional[str],
events: list[Event],
agent_name: str = '',
*,
preserve_function_call_ids: bool = False,
) -> list[types.Content]:
"""Get contents for the current turn only (no conversation history).
@@ -491,6 +511,7 @@ def _get_current_turn_contents(
current_branch: The current branch of the agent.
events: A list of all session events.
agent_name: The name of the agent.
preserve_function_call_ids: Whether to preserve function call ids.
Returns:
A list of contents for the current turn only, preserving context needed
@@ -502,7 +523,12 @@ def _get_current_turn_contents(
if _should_include_event_in_context(current_branch, event) and (
event.author == 'user' or _is_other_agent_reply(agent_name, event)
):
return _get_contents(current_branch, events[i:], agent_name)
return _get_contents(
current_branch,
events[i:],
agent_name,
preserve_function_call_ids=preserve_function_call_ids,
)
return []
@@ -19,6 +19,7 @@ from google.adk.flows.llm_flows import contents
from google.adk.flows.llm_flows.contents import request_processor
from google.adk.flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
from google.adk.flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME
from google.adk.models.google_llm import Gemini
from google.adk.models.llm_request import LlmRequest
from google.genai import types
import pytest
@@ -931,3 +932,139 @@ async def test_function_response_with_thought_not_filtered():
fr_parts = [p for p in fr_content.parts if p.function_response]
assert len(fr_parts) == 1
assert fr_parts[0].function_response.name == "calc_tool"
@pytest.mark.asyncio
async def test_adk_function_call_ids_are_stripped_for_non_interactions_model():
"""Test ADK generated ids are removed for non-interactions requests."""
agent = Agent(model="gemini-2.5-flash", name="test_agent")
llm_request = LlmRequest(model="gemini-2.5-flash")
invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
function_call_id = "adk-test-call-id"
events = [
Event(
invocation_id="inv1",
author="user",
content=types.UserContent("Call the tool"),
),
Event(
invocation_id="inv2",
author="test_agent",
content=types.Content(
role="model",
parts=[
types.Part(
function_call=types.FunctionCall(
id=function_call_id,
name="test_tool",
args={"x": 1},
)
)
],
),
),
Event(
invocation_id="inv3",
author="test_agent",
content=types.Content(
role="user",
parts=[
types.Part(
function_response=types.FunctionResponse(
id=function_call_id,
name="test_tool",
response={"result": 2},
)
)
],
),
),
]
invocation_context.session.events = events
async for _ in contents.request_processor.run_async(
invocation_context, llm_request
):
pass
model_fc_part = llm_request.contents[1].parts[0]
assert model_fc_part.function_call is not None
assert model_fc_part.function_call.id is None
user_fr_part = llm_request.contents[2].parts[0]
assert user_fr_part.function_response is not None
assert user_fr_part.function_response.id is None
@pytest.mark.asyncio
async def test_adk_function_call_ids_preserved_for_interactions_model():
"""Test ADK generated ids are preserved for interactions requests."""
agent = Agent(
model=Gemini(
model="gemini-2.5-flash",
use_interactions_api=True,
),
name="test_agent",
)
llm_request = LlmRequest(model="gemini-2.5-flash")
invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
function_call_id = "adk-test-call-id"
events = [
Event(
invocation_id="inv1",
author="user",
content=types.UserContent("Call the tool"),
),
Event(
invocation_id="inv2",
author="test_agent",
content=types.Content(
role="model",
parts=[
types.Part(
function_call=types.FunctionCall(
id=function_call_id,
name="test_tool",
args={"x": 1},
)
)
],
),
),
Event(
invocation_id="inv3",
author="test_agent",
content=types.Content(
role="user",
parts=[
types.Part(
function_response=types.FunctionResponse(
id=function_call_id,
name="test_tool",
response={"result": 2},
)
)
],
),
),
]
invocation_context.session.events = events
async for _ in contents.request_processor.run_async(
invocation_context, llm_request
):
pass
model_fc_part = llm_request.contents[1].parts[0]
assert model_fc_part.function_call is not None
assert model_fc_part.function_call.id == function_call_id
user_fr_part = llm_request.contents[2].parts[0]
assert user_fr_part.function_response is not None
assert user_fr_part.function_response.id == function_call_id