Skip to main content
Databricks Lakehouse 平台在一个平台上统一了数据、分析和 AI。 本指南提供了一个快速概览,介绍如何使用 Databricks 聊天模型

概述

ChatDatabricks 类封装了一个托管在Databricks Model Serving上的聊天模型端点。此示例笔记本展示了如何封装服务端点并在您的 LangChain 应用程序中将其用作聊天模型。

集成细节

可序列化下载量版本
ChatDatabricksdatabricks-langchainbetaPyPI - DownloadsPyPI - Version

模型功能

工具调用结构化输出图像输入音频输入视频输入标记级别流式传输原生异步标记使用情况概率

支持的方法

ChatDatabricks 支持 ChatModel 的所有方法,包括异步 API。

端点要求

ChatDatabricks 所封装的服务端点必须具有与 OpenAI 兼容的聊天输入/输出格式(参考)。只要输入格式兼容,ChatDatabricks 就可以用于托管在 Databricks Model Serving 上的任何端点类型:
  1. 基础模型 - 一个包含最新基础模型(如 DRBX、Llama3、Mixtral-8x7B 等)的精选列表。这些端点在您的 Databricks 工作区中可以立即使用,无需进行任何设置。
  2. 自定义模型 - 您还可以通过 MLflow 部署自定义模型到服务端点,并选择框架(如 LangChain、Pytorch、Transformers 等)。
  3. 外部模型 - Databricks 端点可以代理托管在 Databricks 之外的模型,例如专有的模型服务(如 OpenAI GPT4)。

配置

要访问 Databricks 模型,您需要创建一个 Databricks 账户、设置凭据(仅限于不在 Databricks 工作区外运行),并安装所需的包。

凭据(仅限于不在 Databricks 之外)

如果您在 Databricks 内运行 LangChain 应用程序,可以跳过此步骤。 否则,您需要手动将 Databricks 工作空间主机名和个人访问令牌设置为 DATABRICKS_HOSTDATABRICKS_TOKEN 环境变量。有关获取访问令牌的详细信息,请参阅 身份验证文档
import getpass
import os

os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
if "DATABRICKS_TOKEN" not in os.environ:
    os.environ["DATABRICKS_TOKEN"] = getpass.getpass(
        "输入您的 Databricks 访问令牌: "
    )
输入您的 Databricks 访问令牌:  ········

安装

Databricks 集成在 databricks-langchain 包中。
pip install -qU databricks-langchain
我们首先演示如何使用 ChatDatabricks 查询托管为基础模型端点的 DBRX-instruct 模型。 对于其他类型的端点,设置端点的方式有所不同。但是,一旦端点就绪,使用 ChatDatabricks 查询它就没有区别了。请参阅此笔记本底部以获取其他类型端点的示例。

实例化

from databricks_langchain import ChatDatabricks

chat_model = ChatDatabricks(
    endpoint="databricks-dbrx-instruct",
    temperature=0.1,
    max_tokens=256,
)

调用

chat_model.invoke("什么是 MLflow?")
AIMessage(content='MLflow 是一个开源平台,用于管理机器学习工作流的端到端流程。它于 2018 年由 Databricks 引入。MLflow 提供了跟踪实验、打包和分享代码以及部署模型的工具。它可以与任何机器学习库一起使用,并可以在本地机器、虚拟机或基于云的集群等多种环境中运行。MLflow 的目标是简化机器学习开发生命周期,使数据科学家和工程师能够更轻松地协作并将其模型部署到生产环境。', response_metadata={'prompt_tokens': 229, 'completion_tokens': 104, 'total_tokens': 333}, id='run-d3fb4d06-3e10-4471-83c9-c282cc62b74d-0')
# 您也可以传递消息列表
messages = [
    ("system", "您是一个可以回答有关 Databricks 的问题的聊天机器人。"),
    ("user", "什么是 Databricks Model Serving?"),
]
chat_model.invoke(messages)
AIMessage(content='Databricks Model Serving 是 Databricks 平台的一个功能,允许数据科学家和工程师轻松将机器学习模型部署到生产环境中。通过 Model Serving,您可以托管、管理和作为 API 提供机器学习模型,使其易于集成到应用程序和业务流程中。它支持多种流行的机器学习框架,包括 TensorFlow、PyTorch 和 scikit-learn,并提供了监控和管理部署模型性能的工具。Model Serving 设计为可扩展、安全且易于使用,是希望快速高效地将机器学习模型部署到生产环境中的组织的理想选择。', response_metadata={'prompt_tokens': 35, 'completion_tokens': 130, 'total_tokens': 165}, id='run-b3feea21-223e-4105-8627-41d647d5ccab-0')

链接

类似于其他聊天模型,ChatDatabricks 可以作为复杂链的一部分使用。
from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "您是一个可以回答有关 {topic} 的问题的聊天机器人。",
        ),
        ("user", "{question}"),
    ]
)

