aidial_adapter_vertexai/chat/gemini/finish_reason.py

from typing import assert_never

import vertexai.preview.generative_models as generative_models
from aidial_sdk.chat_completion import FinishReason
from google.genai.types import FinishReason as GenAIFinishReason

from aidial_adapter_vertexai.chat.gemini.error import FinishReasonOtherError

GenFinishReason = generative_models.FinishReason


_EARLY_TERMINATION_ERROR = "The model terminated generation unexpectedly"
_INVALID_FUNCTION_CALL_ERROR = (
    "The function call generated by the model is invalid"
)


def to_openai_finish_reason(
    finish_reason: GenFinishReason, retriable: bool
) -> FinishReason | None:
    match finish_reason:
        case GenFinishReason.FINISH_REASON_UNSPECIFIED:
            return None
        case GenFinishReason.MAX_TOKENS:
            return FinishReason.LENGTH
        case GenFinishReason.STOP:
            return FinishReason.STOP
        case (
            GenFinishReason.SAFETY
            | GenFinishReason.RECITATION
            | GenFinishReason.BLOCKLIST
            | GenFinishReason.PROHIBITED_CONTENT
            | GenFinishReason.SPII
        ):
            return FinishReason.CONTENT_FILTER

        # The following finish reasons could be usually fixed with a retry
        case GenFinishReason.OTHER:
            raise FinishReasonOtherError(
                msg=_EARLY_TERMINATION_ERROR,
                retriable=retriable,
            )
        case GenFinishReason.MALFORMED_FUNCTION_CALL:
            raise FinishReasonOtherError(
                msg=_INVALID_FUNCTION_CALL_ERROR,
                retriable=retriable,
            )
        case _:
            assert_never(finish_reason)


def genai_to_openai_finish_reason(
    finish_reason: GenAIFinishReason | None, retriable: bool
) -> FinishReason | None:
    if not finish_reason:
        return None
    match finish_reason:
        case "FINISH_REASON_UNSPECIFIED":
            return None
        case "MAX_TOKENS":
            return FinishReason.LENGTH
        case "STOP":
            return FinishReason.STOP
        case (
            "SAFETY"
            | "RECITATION"
            | "BLOCKLIST"
            | "PROHIBITED_CONTENT"
            | "SPII"
        ):
            return FinishReason.CONTENT_FILTER
        case "OTHER":
            raise FinishReasonOtherError(
                msg=_EARLY_TERMINATION_ERROR,
                retriable=retriable,
            )
        case "MALFORMED_FUNCTION_CALL":
            raise FinishReasonOtherError(
                msg=_INVALID_FUNCTION_CALL_ERROR,
                retriable=retriable,
            )
        case _:
            assert_never(finish_reason)