mlflow--mlflow/mlflow/langchain/utils/chat.py

import json
import logging
import time
from collections import defaultdict
from collections.abc import Iterator
from typing import Any

import pydantic
from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    FunctionMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langchain_core.messages import (
    ChatMessage as LangChainChatMessage,
)
from langchain_core.outputs import ChatGenerationChunk
from langchain_core.outputs.generation import Generation

from mlflow.environment_variables import MLFLOW_CONVERT_MESSAGES_DICT_FOR_LANGCHAIN
from mlflow.exceptions import MlflowException
from mlflow.tracing.constant import TokenUsageKey
from mlflow.types.chat import (
    AudioContentPart,
    ChatChoice,
    ChatChoiceDelta,
    ChatChunkChoice,
    ChatCompletionChunk,
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatMessage,
    ChatUsage,
    InputAudio,
)

_logger = logging.getLogger(__name__)


_TOKEN_USAGE_KEY_MAPPING = {
    # OpenAI
    "prompt_tokens": TokenUsageKey.INPUT_TOKENS,
    "completion_tokens": TokenUsageKey.OUTPUT_TOKENS,
    "total_tokens": TokenUsageKey.TOTAL_TOKENS,
    # OpenAI Streaming, Anthropic, etc.
    "input_tokens": TokenUsageKey.INPUT_TOKENS,
    "output_tokens": TokenUsageKey.OUTPUT_TOKENS,
    # Anthropic
    "cache_read_input_tokens": TokenUsageKey.CACHE_READ_INPUT_TOKENS,
    "cache_creation_input_tokens": TokenUsageKey.CACHE_CREATION_INPUT_TOKENS,
    # Gemini
    "cached_content_token_count": TokenUsageKey.CACHE_READ_INPUT_TOKENS,
}


# Maps MIME subtypes to the formats InputAudio accepts: Literal["wav", "mp3"].
# Identity mappings (e.g. "wav" -> "wav") are handled by the fallback in _normalize_content.
_MIME_TO_AUDIO_FORMAT: dict[str, str] = {
    "x-wav": "wav",
    "mpeg": "mp3",
}


def _normalize_content(
    content: str | list[dict[str, Any]],
) -> str | list[dict[str, Any]]:
    """
    Normalize multi-modal content blocks from LangChain's format to MLflow's expected format.

    LangChain uses:

        {"type": "audio", "source_type": "base64", "data": "...", "mime_type": "audio/wav"}

    while MLflow expects:

        {"type": "input_audio", "input_audio": {"data": "...", "format": "wav"}}

    This function converts audio blocks to MLflow's format and returns the normalized content
    so that it can be validated by :class:`~mlflow.types.chat.ChatMessage`.
    """
    if isinstance(content, str):
        return content

    normalized = []
    for block in content:
        match block:
            case {
                "type": "audio",
                "source_type": "base64",
                "mime_type": str(mime_type),
                "data": str(data),
            }:
                # Extract and normalize format from mime_type (e.g. "audio/wav" -> "wav",
                # "audio/mpeg" -> "mp3"). Strip parameters like "; codecs=..."
                raw_subtype = mime_type.rsplit("/", 1)[-1].split(";")[0].strip()
                audio_format = _MIME_TO_AUDIO_FORMAT.get(raw_subtype, raw_subtype)

                try:
                    audio_part = AudioContentPart(
                        type="input_audio",
                        input_audio=InputAudio(
                            data=data,
                            format=audio_format,
                        ),
                    )
                except pydantic.ValidationError as e:
                    raise MlflowException.invalid_parameter_value(
                        f"Unsupported audio format {audio_format!r} derived from "
                        f"mime_type {mime_type!r}. Supported formats: 'wav', 'mp3'."
                    ) from e
                normalized.append(audio_part.model_dump())
            case {"type": "audio"}:
                raise MlflowException.invalid_parameter_value(
                    "Unsupported LangChain audio content. Only base64-encoded audio with a valid "
                    "mime_type is supported for conversion to MLflow chat messages."
                )
            case _:
                normalized.append(block)

    return normalized


