Skip to content

Commit 868380a

Browse files
authored
feat(genai): handle all finish messages for AQA (#1201)
`_get_finish_message` only handled 3 of 13 possible finish reasons, leading to other reasons to finish with "Unexpected generation error" Addresses TODO; looks to handle directly from field if possible, otherwise falling back (like before) but with all reasons mapped.
1 parent a09f070 commit 868380a

File tree

1 file changed

+25
-6
lines changed

1 file changed

+25
-6
lines changed

libs/genai/langchain_google_genai/_genai_extension.py

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -632,22 +632,41 @@ def generate_answer(
632632
)
633633

634634

635-
# TODO: Use candidate.finish_message when that field is launched.
636-
# For now, we derive this message from other existing fields.
637635
def _get_finish_message(candidate: genai.Candidate) -> str:
636+
"""Get a human-readable finish message from the candidate.
637+
638+
Uses the official finish_message field if available, otherwise falls back
639+
to a manual mapping of finish reasons to descriptive messages.
640+
"""
641+
# Use the official field when available
642+
if hasattr(candidate, "finish_message") and candidate.finish_message:
643+
return candidate.finish_message
644+
645+
# Fallback to manual mapping for all known finish reasons
638646
finish_messages: Dict[int, str] = {
647+
genai.Candidate.FinishReason.STOP: "Generation completed successfully",
639648
genai.Candidate.FinishReason.MAX_TOKENS: (
640649
"Maximum token in context window reached"
641650
),
642651
genai.Candidate.FinishReason.SAFETY: "Blocked because of safety",
643652
genai.Candidate.FinishReason.RECITATION: "Blocked because of recitation",
653+
genai.Candidate.FinishReason.LANGUAGE: "Unsupported language detected",
654+
genai.Candidate.FinishReason.BLOCKLIST: "Content hit forbidden terms",
655+
genai.Candidate.FinishReason.PROHIBITED_CONTENT: (
656+
"Inappropriate content detected"
657+
),
658+
genai.Candidate.FinishReason.SPII: "Sensitive personal information detected",
659+
genai.Candidate.FinishReason.IMAGE_SAFETY: "Image safety violation",
660+
genai.Candidate.FinishReason.MALFORMED_FUNCTION_CALL: "Malformed function call",
661+
genai.Candidate.FinishReason.UNEXPECTED_TOOL_CALL: "Unexpected tool call",
662+
genai.Candidate.FinishReason.OTHER: "Other generation issue",
663+
genai.Candidate.FinishReason.FINISH_REASON_UNSPECIFIED: (
664+
"Unspecified finish reason"
665+
),
644666
}
645667

646668
finish_reason = candidate.finish_reason
647-
if finish_reason not in finish_messages:
648-
return "Unexpected generation error"
649-
650-
return finish_messages[finish_reason]
669+
return finish_messages.get(finish_reason, "Unexpected generation error")
651670

652671

653672
def _convert_to_metadata(metadata: Dict[str, Any]) -> List[genai.CustomMetadata]:

0 commit comments

Comments
 (0)