feat: O(1) function call lookups

e2e_test_index.py

@@ -0,0 +1,67 @@
+import time
+from typing import List
+
+from google.genai import types
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events.event import Event
+from google.adk.sessions.session import Session
+from unittest.mock import Mock
+
+
+def main():
+    print("Running E2E performance test for function call index...")
+    num_events = 10000
+
+    events: List[Event] = []
+
+    # 1. Create 10000 events, each with a function call
+    start_time = time.perf_counter()
+    for i in range(num_events):
+        fc = types.FunctionCall(name='some_tool', args={})
+        fc.id = f'call_{i}'
+        fc_part = types.Part(function_call=fc)
+
+        event = Event(
+            invocation_id='inv_1',
+            author='dummy_agent',
+            content=types.Content(parts=[fc_part], role='model'),
+        )
+        events.append(event)
+
+    print(f"Created {num_events} events in {time.perf_counter() - start_time:.4f}s")
+
+    session = Session(id="test_session", app_name="test_app", user_id="test_user", events=events)
+
+    # Target function calls to find
+    targets = [f'call_{0}', f'call_{num_events // 2}', f'call_{num_events - 1}']
+
+    from google.adk.flows.llm_flows.functions import find_matching_function_call
+
+    for target_id in targets:
+        fr = types.FunctionResponse(name='some_tool', response={'result': 'ok'})
+        fr.id = target_id
+        fr_part = types.Part(function_response=fr)
+        fr_event = Event(
+            invocation_id='inv_1',
+            author='dummy_agent',
+            content=types.Content(parts=[fr_part], role='user'),
+        )
+
+        session.events.append(fr_event)
+
+        t0 = time.perf_counter()
+        match = find_matching_function_call(session.events, session.function_call_index)
+        t1 = time.perf_counter()
+
+        assert match is not None, f"Match not found for {target_id}"
+        found_fc = match.get_function_calls()[0]
+        assert found_fc.id == target_id, f"Mismatched ID. Expected {target_id}, got {found_fc.id}"
+        print(f"Found {target_id} in {t1 - t0:.6f}s")
+
+
+    print("Success! Performance scales phenomenally.")
+
+
+if __name__ == '__main__':
+    main()

implementation_plan.md

@@ -0,0 +1,63 @@
+## Task
+
+### User intent with respect to ADK
+Refactor Python's `find_matching_function_call` and `find_event_by_function_call_id` functions to use an index for O(1) lookups instead of O(N) backward searching. This brings Python's performance semantics to parity with an expected `adk-js` refactor.
+
+### Feature Description
+Currently, `find_event_by_function_call_id` in `adk-python/src/google/adk/flows/llm_flows/functions.py` iterates backwards over the `events` list to find the matching function call for a function response. 
+The plan is to introduce a mapped index (e.g., `dict[str, Event]` where the key is the `function_call_id`) that is maintained as events are added via `InvocationContext` and `Session`, allowing `find_event_by_function_call_id` to perform lookup in O(1) time.
+
+### Use Cases & Examples
+- Long-running sessions with many agent iterations and function calls currently incur O(N) lookup penalties for every function response. This refactoring ensures constant time resolution regardless of session scale.
+
+## Context
+
+### ADK Context
+- Documentation context: `InvocationContext` and runners currently assume a list of `events` is enough, but maintaining state efficiently is paramount.
+- Reference context: Parity request with `adk-js` implies aligning the core function logic so it doesn't do a manual iterative search through `events` backwards.
+- General context: Functions to update include `find_matching_function_call` and `find_event_by_function_call_id`. These are utilized in `remote_a2a_agent.py`, `llm_agent.py`, `invocation_context.py`, and `runners.py`.
+
+### Language Specific Context
+- Target language: Python
+- Target repo: `adk-python`
+- General context: Python `dict` provides O(1) time complexity for fetching values by key.
+
+## Definition
+
+### Data Models 
+Add an internal index `_function_call_index: dict[str, Event]` parameter (or context state variable) that maps `function_call.id` to its parent `Event`.
+
+### Inputs
+Function signatures for `find_matching_function_call` and `find_event_by_function_call_id` will need to accept the index (or access it transparently via context classes if moved into them).
+
+### Outputs
+Returns `Optional[Event]` in O(1).
+
+### Side Effects
+The indexing dictionary needs to be kept in sync with `events` arrays.
+
+## Constraints
+
+### Invariants
+The dictionary must contain keys for all function calls present in the session history.
+
+### Preconditions
+When an event containing `function_calls` is processed or appended, the index dictionary must be updated.
+
+### Postconditions
+The lookup function matches the same event it successfully matched previously but without looping.
+
+### Error Handling Protocols
+If a `function_call_id` is missing in the index, standard `None` returns applies.
+
+### Breaking Change Analysis
+This is an internal refactoring for efficiency; if `find_matching_function_call` signature changes, all internal call sites must be updated accordingly.
+
+### Testing
+
+- #### Unit tests with >=95% New Line Coverage
+Ensure existing tests in `test_invocation_context.py` and `test_functions_simple.py` continue to pass.
+- #### Integration tests
+Agent flows with multiple rounds of tool use (e.g., `test_remote_a2a_agent.py`) must be fully validated.
+- #### Manual e2e test
+A test with hundreds of tool calls confirming performance scaling.