def _extract_nested_token_details(d: dict[str, Any]) -> Iterator[tuple[str, int]]:
    """Extract cached token counts from nested detail dicts."""
    match d:
        case {"input_token_details": {"cache_read": int(tokens)}}:
            yield (TokenUsageKey.CACHE_READ_INPUT_TOKENS, tokens)
    match d:
        case {"input_token_details": {"cache_creation": int(tokens)}}:
            yield (TokenUsageKey.CACHE_CREATION_INPUT_TOKENS, tokens)
    match d:
        case {"prompt_tokens_details": {"cached_tokens": int(tokens)}}:
            yield (TokenUsageKey.CACHE_READ_INPUT_TOKENS, tokens)


def convert_lc_message_to_chat_message(lc_message: BaseMessage) -> ChatMessage:
    """
    Convert LangChain's message format to the MLflow's standard chat message format.
    """
    if isinstance(lc_message, AIMessage):
        if tool_calls := _get_tool_calls_from_ai_message(lc_message):
            content = lc_message.content
            # For Anthropic model tool calls are returned twice so we need to filter them out
            if isinstance(content, list):
                content = [c for c in content if c["type"] != "tool_use"]
            content = _normalize_content(content)
            return ChatMessage(
                role="assistant",
                # If tool calls present, content null value should be None not empty string
                # according to the OpenAI spec, which ChatMessage is following
                # Ref: https://github.com/langchain-ai/langchain/blob/32917a0b98cb8edcfb8d0e84f0878434e1c3f192/libs/partners/openai/langchain_openai/chat_models/base.py#L116-L117
                content=content or None,
                tool_calls=tool_calls,
            )
        else:
            return ChatMessage(role="assistant", content=_normalize_content(lc_message.content))
    elif isinstance(lc_message, LangChainChatMessage):
        return ChatMessage(role=lc_message.role, content=_normalize_content(lc_message.content))
    elif isinstance(lc_message, FunctionMessage):
        return ChatMessage(role="function", content=_normalize_content(lc_message.content))
    elif isinstance(lc_message, ToolMessage):
        return ChatMessage(
            role="tool",
            content=_normalize_content(lc_message.content),
            tool_call_id=lc_message.tool_call_id,
        )
    elif isinstance(lc_message, HumanMessage):
        return ChatMessage(role="user", content=_normalize_content(lc_message.content))
    elif isinstance(lc_message, SystemMessage):
        return ChatMessage(role="system", content=_normalize_content(lc_message.content))
    else:
        raise MlflowException.invalid_parameter_value(
            f"Unexpected message type. Expected a BaseMessage subclass, but got: {type(lc_message)}"
        )


def _chat_model_to_langchain_message(message: ChatMessage) -> BaseMessage:
    """
    Convert the MLflow's standard chat message format to LangChain's message format.
    """
    if message.role == "system":
        return SystemMessage(content=message.content)
    elif message.role == "assistant":
        return AIMessage(content=message.content)
    elif message.role == "user":
        return HumanMessage(content=message.content)
    elif message.role == "tool":
        return ToolMessage(content=message.content, tool_call_id=message.tool_call_id)
    elif message.role == "function":
        return FunctionMessage(content=message.content)
    else:
        raise MlflowException.invalid_parameter_value(
            f"Unrecognized chat message role: {message.role}"
        )


def _get_tool_calls_from_ai_message(message: AIMessage) -> list[dict[str, Any]]:
    # Extract tool calls from AIMessage
    tool_calls = [
        {
            "type": "function",
            "id": tc["id"],
            "function": {
                "name": tc["name"],
                "arguments": json.dumps(tc["args"]),
            },
        }
        for tc in message.tool_calls
    ]

    invalid_tool_calls = [
        {
            "type": "function",
            "id": tc["id"],
            "function": {
                "name": tc["name"],
                "arguments": tc["args"],
            },
        }
        for tc in message.invalid_tool_calls
    ]

    if tool_calls or invalid_tool_calls:
        return tool_calls + invalid_tool_calls

    # Get tool calls from additional kwargs if present.
    return [
        {
            k: v
            for k, v in tool_call.items()  # type: ignore[union-attr]
            if k in {"id", "type", "function"}
        }
        for tool_call in message.additional_kwargs.get("tool_calls", [])
    ]


