MlflowException: Error when attempting to pickle the AgentExecutor tools. This model likely does not support serialization.

[Teezq]

I have tried to log a langchain model build with SQL Agent, but have gotten an error as below. Would appreciate if anyone could help me on this

Code:

#Model Building:

from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.sql_database import SQLDatabase
from langchain import OpenAI

from langchain.agents import AgentExecutor

from langchain.agents.agent_types import AgentType

db = SQLDatabase.from_databricks(catalog="alpha_assistant", schema="default")
llm = OpenAI(temperature=1)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True, agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION)

#Model Logging using mlflow
import os

from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.llms import OpenAI

import mlflow

assert "OPENAI_API_KEY" in os.environ, "Please set the OPENAI_API_KEY environment variable."

with mlflow.start_run():
logged_model = mlflow.langchain.log_model(agent, "langchain_model")

loaded_model = mlflow.pyfunc.load_model(logged_model.model_uri)
print(
loaded_model.predict(
[
{
"input": "What are the number of employees in the company?"
}
]
)
)

Error:
MlflowException: Error when attempting to pickle the AgentExecutor tools. This model likely does not support serialization.

[dosubot[bot]]

🤖

Hello @Teezq,

Nice to meet you! I'm here to help with your LangChain-related questions, bugs, or any other issues you might have. I'm a bot designed to assist while we wait for a human maintainer. Don't worry, together we'll figure out what's happening with your code.

Based on the error message you're encountering, it seems like the issue is related to the serialization of the AgentExecutor tools. The AgentExecutor in LangChain is a class that uses an agent and a set of tools to execute a series of actions. It is responsible for running the agent's plan, handling any errors that occur during the execution, and determining when to stop the execution based on the specified maximum number of iterations or maximum execution time.

The mlflow.langchain.log_model method you're using is likely a part of the mlflow library, which is used for experiment tracking and model management in machine learning. It seems like this method is trying to serialize the AgentExecutor tools, which might not be supported.

There was a similar issue reported in the LangChain repository: Trying to initialize an agent with Cohere and LLM-requiring tool causes a ValidationError. The issue was related to the compatibility of Pydantic version with Replit, and updating Pydantic to a compatible version with Replit solved the problem. Although your issue is not exactly the same, it might be worth checking the versions of the libraries you're using to see if there are any compatibility issues.

If the issue persists, it would be helpful to have more information about the specific tools you're using with the AgentExecutor and the versions of the libraries you're using. This could help identify if there are any known issues with these tools or libraries when it comes to serialization.

Sources

• Trying to initialize an agent with Cohere and LLM-requiring tool causes a ValidationError
• libs/langchain/langchain/agents/agent.py
• libs/langchain/tests/unit_tests/agents/test_agent.py

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[selewaut]

@dosu-bot, @Teezq I've been having same problem when pickling and agent that uses the sparkSQl toolkit, any tips on how to be able to pickle? I'm quite new to langchain. I want to deploy model with mlflow as databricks endpoint. and saving the model is a crucial part.

[prithvikannan]

Hey @Teezq, MLflow maintainer here -- I would recommend you use MLflow's log model as code feature which does not require serialization. Hope this helps!