star

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

@@ -0,0 +1,282 @@
+import contextlib
+import inspect
+import logging
+import uuid
+import warnings
+from copy import deepcopy
+
+from packaging.version import Version
+
+import mlflow
+from mlflow.entities import RunTag
+from mlflow.entities.run_status import RunStatus
+from mlflow.exceptions import MlflowException
+from mlflow.langchain.runnables import get_runnable_steps
+from mlflow.tracking.context import registry as context_registry
+from mlflow.utils import name_utils
+from mlflow.utils.autologging_utils import get_autologging_config
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+from mlflow.utils.autologging_utils.safety import _resolve_extra_tags
+
+_logger = logging.getLogger(__name__)
+
+UNSUPPORTED_LOG_MODEL_MESSAGE = (
+    "MLflow autologging does not support logging models containing BaseRetriever because "
+    "logging the model requires `loader_fn` and `persist_dir`. Please log the model manually "
+    "using `mlflow.langchain.log_model(model, artifact_path, loader_fn=..., persist_dir=...)`"
+)
+INFERENCE_FILE_NAME = "inference_inputs_outputs.json"
+
+
+# A *global* state that indicates whether MLflow should patch the inference method
+# for artifact auto-logging (model, signature input example). This disablement
+# is global across threads, as single model inference can trigger multiple threads,
+# for example, LangChain's batch()/abatch() API processes each request in a child thread.
+IS_PATCHING_DISABLED_FOR_ARTIFACTS = False
+
+
+@contextlib.contextmanager
+def disable_patching():
+    """
+    Temporarily disable auto-logging for optional artifacts (model, signature, input
+    examples) to avoid "double-logging" when invoking the patched chain. Without this
+    disablement applied, the patched inference method calls child components that may
+    also be patched, leading to redundant logging.
+    """
+    global IS_PATCHING_DISABLED_FOR_ARTIFACTS
+    original_artifact_flag = IS_PATCHING_DISABLED_FOR_ARTIFACTS
+    IS_PATCHING_DISABLED_FOR_ARTIFACTS = True
+
+    try:
+        yield
+    finally:
+        IS_PATCHING_DISABLED_FOR_ARTIFACTS = original_artifact_flag
+
+
+def patched_inference(func_name, original, self, *args, **kwargs):
+    """
+    A patched implementation of langchain models inference process which enables
+    logging the traces, and other optional artifacts like model, input examples, etc.
+
+    We patch inference functions for different models based on their usage.
+    """
+
+    def _invoke(self, *args, **kwargs):
+        with disable_patching():
+            return original(self, *args, **kwargs)
+
+    config = AutoLoggingConfig.init(mlflow.langchain.FLAVOR_NAME)
+    if not IS_PATCHING_DISABLED_FOR_ARTIFACTS and config.should_log_optional_artifacts():
+        with _setup_autolog_run(config, self) as run_id:
+            result = _invoke(self, *args, **kwargs)
+            _log_optional_artifacts(config, run_id, result, self, func_name, *args, **kwargs)
+    else:
+        result = _invoke(self, *args, **kwargs)
+    return result
+
+
+@contextlib.contextmanager
+def _setup_autolog_run(config, model):
+    """Set up autologging run and return the run ID.
+
+    This function only creates a run when there is no active run and the model does not have
+    a run ID attribute propagated from the previous call. Iff it creates a new run, MLflow should
+    terminate the run at the end of the inference.
+
+    Args:
+        config: AutoLoggingConfig: The autologging configuration.
+        model: Any: The LangChain model instance that runs the inference.
+
+    Returns: yields the run IDs
+    """
+    if propagated_run_id := getattr(model, "run_id", None):
+        # When model has "run_id" attribute, it means the model is already invoked once with autolog
+        # enabled and the run_id is propagated from the previous call, so we don't create a new run.
+        run_id = propagated_run_id
+        # The run should be already terminated at the end of the previous call.
+        should_terminate_run = False
+
+    elif active_run := mlflow.active_run():
+        run_id = active_run.info.run_id
+        tags = _resolve_tags(config.extra_tags, active_run)
+        mlflow.MlflowClient().log_batch(run_id, tags=[RunTag(k, str(v)) for k, v in tags.items()])
+        should_terminate_run = False
+    else:
+        from mlflow.tracking.fluent import _get_experiment_id
+
+        run = mlflow.MlflowClient().create_run(
+            experiment_id=_get_experiment_id(),
+            run_name="langchain-" + name_utils._generate_random_name(),
+            tags=_resolve_tags(config.extra_tags),
+        )
+        run_id = run.info.run_id
+        should_terminate_run = True
+
+    run_status = None
+    try:
+        yield run_id
+    except Exception:
+        run_status = RunStatus.to_string(RunStatus.FAILED)
+        raise
+    finally:
+        if should_terminate_run:
+            mlflow.MlflowClient().set_terminated(run_id, status=run_status)
+
+
+def _resolve_tags(extra_tags, active_run=None):
+    resolved_tags = context_registry.resolve_tags(extra_tags)
+    tags = _resolve_extra_tags(mlflow.langchain.FLAVOR_NAME, resolved_tags)
+    if active_run:
+        # Some context tags like mlflow.runName are immutable once logged, but they might be already
+        # set when the run is created, then we should avoid updating them.
+        excluded_tags = {tag for tag in active_run.data.tags.keys() if tag.startswith("mlflow.")}
+        tags = {k: v for k, v in tags.items() if k not in excluded_tags}
+    return tags
+
+
+def _get_input_data_from_function(func_name, model, args, kwargs):
+    func_param_name_mapping = {
+        "invoke": "input",
+        "batch": "inputs",
+        "stream": "input",
+    }
+    input_example_exc = None
+    if param_name := func_param_name_mapping.get(func_name):
+        inference_func = getattr(model, func_name)
+        # A guard to make sure `param_name` is the first argument of inference function
+        if next(iter(inspect.signature(inference_func).parameters.keys())) != param_name:
+            input_example_exc = MlflowException(
+                "Inference function signature changes, please contact MLflow team to "
+                "fix langchain autologging.",
+            )
+        else:
+            return args[0] if len(args) > 0 else kwargs.get(param_name)
+    else:
+        input_example_exc = MlflowException(
+            f"Unsupported inference function. Only support {list(func_param_name_mapping.keys())}."
+        )
+    _logger.warning(
+        f"Failed to gather input example of model {model.__class__.__name__} "
+        f"due to {input_example_exc}."
+    )
+
+
+def _convert_data_to_dict(data, key):
+    if isinstance(data, dict):
+        return {f"{key}-{k}": v for k, v in data.items()}
+    if isinstance(data, list):
+        return {key: data}
+    if isinstance(data, str):
+        return {key: [data]}
+    raise MlflowException("Unsupported data type.")
+
+
+def _update_langchain_model_config(model):
+    # Langchain models are Pydantic models, and the value for extra is
+    # ignored, we need to set it to allow so as to set attributes on
+    # the model to keep track of logging status
+    import langchain
+
+    try:
+        # LangChain 0.3.0 and above is fully migrated to Pydantic v2
+        if Version(langchain.__version__) >= Version("0.3.0"):
+            if hasattr(model, "model_config") and model.model_config is not None:
+                model.model_config["extra"] = "allow"
+                model.__pydantic_extra__ = {}
+                return True
+        else:
+            from langchain_core.pydantic_v1 import Extra
+
+            if hasattr(model, "__config__"):
+                model.__config__.extra = Extra.allow
+            return True
+    except Exception as e:
+        warnings.warn(
+            "Failed to set extra attribute on the model for keeping track of logging status. "
+            f"MLflow langchain autologging might log model several times. Error: {e}"
+        )
+        return False
+
+
+def _runnable_with_retriever(model):
+    from langchain.schema import BaseRetriever
+
+    with contextlib.suppress(ImportError):
+        from langchain.schema.runnable import RunnableBranch, RunnableParallel, RunnableSequence
+        from langchain.schema.runnable.passthrough import RunnableAssign
+
+        if isinstance(model, RunnableBranch):
+            return any(_runnable_with_retriever(runnable) for _, runnable in model.branches)
+
+        if isinstance(model, RunnableParallel):
+            return any(
+                _runnable_with_retriever(runnable)
+                for runnable in get_runnable_steps(model).values()
+            )
+
+        if isinstance(model, RunnableSequence):
+            return any(_runnable_with_retriever(runnable) for runnable in get_runnable_steps(model))
+
+        if isinstance(model, RunnableAssign):
+            return _runnable_with_retriever(model.mapper)
+
+    return isinstance(model, BaseRetriever)
+
+
+def _chain_with_retriever(model):
+    with contextlib.suppress(ImportError):
+        from langchain.chains import RetrievalQA
+
+        return isinstance(model, RetrievalQA)
+    return False
+
+
+def _log_optional_artifacts(autolog_config, run_id, result, self, func_name, *args, **kwargs):
+    input_example = None
+    if autolog_config.log_models and not hasattr(self, "_mlflow_model_logged"):
+        if _runnable_with_retriever(self) or _chain_with_retriever(self):
+            _logger.info(UNSUPPORTED_LOG_MODEL_MESSAGE)
+        else:
+            # warn user in case we did't capture some cases where retriever is used
+            warnings.warn(UNSUPPORTED_LOG_MODEL_MESSAGE)
+            if autolog_config.log_input_examples:
+                input_example = deepcopy(
+                    _get_input_data_from_function(func_name, self, args, kwargs)
+                )
+                if not autolog_config.log_model_signatures:
+                    _logger.info(
+                        "Signature is automatically generated for logged model if "
+                        "input_example is provided. To disable log_model_signatures, "
+                        "please also disable log_input_examples."
+                    )
+
+            registered_model_name = get_autologging_config(
+                mlflow.langchain.FLAVOR_NAME, "registered_model_name", None
+            )
+            try:
+                with disable_patching():
+                    mlflow.langchain.log_model(
+                        self,
+                        "model",
+                        input_example=input_example,
+                        registered_model_name=registered_model_name,
+                        run_id=run_id,
+                    )
+            except Exception as e:
+                _logger.warning(f"Failed to log model due to error {e}.")
+        # only try logging model once, even if it can't be logged
+        # we don't want to spam the user with warnings/infos
+        if _update_langchain_model_config(self):
+            self._mlflow_model_logged = True
+
+    # Even if the model is not logged, we keep a single run per model
+    if _update_langchain_model_config(self):
+        # NB: We have to set these attributes AFTER the model is logged, otherwise those extra
+        # attributes will be logged as a part of the pickled model and pollute the loaded model.
+        if not hasattr(self, "run_id"):
+            self.run_id = run_id
+        if not hasattr(self, "session_id"):
+            self.session_id = uuid.uuid4().hex
+        self.inference_id = getattr(self, "inference_id", 0) + 1
+
+    return result

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

