diff --git a/e2e_test_index.py b/e2e_test_index.py new file mode 100644 index 00000000000..60ae7bf11d6 --- /dev/null +++ b/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() diff --git a/implementation_plan.md b/implementation_plan.md new file mode 100644 index 00000000000..2f816dbf3c0 --- /dev/null +++ b/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. diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index a27fc1ded03..7574b21c049 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/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.""" diff --git a/src/google/adk/agents/remote_a2a_agent.py b/src/google/adk/agents/remote_a2a_agent.py index 313aa625d39..8684e0ca44f 100644 --- a/src/google/adk/agents/remote_a2a_agent.py +++ b/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 diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index adfd7e84e8f..36b0de8688d 100644 --- a/src/google/adk/runners.py +++ b/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 ) diff --git a/src/google/adk/sessions/session.py b/src/google/adk/sessions/session.py index dab5476ce31..75bf1d608d0 100644 --- a/src/google/adk/sessions/session.py +++ b/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) diff --git a/tests/unittests/agents/test_invocation_context.py b/tests/unittests/agents/test_invocation_context.py index a7bfd87bd2f..a01340ce912 100644 --- a/tests/unittests/agents/test_invocation_context.py +++ b/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 diff --git a/tests/unittests/agents/test_remote_a2a_agent.py b/tests/unittests/agents/test_remote_a2a_agent.py index 8caa290983d..24a4bd0e294 100644 --- a/tests/unittests/agents/test_remote_a2a_agent.py +++ b/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 diff --git a/tests/unittests/sessions/test_session.py b/tests/unittests/sessions/test_session.py new file mode 100644 index 00000000000..b0194f32140 --- /dev/null +++ b/tests/unittests/sessions/test_session.py @@ -0,0 +1,114 @@ +from typing import Any +import pytest +from google.genai import types +from google.adk.events.event import Event +from google.adk.sessions.session import Session + + +def test_function_call_index_incremental(): + """Tests that function calls are incrementally indexed.""" + session = Session(id="1", app_name="test", user_id="user") + + # 1. Add event with function call + fc = types.FunctionCall(name='test_tool', args={}) + fc.id = "call_1" + event1 = Event( + invocation_id="inv1", + author="agent", + content=types.Content(parts=[types.Part(function_call=fc)], role="model") + ) + session.events.append(event1) + + assert session.function_call_index == {"call_1": event1} + assert session._indexed_events_count == 1 + + # 2. Add another event + fc2 = types.FunctionCall(name='test_tool', args={}) + fc2.id = "call_2" + event2 = Event( + invocation_id="inv1", + author="agent", + content=types.Content(parts=[types.Part(function_call=fc2)], role="model") + ) + session.events.append(event2) + + assert session.function_call_index == {"call_1": event1, "call_2": event2} + assert session._indexed_events_count == 2 + +def test_function_call_index_truncation(): + """Tests that the index rebuilds when events are truncated.""" + session = Session(id="1", app_name="test", user_id="user") + + fc = types.FunctionCall(name='test_tool', args={}) + fc.id = "call_1" + event1 = Event( + invocation_id="inv1", + author="agent", + content=types.Content(parts=[types.Part(function_call=fc)], role="model") + ) + session.events.append(event1) + + # Add a second event + fc2 = types.FunctionCall(name='test_tool', args={}) + fc2.id = "call_2" + event2 = Event( + invocation_id="inv1", + author="agent", + content=types.Content(parts=[types.Part(function_call=fc2)], role="model") + ) + session.events.append(event2) + + # Access index to cache it + _ = session.function_call_index + assert session._indexed_events_count == 2 + + # Remove one event to simulate truncation + session.events.pop() + + # Index should rebuild for remaining events + assert session.function_call_index == {"call_1": event1} + assert session._indexed_events_count == 1 + +def test_get_matching_function_call(): + """Tests getting a matching function call based on a function response.""" + session = Session(id="1", app_name="test", user_id="user") + + # Test empty events + assert session.get_matching_function_call() is None + + # 1. Add event with function call + fc = types.FunctionCall(name='test_tool', args={}) + fc.id = "call_1" + event1 = Event( + invocation_id="inv1", + author="agent", + content=types.Content(parts=[types.Part(function_call=fc)], role="model") + ) + session.events.append(event1) + + # Test last event has no function response + assert session.get_matching_function_call() is None + + # 2. Add event with function response + fr = types.FunctionResponse(name='test_tool', response={"status": "ok"}) + fr.id = "call_1" + event2 = Event( + invocation_id="inv1", + author="user", + content=types.Content(parts=[types.Part(function_response=fr)], role="user") + ) + session.events.append(event2) + + # Test it finds the function call + assert session.get_matching_function_call() == event1 + + # 3. Add response with no ID + fr2 = types.FunctionResponse(name='test_tool', response={"status": "ok"}) + fr2.id = "" + event3 = Event( + invocation_id="inv1", + author="user", + content=types.Content(parts=[types.Part(function_response=fr2)], role="user") + ) + session.events.append(event3) + assert session.get_matching_function_call() is None