Skip to content

Commit 61a3933

Browse files
committed
fix: Support generalized history config injection for Gemini 3.1 Live on Vertex AI
- Exposed history_config in RunConfig. - Mapped history_config to LLM live connect request configuration. - Generalized history connection logic to automatically inject `initial_history_in_client_content = True` when seeding history on a fresh connection for both Gemini API and Vertex AI backends. - Updated and added comprehensive unit tests to verify history configuration behaviour. TAG=agy CONV=822f8c76-9099-4f01-a2b8-10a7de0d61a2 Change-Id: Ib532626d5d7d887b17664567aed94ba09ad90b33
1 parent ded6dbb commit 61a3933

4 files changed

Lines changed: 92 additions & 25 deletions

File tree

src/google/adk/agents/run_config.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,9 @@ class RunConfig(BaseModel):
247247
session_resumption: Optional[types.SessionResumptionConfig] = None
248248
"""Configures session resumption mechanism. Only support transparent session resumption mode now."""
249249

250+
history_config: Optional[types.HistoryConfig] = None
251+
"""Configures the exchange of history between the client and the server."""
252+
250253
context_window_compression: Optional[types.ContextWindowCompressionConfig] = (
251254
None
252255
)

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

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -553,18 +553,26 @@ async def run_live(
553553
if session_resumption.transparent is None:
554554
session_resumption.transparent = True
555555

556+
# When seeding a fresh connection with prior conversation history, set
557+
# initial_history_in_client_content to True. This tells the Live server
558+
# that the provided history already includes the model's past responses,
559+
# preventing the server from generating duplicate responses for those replayed turns.
556560
if (
557-
isinstance(llm, Gemini)
558-
and llm._api_backend == GoogleLLMVariant.GEMINI_API
559-
and model_name_utils.is_gemini_3_1_flash_live(llm_request.model)
560-
and llm_request.contents
561+
llm_request.contents
561562
and not invocation_context.live_session_resumption_handle
562563
):
563-
if llm_request.live_connect_config is None:
564+
if not llm_request.live_connect_config:
564565
llm_request.live_connect_config = types.LiveConnectConfig()
565-
if llm_request.live_connect_config.history_config is None:
566-
llm_request.live_connect_config.history_config = types.HistoryConfig(
567-
initial_history_in_client_content=True
566+
if not llm_request.live_connect_config.history_config:
567+
llm_request.live_connect_config.history_config = (
568+
types.HistoryConfig()
569+
)
570+
if (
571+
llm_request.live_connect_config.history_config.initial_history_in_client_content
572+
is None
573+
):
574+
llm_request.live_connect_config.history_config.initial_history_in_client_content = (
575+
True
568576
)
569577

570578
logger.info(

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,9 @@ def _build_basic_request(
9494
llm_request.live_connect_config.session_resumption = (
9595
invocation_context.run_config.session_resumption
9696
)
97+
llm_request.live_connect_config.history_config = (
98+
invocation_context.run_config.history_config
99+
)
97100
llm_request.live_connect_config.context_window_compression = (
98101
invocation_context.run_config.context_window_compression
99102
)

tests/unittests/flows/llm_flows/test_base_llm_flow.py

Lines changed: 70 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1390,16 +1390,14 @@ async def mock_receive_2():
13901390

13911391
@pytest.mark.asyncio
13921392
@pytest.mark.parametrize(
1393-
"api_backend,should_have_history_config",
1393+
"api_backend",
13941394
[
1395-
(GoogleLLMVariant.GEMINI_API, True),
1396-
(GoogleLLMVariant.VERTEX_AI, False),
1395+
GoogleLLMVariant.GEMINI_API,
1396+
GoogleLLMVariant.VERTEX_AI,
13971397
],
13981398
)
1399-
async def test_run_live_history_config_gated_by_backend(
1400-
api_backend, should_have_history_config
1401-
):
1402-
"""Test that run_live only sets history_config for Gemini API backend."""
1399+
async def test_run_live_history_config_set_for_all_backends(api_backend):
1400+
"""Test that run_live sets history_config for all backends."""
14031401

14041402
real_model = Gemini(model='gemini-3.1-flash-live-preview')
14051403
mock_connection = mock.AsyncMock()
@@ -1448,13 +1446,68 @@ async def mock_preprocess(ctx, req):
14481446

14491447
assert mock_connect.call_count == 1
14501448
called_req = mock_connect.call_args[0][0]
1451-
if should_have_history_config:
1452-
assert called_req.live_connect_config is not None
1453-
assert called_req.live_connect_config.history_config is not None
1454-
assert (
1455-
called_req.live_connect_config.history_config.initial_history_in_client_content
1456-
is True
1457-
)
1458-
else:
1459-
if called_req.live_connect_config:
1460-
assert called_req.live_connect_config.history_config is None
1449+
assert called_req.live_connect_config is not None
1450+
assert called_req.live_connect_config.history_config is not None
1451+
assert (
1452+
called_req.live_connect_config.history_config.initial_history_in_client_content
1453+
is True
1454+
)
1455+
1456+
1457+
@pytest.mark.asyncio
1458+
async def test_run_live_respects_explicit_initial_history_in_client_content_false():
1459+
"""Test that run_live respects explicit initial_history_in_client_content=False in RunConfig."""
1460+
1461+
real_model = Gemini()
1462+
mock_connection = mock.AsyncMock()
1463+
1464+
agent = Agent(name='test_agent', model=real_model)
1465+
invocation_context = await testing_utils.create_invocation_context(
1466+
agent=agent
1467+
)
1468+
invocation_context.live_request_queue = LiveRequestQueue()
1469+
run_config = RunConfig(
1470+
history_config=types.HistoryConfig(initial_history_in_client_content=False)
1471+
)
1472+
invocation_context.run_config = run_config
1473+
1474+
flow = BaseLlmFlowForTesting()
1475+
1476+
async def mock_preprocess(ctx, req):
1477+
req.contents = [types.Content(parts=[types.Part.from_text(text='history')])]
1478+
from google.adk.flows.llm_flows.basic import _build_basic_request
1479+
_build_basic_request(ctx, req)
1480+
yield Event(id=Event.new_id(), author='test')
1481+
1482+
with mock.patch.object(
1483+
flow, '_preprocess_async', side_effect=mock_preprocess
1484+
):
1485+
with mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock):
1486+
1487+
class StopTestError(Exception):
1488+
pass
1489+
1490+
async def mock_receive():
1491+
yield LlmResponse(
1492+
content=types.Content(parts=[types.Part.from_text(text='hi')])
1493+
)
1494+
raise StopTestError('stop')
1495+
1496+
mock_connection.receive = mock.Mock(side_effect=mock_receive)
1497+
1498+
with mock.patch(
1499+
'google.adk.models.google_llm.Gemini.connect'
1500+
) as mock_connect:
1501+
mock_connect.return_value.__aenter__.return_value = mock_connection
1502+
1503+
try:
1504+
async for _ in flow.run_live(invocation_context):
1505+
pass
1506+
except StopTestError:
1507+
pass
1508+
1509+
assert mock_connect.call_count == 1
1510+
call_req = mock_connect.call_args[0][0]
1511+
assert call_req.live_connect_config.history_config is not None
1512+
assert call_req.live_connect_config.history_config.initial_history_in_client_content is False
1513+

0 commit comments

Comments
 (0)