databricks_langchain_chat_model_quickstart_with_dbrx.py

1import getpass
2import os
3
4# Install the package if you haven't already:
5# %pip install -qU langchain-databricks
6
7# Set up Databricks credentials if not running inside a Databricks workspace
8if "DATABRICKS_HOST" not in os.environ:
9    os.environ["DATABRICKS_HOST"] = input("Enter your Databricks host (e.g., https://adb-xxx.x.azuredatabricks.net): ")
10if "DATABRICKS_TOKEN" not in os.environ:
11    os.environ["DATABRICKS_TOKEN"] = getpass.getpass("Enter your Databricks personal access token: ")
12
13from langchain_databricks import ChatDatabricks
14from langchain_core.messages import HumanMessage, SystemMessage
15
16# Initialize the Chat Model
17# Replace 'databricks-dbrx-instruct' with the name of your served endpoint
18chat_model = ChatDatabricks(
19    endpoint="databricks-dbrx-instruct",
20    temperature=0.1,
21    max_tokens=256,
22)
23
24# Define messages
25messages = [
26    SystemMessage(content="You are a helpful assistant."),
27    HumanMessage(content="What is Databricks?"),
28]
29
30# Invoke the model
31response = chat_model.invoke(messages)
32
33# Print the result
34print(response.content)