mlflow--mlflow/mlflow/langchain/model.py

"""
The ``mlflow.langchain`` module provides an API for logging and loading LangChain models.
This module exports multivariate LangChain models in the langchain flavor and univariate
LangChain models in the pyfunc flavor:

LangChain (native) format
    This is the main flavor that can be accessed with LangChain APIs.
:py:mod:`mlflow.pyfunc`
    Produced for use by generic pyfunc-based deployment tools and for batch inference.

.. _LangChain:
    https://python.langchain.com/en/latest/index.html
"""

import logging
import os
import tempfile
import warnings
from typing import Any, Iterator

import cloudpickle
import pandas as pd
import yaml
from packaging.version import Version

import mlflow
from mlflow import pyfunc
from mlflow.entities.model_registry.prompt import Prompt
from mlflow.environment_variables import MLFLOW_ALLOW_PICKLE_DESERIALIZATION
from mlflow.exceptions import MlflowException
from mlflow.langchain.constants import FLAVOR_NAME
from mlflow.langchain.databricks_dependencies import _detect_databricks_dependencies
from mlflow.langchain.runnables import _load_runnables, _save_runnables
from mlflow.langchain.utils.logging import (
    _BASE_LOAD_KEY,
    _MODEL_LOAD_KEY,
    _RUNNABLE_LOAD_KEY,
    _load_base_lcs,
    _save_base_lcs,
    _validate_and_prepare_lc_model_or_path,
    lc_runnables_types,
    patch_langchain_type_to_cls_dict,
)
from mlflow.models import Model, ModelInputExample, ModelSignature
from mlflow.models.dependencies_schemas import (
    _clear_dependencies_schemas,
    _get_dependencies_schema_from_model,
    _get_dependencies_schemas,
)
from mlflow.models.model import (
    MLMODEL_FILE_NAME,
    MODEL_CODE_PATH,
    MODEL_CONFIG,
    _update_active_model_id_based_on_mlflow_model,
)
from mlflow.models.resources import DatabricksFunction, Resource, _ResourceBuilder
from mlflow.models.signature import _infer_signature_from_input_example
from mlflow.models.utils import (
    _convert_llm_input_data,
    _load_model_code_path,
    _save_example,
)
from mlflow.pyfunc import FLAVOR_NAME as PYFUNC_FLAVOR_NAME
from mlflow.pyfunc.context import Context
from mlflow.tracing.provider import trace_disabled
from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
from mlflow.types.schema import ColSpec, DataType, Schema
from mlflow.utils.databricks_utils import (
    _get_databricks_serverless_env_vars,
    is_in_databricks_model_serving_environment,
    is_in_databricks_runtime,
    is_in_databricks_serverless_runtime,
    is_mlflow_tracing_enabled_in_model_serving,
)
from mlflow.utils.docstring_utils import (
    LOG_MODEL_PARAM_DOCS,
    docstring_version_compatibility_warning,
    format_docstring,
)
from mlflow.utils.environment import (
    _CONDA_ENV_FILE_NAME,
    _CONSTRAINTS_FILE_NAME,
    _PYTHON_ENV_FILE_NAME,
    _REQUIREMENTS_FILE_NAME,
    _mlflow_conda_env,
    _process_conda_env,
    _process_pip_requirements,
    _PythonEnv,
    _validate_env_arguments,
)
from mlflow.utils.file_utils import get_total_file_size, write_to
from mlflow.utils.model_utils import (
    _add_code_from_conf_to_system_path,
    _get_flavor_configuration,
    _validate_and_copy_code_paths,
    _validate_and_copy_file_to_directory,
    _validate_and_get_model_config_from_file,
    _validate_and_prepare_target_save_path,
)
from mlflow.utils.requirements_utils import _get_pinned_requirement

logger = logging.getLogger(mlflow.__name__)

_MODEL_TYPE_KEY = "model_type"


def get_default_pip_requirements():
    """
    Returns:
        A list of default pip requirements for MLflow Models produced by this flavor.
        Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment
        that, at a minimum, contains these requirements.
    """
    # pin pydantic and cloudpickle version as they are used in langchain
    # model saving and loading
    return list(map(_get_pinned_requirement, ["langchain", "pydantic", "cloudpickle"]))


