star

venv/lib/python3.9/site-packages/mlflow/langchain/utils/__init__.py

@@ -0,0 +1,664 @@
+"""Utility functions for mlflow.langchain."""
+
+import contextlib
+import functools
+import importlib
+import json
+import logging
+import os
+import re
+import shutil
+import types
+import warnings
+from functools import lru_cache
+from importlib.util import find_spec
+from typing import Callable, NamedTuple
+
+import cloudpickle
+import yaml
+from packaging import version
+from packaging.version import Version
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _validate_and_get_model_code_path
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.utils.class_utils import _get_class_from_string
+
+_AGENT_PRIMITIVES_FILE_NAME = "agent_primitive_args.json"
+_AGENT_PRIMITIVES_DATA_KEY = "agent_primitive_data"
+_AGENT_DATA_FILE_NAME = "agent.yaml"
+_AGENT_DATA_KEY = "agent_data"
+_TOOLS_DATA_FILE_NAME = "tools.pkl"
+_TOOLS_DATA_KEY = "tools_data"
+_LOADER_FN_FILE_NAME = "loader_fn.pkl"
+_LOADER_FN_KEY = "loader_fn"
+_LOADER_ARG_KEY = "loader_arg"
+_PERSIST_DIR_NAME = "persist_dir_data"
+_PERSIST_DIR_KEY = "persist_dir"
+_MODEL_DATA_YAML_FILE_NAME = "model.yaml"
+_MODEL_DATA_PKL_FILE_NAME = "model.pkl"
+_MODEL_DATA_FOLDER_NAME = "model"
+_MODEL_DATA_KEY = "model_data"
+_MODEL_TYPE_KEY = "model_type"
+_RUNNABLE_LOAD_KEY = "runnable_load"
+_BASE_LOAD_KEY = "base_load"
+_CONFIG_LOAD_KEY = "config_load"
+_PICKLE_LOAD_KEY = "pickle_load"
+_MODEL_LOAD_KEY = "model_load"
+_UNSUPPORTED_MODEL_WARNING_MESSAGE = (
+    "MLflow does not guarantee support for Chains outside of the subclasses of LLMChain, found %s"
+)
+_UNSUPPORTED_LLM_WARNING_MESSAGE = (
+    "MLflow does not guarantee support for LLMs outside of HuggingFacePipeline and OpenAI, found %s"
+)
+
+
+_CHAT_MODELS_ERROR_MSG = re.compile("Loading (openai-chat|azure-openai-chat) LLM not supported")
+
+
+try:
+    import langchain_community
+
+    # Since langchain-community 0.0.27, saving or loading a module that relies on the pickle
+    # deserialization requires passing `allow_dangerous_deserialization=True`.
+    IS_PICKLE_SERIALIZATION_RESTRICTED = Version(langchain_community.__version__) >= Version(
+        "0.0.27"
+    )
+except ImportError:
+    IS_PICKLE_SERIALIZATION_RESTRICTED = False
+
+logger = logging.getLogger(__name__)
+
+
+@lru_cache
+def base_lc_types():
+    import langchain.agents.agent
+    import langchain.chains.base
+    import langchain.schema
+
+    return (
+        langchain.chains.base.Chain,
+        langchain.agents.agent.AgentExecutor,
+        langchain.schema.BaseRetriever,
+    )
+
+
+@lru_cache
+def picklable_runnable_types():
+    """
+    Runnable types that can be pickled and unpickled by cloudpickle.
+    """
+    from langchain.chat_models.base import SimpleChatModel
+    from langchain.prompts import ChatPromptTemplate
+    from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
+
+    return (
+        SimpleChatModel,
+        ChatPromptTemplate,
+        RunnablePassthrough,
+        RunnableLambda,
+    )
+
+
+@lru_cache
+def lc_runnable_with_steps_types():
+    from langchain.schema.runnable import RunnableParallel, RunnableSequence
+
+    return (RunnableParallel, RunnableSequence)
+
+
+def lc_runnable_assign_types():
+    from langchain.schema.runnable.passthrough import RunnableAssign
+
+    return (RunnableAssign,)
+
+
+def lc_runnable_branch_types():
+    from langchain.schema.runnable import RunnableBranch
+
+    return (RunnableBranch,)
+
+
+def lc_runnable_binding_types():
+    from langchain.schema.runnable import RunnableBinding
+
+    return (RunnableBinding,)
+
+
+def lc_runnables_types():
+    return (
+        picklable_runnable_types()
+        + lc_runnable_with_steps_types()
+        + lc_runnable_branch_types()
+        + lc_runnable_assign_types()
+        + lc_runnable_binding_types()
+    )
+
+
+def langgraph_types():
+    try:
+        from langgraph.graph.graph import CompiledGraph
+
+        return (CompiledGraph,)
+    except ImportError:
+        return ()
+
+
+def supported_lc_types():
+    return base_lc_types() + lc_runnables_types() + langgraph_types()
+
+
+# Wrapping as a function to avoid callign supported_lc_types() at import time
+def get_unsupported_model_message(model_type):
+    return (
+        "MLflow langchain flavor only supports subclasses of "
+        f"{supported_lc_types()}, found {model_type}."
+    )
+
+
+@lru_cache
+def custom_type_to_loader_dict():
+    # helper function to load output_parsers from config
+    def _load_output_parser(config: dict) -> dict:
+        """Load output parser."""
+        from langchain.schema.output_parser import StrOutputParser
+
+        output_parser_type = config.pop("_type", None)
+        if output_parser_type == "default":
+            return StrOutputParser(**config)
+        else:
+            raise ValueError(f"Unsupported output parser {output_parser_type}")
+
+    return {"default": _load_output_parser}
+
+
+class _SpecialChainInfo(NamedTuple):
+    loader_arg: str
+
+
+def _get_special_chain_info_or_none(chain):
+    for (
+        special_chain_class,
+        loader_arg,
+    ) in _get_map_of_special_chain_class_to_loader_arg().items():
+        if isinstance(chain, special_chain_class):
+            return _SpecialChainInfo(loader_arg=loader_arg)
+
+
+@lru_cache
+def _get_map_of_special_chain_class_to_loader_arg():
+    import langchain
+
+    from mlflow.langchain.retriever_chain import _RetrieverChain
+
+    class_name_to_loader_arg = {
+        "langchain.chains.RetrievalQA": "retriever",
+        "langchain.chains.APIChain": "requests_wrapper",
+        "langchain.chains.HypotheticalDocumentEmbedder": "embeddings",
+    }
+    # NB: SQLDatabaseChain was migrated to langchain_experimental beginning with version 0.0.247
+    if version.parse(langchain.__version__) <= version.parse("0.0.246"):
+        class_name_to_loader_arg["langchain.chains.SQLDatabaseChain"] = "database"
+    else:
+        if find_spec("langchain_experimental"):
+            # Add this entry only if langchain_experimental is installed
+            class_name_to_loader_arg["langchain_experimental.sql.SQLDatabaseChain"] = "database"
+
+    class_to_loader_arg = {
+        _RetrieverChain: "retriever",
+    }
+    for class_name, loader_arg in class_name_to_loader_arg.items():
+        try:
+            cls = _get_class_from_string(class_name)
+            class_to_loader_arg[cls] = loader_arg
+        except Exception:
+            logger.warning(
+                "Unexpected import failure for class '%s'. Please file an issue at"
+                " https://github.com/mlflow/mlflow/issues/.",
+                class_name,
+                exc_info=True,
+            )
+
+    return class_to_loader_arg
+
+
+@lru_cache
+def _get_supported_llms():
+    supported_llms = set()
+
+    def try_adding_llm(module, class_name):
+        if cls := getattr(module, class_name, None):
+            supported_llms.add(cls)
+
+    def safe_import_and_add(module_name, class_name):
+        """Add conditional support for `partner` and `community` APIs in langchain"""
+        try:
+            module = importlib.import_module(module_name)
+            try_adding_llm(module, class_name)
+        except ImportError:
+            pass
+
+    safe_import_and_add("langchain.llms.openai", "OpenAI")
+    # HuggingFacePipeline is moved to langchain_huggingface since langchain 0.2.0
+    safe_import_and_add("langchain.llms", "HuggingFacePipeline")
+    safe_import_and_add("langchain.langchain_huggingface", "HuggingFacePipeline")
+    safe_import_and_add("langchain_openai", "OpenAI")
+    safe_import_and_add("langchain_databricks", "ChatDatabricks")
+
+    for llm_name in ["Databricks", "Mlflow"]:
+        safe_import_and_add("langchain.llms", llm_name)
+
+    for chat_model_name in [
+        "ChatDatabricks",
+        "ChatMlflow",
+        "ChatOpenAI",
+        "AzureChatOpenAI",
+    ]:
+        safe_import_and_add("langchain.chat_models", chat_model_name)
+
+    return supported_llms
+
+
+def _agent_executor_contains_unsupported_llm(lc_model, _SUPPORTED_LLMS):
+    import langchain.agents.agent
+
+    return (
+        isinstance(lc_model, langchain.agents.agent.AgentExecutor)
+        # 'RunnableMultiActionAgent' object has no attribute 'llm_chain'
+        and hasattr(lc_model.agent, "llm_chain")
+        and not any(
+            isinstance(lc_model.agent.llm_chain.llm, supported_llm)
+            for supported_llm in _SUPPORTED_LLMS
+        )
+    )
+
+
+# temp_dir is only required when lc_model could be a file path
+def _validate_and_prepare_lc_model_or_path(lc_model, loader_fn, temp_dir=None):
+    import langchain.agents.agent
+    import langchain.chains.base
+    import langchain.chains.llm
+    import langchain.llms.huggingface_hub
+    import langchain.llms.openai
+    import langchain.schema
+
+    # lc_model is a file path
+    if isinstance(lc_model, str):
+        return _validate_and_get_model_code_path(lc_model, temp_dir)
+
+    if not isinstance(lc_model, supported_lc_types()):
+        raise mlflow.MlflowException.invalid_parameter_value(
+            get_unsupported_model_message(type(lc_model).__name__)
+        )
+
+    _SUPPORTED_LLMS = _get_supported_llms()
+    if isinstance(lc_model, langchain.chains.llm.LLMChain) and not any(
+        isinstance(lc_model.llm, supported_llm) for supported_llm in _SUPPORTED_LLMS
+    ):
+        logger.warning(
+            _UNSUPPORTED_LLM_WARNING_MESSAGE,
+            type(lc_model.llm).__name__,
+        )
+
+    if _agent_executor_contains_unsupported_llm(lc_model, _SUPPORTED_LLMS):
+        logger.warning(
+            _UNSUPPORTED_LLM_WARNING_MESSAGE,
+            type(lc_model.agent.llm_chain.llm).__name__,
+        )
+
+    if special_chain_info := _get_special_chain_info_or_none(lc_model):
+        if loader_fn is None:
+            raise mlflow.MlflowException.invalid_parameter_value(
+                f"For {type(lc_model).__name__} models, a `loader_fn` must be provided."
+            )
+        if not isinstance(loader_fn, types.FunctionType):
+            raise mlflow.MlflowException.invalid_parameter_value(
+                "The `loader_fn` must be a function that returns a {loader_arg}.".format(
+                    loader_arg=special_chain_info.loader_arg
+                )
+            )
+
+    # If lc_model is a retriever, wrap it in a _RetrieverChain
+    if isinstance(lc_model, langchain.schema.BaseRetriever):
+        from mlflow.langchain.retriever_chain import _RetrieverChain
+
+        if loader_fn is None:
+            raise mlflow.MlflowException.invalid_parameter_value(
+                f"For {type(lc_model).__name__} models, a `loader_fn` must be provided."
+            )
+        if not isinstance(loader_fn, types.FunctionType):
+            raise mlflow.MlflowException.invalid_parameter_value(
+                "The `loader_fn` must be a function that returns a retriever."
+            )
+        lc_model = _RetrieverChain(retriever=lc_model)
+
+    return lc_model
+
+
+def _save_base_lcs(model, path, loader_fn=None, persist_dir=None):
+    from langchain.agents.agent import AgentExecutor
+    from langchain.chains.base import Chain
+    from langchain.chains.llm import LLMChain
+    from langchain.chat_models.base import BaseChatModel
+
+    model_data_path = os.path.join(path, _MODEL_DATA_YAML_FILE_NAME)
+    model_data_kwargs = {
+        _MODEL_DATA_KEY: _MODEL_DATA_YAML_FILE_NAME,
+        _MODEL_LOAD_KEY: _BASE_LOAD_KEY,
+    }
+
+    if isinstance(model, (LLMChain, BaseChatModel)):
+        model.save(model_data_path)
+    elif isinstance(model, AgentExecutor):
+        if model.agent and getattr(model.agent, "llm_chain", None):
+            model.agent.llm_chain.save(model_data_path)
+
+        if model.agent:
+            agent_data_path = os.path.join(path, _AGENT_DATA_FILE_NAME)
+            model.save_agent(agent_data_path)
+            model_data_kwargs[_AGENT_DATA_KEY] = _AGENT_DATA_FILE_NAME
+
+        if model.tools:
+            tools_data_path = os.path.join(path, _TOOLS_DATA_FILE_NAME)
+            try:
+                with open(tools_data_path, "wb") as f:
+                    cloudpickle.dump(model.tools, f)
+            except Exception as e:
+                raise mlflow.MlflowException(
+                    "Error when attempting to pickle the AgentExecutor tools. "
+                    "This model likely does not support serialization."
+                ) from e
+            model_data_kwargs[_TOOLS_DATA_KEY] = _TOOLS_DATA_FILE_NAME
+        else:
+            raise mlflow.MlflowException.invalid_parameter_value(
+                "For initializing the AgentExecutor, tools must be provided."
+            )
+
+        key_to_ignore = ["llm_chain", "agent", "tools", "callback_manager"]
+        temp_dict = {k: v for k, v in model.__dict__.items() if k not in key_to_ignore}
+
+        agent_primitive_path = os.path.join(path, _AGENT_PRIMITIVES_FILE_NAME)
+        with open(agent_primitive_path, "w") as config_file:
+            json.dump(temp_dict, config_file, indent=4)
+
+        model_data_kwargs[_AGENT_PRIMITIVES_DATA_KEY] = _AGENT_PRIMITIVES_FILE_NAME
+
+    elif special_chain_info := _get_special_chain_info_or_none(model):
+        # Save loader_fn by pickling
+        loader_fn_path = os.path.join(path, _LOADER_FN_FILE_NAME)
+        with open(loader_fn_path, "wb") as f:
+            cloudpickle.dump(loader_fn, f)
+        model_data_kwargs[_LOADER_FN_KEY] = _LOADER_FN_FILE_NAME
+        model_data_kwargs[_LOADER_ARG_KEY] = special_chain_info.loader_arg
+
+        if persist_dir is not None:
+            if os.path.exists(persist_dir):
+                # Save persist_dir by copying into subdir _PERSIST_DIR_NAME
+                persist_dir_data_path = os.path.join(path, _PERSIST_DIR_NAME)
+                shutil.copytree(persist_dir, persist_dir_data_path)
+                model_data_kwargs[_PERSIST_DIR_KEY] = _PERSIST_DIR_NAME
+            else:
+                raise mlflow.MlflowException.invalid_parameter_value(
+                    "The directory provided for persist_dir does not exist."
+                )
+
+        # Save model
+        model.save(model_data_path)
+    elif isinstance(model, Chain):
+        logger.warning(get_unsupported_model_message(type(model).__name__))
+        model.save(model_data_path)
+    else:
+        raise mlflow.MlflowException.invalid_parameter_value(
+            get_unsupported_model_message(type(model).__name__)
+        )
+
+    return model_data_kwargs
+
+
+def _load_from_pickle(path):
+    with open(path, "rb") as f:
+        return cloudpickle.load(f)
+
+
+def _load_from_json(path):
+    with open(path) as f:
+        return json.load(f)
+
+
+def _load_from_yaml(path):
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def _get_path_by_key(root_path, key, conf):
+    key_path = conf.get(key)
+    return os.path.join(root_path, key_path) if key_path else None
+
+
+def _patch_loader(loader_func: Callable) -> Callable:
+    """
+    Patch LangChain loader function like load_chain() to handle the breaking change introduced in
+    LangChain 0.1.12.
+
+    Since langchain-community 0.0.27, loading a module that relies on the pickle deserialization
+    requires the `allow_dangerous_deserialization` flag to be set to True, for security reasons.
+    However, this flag could not be specified via the LangChain's loading API like load_chain(),
+    load_llm(), until LangChain 0.1.14. As a result, such module cannot be loaded with MLflow
+    with earlier version of LangChain and we have to tell the user to upgrade LangChain to 0.0.14
+    or above.
+
+    Args:
+        loader_func: The LangChain loader function to be patched e.g. load_chain().
+
+    Returns:
+        The patched loader function.
+    """
+    if not IS_PICKLE_SERIALIZATION_RESTRICTED:
+        return loader_func
+
+    import langchain
+
+    if Version(langchain.__version__) >= Version("0.1.14"):
+        # For LangChain 0.1.14 and above, we can pass `allow_dangerous_deserialization` flag
+        # via the loader APIs. Since the model is serialized by the user (or someone who has
+        # access to the tracking server), it is safe to set this flag to True.
+        def patched_loader(*args, **kwargs):
+            return loader_func(*args, **kwargs, allow_dangerous_deserialization=True)
+    else:
+
+        def patched_loader(*args, **kwargs):
+            try:
+                return loader_func(*args, **kwargs)
+            except ValueError as e:
+                if "This code relies on the pickle module" in str(e):
+                    raise MlflowException(
+                        "Since langchain-community 0.0.27, loading a module that relies on "
+                        "the pickle deserialization requires the `allow_dangerous_deserialization` "
+                        "flag to be set to True when loading. However, this flag is not supported "
+                        "by the installed version of LangChain. Please upgrade LangChain to 0.1.14 "
+                        "or above by running `pip install langchain>=0.1.14`.",
+                        error_code=INTERNAL_ERROR,
+                    ) from e
+                else:
+                    raise
+
+    return patched_loader
+
+
+def _load_base_lcs(
+    local_model_path,
+    conf,
+):
+    lc_model_path = os.path.join(
+        local_model_path, conf.get(_MODEL_DATA_KEY, _MODEL_DATA_YAML_FILE_NAME)
+    )
+
+    agent_path = _get_path_by_key(local_model_path, _AGENT_DATA_KEY, conf)
+    tools_path = _get_path_by_key(local_model_path, _TOOLS_DATA_KEY, conf)
+    agent_primitive_path = _get_path_by_key(local_model_path, _AGENT_PRIMITIVES_DATA_KEY, conf)
+    loader_fn_path = _get_path_by_key(local_model_path, _LOADER_FN_KEY, conf)
+    persist_dir = _get_path_by_key(local_model_path, _PERSIST_DIR_KEY, conf)
+
+    model_type = conf.get(_MODEL_TYPE_KEY)
+    loader_arg = conf.get(_LOADER_ARG_KEY)
+
+    from langchain.chains.loading import load_chain
+
+    from mlflow.langchain.retriever_chain import _RetrieverChain
+
+    if loader_arg is not None:
+        if loader_fn_path is None:
+            raise mlflow.MlflowException.invalid_parameter_value(
+                "Missing file for loader_fn which is required to build the model."
+            )
+        loader_fn = _load_from_pickle(loader_fn_path)
+        kwargs = {loader_arg: loader_fn(persist_dir)}
+        if model_type == _RetrieverChain.__name__:
+            model = _RetrieverChain.load(lc_model_path, **kwargs).retriever
+        else:
+            model = _patch_loader(load_chain)(lc_model_path, **kwargs)
+    elif agent_path is None and tools_path is None:
+        model = _patch_loader(load_chain)(lc_model_path)
+    else:
+        from langchain.agents import initialize_agent
+
+        llm = _patch_loader(load_chain)(lc_model_path)
+        tools = []
+        kwargs = {}
+
+        if os.path.exists(tools_path):
+            tools = _load_from_pickle(tools_path)
+        else:
+            raise mlflow.MlflowException(
+                "Missing file for tools which is required to build the AgentExecutor object."
+            )
+
+        if os.path.exists(agent_primitive_path):
+            kwargs = _load_from_json(agent_primitive_path)
+
+        model = initialize_agent(tools=tools, llm=llm, agent_path=agent_path, **kwargs)
+    return model
+
+
+def patch_langchain_type_to_cls_dict(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        def _load_chat_openai():
+            from langchain_community.chat_models import ChatOpenAI
+
+            return ChatOpenAI
+
+        def _load_azure_chat_openai():
+            from langchain_community.chat_models import AzureChatOpenAI
+
+            return AzureChatOpenAI
+
+        def _load_chat_databricks():
+            from langchain_databricks import ChatDatabricks
+
+            return ChatDatabricks
+
+        def _patched_get_type_to_cls_dict(original):
+            def _wrapped():
+                return {
+                    **original(),
+                    "openai-chat": _load_chat_openai,
+                    "azure-openai-chat": _load_azure_chat_openai,
+                    "chat-databricks": _load_chat_databricks,
+                }
+
+            return _wrapped
+
+        modules_to_patch = ["langchain.llms", "langchain_community.llms.loading"]
+        originals = {}
+        for name in modules_to_patch:
+            try:
+                module = importlib.import_module(name)
+                originals[name] = module.get_type_to_cls_dict  # Record original impl for cleanup
+            except (ImportError, AttributeError):
+                continue
+            module.get_type_to_cls_dict = _patched_get_type_to_cls_dict(originals[name])
+
+        try:
+            return func(*args, **kwargs)
+        except ValueError as e:
+            if m := _CHAT_MODELS_ERROR_MSG.search(str(e)):
+                model_name = "ChatOpenAI" if m.group(1) == "openai-chat" else "AzureChatOpenAI"
+                raise mlflow.MlflowException(
+                    f"Loading {model_name} chat model is not supported in MLflow with the "
+                    "current version of LangChain. Please upgrade LangChain to 0.0.307 or above "
+                    "by running `pip install langchain>=0.0.307`."
+                ) from e
+            else:
+                raise
+        finally:
+            # Clean up the patch
+            for module_name, original_impl in originals.items():
+                module = importlib.import_module(module_name)
+                module.get_type_to_cls_dict = original_impl
+
+    return wrapper
+
+
+def register_pydantic_serializer():
+    """
+    Helper function to pickle pydantic fields for pydantic v1.
+    Pydantic's Cython validators are not serializable.
+    https://github.com/cloudpipe/cloudpickle/issues/408
+    """
+    import pydantic
+
+    if Version(pydantic.__version__) >= Version("2.0.0"):
+        return
+
+    import pydantic.fields
+
+    def custom_serializer(obj):
+        return {
+            "name": obj.name,
+            # outer_type_ is the original type for ModelFields,
+            # while type_ can be updated later with the nested type
+            # like int for List[int].
+            "type_": obj.outer_type_,
+            "class_validators": obj.class_validators,
+            "model_config": obj.model_config,
+            "default": obj.default,
+            "default_factory": obj.default_factory,
+            "required": obj.required,
+            "final": obj.final,
+            "alias": obj.alias,
+            "field_info": obj.field_info,
+        }
+
+    def custom_deserializer(kwargs):
+        return pydantic.fields.ModelField(**kwargs)
+
+    def _CloudPicklerReducer(obj):
+        return custom_deserializer, (custom_serializer(obj),)
+
+    warnings.warn(
+        "Using custom serializer to pickle pydantic.fields.ModelField classes, "
+        "this might miss some fields and validators. To avoid this, "
+        "please upgrade pydantic to v2 using `pip install pydantic -U` with "
+        "langchain 0.0.267 and above."
+    )
+    cloudpickle.CloudPickler.dispatch[pydantic.fields.ModelField] = _CloudPicklerReducer
+
+
+def unregister_pydantic_serializer():
+    import pydantic
+
+    if Version(pydantic.__version__) >= Version("2.0.0"):
+        return
+
+    cloudpickle.CloudPickler.dispatch.pop(pydantic.fields.ModelField, None)
+
+
+@contextlib.contextmanager
+def register_pydantic_v1_serializer_cm():
+    try:
+        register_pydantic_serializer()
+        yield
+    finally:
+        unregister_pydantic_serializer()

venv/lib/python3.9/site-packages/mlflow/langchain/utils/chat.py

@@ -0,0 +1,347 @@
+import json
+import logging
+import time
+from typing import Any, Union
+
+import pydantic
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    ChatMessage,
+    FunctionMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolMessage,
+)
+from langchain_core.outputs.chat_generation import ChatGeneration
+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.types.chat import (
+    ChatChoice,
+    ChatChoiceDelta,
+    ChatChunkChoice,
+    ChatCompletionChunk,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatMessage,
+    ChatUsage,
+)
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+
+_logger = logging.getLogger(__name__)
+
+
+def convert_lc_message_to_chat_message(lc_message: Union[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):
+            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=lc_message.content or None,
+                tool_calls=tool_calls,
+            )
+        else:
+            return ChatMessage(role="assistant", content=lc_message.content)
+    elif isinstance(lc_message, ChatMessage):
+        return ChatMessage(role=lc_message.role, content=lc_message.content)
+    elif isinstance(lc_message, FunctionMessage):
+        return ChatMessage(role="function", content=lc_message.content)
+    elif isinstance(lc_message, ToolMessage):
+        return ChatMessage(
+            role="tool",
+            content=lc_message.content,
+            tool_call_id=lc_message.tool_call_id,
+        )
+    elif isinstance(lc_message, HumanMessage):
+        return ChatMessage(role="user", content=lc_message.content)
+    elif isinstance(lc_message, SystemMessage):
+        return ChatMessage(role="system", 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]:
+    # AIMessage does not have tool_calls field in LangChain < 0.1.0.
+    if not hasattr(message, "tool_calls"):
+        return []
+
+    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 convert_lc_generation_to_chat_message(lc_gen: Generation) -> ChatMessage:
+    """
+    Convert LangChain's generation format to the MLflow's standard chat message format.
+    """
+    if isinstance(lc_gen, ChatGeneration):
+        try:
+            return convert_lc_message_to_chat_message(lc_gen.message)
+        except Exception as e:
+            # When failed to convert the message, return as assistant message
+            _logger.debug(
+                f"Failed to convert the message from ChatGeneration to ResponseMessage: {e}",
+                exc_info=True,
+            )
+
+    return ChatMessage(role="assistant", content=lc_gen.text)
+
+
+def try_transform_response_to_chat_format(response: Any) -> dict:
+    """
+    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,
+            ),
+        )
+        if IS_PYDANTIC_V2_OR_NEWER:
+            return transformed_response.model_dump(mode="json", exclude_unset=True)
+        else:
+            return json.loads(transformed_response.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,
+            created=int(time.time()),
+            model="",
+            choices=[
+                ChatChunkChoice(
+                    index=0,
+                    delta=ChatChoiceDelta(
+                        role="assistant",
+                        content=message_content,
+                    ),
+                    finish_reason=finish_reason,
+                )
+            ],
+        )
+
+        if IS_PYDANTIC_V2_OR_NEWER:
+            return transformed_response.model_dump(mode="json")
+        else:
+            return json.loads(transformed_response.json())
+
+    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[Union[BaseMessage]]:
+    model = ChatCompletionRequest.validate_compat(chat_request)
+    return [_chat_model_to_langchain_message(message) for message in model.messages]
+
+
+def _convert_chat_request(chat_request: Union[dict, list[dict]]):
+    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.__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):
+    # 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 langchain.agents import AgentExecutor
+
+    if 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]):
+        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

venv/lib/python3.9/site-packages/mlflow/langchain/utils/serialization.py

@@ -0,0 +1,36 @@
+import inspect
+
+from packaging.version import Version
+
+
+def convert_to_serializable(response):
+    """
+    Convert the response to a JSON serializable format.
+
+    LangChain response objects often contains Pydantic objects, which causes an serialization
+    error when the model is served behind REST endpoint.
+    """
+    import langchain
+
+    # LangChain >= 0.3.0 uses Pydantic 2.x while < 0.3.0 is based on Pydantic 1.x.
+    if Version(langchain.__version__) >= Version("0.3.0"):
+        from pydantic import BaseModel
+
+        if isinstance(response, BaseModel):
+            return response.model_dump()
+    else:
+        from langchain_core.pydantic_v1 import BaseModel as LangChainBaseModel
+
+        if isinstance(response, LangChainBaseModel):
+            return response.dict()
+
+    if inspect.isgenerator(response):
+        return (convert_to_serializable(chunk) for chunk in response)
+    elif isinstance(response, dict):
+        return {k: convert_to_serializable(v) for k, v in response.items()}
+    elif isinstance(response, list):
+        return [convert_to_serializable(v) for v in response]
+    elif isinstance(response, tuple):
+        return tuple(convert_to_serializable(v) for v in response)
+
+    return response