MLflow LangChain Model Serving Fails: _UNSUPPORTED_MODEL_ERROR_MESSAGE ImportError

[Darshan13]

I’m trying to log and serve a LangChain-based RFI chatbot model using MLflow in Databricks, but I’m encountering the following error when loading the model:
ImportError: cannot import name '_UNSUPPORTED_MODEL_ERROR_MESSAGE' from 'mlflow.langchain.utils'

Library I am using: 

%pip install \

  "mlflow[databricks]==2.14.1" \

  "langchain==0.3.27" \

  "langchain-community==0.3.27" \

  "langchain-databricks==0.1.2" \

  "databricks-vectorsearch==0.40" \

  "databricks-sdk==0.58.0" \

  "pydantic==2.12.4" \

  "pyspark==3.5.0" \

  "grpcio==1.60.0" \

def get_retriever(persist_dir: str = None😞

    host = os.getenv("DATABRICKS_HOST") or "https://" + spark.conf.get("spark.databricks.workspaceUrl")

    # Try to read token from env first

    token = os.getenv("DATABRICKS_TOKEN")

    if token is None:

        # WARNING: Only for notebook testing.

        token = "************"

    VECTOR_SEARCH_ENDPOINT_NAME = os.getenv("VECTOR_SEARCH_ENDPOINT_NAME", "doc_vector_endpoint")

    INDEX_NAME = os.getenv("VECTOR_SEARCH_INDEX", "*********")

    embedding_model = DatabricksEmbeddings(endpoint="databricks-gte-large-en")

    vsc = VectorSearchClient(workspace_url=host, personal_access_token=token, disable_notice=True)

    vs_index = vsc.get_index(endpoint_name=VECTOR_SEARCH_ENDPOINT_NAME, index_name=INDEX_NAME)

    vectorstore = DatabricksVectorSearch(vs_index, text_column="chunk_text", embedding=embedding_model)

    return vectorstore.as_retriever()

from langchain.chains import RetrievalQA

from langchain.prompts import PromptTemplate

from langchain_community.chat_models import ChatDatabricks

# Initialize the chat model

chat_model = ChatDatabricks(

    endpoint="databricks-meta-llama-3-1-8b-instruct",

    temperature=0.0,

    max_tokens=200

)

# Updated RFI-specific prompt

RFI_TEMPLATE = """

You are an expert RFI (Request for Information) answering assistant. Use only the provided

documents to answer the user's question accurately and concisely.

If the information is not available in the provided context, respond:

"I could not find that information in the available documents."

Include any metadata (such as file name, page number, or sheet name) if it helps the user

understand the source of your answer.

Context (includes excerpts and metadata):

{context}

Question:

{question}

Please provide your answer in the following format:

**Answer:** <Your concise, factual response>

Avoid speculation or assumptions. Respond in a clear and professional tone.

"""

prompt = PromptTemplate(template=RFI_TEMPLATE, input_variables=["context", "question"])

# Create retriever (make sure your get_retriever() returns the vector retriever with metadata)

retriever = get_retriever()

# Create the RFI retrieval chain

rfi_chain = RetrievalQA.from_chain_type(

    llm=chat_model,

    chain_type='stuff',

    retriever=retriever,

    chain_type_kwargs={"prompt": prompt}

)

from mlflow.models import infer_signature

import mlflow

import langchain

import mlflow.langchain

import os

# ---------------------------------------------------------------------

# Environment + registry setup

# ---------------------------------------------------------------------

os.environ["MLFLOW_HOME"] = "/databricks/mlflow"

mlflow.set_registry_uri("databricks-uc")

model_name = "*************"

# ---------------------------------------------------------------------

# Start MLflow run

# ---------------------------------------------------------------------

with mlflow.start_run(run_name="rfi_chatbot") as run:

    question = "What was xyz's total revenue in 2024 Q3?"

    result = rfi_chain({"query": question})  # your chain must return dict {"result": ...}

    # Create MLflow model signature

    signature = infer_signature(

        {"query": question},

        {"result": result["result"]}

    )

    # -----------------------------------------------------------------

    # Log the LangChain model to Unity Catalog

    # -----------------------------------------------------------------

    model_info = mlflow.langchain.log_model(

        lc_model=rfi_chain,

        loader_fn=get_retriever,  

        artifact_path="model",

        registered_model_name=model_name,

        signature=signature,

        input_example={"query": question},

        pip_requirements=[

    "mlflow[databricks]==2.14.1",

    "langchain==0.3.27",

    "langchain-community==0.3.27",

    "langchain-databricks==0.1.2",

    "databricks-vectorsearch==0.40",

    "databricks-sdk==0.58.0",

    "pydantic>=2.0.0,<3.0.0",

    "pyspark==3.5.0",

    "grpcio==1.60.0"

]

    )

print("Model successfully logged to Unity Catalog as:", model_name)

print("Run ID:", run.info.run_id)

so its registered my model in unity catalog successfully but when i create a endpoint serving its not working and throwing error: ImportError: cannot import name '_UNSUPPORTED_MODEL_ERROR_MESSAGE' from 'mlflow.langchain.utils'

[KaushalVachhani]

@Darshan13 , Based on the error, It looks like when Databricks Model Serving tries to load your model, MLflow's internal code attempts to import _UNSUPPORTED_MODEL_ERROR_MESSAGE which doesn't exist in the LangChain utils structure.

Are you successfully loading the same model from UC and performing predictions within the notebook (Excluding model serving)? 

I still feel that this might be a version compatibility issue between MLflow 2.14.1 and LangChain 0.3.27. Mlflow version is comparatively older than LangChain. I suggest giving it a try by upgrading MLflow and seeing how it goes.