def get_default_conda_env():
    """
    Returns:
        The default Conda environment for MLflow Models produced by calls to
        :func:`save_model()` and :func:`log_model()`.
    """
    return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements())


@format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
@docstring_version_compatibility_warning(FLAVOR_NAME)
@trace_disabled  # Suppress traces for internal predict calls while saving model
def save_model(
    lc_model,
    path,
    conda_env=None,
    code_paths=None,
    mlflow_model=None,
    signature: ModelSignature = None,
    input_example: ModelInputExample = None,
    pip_requirements=None,
    extra_pip_requirements=None,
    metadata=None,
    loader_fn=None,
    persist_dir=None,
    model_config=None,
    streamable: bool | None = None,
):
    """
    Save a LangChain model to a path on the local file system.

    Args:
        lc_model: A LangChain model, which could be a
            `Chain <https://python.langchain.com/docs/modules/chains/>`_,
            `Agent <https://python.langchain.com/docs/modules/agents/>`_,
            `retriever <https://python.langchain.com/docs/modules/data_connection/retrievers/>`_,
            or `RunnableSequence <https://python.langchain.com/docs/modules/chains/foundational/sequential_chains#using-lcel>`_,
            or a path containing the `LangChain model code <https://github.com/mlflow/mlflow/blob/master/examples/langchain/chain_as_code_driver.py>`
            for the above types. When using model as path, make sure to set the model
            by using :func:`mlflow.models.set_model()`.

            .. Note:: Experimental: Using model as path may change or be removed in a future
                        release without warning.
        path: Local path where the serialized model (as YAML) is to be saved.
        conda_env: {{ conda_env }}
        code_paths: {{ code_paths }}
        mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
        signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
            describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
            If not specified, the model signature would be set according to
            `lc_model.input_keys` and `lc_model.output_keys` as columns names, and
            `DataType.string` as the column type.
            Alternatively, you can explicitly specify the model signature.
            The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
            from datasets with valid model input (e.g. the training dataset with target
            column omitted) and valid model output (e.g. model predictions generated on
            the training dataset), for example:

            .. code-block:: python

                from mlflow.models import infer_signature

                chain = LLMChain(llm=llm, prompt=prompt)
                prediction = chain.run(input_str)
                input_columns = [
                    {"type": "string", "name": input_key} for input_key in chain.input_keys
                ]
                signature = infer_signature(input_columns, predictions)

        input_example: {{ input_example }}
        pip_requirements: {{ pip_requirements }}
        extra_pip_requirements: {{ extra_pip_requirements }}
        metadata: {{ metadata }}
        loader_fn: A function that's required for models containing objects that aren't natively
            serialized by LangChain.
            This function takes a string `persist_dir` as an argument and returns the
            specific object that the model needs. Depending on the model,
            this could be a retriever, vectorstore, requests_wrapper, embeddings, or
            database. For RetrievalQA Chain and retriever models, the object is a
            (`retriever <https://python.langchain.com/docs/modules/data_connection/retrievers/>`_).
            For APIChain models, it's a
            (`requests_wrapper <https://python.langchain.com/docs/modules/agents/tools/integrations/requests>`_).
            For HypotheticalDocumentEmbedder models, it's an
            (`embeddings <https://python.langchain.com/docs/modules/data_connection/text_embedding/>`_).
            For SQLDatabaseChain models, it's a
            (`database <https://python.langchain.com/docs/modules/agents/toolkits/sql_database>`_).
        persist_dir: The directory where the object is stored. The `loader_fn`
            takes this string as the argument to load the object.
            This is optional for models containing objects that aren't natively
            serialized by LangChain. MLflow logs the content in this directory as
            artifacts in the subdirectory named `persist_dir_data`.

            Here is the code snippet for logging a RetrievalQA chain with `loader_fn`
            and `persist_dir`:

            .. Note:: In langchain_community >= 0.0.27, loading pickled data requires providing the
                ``allow_dangerous_deserialization`` argument.

            .. code-block:: python

                qa = RetrievalQA.from_llm(llm=OpenAI(), retriever=db.as_retriever())


                def load_retriever(persist_directory):
                    embeddings = OpenAIEmbeddings()
                    vectorstore = FAISS.load_local(
                        persist_directory,
                        embeddings,
                        # you may need to add the line below
                        # for langchain_community >= 0.0.27
                        allow_dangerous_deserialization=True,
                    )
                    return vectorstore.as_retriever()


                with mlflow.start_run() as run:
                    logged_model = mlflow.langchain.log_model(
                        qa,
                        name="retrieval_qa",
                        loader_fn=load_retriever,
                        persist_dir=persist_dir,
                    )

            See a complete example in examples/langchain/retrieval_qa_chain.py.
        model_config: The model configuration to apply to the model if saving model from code. This
            configuration is available during model loading.

            .. Note:: Experimental: This parameter may change or be removed in a future
                                    release without warning.
        streamable: A boolean value indicating if the model supports streaming prediction. If
            True, the model must implement `stream` method. If None, streamable is
            set to True if the model implements `stream` method. Default to `None`.
    """
    import langchain

    with tempfile.TemporaryDirectory() as temp_dir:
        from mlflow.langchain._compat import import_base_retriever

        BaseRetriever = import_base_retriever()

        lc_model_or_path = _validate_and_prepare_lc_model_or_path(lc_model, loader_fn, temp_dir)

        _validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)

        path = os.path.abspath(path)
        _validate_and_prepare_target_save_path(path)

        if isinstance(model_config, str):
            model_config = _validate_and_get_model_config_from_file(model_config)

        model_code_path = None
        if isinstance(lc_model_or_path, str):
            # The LangChain model is defined as Python code located in the file at the path
            # specified by `lc_model`. Verify that the path exists and, if so, copy it to the
            # model directory along with any other specified code modules
            model_code_path = lc_model_or_path

            lc_model = _load_model_code_path(model_code_path, model_config)
            _validate_and_copy_file_to_directory(model_code_path, path, "code")
        else:
            lc_model = lc_model_or_path

    code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)

    if mlflow_model is None:
        mlflow_model = Model()
    saved_example = _save_example(mlflow_model, input_example, path)

    if signature is None:
        if saved_example is not None:
            wrapped_model = _LangChainModelWrapper(lc_model)
            signature = _infer_signature_from_input_example(saved_example, wrapped_model)
        else:
            if hasattr(lc_model, "input_keys"):
                input_columns = [
                    ColSpec(type=DataType.string, name=input_key)
                    for input_key in lc_model.input_keys
                ]
                input_schema = Schema(input_columns)
            else:
                input_schema = None
            if (
                hasattr(lc_model, "output_keys")
                and len(lc_model.output_keys) == 1
                and not isinstance(lc_model, BaseRetriever)
            ):
                output_columns = [
                    ColSpec(type=DataType.string, name=output_key)
                    for output_key in lc_model.output_keys
                ]
                output_schema = Schema(output_columns)
            else:
                # TODO: empty output schema if multiple output_keys or is a retriever. fix later!
                # https://databricks.atlassian.net/browse/ML-34706
                output_schema = None

            signature = (
                ModelSignature(input_schema, output_schema)
                if input_schema or output_schema
                else None
            )

    if signature is not None:
        mlflow_model.signature = signature
    if metadata is not None:
        mlflow_model.metadata = metadata

    with _get_dependencies_schemas() as dependencies_schemas:
        schema = dependencies_schemas.to_dict()
        if schema is not None:
            if mlflow_model.metadata is None:
                mlflow_model.metadata = {}
            mlflow_model.metadata.update(schema)

    if streamable is None:
        streamable = hasattr(lc_model, "stream")

    model_data_kwargs = {}
    flavor_conf = {}
    if not isinstance(model_code_path, str):
        if Version(langchain.__version__).major >= 1:
            raise MlflowException.invalid_parameter_value(
                "LangChain v1 onward only supports models-from-code, i.e., the 'lc_model' "
                "argument value must be a path containing the `LangChain` model code. "
                "You can refer to documentation at "
                "https://mlflow.org/docs/latest/ml/model/models-from-code/#examples-and-patterns "
                "for example code."
            )
        else:
            logger.warning(
                "Saving langchain model in the cloudpickle format requires exercising "
                "caution because these formats rely on Python's object serialization mechanism, "
                "which can execute arbitrary code during deserialization."
                "The recommended alternative is to save it as 'models-from-code' artifacts."
                "You can refer to documentation at "
                "https://mlflow.org/docs/latest/ml/model/models-from-code/#examples-and-patterns "
                "for example code.",
            )

        model_data_kwargs = _save_model(lc_model, path, loader_fn, persist_dir)
        flavor_conf = {
            _MODEL_TYPE_KEY: lc_model.__class__.__name__,
            **model_data_kwargs,
        }

    pyfunc.add_to_model(
        mlflow_model,
        loader_module="mlflow.langchain",
        conda_env=_CONDA_ENV_FILE_NAME,
        python_env=_PYTHON_ENV_FILE_NAME,
        code=code_dir_subpath,
        predict_stream_fn="predict_stream",
        streamable=streamable,
        model_code_path=model_code_path,
        model_config=model_config,
        **model_data_kwargs,
    )

    needs_databricks_auth = False
    if mlflow_model.resources is None:
        if databricks_resources := _detect_databricks_dependencies(lc_model):
            logger.info(
                "Attempting to auto-detect Databricks resource dependencies for the "
                "current langchain model. Dependency auto-detection is "
                "best-effort and may not capture all dependencies of your langchain "
                "model, resulting in authorization errors when serving or querying "
                "your model. We recommend that you explicitly pass `resources` "
                "to mlflow.langchain.log_model() to ensure authorization to "
                "dependent resources succeeds when the model is deployed."
            )
            serialized_databricks_resources = _ResourceBuilder.from_resources(databricks_resources)
            mlflow_model.resources = serialized_databricks_resources
            needs_databricks_auth = any(
                isinstance(r, DatabricksFunction) for r in databricks_resources
            )

    mlflow_model.add_flavor(
        FLAVOR_NAME,
        langchain_version=langchain.__version__,
        code=code_dir_subpath,
        streamable=streamable,
        **flavor_conf,
    )
    if size := get_total_file_size(path):
        mlflow_model.model_size_bytes = size
    mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))

    if conda_env is None:
        if pip_requirements is None:
            default_reqs = get_default_pip_requirements()
            extra_env_vars = (
                _get_databricks_serverless_env_vars()
                if needs_databricks_auth and is_in_databricks_serverless_runtime()
                else None
            )
            inferred_reqs = mlflow.models.infer_pip_requirements(
                str(path), FLAVOR_NAME, fallback=default_reqs, extra_env_vars=extra_env_vars
            )
            default_reqs = sorted(set(inferred_reqs).union(default_reqs))
        else:
            default_reqs = None
        conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
            default_reqs, pip_requirements, extra_pip_requirements
        )
    else:
        conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)

    with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
        yaml.safe_dump(conda_env, stream=f, default_flow_style=False)

    if pip_constraints:
        write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))

    write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))

    _PythonEnv.current().to_yaml(os.path.join(path, _PYTHON_ENV_FILE_NAME))


