Repository: mlflow/mlflow
Issue category: Bug Report

[BUG] MLflow langchain does not support logging RunnableWithMessageHistory
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Where did you encounter this bug?

Databricks

### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

- Client: 2.16.2


### System information

- **OS Platform and Distribution**: Linux (5.4.0-1135-azure-fips)
- **Python version**: 3.11.0


### Describe the problem

I am trying to log a Langchain chain for conversational RAG with memory using Langchains RunnableWithMessageHistory. However, I get an error that says that this flavor is not supported. Is there a workaround for this?

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
System information: Linux #142+fips1-Ubuntu SMP Tue Jul 30 21:00:25 UTC 2024
Python version: 3.11.0rc1
MLflow version: 2.16.2
MLflow module location: /local_disk0/.ephemeral_nfs/envs/pythonEnv-ada40d50-865d-42f8-8ded-3ea7c11f55a6/lib/python3.11/site-packages/mlflow/__init__.py
Tracking URI: databricks
Registry URI: databricks-uc
Databricks runtime version: 15.4
MLflow environment variables: 
  MLFLOW_CONDA_HOME: /databricks/conda
  MLFLOW_DEPLOYMENTS_TARGET: databricks
  MLFLOW_GATEWAY_URI: databricks
  MLFLOW_PYTHON_EXECUTABLE: /databricks/spark/scripts/mlflow_python.sh
  MLFLOW_REGISTRY_URI: databricks-uc
  MLFLOW_TRACKING_URI: databricks
MLflow dependencies: 
  Flask: 2.2.5
  Jinja2: 3.1.2
  aiohttp: 3.10.5
  alembic: 1.13.3
  azure-storage-file-datalake: 12.14.0
  boto3: 1.34.39
  botocore: 1.34.39
  docker: 7.1.0
  fastapi: 0.115.0
  google-cloud-storage: 2.10.0
  graphene: 3.3
  gunicorn: 20.1.0
  kubernetes: 31.0.0
  langchain: 0.3.0
  markdown: 3.4.1
  matplotlib: 3.7.2
  mlflow-skinny: 2.16.2
  numpy: 1.23.5
  pandas: 1.5.3
  pyarrow: 14.0.1
  pydantic: 2.9.2
  scikit-learn: 1.3.0
  scipy: 1.11.1
  sqlalchemy: 2.0.35
  tiktoken: 0.7.0
  uvicorn: 0.30.6
  virtualenv: 20.24.2
  watchfiles: 0.24.0
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
from langchain_openai import AzureChatOpenAI
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

# fill with details
llm = AzureChatOpenAI()
vector_search_as_retriever = DatabricksVectorSearch().as_retriever()

contextualize_q_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

history_aware_retriever = create_history_aware_retriever(
    llm, vector_search_as_retriever, contextualize_q_prompt
)

qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)

rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


conversational_rag_chain = RunnableWithMessageHistory(
    rag_chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer",
) 

# Error
with mlflow.start_run(run_name="test"):
    mlflow.set_tag("type", "chain")
    logged_chain_info = mlflow.langchain.log_model(
        lc_model=conversational_rag_chain,  
        artifact_path="chain"
    )
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
MlflowException: MLflow langchain flavor only supports subclasses of (<class 'langchain.chains.base.Chain'>, <class 'langchain.agents.agent.AgentExecutor'>, <class 'langchain_core.retrievers.BaseRetriever'>, <class 'langchain_core.language_models.chat_models.SimpleChatModel'>, <class 'langchain_core.prompts.chat.ChatPromptTemplate'>, <class 'langchain_core.runnables.passthrough.RunnablePassthrough'>, <class 'langchain_core.runnables.base.RunnableLambda'>, <class 'langchain_core.runnables.base.RunnableParallel'>, <class 'langchain_core.runnables.base.RunnableSequence'>, <class 'langchain_core.runnables.branch.RunnableBranch'>, <class 'langchain_core.runnables.passthrough.RunnableAssign'>, <class 'langchain_core.runnables.base.RunnableBinding'>), found RunnableWithMessageHistory.
File <command-2576690084880631>, line 5
      3 with mlflow.start_run(run_name=f"dbdemos_rag_azure"):
      4     mlflow.set_tag("type", "chain")
----> 5     logged_chain_info = mlflow.langchain.log_model(
      6         lc_model=conversational_rag_chain,  # Chain code file e.g., /path/to/the/chain.py 
      7         model_config='rag_multi_chain_config.yaml',  # Chain configuration 
      8         artifact_path="chain"
      9     )
     11 # Test the chain locally
     12 chain = mlflow.langchain.load_model(logged_chain_info.model_uri)

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ada40d50-865d-42f8-8ded-3ea7c11f55a6/lib/python3.11/site-packages/mlflow/tracing/provider.py:268, in trace_disabled.<locals>.wrapper(*args, **kwargs)
    266 disable()
    267 try:
--> 268     is_func_called, result = True, f(*args, **kwargs)
    269 finally:
    270     enable()
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ada40d50-865d-42f8-8ded-3ea7c11f55a6/lib/python3.11/site-packages/mlflow/langchain/__init__.py:549, in log_model(lc_model, artifact_path, conda_env, code_paths, registered_model_name, signature, input_example, await_registration_for, pip_requirements, extra_pip_requirements, metadata, loader_fn, persist_dir, example_no_conversion, run_id, model_config, streamable)
    403 @experimental
    404 @format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
    405 @docstring_version_compatibility_warning(FLAVOR_NAME)
   (...)
    424     streamable=None,
    425 ):
    426     """
    427     Log a LangChain model as an MLflow artifact for the current run.
    428 
   (...)
    547         metadata of the logged model.
    548     """
--> 549     return Model.log(
    550         artifact_path=artifact_path,
    551         flavor=mlflow.langchain,
    552         registered_model_name=registered_model_name,
    553         lc_model=lc_model,
    554         conda_env=conda_env,
    555         code_paths=code_paths,
    556         signature=signature,
    557         input_example=input_example,
    558         await_registration_for=await_registration_for,
    559         pip_requirements=pip_requirements,
    560         extra_pip_requirements=extra_pip_requirements,
    561         metadata=metadata,
    562         loader_fn=loader_fn,
    563         persist_dir=persist_dir,
    564         example_no_conversion=example_no_conversion,
    565         run_id=run_id,
    566         model_config=model_config,
    567         streamable=streamable,
    568     )
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ada40d50-865d-42f8-8ded-3ea7c11f55a6/lib/python3.11/site-packages/mlflow/models/model.py:725, in Model.log(cls, artifact_path, flavor, registered_model_name, await_registration_for, metadata, run_id, resources, **kwargs)
    721     run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id
    722 mlflow_model = cls(
    723     artifact_path=artifact_path, run_id=run_id, metadata=metadata, resources=resources
    724 )
--> 725 flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
    726 # `save_model` calls `load_model` to infer the model requirements, which may result in
    727 # __pycache__ directories being created in the model directory.
    728 for pycache in Path(local_path).rglob("__pycache__"):
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ada40d50-865d-42f8-8ded-3ea7c11f55a6/lib/python3.11/site-packages/mlflow/tracing/provider.py:272, in trace_disabled.<locals>.wrapper(*args, **kwargs)
    270             enable()
    271     else:
--> 272         is_func_called, result = True, f(*args, **kwargs)
    273 # We should only catch the exception from disable() and enable()
    274 # and let other exceptions propagate.
    275 except MlflowTracingException as e:
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ada40d50-865d-42f8-8ded-3ea7c11f55a6/lib/python3.11/site-packages/mlflow/langchain/__init__.py:262, in save_model(lc_model, path, conda_env, code_paths, mlflow_model, signature, input_example, pip_requirements, extra_pip_requirements, metadata, loader_fn, persist_dir, example_no_conversion, model_config, streamable)
    259 import langchain
    260 from langchain.schema import BaseRetriever
--> 262 lc_model_or_path = _validate_and_prepare_lc_model_or_path(lc_model, loader_fn, temp_dir)
    264 _validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
    266 path = os.path.abspath(path)
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ada40d50-865d-42f8-8ded-3ea7c11f55a6/lib/python3.11/site-packages/mlflow/langchain/utils/__init__.py:293, in _validate_and_prepare_lc_model_or_path(lc_model, loader_fn, temp_dir)
    290     return _validate_and_get_model_code_path(lc_model, temp_dir)
    292 if not isinstance(lc_model, supported_lc_types()):
--> 293     raise mlflow.MlflowException.invalid_parameter_value(
    294         get_unsupported_model_message(type(lc_model).__name__)
    295     )
    297 _SUPPORTED_LLMS = _get_supported_llms()
    298 if isinstance(lc_model, langchain.chains.llm.LLMChain) and not any(
    299     isinstance(lc_model.llm, supported_llm) for supported_llm in _SUPPORTED_LLMS
    300 ):

```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Functions to edit:
mlflow/models/model.py:_validate_langchain_model

Patch:
diff --git a/mlflow/models/model.py b/mlflow/models/model.py
index 2326c3df57402..7ae1fbede42db 100644
--- a/mlflow/models/model.py
+++ b/mlflow/models/model.py
@@ -1116,9 +1116,20 @@ def update_model_requirements(
 
 
 def _validate_langchain_model(model):
-    from mlflow.langchain import _validate_and_prepare_lc_model_or_path
+    from langchain_core.runnables.base import Runnable
 
-    return _validate_and_prepare_lc_model_or_path(model, None)
+    from mlflow.models.utils import _validate_and_get_model_code_path
+
+    if isinstance(model, str):
+        return _validate_and_get_model_code_path(model, None)
+
+    if not isinstance(model, Runnable):
+        raise MlflowException.invalid_parameter_value(
+            "Model must be a Langchain Runnable type or path to a Langchain model, "
+            f"got {type(model)}"
+        )
+
+    return model
 
 
 def _validate_llama_index_model(model):