feat(live): Add live_session_id to LlmResponse

This change introduces a `live_session_id` field to the `LlmResponse` dataclass and populates it in all responses generated by the `GeminiLlmConnection` when using a Live session. This allows tracking which Live session each response belongs to.

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 900354397

Author: wuliang229

src/google/adk/models/gemini_llm_connection.py

@@ -176,6 +176,7 @@ def __build_full_text_response(self, text: str):
             role='model',
             parts=[types.Part.from_text(text=text)],
         ),
+        live_session_id=self._gemini_session.session_id,
     )
 
   async def receive(self) -> AsyncGenerator[LlmResponse, None]:
@@ -192,11 +193,13 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
       # partial content and emit responses as needed.
       async for message in agen:
         logger.debug('Got LLM Live message: %s', message)
+        live_session_id = self._gemini_session.session_id
         if message.usage_metadata:
           # Tracks token usage data per model.
           yield LlmResponse(
               usage_metadata=message.usage_metadata,
               model_version=self._model_version,
+              live_session_id=live_session_id,
           )
         if message.server_content:
           content = message.server_content.model_turn
@@ -211,13 +214,15 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
                 grounding_metadata=message.server_content.grounding_metadata,
                 interrupted=message.server_content.interrupted,
                 model_version=self._model_version,
+                live_session_id=live_session_id,
             )
 
           if content and content.parts:
             llm_response = LlmResponse(
                 content=content,
                 interrupted=message.server_content.interrupted,
                 model_version=self._model_version,
+                live_session_id=live_session_id,
             )
             # grounding_metadata is yielded again at turn_complete,
             # so avoid duplicating it here if turn_complete is true.
@@ -248,6 +253,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
                   ),
                   partial=True,
                   model_version=self._model_version,
+                  live_session_id=live_session_id,
               )
             # finished=True and partial transcription may happen in the same
             # message.
@@ -259,6 +265,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
                   ),
                   partial=False,
                   model_version=self._model_version,
+                  live_session_id=live_session_id,
               )
               self._input_transcription_text = ''
           if message.server_content.output_transcription:
@@ -273,6 +280,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
                   ),
                   partial=True,
                   model_version=self._model_version,
+                  live_session_id=live_session_id,
               )
             if message.server_content.output_transcription.finished:
               yield LlmResponse(
@@ -282,6 +290,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
                   ),
                   partial=False,
                   model_version=self._model_version,
+                  live_session_id=live_session_id,
               )
               self._output_transcription_text = ''
           # The Gemini API might not send a transcription finished signal.
@@ -300,6 +309,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
                   ),
                   partial=False,
                   model_version=self._model_version,
+                  live_session_id=live_session_id,
               )
               self._input_transcription_text = ''
             if self._output_transcription_text:
@@ -310,6 +320,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
                   ),
                   partial=False,
                   model_version=self._model_version,
+                  live_session_id=live_session_id,
               )
               self._output_transcription_text = ''
           if message.server_content.turn_complete:
@@ -321,13 +332,15 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
               yield LlmResponse(
                   content=types.Content(role='model', parts=tool_call_parts),
                   model_version=self._model_version,
+                  live_session_id=live_session_id,
               )
               tool_call_parts = []
             yield LlmResponse(
                 turn_complete=True,
                 interrupted=message.server_content.interrupted,
                 grounding_metadata=message.server_content.grounding_metadata,
                 model_version=self._model_version,
+                live_session_id=live_session_id,
             )
             break
           # in case of empty content or parts, we still surface it
@@ -342,6 +355,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
               yield LlmResponse(
                   interrupted=message.server_content.interrupted,
                   model_version=self._model_version,
+                  live_session_id=live_session_id,
               )
         if message.tool_call:
           logger.debug('Received tool call: %s', message.tool_call)
@@ -358,20 +372,23 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
               LlmResponse(
                   live_session_resumption_update=message.session_resumption_update,
                   model_version=self._model_version,
+                  live_session_id=live_session_id,
               )
           )
         if message.go_away:
           logger.debug('Received GoAway message: %s', message.go_away)
           yield LlmResponse(
               go_away=message.go_away,
               model_version=self._model_version,
+              live_session_id=live_session_id,
           )
 
       if tool_call_parts:
         logger.debug('Exited loop with pending tool_call_parts')
         yield LlmResponse(
             content=types.Content(role='model', parts=tool_call_parts),
             model_version=self._model_version,
+            live_session_id=self._gemini_session.session_id,
         )
 
   async def close(self):

src/google/adk/models/llm_response.py

@@ -110,6 +110,9 @@ class LlmResponse(BaseModel):
   ] = None
   """The session resumption update of the LlmResponse"""
 
+  live_session_id: Optional[str] = None
+  """The session ID of the Live session."""
+
   go_away: Optional[types.LiveServerGoAway] = None
   """The GoAway signal from the Live model."""
 

tests/unittests/models/test_gemini_llm_connection.py

@@ -25,7 +25,9 @@
 @pytest.fixture
 def mock_gemini_session():
   """Mock Gemini session for testing."""
-  return mock.AsyncMock()
+  mock_session = mock.AsyncMock()
+  mock_session.session_id = 'test-session-id'
+  return mock_session
 
 
 @pytest.fixture
@@ -247,6 +249,41 @@ async def mock_receive_generator():
   assert content_response.content == mock_content
 
 
+async def test_receive_populates_live_session_id(
+    gemini_connection, mock_gemini_session
+):
+  """Test that receive populates live_session_id in LlmResponse."""
+  mock_message = mock.AsyncMock()
+  mock_message.usage_metadata = None
+  mock_message.server_content = None
+  mock_message.tool_call = None
+  mock_message.session_resumption_update = None
+  mock_message.go_away = None
+
+  mock_server_content = mock.Mock()
+  mock_server_content.model_turn = types.Content(
+      role='model', parts=[types.Part.from_text(text='text')]
+  )
+  mock_server_content.interrupted = False
+  mock_server_content.input_transcription = None
+  mock_server_content.output_transcription = None
+  mock_server_content.turn_complete = False
+  mock_server_content.grounding_metadata = None
+
+  mock_message.server_content = mock_server_content
+
+  async def mock_receive_generator():
+    yield mock_message
+
+  mock_gemini_session.receive = mock.Mock(return_value=mock_receive_generator())
+
+  responses = [resp async for resp in gemini_connection.receive()]
+
+  assert responses
+  for resp in responses:
+    assert resp.live_session_id == 'test-session-id'
+
+
 @pytest.mark.asyncio
 async def test_receive_transcript_finished_on_interrupt(
     gemini_api_connection,