style(gemini): drop the narrative comment and type the test helpers

Both from the Greptile review: AGENTS.md limits comments to complex logic,
tool inputs and TODOs, and requires fully typed code. The reasoning that was
in the comment lives in the PR description and the commit message instead
This commit is contained in:
basil-k-aji-dev 2026-09-20 17:27:49 +05:30
parent 874e9bb67f
commit 75fe898908
2 changed files with 10 additions and 22 deletions

View file

@ -2289,14 +2289,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
is_function_call=is_function_call(standard_optional_params),
)
# Outside the "content"/"parts" branch on purpose. Gemini does not
# promise to put groundingMetadata on a chunk that also carries text:
# streaming commonly delivers it on the final candidate, which has a
# finishReason and no parts at all. Converting only when parts were
# present dropped the citations for exactly those responses, which is
# why the same prompt produced annotations on some runs and not
# others while non-streaming -- one candidate, always with parts --
# produced them every time.
if candidate_grounding_metadata:
annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations(
grounding_metadata=candidate_grounding_metadata,

View file

@ -6376,7 +6376,7 @@ _GROUNDING_URIS: Final = [
]
def _gemini_stream_iterator():
def _gemini_stream_iterator() -> "ModelResponseIterator":
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ModelResponseIterator
logging_obj = MagicMock()
@ -6384,8 +6384,8 @@ def _gemini_stream_iterator():
return ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj)
def _delta_annotations(model_response):
found: list = []
def _delta_annotations(model_response: object) -> List[dict]:
found: List[dict] = []
for choice in model_response.choices:
annotations = getattr(getattr(choice, "delta", None), "annotations", None)
if annotations:
@ -6393,12 +6393,11 @@ def _delta_annotations(model_response):
return found
def _citation_urls(annotations):
def _citation_urls(annotations: List[dict]) -> List[str]:
return [annotation["url_citation"]["url"] for annotation in annotations]
def test_streaming_grounding_on_the_final_chunk_produces_annotations():
"""Gemini often sends groundingMetadata on a candidate with finishReason and no parts"""
def test_streaming_grounding_on_the_final_chunk_produces_annotations() -> None:
iterator = _gemini_stream_iterator()
iterator.chunk_parser({"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": _GROUNDED_TEXT}]}}]})
@ -6409,8 +6408,7 @@ def test_streaming_grounding_on_the_final_chunk_produces_annotations():
assert _citation_urls(_delta_annotations(final)) == _GROUNDING_URIS
def test_streaming_grounding_alongside_text_produces_annotations():
"""The shape that already worked keeps working"""
def test_streaming_grounding_alongside_text_produces_annotations() -> None:
iterator = _gemini_stream_iterator()
chunk = iterator.chunk_parser(
@ -6432,7 +6430,7 @@ def test_streaming_grounding_alongside_text_produces_annotations():
assert annotations[0]["url_citation"]["start_index"] == 0
def test_streaming_without_grounding_carries_no_annotations():
def test_streaming_without_grounding_carries_no_annotations() -> None:
iterator = _gemini_stream_iterator()
chunk = iterator.chunk_parser(
@ -6442,8 +6440,7 @@ def test_streaming_without_grounding_carries_no_annotations():
assert _delta_annotations(chunk) == []
def test_streaming_grounding_without_web_chunks_carries_no_annotations():
"""Maps grounding has no web URI to cite, so it must not invent one"""
def test_streaming_grounding_without_web_chunks_carries_no_annotations() -> None:
iterator = _gemini_stream_iterator()
chunk = iterator.chunk_parser(
@ -6466,7 +6463,7 @@ def test_streaming_grounding_without_web_chunks_carries_no_annotations():
assert _delta_annotations(chunk) == []
def test_non_streaming_grounding_annotations_are_unchanged():
def test_non_streaming_grounding_annotations_are_unchanged() -> None:
model_response = ModelResponse()
VertexGeminiConfig._process_candidates(
[
@ -6485,8 +6482,7 @@ def test_non_streaming_grounding_annotations_are_unchanged():
assert _citation_urls(annotations or []) == _GROUNDING_URIS
def test_streamed_grounding_survives_reassembly():
"""A client rebuilding the stream ends up with the citations a non-streaming call returns"""
def test_streamed_grounding_survives_reassembly() -> None:
iterator = _gemini_stream_iterator()
chunks = [
iterator.chunk_parser(