databricks_ai_bridge_chat_completions_quickstart.py

1from databricks.ai_bridge import AIBridge
2
3# Initialize the AI Bridge
4# This assumes you have configured your Databricks CLI credentials 
5# or set DATABRICKS_HOST and DATABRICKS_TOKEN environment variables.
6bridge = AIBridge()
7
8# Example: Chat with a Model Serving endpoint
9response = bridge.chat.completions.create(
10    model="databricks-llama-2-70b-chat",
11    messages=[
12        {"role": "system", "content": "You are a helpful assistant."},
13        {"role": "user", "content": "What is Databricks AI Bridge?"}
14    ]
15)
16
17print(response.choices[0].message.content)