def try_transform_response_to_chat_format(response: Any) -> dict[str, Any]:
    """
    Try to convert the response to the standard chat format and return its dict representation.

    If the response is not one of the supported types, return the response as-is.
    """
    if isinstance(response, (str, AIMessage)):
        if isinstance(response, str):
            message_id = None
            message = ChatMessage(role="assistant", content=response)
        else:
            message_id = getattr(response, "id", None)
            message = convert_lc_message_to_chat_message(response)

        transformed_response = ChatCompletionResponse(
            id=message_id,
            created=int(time.time()),
            model="",
            object="chat.completion",
            choices=[
                ChatChoice(
                    index=0,
                    message=message,
                    finish_reason=None,
                )
            ],
            usage=ChatUsage(
                prompt_tokens=None,
                completion_tokens=None,
                total_tokens=None,
            ),
        )
        return transformed_response.model_dump(mode="json", exclude_unset=True)
    else:
        return response


def try_transform_response_iter_to_chat_format(chunk_iter):
    from langchain_core.messages.ai import AIMessageChunk

    def _gen_converted_chunk(message_content, message_id, finish_reason):
        transformed_response = ChatCompletionChunk(
            id=message_id,
            object="chat.completion.chunk",
            created=int(time.time()),
            model="",
            choices=[
                ChatChunkChoice(
                    index=0,
                    delta=ChatChoiceDelta(
                        role="assistant",
                        content=message_content,
                    ),
                    finish_reason=finish_reason,
                )
            ],
        )

        return transformed_response.model_dump(mode="json", exclude_unset=True)

    def _convert(chunk):
        if isinstance(chunk, str):
            message_content = chunk
            message_id = None
            finish_reason = None
        elif isinstance(chunk, AIMessageChunk):
            message_content = chunk.content
            message_id = getattr(chunk, "id", None)

            if response_metadata := getattr(chunk, "response_metadata", None):
                finish_reason = response_metadata.get("finish_reason")
            else:
                finish_reason = None
        elif isinstance(chunk, AIMessage):
            # The langchain chat model does not support stream
            # so `model.stream` returns the whole result.
            message_content = chunk.content
            message_id = getattr(chunk, "id", None)
            finish_reason = "stop"
        else:
            return chunk
        return _gen_converted_chunk(
            message_content,
            message_id=message_id,
            finish_reason=finish_reason,
        )

    return map(_convert, chunk_iter)


def _convert_chat_request_or_throw(
    chat_request: dict[str, Any],
) -> list[BaseMessage]:
    model = ChatCompletionRequest.model_validate(chat_request)
    return [_chat_model_to_langchain_message(message) for message in model.messages]


def _convert_chat_request(chat_request: dict[str, Any] | list[dict[str, Any]]):
    if isinstance(chat_request, list):
        return [_convert_chat_request_or_throw(request) for request in chat_request]
    else:
        return _convert_chat_request_or_throw(chat_request)


def _get_lc_model_input_fields(lc_model) -> set[str]:
    try:
        if hasattr(lc_model, "input_schema"):
            return set(lc_model.input_schema.model_fields)
    except Exception as e:
        _logger.debug(
            f"Unexpected exception while checking LangChain input schema for"
            f" request transformation: {e}"
        )

    return set()


def _should_transform_request_json_for_chat(lc_model):
    # Don't convert the request to LangChain's Message format for LangGraph models.
    # Inputs may have key like "messages", but they are graph state fields, not OAI chat format.
    try:
        from langgraph.graph.state import CompiledStateGraph

        if isinstance(lc_model, CompiledStateGraph):
            return False
    except ImportError:
        pass

    # Avoid converting the request to LangChain's Message format if the chain
    # is an AgentExecutor, as LangChainChatMessage might not be accepted by the chain
    from mlflow.langchain._compat import try_import_agent_executor

    AgentExecutor = try_import_agent_executor()
    if AgentExecutor and isinstance(lc_model, AgentExecutor):
        return False

    input_fields = _get_lc_model_input_fields(lc_model)
    if "messages" in input_fields:
        # If the chain accepts a "messages" field directly, don't attempt to convert
        # the request to LangChain's Message format automatically. Assume that the chain
        # is handling the "messages" field by itself
        return False

    return True