src/google/adk/agents/invocation_context.py

@@ -543,21 +543,20 @@ def _find_matching_function_call(
       self, function_response_event: Event
   ) -> Optional[Event]:
     """Finds the function call event in the current invocation that matches the function response id."""
-    from ..flows.llm_flows.functions import find_event_by_function_call_id
-
     function_responses = function_response_event.get_function_responses()
     if not function_responses:
       return None
 
-    events = self._get_events(current_invocation=True)
-    if events and events[-1].id == function_response_event.id:
-      search_space = events[:-1]
-    else:
-      search_space = events
+    fc_event = self.session.function_call_index.get(function_responses[0].id)
+    if not fc_event and hasattr(self, '_local_state') and getattr(self._local_state, 'events', None):
+      search_session = self._generate_session_with_injected_local_events(
+          function_response_event
+      )
+      fc_event = search_session.function_call_index.get(function_responses[0].id)
 
-    return find_event_by_function_call_id(
-        search_space, function_responses[0].id
-    )
+    if fc_event and fc_event.invocation_id == self.invocation_id:
+      return fc_event
+    return None
 
   def stamp_event_branch_context(self, event: Event) -> None:
     """Stamps the event with the branch and isolation scope of its matching function call."""

src/google/adk/agents/remote_a2a_agent.py

@@ -67,7 +67,7 @@
 from ..events.event import Event
 from ..flows.llm_flows.contents import _is_other_agent_reply
 from ..flows.llm_flows.contents import _present_other_agent_message