@format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
@docstring_version_compatibility_warning(FLAVOR_NAME)
@trace_disabled  # Suppress traces for internal predict calls while logging model
def log_model(
    lc_model,
    artifact_path: str | None = None,
    conda_env=None,
    code_paths=None,
    registered_model_name=None,
    signature: ModelSignature = None,
    input_example: ModelInputExample = None,
    await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
    pip_requirements=None,
    extra_pip_requirements=None,
    metadata=None,
    loader_fn=None,
    persist_dir=None,
    run_id=None,
    model_config=None,
    streamable=None,
    resources: list[Resource] | str | None = None,
    prompts: list[str | Prompt] | None = None,
    name: str | None = None,
    params: dict[str, Any] | None = None,
    tags: dict[str, Any] | None = None,
    model_type: str | None = None,
    step: int = 0,
    model_id: str | None = None,
):
    """
    Log a LangChain model as an MLflow artifact for the current run.

    Args:
        lc_model: A LangChain model, which could be a
            `Chain <https://python.langchain.com/docs/modules/chains/>`_,
            `Agent <https://python.langchain.com/docs/modules/agents/>`_, or
            `retriever <https://python.langchain.com/docs/modules/data_connection/retrievers/>`_
            or a path containing the `LangChain model code <https://github.com/mlflow/mlflow/blob/master/examples/langchain/chain_as_code_driver.py>`
            for the above types. When using model as path, make sure to set the model
            by using :func:`mlflow.models.set_model()`.

            .. Note:: Experimental: Using model as path may change or be removed in a future
                                    release without warning.
        artifact_path: Deprecated. Use `name` instead.
        conda_env: {{ conda_env }}
        code_paths: {{ code_paths }}
        registered_model_name: If given, create a model
            version under ``registered_model_name``, also creating a
            registered model if one with the given name does not exist.
        signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
            describes model input and output
            :py:class:`Schema <mlflow.types.Schema>`.
            If not specified, the model signature would be set according to
            `lc_model.input_keys` and `lc_model.output_keys` as columns names, and
            `DataType.string` as the column type.
            Alternatively, you can explicitly specify the model signature.
            The model signature can be :py:func:`inferred
            <mlflow.models.infer_signature>` from datasets with valid model input
            (e.g. the training dataset with target column omitted) and valid model
            output (e.g. model predictions generated on the training dataset),
            for example:

            .. code-block:: python

                from mlflow.models import infer_signature

                chain = LLMChain(llm=llm, prompt=prompt)
                prediction = chain.run(input_str)
                input_columns = [
                    {"type": "string", "name": input_key} for input_key in chain.input_keys
                ]
                signature = infer_signature(input_columns, predictions)

        input_example: {{ input_example }}
        await_registration_for: Number of seconds to wait for the model version
            to finish being created and is in ``READY`` status.
            By default, the function waits for five minutes.
            Specify 0 or None to skip waiting.
        pip_requirements: {{ pip_requirements }}
        extra_pip_requirements: {{ extra_pip_requirements }}
        metadata: {{ metadata }}
        loader_fn: A function that's required for models containing objects that aren't natively
            serialized by LangChain.
            This function takes a string `persist_dir` as an argument and returns the
            specific object that the model needs. Depending on the model,
            this could be a retriever, vectorstore, requests_wrapper, embeddings, or
            database. For RetrievalQA Chain and retriever models, the object is a
            (`retriever <https://python.langchain.com/docs/modules/data_connection/retrievers/>`_).
            For APIChain models, it's a
            (`requests_wrapper <https://python.langchain.com/docs/modules/agents/tools/integrations/requests>`_).
            For HypotheticalDocumentEmbedder models, it's an
            (`embeddings <https://python.langchain.com/docs/modules/data_connection/text_embedding/>`_).
            For SQLDatabaseChain models, it's a
            (`database <https://python.langchain.com/docs/modules/agents/toolkits/sql_database>`_).
        persist_dir: The directory where the object is stored. The `loader_fn`
            takes this string as the argument to load the object.
            This is optional for models containing objects that aren't natively
            serialized by LangChain. MLflow logs the content in this directory as
            artifacts in the subdirectory named `persist_dir_data`.

            Here is the code snippet for logging a RetrievalQA chain with `loader_fn`
            and `persist_dir`:

            .. Note:: In langchain_community >= 0.0.27, loading pickled data requires providing the
                ``allow_dangerous_deserialization`` argument.

            .. code-block:: python

                qa = RetrievalQA.from_llm(llm=OpenAI(), retriever=db.as_retriever())


                def load_retriever(persist_directory):
                    embeddings = OpenAIEmbeddings()
                    vectorstore = FAISS.load_local(
                        persist_directory,
                        embeddings,
                        # you may need to add the line below
                        # for langchain_community >= 0.0.27
                        allow_dangerous_deserialization=True,
                    )
                    return vectorstore.as_retriever()


                with mlflow.start_run() as run:
                    logged_model = mlflow.langchain.log_model(
                        qa,
                        name="retrieval_qa",
                        loader_fn=load_retriever,
                        persist_dir=persist_dir,
                    )

            See a complete example in examples/langchain/retrieval_qa_chain.py.
        run_id: run_id to associate with this model version. If specified, we resume the
                run and log the model to that run. Otherwise, a new run is created.
                Default to None.
        model_config: The model configuration to apply to the model if saving model from code. This
            configuration is available during model loading.

            .. Note:: Experimental: This parameter may change or be removed in a future
                                    release without warning.
        streamable: A boolean value indicating if the model supports streaming prediction. If
            True, the model must implement `stream` method. If None, If None, streamable is
            set to True if the model implements `stream` method. Default to `None`.
        resources: A list of model resources or a resources.yaml file containing a list of
            resources required to serve the model. If logging a LangChain model with dependencies
            (e.g. on LLM model serving endpoints), we encourage explicitly passing dependencies
            via this parameter. Otherwise, ``log_model`` will attempt to infer dependencies,
            but dependency auto-inference is best-effort and may miss some dependencies.
        prompts: {{ prompts }}

        name: {{ name }}
        params: {{ params }}
        tags: {{ tags }}
        model_type: {{ model_type }}
        step: {{ step }}
        model_id: {{ model_id }}

    Returns:
        A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
        metadata of the logged model.
    """
    return Model.log(
        artifact_path=artifact_path,
        name=name,
        flavor=mlflow.langchain,
        registered_model_name=registered_model_name,
        lc_model=lc_model,
        conda_env=conda_env,
        code_paths=code_paths,
        signature=signature,
        input_example=input_example,
        await_registration_for=await_registration_for,
        pip_requirements=pip_requirements,
        extra_pip_requirements=extra_pip_requirements,
        metadata=metadata,
        loader_fn=loader_fn,
        persist_dir=persist_dir,
        run_id=run_id,
        model_config=model_config,
        streamable=streamable,
        resources=resources,
        prompts=prompts,
        params=params,
        tags=tags,
        model_type=model_type,
        step=step,
        model_id=model_id,
    )


