Skip to main content

Databricks

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

This notebook provides a quick overview for getting started with Databricks LLM models. For detailed documentation of all features and configurations head to the API reference.

Overviewโ€‹

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.

Setupโ€‹

To access Databricks models you'll need to create a 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"), # noqa: F821
)
API Reference:Databricks

Installationโ€‹

The LangChain Databricks integration lives in the 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:โ€‹

The expected MLflow model signature is:

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

Invocationโ€‹

from langchain_community.llms import Databricks

llm = Databricks(endpoint_name="YOUR_ENDPOINT_NAME")
llm.invoke("How are you?")
API Reference:Databricks
'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 imcompatible 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.'

Was this page helpful?