Python: fix(vertex-ai): filter out thought parts from chat message content#14022
Python: fix(vertex-ai): filter out thought parts from chat message content#14022rmotgi1227 wants to merge 1 commit into
Conversation
When Gemini is used via Vertex AI with include_thoughts=True, thought
parts have {"text": "...", "thought": True} in their dict representation.
Without the guard, thought text leaks into TextContent / StreamingTextContent
items in both _create_chat_message_content and _create_streaming_chat_message_content.
Add 'not part_dict.get("thought")' to the text-part check at both
call sites, mirroring the fix applied to GoogleAIChatCompletion for
the same issue (microsoft#13710).
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Prevents Gemini “thought”/reasoning text parts (when include_thoughts=True) from being surfaced as normal assistant text output in the Vertex AI chat completion connector.
Changes:
- Filter out text parts marked with
thought=Truein non-streaming chat message content creation. - Apply the same thought filtering in streaming chat message content creation.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| for idx, part in enumerate(candidate.content.parts): | ||
| part_dict = part.to_dict() | ||
| if "text" in part_dict: | ||
| if "text" in part_dict and not part_dict.get("thought"): |
| for idx, part in enumerate(candidate.content.parts): | ||
| part_dict = part.to_dict() | ||
| if "text" in part_dict: | ||
| if "text" in part_dict and not part_dict.get("thought"): |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Automated Code Review
Reviewers: 4 | Confidence: 92%
✓ Correctness
The fix is logically correct. The guard
not part_dict.get("thought")properly excludes thought parts (whereto_dict()returns{"text": "..", "thought": True}) from being emitted asTextContent/StreamingTextContent, while allowing regular text parts (where thethoughtkey is absent, so.get("thought")returnsNoneandnot NoneisTrue) to pass through. The change is minimal, precise, and mirrors the existingthought_signaturehandling pattern already present in the same loops. No correctness issues found.
✓ Security Reliability
The fix correctly filters out thought parts from Vertex AI responses to prevent chain-of-thought leakage. The implementation using
not part_dict.get("thought")is safe—it returns None (falsy) when the key is absent, preserving regular text parts, and correctly excludes parts withthought: True. The pattern mirrors the existingthought_signaturehandling already present in the same loops. No additional security or reliability concerns beyond the already-flaged missing test coverage.
✓ Test Coverage
The fix correctly filters thought parts from leaking into chat message content, but lacks unit test coverage for this security-sensitive behavior. Both
_create_chat_message_contentand_create_streaming_chat_message_contenthave no tests verifying that parts with{"text": "...", "thought": True}are excluded from the output items. The existing tests in the file only coverthought_signaturemetadata on function call parts, not the thought-text filtering added here. No tests exist anywhere in the Google connector test suite for the"thought"key filtering pattern.
✓ Design Approach
I did not find a design-approach issue in this diff. The new guard is narrowly targeted at the described leakage path, and within the bounded context I inspected it is consistent with the existing Vertex AI handling of part metadata such as
thought_signaturein both non-streaming and streaming deserialization.
Automated review by rmotgi1227's agents
| for idx, part in enumerate(candidate.content.parts): | ||
| part_dict = part.to_dict() | ||
| if "text" in part_dict: | ||
| if "text" in part_dict and not part_dict.get("thought"): |
There was a problem hiding this comment.
This security-sensitive guard needs a dedicated unit test. Please add a test to test_vertex_ai_chat_completion.py that creates a mock part with to_dict() returning {"text": "thinking..", "thought": True} alongside a normal text part ({"text": "visible"}), calls _create_chat_message_content, and asserts only the non-thought text appears in the result items.
| for idx, part in enumerate(candidate.content.parts): | ||
| part_dict = part.to_dict() | ||
| if "text" in part_dict: | ||
| if "text" in part_dict and not part_dict.get("thought"): |
There was a problem hiding this comment.
Same test coverage gap for the streaming path. Please add a test that verifies _create_streaming_chat_message_content excludes thought parts (parts whose to_dict() includes "thought": True) while still including regular text parts as StreamingTextContent.
github-actions
Bot
changed the title
Summary
When Gemini is used via Vertex AI with
include_thoughts=True, the model returns reasoning/thought parts alongside regular text parts. Thought parts have{"text": "...", "thought": True}in theirto_dict()representation._create_chat_message_contentand_create_streaming_chat_message_contentinVertexAIChatCompletionchecked onlyif "text" in part_dict, so thought text leaked intoTextContent/StreamingTextContentitems — exposing internal chain-of-thought reasoning as regular assistant output.Fix
Add
and not part_dict.get("thought")to the text-part guard at both call sites (non-streaming line 252, streaming line 307):This matches the fix already applied to
GoogleAIChatCompletionfor the same issue (#13710), and mirrors thethought_signaturehandling already present in the same loops.Related
GoogleAIChatCompletion(fixed by Python: fix(python/google): filter thinking text parts from chat completion responses #13711 / Python: fix: filter google ai thought text parts #13890, which do not touchVertexAIChatCompletion)