# patch_langchain_type_to_cls_dict here as we attempt to load model
# if it's saved by `dict` method
@patch_langchain_type_to_cls_dict
def _save_model(model, path, loader_fn, persist_dir):
    if Version(cloudpickle.__version__) < Version("2.1.0"):
        warnings.warn(
            "If you are constructing a custom LangChain model, "
            "please upgrade cloudpickle to version 2.1.0 or later "
            "using `pip install cloudpickle>=2.1.0` "
            "to ensure the model can be loaded correctly."
        )

    if isinstance(model, lc_runnables_types()):
        return _save_runnables(model, path, loader_fn=loader_fn, persist_dir=persist_dir)
    else:
        return _save_base_lcs(model, path, loader_fn, persist_dir)


@patch_langchain_type_to_cls_dict
def _load_model(local_model_path, flavor_conf):
    # model_type is not accurate as the class can be subclass
    # of supported types, we define _MODEL_LOAD_KEY to ensure
    # which load function to use
    model_load_fn = flavor_conf.get(_MODEL_LOAD_KEY)
    if model_load_fn == _RUNNABLE_LOAD_KEY:
        model = _load_runnables(local_model_path, flavor_conf)
    elif model_load_fn == _BASE_LOAD_KEY:
        model = _load_base_lcs(local_model_path, flavor_conf)
    else:
        raise mlflow.MlflowException(
            "Failed to load LangChain model. Unknown model type: "
            f"{flavor_conf.get(_MODEL_TYPE_KEY)}"
        )
    return model