@@ -0,0 +1,329 @@
+# Based ons: https://github.com/openai/openai-cookbook/blob/6df6ceff470eeba26a56de131254e775292eac22/examples/api_request_parallel_processor.py
+# Several changes were made to make it work with MLflow.
+# Currently, only chat completion is supported.
+
+"""
+API REQUEST PARALLEL PROCESSOR
+
+Using the LangChain API to process lots of text quickly takes some care.
+If you trickle in a million API requests one by one, they'll take days to complete.
+This script parallelizes requests using LangChain API.
+
+Features:
+- Streams requests from file, to avoid running out of memory for giant jobs
+- Makes requests concurrently, to maximize throughput
+- Logs errors, to diagnose problems with requests
+"""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+import time
+import traceback
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+import langchain.chains
+from langchain.callbacks.base import BaseCallbackHandler
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.langchain.utils.chat import (
+    transform_request_json_for_chat_if_necessary,
+    try_transform_response_iter_to_chat_format,
+    try_transform_response_to_chat_format,
+)
+from mlflow.langchain.utils.serialization import convert_to_serializable
+from mlflow.pyfunc.context import (
+    Context,
+    get_prediction_context,
+    maybe_set_prediction_context,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class StatusTracker:
+    """
+    Stores metadata about the script's progress. Only one instance is created.
+    """
+
+    num_tasks_started: int = 0
+    num_tasks_in_progress: int = 0  # script ends when this reaches 0
+    num_tasks_succeeded: int = 0
+    num_tasks_failed: int = 0
+    num_api_errors: int = 0  # excluding rate limit errors, counted above
+    lock: threading.Lock = threading.Lock()
+
+    def start_task(self):
+        with self.lock:
+            self.num_tasks_started += 1
+            self.num_tasks_in_progress += 1
+
+    def complete_task(self, *, success: bool):
+        with self.lock:
+            self.num_tasks_in_progress -= 1
+            if success:
+                self.num_tasks_succeeded += 1
+            else:
+                self.num_tasks_failed += 1
+
+    def increment_num_api_errors(self):
+        with self.lock:
+            self.num_api_errors += 1
+
+
+@dataclass
+class APIRequest:
+    """
+    Stores an API request's inputs, outputs, and other metadata. Contains a method to make an API
+    call.
+
+    Args:
+        index: The request's index in the tasks list
+        lc_model: The LangChain model to call
+        request_json: The request's input data
+        results: The list to append the request's output data to, it's a list of tuples
+            (index, response)
+        errors: A dictionary to store any errors that occur
+        convert_chat_responses: Whether to convert the model's responses to chat format
+        did_perform_chat_conversion: Whether the input data was converted to chat format
+            based on the model's type and input data.
+        stream: Whether the request is a stream request
+        prediction_context: The prediction context to use for the request
+    """
+
+    index: int
+    lc_model: langchain.chains.base.Chain
+    request_json: dict
+    results: list[tuple[int, str]]
+    errors: dict
+    convert_chat_responses: bool
+    did_perform_chat_conversion: bool
+    stream: bool
+    params: dict[str, Any]
+    prediction_context: Optional[Context] = None
+
+    def _predict_single_input(self, single_input, callback_handlers, **kwargs):
+        config = kwargs.pop("config", {})
+        config["callbacks"] = config.get("callbacks", []) + (callback_handlers or [])
+        if self.stream:
+            return self.lc_model.stream(single_input, config=config, **kwargs)
+        if hasattr(self.lc_model, "invoke"):
+            return self.lc_model.invoke(single_input, config=config, **kwargs)
+        else:
+            # for backwards compatibility, __call__ is deprecated and will be removed in 0.3.0
+            # kwargs shouldn't have config field if invoking with __call__
+            return self.lc_model(single_input, callbacks=callback_handlers, **kwargs)
+
+    def _try_convert_response(self, response):
+        if self.stream:
+            return try_transform_response_iter_to_chat_format(response)
+        else:
+            return try_transform_response_to_chat_format(response)
+
+    def single_call_api(self, callback_handlers: Optional[list[BaseCallbackHandler]]):
+        from langchain.schema import BaseRetriever
+
+        from mlflow.langchain.utils import langgraph_types, lc_runnables_types
+
+        if isinstance(self.lc_model, BaseRetriever):
+            # Retrievers are invoked differently than Chains
+            response = self.lc_model.get_relevant_documents(
+                **self.request_json, callbacks=callback_handlers, **self.params
+            )
+        elif isinstance(self.lc_model, lc_runnables_types() + langgraph_types()):
+            if isinstance(self.request_json, dict):
+                # This is a temporary fix for the case when spark_udf converts
+                # input into pandas dataframe with column name, while the model
+                # does not accept dictionaries as input, it leads to errors like
+                # Expected Scalar value for String field 'query_text'
+                try:
+                    response = self._predict_single_input(
+                        self.request_json, callback_handlers, **self.params
+                    )
+                except TypeError as e:
+                    _logger.debug(
+                        f"Failed to invoke {self.lc_model.__class__.__name__} "
+                        f"with {self.request_json}. Error: {e!r}. Trying to "
+                        "invoke with the first value of the dictionary."
+                    )
+                    self.request_json = next(iter(self.request_json.values()))
+                    (
+                        prepared_request_json,
+                        did_perform_chat_conversion,
+                    ) = transform_request_json_for_chat_if_necessary(
+                        self.request_json, self.lc_model
+                    )
+                    self.did_perform_chat_conversion = did_perform_chat_conversion
+
+                    response = self._predict_single_input(
+                        prepared_request_json, callback_handlers, **self.params
+                    )
+            else:
+                response = self._predict_single_input(
+                    self.request_json, callback_handlers, **self.params
+                )
+
+            if self.did_perform_chat_conversion or self.convert_chat_responses:
+                response = self._try_convert_response(response)
+        else:
+            # return_only_outputs is invalid for stream call
+            if isinstance(self.lc_model, langchain.chains.base.Chain) and not self.stream:
+                kwargs = {"return_only_outputs": True}
+            else:
+                kwargs = {}
+            kwargs.update(**self.params)
+            response = self._predict_single_input(self.request_json, callback_handlers, **kwargs)
+
+            if self.did_perform_chat_conversion or self.convert_chat_responses:
+                response = self._try_convert_response(response)
+            elif isinstance(response, dict) and len(response) == 1:
+                # to maintain existing code, single output chains will still return
+                # only the result
+                response = response.popitem()[1]
+
+        return convert_to_serializable(response)
+
+    def call_api(
+        self, status_tracker: StatusTracker, callback_handlers: Optional[list[BaseCallbackHandler]]
+    ):
+        """
+        Calls the LangChain API and stores results.
+        """
+        _logger.debug(f"Request #{self.index} started with payload: {self.request_json}")
+
+        try:
+            with maybe_set_prediction_context(self.prediction_context):
+                response = self.single_call_api(callback_handlers)
+            _logger.debug(f"Request #{self.index} succeeded with response: {response}")
+            self.results.append((self.index, response))
+            status_tracker.complete_task(success=True)
+        except Exception as e:
+            self.errors[self.index] = (
+                f"error: {e!r} {traceback.format_exc()}\n request payload: {self.request_json}"
+            )
+            status_tracker.increment_num_api_errors()
+            status_tracker.complete_task(success=False)
+
+
+def process_api_requests(
+    lc_model,
+    requests: Optional[list[Union[Any, dict[str, Any]]]] = None,
+    max_workers: int = 10,
+    callback_handlers: Optional[list[BaseCallbackHandler]] = None,
+    convert_chat_responses: bool = False,
+    params: Optional[dict[str, Any]] = None,
+):
+    """
+    Processes API requests in parallel.
+    """
+
+    # initialize trackers
+    retry_queue = queue.Queue()
+    status_tracker = StatusTracker()  # single instance to track a collection of variables
+    next_request = None  # variable to hold the next request to call
+
+    results = []
+    errors = {}
+
+    # Note: we should call `transform_request_json_for_chat_if_necessary`
+    # for the whole batch data, because the conversion should obey the rule
+    # that if any record in the batch can't be converted, then all the record
+    # in this batch can't be converted.
+    (
+        converted_chat_requests,
+        did_perform_chat_conversion,
+    ) = transform_request_json_for_chat_if_necessary(requests, lc_model)
+
+    requests_iter = enumerate(converted_chat_requests)
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        while True:
+            # get next request (if one is not already waiting for capacity)
+            if not retry_queue.empty():
+                next_request = retry_queue.get_nowait()
+                _logger.warning(f"Retrying request {next_request.index}: {next_request}")
+            elif req := next(requests_iter, None):
+                # get new request
+                index, converted_chat_request_json = req
+                next_request = APIRequest(
+                    index=index,
+                    lc_model=lc_model,
+                    request_json=converted_chat_request_json,
+                    results=results,
+                    errors=errors,
+                    convert_chat_responses=convert_chat_responses,
+                    did_perform_chat_conversion=did_perform_chat_conversion,
+                    stream=False,
+                    prediction_context=get_prediction_context(),
+                    params=params,
+                )
+                status_tracker.start_task()
+            else:
+                next_request = None
+
+            # if enough capacity available, call API
+            if next_request:
+                # call API
+                executor.submit(
+                    next_request.call_api,
+                    status_tracker=status_tracker,
+                    callback_handlers=callback_handlers,
+                )
+
+            # if all tasks are finished, break
+            # check next_request to avoid terminating the process
+            # before extra requests need to be processed
+            if status_tracker.num_tasks_in_progress == 0 and next_request is None:
+                break
+
+            time.sleep(0.001)  # avoid busy waiting
+
+        # after finishing, log final status
+        if status_tracker.num_tasks_failed > 0:
+            raise mlflow.MlflowException(
+                f"{status_tracker.num_tasks_failed} tasks failed. Errors: {errors}"
+            )
+
+        return [res for _, res in sorted(results)]
+
+
+def process_stream_request(
+    lc_model,
+    request_json: Union[Any, dict[str, Any]],
+    callback_handlers: Optional[list[BaseCallbackHandler]] = None,
+    convert_chat_responses: bool = False,
+    params: Optional[dict[str, Any]] = None,
+):
+    """
+    Process single stream request.
+    """
+    if not hasattr(lc_model, "stream"):
+        raise MlflowException(
+            f"Model {lc_model.__class__.__name__} does not support streaming prediction output. "
+            "No `stream` method found."
+        )
+
+    (
+        converted_chat_requests,
+        did_perform_chat_conversion,
+    ) = transform_request_json_for_chat_if_necessary(request_json, lc_model)
+
+    api_request = APIRequest(
+        index=0,
+        lc_model=lc_model,
+        request_json=converted_chat_requests,
+        results=None,
+        errors=None,
+        convert_chat_responses=convert_chat_responses,
+        did_perform_chat_conversion=did_perform_chat_conversion,
+        stream=True,
+        prediction_context=get_prediction_context(),
+        params=params,
+    )
+    with maybe_set_prediction_context(api_request.prediction_context):
+        return api_request.single_call_api(callback_handlers)

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

@@ -0,0 +1,337 @@
+import importlib.metadata
+import json
+from typing import Annotated, Any, Optional, TypedDict, Union
+from uuid import uuid4
+
+from packaging.version import Version
+
+try:
+    from langchain_core.messages import AnyMessage, BaseMessage, convert_to_messages
+    from langchain_core.runnables import RunnableConfig
+    from langchain_core.runnables.utils import Input
+
+    try:
+        # LangGraph >= 0.3
+        from langgraph.prebuilt import ToolNode
+    except ImportError as e:
+        # If LangGraph 0.3.x is installed but langgraph_prebuilt is not,
+        # show a friendlier error message
+        if Version(importlib.metadata("langgraph").version) >= Version("0.3.0"):
+            raise ImportError(
+                "Please install `langgraph-prebuilt>=0.1.2` to use MLflow LangGraph ChatAgent "
+                "helpers with LangGraph 0.3.x.\n"
+                "If you already have the proper versions installed, please try running "
+                "`pip install --force-reinstall langgraph`. This is a known issue. See: "
+                "https://github.com/langchain-ai/langgraph/issues/3662"
+            ) from e
+
+        # LangGraph < 0.3
+        from langgraph.prebuilt.tool_node import ToolNode
+
+except ImportError as e:
+    raise ImportError(
+        "Please install `langchain>=0.2.17` and `langgraph>=0.2.0` to use LangGraph ChatAgent"
+        "helpers."
+    ) from e
+
+
+from mlflow.langchain.utils.chat import convert_lc_message_to_chat_message
+from mlflow.types.agent import ChatAgentMessage
+from mlflow.utils.annotations import experimental
+
+
+def _add_agent_messages(left: Union[dict, list[dict]], right: Union[dict, list[dict]]):
+    if not isinstance(left, list):
+        left = [left]
+    if not isinstance(right, list):
+        right = [right]
+    # assign missing ids
+    for i, m in enumerate(left):
+        if isinstance(m, BaseMessage):
+            left[i] = parse_message(m)
+        if left[i].get("id") is None:
+            left[i]["id"] = str(uuid4())
+
+    for i, m in enumerate(right):
+        if isinstance(m, BaseMessage):
+            right[i] = parse_message(m)
+        if right[i].get("id") is None:
+            right[i]["id"] = str(uuid4())
+
+    # merge
+    left_idx_by_id = {m.get("id"): i for i, m in enumerate(left)}
+    merged = left.copy()
+    for m in right:
+        if (existing_idx := left_idx_by_id.get(m.get("id"))) is not None:
+            merged[existing_idx] = m
+        else:
+            merged.append(m)
+    return merged
+
+
+@experimental
+class ChatAgentState(TypedDict):
+    """
+    Helper class that enables building a LangGraph agent that produces ChatAgent-compatible
+    messages as state is updated. Other ChatAgent request fields (custom_inputs, context) and
+    response fields (custom_outputs) are also exposed within the state so they can be used and
+    updated over the course of agent execution. Use this class with
+    :py:class:`ChatAgentToolNode <mlflow.langchain.chat_agent_langgraph.ChatAgentToolNode>`.
+
+    **LangGraph ChatAgent Example**
+
+    This example has been tested to work with LangGraph 0.2.70.
+
+    Step 1: Create the LangGraph Agent
+
+    This example is adapted from LangGraph's
+    `create_react_agent <https://langchain-ai.github.io/langgraph/how-tos/create-react-agent/>`__
+    documentation. The notable differences are changes to be ChatAgent compatible. They include:
+
+    - We use :py:class:`ChatAgentState <mlflow.langchain.chat_agent_langgraph.ChatAgentState>`,
+      which has an internal state of
+      :py:class:`ChatAgentMessage <mlflow.types.agent.ChatAgentMessage>`
+      objects and a ``custom_outputs`` attribute under the hood
+    - We use :py:class:`ChatAgentToolNode <mlflow.langchain.chat_agent_langgraph.ChatAgentToolNode>`
+      instead of LangGraph's ToolNode to enable returning attachments and custom_outputs from
+      LangChain and UnityCatalog Tools
+
+    .. code-block:: python
+
+        from typing import Optional, Sequence, Union
+
+        from langchain_core.language_models import LanguageModelLike
+        from langchain_core.runnables import RunnableConfig, RunnableLambda
+        from langchain_core.tools import BaseTool
+        from langgraph.graph import END, StateGraph
+        from langgraph.graph.graph import CompiledGraph
+        from langgraph.prebuilt import ToolNode
+        from mlflow.langchain.chat_agent_langgraph import ChatAgentState, ChatAgentToolNode
+
+
+        def create_tool_calling_agent(
+            model: LanguageModelLike,
+            tools: Union[ToolNode, Sequence[BaseTool]],
+            agent_prompt: Optional[str] = None,
+        ) -> CompiledGraph:
+            model = model.bind_tools(tools)
+
+            def routing_logic(state: ChatAgentState):
+                last_message = state["messages"][-1]
+                if last_message.get("tool_calls"):
+                    return "continue"
+                else:
+                    return "end"
+
+            if agent_prompt:
+                system_message = {"role": "system", "content": agent_prompt}
+                preprocessor = RunnableLambda(
+                    lambda state: [system_message] + state["messages"]
+                )
+            else:
+                preprocessor = RunnableLambda(lambda state: state["messages"])
+            model_runnable = preprocessor | model
+
+            def call_model(
+                state: ChatAgentState,
+                config: RunnableConfig,
+            ):
+                response = model_runnable.invoke(state, config)
+
+                return {"messages": [response]}
+
+            workflow = StateGraph(ChatAgentState)
+
+            workflow.add_node("agent", RunnableLambda(call_model))
+            workflow.add_node("tools", ChatAgentToolNode(tools))
+
+            workflow.set_entry_point("agent")
+            workflow.add_conditional_edges(
+                "agent",
+                routing_logic,
+                {
+                    "continue": "tools",
+                    "end": END,
+                },
+            )
+            workflow.add_edge("tools", "agent")
+
+            return workflow.compile()
+
+    Step 2: Define the LLM and your tools
+
+    If you want to return attachments and custom_outputs from your tool, you can return a
+    dictionary with keys “content”, “attachments”, and “custom_outputs”. This dictionary will be
+    parsed out by the ChatAgentToolNode and properly stored in your LangGraph's state.
+
+
+    .. code-block:: python
+
+        from random import randint
+        from typing import Any
+
+        from databricks_langchain import ChatDatabricks
+        from langchain_core.tools import tool
+
+
+        @tool
+        def generate_random_ints(min: int, max: int, size: int) -> dict[str, Any]:
+            \"""Generate size random ints in the range [min, max].\"""
+            attachments = {"min": min, "max": max}
+            custom_outputs = [randint(min, max) for _ in range(size)]
+            content = f"Successfully generated array of {size} random ints in [{min}, {max}]."
+            return {
+                "content": content,
+                "attachments": attachments,
+                "custom_outputs": {"random_nums": custom_outputs},
+            }
+
+
+        mlflow.langchain.autolog()
+        tools = [generate_random_ints]
+        llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct")
+        langgraph_agent = create_tool_calling_agent(llm, tools)
+
+
+    Step 3: Wrap your LangGraph agent with ChatAgent
+
+    This makes your agent easily loggable and deployable with the PyFunc flavor in serving.
+
+    .. code-block:: python
+
+        from typing import Any, Generator, Optional
+
+        from langgraph.graph.state import CompiledStateGraph
+        from mlflow.pyfunc import ChatAgent
+        from mlflow.types.agent import (
+            ChatAgentChunk,
+            ChatAgentMessage,
+            ChatAgentResponse,
+            ChatContext,
+        )
+
+
+        class LangGraphChatAgent(ChatAgent):
+            def __init__(self, agent: CompiledStateGraph):
+                self.agent = agent
+
+            def predict(
+                self,
+                messages: list[ChatAgentMessage],
+                context: Optional[ChatContext] = None,
+                custom_inputs: Optional[dict[str, Any]] = None,
+            ) -> ChatAgentResponse:
+                request = {"messages": self._convert_messages_to_dict(messages)}
+
+                messages = []
+                for event in self.agent.stream(request, stream_mode="updates"):
+                    for node_data in event.values():
+                        messages.extend(
+                            ChatAgentMessage(**msg) for msg in node_data.get("messages", [])
+                        )
+                return ChatAgentResponse(messages=messages)
+
+            def predict_stream(
+                self,
+                messages: list[ChatAgentMessage],
+                context: Optional[ChatContext] = None,
+                custom_inputs: Optional[dict[str, Any]] = None,
+            ) -> Generator[ChatAgentChunk, None, None]:
+                request = {"messages": self._convert_messages_to_dict(messages)}
+                for event in self.agent.stream(request, stream_mode="updates"):
+                    for node_data in event.values():
+                        yield from (
+                            ChatAgentChunk(**{"delta": msg}) for msg in node_data["messages"]
+                        )
+
+
+        chat_agent = LangGraphChatAgent(langgraph_agent)
+
+    Step 4: Test out your model
+
+    Call ``.predict()`` and ``.predict_stream`` with dictionaries with the ChatAgentRequest schema.
+
+    .. code-block:: python
+
+        chat_agent.predict({"messages": [{"role": "user", "content": "What is 10 + 10?"}]})
+
+        for event in chat_agent.predict_stream(
+            {"messages": [{"role": "user", "content": "Generate me a few random nums"}]}
+        ):
+            print(event)
+
+    This LangGraph ChatAgent can be logged with the logging code described in the "Logging a
+    ChatAgent" section of the docstring of :py:class:`ChatAgent <mlflow.pyfunc.ChatAgent>`.
+    """
+
+    messages: Annotated[list, _add_agent_messages]
+    context: Optional[dict[str, Any]]
+    custom_inputs: Optional[dict[str, Any]]
+    custom_outputs: Optional[dict[str, Any]]
+
+
+def parse_message(
+    msg: AnyMessage, name: Optional[str] = None, attachments: Optional[dict] = None
+) -> dict[str, Any]:
+    """
+    Parse different LangChain message types into their ChatAgentMessage schema dict equivalents
+    """
+    chat_message_dict = convert_lc_message_to_chat_message(msg).model_dump_compat()
+    chat_message_dict["attachments"] = attachments
+    chat_message_dict["name"] = msg.name or name
+    chat_message_dict["id"] = msg.id
+    # _convert_to_message from langchain_core.messages.utils expects an empty string instead of None
+    if not chat_message_dict.get("content"):
+        chat_message_dict["content"] = ""
+
+    chat_agent_msg = ChatAgentMessage(**chat_message_dict)
+    return chat_agent_msg.model_dump_compat(exclude_none=True)
+
+
+@experimental
+class ChatAgentToolNode(ToolNode):
+    """
+    Helper class to make ToolNodes be compatible with
+    :py:class:`ChatAgentState <mlflow.langchain.chat_agent_langgraph.ChatAgentState>`.
+    Parse ``attachments`` and ``custom_outputs`` keys from the string output of a
+    LangGraph tool.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def invoke(self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any) -> Any:
+        """
+        Wraps the standard ToolNode invoke method to:
+        - Parse ChatAgentState into LangChain messages
+        - Parse dictionary string outputs from both UC function and standard LangChain python tools
+          that include keys ``content``, ``attachments``, and ``custom_outputs``.
+        """
+        messages = input["messages"]
+        for msg in messages:
+            for tool_call in msg.get("tool_calls", []):
+                tool_call["name"] = tool_call["function"]["name"]
+                tool_call["args"] = json.loads(tool_call["function"]["arguments"])
+        input["messages"] = convert_to_messages(messages)
+
+        result = super().invoke(input, config, **kwargs)
+
+        messages = []
+        custom_outputs = None
+        for m in result["messages"]:
+            try:
+                return_obj = json.loads(m.content)
+                if all(key in return_obj for key in ("format", "value", "truncated")):
+                    # Dictionary output with custom_outputs and attachments from a UC function
+                    try:
+                        return_obj = json.loads(return_obj["value"])
+                    except Exception:
+                        pass
+                if "custom_outputs" in return_obj:
+                    custom_outputs = return_obj["custom_outputs"]
+                if m.id is None:
+                    m.id = str(uuid4())
+                messages.append(parse_message(m, attachments=return_obj.get("attachments")))
+            except Exception:
+                messages.append(parse_message(m))
+        return {"messages": messages, "custom_outputs": custom_outputs}

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

@@ -0,0 +1,429 @@
+import importlib
+import inspect
+import logging
+import warnings
+from typing import Any, Generator, Optional
+
+from packaging import version
+
+from mlflow.models.resources import (
+    DatabricksFunction,
+    DatabricksServingEndpoint,
+    DatabricksSQLWarehouse,
+    DatabricksVectorSearchIndex,
+    Resource,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+def _get_embedding_model_endpoint_names(index):
+    embedding_model_endpoint_names = []
+    desc = index.describe()
+    delta_sync_index_spec = desc.get("delta_sync_index_spec", {})
+    embedding_source_columns = delta_sync_index_spec.get("embedding_source_columns", [])
+    for column in embedding_source_columns:
+        embedding_model_endpoint_name = column.get("embedding_model_endpoint_name", None)
+        if embedding_model_endpoint_name:
+            embedding_model_endpoint_names.append(embedding_model_endpoint_name)
+    return embedding_model_endpoint_names
+
+
+def _get_vectorstore_from_retriever(retriever) -> Generator[Resource, None, None]:
+    vectorstore = getattr(retriever, "vectorstore", None)
+    if _isinstance_with_multiple_modules(
+        vectorstore,
+        "DatabricksVectorSearch",
+        ["langchain_databricks", "langchain_community.vectorstores", "langchain.vectorstores"],
+    ):
+        index = vectorstore.index
+        yield DatabricksVectorSearchIndex(index_name=index.name)
+        for embedding_endpoint in _get_embedding_model_endpoint_names(index):
+            yield DatabricksServingEndpoint(endpoint_name=embedding_endpoint)
+
+    embeddings = getattr(vectorstore, "embeddings", None)
+    if _isinstance_with_multiple_modules(
+        embeddings,
+        "DatabricksEmbeddings",
+        ["langchain_databricks", "langchain_community.embeddings", "langchain.embeddings"],
+    ):
+        yield DatabricksServingEndpoint(endpoint_name=embeddings.endpoint)
+
+
+def _is_langchain_community_uc_function_toolkit(obj):
+    try:
+        from langchain_community.tools.databricks import UCFunctionToolkit
+    except Exception:
+        return False
+
+    return isinstance(obj, UCFunctionToolkit)
+
+
+def _is_unitycatalog_tool(obj):
+    try:
+        from unitycatalog.ai.langchain.toolkit import UnityCatalogTool
+    except Exception:
+        return False
+
+    return isinstance(obj, UnityCatalogTool)
+
+
+def _extract_databricks_dependencies_from_tools(tools) -> Generator[Resource, None, None]:
+    if isinstance(tools, list):
+        warehouse_ids = set()
+        for tool in tools:
+            if _isinstance_with_multiple_modules(
+                tool, "BaseTool", ["langchain_core.tools", "langchain_community.tools"]
+            ):
+                # Handle Retriever tools
+                if hasattr(tool.func, "keywords") and "retriever" in tool.func.keywords:
+                    retriever = tool.func.keywords.get("retriever")
+                    yield from _get_vectorstore_from_retriever(retriever)
+                elif _is_unitycatalog_tool(tool):
+                    if warehouse_id := tool.client_config.get("warehouse_id"):
+                        warehouse_ids.add(warehouse_id)
+                    yield DatabricksFunction(function_name=tool.uc_function_name)
+                else:
+                    # Tools here are a part of the BaseTool and have no attribute of a
+                    # WarehouseID Extract the global variables of the function defined
+                    # in the tool to get the UCFunctionToolkit Constants
+                    nonlocal_vars = inspect.getclosurevars(tool.func).nonlocals
+                    if "self" in nonlocal_vars and _is_langchain_community_uc_function_toolkit(
+                        nonlocal_vars.get("self")
+                    ):
+                        uc_function_toolkit = nonlocal_vars.get("self")
+                        # As we are iterating through each tool, adding a warehouse id everytime
+                        # is a duplicative resource. Use a set to dedup warehouse ids and add
+                        # them in the end
+                        warehouse_ids.add(uc_function_toolkit.warehouse_id)
+
+                        # In langchain the names of the tools are modified to have underscores:
+                        # main.catalog.test_func -> main_catalog_test_func
+                        # The original name of the tool is stored as the key in the tools
+                        # dictionary. This code finds the correct tool and extract the key
+                        langchain_tool_name = tool.name
+                        filtered_tool_names = [
+                            tool_name
+                            for tool_name, uc_tool in uc_function_toolkit.tools.items()
+                            if uc_tool.name == langchain_tool_name
+                        ]
+                        # This should always have the length 1
+                        for tool_name in filtered_tool_names:
+                            yield DatabricksFunction(function_name=tool_name)
+        # Add the deduped warehouse ids
+        for warehouse_id in warehouse_ids:
+            yield DatabricksSQLWarehouse(warehouse_id=warehouse_id)
+
+
+def _extract_databricks_dependencies_from_retriever(retriever) -> Generator[Resource, None, None]:
+    # ContextualCompressionRetriever uses attribute "base_retriever"
+    if hasattr(retriever, "base_retriever"):
+        retriever = getattr(retriever, "base_retriever", None)
+
+    # Most other retrievers use attribute "retriever"
+    if hasattr(retriever, "retriever"):
+        retriever = getattr(retriever, "retriever", None)
+
+    # EnsembleRetriever uses attribute "retrievers" for multiple retrievers
+    if hasattr(retriever, "retrievers"):
+        retriever = getattr(retriever, "retrievers", None)
+
+    # If there are multiple retrievers, we iterate over them to get dependencies from each of them
+    if isinstance(retriever, list):
+        for single_retriever in retriever:
+            yield from _get_vectorstore_from_retriever(single_retriever)
+    else:
+        yield from _get_vectorstore_from_retriever(retriever)
+
+
+def _extract_databricks_dependencies_from_llm(llm) -> Generator[Resource, None, None]:
+    if _isinstance_with_multiple_modules(
+        llm, "Databricks", ["langchain.llms", "langchain_community.llms"]
+    ):
+        yield DatabricksServingEndpoint(endpoint_name=llm.endpoint_name)
+
+
+def _extract_databricks_dependencies_from_chat_model(chat_model) -> Generator[Resource, None, None]:
+    if _isinstance_with_multiple_modules(
+        chat_model,
+        "ChatDatabricks",
+        ["langchain_databricks", "langchain.chat_models", "langchain_community.chat_models"],
+    ):
+        yield DatabricksServingEndpoint(endpoint_name=chat_model.endpoint)
+
+
+def _extract_databricks_dependencies_from_tool_nodes(tool_node) -> Generator[Resource, None, None]:
+    try:
+        try:
+            # LangGraph >= 0.3
+            from langgraph.prebuilt import ToolNode
+        except ImportError:
+            # LangGraph < 0.3
+            from langgraph.prebuilt.tool_node import ToolNode
+
+        if isinstance(tool_node, ToolNode):
+            yield from _extract_databricks_dependencies_from_tools(
+                list(tool_node.tools_by_name.values())
+            )
+    except ImportError:
+        pass
+
+
+def _isinstance_with_multiple_modules(
+    object: Any, class_name: str, from_modules: list[str]
+) -> bool:
+    """
+    Databricks components are defined in different modules in LangChain e.g.
+    langchain, langchain_community, langchain_databricks due to historical migrations.
+    To keep backward compatibility, we need to check if the object is an instance of the
+    class defined in any of those different modules.
+
+    Args:
+        object: The object to check
+        class_name: The name of the class to check
+        from_modules: The list of modules to import the class from.
+    """
+    # Suppress LangChainDeprecationWarning for old imports
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", DeprecationWarning)
+
+        for module_path in from_modules:
+            try:
+                module = importlib.import_module(module_path)
+                cls = getattr(module, class_name)
+
+                if cls is not None and isinstance(object, cls):
+                    return True
+            except (ImportError, AttributeError):
+                pass
+
+    return False
+
+
+_LEGACY_MODEL_ATTR_SET = {
+    "llm",  # LLMChain
+    "retriever",  # RetrievalQA
+    "llm_chain",  # StuffDocumentsChain, MapRerankDocumentsChain, MapReduceDocumentsChain
+    "question_generator",  # BaseConversationalRetrievalChain
+    "initial_llm_chain",  # RefineDocumentsChain
+    "refine_llm_chain",  # RefineDocumentsChain
+    "combine_documents_chain",  # RetrievalQA, ReduceDocumentsChain
+    "combine_docs_chain",  # BaseConversationalRetrievalChain
+    "collapse_documents_chain",  # ReduceDocumentsChain,
+    "agent",  # Agent,
+    "tools",  # Tools
+}
+
+
+def _extract_dependency_list_from_lc_model(lc_model) -> Generator[Resource, None, None]:
+    """
+    This function contains the logic to examine a non-Runnable component of a langchain model.
+    The logic here does not cover all legacy chains. If you need to support a custom chain,
+    you need to monkey patch this function.
+    """
+    if lc_model is None:
+        return
+
+    # leaf node
+    yield from _extract_databricks_dependencies_from_chat_model(lc_model)
+    yield from _extract_databricks_dependencies_from_retriever(lc_model)
+    yield from _extract_databricks_dependencies_from_llm(lc_model)
+    yield from _extract_databricks_dependencies_from_tools(lc_model)
+    yield from _extract_databricks_dependencies_from_tool_nodes(lc_model)
+
+    # recursively inspect legacy chain
+    for attr_name in _LEGACY_MODEL_ATTR_SET:
+        yield from _extract_dependency_list_from_lc_model(getattr(lc_model, attr_name, None))
+
+
+def _traverse_runnable(
+    lc_model,
+    visited: Optional[set[int]] = None,
+) -> Generator[Resource, None, None]:
+    """
+    This function contains the logic to traverse a langchain_core.runnables.RunnableSerializable
+    object. It first inspects the current object using _extract_dependency_list_from_lc_model
+    and then, if the current object is a Runnable, it recursively inspects its children returned
+    by lc_model.get_graph().nodes.values().
+    This function supports arbitrary LCEL chain.
+    """
+    import pydantic
+    from langchain_core.runnables import Runnable, RunnableLambda
+
+    visited = visited or set()
+    current_object_id = id(lc_model)
+    if current_object_id in visited:
+        return
+
+    # Visit the current object
+    visited.add(current_object_id)
+    yield from _extract_dependency_list_from_lc_model(lc_model)
+
+    if isinstance(lc_model, Runnable):
+        # Visit the returned graph
+        if isinstance(lc_model, RunnableLambda) and version.parse(
+            pydantic.version.VERSION
+        ) >= version.parse("2.0"):
+            nodes = _get_nodes_from_runnable_lambda(lc_model)
+        else:
+            nodes = _get_nodes_from_runnable_callable(lc_model)
+            # If no nodes are found continue with the default behaviour
+            if len(nodes) == 0:
+                nodes = lc_model.get_graph().nodes.values()
+
+        for node in nodes:
+            yield from _traverse_runnable(node.data, visited)
+    else:
+        # No-op for non-runnable, if any
+        pass
+
+
+def _get_deps_from_closures(lc_model):
+    """
+    In some cases, the dependency extraction of Runnable Lambda fails because the call
+    `inspect.getsource(func)` can fail. This causes deps of RunnableLambda to be empty.
+    Therefore this method adds an additional way of getting dependencies through
+    closure variables.
+
+    TODO: Remove when issue gets resolved: https://github.com/langchain-ai/langchain/issues/27970
+    """
+    if not hasattr(lc_model, "func"):
+        return []
+
+    try:
+        from langchain_core.runnables import Runnable
+
+        closure = inspect.getclosurevars(lc_model.func)
+        candidates = {**closure.globals, **closure.nonlocals}
+        deps = []
+
+        # This code is taken from Langchain deps here: https://github.com/langchain-ai/langchain/blob/14f182795312f01985344576b5199681683641e1/libs/core/langchain_core/runnables/base.py#L4481
+        for _, v in candidates.items():
+            if isinstance(v, Runnable):
+                deps.append(v)
+            elif isinstance(getattr(v, "__self__", None), Runnable):
+                deps.append(v.__self__)
+
+        return deps
+    except Exception:
+        return []
+
+
+def _get_nodes_from_runnable_lambda(lc_model):
+    """
+    This is a workaround for the LangGraph issue: https://github.com/langchain-ai/langgraph/issues/1856
+
+    For RunnableLambda, we calling lc_model.get_graph() to get the nodes, which inspect
+    the input and output schema using wrapped function's type annotation. However, the
+    prebuilt graph (e.g. create_react_agent) from LangGraph uses typing.TypeDict annotation,
+    which is not supported by Pydantic V2 on Python < 3.12. If we try to inspect such
+    function, it will raise the following error:
+
+        pydantic.errors.PydanticUserError: Please use `typing_extensions.TypedDict`
+        instead of`typing.TypedDict` on Python < 3.12. For further information visit
+        https://errors.pydantic.dev/2.9/u/typed-dict-version
+
+    Therefore, we cannot use get_graph() for RunnableLambda until LangGraph fixes this issue.
+    Luckily, we are not interested in the input/output nodes for extracting databricks
+    dependencies. We only care about lc_models.deps, which contains the components that
+    the RunnableLambda depends on. Therefore, this function extracts the necessary parts
+    from the original get_graph() function, dropping the input/output related logic.
+    https://github.com/langchain-ai/langchain/blob/2ea5f60cc5747a334550273a5dba1b70b11414c1/libs/core/langchain_core/runnables/base.py#L4493C1-L4512C46
+    """
+
+    if deps := lc_model.deps or _get_deps_from_closures(lc_model):
+        nodes = []
+        for dep in deps:
+            dep_graph = dep.get_graph()
+            dep_graph.trim_first_node()
+            dep_graph.trim_last_node()
+            nodes.extend(dep_graph.nodes.values())
+    else:
+        nodes = lc_model.get_graph().nodes.values()
+    return nodes
+
+
+def _get_nodes_from_runnable_callable(lc_model):
+    """
+    RunnableLambda has a `deps` property which goes through the function and extracts a
+    ny dependencies. RunnableCallable does not have this property so we cannot derive all
+    the dependencies from the function. This helper method also looks into the function of the
+    callable to retrieve these dependencies.
+
+    The code here is from: https://github.com/langchain-ai/langchain/blob/12fea5b868edd12b0d576e7f8bfc922d0167eeab/libs/core/langchain_core/runnables/base.py#L4467
+    """
+
+    # If Runnable Callable is not importable or if the lc_model is not an instance
+    # of RunnableCallable return early
+    try:
+        from langchain_core.runnables import Runnable
+        from langchain_core.runnables.utils import get_function_nonlocals
+        from langgraph.utils.runnable import RunnableCallable
+
+        if not isinstance(lc_model, RunnableCallable):
+            return []
+    except ImportError:
+        return []
+
+    if hasattr(lc_model, "func"):
+        objects = get_function_nonlocals(lc_model.func)
+    elif hasattr(lc_model, "afunc"):
+        objects = get_function_nonlocals(lc_model.afunc)
+    else:
+        objects = []
+
+    deps = []
+    for obj in objects:
+        if isinstance(obj, Runnable):
+            deps.append(obj)
+        elif isinstance(getattr(obj, "__self__", None), Runnable):
+            deps.append(obj.__self__)
+
+    nodes = []
+    for dep in deps:
+        dep_graph = dep.get_graph()
+        dep_graph.trim_first_node()
+        dep_graph.trim_last_node()
+        nodes.extend(dep_graph.nodes.values())
+    return nodes
+
+
+def _detect_databricks_dependencies(lc_model, log_errors_as_warnings=True) -> list[Resource]:
+    """
+    Detects the databricks dependencies of a langchain model and returns a list of
+    detected endpoint names and index names.
+
+    lc_model can be an arbitrary `chain that is built with LCEL <https://python.langchain.com/docs/modules/chains#lcel-chains>`_,
+    which is a langchain_core.runnables.RunnableSerializable.
+    `Legacy chains <https://python.langchain.com/docs/modules/chains#legacy-chains>`_ have limited
+    support. Only RetrievalQA, StuffDocumentsChain, ReduceDocumentsChain, RefineDocumentsChain,
+    MapRerankDocumentsChain, MapReduceDocumentsChain, BaseConversationalRetrievalChain are
+    supported. If you need to support a custom chain, you need to monkey patch
+    the function mlflow.langchain.databricks_dependencies._extract_dependency_list_from_lc_model().
+
+    For an LCEL chain, all the langchain_core.runnables.RunnableSerializable nodes will be
+    traversed.
+
+    If a retriever is found, it will be used to extract the databricks vector search and embeddings
+    dependencies.
+    If an llm is found, it will be used to extract the databricks llm dependencies.
+    If a chat_model is found, it will be used to extract the databricks chat dependencies.
+    """
+    try:
+        dependency_list = list(_traverse_runnable(lc_model))
+        # Filter out duplicate dependencies so same dependencies are not added multiple times
+        # We can't use set here as the object is not hashable so we need to filter it out manually.
+        unique_dependencies = []
+        for dependency in dependency_list:
+            if dependency not in unique_dependencies:
+                unique_dependencies.append(dependency)
+        return unique_dependencies
+    except Exception:
+        if log_errors_as_warnings:
+            _logger.warning(
+                "Unable to detect Databricks dependencies. "
+                "Set logging level to DEBUG to see the full traceback."
+            )
+            _logger.debug("", exc_info=True)
+            return []
+        raise

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

@@ -0,0 +1,662 @@
+import ast
+import logging
+from contextvars import ContextVar
+from typing import Any, Optional, Sequence, Union
+from uuid import UUID
+
+import pydantic
+from langchain_core.agents import AgentAction, AgentFinish
+from langchain_core.callbacks.base import BaseCallbackHandler
+from langchain_core.documents import Document
+from langchain_core.load.dump import dumps
+from langchain_core.messages import BaseMessage
+from langchain_core.outputs import (
+    ChatGenerationChunk,
+    GenerationChunk,
+    LLMResult,
+)
+from tenacity import RetryCallState
+
+import mlflow
+from mlflow import MlflowClient
+from mlflow.entities import Document as MlflowDocument
+from mlflow.entities import LiveSpan, SpanEvent, SpanStatus, SpanStatusCode, SpanType
+from mlflow.exceptions import MlflowException
+from mlflow.langchain.utils.chat import (
+    convert_lc_generation_to_chat_message,
+    convert_lc_message_to_chat_message,
+)
+from mlflow.pyfunc.context import Context, maybe_set_prediction_context
+from mlflow.tracing.constant import SpanAttributeKey
+from mlflow.tracing.provider import detach_span_from_context, set_span_in_context
+from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
+from mlflow.tracing.utils.token import SpanWithToken
+from mlflow.types.chat import ChatMessage, ChatTool, FunctionToolDefinition
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+from mlflow.utils.autologging_utils import ExceptionSafeAbstractClass
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+
+_logger = logging.getLogger(__name__)
+
+_should_attach_span_to_context = ContextVar("should_attach_span_to_context", default=True)
+
+
+def patched_callback_manager_init(original, self, *args, **kwargs):
+    original(self, *args, **kwargs)
+
+    if not AutoLoggingConfig.init(mlflow.langchain.FLAVOR_NAME).log_traces:
+        return
+
+    for handler in self.inheritable_handlers:
+        if isinstance(handler, MlflowLangchainTracer):
+            return
+
+    _handler = MlflowLangchainTracer()
+    self.add_handler(_handler, inherit=True)
+
+
+def patched_callback_manager_merge(original, self, *args, **kwargs):
+    """
+    Patch BaseCallbackManager.merge to avoid a duplicated callback issue.
+
+    In the above patched __init__, we check `inheritable_handlers` to see if the MLflow tracer
+    is already propagated. This works when the `inheritable_handlers` is specified as constructor
+    arguments. However, in the `merge` method, LangChain does not use constructor but set
+    callbacks via the setter method. This causes duplicated callbacks injection.
+    https://github.com/langchain-ai/langchain/blob/d9a069c414a321e7a3f3638a32ecf8a37ec2d188/libs/core/langchain_core/callbacks/base.py#L962-L982
+    """
+    # Get the MLflow callback inherited from parent
+    inherited = self.inheritable_handlers + args[0].inheritable_handlers
+    inherited_mlflow_cb = next(
+        (cb for cb in inherited if isinstance(cb, MlflowLangchainTracer)), None
+    )
+
+    if not inherited_mlflow_cb:
+        return original(self, *args, **kwargs)
+
+    merged = original(self, *args, **kwargs)
+    # If a new MLflow callback is generated inside __init__, remove it
+    duplicate_mlflow_cbs = [
+        cb
+        for cb in merged.inheritable_handlers
+        if isinstance(cb, MlflowLangchainTracer) and cb != inherited_mlflow_cb
+    ]
+    for cb in duplicate_mlflow_cbs:
+        merged.remove_handler(cb)
+
+    return merged
+
+
+def patched_runnable_sequence_batch(original, self, *args, **kwargs):
+    """
+    Patch to terminate span context attachment during batch execution.
+
+    RunnableSequence's batch() methods are implemented in a peculiar way
+    that iterates on steps->items sequentially within the same thread. For example, if a
+    sequence has 2 steps and the batch size is 3, the execution flow will be:
+      - Step 1 for item 1
+      - Step 1 for item 2
+      - Step 1 for item 3
+      - Step 2 for item 1
+      - Step 2 for item 2
+      - Step 2 for item 3
+    Due to this behavior, we cannot attach the span to the context for this particular
+    API, otherwise spans for different inputs will be mixed up.
+    """
+    original_state = _should_attach_span_to_context.get()
+    _should_attach_span_to_context.set(False)
+    try:
+        return original(self, *args, **kwargs)
+    finally:
+        _should_attach_span_to_context.set(original_state)
+
+
+class MlflowLangchainTracer(BaseCallbackHandler, metaclass=ExceptionSafeAbstractClass):
+    """
+    Callback for auto-logging traces.
+    We need to inherit ExceptionSafeAbstractClass to avoid invalid new
+    input arguments added to original function call.
+
+    Args:
+        prediction_context: Optional prediction context object to be set for the
+            thread-local context. Occasionally this has to be passed manually because
+            the callback may be invoked asynchronously and Langchain doesn't correctly
+            propagate the thread-local context.
+    """
+
+    def __init__(
+        self,
+        prediction_context: Optional[Context] = None,
+    ):
+        # NB: The tracer can handle multiple traces in parallel under multi-threading scenarios.
+        # DO NOT use instance variables to manage the state of single trace.
+        super().__init__()
+        self._mlflow_client = MlflowClient()
+        # run_id: (LiveSpan, OTel token)
+        self._run_span_mapping: dict[str, SpanWithToken] = {}
+        self._prediction_context = prediction_context
+
+    def _get_span_by_run_id(self, run_id: UUID) -> Optional[LiveSpan]:
+        if span_with_token := self._run_span_mapping.get(str(run_id), None):
+            return span_with_token.span
+        raise MlflowException(f"Span for run_id {run_id!s} not found.")
+
+    def _serialize_invocation_params(
+        self, attributes: Optional[dict[str, Any]]
+    ) -> Optional[dict[str, Any]]:
+        """
+        Serialize the 'invocation_params' in the attributes dictionary.
+        If 'invocation_params' contains a key 'response_format' whose value is a subclass
+        of pydantic.BaseModel, replace it with its JSON schema.
+        """
+        if not attributes:
+            return attributes
+
+        invocation_params = attributes.get("invocation_params")
+        if not isinstance(invocation_params, dict):
+            return attributes
+
+        response_format = invocation_params.get("response_format")
+        if isinstance(response_format, type) and issubclass(response_format, pydantic.BaseModel):
+            try:
+                invocation_params["response_format"] = (
+                    response_format.model_json_schema()
+                    if IS_PYDANTIC_V2_OR_NEWER
+                    else response_format.schema()
+                )
+            except Exception as e:
+                _logger.error(
+                    "Failed to generate JSON schema for response_format: %s", e, exc_info=True
+                )
+        return attributes
+
+    def _start_span(
+        self,
+        span_name: str,
+        parent_run_id: Optional[UUID],
+        span_type: str,
+        run_id: UUID,
+        inputs: Optional[Union[str, dict[str, Any]]] = None,
+        attributes: Optional[dict[str, Any]] = None,
+    ) -> LiveSpan:
+        """Start MLflow Span (or Trace if it is root component)"""
+        serialized_attributes = self._serialize_invocation_params(attributes)
+        with maybe_set_prediction_context(self._prediction_context):
+            parent = self._get_parent_span(parent_run_id)
+            if parent:
+                span = self._mlflow_client.start_span(
+                    name=span_name,
+                    request_id=parent.request_id,
+                    parent_id=parent.span_id,
+                    span_type=span_type,
+                    inputs=inputs,
+                    attributes=serialized_attributes,
+                )
+            else:
+                # When parent_run_id is None, this is root component so start trace
+                dependencies_schemas = (
+                    self._prediction_context.dependencies_schemas
+                    if self._prediction_context
+                    else None
+                )
+                span = self._mlflow_client.start_trace(
+                    name=span_name,
+                    span_type=span_type,
+                    inputs=inputs,
+                    attributes=serialized_attributes,
+                    tags=dependencies_schemas,
+                )
+
+            # Attach the span to the current context to mark it "active"
+            token = set_span_in_context(span) if _should_attach_span_to_context.get() else None
+            self._run_span_mapping[str(run_id)] = SpanWithToken(span, token)
+        return span
+
+    def _get_parent_span(self, parent_run_id) -> Optional[LiveSpan]:
+        """
+        Get parent span from multiple sources:
+        1. If there is an active span in current context, use it as parent span
+        2. If parent_run_id is provided, get the corresponding span from the run -> span mapping
+        3. If none of the above, return None
+        """
+        if active_span := mlflow.get_current_active_span():
+            return active_span
+        elif parent_run_id:
+            return self._get_span_by_run_id(parent_run_id)
+        return None
+
+    def _end_span(
+        self,
+        run_id: UUID,
+        span: LiveSpan,
+        outputs=None,
+        attributes=None,
+        status=SpanStatus(SpanStatusCode.OK),
+    ):
+        """Close MLflow Span (or Trace if it is root component)"""
+        try:
+            with maybe_set_prediction_context(self._prediction_context):
+                self._mlflow_client.end_span(
+                    request_id=span.request_id,
+                    span_id=span.span_id,
+                    outputs=outputs,
+                    attributes=attributes,
+                    status=status,
+                )
+        finally:
+            # Span should be detached from the context even when the client.end_span fails
+            st = self._run_span_mapping.pop(str(run_id), None)
+            if _should_attach_span_to_context.get():
+                if st.token is None:
+                    raise MlflowException(
+                        f"Token for span {st.span} is not found. "
+                        "Cannot detach the span from context."
+                    )
+                detach_span_from_context(st.token)
+
+    def flush(self):
+        """Flush the state of the tracer."""
+        # Ideally, all spans should be popped and ended. However, LangChain sometimes
+        # does not trigger the end event properly and some spans may be left open.
+        # To avoid leaking tracing context, we remove all spans from the mapping.
+        for st in self._run_span_mapping.values():
+            if st.token:
+                _logger.debug(f"Found leaked span {st.span}. Force ending it.")
+                detach_span_from_context(st.token)
+
+        self._run_span_mapping = {}
+
+    def _assign_span_name(self, serialized: dict[str, Any], default_name="unknown") -> str:
+        return serialized.get("name", serialized.get("id", [default_name])[-1])
+
+    def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[BaseMessage]],
+        *,
+        run_id: UUID,
+        tags: Optional[list[str]] = None,
+        parent_run_id: Optional[UUID] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        name: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        """Run when a chat model starts running."""
+
+        if metadata:
+            kwargs.update({"metadata": metadata})
+
+        span = self._start_span(
+            span_name=name or self._assign_span_name(serialized, "chat model"),
+            parent_run_id=parent_run_id,
+            span_type=SpanType.CHAT_MODEL,
+            run_id=run_id,
+            inputs=messages,
+            attributes=kwargs,
+        )
+
+        mlflow_messages = [
+            convert_lc_message_to_chat_message(msg)
+            for message_list in messages
+            for msg in message_list
+        ]
+        set_span_chat_messages(span, mlflow_messages)
+
+        if tools := self._extract_tool_definitions(kwargs):
+            set_span_chat_tools(span, tools)
+
+    def on_llm_start(
+        self,
+        serialized: dict[str, Any],
+        prompts: list[str],
+        *,
+        run_id: UUID,
+        tags: Optional[list[str]] = None,
+        parent_run_id: Optional[UUID] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        name: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Run when LLM (non-chat models) starts running."""
+        if metadata:
+            kwargs.update({"metadata": metadata})
+
+        span = self._start_span(
+            span_name=name or self._assign_span_name(serialized, "llm"),
+            parent_run_id=parent_run_id,
+            span_type=SpanType.LLM,
+            run_id=run_id,
+            inputs=prompts,
+            attributes=kwargs,
+        )
+
+        mlflow_messages = [ChatMessage(role="user", content=prompt) for prompt in prompts]
+        set_span_chat_messages(span, mlflow_messages)
+
+        if tools := self._extract_tool_definitions(kwargs):
+            set_span_chat_tools(span, tools)
+
+    def _extract_tool_definitions(self, kwargs: dict[str, Any]) -> list[ChatTool]:
+        raw_tools = kwargs.get("invocation_params", {}).get("tools", [])
+        tools = []
+        for raw_tool in raw_tools:
+            # First, try to parse the raw tool dictionary as OpenAI-style tool
+            try:
+                tool = ChatTool.validate_compat(raw_tool)
+                tools.append(tool)
+            except pydantic.ValidationError:
+                # If not OpenAI style, just try to extract the name and descriptions.
+                if name := raw_tool.get("name"):
+                    tool = ChatTool(
+                        type="function",
+                        function=FunctionToolDefinition(
+                            name=name, description=raw_tool.get("description")
+                        ),
+                    )
+                    tools.append(tool)
+                else:
+                    _logger.warning(f"Failed to parse tool definition for tracing: {raw_tool}.")
+
+        return tools
+
+    def on_llm_new_token(
+        self,
+        token: str,
+        *,
+        chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ):
+        """Run on new LLM token. Only available when streaming is enabled."""
+        llm_span = self._get_span_by_run_id(run_id)
+        event_kwargs = {"token": token}
+        if chunk:
+            event_kwargs["chunk"] = dumps(chunk)
+        llm_span.add_event(
+            SpanEvent(
+                name="new_token",
+                attributes=event_kwargs,
+            )
+        )
+
+    def on_retry(
+        self,
+        retry_state: RetryCallState,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ):
+        """Run on a retry event."""
+        span = self._get_span_by_run_id(run_id)
+        retry_d: dict[str, Any] = {
+            "slept": retry_state.idle_for,
+            "attempt": retry_state.attempt_number,
+        }
+        if retry_state.outcome is None:
+            retry_d["outcome"] = "N/A"
+        elif retry_state.outcome.failed:
+            retry_d["outcome"] = "failed"
+            exception = retry_state.outcome.exception()
+            retry_d["exception"] = str(exception)
+            retry_d["exception_type"] = exception.__class__.__name__
+        else:
+            retry_d["outcome"] = "success"
+            retry_d["result"] = str(retry_state.outcome.result())
+        span.add_event(
+            SpanEvent(
+                name="retry",
+                attributes=retry_d,
+            )
+        )
+
+    def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any):
+        """End the span for an LLM run."""
+        llm_span = self._get_span_by_run_id(run_id)
+
+        # Record the chat messages attribute
+        input_messages = llm_span.get_attribute(SpanAttributeKey.CHAT_MESSAGES) or []
+        output_messages = [
+            convert_lc_generation_to_chat_message(gen)
+            for gen_list in response.generations
+            for gen in gen_list
+        ]
+        set_span_chat_messages(llm_span, input_messages + output_messages)
+        self._end_span(run_id, llm_span, outputs=response)
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ):
+        """Handle an error for an LLM run."""
+        llm_span = self._get_span_by_run_id(run_id)
+        llm_span.add_event(SpanEvent.from_exception(error))
+        self._end_span(run_id, llm_span, status=SpanStatus(SpanStatusCode.ERROR, str(error)))
+
+    def on_chain_start(
+        self,
+        serialized: dict[str, Any],
+        inputs: Union[dict[str, Any], Any],
+        *,
+        run_id: UUID,
+        tags: Optional[list[str]] = None,
+        parent_run_id: Optional[UUID] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        run_type: Optional[str] = None,
+        name: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        """Start span for a chain run."""
+        if metadata:
+            kwargs.update({"metadata": metadata})
+        # not considering streaming events for now
+        self._start_span(
+            span_name=name or self._assign_span_name(serialized, "chain"),
+            parent_run_id=parent_run_id,
+            span_type=SpanType.CHAIN,
+            run_id=run_id,
+            inputs=inputs,
+            attributes=kwargs,
+        )
+
+    def on_chain_end(
+        self,
+        outputs: dict[str, Any],
+        *,
+        run_id: UUID,
+        inputs: Optional[Union[dict[str, Any], Any]] = None,
+        **kwargs: Any,
+    ):
+        """Run when chain ends running."""
+        chain_span = self._get_span_by_run_id(run_id)
+        if inputs:
+            chain_span.set_inputs(inputs)
+        self._end_span(run_id, chain_span, outputs=outputs)
+
+    def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        inputs: Optional[Union[dict[str, Any], Any]] = None,
+        run_id: UUID,
+        **kwargs: Any,
+    ):
+        """Run when chain errors."""
+        chain_span = self._get_span_by_run_id(run_id)
+        if inputs:
+            chain_span.set_inputs(inputs)
+        chain_span.add_event(SpanEvent.from_exception(error))
+        self._end_span(run_id, chain_span, status=SpanStatus(SpanStatusCode.ERROR, str(error)))
+
+    def on_tool_start(
+        self,
+        serialized: dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        tags: Optional[list[str]] = None,
+        parent_run_id: Optional[UUID] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        name: Optional[str] = None,
+        # We don't use inputs here because LangChain override the original inputs
+        # with None for some cases. In order to avoid losing the original inputs,
+        # we try to parse the input_str instead.
+        # https://github.com/langchain-ai/langchain/blob/master/libs/core/langchain_core/tools/base.py#L636-L640
+        inputs: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """Start span for a tool run."""
+        if metadata:
+            kwargs.update({"metadata": metadata})
+
+        # For function calling, input_str can be a stringified dictionary
+        # like "{'key': 'value'}". We try parsing it for better rendering,
+        # but conservatively fallback to original if it fails.
+        try:
+            inputs = ast.literal_eval(input_str)
+        except Exception:
+            inputs = input_str
+
+        self._start_span(
+            span_name=name or self._assign_span_name(serialized, "tool"),
+            parent_run_id=parent_run_id,
+            span_type=SpanType.TOOL,
+            run_id=run_id,
+            inputs=inputs,
+            attributes=kwargs,
+        )
+
+    def on_tool_end(self, output: Any, *, run_id: UUID, **kwargs: Any):
+        """Run when tool ends running."""
+        tool_span = self._get_span_by_run_id(run_id)
+        self._end_span(run_id, tool_span, outputs=output)
+
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ):
+        """Run when tool errors."""
+        tool_span = self._get_span_by_run_id(run_id)
+        tool_span.add_event(SpanEvent.from_exception(error))
+        self._end_span(run_id, tool_span, status=SpanStatus(SpanStatusCode.ERROR, str(error)))
+
+    def on_retriever_start(
+        self,
+        serialized: dict[str, Any],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[list[str]] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        name: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        """Run when Retriever starts running."""
+        if metadata:
+            kwargs.update({"metadata": metadata})
+        self._start_span(
+            span_name=name or self._assign_span_name(serialized, "retriever"),
+            parent_run_id=parent_run_id,
+            span_type=SpanType.RETRIEVER,
+            run_id=run_id,
+            inputs=query,
+            attributes=kwargs,
+        )
+
+    def on_retriever_end(self, documents: Sequence[Document], *, run_id: UUID, **kwargs: Any):
+        """Run when Retriever ends running."""
+        retriever_span = self._get_span_by_run_id(run_id)
+        try:
+            # attempt to convert documents to MlflowDocument
+            documents = [MlflowDocument.from_langchain_document(doc) for doc in documents]
+        except Exception as e:
+            _logger.debug(
+                f"Failed to convert LangChain Document to MLflow Document: {e}",
+                exc_info=True,
+            )
+        self._end_span(
+            run_id,
+            retriever_span,
+            outputs=documents,
+        )
+
+    def on_retriever_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ):
+        """Run when Retriever errors."""
+        retriever_span = self._get_span_by_run_id(run_id)
+        retriever_span.add_event(SpanEvent.from_exception(error))
+        self._end_span(run_id, retriever_span, status=SpanStatus(SpanStatusCode.ERROR, str(error)))
+
+    def on_agent_action(
+        self,
+        action: AgentAction,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Run on agent action.
+
+        NB: Agent action doesn't create a new LangChain Run, so instead of creating a new span,
+        an action will be recorded as an event of the existing span created by a parent chain.
+        """
+        span = self._get_span_by_run_id(run_id)
+        span.add_event(
+            SpanEvent(
+                name="agent_action",
+                attributes={
+                    "tool": action.tool,
+                    "tool_input": dumps(action.tool_input),
+                    "log": action.log,
+                },
+            )
+        )
+
+    def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
+        """Run on agent end."""
+        span = self._get_span_by_run_id(run_id)
+        span.add_event(
+            SpanEvent(
+                name="agent_finish",
+                attributes={"return_values": dumps(finish.return_values), "log": finish.log},
+            )
+        )
+
+    def on_text(
+        self,
+        text: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run on arbitrary text."""
+        try:
+            span = self._get_span_by_run_id(run_id)
+        except MlflowException:
+            _logger.warning("Span not found for text event. Skipping text event logging.")
+        else:
+            span.add_event(
+                SpanEvent(
+                    "text",
+                    attributes={"text": text},
+                )
+            )

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

@@ -0,0 +1,142 @@
+from dataclasses import asdict
+from typing import Any, Iterator
+from uuid import uuid4
+
+from langchain_core.messages.base import BaseMessage
+from langchain_core.output_parsers.transform import BaseTransformOutputParser
+
+from mlflow.models.rag_signatures import (
+    ChainCompletionChoice,
+    Message,
+    StringResponse,
+)
+from mlflow.models.rag_signatures import (
+    ChatCompletionResponse as RagChatCompletionResponse,
+)
+from mlflow.types.agent import ChatAgentChunk, ChatAgentMessage, ChatAgentResponse
+from mlflow.types.llm import (
+    ChatChoice,
+    ChatChoiceDelta,
+    ChatChunkChoice,
+    ChatCompletionChunk,
+    ChatCompletionResponse,
+    ChatMessage,
+)
+from mlflow.utils.annotations import deprecated, experimental
+
+
+@deprecated("mlflow.langchain.output_parser.ChatCompletionOutputParser")
+class ChatCompletionsOutputParser(BaseTransformOutputParser[dict[str, Any]]):
+    """
+    OutputParser that wraps the string output into a dictionary representation of a
+    :py:class:`ChatCompletionResponse`
+    """
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this class is serializable."""
+        return True
+
+    @property
+    def _type(self) -> str:
+        """Return the output parser type for serialization."""
+        return "mlflow_simplified_chat_completions"
+
+    def parse(self, text: str) -> dict[str, Any]:
+        return asdict(
+            RagChatCompletionResponse(
+                choices=[ChainCompletionChoice(message=Message(role="assistant", content=text))],
+                object="chat.completion",
+            )
+        )
+
+
+class ChatCompletionOutputParser(BaseTransformOutputParser[str]):
+    """
+    OutputParser that wraps the string output into a dictionary representation of a
+    :py:class:`ChatCompletionResponse` or :py:class:`ChatCompletionChunk`
+    when streaming
+    """
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this class is serializable."""
+        return True
+
+    @property
+    def _type(self) -> str:
+        """Return the output parser type for serialization."""
+        return "mlflow_chat_completion"
+
+    def parse(self, text: str) -> dict[str, Any]:
+        """Returns the input text as a ChatCompletionResponse with no changes."""
+        return ChatCompletionResponse(
+            choices=[ChatChoice(message=ChatMessage(role="assistant", content=text))]
+        ).to_dict()
+
+    def transform(self, input: Iterator[BaseMessage], config, **kwargs) -> Iterator[dict[str, Any]]:
+        """Returns a generator of ChatCompletionChunk objects"""
+        for chunk in input:
+            yield ChatCompletionChunk(
+                choices=[ChatChunkChoice(delta=ChatChoiceDelta(content=chunk.content))]
+            ).to_dict()
+
+
+@deprecated("mlflow.langchain.output_parser.ChatCompletionOutputParser")
+class StringResponseOutputParser(BaseTransformOutputParser[dict[str, Any]]):
+    """
+    OutputParser that wraps the string output into an dictionary representation of a
+    :py:class:`StringResponse`
+    """
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this class is serializable."""
+        return True
+
+    @property
+    def _type(self) -> str:
+        """Return the output parser type for serialization."""
+        return "mlflow_simplified_str_object"
+
+    def parse(self, text: str) -> dict[str, Any]:
+        return asdict(StringResponse(content=text))
+
+
+@experimental
+class ChatAgentOutputParser(BaseTransformOutputParser[str]):
+    """
+    OutputParser that wraps the string output into a dictionary representation of a
+    :py:class:`ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>` or a
+    :py:class:`ChatAgentChunk <mlflow.types.agent.ChatAgentChunk>` for easy interoperability.
+    """
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this class is serializable."""
+        return True
+
+    @property
+    def _type(self) -> str:
+        """Return the output parser type for serialization."""
+        return "mlflow_chat_agent"
+
+    def parse(self, text: str) -> dict[str, Any]:
+        """
+        Returns the output text as a dictionary representation of a
+        :py:class:`ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>`.
+        """
+        return ChatAgentResponse(
+            messages=[ChatAgentMessage(content=text, role="assistant", id=str(uuid4()))]
+        ).model_dump_compat(exclude_none=True)
+
+    def transform(self, input: Iterator[BaseMessage], config, **kwargs) -> Iterator[dict[str, Any]]:
+        """
+        Returns a generator of
+        :py:class:`ChatAgentChunk <mlflow.types.agent.ChatAgentChunk>` objects
+        """
+        for chunk in input:
+            if chunk.content:
+                yield ChatAgentChunk(
+                    delta=ChatAgentMessage(content=chunk.content, role="assistant", id=chunk.id)
+                ).model_dump_compat(exclude_none=True)

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

@@ -0,0 +1,156 @@
+"""Chain for wrapping a retriever."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import yaml
+from langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun
+from langchain.chains.base import Chain
+from langchain.schema import BaseRetriever, Document
+from pydantic import Extra, Field
+
+from mlflow.utils.annotations import experimental
+
+
+@experimental
+class _RetrieverChain(Chain):
+    """
+    Chain that wraps a retriever for use with MLflow.
+
+    The MLflow ``langchain`` flavor provides the functionality to log a retriever object and
+    evaluate it individually. This is useful if you want to evaluate the quality of the
+    relevant documents returned by a retriever object without directing these documents
+    through a large language model (LLM) to yield a summarized response.
+
+    In order to log the retriever object in the ``langchain`` flavor, the retriever object
+    needs to be wrapped within a ``_RetrieverChain``.
+
+    See ``examples/langchain/retriever_chain.py`` for how to log the ``_RetrieverChain``.
+
+    Args:
+        retriever: The retriever to wrap.
+    """
+
+    input_key: str = "query"
+    output_key: str = "source_documents"
+    retriever: BaseRetriever = Field(exclude=True)
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
+    @property
+    def input_keys(self) -> list[str]:
+        """Return the input keys."""
+        return [self.input_key]
+
+    @property
+    def output_keys(self) -> list[str]:
+        """Return the output keys."""
+        return [self.output_key]
+
+    def _get_docs(self, question: str) -> list[Document]:
+        """Get documents from the retriever."""
+        return self.retriever.get_relevant_documents(question)
+
+    def _call(
+        self,
+        inputs: dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> dict[str, Any]:
+        """Run _get_docs on input query.
+        Returns the retrieved documents under the key 'source_documents'.
+
+        Example:
+
+        .. code-block:: python
+
+            chain = _RetrieverChain(retriever=...)
+            res = chain({"query": "This is my query"})
+            docs = res["source_documents"]
+        """
+        question = inputs[self.input_key]
+        docs = self._get_docs(question)
+        list_of_str_page_content = [doc.page_content for doc in docs]
+        return {self.output_key: json.dumps(list_of_str_page_content)}
+
+    async def _aget_docs(self, question: str) -> list[Document]:
+        """Get documents from the retriever."""
+        return await self.retriever.aget_relevant_documents(question)
+
+    async def _acall(
+        self,
+        inputs: dict[str, Any],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> dict[str, Any]:
+        """Run _get_docs on input query.
+        Returns the retrieved documents under the key 'source_documents'.
+
+        Example:
+
+        .. code-block:: python
+
+            chain = _RetrieverChain(retriever=...)
+            res = chain({"query": "This is my query"})
+            docs = res["source_documents"]
+        """
+        question = inputs[self.input_key]
+        docs = await self._aget_docs(question)
+        list_of_str_page_content = [doc.page_content for doc in docs]
+        return {self.output_key: json.dumps(list_of_str_page_content)}
+
+    @property
+    def _chain_type(self) -> str:
+        """Return the chain type."""
+        return "retriever_chain"
+
+    @classmethod
+    def load(cls, file: Union[str, Path], **kwargs: Any) -> _RetrieverChain:
+        """Load a _RetrieverChain from a file."""
+        # Convert file to Path object.
+        file_path = Path(file) if isinstance(file, str) else file
+        # Load from either json or yaml.
+        if file_path.suffix == ".json":
+            with open(file_path) as f:
+                config = json.load(f)
+        elif file_path.suffix in (".yaml", ".yml"):
+            with open(file_path) as f:
+                # This is to ignore certain tags that are not supported
+                # with pydantic >= 2.0
+                yaml.add_multi_constructor(
+                    "tag:yaml.org,2002:python/object",
+                    lambda loader, suffix, node: None,
+                    Loader=yaml.SafeLoader,
+                )
+                config = yaml.load(f, yaml.SafeLoader)
+        else:
+            raise ValueError("File type must be json or yaml")
+
+        # Override default 'verbose' and 'memory' for the chain
+        if verbose := kwargs.pop("verbose", None):
+            config["verbose"] = verbose
+        if memory := kwargs.pop("memory", None):
+            config["memory"] = memory
+
+        if "_type" not in config:
+            raise ValueError("Must specify a chain Type in config")
+        config_type = config.pop("_type")
+
+        if config_type != "retriever_chain":
+            raise ValueError(f"Loading {config_type} chain not supported")
+
+        retriever = kwargs.pop("retriever", None)
+        if retriever is None:
+            raise ValueError("`retriever` must be present.")
+
+        config.pop("retriever", None)
+
+        return cls(
+            retriever=retriever,
+            **config,
+        )

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

@@ -0,0 +1,527 @@
+from __future__ import annotations
+
+import os
+import re
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Union
+
+import cloudpickle
+import yaml
+
+from mlflow.exceptions import MlflowException
+from mlflow.langchain.utils import (
+    _BASE_LOAD_KEY,
+    _CONFIG_LOAD_KEY,
+    _MODEL_DATA_FOLDER_NAME,
+    _MODEL_DATA_KEY,
+    _MODEL_DATA_PKL_FILE_NAME,
+    _MODEL_DATA_YAML_FILE_NAME,
+    _MODEL_LOAD_KEY,
+    _MODEL_TYPE_KEY,
+    _PICKLE_LOAD_KEY,
+    _RUNNABLE_LOAD_KEY,
+    _load_base_lcs,
+    _load_from_json,
+    _load_from_pickle,
+    _load_from_yaml,
+    _patch_loader,
+    _save_base_lcs,
+    _validate_and_prepare_lc_model_or_path,
+    base_lc_types,
+    custom_type_to_loader_dict,
+    get_unsupported_model_message,
+    lc_runnable_assign_types,
+    lc_runnable_binding_types,
+    lc_runnable_branch_types,
+    lc_runnable_with_steps_types,
+    lc_runnables_types,
+    patch_langchain_type_to_cls_dict,
+    picklable_runnable_types,
+)
+
+if TYPE_CHECKING:
+    from langchain.schema.runnable import Runnable
+
+_STEPS_FOLDER_NAME = "steps"
+_RUNNABLE_STEPS_FILE_NAME = "steps.yaml"
+_BRANCHES_FOLDER_NAME = "branches"
+_MAPPER_FOLDER_NAME = "mapper"
+_RUNNABLE_BRANCHES_FILE_NAME = "branches.yaml"
+_DEFAULT_BRANCH_NAME = "default"
+_RUNNABLE_BINDING_CONF_FILE_NAME = "binding_conf.yaml"
+
+
+@patch_langchain_type_to_cls_dict
+def _load_model_from_config(path, model_config):
+    from langchain.chains.loading import type_to_loader_dict as chains_type_to_loader_dict
+    from langchain.llms import get_type_to_cls_dict as llms_get_type_to_cls_dict
+
+    try:
+        from langchain.prompts.loading import type_to_loader_dict as prompts_types
+    except ImportError:
+        prompts_types = {"prompt", "few_shot_prompt"}
+
+    config_path = os.path.join(path, model_config.get(_MODEL_DATA_KEY, _MODEL_DATA_YAML_FILE_NAME))
+    # Load runnables from config file
+    if config_path.endswith(".yaml"):
+        config = _load_from_yaml(config_path)
+    elif config_path.endswith(".json"):
+        config = _load_from_json(config_path)
+    else:
+        raise MlflowException(
+            f"Cannot load runnable without a config file. Got path {config_path}."
+        )
+    _type = config.get("_type")
+    if _type in chains_type_to_loader_dict:
+        from langchain.chains.loading import load_chain
+
+        return _patch_loader(load_chain)(config_path)
+    elif _type in prompts_types:
+        from langchain.prompts.loading import load_prompt
+
+        return load_prompt(config_path)
+    elif _type in llms_get_type_to_cls_dict():
+        from langchain_community.llms.loading import load_llm
+
+        return _patch_loader(load_llm)(config_path)
+    elif _type in custom_type_to_loader_dict():
+        return custom_type_to_loader_dict()[_type](config)
+    raise MlflowException(f"Unsupported type {_type} for loading.")
+
+
+def _load_model_from_path(path: str, model_config=None):
+    model_load_fn = model_config.get(_MODEL_LOAD_KEY)
+    if model_load_fn == _RUNNABLE_LOAD_KEY:
+        return _load_runnables(path, model_config)
+    if model_load_fn == _BASE_LOAD_KEY:
+        return _load_base_lcs(path, model_config)
+    if model_load_fn == _CONFIG_LOAD_KEY:
+        return _load_model_from_config(path, model_config)
+    if model_load_fn == _PICKLE_LOAD_KEY:
+        return _load_from_pickle(os.path.join(path, model_config.get(_MODEL_DATA_KEY)))
+    raise MlflowException(f"Unsupported model load key {model_load_fn}")
+
+
+def _validate_path(file_path: Union[str, Path]):
+    load_path = Path(file_path)
+    if not load_path.exists() or not load_path.is_dir():
+        raise MlflowException(
+            f"Path {load_path} must be an existing directory in order to load model."
+        )
+    return load_path
+
+
+def _load_runnable_with_steps(file_path: Union[Path, str], model_type: str):
+    """Load the model
+
+    Args:
+        file_path: Path to file to load the model from.
+        model_type: Type of the model to load.
+    """
+    from langchain.schema.runnable import RunnableParallel, RunnableSequence
+
+    load_path = _validate_path(file_path)
+
+    steps_conf_file = load_path / _RUNNABLE_STEPS_FILE_NAME
+    if not steps_conf_file.exists():
+        raise MlflowException(
+            f"File {steps_conf_file} must exist in order to load runnable with steps."
+        )
+    steps_conf = _load_from_yaml(steps_conf_file)
+    steps_path = load_path / _STEPS_FOLDER_NAME
+    _validate_path(steps_path)
+
+    steps = {}
+    # ignore hidden files
+    for step in (f for f in os.listdir(steps_path) if not f.startswith(".")):
+        config = steps_conf.get(step)
+        # load model from the folder of the step
+        runnable = _load_model_from_path(os.path.join(steps_path, step), config)
+        steps[step] = runnable
+
+    if model_type == RunnableSequence.__name__:
+        steps = [value for _, value in sorted(steps.items(), key=lambda item: int(item[0]))]
+        return runnable_sequence_from_steps(steps)
+    if model_type == RunnableParallel.__name__:
+        return RunnableParallel(steps)
+
+
+def runnable_sequence_from_steps(steps):
+    """Construct a RunnableSequence from steps.
+
+    Args:
+        steps: List of steps to construct the RunnableSequence from.
+    """
+    from langchain.schema.runnable import RunnableSequence
+
+    if len(steps) < 2:
+        raise ValueError(f"RunnableSequence must have at least 2 steps, got {len(steps)}.")
+
+    first, *middle, last = steps
+    return RunnableSequence(first=first, middle=middle, last=last)
+
+
+def _load_runnable_branch(file_path: Union[Path, str]):
+    """Load the model
+
+    Args:
+        file_path: Path to file to load the model from.
+    """
+    from langchain.schema.runnable import RunnableBranch
+
+    load_path = _validate_path(file_path)
+
+    branches_conf_file = load_path / _RUNNABLE_BRANCHES_FILE_NAME
+    if not branches_conf_file.exists():
+        raise MlflowException(
+            f"File {branches_conf_file} must exist in order to load runnable with steps."
+        )
+    branches_conf = _load_from_yaml(branches_conf_file)
+    branches_path = load_path / _BRANCHES_FOLDER_NAME
+    _validate_path(branches_path)
+
+    branches = []
+    for branch in os.listdir(branches_path):
+        # load model from the folder of the branch
+        if branch == _DEFAULT_BRANCH_NAME:
+            default_branch_path = branches_path / _DEFAULT_BRANCH_NAME
+            default = _load_model_from_path(
+                default_branch_path, branches_conf.get(_DEFAULT_BRANCH_NAME)
+            )
+        else:
+            branch_tuple = []
+            for i in range(2):
+                config = branches_conf.get(f"{branch}-{i}")
+                runnable = _load_model_from_path(
+                    os.path.join(branches_path, branch, str(i)), config
+                )
+                branch_tuple.append(runnable)
+            branches.append(tuple(branch_tuple))
+
+    # default branch must be the last branch
+    branches.append(default)
+
+    return RunnableBranch(*branches)
+
+
+def _load_runnable_assign(file_path: Union[Path, str]):
+    """Load the model
+
+    Args:
+        file_path: Path to file to load the model from.
+    """
+    from langchain.schema.runnable.passthrough import RunnableAssign
+
+    load_path = _validate_path(file_path)
+
+    mapper_file = load_path / _MAPPER_FOLDER_NAME
+    _validate_path(mapper_file)
+    mapper = _load_runnable_with_steps(mapper_file, "RunnableParallel")
+    return RunnableAssign(mapper)
+
+
+def _load_runnable_binding(file_path: Union[Path, str]):
+    """
+    Load runnable binding model from the path
+    """
+    from langchain.schema.runnable import RunnableBinding
+
+    load_path = _validate_path(file_path)
+
+    model_conf = _load_from_yaml(load_path / _RUNNABLE_BINDING_CONF_FILE_NAME)
+    for field, value in model_conf.items():
+        if _is_json_primitive(value):
+            model_conf[field] = value
+        # value is dictionary
+        else:
+            model_conf[field] = _load_model_from_path(load_path, value)
+    return RunnableBinding(**model_conf)
+
+
+def _save_internal_runnables(runnable, path, loader_fn, persist_dir):
+    conf = {}
+    if isinstance(runnable, lc_runnables_types()):
+        conf[_MODEL_TYPE_KEY] = runnable.__class__.__name__
+        conf.update(_save_runnables(runnable, path, loader_fn, persist_dir))
+    elif isinstance(runnable, base_lc_types()):
+        lc_model = _validate_and_prepare_lc_model_or_path(runnable, loader_fn)
+        conf[_MODEL_TYPE_KEY] = lc_model.__class__.__name__
+        conf.update(_save_base_lcs(lc_model, path, loader_fn, persist_dir))
+    else:
+        conf = {
+            _MODEL_TYPE_KEY: runnable.__class__.__name__,
+            _MODEL_DATA_KEY: _MODEL_DATA_YAML_FILE_NAME,
+            _MODEL_LOAD_KEY: _CONFIG_LOAD_KEY,
+        }
+        model_path = path / _MODEL_DATA_YAML_FILE_NAME
+
+        _warning_if_imported_from_lc_partner_pkg(runnable)
+
+        # Save some simple runnables that langchain natively supports.
+        if hasattr(runnable, "save"):
+            runnable.save(model_path)
+        elif hasattr(runnable, "dict"):
+            runnable_dict = runnable.dict()
+            with open(model_path, "w") as f:
+                yaml.dump(runnable_dict, f, default_flow_style=False)
+            # if the model cannot be loaded back, then `dict` is not enough for saving.
+            _load_model_from_config(path, conf)
+        else:
+            raise Exception("Cannot save runnable without `save` or `dict` methods.")
+    return conf
+
+
+_LC_PARTNER_MODULE_PATTERN = re.compile(
+    r"langchain_(?!core|community|experimental|cli|text-splitters)([a-z0-9-]+)$"
+)
+
+
+def _warning_if_imported_from_lc_partner_pkg(runnable):
+    """
+    Issues a warning if the model contains LangChain partner packages in its requirements.
+
+    Popular integrations like OpenAI have been migrated from the central langchain-community
+    package to their own partner packages (e.g. langchain-openai). However, the class loading
+    mechanism in MLflow does not handle partner packages and always loads the community version.
+    This can lead to unexpected behavior because the community version is no longer maintained.
+    """
+    module = runnable.__module__
+    root_module = module.split(".")[0]
+    if m := _LC_PARTNER_MODULE_PATTERN.match(root_module):
+        warnings.warn(
+            "Your model contains a class imported from the LangChain partner package "
+            f"`langchain-{m.group(1)}`. When loading the model back, MLflow will use the "
+            "community version of the classes instead of the partner packages, which may "
+            "lead to unexpected behavior. To ensure that the model is loaded correctly, "
+            "it is recommended to save the model with the 'model-from-code' method "
+            "instead: https://mlflow.org/docs/latest/models.html#models-from-code"
+        )
+
+
+def _save_runnable_with_steps(model, file_path: Union[Path, str], loader_fn=None, persist_dir=None):  # noqa: D417
+    """Save the model with steps. Currently it supports saving RunnableSequence and
+    RunnableParallel.
+
+    If saving a RunnableSequence, steps is a list of Runnable objects. We save each step to the
+    subfolder named by the step index.
+    e.g.  - model
+            - steps
+              - 0
+                - model.yaml
+              - 1
+                - model.pkl
+            - steps.yaml
+    If saving a RunnableParallel, steps is a dictionary of key-Runnable pairs. We save each step to
+    the subfolder named by the key.
+    e.g.  - model
+            - steps
+              - context
+                - model.yaml
+              - question
+                - model.pkl
+            - steps.yaml
+
+    We save steps.yaml file to the model folder. It contains each step's model's configuration.
+
+    Args:
+        model: Runnable to be saved.
+        file_path: Path to file to save the model to.
+    """
+    # Convert file to Path object.
+    save_path = Path(file_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+
+    # Save steps into a folder
+    steps_path = save_path / _STEPS_FOLDER_NAME
+    steps_path.mkdir()
+
+    steps = get_runnable_steps(model)
+    if isinstance(steps, list):
+        generator = enumerate(steps)
+    elif isinstance(steps, dict):
+        generator = steps.items()
+    else:
+        raise MlflowException(
+            f"Runnable {model} steps attribute must be either a list or a dictionary. "
+            f"Got {type(steps).__name__}."
+        )
+    unsaved_runnables = {}
+    steps_conf = {}
+    for key, runnable in generator:
+        step = str(key)
+        # Save each step into a subfolder named by step
+        save_runnable_path = steps_path / step
+        save_runnable_path.mkdir()
+        try:
+            steps_conf[step] = _save_internal_runnables(
+                runnable, save_runnable_path, loader_fn, persist_dir
+            )
+        except Exception as e:
+            unsaved_runnables[step] = f"{runnable.get_name()} -- {e}"
+
+    if unsaved_runnables:
+        raise MlflowException(f"Failed to save runnable sequence: {unsaved_runnables}.")
+
+    # save steps configs
+    with save_path.joinpath(_RUNNABLE_STEPS_FILE_NAME).open("w") as f:
+        yaml.dump(steps_conf, f, default_flow_style=False)
+
+
+def _save_runnable_branch(model, file_path, loader_fn, persist_dir):
+    """
+    Save runnable branch in to path.
+    """
+    save_path = Path(file_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+    # save branches into a folder
+    branches_path = save_path / _BRANCHES_FOLDER_NAME
+    branches_path.mkdir()
+
+    unsaved_runnables = {}
+    branches_conf = {}
+    for index, branch_tuple in enumerate(model.branches):
+        # Save each branch into a subfolder named by index
+        # and save condition and runnable into subfolder
+        for i, runnable in enumerate(branch_tuple):
+            save_runnable_path = branches_path / str(index) / str(i)
+            save_runnable_path.mkdir(parents=True)
+            branches_conf[f"{index}-{i}"] = {}
+
+            try:
+                branches_conf[f"{index}-{i}"] = _save_internal_runnables(
+                    runnable, save_runnable_path, loader_fn, persist_dir
+                )
+            except Exception as e:
+                unsaved_runnables[f"{index}-{i}"] = f"{runnable.get_name()} -- {e}"
+
+    # save default branch
+    default_branch_path = branches_path / _DEFAULT_BRANCH_NAME
+    default_branch_path.mkdir()
+    try:
+        branches_conf[_DEFAULT_BRANCH_NAME] = _save_internal_runnables(
+            model.default, default_branch_path, loader_fn, persist_dir
+        )
+    except Exception as e:
+        unsaved_runnables[_DEFAULT_BRANCH_NAME] = f"{model.default.get_name()} -- {e}"
+    if unsaved_runnables:
+        raise MlflowException(f"Failed to save runnable branch: {unsaved_runnables}.")
+
+    # save branches configs
+    with save_path.joinpath(_RUNNABLE_BRANCHES_FILE_NAME).open("w") as f:
+        yaml.dump(branches_conf, f, default_flow_style=False)
+
+
+def _save_runnable_assign(model, file_path, loader_fn=None, persist_dir=None):
+    from langchain.schema.runnable import RunnableParallel
+
+    save_path = Path(file_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+    # save mapper into a folder
+    mapper_path = save_path / _MAPPER_FOLDER_NAME
+    mapper_path.mkdir()
+
+    if not isinstance(model.mapper, RunnableParallel):
+        raise MlflowException(
+            f"Failed to save model {model} with type {model.__class__.__name__}. "
+            "RunnableAssign's mapper must be a RunnableParallel."
+        )
+    _save_runnable_with_steps(model.mapper, mapper_path, loader_fn, persist_dir)
+
+
+def _is_json_primitive(value):
+    return (
+        value is None
+        or isinstance(value, (str, int, float, bool))
+        or (isinstance(value, list) and all(_is_json_primitive(v) for v in value))
+    )
+
+
+def _save_runnable_binding(model, file_path, loader_fn=None, persist_dir=None):
+    save_path = Path(file_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+    model_config = {}
+
+    # runnableBinding bound is the real runnable to be invoked
+    model_config["bound"] = _save_internal_runnables(model.bound, save_path, loader_fn, persist_dir)
+
+    # save other fields
+    for field, value in model.dict().items():
+        if _is_json_primitive(value):
+            model_config[field] = value
+        elif field != "bound":
+            model_config[field] = {
+                _MODEL_LOAD_KEY: _PICKLE_LOAD_KEY,
+                _MODEL_DATA_KEY: f"{field}.pkl",
+            }
+            _pickle_object(value, os.path.join(save_path, f"{field}.pkl"))
+
+    # save fields configs
+    with save_path.joinpath(_RUNNABLE_BINDING_CONF_FILE_NAME).open("w") as f:
+        yaml.dump(model_config, f, default_flow_style=False)
+
+
+def _pickle_object(model, path: str):
+    if not path.endswith(".pkl"):
+        raise ValueError(f"File path must end with .pkl, got {path}.")
+    with open(path, "wb") as f:
+        cloudpickle.dump(model, f)
+
+
+def _save_runnables(model, path, loader_fn=None, persist_dir=None):
+    model_data_kwargs = {
+        _MODEL_LOAD_KEY: _RUNNABLE_LOAD_KEY,
+        _MODEL_TYPE_KEY: model.__class__.__name__,
+    }
+    if isinstance(model, lc_runnable_with_steps_types()):
+        model_data_path = _MODEL_DATA_FOLDER_NAME
+        _save_runnable_with_steps(
+            model, os.path.join(path, model_data_path), loader_fn, persist_dir
+        )
+    elif isinstance(model, picklable_runnable_types()):
+        model_data_path = _MODEL_DATA_PKL_FILE_NAME
+        _pickle_object(model, os.path.join(path, model_data_path))
+    elif isinstance(model, lc_runnable_branch_types()):
+        model_data_path = _MODEL_DATA_FOLDER_NAME
+        _save_runnable_branch(model, os.path.join(path, model_data_path), loader_fn, persist_dir)
+    elif isinstance(model, lc_runnable_assign_types()):
+        model_data_path = _MODEL_DATA_FOLDER_NAME
+        _save_runnable_assign(model, os.path.join(path, model_data_path), loader_fn, persist_dir)
+    elif isinstance(model, lc_runnable_binding_types()):
+        model_data_path = _MODEL_DATA_FOLDER_NAME
+        _save_runnable_binding(model, os.path.join(path, model_data_path), loader_fn, persist_dir)
+    else:
+        raise MlflowException.invalid_parameter_value(
+            get_unsupported_model_message(type(model).__name__)
+        )
+    model_data_kwargs[_MODEL_DATA_KEY] = model_data_path
+    return model_data_kwargs
+
+
+def _load_runnables(path, conf):
+    model_type = conf.get(_MODEL_TYPE_KEY)
+    model_data = conf.get(_MODEL_DATA_KEY, _MODEL_DATA_YAML_FILE_NAME)
+    if model_type in (x.__name__ for x in lc_runnable_with_steps_types()):
+        return _load_runnable_with_steps(os.path.join(path, model_data), model_type)
+    if (
+        model_type in (x.__name__ for x in picklable_runnable_types())
+        or model_data == _MODEL_DATA_PKL_FILE_NAME
+    ):
+        return _load_from_pickle(os.path.join(path, model_data))
+    if model_type in (x.__name__ for x in lc_runnable_branch_types()):
+        return _load_runnable_branch(os.path.join(path, model_data))
+    if model_type in (x.__name__ for x in lc_runnable_assign_types()):
+        return _load_runnable_assign(os.path.join(path, model_data))
+    if model_type in (x.__name__ for x in lc_runnable_binding_types()):
+        return _load_runnable_binding(os.path.join(path, model_data))
+    raise MlflowException.invalid_parameter_value(get_unsupported_model_message(model_type))
+
+
+def get_runnable_steps(model: Runnable):
+    try:
+        return model.steps
+    except AttributeError:
+        # RunnableParallel stores steps as `steps__` attribute since version 0.16.0, while it was
+        # stored as `steps` attribute before that and other runnables like RunnableSequence still
+        # has `steps` property.
+        return model.steps__

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