Skip to content

Commit a1c0938

Browse files
aperepellizzij
authored andcommitted
feat: map LiteLLM finish_reason strings to FinishReason enum
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
1 parent 7996112 commit a1c0938

3 files changed

Lines changed: 20 additions & 15 deletions

File tree

src/google/adk/models/lite_llm.py

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,19 @@
6464
_NEW_LINE = "\n"
6565
_EXCLUDED_PART_FIELD = {"inline_data": {"data"}}
6666

67+
# Mapping of LiteLLM finish_reason strings to FinishReason enum values
68+
# Note: tool_calls/function_call map to STOP because:
69+
# 1. FinishReason.TOOL_CALL enum does not exist (as of google-genai 0.8.0)
70+
# 2. Tool calls represent normal completion (model stopped to invoke tools)
71+
# 3. Gemini native responses use STOP for tool calls (see lite_llm.py:910)
72+
_FINISH_REASON_MAPPING = {
73+
"length": types.FinishReason.MAX_TOKENS,
74+
"stop": types.FinishReason.STOP,
75+
"tool_calls": types.FinishReason.STOP, # Normal completion with tool invocation
76+
"function_call": types.FinishReason.STOP, # Legacy function call variant
77+
"content_filter": types.FinishReason.SAFETY,
78+
}
79+
6780

6881
class ChatCompletionFileUrlObject(TypedDict, total=False):
6982
file_data: str
@@ -508,18 +521,9 @@ def _model_response_to_generate_content_response(
508521
# Map LiteLLM finish_reason strings to FinishReason enum
509522
# This provides type consistency with Gemini native responses and avoids warnings
510523
finish_reason_str = str(finish_reason).lower()
511-
if finish_reason_str == "length":
512-
llm_response.finish_reason = types.FinishReason.MAX_TOKENS
513-
elif finish_reason_str == "stop":
514-
llm_response.finish_reason = types.FinishReason.STOP
515-
elif "tool" in finish_reason_str or "function" in finish_reason_str:
516-
# Handle tool_calls, function_call variants
517-
llm_response.finish_reason = types.FinishReason.STOP
518-
elif finish_reason_str == "content_filter":
519-
llm_response.finish_reason = types.FinishReason.SAFETY
520-
else:
521-
# For unknown reasons, use OTHER
522-
llm_response.finish_reason = types.FinishReason.OTHER
524+
llm_response.finish_reason = _FINISH_REASON_MAPPING.get(
525+
finish_reason_str, types.FinishReason.OTHER
526+
)
523527
if response.get("usage", None):
524528
llm_response.usage_metadata = types.GenerateContentResponseUsageMetadata(
525529
prompt_token_count=response["usage"].get("prompt_tokens", 0),

src/google/adk/models/llm_response.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,10 +78,11 @@ class LlmResponse(BaseModel):
7878
Only used for streaming mode.
7979
"""
8080

81-
finish_reason: Optional[Union[types.FinishReason, str]] = None
81+
finish_reason: Optional[types.FinishReason] = None
8282
"""The finish reason of the response.
8383
84-
Can be either a types.FinishReason enum (from Gemini) or a string (from LiteLLM).
84+
Always a types.FinishReason enum. String values from underlying model providers
85+
are mapped to corresponding enum values (with fallback to OTHER for unknown values).
8586
"""
8687

8788
error_code: Optional[str] = None

src/google/adk/telemetry/tracing.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -306,7 +306,7 @@ def trace_call_llm(
306306
if isinstance(llm_response.finish_reason, types.FinishReason):
307307
finish_reason_str = llm_response.finish_reason.name.lower()
308308
else:
309-
# Fallback for string values (should not occur with LiteLLM after enum mapping)
309+
# Defensive fallback for string values (should never occur - all values mapped to enum)
310310
finish_reason_str = str(llm_response.finish_reason).lower()
311311
span.set_attribute(
312312
'gen_ai.response.finish_reasons',

0 commit comments

Comments
 (0)