Databricks 集成

使用 LangChain Python 集成 Databricks 大语言模型 (LLM)。

Databricks Lakehouse Platform unifies data, analytics, and AI on one platform.

本指南提供了开始使用 Databricks LLM models. 有关所有 功能和配置的详细文档，请前往 API reference.

概述

Databricks LLM class wraps a completion endpoint hosted as either of these two endpoint types:

• Databricks Model Serving, recommended for production and development,
• Cluster driver proxy app, recommended for interactive development.

This example notebook shows how to wrap your LLM endpoint and use it as an LLM in your LangChain application.

Limitations

The Databricks LLM class is legacy implementation and has several limitations in the feature compatibility.

• Only supports synchronous invocation. Streaming or async APIs are not supported.
• batch API is not supported.

To use those features, please use the new ChatDatabricks class instead. ChatDatabricks supports all APIs of ChatModel including streaming, async, batch, etc.

设置

要访问 Databricks 模型，您需要创建一个 Databricks account, set up credentials (only if you are outside Databricks workspace), and install required packages.

Credentials (only if you are outside databricks)

If you are running LangChain app inside Databricks, you can skip this step.

Otherwise, you need manually set the Databricks workspace hostname and personal access token to DATABRICKS_HOST and DATABRICKS_TOKEN environment variables, respectively. See Authentication Documentation for how to get an access 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(
        "Enter your Databricks access token: "
    )

Alternatively, you can pass those parameters when initializing the Databricks class.

from langchain_community.llms import Databricks

databricks = Databricks(
    host="https://your-workspace.cloud.databricks.com",
    # We strongly recommend NOT to hardcode your access token in your code, instead use secret management tools
    # or environment variables to store your access token securely. The following example uses Databricks Secrets
    # to retrieve the access token that is available within the Databricks notebook.
    token=dbutils.secrets.get(scope="YOUR_SECRET_SCOPE", key="databricks-token"),
)

安装

LangChain 的 Databricks 集成位于 langchain-community package. Also, mlflow >= 2.9 is required to run the code in this notebook.

pip install -qU langchain-community mlflow>=2.9.0

Wrapping model serving endpoint

Prerequisites

• An LLM was registered and deployed to a Databricks serving endpoint.
• You have “Can Query” permission to the endpoint.

The expected MLflow model signature is:

• inputs: [{"name": "prompt", "type": "string"}, {"name": "stop", "type": "list[string]"}]
• outputs: [{"type": "string"}]

调用

from langchain_community.llms import Databricks

llm = Databricks(endpoint_name="YOUR_ENDPOINT_NAME")
llm.invoke("How are you?")

'I am happy to hear that you are in good health and as always, you are appreciated.'

llm.invoke("How are you?", stop=["."])

'Good'

Transform input and output

Sometimes you may want to wrap a serving endpoint that has incompatible model signature or you want to insert extra configs. You can use the transform_input_fn and transform_output_fn arguments to define additional pre/post process.

# Use `transform_input_fn` and `transform_output_fn` if the serving endpoint
# expects a different input schema and does not return a JSON string,
# respectively, or you want to apply a prompt template on top.


def transform_input(**request):
    full_prompt = f"""{request["prompt"]}
    Be Concise.
    """
    request["prompt"] = full_prompt
    return request


def transform_output(response):
    return response.upper()


llm = Databricks(
    endpoint_name="YOUR_ENDPOINT_NAME",
    transform_input_fn=transform_input,
    transform_output_fn=transform_output,
)

llm.invoke("How are you?")

'I AM DOING GREAT THANK YOU.'

---

Connect these docs to Claude, VSCode, and more via MCP for real-time answers.

Edit this page on GitHub or file an issue.