From 91a8fd4b8956d10afa243f8f9d23771b15e4c4f9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 12 May 2026 16:10:41 +0000 Subject: [PATCH] fix(vertex_ai): preserve uppercase Gemini finish_reason so SAFETY maps to content_filter Lowercasing hard-stop finish reasons (SAFETY, RECITATION, BLOCKLIST, PROHIBITED_CONTENT, SPII) before handing them to StreamingChoices caused map_finish_reason to miss the _FINISH_REASON_MAP lookup (which is keyed on the uppercase Gemini values) and silently fall back to 'stop'. Pass the raw uppercase finish_reason through so safety-blocked Agent Engine responses surface as 'content_filter' as intended. --- litellm/llms/vertex_ai/agent_engine/sse_iterator.py | 8 +++++--- .../llms/vertex_ai/agent_engine/test_transformation.py | 9 +++++---- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/litellm/llms/vertex_ai/agent_engine/sse_iterator.py b/litellm/llms/vertex_ai/agent_engine/sse_iterator.py index 23be4248c80..66c5b220870 100644 --- a/litellm/llms/vertex_ai/agent_engine/sse_iterator.py +++ b/litellm/llms/vertex_ai/agent_engine/sse_iterator.py @@ -119,9 +119,11 @@ class VertexAgentEngineResponseIterator(BaseModelResponseIterator): if raw_finish_reason and ( text is not None or raw_finish_reason in self._HARD_STOP_FINISH_REASONS ): - finish_reason = ( - "stop" if raw_finish_reason == "STOP" else raw_finish_reason.lower() - ) + # Pass the raw Gemini finish reason (uppercase) through so that + # downstream ``map_finish_reason`` can map hard-stop signals like + # ``SAFETY`` / ``RECITATION`` to ``content_filter`` instead of + # silently falling back to ``stop``. + finish_reason = raw_finish_reason usage = None usage_metadata = chunk.get("usage_metadata") or {} diff --git a/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py b/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py index d60a9d3f7bd..3e7fbc28872 100644 --- a/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py +++ b/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py @@ -216,10 +216,11 @@ class TestVertexAgentEngineChunkParser: result = self._iterator().chunk_parser(chunk) - # StreamingChoices normalizes "safety" → OpenAI "stop"; the key - # behavior is that *some* terminal finish_reason is surfaced - # instead of being dropped (which is what happens for plain STOP). - assert result.choices[0].finish_reason is not None + # StreamingChoices normalizes Gemini "SAFETY" → OpenAI "content_filter" + # via map_finish_reason; the raw uppercase value must be passed through + # so that downstream consumers can distinguish a safety block from a + # plain stop. + assert result.choices[0].finish_reason == "content_filter" def test_chunk_parser_surfaces_max_tokens_finish_reason_without_content(self): chunk = {