fix: N sized sliding window

This fix allows for more efficient interaction with Context Caching, as Context Windows are removed when they exceed N, rather than immediately when they exceed the desired number.

Adapted to the new baseline (using invocation_start_indices instead of num_model_turns).

Merges #3271

Co-authored-by: Shangjie Chen <deanchen@google.com>
COPYBARA_INTEGRATE_REVIEW=#3271 from UlookEE:n_sized_sliding_window 25429aa
PiperOrigin-RevId: 936165882

Author: 2 people

src/google/adk/plugins/context_filter_plugin.py

@@ -107,6 +107,7 @@ def __init__(
           Callable[[list[types.Content]], list[types.Content]]
       ] = None,
       name: str = "context_filter_plugin",
+      remove_amount: int = 1,
   ):
     """Initializes the context management plugin.
 
@@ -117,10 +118,15 @@ def __init__(
         message starts a new invocation.
       custom_filter: A function to filter the context.
       name: The name of the plugin instance.
+      remove_amount: The number of invocations to remove when the context
+        exceeds the limit.
     """
+    if remove_amount < 1:
+      raise ValueError("remove_amount must be at least 1")
     super().__init__(name)
     self._num_invocations_to_keep = num_invocations_to_keep
     self._custom_filter = custom_filter
+    self._remove_amount = remove_amount
 
   async def before_model_callback(
       self, *, callback_context: CallbackContext, llm_request: LlmRequest
@@ -134,7 +140,10 @@ async def before_model_callback(
           and self._num_invocations_to_keep > 0
       ):
         invocation_start_indices = _get_invocation_start_indices(contents)
-        if len(invocation_start_indices) > self._num_invocations_to_keep:
+        if (
+            len(invocation_start_indices)
+            >= self._num_invocations_to_keep + self._remove_amount
+        ):
           split_index = invocation_start_indices[-self._num_invocations_to_keep]
 
           # Adjust split_index to avoid orphaned function_responses.

tests/unittests/plugins/test_context_filtering_plugin.py

@@ -343,3 +343,126 @@ async def test_last_invocation_with_tool_call_keeps_user_prompt():
 
   assert "user_prompt_2" in texts
   assert "final_answer_2" in texts
+
+
+@pytest.mark.asyncio
+async def test_filter_with_remove_amount():
+  """Tests that remove_amount correctly removes additional invocations."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=2, remove_amount=1)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+      _create_content("user", "user_prompt_3"),
+      _create_content("model", "model_response_3"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # With num_invocations_to_keep=2 and remove_amount=1, keeps last 2.
+  assert len(llm_request.contents) == 4
+  assert llm_request.contents[0].parts[0].text == "user_prompt_2"
+  assert llm_request.contents[1].parts[0].text == "model_response_2"
+  assert llm_request.contents[2].parts[0].text == "user_prompt_3"
+  assert llm_request.contents[3].parts[0].text == "model_response_3"
+
+
+@pytest.mark.asyncio
+async def test_filter_with_higher_remove_amount():
+  """Tests remove_amount with a higher value to remove more invocations."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=3, remove_amount=2)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+      _create_content("user", "user_prompt_3"),
+      _create_content("model", "model_response_3"),
+      _create_content("user", "user_prompt_4"),
+      _create_content("model", "model_response_4"),
+      _create_content("user", "user_prompt_5"),
+      _create_content("model", "model_response_5"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # With num_invocations_to_keep=3 and remove_amount=2, keeps last 3.
+  assert len(llm_request.contents) == 6
+  assert llm_request.contents[0].parts[0].text == "user_prompt_3"
+  assert llm_request.contents[1].parts[0].text == "model_response_3"
+  assert llm_request.contents[2].parts[0].text == "user_prompt_4"
+  assert llm_request.contents[3].parts[0].text == "model_response_4"
+  assert llm_request.contents[4].parts[0].text == "user_prompt_5"
+  assert llm_request.contents[5].parts[0].text == "model_response_5"
+
+
+def test_invalid_remove_amount():
+  """Tests that initializing with remove_amount < 1 raises ValueError."""
+  with pytest.raises(ValueError, match="remove_amount must be at least 1"):
+    ContextFilterPlugin(num_invocations_to_keep=1, remove_amount=0)
+
+  with pytest.raises(ValueError, match="remove_amount must be at least 1"):
+    ContextFilterPlugin(num_invocations_to_keep=1, remove_amount=-1)
+
+
+@pytest.mark.asyncio
+async def test_filter_remove_amount_with_multiple_user_turns():
+  """Tests remove_amount with multiple user turns in invocations."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=2, remove_amount=1)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2a"),
+      _create_content("user", "user_prompt_2b"),
+      _create_content("model", "model_response_2"),
+      _create_content("user", "user_prompt_3"),
+      _create_content("model", "model_response_3"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # Should keep last 2 invocations including multiple user turns
+  assert len(llm_request.contents) == 5
+  assert llm_request.contents[0].parts[0].text == "user_prompt_2a"
+  assert llm_request.contents[1].parts[0].text == "user_prompt_2b"
+  assert llm_request.contents[2].parts[0].text == "model_response_2"
+  assert llm_request.contents[3].parts[0].text == "user_prompt_3"
+  assert llm_request.contents[4].parts[0].text == "model_response_3"
+
+
+@pytest.mark.asyncio
+async def test_filter_bypass_when_under_remove_threshold():
+  """Tests that filtering is bypassed when total invocations are between keep limit and keep+remove limit."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=2, remove_amount=2)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+      _create_content("user", "user_prompt_3"),
+      _create_content("model", "model_response_3"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+  original_contents = list(llm_request.contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # With num_invocations_to_keep=2 and remove_amount=2, threshold is 4.
+  # We have 3 invocations, so no filtering should occur.
+  assert llm_request.contents == original_contents