def transform_request_json_for_chat_if_necessary(request_json, lc_model):
    """
    Convert the input request JSON to LangChain's Message format if the LangChain model
    accepts ChatMessage objects (e.g. AIMessage, HumanMessage, SystemMessage) as input.

    Args:
        request_json: The input request JSON.
        lc_model: The LangChain model.

    Returns:
        A 2-element tuple containing:

            1. The new request.
            2. A boolean indicating whether or not the request was transformed from the OpenAI
            chat format.
    """

    def json_dict_might_be_chat_request(json_message):
        return (
            isinstance(json_message, dict)
            and "messages" in json_message
            and
            # Additional keys can't be specified when calling LangChain invoke() / batch()
            # with chat messages
            len(json_message) == 1
            # messages field should be a list
            and isinstance(json_message["messages"], list)
        )

    def is_list_of_chat_messages(json_message: list[dict[str, Any]]):
        return isinstance(json_message, list) and all(
            json_dict_might_be_chat_request(message) for message in json_message
        )

    should_convert = MLFLOW_CONVERT_MESSAGES_DICT_FOR_LANGCHAIN.get()
    if should_convert is None:
        should_convert = _should_transform_request_json_for_chat(lc_model) and (
            json_dict_might_be_chat_request(request_json) or is_list_of_chat_messages(request_json)
        )
        if should_convert:
            _logger.debug(
                "Converting the request JSON to LangChain's Message format. "
                "To disable this conversion, set the environment variable "
                f"`{MLFLOW_CONVERT_MESSAGES_DICT_FOR_LANGCHAIN}` to 'false'."
            )

    if should_convert:
        try:
            return _convert_chat_request(request_json), True
        except pydantic.ValidationError:
            _logger.debug(
                "Failed to convert the request JSON to LangChain's Message format. "
                "The request will be passed to the LangChain model as-is. ",
                exc_info=True,
            )
            return request_json, False
    else:
        return request_json, False


def parse_token_usage(
    lc_generations: list[Generation],
) -> dict[str, int] | None:
    """Parse the token usage from the LangChain generations."""

    # Check if this is streaming (contains ChatGenerationChunk)
    is_streaming = any(isinstance(gen, ChatGenerationChunk) for gen in lc_generations)

    if is_streaming:
        # Streaming mode: collect all generations with usage, use only the last one
        # (which contains the final cumulative token counts)
        generations_with_usage = [
            token_usage
            for generation in lc_generations
            if (token_usage := _parse_token_usage_from_generation(generation))
        ]

        if generations_with_usage:
            return generations_with_usage[-1]
        return None

    # Non-streaming mode: existing behavior (sum all generations)
    aggregated = defaultdict(int)
    for generation in lc_generations:
        if token_usage := _parse_token_usage_from_generation(generation):
            for key in token_usage:
                aggregated[key] += token_usage[key]

    return dict(aggregated) if aggregated else None


def _parse_token_usage_from_generation(
    generation: Generation,
) -> dict[str, int] | None:
    message = getattr(generation, "message", None)
    if not message:
        return None

    metadata = (
        message.usage_metadata
        or message.response_metadata.get("usage")
        or message.response_metadata.get("token_usage")
    )
    return _parse_token_counts(metadata) if metadata else None


def _parse_token_counts(usage_metadata: dict[str, Any]) -> dict[str, int]:
    """Standardize token usage metadata keys to MLflow's token usage keys."""
    usage = {}
    for key, value in usage_metadata.items():
        if usage_key := _TOKEN_USAGE_KEY_MAPPING.get(key):
            usage[usage_key] = value

    # Extract from nested detail dicts (e.g. input_token_details.cache_read).
    # Uses setdefault so flat keys above take priority.
    for usage_key, value in _extract_nested_token_details(usage_metadata):
        usage.setdefault(usage_key, value)

    # If the total tokens are not present, calculate it from the input and output tokens
    if usage and usage.get(TokenUsageKey.TOTAL_TOKENS) is None:
        usage[TokenUsageKey.TOTAL_TOKENS] = usage.get(TokenUsageKey.INPUT_TOKENS, 0) + usage.get(
            TokenUsageKey.OUTPUT_TOKENS, 0
        )

    return usage