mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
78e74b5bf2
commit
b650181384
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user