代码可创建预测模型，Databricks模型服务部署遇依赖错误求助

我编写了如下代码：

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

llm = OpenAI(temperature=0)

class UCBot():

    def __init__(self, llm):
        self.llm = llm
        self.toolkit = SQLDatabaseToolkit(db=SQLDatabase.from_databricks(catalog="samples", schema="nyctaxi"), llm=llm)
        self.agent = create_sql_agent(llm=self.llm, toolkit=self.toolkit, verbose=True, top_k=1)

    def get_answer(self, question):
        return self.agent.run(question)
    
class MLflowUCBot(mlflow.pyfunc.PythonModel):
    def __init__(self, llm):
        self.llm = llm

    def predict(self, context, input):
        ucbot = UCBot(self.llm)
        return ucbot.get_answer(input)


# Persist model to mlflow
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        python_model=MLflowUCBot(llm),
        extra_pip_requirements=['langchain', 'databricks-sql-connector', 'sqlalchemy', 'openai'],
        artifact_path='model',
        registered_model_name="my_model",
        input_example={"input":"how many tables?"}
    )

该代码可成功创建模型并执行预测，但尝试通过Databricks模型服务进行封装部署时，先后出现以下错误：

1. 加载模型时提示：An error occurred while loading the model. No module named 'openai'；
2. 添加openai依赖后，仍报错：An error occurred while loading the model. No module named 'openai.api_resources'。

请问我的代码能否在Databricks模型服务中封装部署？如何解决上述错误？

你的代码可以在Databricks模型服务中部署，问题根源在于依赖版本兼容性及模型初始化逻辑，以下是具体解决步骤：

1. 锁定兼容的OpenAI版本 #

openai.api_resources不存在的错误，是因为OpenAI v1.x版本重构了包结构，移除了该模块，而你使用的LangChain版本依赖旧版OpenAI（v0.x）。需要指定明确的兼容版本：

• 修改extra_pip_requirements，将openai替换为具体的旧版版本号，比如：

  extra_pip_requirements=['langchain', 'databricks-sql-connector', 'sqlalchemy', 'openai==0.28.1']
  

2. 调整模型初始化逻辑 #

当前代码在MLflowUCBot初始化时传入已实例化的llm，且每次预测都重新创建UCBot实例，容易导致模型加载时因环境未就绪（如DB连接、API密钥未配置）出错。建议改用load_context方法延迟初始化：

class MLflowUCBot(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # 在模型加载阶段完成LLM和Agent的初始化
        self.llm = OpenAI(temperature=0)
        self.toolkit = SQLDatabaseToolkit(db=SQLDatabase.from_databricks(catalog="samples", schema="nyctaxi"), llm=self.llm)
        self.agent = create_sql_agent(llm=self.llm, toolkit=self.toolkit, verbose=True, top_k=1)

    def predict(self, context, input):
        return self.agent.run(input["input"])

同时修改模型日志代码，不再传入已初始化的llm：

mlflow.pyfunc.log_model(
    python_model=MLflowUCBot(),
    extra_pip_requirements=['langchain', 'databricks-sql-connector', 'sqlalchemy', 'openai==0.28.1'],
    artifact_path='model',
    registered_model_name="my_model",
    input_example={"input":"how many tables?"}
)

3. 配置Databricks环境权限与变量 #

• 确保模型服务使用的服务账号拥有访问samples.nyctaxi目录的权限；
• 在模型服务的环境变量中配置OPENAI_API_KEY，避免硬编码密钥。

内容的提问来源于stack exchange，提问作者John Francis