-from ..flows.llm_flows.functions import find_matching_function_call
+
 from .base_agent import BaseAgent
 
 __all__ = [
@@ -356,7 +356,7 @@ def _create_a2a_request_for_user_function_response(
     """
     if not ctx.session.events or ctx.session.events[-1].author != "user":
       return None
-    function_call_event = find_matching_function_call(ctx.session.events)
+    function_call_event = ctx.session.get_matching_function_call()
     if not function_call_event:
       return None
 

src/google/adk/runners.py

@@ -47,8 +47,7 @@
 from .events.event import Event
 from .events.event import EventActions
 from .flows.llm_flows import contents
-from .flows.llm_flows.functions import find_event_by_function_call_id
-from .flows.llm_flows.functions import find_matching_function_call
+
 from .memory.base_memory_service import BaseMemoryService
 from .platform.thread import create_thread
 from .plugins.base_plugin import BasePlugin
@@ -439,9 +438,7 @@ def _resolve_invocation_id(
     if not function_responses:
       return invocation_id
 
-    fc_event = find_event_by_function_call_id(
-        session.events, function_responses[0].id
-    )
+    fc_event = session.function_call_index.get(function_responses[0].id)
     if not fc_event:
       raise ValueError(
           'Function call event not found for function response id:'
@@ -1768,7 +1765,7 @@ def _find_agent_to_run(
     # the agent that returned the corresponding function call regardless the
     # type of the agent. e.g. a remote a2a agent may surface a credential
     # request as a special long-running function tool call.
-    event = find_matching_function_call(session.events)
+    event = session.get_matching_function_call()
     is_resumable = (
         self.resumability_config and self.resumability_config.is_resumable
     )

src/google/adk/sessions/session.py

@@ -71,3 +71,32 @@ class Session(BaseModel):
 
   _storage_update_marker: str | None = PrivateAttr(default=None)
   """Internal storage revision marker used for stale-session detection."""
+
+  _function_call_index: dict[str, Event] = PrivateAttr(default_factory=dict)
+  _indexed_events_count: int = PrivateAttr(default=0)
+
+  @property
+  def function_call_index(self) -> dict[str, Event]:
+    current_length = len(self.events)
+    if self._indexed_events_count < current_length:
+      for event in self.events[self._indexed_events_count:]:
+        for fc in event.get_function_calls():
+          if fc.id:
+            self._function_call_index[fc.id] = event
+      self._indexed_events_count = current_length
+    elif self._indexed_events_count > current_length:
+      self._function_call_index.clear()
+      for event in self.events:
+        for fc in event.get_function_calls():
+          if fc.id:
+            self._function_call_index[fc.id] = event
+      self._indexed_events_count = current_length
+    return self._function_call_index
+
+  def get_matching_function_call(self) -> Event | None:
+    if not self.events:
+      return None
+    responses = self.events[-1].get_function_responses()
+    if not responses or not responses[0].id:
+      return None
+    return self.function_call_index.get(responses[0].id)

tests/unittests/agents/test_invocation_context.py

@@ -540,7 +540,7 @@ def _create_invocation_context(events):
           session_service=Mock(spec=BaseSessionService),
           agent=Mock(spec=BaseAgent, name='agent'),
           invocation_id='inv_1',
-          session=Mock(spec=Session, events=events),
+          session=Session(id="test", app_name="test", user_id="test", events=events),
       )
 
     return _create_invocation_context

tests/unittests/agents/test_remote_a2a_agent.py

@@ -624,16 +624,13 @@ def setup_method(self):
 
   def test_create_a2a_request_for_user_function_response_no_function_call(self):
     """Test function response request creation when no function call exists."""
-    with patch(
-        "google.adk.agents.remote_a2a_agent.find_matching_function_call"
-    ) as mock_find:
-      mock_find.return_value = None
+    self.mock_session.get_matching_function_call.return_value = None
 
-      result = self.agent._create_a2a_request_for_user_function_response(
-          self.mock_context
-      )
+    result = self.agent._create_a2a_request_for_user_function_response(
+        self.mock_context
+    )
 
-      assert result is None
+    assert result is None
 
   def test_create_a2a_request_for_user_function_response_success(self):
     """Test successful function response request creation."""
@@ -651,12 +648,9 @@ def test_create_a2a_request_for_user_function_response_success(self):
     mock_latest_event.author = "user"
     self.mock_session.events = [mock_latest_event]
 
-    with patch(
-        "google.adk.agents.remote_a2a_agent.find_matching_function_call"
-    ) as mock_find:
-      mock_find.return_value = mock_function_event
+    self.mock_session.get_matching_function_call.return_value = mock_function_event
 
-      with patch(
+    with patch(
           "google.adk.agents.remote_a2a_agent.convert_event_to_a2a_message"
       ) as mock_convert:
         # Create a proper mock A2A message
@@ -1530,16 +1524,13 @@ def setup_method(self):
 
   def test_create_a2a_request_for_user_function_response_no_function_call(self):
     """Test function response request creation when no function call exists."""
-    with patch(
-        "google.adk.agents.remote_a2a_agent.find_matching_function_call"
-    ) as mock_find:
-      mock_find.return_value = None
+    self.mock_session.get_matching_function_call.return_value = None
 
-      result = self.agent._create_a2a_request_for_user_function_response(
-          self.mock_context
-      )
+    result = self.agent._create_a2a_request_for_user_function_response(
+        self.mock_context
+    )
 
-      assert result is None
+    assert result is None
 
   def test_create_a2a_request_for_user_function_response_success(self):
     """Test successful function response request creation."""
@@ -1557,12 +1548,9 @@ def test_create_a2a_request_for_user_function_response_success(self):
     mock_latest_event.author = "user"
     self.mock_session.events = [mock_latest_event]
 
-    with patch(
-        "google.adk.agents.remote_a2a_agent.find_matching_function_call"
-    ) as mock_find:
-      mock_find.return_value = mock_function_event
+    self.mock_session.get_matching_function_call.return_value = mock_function_event
 
-      with patch(
+    with patch(
           "google.adk.agents.remote_a2a_agent.convert_event_to_a2a_message"
       ) as mock_convert:
         # Create a proper mock A2A message