如何用MLflow/OpenAI包装Azure Databricks部署的LLM适配LangChain？

完全可以实现——你可以通过MLflow将Databricks部署的LLM包装成兼容OpenAI API规范的端点，之后就能直接用LangChain的ChatOpenAI类调用，以下是针对新手的分步操作指南：

1. 在Databricks中用MLflow包装模型为OpenAI兼容格式 #

首先在Databricks Notebook中运行以下代码，将你的现有LLM模型包装成符合OpenAI API格式的MLflow模型：

import mlflow
from mlflow.models import ModelSignature
from mlflow.types import DataType, Schema, ColSpec

# 定义匹配OpenAI ChatCompletion的输入输出Schema
input_schema = Schema([
    ColSpec(DataType.string, "messages"),
    ColSpec(DataType.integer, "max_tokens", optional=True),
    ColSpec(DataType.float, "temperature", optional=True)
])
output_schema = Schema([ColSpec(DataType.string, "choices")])
signature = ModelSignature(inputs=input_schema, outputs=output_schema)

class OpenAICompatibleLLM(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # 初始化Databricks客户端
        from databricks.sdk import WorkspaceClient
        self.dbr_client = WorkspaceClient()
        self.target_endpoint = "my-code-llama"  # 你的原始LLM服务端点名称

    def predict(self, context, model_input):
        # 转换输入格式为Databricks Serving接受的结构
        req_payload = {
            "messages": model_input["messages"].tolist()[0],
            "max_tokens": model_input.get("max_tokens", [1000])[0],
            "temperature": model_input.get("temperature", [0.7])[0]
        }
        # 调用原始Databricks LLM端点
        dbr_response = self.dbr_client.serving_endpoints.invoke(
            endpoint_name=self.target_endpoint,
            request=req_payload
        )
        # 转换响应为OpenAI兼容格式
        openai_style_response = {
            "choices": [
                {
                    "message": {"content": dbr_response["choices"][0]["message"]["content"]}
                }
            ]
        }
        return openai_style_response

# 保存包装后的模型到MLflow
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        artifact_path="openai-compatible-llm",
        python_model=OpenAICompatibleLLM(),
        signature=signature,
        pip_requirements=["databricks-sdk>=0.20.0"]
    )

2. 部署兼容OpenAI API的服务端点 #

1. 进入Databricks的MLflow模型注册表，找到刚才保存的openai-compatible-llm模型
2. 点击部署到服务端点，配置计算资源（新手建议选择基础型实例）
3. 部署完成后，复制该端点的invocations地址（格式和你原始的LLM端点类似）

3. 在LangChain中用ChatOpenAI类调用 #

现在可以直接用你期望的方式调用了，注意替换对应的参数：

from langchain.chat_models import ChatOpenAI
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

# 初始化ChatOpenAI实例，指向Databricks的兼容端点
llm = ChatOpenAI(
    openai_api_base="https://adb-17272728282828282.1.azuredatabricks.net/serving-endpoints/[你的新兼容端点名称]/invocations",
    openai_api_key="[你的Databricks个人访问令牌]",  # 生成方式：Databricks右上角用户设置→生成令牌
    model_name="my-code-llama",
    max_tokens=1000,
    streaming=True,
    callbacks=[StreamingStdOutCallbackHandler()]
)

# 测试调用
response = llm.invoke("写一段简单的Python排序代码")
print("\n最终响应：", response.content)

新手必看注意事项 #

• 个人访问令牌（PAT）：必须生成有serving_endpoints:invoke权限的令牌，否则无法调用端点
• 流式支持：如果需要流式输出，确保你的原始Databricks LLM端点支持流式响应，且部署兼容端点时开启流式配置
• 格式匹配：如果调用时报格式错误，检查包装类中的输入输出转换逻辑，确保和OpenAI的ChatCompletion规范完全一致

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