MLflow LangChain Model Serving Fails: _UNSUPPORTED_MODEL_ERROR_MESSAGE ImportError

Darshan13
New Contributor

ā€Ž11-11-2025 01:40 AM

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'






0 Kudos

KaushalVachhani
Databricks Employee
Databricks Employee

ā€Ž11-11-2025 05:38 AM

@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.