Chore: address review feedback for grounding metadata extraction

- Applied formatting using pyink and isort
- Refactored grounding metadata unit tests to use a helper for mock message creation

Author: VedantMadane

contributing/samples/gepa/experiment.py

@@ -43,7 +43,6 @@
 from tau_bench.types import EnvRunResult
 from tau_bench.types import RunConfig
 import tau_bench_agent as tau_bench_agent_lib
-
 import utils
 
 

contributing/samples/gepa/run_experiment.py

@@ -25,7 +25,6 @@
 from absl import flags
 import experiment
 from google.genai import types
-
 import utils
 
 _OUTPUT_DIR = flags.DEFINE_string(

tests/unittests/models/test_gemini_llm_connection.py

@@ -783,11 +783,36 @@ async def test_send_history_filters_various_audio_mime_types(
   mock_gemini_session.send.assert_not_called()
 
 
+def _create_mock_receive_message(
+    model_turn=None,
+    grounding_metadata=None,
+    interrupted=False,
+    turn_complete=False,
+    tool_call=None,
+):
+  """Helper to create a mock message from the Gemini API."""
+  mock_server_content = mock.Mock()
+  mock_server_content.model_turn = model_turn
+  mock_server_content.interrupted = interrupted
+  mock_server_content.input_transcription = None
+  mock_server_content.output_transcription = None
+  mock_server_content.turn_complete = turn_complete
+  mock_server_content.generation_complete = False
+  mock_server_content.grounding_metadata = grounding_metadata
+
+  mock_message = mock.Mock()
+  mock_message.usage_metadata = None
+  mock_message.server_content = mock_server_content
+  mock_message.tool_call = tool_call
+  mock_message.session_resumption_update = None
+  return mock_message
+
+
 @pytest.mark.asyncio
 async def test_receive_extracts_grounding_metadata(
     gemini_connection, mock_gemini_session
 ):
-  """Test that grounding_metadata is extracted from server_content and included in LlmResponse."""
+  """Test that grounding_metadata is extracted and included in LlmResponse."""
   mock_content = types.Content(
       role='model', parts=[types.Part.from_text(text='response text')]
   )
@@ -796,20 +821,11 @@ async def test_receive_extracts_grounding_metadata(
       web_search_queries=['web search query'],
   )
 
-  mock_server_content = mock.Mock()
-  mock_server_content.model_turn = mock_content
-  mock_server_content.interrupted = False
-  mock_server_content.input_transcription = None
-  mock_server_content.output_transcription = None
-  mock_server_content.turn_complete = True
-  mock_server_content.generation_complete = False
-  mock_server_content.grounding_metadata = mock_grounding_metadata
-
-  mock_message = mock.Mock()
-  mock_message.usage_metadata = None
-  mock_message.server_content = mock_server_content
-  mock_message.tool_call = None
-  mock_message.session_resumption_update = None
+  mock_message = _create_mock_receive_message(
+      model_turn=mock_content,
+      grounding_metadata=mock_grounding_metadata,
+      turn_complete=True,
+  )
 
   async def mock_receive_generator():
     yield mock_message
@@ -842,36 +858,12 @@ async def test_receive_grounding_metadata_at_turn_complete(
   )
 
   # First message with grounding but no content
-  mock_server_content1 = mock.Mock()
-  mock_server_content1.model_turn = None
-  mock_server_content1.interrupted = False
-  mock_server_content1.input_transcription = None
-  mock_server_content1.output_transcription = None
-  mock_server_content1.turn_complete = False
-  mock_server_content1.generation_complete = False
-  mock_server_content1.grounding_metadata = mock_grounding_metadata
-
-  message1 = mock.Mock()
-  message1.usage_metadata = None
-  message1.server_content = mock_server_content1
-  message1.tool_call = None
-  message1.session_resumption_update = None
+  message1 = _create_mock_receive_message(
+      grounding_metadata=mock_grounding_metadata
+  )
 
   # Second message with turn_complete
-  mock_server_content2 = mock.Mock()
-  mock_server_content2.model_turn = None
-  mock_server_content2.interrupted = False
-  mock_server_content2.input_transcription = None
-  mock_server_content2.output_transcription = None
-  mock_server_content2.turn_complete = True
-  mock_server_content2.generation_complete = False
-  mock_server_content2.grounding_metadata = None
-
-  message2 = mock.Mock()
-  message2.usage_metadata = None
-  message2.server_content = mock_server_content2
-  message2.tool_call = None
-  message2.session_resumption_update = None
+  message2 = _create_mock_receive_message(turn_complete=True)
 
   async def mock_receive_generator():
     yield message1
@@ -902,20 +894,11 @@ async def test_receive_grounding_metadata_with_text_and_turn_complete(
   )
 
   # Message with both content and grounding, followed by turn_complete
-  mock_server_content = mock.Mock()
-  mock_server_content.model_turn = mock_content
-  mock_server_content.interrupted = False
-  mock_server_content.input_transcription = None
-  mock_server_content.output_transcription = None
-  mock_server_content.turn_complete = True
-  mock_server_content.generation_complete = False
-  mock_server_content.grounding_metadata = mock_grounding_metadata
-
-  mock_message = mock.Mock()
-  mock_message.usage_metadata = None
-  mock_message.server_content = mock_server_content
-  mock_message.tool_call = None
-  mock_message.session_resumption_update = None
+  mock_message = _create_mock_receive_message(
+      model_turn=mock_content,
+      grounding_metadata=mock_grounding_metadata,
+      turn_complete=True,
+  )
 
   async def mock_receive_generator():
     yield mock_message
@@ -946,20 +929,9 @@ async def test_receive_grounding_metadata_with_tool_call(
   )
 
   # First message with grounding metadata
-  mock_server_content1 = mock.Mock()
-  mock_server_content1.model_turn = None
-  mock_server_content1.interrupted = False
-  mock_server_content1.input_transcription = None
-  mock_server_content1.output_transcription = None
-  mock_server_content1.turn_complete = False
-  mock_server_content1.generation_complete = False
-  mock_server_content1.grounding_metadata = mock_grounding_metadata
-
-  message1 = mock.Mock()
-  message1.usage_metadata = None
-  message1.server_content = mock_server_content1
-  message1.tool_call = None
-  message1.session_resumption_update = None
+  message1 = _create_mock_receive_message(
+      grounding_metadata=mock_grounding_metadata
+  )
 
   # Second message with tool_call
   mock_function_call = types.FunctionCall(
@@ -968,11 +940,8 @@ async def test_receive_grounding_metadata_with_tool_call(
   mock_tool_call = mock.Mock()
   mock_tool_call.function_calls = [mock_function_call]
 
-  message2 = mock.Mock()
-  message2.usage_metadata = None
+  message2 = _create_mock_receive_message(tool_call=mock_tool_call)
   message2.server_content = None
-  message2.tool_call = mock_tool_call
-  message2.session_resumption_update = None
 
   async def mock_receive_generator():
     yield message1
@@ -1006,55 +975,18 @@ async def test_receive_interrupted_with_pending_text_preserves_flag(
   mock_content1 = types.Content(
       role='model', parts=[types.Part.from_text(text='partial')]
   )
-  mock_server_content1 = mock.Mock()
-  mock_server_content1.model_turn = mock_content1
-  mock_server_content1.interrupted = False
-  mock_server_content1.input_transcription = None
-  mock_server_content1.output_transcription = None
-  mock_server_content1.turn_complete = False
-  mock_server_content1.generation_complete = False
-  mock_server_content1.grounding_metadata = mock_grounding_metadata
-
-  message1 = mock.Mock()
-  message1.usage_metadata = None
-  message1.server_content = mock_server_content1
-  message1.tool_call = None
-  message1.session_resumption_update = None
+  message1 = _create_mock_receive_message(
+      model_turn=mock_content1, grounding_metadata=mock_grounding_metadata
+  )
 
   # Second message with more text
   mock_content2 = types.Content(
       role='model', parts=[types.Part.from_text(text=' text')]
   )
-  mock_server_content2 = mock.Mock()
-  mock_server_content2.model_turn = mock_content2
-  mock_server_content2.interrupted = False
-  mock_server_content2.input_transcription = None
-  mock_server_content2.output_transcription = None
-  mock_server_content2.turn_complete = False
-  mock_server_content2.generation_complete = False
-  mock_server_content2.grounding_metadata = None
-
-  message2 = mock.Mock()
-  message2.usage_metadata = None
-  message2.server_content = mock_server_content2
-  message2.tool_call = None
-  message2.session_resumption_update = None
+  message2 = _create_mock_receive_message(model_turn=mock_content2)
 
   # Third message with interrupted signal
-  mock_server_content3 = mock.Mock()
-  mock_server_content3.model_turn = None
-  mock_server_content3.interrupted = True
-  mock_server_content3.input_transcription = None
-  mock_server_content3.output_transcription = None
-  mock_server_content3.turn_complete = False
-  mock_server_content3.generation_complete = False
-  mock_server_content3.grounding_metadata = None
-
-  message3 = mock.Mock()
-  message3.usage_metadata = None
-  message3.server_content = mock_server_content3
-  message3.tool_call = None
-  message3.session_resumption_update = None
+  message3 = _create_mock_receive_message(interrupted=True)
 
   async def mock_receive_generator():
     yield message1