feat: Implement Live Session Resumption

Previous implementation doesn't pass the actual handle to server. Now we cache the handle and pass it over when reconnection happens.

To enable:
    run_config = RunConfig(
        session_resumption=types.SessionResumptionConfig(transparent=True)
    )

PiperOrigin-RevId: 791308462
This commit is contained in:
Hangfei Lin
2025-08-05 11:50:54 -07:00
committed by Copybara-Service
parent 423542a43f
commit 71fbc9275b
9 changed files with 333 additions and 76 deletions
@@ -219,6 +219,13 @@ class GeminiLlmConnection(BaseLlmConnection):
for function_call in message.tool_call.function_calls
]
yield LlmResponse(content=types.Content(role='model', parts=parts))
if message.session_resumption_update:
logger.info('Redeived session reassumption message: %s', message)
yield (
LlmResponse(
live_session_resumption_update=message.session_resumption_update
)
)
async def close(self):
"""Closes the llm server connection."""
+1
View File
@@ -289,6 +289,7 @@ class Gemini(BaseLlm):
],
)
llm_request.live_connect_config.tools = llm_request.config.tools
logger.info('Connecting to live with llm_request:%s', llm_request)
async with self._live_api_client.aio.live.connect(
model=llm_request.model, config=llm_request.live_connect_config
) as live_session:
+5
View File
@@ -89,6 +89,11 @@ class LlmResponse(BaseModel):
usage_metadata: Optional[types.GenerateContentResponseUsageMetadata] = None
"""The usage metadata of the LlmResponse"""
live_session_resumption_update: Optional[
types.LiveServerSessionResumptionUpdate
] = None
"""The session resumption update of the LlmResponse"""
@staticmethod
def create(
generate_content_response: types.GenerateContentResponse,