class _LangChainModelWrapper:
    def __init__(self, lc_model, model_path=None):
        self.lc_model = lc_model
        self.model_path = model_path

    def get_raw_model(self):
        """
        Returns the underlying model.
        """
        return self.lc_model

    def predict(
        self,
        data: pd.DataFrame | list[str | dict[str, Any]] | Any,
        params: dict[str, Any] | None = None,
    ) -> list[str | dict[str, Any]]:
        """
        Args:
            data: Model input data.
            params: Additional parameters to pass to the model for inference.

        Returns:
            Model predictions.
        """
        # TODO: We don't automatically turn tracing on in OSS model serving, because we haven't
        # implemented storage option for traces in OSS model serving (counterpart to the
        # Inference Table in Databricks model serving).
        if (
            is_in_databricks_model_serving_environment()
            # TODO: This env var was once used for controlling whether or not to inject the
            #   tracer in Databricks model serving. However, now we have the new env var
            #   `ENABLE_MLFLOW_TRACING` to control that. We don't remove this condition
            #   right now in the interest of caution, but we should remove this condition
            #   after making sure that the functionality is stable.
            and os.environ.get("MLFLOW_ENABLE_TRACE_IN_SERVING", "false").lower() == "true"
            # if this is False, tracing is disabled and we shouldn't inject the tracer
            and is_mlflow_tracing_enabled_in_model_serving()
        ):
            from mlflow.langchain.langchain_tracer import MlflowLangchainTracer

            callbacks = [MlflowLangchainTracer()]
        else:
            callbacks = None

        return self._predict_with_callbacks(data, params, callback_handlers=callbacks)

    def _update_dependencies_schemas_in_prediction_context(
        self, callback_handlers
    ) -> Context | None:
        from mlflow.langchain.langchain_tracer import MlflowLangchainTracer

        if (
            callback_handlers
            and (
                tracer := next(
                    (c for c in callback_handlers if isinstance(c, MlflowLangchainTracer)), None
                )
            )
            and self.model_path
        ):
            model = Model.load(self.model_path)
            context = tracer._prediction_context
            if context and (schema := _get_dependencies_schema_from_model(model)):
                context.update(**schema)
            return context

    def _predict_with_callbacks(
        self,
        data: pd.DataFrame | list[str | dict[str, Any]] | Any,
        params: dict[str, Any] | None = None,
        callback_handlers=None,
        convert_chat_responses=False,
    ) -> list[str | dict[str, Any]]:
        """
        Args:
            data: Model input data.
            params: Additional parameters to pass to the model for inference.
            callback_handlers: Callback handlers to pass to LangChain.
            convert_chat_responses: If true, forcibly convert response to chat model
                response format.

        Returns:
            Model predictions.
        """
        from mlflow.langchain.api_request_parallel_processor import process_api_requests

        context = self._update_dependencies_schemas_in_prediction_context(callback_handlers)
        messages, return_first_element = self._prepare_predict_messages(data)
        results = process_api_requests(
            lc_model=self.lc_model,
            requests=messages,
            callback_handlers=callback_handlers,
            convert_chat_responses=convert_chat_responses,
            params=params or {},
            context=context,
        )
        return results[0] if return_first_element else results

    def _prepare_predict_messages(self, data):
        """
        Return a tuple of (preprocessed_data, return_first_element)
        `preprocessed_data` is always a list,
        and `return_first_element` means if True, we should return the first element
        of inference result, otherwise we should return the whole inference result.
        """
        data = _convert_llm_input_data(data)

        if not isinstance(data, list):
            # if the input data is not a list (i.e. single input),
            # we still need to convert it to a one-element list `[data]`
            # because `process_api_requests` only accepts list as valid input.
            # and in this case,
            # we should return the first element of the inference result
            # because we change input `data` to `[data]`
            return [data], True
        if isinstance(data, list):
            return data, False
        raise mlflow.MlflowException.invalid_parameter_value(
            "Input must be a pandas DataFrame or a list "
            f"for model {self.lc_model.__class__.__name__}"
        )

    def _prepare_predict_stream_messages(self, data):
        data = _convert_llm_input_data(data)

        if isinstance(data, list):
            # `predict_stream` only accepts single input.
            # but `enforce_schema` might convert single input into a list like `[single_input]`
            # so extract the first element in the list.
            if len(data) != 1:
                raise MlflowException(
                    f"'predict_stream' requires single input, but it got input data {data}"
                )
            return data[0]
        return data

    def predict_stream(
        self,
        data: Any,
        params: dict[str, Any] | None = None,
    ) -> Iterator[str | dict[str, Any]]:
        """
        Args:
            data: Model input data, only single input is allowed.
            params: Additional parameters to pass to the model for inference.

        Returns:
            An iterator of model prediction chunks.
        """
        from mlflow.langchain.api_request_parallel_processor import (
            process_stream_request,
        )

        data = self._prepare_predict_stream_messages(data)
        return process_stream_request(
            lc_model=self.lc_model,
            request_json=data,
            params=params or {},
        )

    def _predict_stream_with_callbacks(
        self,
        data: Any,
        params: dict[str, Any] | None = None,
        callback_handlers=None,
        convert_chat_responses=False,
    ) -> Iterator[str | dict[str, Any]]:
        """
        Args:
            data: Model input data, only single input is allowed.
            params: Additional parameters to pass to the model for inference.
            callback_handlers: Callback handlers to pass to LangChain.
            convert_chat_responses: If true, forcibly convert response to chat model
                response format.

        Returns:
            An iterator of model prediction chunks.
        """
        from mlflow.langchain.api_request_parallel_processor import (
            process_stream_request,
        )

        self._update_dependencies_schemas_in_prediction_context(callback_handlers)
        data = self._prepare_predict_stream_messages(data)
        return process_stream_request(
            lc_model=self.lc_model,
            request_json=data,
            callback_handlers=callback_handlers,
            convert_chat_responses=convert_chat_responses,
            params=params or {},
        )


