Skip to content

Commit b650181

Browse files
ac-machachecopybara-github
authored andcommitted
feat: add support for ContetxtWindowCompressionConfig in RunConfig
Merge #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=#2206 from ac-machache:support/add-context-compression-config c8a5b15 PiperOrigin-RevId: 819855786
1 parent 78e74b5 commit b650181

File tree

3 files changed

+64
-0
lines changed

3 files changed

+64
-0
lines changed

src/google/adk/agents/run_config.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,11 @@ class RunConfig(BaseModel):
9494
session_resumption: Optional[types.SessionResumptionConfig] = None
9595
"""Configures session resumption mechanism. Only support transparent session resumption mode now."""
9696

97+
context_window_compression: Optional[types.ContextWindowCompressionConfig] = (
98+
None
99+
)
100+
"""Configuration for context window compression. If set, this will enable context window compression for LLM input."""
101+
97102
save_live_audio: bool = False
98103
"""Saves live video and audio data to session and artifact service.
99104

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,9 @@ async def run_async(
7979
llm_request.live_connect_config.session_resumption = (
8080
invocation_context.run_config.session_resumption
8181
)
82+
llm_request.live_connect_config.context_window_compression = (
83+
invocation_context.run_config.context_window_compression
84+
)
8285

8386
# TODO: handle tool append here, instead of in BaseTool.process_llm_request.
8487

tests/unittests/streaming/test_live_streaming_configs.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -586,3 +586,59 @@ def test_streaming_with_session_resumption_config():
586586
llm_request_sent_to_mock.live_connect_config.session_resumption.transparent
587587
is True
588588
)
589+
590+
591+
def test_streaming_with_context_window_compression_config():
592+
"""Test streaming with context window compression config."""
593+
response = LlmResponse(turn_complete=True)
594+
595+
mock_model = testing_utils.MockModel.create([response])
596+
597+
root_agent = Agent(
598+
name='root_agent',
599+
model=mock_model,
600+
tools=[],
601+
)
602+
603+
runner = testing_utils.InMemoryRunner(
604+
root_agent=root_agent, response_modalities=['AUDIO']
605+
)
606+
607+
# Create run config with context window compression
608+
run_config = RunConfig(
609+
context_window_compression=types.ContextWindowCompressionConfig(
610+
trigger_tokens=1000,
611+
sliding_window=types.SlidingWindow(target_tokens=500),
612+
)
613+
)
614+
615+
live_request_queue = LiveRequestQueue()
616+
live_request_queue.send_realtime(
617+
blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm')
618+
)
619+
620+
res_events = runner.run_live(live_request_queue, run_config)
621+
622+
assert res_events is not None, 'Expected a list of events, got None.'
623+
assert (
624+
len(res_events) > 0
625+
), 'Expected at least one response, but got an empty list.'
626+
assert len(mock_model.requests) == 1
627+
628+
# Get the request that was captured
629+
llm_request_sent_to_mock = mock_model.requests[0]
630+
631+
# Assert that the request contained the correct configuration
632+
assert llm_request_sent_to_mock.live_connect_config is not None
633+
assert (
634+
llm_request_sent_to_mock.live_connect_config.context_window_compression
635+
is not None
636+
)
637+
assert (
638+
llm_request_sent_to_mock.live_connect_config.context_window_compression.trigger_tokens
639+
== 1000
640+
)
641+
assert (
642+
llm_request_sent_to_mock.live_connect_config.context_window_compression.sliding_window.target_tokens
643+
== 500
644+
)

0 commit comments

Comments
 (0)