chain = prompt | chat_model
chain.invoke(
    {
        "topic": "Databricks",
        "question": "什么是 Unity Catalog?"
    }
)
AIMessage(content="Unity Catalog 是 Databricks 中的一个新数据目录功能,允许您在您的数据景观中(包括数据湖、数据仓库和数据集市)发现、管理和治理所有数据资产。它提供了一个集中式存储库来存储和管理元数据、数据血缘关系以及对所有数据资产的访问控制。Unity Catalog 使数据团队能够轻松地发现并访问所需的任何数据,同时确保遵守数据隐私和安全法规。它设计为与 Databricks 的 Lakehouse 平台无缝集成,提供了一个统一的数据管理和分析体验。", response_metadata={'prompt_tokens': 32, 'completion_tokens': 118, 'total_tokens': 150}, id='run-82d72624-f8df-4c0d-a976-919feec09a55-0')

流式调用

for chunk in chat_model.stream("你好吗?"):
    print(chunk.content, end="|")
I|'m| 一个| AI| 而不| 感受| 情感|,| 但| 我| 在| 这里| 并| 准备好| 协助| 你|。| 今天| 我能| 帮忙| 什么|?||

异步调用

import asyncio

country = ["日本", "意大利", "澳大利亚"]
futures = [chat_model.ainvoke(f"哪里是 {c} 的首都?") for c in country]
await asyncio.gather(*futures)

工具调用

ChatDatabricks 支持与 OpenAI 兼容的工具调用 API,允许您描述工具及其参数,并让模型返回一个包含要调用的工具和该工具输入的 JSON 对象。工具调用对于构建使用工具的链和代理以及从模型获取结构化输出非常有用。 通过 ChatDatabricks.bind_tools,我们可以轻松地将 Pydantic 类、字典模式、LangChain 工具或甚至是函数作为工具传递给模型。在幕后,这些会转换为 OpenAI 兼容的工具模式,如下所示:
{
    "name": "...",
    "description": "...",
    "parameters": {...}  # JSONSchema
}
并在每次调用模型时传入。
from pydantic import BaseModel, Field


class GetWeather(BaseModel):
    """获取给定位置的当前天气"""

    location: str = Field(description="城市和州,例如:旧金山,CA")


class GetPopulation(BaseModel):
    """获取给定位置的当前人口"""

    location: str = Field(description="城市和州,例如:旧金山,CA")


llm_with_tools = chat_model.bind_tools([GetWeather, GetPopulation])
ai_msg = llm_with_tools.invoke(
    "洛杉矶比纽约哪个更热?哪个更大?"
)
print(ai_msg.tool_calls)

包装自定义模型端点

先决条件: 一旦端点就绪,其使用模式与基础模型相同。
chat_model_custom = ChatDatabricks(
    endpoint="YOUR_ENDPOINT_NAME",
    temperature=0.1,
    max_tokens=256,
)

chat_model_custom.invoke("你好吗?")

包装外部模型

先决条件:创建代理端点 首先,创建一个新的 Databricks 服务端点以将请求代理到目标外部模型。对于代理外部模型来说,端点的创建通常很快。 这需要在 Databricks 密钥管理器中注册您的 OpenAI API 密钥如下:
# 将 `<scope>` 替换为您的范围
databricks secrets create-scope <scope>
databricks secrets put-secret <scope> openai-api-key --string-value $OPENAI_API_KEY
有关如何设置 Databricks CLI 和管理密钥的详细信息,请参阅 docs.databricks.com/en/security/secrets/secrets.html
from mlflow.deployments import get_deploy_client

client = get_deploy_client("databricks")

secret = "secrets/<scope>/openai-api-key"  # 将 `<scope>` 替换为您的范围
endpoint_name = "my-chat"  # 如果 my-chat 已存在,请更改此名称
client.create_endpoint(
    name=endpoint_name,
    config={
        "served_entities": [
            {
                "name": "my-chat",
                "external_model": {
                    "name": "gpt-3.5-turbo",
                    "provider": "openai",
                    "task": "llm/v1/chat",
                    "openai_config": {
                        "openai_api_key": "{{" + secret + "}}",
                    },
                },
            }
        ],
    },
)
一旦端点状态变为“就绪”,您就可以像其他类型端点一样查询该端点。
chat_model_external = ChatDatabricks(
    endpoint=endpoint_name,
    temperature=0.1,
    max_tokens=256,
)
chat_model_external.invoke("如何使用 Databricks?")

在 Databricks 上调用函数

Databricks 函数调用是 OpenAI 兼容的,仅在模型服务作为基础模型 API 的一部分可用。 有关支持的模型,请参阅 Databricks 函数调用介绍
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-70b-instruct")
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "获取给定位置的当前天气",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "城市和州,例如:旧金山,CA",
                    },
                    "unit": {"type": "string", "enum": ["摄氏度", "华氏度"]},
                },
            },
        },
    }
]

# 支持的 tool_choice 值:"auto"、"required"、"none"、函数名称字符串格式或字典,如 {"type": "function", "function": {"name": <<tool_name>>}}
model = llm.bind_tools(tools, tool_choice="auto")

messages = [{"role": "user", "content": "芝加哥当前的温度是多少?"}]
print(model.invoke(messages))
有关如何在链中使用 UC 函数,请参阅 Databricks Unity Catalog

API 参考

要详细了解 ChatDatabricks 的所有功能和配置,请参阅 API 参考:api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_langchain.html#databricks_langchain.ChatDatabricks