fix: propagate finish_reason from LiteLLM responses #3114
Conversation
Fixes google#3109 This change ensures that the finish_reason field from LiteLLM responses is properly propagated to LlmResponse objects, enabling callbacks to detect completion conditions like max_tokens truncation. Changes: - Extract finish_reason from LiteLLM response in lite_llm.py - Update tracing.py to handle both enum (Gemini) and string (LiteLLM) finish_reason values - Add comprehensive unit tests for finish_reason propagation The fix allows after_model_callback functions to properly detect: - "length": max_tokens limit reached - "stop": natural completion - "tool_calls": tool invocations - "content_filter": filtered content
Summary of ChangesHello @aperepel, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This PR addresses issue #3109 by implementing the propagation of the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This is a great pull request that addresses an important issue with propagating the finish_reason. The changes in lite_llm.py and the new tests are well-targeted. I have a few suggestions to improve robustness and maintainability, particularly in tracing.py and the new tests.
- Use .name instead of .value for enum finish_reason (more robust for IntEnum) - Extract first choice using walrus operator for better readability - Consolidate tests using @pytest.mark.parametrize to reduce duplication - Strengthen test assertions to verify response content All 53 tests pass.
|
/gemini review |
There was a problem hiding this comment.
Code Review
Thank you for this contribution. This pull request effectively addresses the propagation of finish_reason from LiteLLM, which is a valuable enhancement for observing model completion states. I've identified a critical issue and a couple of areas for improvement to enhance the robustness and correctness of the code. Please see my detailed comments below.
Addressing review comments Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Addressing review comments Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Address type safety issue where finish_reason can be either: - types.FinishReason enum (from Gemini responses) - str (from LiteLLM responses) Updated LlmResponse.finish_reason type hint to: Optional[Union[types.FinishReason, str]] This ensures type checkers correctly validate the dual nature of this field across different model providers. All 53 tests pass.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This is a great pull request that addresses an important issue with propagating the finish_reason. The changes are well-implemented and the addition of comprehensive unit tests is excellent. I have one suggestion regarding the data model for LlmResponse to improve long-term maintainability.
- Map finish_reason strings to proper FinishReason enum values in lite_llm.py - 'length' -> FinishReason.MAX_TOKENS - 'stop' -> FinishReason.STOP - 'tool_calls'/'function_call' -> FinishReason.STOP - 'content_filter' -> FinishReason.SAFETY - unknown values -> FinishReason.OTHER - Add clarifying comment in tracing.py for string fallback path - Update test_litellm.py to verify enum mapping: - Assert finish_reason is FinishReason enum instance - Verify correct enum values for each finish_reason string - Add test for unknown finish_reason mapping to OTHER Benefits: - Type consistency with Gemini native responses - Avoids runtime warnings from string finish_reason - Enables proper instanceof checks in callbacks - Better integration with ADK telemetry
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request effectively propagates the finish_reason from LiteLLM responses by updating the LlmResponse model, mapping string reasons to the FinishReason enum, and enhancing telemetry tracing to handle both types. The changes are well-supported by comprehensive new unit tests. I have a couple of suggestions to improve code clarity and documentation accuracy.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request does a good job of propagating the finish_reason from LiteLLM responses, which is crucial for observing and handling different model completion scenarios. The changes are well-structured and include comprehensive tests. I've included two suggestions: one to use a more specific enum value for tool calls to avoid information loss, and another to refactor the new tests for improved maintainability.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request does a great job of propagating the finish_reason from LiteLLM responses and mapping them to the standard FinishReason enum. This is a valuable fix that improves observability and allows for better handling of different completion scenarios. The changes are well-implemented, and the addition of comprehensive unit tests ensures the new logic is robust. I have a couple of suggestions to further improve code clarity and maintainability.
Maps LiteLLM finish_reason string values to proper FinishReason enum for type consistency with Gemini native responses. Changes: - Add _FINISH_REASON_MAPPING dictionary for string->enum conversion - "length" -> FinishReason.MAX_TOKENS - "stop" -> FinishReason.STOP - "tool_calls"/"function_call" -> FinishReason.STOP - "content_filter" -> FinishReason.SAFETY - Unknown values -> FinishReason.OTHER (fallback) - Update finish_reason type hint to Optional[FinishReason] (no Union needed) - Update telemetry tracing to use .name for enum serialization - Add explanatory comments: - Why tool_calls maps to STOP (no TOOL_CALL enum exists) - Docstring clarifies mapping applies to all model providers Tests: - test_finish_reason_propagation: verifies enum mapping for all values - test_finish_reason_unknown_maps_to_other: verifies fallback behavior Benefits: - Type consistency: finish_reason is always FinishReason enum - No runtime warnings from mixed types - Enables proper isinstance() checks in callbacks - Dictionary mapping improves maintainability - Better integration with ADK telemetry
aperepel
force-pushed
the
fix/litellm-finish-reason-3109
branch
from
b2e2d7b to
beadc63
Compare
|
The bot seems to be reviewing only the last commit, not the whole PR. Squashed changes into one now. /gemini review |
There was a problem hiding this comment.
Code Review
This pull request effectively addresses the issue of propagating finish_reason from LiteLLM responses by introducing a mapping to the standard FinishReason enum. The changes are well-implemented across the model, telemetry, and response layers, and are supported by comprehensive new unit tests. My feedback includes a couple of suggestions to improve code maintainability and simplify logic by removing now-redundant defensive code.
- Simplify tracing.py by removing isinstance check (always enum now) - Refactor test assertions to use dictionary mapping instead of if/elif - Reduce code duplication and improve readability Addresses Gemini Code Assist bot suggestions: - tracing.py: Direct .name access since finish_reason is always enum - test_litellm.py: Dictionary mapping for cleaner test assertions
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request correctly propagates the finish_reason from LiteLLM responses, ensuring it's available as a types.FinishReason enum in LlmResponse objects. The changes are well-implemented across lite_llm.py and tracing.py, and the updated docstring in llm_response.py adds clarity. The new unit tests in test_litellm.py are comprehensive and cover various scenarios, including a fallback for unknown reasons. I've added one suggestion to improve the maintainability of the new tests by reducing code duplication.
Import and use the actual _FINISH_REASON_MAPPING from lite_llm instead of duplicating it in tests. This ensures tests stay in sync with implementation changes automatically. Benefits: - Single source of truth for finish_reason mappings - Tests automatically reflect any future mapping changes - Reduced code duplication - Better maintainability Addresses review comment: google#3114 (review)
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request effectively addresses the issue of propagating finish_reason from LiteLLM responses. The introduction of a mapping to FinishReason enums standardizes the field, which is a solid approach for consistency across different models. The changes are well-tested with new unit tests covering various scenarios. The fix in the telemetry code to correctly handle the enum is also a valuable improvement. Overall, this is a well-executed change that enhances the framework's robustness.
The Union type is no longer needed since finish_reason is always a FinishReason enum (never a string after our mapping). Addresses review comment: google#3114 (comment)
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request effectively propagates the finish_reason from LiteLLM responses, ensuring consistency by mapping string values to the FinishReason enum. The changes are well-implemented across the model, telemetry, and response modules. The addition of comprehensive unit tests is excellent and covers various scenarios, including the fallback for unknown reasons. I've identified one potential issue regarding list handling that could lead to an IndexError.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request effectively addresses the propagation of finish_reason from LiteLLM responses by mapping them to the standard types.FinishReason enum. The implementation is robust, handling various completion scenarios and unknown reasons gracefully, which ensures type consistency across different model providers. The corresponding fix in the telemetry code to correctly process the enum value is a crucial improvement. The addition of comprehensive unit tests is excellent, providing strong validation for the new logic. Overall, this is a high-quality change that enhances the framework's reliability and observability.
3 participants
Summary
Fixes #3109
This PR ensures that the
finish_reasonfield from LiteLLM responses is properly propagated toLlmResponseobjects, enablingafter_model_callbackfunctions to detect completion conditions like max_tokens truncation.Changes
src/google/adk/models/lite_llm.py
_model_response_to_generate_content_response()to extractfinish_reasonfrom LiteLLM responsellm_response.finish_reasonwhen present in the responsesrc/google/adk/telemetry/tracing.py
trace_call_llm()to handle both enum (Gemini) and string (LiteLLM) finish_reason valueshasattr()check to detect enum vs string typetests/unittests/models/test_litellm.py
Test Results
All tests pass:
Impact
This fix allows
after_model_callbackfunctions to properly detect:"length": max_tokens limit reached"stop": natural completion"tool_calls": tool invocations"content_filter": filtered contentThis enables implementing retry logic for incomplete responses, logging completion statistics, and handling different completion conditions appropriately.