Repository: mlflow/mlflow

Issue:
[FR] Support for other APIs when using LangChain
### Willingness to contribute

No. I cannot contribute this feature at this time.

### Proposal Summary

I realized that logging/loading is only for HuggingFace and OpenAI when using the langchain flavor. It would be nice to have support for other APIs such as ChatOpenAI and Azure OpenAI.

```
llm = ChatOpenAI(temperature=0.1, max_tokens=1000)
chain = LLMChain(llm=llm, prompt=chat_prompt)
with mlflow.start_run():
    model_info = mlflow.langchain.log_model(chain, "langchain_model")
```
```
2024/02/05 09:28:12 WARNING mlflow.langchain.utils: MLflow does not guarantee support for LLMs outside of HuggingFaceHub and OpenAI, found ChatOpenAI
```

I am not sure how we can use custom pyfuncs to add functionality to the existing langchain flavor.

### Motivation

> #### What is the use case for this feature?
To reproduce experiments that involve APIs instead of HuggingFace

> #### Why is this use case valuable to support for MLflow users in general?
When logging only the evaluation results (with a function `model(input_df)` instead of a `logged_model`), we miss out on features such as the ability to register the model that performed best

> #### Why is this use case valuable to support for your project(s) or organization?
Keep track of the prompts and models that performed best by clicking on "Register model"

> #### Why is it currently difficult to achieve this use case?
Unclear how to add functionality to existing langchain flavor


### Details

_No response_

### What component(s) does this bug affect?

- [x] `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
- [x] `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


Discussion:
I have the same problem doing this:
```py
from langchain_openai import ChatOpenAI


llm = ChatOpenAI(model="gpt-4", temperature=0) 
chain = RetrievalQA.from_llm(llm=llm, retriever=db.as_retriever(), prompt=prompt)

with tempfile.TemporaryDirectory() as temp_dir:
    persist_dir = os.path.join(temp_dir, "faiss_index")

    # Persist the db to a local fs folder
    db.save_local(persist_dir)

    def load_retriever(persist_directory):
        db = FAISS.load_local(persist_directory, OpenAIEmbeddings())
        return db.as_retriever()


    with mlflow.start_run(run_name="gpt-4"):
        info = mlflow.langchain.log_model(chain,
                                        artifact_path="retrieval_qa",
                                        loader_fn=load_retriever,
                                        persist_dir=persist_dir)
        
        model = mlflow.pyfunc.load_model(info.model_uri)
```
which breaks with the following error:
```
[python3.10/site-packages/mlflow/langchain/__init__.py:648](.venv/lib/python3.10/site-packages/mlflow/langchain/__init__.py:648), in _load_pyfunc(path)

ValueError: Loading openai-chat LLM not supported
```
Thank you for reporting the issue, @ai-learner-00, @rabah-khalek  !

