Skip to content

Python: fix Gemini finish reason fallback and usage-attach cascade - #7837

Open
Jeremy Schoemaker (shoemoney) wants to merge 1 commit into
microsoft:mainfrom
shoemoney:fix/gemini-unmapped-finish-reason
Open

Python: fix Gemini finish reason fallback and usage-attach cascade#7837
Jeremy Schoemaker (shoemoney) wants to merge 1 commit into
microsoft:mainfrom
shoemoney:fix/gemini-unmapped-finish-reason

Conversation

@shoemoney

Copy link
Copy Markdown

Motivation & Context

#7105 gave chat clients a fallback so an unmapped provider finish reason value passes through as the raw string instead of vanishing to None. It touched ag, bedrock, claude, core, github_copilot, ollama, and openai. It did not touch gemini.

Bedrock's version of the same method after #7105:

def _map_finish_reason(self, reason: str | None) -> str | None:
    if not reason:
        return None
    return FINISH_REASON_MAP.get(reason.lower(), reason)

Gemini's still has no fallback:

def _map_finish_reason(self, reason: str | None) -> FinishReasonLiteral | None:
    if not reason:
        return None
    return _FINISH_REASON_MAP.get(reason)

_FINISH_REASON_MAP covers 13 of the 18 members of google.genai.types.FinishReason in the package's own pinned dependency range (google-genai>=1.69.0,<3.0.0). Missing: OTHER, TOO_MANY_TOOL_CALLS, NO_IMAGE, IMAGE_OTHER (plus FINISH_REASON_UNSPECIFIED, which is a legitimate absent case and correctly stays None).

The bug cascades. In _process_chunk:

if finish_reason and (usage := self._parse_usage(chunk.usage_metadata)):

Usage is attached to a streamed chunk only when finish_reason is truthy, so an unmapped reason drops the finish reason and the whole turn's token/billing accounting. ChatTelemetryLayer in observability.py skips recording the terminal span state the same way.

Concrete scenario: a Gemini call inside an agentic tool loop trips Gemini's own tool-call-count guardrail and returns finish_reason=TOO_MANY_TOOL_CALLS. The caller sees finish_reason=None, no usage, and no way to tell the run stopped abnormally instead of completing normally.

Description & Review Guide

  • What are the major changes?
    • _map_finish_reason now falls back to the raw reason string, wrapped as FinishReason(reason), instead of None. FINISH_REASON_UNSPECIFIED and an absent reason still map to None.
    • The return type widens from FinishReasonLiteral | None to FinishReasonLiteral | FinishReason | None. This is not the same pattern bedrock uses (str | None with no wrapping) — bedrock's file carries a blanket # type: ignore at the top of the module, gemini's does not, so a bare str fails this package's strict Pyright config. FinishReason(reason) is the same construct ollama and openai used for the same fallback in Python: Normalize chat finish reasons #7105.
  • What is the impact of these changes?
    • An unmapped-but-real Gemini finish reason (currently OTHER, TOO_MANY_TOOL_CALLS, NO_IMAGE, IMAGE_OTHER) is now surfaced to the caller instead of silently disappearing, and the final streamed chunk's usage/token accounting is no longer dropped alongside it.
  • What do you want reviewers to focus on?
    • Whether FINISH_REASON_UNSPECIFIED is the only value that should keep mapping to None, or whether other values should be excluded from the fallback too.

Related Issue

Fixes #7836

Testing

Ran the gemini package's unit test suite only (not the full monorepo suite):

python -m pytest packages/gemini/tests -m "not integration" -v

GREEN (with the fix): 157 passed, 8 deselected (integration tests, no credentials configured).

Reverted only the source change (kept the new/updated tests) to confirm RED:

5 failed, 152 passed, 8 deselected
FAILED test_finish_reason_mapping[OTHER-OTHER]
FAILED test_finish_reason_mapping[TOO_MANY_TOOL_CALLS-TOO_MANY_TOOL_CALLS]
FAILED test_finish_reason_mapping[NO_IMAGE-NO_IMAGE]
FAILED test_finish_reason_mapping[IMAGE_OTHER-IMAGE_OTHER]
FAILED test_unmapped_finish_reason_still_attaches_usage_on_streamed_final_chunk

Restored the fix, reran, back to GREEN (157 passed).

Also ran pyright and ruff check/ruff format --check against the changed files only; all clean. Did not run the .NET suite, the full Python monorepo test suite, or integration tests (no Gemini/Vertex credentials in this environment).

Contribution Checklist

  • The code builds clean without any errors or warnings
  • All unit tests pass, and I have added new tests where possible
  • The PR follows the Contribution Guidelines
  • This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
  • This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

microsoft#7105 added a fallback so unmapped provider finish reasons pass
through as the raw string instead of being dropped to None. It
covered ag, bedrock, claude, core, github_copilot, ollama, and
openai, but not gemini.

_FINISH_REASON_MAP covers 13 of the 18 members of
google.genai.types.FinishReason in the package's pinned dependency
range. OTHER, TOO_MANY_TOOL_CALLS, NO_IMAGE, and IMAGE_OTHER fell
through to None. FINISH_REASON_UNSPECIFIED still correctly maps to
None.

_process_chunk only attaches usage to a streamed chunk when
finish_reason is truthy, so an unmapped reason silently dropped both
the finish reason and the whole turn's usage/token accounting.

Mirrors the ollama/openai pattern of wrapping the raw fallback in
FinishReason(...) rather than bedrock's plain str return, since
bedrock's file carries a blanket type: ignore that gemini's does
not.

Fixes microsoft#7836

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR fixes a bug in the Gemini Python chat client where unmapped provider finish reasons were silently dropped to None. #7105 added a raw-string fallback to the other chat clients (ag, bedrock, claude, core, github_copilot, ollama, openai) but missed gemini. Because streamed-chunk usage is attached only when finish_reason is truthy, an unmapped reason (e.g. TOO_MANY_TOOL_CALLS) dropped both the finish reason and the turn's token/billing accounting. The fix mirrors the established ollama/openai fallback pattern, wrapping the raw reason as FinishReason(reason) so it type-checks under this package's strict Pyright config.

Changes:

  • _map_finish_reason now falls back to FinishReason(reason) for values absent from _FINISH_REASON_MAP, while FINISH_REASON_UNSPECIFIED and absent reasons still map to None; return type widened to FinishReasonLiteral | FinishReason | None.
  • Added the FinishReason import and updated the docstring to explain the fallback behavior.
  • Expanded finish-reason mapping tests and added a regression test confirming usage is still attached to the final streamed chunk under an unmapped reason.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File Description
python/packages/gemini/agent_framework_gemini/_chat_client.py Adds raw-string fallback in _map_finish_reason, imports FinishReason, widens the return type, and updates the docstring.
python/packages/gemini/tests/test_gemini_client.py Expands parametrized mapping cases (mapped, unmapped-passthrough, None/UNSPECIFIED) and adds a streamed-chunk usage-cascade regression test.

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

python Usage: [Issues, PRs], Target: Python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Python: [Bug]: Gemini chat client drops unmapped finish reasons and their usage/token accounting

2 participants