feat: add support for ContetxtWindowCompressionConfig in RunConfig

Merge https://github.com/google/adk-python/pull/2206

### Summary

This PR adds support for `ContextWindowCompressionConfig` in `RunConfig`.
This enables context window compression using a `trigger_tokens` threshold and a sliding window with a `target_tokens` limit.

This feature is useful for managing long-running audio inputs.

### Related Issue

Closes #2188

### Testing Plan

- Added new unit test: `test_streaming_with_context_window_compression_config`

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2206 from ac-machache:support/add-context-compression-config c8a5b15cae2d2b72f331797d07ae0bbaf977ed3c
PiperOrigin-RevId: 819855786
This commit is contained in:
machache
2025-10-15 12:00:21 -07:00
committed by Copybara-Service
parent 78e74b5bf2
commit b650181384
3 changed files with 64 additions and 0 deletions
+5
View File
@@ -94,6 +94,11 @@ class RunConfig(BaseModel):
session_resumption: Optional[types.SessionResumptionConfig] = None
"""Configures session resumption mechanism. Only support transparent session resumption mode now."""
context_window_compression: Optional[types.ContextWindowCompressionConfig] = (
None
)
"""Configuration for context window compression. If set, this will enable context window compression for LLM input."""
save_live_audio: bool = False
"""Saves live video and audio data to session and artifact service.
+3
View File
@@ -79,6 +79,9 @@ class _BasicLlmRequestProcessor(BaseLlmRequestProcessor):
llm_request.live_connect_config.session_resumption = (
invocation_context.run_config.session_resumption
)
llm_request.live_connect_config.context_window_compression = (
invocation_context.run_config.context_window_compression
)
# TODO: handle tool append here, instead of in BaseTool.process_llm_request.
@@ -586,3 +586,59 @@ def test_streaming_with_session_resumption_config():
llm_request_sent_to_mock.live_connect_config.session_resumption.transparent
is True
)
def test_streaming_with_context_window_compression_config():
"""Test streaming with context window compression config."""
response = LlmResponse(turn_complete=True)
mock_model = testing_utils.MockModel.create([response])
root_agent = Agent(
name='root_agent',
model=mock_model,
tools=[],
)
runner = testing_utils.InMemoryRunner(
root_agent=root_agent, response_modalities=['AUDIO']
)
# Create run config with context window compression
run_config = RunConfig(
context_window_compression=types.ContextWindowCompressionConfig(
trigger_tokens=1000,
sliding_window=types.SlidingWindow(target_tokens=500),
)
)
live_request_queue = LiveRequestQueue()
live_request_queue.send_realtime(
blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm')
)
res_events = runner.run_live(live_request_queue, run_config)
assert res_events is not None, 'Expected a list of events, got None.'
assert (
len(res_events) > 0
), 'Expected at least one response, but got an empty list.'
assert len(mock_model.requests) == 1
# Get the request that was captured
llm_request_sent_to_mock = mock_model.requests[0]
# Assert that the request contained the correct configuration
assert llm_request_sent_to_mock.live_connect_config is not None
assert (
llm_request_sent_to_mock.live_connect_config.context_window_compression
is not None
)
assert (
llm_request_sent_to_mock.live_connect_config.context_window_compression.trigger_tokens
== 1000
)
assert (
llm_request_sent_to_mock.live_connect_config.context_window_compression.sliding_window.target_tokens
== 500
)