Merge branch 'main' into feature/chroma-memory-service

Author: ayman3000

src/google/adk/flows/llm_flows/contents.py

@@ -42,8 +42,16 @@ async def run_async(
       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,13 +63,15 @@ async def run_async(
           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)
       llm_request.contents = _get_current_turn_contents(
           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:
-      remove_client_function_call_id(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 []
 

src/google/adk/memory/vertex_ai_memory_bank_service.py

@@ -27,6 +27,8 @@
 from .memory_entry import MemoryEntry
 
 if TYPE_CHECKING:
+  import vertexai
+
   from ..sessions.session import Session
 
 logger = logging.getLogger('google_adk.' + __name__)
@@ -88,8 +90,8 @@ async def add_session_to_memory(self, session: Session):
             'content': event.content.model_dump(exclude_none=True, mode='json')
         })
     if events:
-      client = self._get_api_client()
-      operation = client.agent_engines.memories.generate(
+      api_client = self._get_api_client()
+      operation = await api_client.agent_engines.memories.generate(
           name='reasoningEngines/' + self._agent_engine_id,
           direct_contents_source={'events': events},
           scope={
@@ -108,22 +110,24 @@ async def search_memory(self, *, app_name: str, user_id: str, query: str):
     if not self._agent_engine_id:
       raise ValueError('Agent Engine ID is required for Memory Bank.')
 
-    client = self._get_api_client()
-    retrieved_memories_iterator = client.agent_engines.memories.retrieve(
-        name='reasoningEngines/' + self._agent_engine_id,
-        scope={
-            'app_name': app_name,
-            'user_id': user_id,
-        },
-        similarity_search_params={
-            'search_query': query,
-        },
+    api_client = self._get_api_client()
+    retrieved_memories_iterator = (
+        await api_client.agent_engines.memories.retrieve(
+            name='reasoningEngines/' + self._agent_engine_id,
+            scope={
+                'app_name': app_name,
+                'user_id': user_id,
+            },
+            similarity_search_params={
+                'search_query': query,
+            },
+        )
     )
 
     logger.info('Search memory response received.')
 
-    memory_events = []
-    for retrieved_memory in retrieved_memories_iterator:
+    memory_events: list[MemoryEntry] = []
+    async for retrieved_memory in retrieved_memories_iterator:
       # TODO: add more complex error handling
       logger.debug('Retrieved memory: %s', retrieved_memory)
       memory_events.append(
@@ -138,21 +142,22 @@ async def search_memory(self, *, app_name: str, user_id: str, query: str):
       )
     return SearchMemoryResponse(memories=memory_events)
 
-  def _get_api_client(self):
+  def _get_api_client(self) -> vertexai.AsyncClient:
     """Instantiates an API client for the given project and location.
 
     It needs to be instantiated inside each request so that the event loop
     management can be properly propagated.
     Returns:
-      An API client for the given project and location or express mode api key.
+      An async API client for the given project and location or express mode api
+      key.
     """
     import vertexai
 
     return vertexai.Client(
         project=self._project,
         location=self._location,
         api_key=self._express_mode_api_key,
-    )
+    ).aio
 
 
 def _should_filter_out_event(content: types.Content) -> bool:

tests/unittests/flows/llm_flows/test_contents.py

@@ -19,6 +19,7 @@
 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