[BUG] langchain with binding failing to load model

[harshilprajapati96]

Issues Policy acknowledgement

• I have read and agree to submit bug reports in accordance with the issues policy

Where did you encounter this bug?

Local machine

Willingness to contribute

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

MLflow version

• Client: 2.12.1
• Tracking server: 2.12.1

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS Sonoma Version 14.4.1
• Python version: 3.11.7
• yarn version, if running the dev UI:

Describe the problem

RunnableSequence with ChatOpenAI model & bind function call is logged with mlflow.langchain.log_model but while deserializing it fails to load the model mlflow.langchain.load_model

Tracking information

System information: Darwin Darwin Kernel Version 23.4.0: Fri Mar 15 00:12:49 PDT 2024; root:xnu-10063.101.17~1/RELEASE_ARM64_T6020
Python version: 3.11.7
MLflow version: 2.12.1
MLflow module location: <project>/.venv/lib/python3.11/site-packages/mlflow/__init__.py
Tracking URI: file://<project>/mlruns
Registry URI: file://<project>/mlruns
MLflow dependencies: 
  Flask: 3.0.3
  Jinja2: 3.1.3
  aiohttp: 3.9.5
  alembic: 1.13.1
  boto3: 1.34.88
  botocore: 1.34.88
  click: 8.1.7
  cloudpickle: 3.0.0
  docker: 7.0.0
  entrypoints: 0.4
  gitpython: 3.1.43
  graphene: 3.3
  gunicorn: 21.2.0
  importlib-metadata: 7.1.0
  markdown: 3.6
  matplotlib: 3.8.4
  numpy: 1.26.4
  packaging: 23.2
  pandas: 2.2.2
  protobuf: 5.26.1
  pyarrow: 15.0.2
  pydantic: 2.7.0
  pytz: 2024.1
  pyyaml: 6.0.1
  querystring-parser: 1.2.4
  requests: 2.31.0
  scikit-learn: 1.4.2
  scipy: 1.13.0
  sqlalchemy: 2.0.29
  sqlparse: 0.5.0
  tiktoken: 0.6.0

Code to reproduce issue

import mlflow
mlflow.set_tracking_uri("http://localhost:8080/")
mlflow.set_experiment("langchain")

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_community.chat_models import ChatOpenAI


function = {
    "name": "solver",
    "description": "Formulates and solves an equation",
    "parameters": {
        "type": "object",
        "properties": {
            "equation": {
                "type": "string",
                "description": "The algebraic expression of the equation",
            },
            "solution": {
                "type": "string",
                "description": "The solution to the equation",
            },
        },
        "required": ["equation", "solution"],
    },
}
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "Write out the following equation using algebraic symbols then solve it.",
        ),
        ("human", "{equation_statement}"),
    ]
)
model = ChatOpenAI(model="gpt-4", temperature=0).bind(function_call={"name": "solver"}, functions=[function])
chain = {"equation_statement": RunnablePassthrough()} | prompt | model
print(chain.invoke("x raised to the third plus seven equals 12"))

with mlflow.start_run():
    model_info = mlflow.langchain.log_model(
        chain,
        "langchain_model",
        code_paths=["./config.yaml"],
        pip_requirements=[
            "boto3==1.34.88",
            "langchain-openai==0.1.3",
            "langchain==0.1.16",
            "mlflow==2.12.1",
            "spacy==3.7.4",
            "textstat==0.7.3",
        ],
    )

loaded_model = mlflow.langchain.load_model(model_info.model_uri)

ref: https://python.langchain.com/docs/expression_language/primitives/binding/#attaching-openai-tools

Stack trace

Traceback (most recent call last):
  File "<project>/bug.py", line 61, in <module>
    loaded_model = mlflow.langchain.load_model(model_info.model_uri)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project>/.venv/lib/python3.11/site-packages/mlflow/langchain/__init__.py", line 873, in load_model
    return _load_model_from_local_fs(local_model_path)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project>/.venv/lib/python3.11/site-packages/mlflow/langchain/__init__.py", line 846, in _load_model_from_local_fs
    return _load_model(local_model_path, flavor_conf)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project>/.venv/lib/python3.11/site-packages/mlflow/langchain/__init__.py", line 561, in _load_model
    model = _load_runnables(local_model_path, flavor_conf)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project>/.venv/lib/python3.11/site-packages/mlflow/langchain/runnables.py", line 430, in _load_runnables
    return _load_runnable_with_steps(os.path.join(path, model_data), model_type)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project>/.venv/lib/python3.11/site-packages/mlflow/langchain/runnables.py", line 133, in _load_runnable_with_steps
    runnable = _load_model_from_path(os.path.join(steps_path, step), config)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project>/.venv/lib/python3.11/site-packages/mlflow/langchain/runnables.py", line 94, in _load_model_from_path
    return _load_model_from_config(path, model_config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project>/.venv/lib/python3.11/site-packages/mlflow/langchain/runnables.py", line 84, in _load_model_from_config
    raise MlflowException(f"Unsupported type {_type} for loading.")
mlflow.exceptions.MlflowException: Unsupported type None for loading.```


### Other info / logs
N/A

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
• 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

[daniellok-db]

cc @serena-ruan could you help me take a look at this one? I'm not sure why _type would be None here

[harshilprajapati96]

cc @serena-ruan could you help me take a look at this one? I'm not sure why _type would be None here

Any update? I can assist with fix for this bug with guidance

[serena-ruan]

This is not supported yet, we've fixed the issue in latest master, so the model shouldn't be able to be saved at first place.
Langchain doesn't support native saving/loading for RunnableBinding model, we internally use dict to save, but this function doesn't contain all the params needed for loading. We have supporting RunnableBinding in our roadmap, but it's not in a short time as we have some more urgent work :)

[github-actions]

@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.