ChatOpenAI model is not supported in MLflow Langchain flavor yet, due to a known limitation of deserialization in Langchain ([ref](https://github.com/langchain-ai/langchain/pull/8164#issuecomment-1673796125)). We reported the issue to Langchain but it may take time to be able to support that.

In the meantime, you can work around the issue by either:

1. Using the legacy [OpenAI](https://python.langchain.com/docs/integrations/llms/openai) class. (this may not work with the latest Langchain/OpenAI version tho).
2. It appears that the new `lancghain.load` deserialization method works for the new Chat model like ([this](https://github.com/langchain-ai/langchain/pull/8164#issuecomment-1665072602)). You could use this method to create a custom pyfunc model. Please refer to [Serving LLMs with MLflow: Leveraging Custom PyFunc](https://mlflow.org/docs/latest/llms/custom-pyfunc-for-llms/notebooks/custom-pyfunc-advanced-llm.html) for the concrete steps for defining and logging the custom pyfunc model. The example uses Transformers model but you should be able to achieve the same thing for langchain by a few modification like `transformers.AutoModelForCausalLM.from_pretrained(...)` -> `langchain.load.loads(...)`

Sorry for the inconvenience, but I hope this answer helps mitigate the issue!
Thanks a lot @B-Step62,

None of these workarounds work:

> 1. Using the legacy [OpenAI](https://python.langchain.com/docs/integrations/llms/openai) class. (this may not work with the latest Langchain/OpenAI version tho).

Indeed as you said, while the `mlflow.langchain.log_model`/`mlflow.pyfunc.load_model` work with the legacy `langchain_community.llms.OpenAI` (or `langchain.llms.OpenAI`) class, this API was removed from the side of openai so `model.predict(...)` fails (see below).
<details> <summary> code + error </summary> 

```py
from langchain_community.llms import OpenAI
# or from langchain.llms import OpenAI

llm = OpenAI(model="gpt-4", temperature=0) 
chain = RetrievalQA.from_llm(llm=llm, retriever=db.as_retriever(), prompt=prompt)

with tempfile.TemporaryDirectory() as temp_dir:
    persist_dir = os.path.join(temp_dir, "faiss_index")

    # Persist the db to a local fs folder
    db.save_local(persist_dir)

    def load_retriever(persist_directory):
        db = FAISS.load_local(persist_directory, OpenAIEmbeddings())
        return db.as_retriever()


    with mlflow.start_run(run_name="gpt-4"):
        info = mlflow.langchain.log_model(chain,
                                        artifact_path="retrieval_qa",
                                        loader_fn=load_retriever,
                                        persist_dir=persist_dir) # --> this works
        
        model = mlflow.pyfunc.load_model(info.model_uri) # --> this works
        
        model.predict(df_example) # --> this fails
```
```
MlflowException: 2 tasks failed. Errors: {0: 'error: APIRemovedInV1(\'\\n\\nYou tried to access openai.ChatCompletion, but this is no longer supported in openai>=1.0.0
```

</details>

> 2. It appears that the new lancghain.load deserialization method works for the new Chat model like (https://github.com/langchain-ai/langchain/pull/8164#issuecomment-1665072602). You could use this method to create a custom pyfunc model. Please refer to [Serving LLMs with MLflow: Leveraging Custom PyFunc](https://mlflow.org/docs/latest/llms/custom-pyfunc-for-llms/notebooks/custom-pyfunc-advanced-llm.html) for the concrete steps for defining and logging the custom pyfunc model. The example uses Transformers model but you should be able to achieve the same thing for langchain by a few modification like transformers.AutoModelForCausalLM.from_pretrained(...) -> langchain.load.loads(...)

This doesn't work for example code above (with `RetrievalQA` powered by `gpt-4`) with the following error:
```
NotImplementedError: Trying to load an object that doesn't implement serialization
```

Indeed, the problem stems from langchain's support... Appreciated your reactiveness and response here: https://github.com/langchain-ai/langchain/issues/18420 and your 3 proposals. I am happy to contribute to help on any of these 3. This problem has been blocking us for a while


<!-- assign-maintainer -->
@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.
@rabah-khalek Thank you for the detailed response!

> this API was removed from the side of openai so model.predict(...) fails (see below).

Yes, this seems to block usage of OpenAI >= 1.0.

> Indeed, the problem stems from langchain's support... Appreciated your reactiveness and response here: https://github.com/langchain-ai/langchain/issues/18420 and your 3 proposals. I am happy to contribute to help on any of these 3. This problem has been blocking us for a while

Thank you so much for the willingness to contribute! I thought the second option is preferable in long term (it seems the first option was tried multiple times but always closed suddenly like [this](https://github.com/langchain-ai/langchain/pull/12331)), but the error `NotImplementedError: Trying to load an object that doesn't implement serialization` error you've got might indicate it takes some effort to make it serializable. Ideally they could provide some short-term workaround but sadly we haven't get any attention to the issue yet. It would be really appreciated if you could spend some time for investigating if it is feasible to serde the new chat models. I will also try digging into this but may not have much bandwidth very soon.

Resolving patch:
diff --git a/mlflow/langchain/__init__.py b/mlflow/langchain/__init__.py
--- a/mlflow/langchain/__init__.py
+++ b/mlflow/langchain/__init__.py
@@ -53,6 +53,7 @@
     _save_base_lcs,
     _validate_and_wrap_lc_model,
     lc_runnables_types,
+    patch_langchain_type_to_cls_dict,
     register_pydantic_v1_serializer_cm,
 )
 from mlflow.models import Model, ModelInputExample, ModelSignature, get_model_info
@@ -849,7 +850,8 @@ def _load_model_from_local_fs(local_model_path):
         return _load_model_code_path(code_path, config_path)
     else:
         _add_code_from_conf_to_system_path(local_model_path, flavor_conf)
-        return _load_model(local_model_path, flavor_conf)
+        with patch_langchain_type_to_cls_dict():
+            return _load_model(local_model_path, flavor_conf)
 
 
 @experimental
diff --git a/mlflow/langchain/utils.py b/mlflow/langchain/utils.py
--- a/mlflow/langchain/utils.py
+++ b/mlflow/langchain/utils.py
@@ -1,9 +1,11 @@
 """Utility functions for mlflow.langchain."""
 
 import contextlib
+import importlib
 import json
 import logging
 import os
+import re
 import shutil
 import types
 import warnings
@@ -55,9 +57,12 @@
 _UNSUPPORTED_LLM_WARNING_MESSAGE = (
     "MLflow does not guarantee support for LLMs outside of HuggingFaceHub and OpenAI, found %s"
 )
-_UNSUPPORTED_LANGCHAIN_VERSION_ERROR_MESSAGE = (
-    "Saving {instance_type} models is only supported in langchain 0.0.194 and above."
-)
+
+# Minimum version of langchain required to support ChatOpenAI and AzureChatOpenAI in MLflow
+# Before this version, our hacky patching to support loading ChatOpenAI and AzureChatOpenAI
+# will not work.
+_LC_MIN_VERSION_SUPPORT_CHAT_OPEN_AI = Version("0.0.307")
+_CHAT_MODELS_ERROR_MSG = re.compile("Loading (openai-chat|azure-openai-chat) LLM not supported")
 
 logger = logging.getLogger(__name__)
 
@@ -238,6 +243,12 @@ def _get_supported_llms():
     if hasattr(langchain.chat_models, "ChatMlflow"):
         llms.add(langchain.chat_models.ChatMlflow)
 
+    if hasattr(langchain.chat_models, "ChatOpenAI"):
+        llms.add(langchain.chat_models.ChatOpenAI)
+
+    if hasattr(langchain.chat_models, "AzureChatOpenAI"):
+        llms.add(langchain.chat_models.AzureChatOpenAI)
+
     return llms
 
 
@@ -280,14 +291,6 @@ def _validate_and_wrap_lc_model(lc_model, loader_fn):
         )
 
     if special_chain_info := _get_special_chain_info_or_none(lc_model):
-        if isinstance(lc_model, langchain.chains.RetrievalQA) and version.parse(
-            langchain.__version__
-        ) < version.parse("0.0.194"):
-            raise mlflow.MlflowException.invalid_parameter_value(
-                _UNSUPPORTED_LANGCHAIN_VERSION_ERROR_MESSAGE.format(
-                    instance_type=type(lc_model).__name__
-                )
-            )
         if loader_fn is None:
             raise mlflow.MlflowException.invalid_parameter_value(
                 f"For {type(lc_model).__name__} models, a `loader_fn` must be provided."
@@ -317,9 +320,10 @@ def _validate_and_wrap_lc_model(lc_model, loader_fn):
 
 
 def _save_base_lcs(model, path, loader_fn=None, persist_dir=None):
-    import langchain.agents.agent
-    import langchain.chains.base
-    import langchain.chains.llm
+    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 = {
@@ -327,9 +331,9 @@ def _save_base_lcs(model, path, loader_fn=None, persist_dir=None):
         _MODEL_LOAD_KEY: _BASE_LOAD_KEY,
     }
 
-    if isinstance(model, langchain.chains.llm.LLMChain):
+    if isinstance(model, (LLMChain, BaseChatModel)):
         model.save(model_data_path)
-    elif isinstance(model, langchain.agents.agent.AgentExecutor):
+    elif isinstance(model, AgentExecutor):
         if model.agent and model.agent.llm_chain:
             model.agent.llm_chain.save(model_data_path)
 
@@ -384,7 +388,7 @@ def _save_base_lcs(model, path, loader_fn=None, persist_dir=None):
 
         # Save model
         model.save(model_data_path)
-    elif isinstance(model, langchain.chains.base.Chain):
+    elif isinstance(model, Chain):
         logger.warning(
             _UNSUPPORTED_MODEL_WARNING_MESSAGE,
             type(model).__name__,
@@ -473,6 +477,75 @@ def _load_base_lcs(
     return model
 
 
+@contextlib.contextmanager
+def patch_langchain_type_to_cls_dict():
+    """Patch LangChain's type_to_cls_dict config to handle unsupported types like ChatOpenAI.
+
+    The type_to_cls_dict is a hard-coded dictionary in LangChain code base that defines the mapping
+    between the LLM type e.g. "openai" to the loader function for the corresponding LLM class.
+    However, this dictionary doesn't contain some chat models like ChatOpenAI, AzureChatOpenAI,
+    which makes it unable to save and load chains with these models. Ideally, the config should
+    be updated in the LangChain code base, but similar requests have been rejected multiple times
+    in the past, because they consider this serde method to be deprecated, and instead prompt
+    users to use their new serde method https://github.com/langchain-ai/langchain/pull/8164#issuecomment-1659723157.
+    However, we can't simply migrate to the new method because it doesn't support common chains
+    like RetrievalQA, AgentExecutor, etc.
+    Therefore, we apply a hacky solution to patch the type_to_cls_dict from our side to support
+    these models, until a better solution is provided by LangChain.
+    """
+
+    def _load_chat_openai():
+        from langchain.chat_models import ChatOpenAI
+
+        return ChatOpenAI
+
+    def _load_azure_chat_openai():
+        from langchain.chat_models import AzureChatOpenAI
+
+        return AzureChatOpenAI
+
+    def _patched_get_type_to_cls_dict(original):
+        def _wrapped():
+            return {
+                **original(),
+                "openai-chat": _load_chat_openai,
+                "azure-openai-chat": _load_azure_chat_openai,
+            }
+
+        return _wrapped
+
+    # NB: get_type_to_cls_dict() method is defined in the following two modules with the same
+    # name but with slight different elements. This is most likely just a mistake in the
+    # LangChain codebase, but we patch them separately to avoid any potential issues.
+    modules_to_patch = ["langchain.llms", "langchain_community.llms.loading"]
+    originals = {}
+    for module_name in modules_to_patch:
+        try:
+            module = importlib.import_module(module_name)
+            originals[module_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[module_name])
+
+    try:
+        yield
+    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
+
+
 def register_pydantic_serializer():
     """
     Helper function to pickle pydantic fields for pydantic v1.


Test patch:
diff --git a/tests/langchain/test_langchain_model_export.py b/tests/langchain/test_langchain_model_export.py
--- a/tests/langchain/test_langchain_model_export.py
+++ b/tests/langchain/test_langchain_model_export.py
@@ -51,6 +51,7 @@
 from mlflow.deployments import PredictionsResponse
 from mlflow.exceptions import MlflowException
 from mlflow.langchain.api_request_parallel_processor import APIRequest
+from mlflow.langchain.utils import _LC_MIN_VERSION_SUPPORT_CHAT_OPEN_AI
 from mlflow.models.signature import ModelSignature, Schema, infer_signature
 from mlflow.types.schema import Array, ColSpec, DataType, Object, Property
 from mlflow.utils.openai_utils import (
@@ -334,6 +335,58 @@ def test_pyfunc_spark_udf_with_langchain_model(spark):
     assert pdf["answer"].tolist() == [TEST_CONTENT, TEST_CONTENT]
 
 
+@pytest.mark.skipif(
+    Version(langchain.__version__) < _LC_MIN_VERSION_SUPPORT_CHAT_OPEN_AI,
+    reason=f"Chat model loading only works for Langchain>={_LC_MIN_VERSION_SUPPORT_CHAT_OPEN_AI}",
+)
+def test_save_and_load_chat_openai(model_path):
+    from langchain.chat_models import ChatOpenAI
+
+    llm = ChatOpenAI(temperature=0.9)
+    prompt = PromptTemplate.from_template("What is a good name for a company that makes {product}?")
+    chain = LLMChain(llm=llm, prompt=prompt)
+    mlflow.langchain.save_model(chain, model_path)
+
+    loaded_model = mlflow.langchain.load_model(model_path)
+    assert loaded_model == chain
+
+    loaded_pyfunc_model = mlflow.pyfunc.load_model(model_path)
+    prediction = loaded_pyfunc_model.predict([{"product": "Mlflow?"}])
+    assert prediction == [TEST_CONTENT]
+
+
+@pytest.mark.skipif(
+    Version(langchain.__version__) >= _LC_MIN_VERSION_SUPPORT_CHAT_OPEN_AI,
+    reason="This test is for non-supported LC version of loading ChatOpenAI model",
+)
+def test_save_and_load_chat_openai_with_unsupported_version_raise_helpful_message(model_path):
+    from langchain.chat_models import ChatOpenAI
+
+    llm = ChatOpenAI(temperature=0.9)
+    prompt = PromptTemplate.from_template("What is a good name for a company that makes {product}?")
+    chain = LLMChain(llm=llm, prompt=prompt)
+    mlflow.langchain.save_model(chain, model_path)
+
+    with pytest.raises(MlflowException, match="Loading ChatOpenAI chat model is not supported"):
+        mlflow.langchain.load_model(model_path)
+
+
+@pytest.mark.skipif(
+    Version(langchain.__version__) < _LC_MIN_VERSION_SUPPORT_CHAT_OPEN_AI,
+    reason=f"Chat model loading only works for Langchain>={_LC_MIN_VERSION_SUPPORT_CHAT_OPEN_AI}",
+)
+def test_save_and_load_azure_chat_openai(model_path):
+    from langchain.chat_models import AzureChatOpenAI
+
+    llm = AzureChatOpenAI(temperature=0.9)
+    prompt = PromptTemplate.from_template("What is a good name for a company that makes {product}?")
+    chain = LLMChain(llm=llm, prompt=prompt)
+    mlflow.langchain.save_model(chain, model_path)
+
+    loaded_model = mlflow.langchain.load_model(model_path)
+    assert loaded_model == chain
+
+
 def test_langchain_log_huggingface_hub_model_metadata(model_path):
     model = create_huggingface_model(model_path)
     with mlflow.start_run():
@@ -422,10 +475,6 @@ def test_langchain_native_log_and_load_qa_with_sources_chain():
     assert model == loaded_model
 
 
-@pytest.mark.skipif(
-    version.parse(langchain.__version__) < version.parse("0.0.194"),
-    reason="Saving RetrievalQA chains requires langchain>=0.0.194",
-)
 def test_log_and_load_retrieval_qa_chain(tmp_path):
     # Create the vector db, persist the db to a local fs folder
     loader = TextLoader("tests/langchain/state_of_the_union.txt")
@@ -487,10 +536,6 @@ def load_retriever(persist_directory):
     )
 
 
-@pytest.mark.skipif(
-    version.parse(langchain.__version__) < version.parse("0.0.194"),
-    reason="Saving RetrievalQA chains requires langchain>=0.0.194",
-)
 def test_log_and_load_retrieval_qa_chain_multiple_output(tmp_path):
     # Create the vector db, persist the db to a local fs folder
     loader = TextLoader("tests/langchain/state_of_the_union.txt")
@@ -1390,6 +1435,29 @@ def test_save_load_complex_runnable_sequence():
     }
 
 
+@pytest.mark.skipif(
+    Version(langchain.__version__) < _LC_MIN_VERSION_SUPPORT_CHAT_OPEN_AI,
+    reason=f"Chat model loading only works for Langchain>={_LC_MIN_VERSION_SUPPORT_CHAT_OPEN_AI}",
+)
+def test_save_load_runnable_sequence_with_chat_openai():
+    from langchain.schema.output_parser import StrOutputParser
+    from langchain.schema.runnable import RunnableSequence
+    from langchain_community.chat_models import ChatOpenAI
+
+    prompt1 = PromptTemplate.from_template("what is the city {person} is from?")
+    llm = ChatOpenAI(temperature=0.9)
+    model = prompt1 | llm | StrOutputParser()
+
+    with mlflow.start_run():
+        model_info = mlflow.langchain.log_model(model, "model_path")
+
+    loaded_model = mlflow.langchain.load_model(model_info.model_uri)
+    assert type(loaded_model) == RunnableSequence
+    assert type(loaded_model.steps[0]) == PromptTemplate
+    assert type(loaded_model.steps[1]) == ChatOpenAI
+    assert type(loaded_model.steps[2]) == StrOutputParser
+
+
 @pytest.mark.skipif(
     Version(langchain.__version__) < Version("0.0.311"), reason="feature not existing"
 )