def _load_pyfunc(path: str, model_config: dict[str, Any] | None = None):
    """Load PyFunc implementation for LangChain. Called by ``pyfunc.load_model``.

    Args:
        path: Local filesystem path to the MLflow Model with the ``langchain`` flavor.
    """
    return _LangChainModelWrapper(_load_model_from_local_fs(path, model_config), path)


def _load_model_from_local_fs(local_model_path, model_config_overrides=None):
    mlflow_model = Model.load(local_model_path)
    flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME)
    pyfunc_flavor_conf = _get_flavor_configuration(
        model_path=local_model_path, flavor_name=PYFUNC_FLAVOR_NAME
    )
    # Add code from the langchain flavor to the system path
    _add_code_from_conf_to_system_path(local_model_path, flavor_conf)
    # The model_code_path and the model_config were previously saved langchain flavor but now we
    # also save them inside the pyfunc flavor. For backwards compatibility of previous models,
    # we need to check both places.
    if MODEL_CODE_PATH in pyfunc_flavor_conf or MODEL_CODE_PATH in flavor_conf:
        model_config = pyfunc_flavor_conf.get(MODEL_CONFIG, flavor_conf.get(MODEL_CONFIG, None))
        if isinstance(model_config, str):
            config_path = os.path.join(
                local_model_path,
                os.path.basename(model_config),
            )
            model_config = _validate_and_get_model_config_from_file(config_path)

        flavor_code_path = pyfunc_flavor_conf.get(
            MODEL_CODE_PATH, flavor_conf.get(MODEL_CODE_PATH, None)
        )
        model_code_path = os.path.join(
            local_model_path,
            os.path.basename(flavor_code_path),
        )
        try:
            model = _load_model_code_path(
                model_code_path, {**(model_config or {}), **(model_config_overrides or {})}
            )
        finally:
            # We would like to clean up the dependencies schema which is set to global
            # after loading the mode to avoid the schema being used in the next model loading
            _clear_dependencies_schemas()
    else:
        if (
            not MLFLOW_ALLOW_PICKLE_DESERIALIZATION.get()
            and not is_in_databricks_runtime()
            and not is_in_databricks_model_serving_environment()
        ):
            raise MlflowException(
                "Deserializing model using pickle is disallowed, but this model is saved "
                "in cloudpickle format. To address this issue, you need to set environment "
                "variable 'MLFLOW_ALLOW_PICKLE_DESERIALIZATION' to 'true'; or save the model as "
                "models-from-code artifacts, you can refer to documentation at "
                "https://mlflow.org/docs/latest/ml/model/models-from-code/#examples-and-patterns "
                "for example code."
            )
        model = _load_model(local_model_path, flavor_conf)
    # set active model after model loading since experiment ID might be set
    # in the model loading process
    _update_active_model_id_based_on_mlflow_model(mlflow_model)
    return model


@docstring_version_compatibility_warning(FLAVOR_NAME)
@trace_disabled  # Suppress traces while loading model
def load_model(model_uri, dst_path=None):
    """
    Load a LangChain model from a local file or a run.

    Args:
        model_uri: The location, in URI format, of the MLflow model. For example:

            - ``/Users/me/path/to/local/model``
            - ``relative/path/to/local/model``
            - ``s3://my_bucket/path/to/model``
            - ``runs:/<mlflow_run_id>/run-relative/path/to/model``

            For more information about supported URI schemes, see
            `Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
            artifact-locations>`_.
        dst_path: The local filesystem path to which to download the model artifact.
            This directory must already exist. If unspecified, a local output
            path will be created.

    Returns:
        A LangChain model instance.
    """
    model_uri = str(model_uri)
    local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
    return _load_model_from_local_fs(local_model_path)