Introduction
Just a few days ago on May 23Databricks dropped a huge update: you can now run Claude Opus 4 and Sonnet 4 right inside your Databricks workspace. In plain terms, that means you can build AI agents that tap directly into your own data lakes, run on secure infrastructure, and follow governance rules you set up. If you’ve ever wanted to give your analytics team a chatty helper that truly understands your data, this is the moment.

In this guide, you’ll learn how to spin up a mini “insights assistant” on Databricks. By the end, you’ll have a Python-powered agent to answer questions like “What were the top five pickup zones by average fare in January?” or “Show me the daily trip count for February.” We’ll walk through every step, from setting up your workspace to deploying a REST API that others can call. Let’s dive in.
1. Environment Setup
First things first: let’s get a workspace ready.
- Sign up for Databricks (if you don’t already have an account). You can try the free trial here.
- Create a cluster using the latest Databricks Runtime. A small Standard node works fine for our demo.
- Install the Databricks Python SDK. In a notebook cell or your local machine, run:
pip install databricks-sdk
- Configure your credentials by setting environment variables. On your local shell:
export DATABRICKS_HOST=https://<your-workspace>.databricks.com
export DATABRICKS_TOKEN=<your-personal-access-token>
Replace <your-workspace>
and <your-personal-access-token>
with your values. With that, the SDK can talk to your workspace without you typing in passwords each time.
2. Ingesting Sample Data
We need some data for our agent to explore. Let’s pick the popular NYC taxi dataset.
We need some data for our agent to explore. Let’s pick the popular NYC taxi dataset.
- Upload the CSV to DBFS (Databricks File System). In a notebook cell:
import os
from databricks import sql
# Path on DBFS
dbfs_path = "/mnt/datasets/nyc_taxi/yellow_tripdata_2023-01.csv"
You can also use the UI’s “Data” tab to upload files, then note the DBFS path.
2. Load into a Delta Lake table. In a notebook cell:
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
# Read the CSV
df = spark.read.option("header", "true").csv(dbfs_path)
# Convert to Delta Lake
df.write.format("delta").mode("overwrite").saveAsTable("nyc_taxi_jan2023")
Now you have a table called nyc_taxi_jan2023
, ready for questions.
3. Initializing Claude Opus 4 in Python
Next, let’s hook into Claude Opus 4 using the Databricks AI client.
- Install the AI package (if not already):
pip install databricks-ai
2. Set up your Python code:
from databricks_ai import Client
# Create a client instance
ai_client = Client()
# Instantiate the Claude Opus 4 model
model = ai_client.get_model("claude-opus-4-preview")
3. Configure parameters for generation:
params = {
"temperature": 0.2, # low randomness
"max_tokens": 512, # max length of reply
}
With that, you’re ready to chat with Claude over your data.
4. Crafting Prompts for Data Queries
To get helpful answers, we need to give Claude the right context.
- Table schema prompt:
# Fetch schema from Spark
fields = [f"{f.name}: {f.dataType.simpleString()}" for f in spark.table("nyc_taxi_jan2023").schema]
schema_str = "\n".join(fields)
base_prompt = f"""
You are a data assistant. The table 'nyc_taxi_jan2023' has these columns:
{schema_str}
When asked a question, write only JSON with keys 'query' and 'explanation'. 'query' is a SQL query to run on this table, 'explanation' is a short human-friendly note.
"""
2. Example query prompt:
user_question = "Show me the top 5 pickup_zone by average fare in January."
prompt = base_prompt + "\nQuestion: " + user_question
3. Send to Claude:
response = model.generate(prompt, **params)
result = response.choices[0].text.strip()
print(result)
You should get something like:
{
"query": "SELECT pickup_zone, AVG(fare_amount) AS avg_fare FROM nyc_taxi_jan2023 GROUP BY pickup_zone ORDER BY avg_fare DESC LIMIT 5;",
"explanation": "This query finds the top 5 zones with the highest average fare in January."
}
5. Building an “Insights Agent” Function
Let’s wrap this logic so we can reuse it.
import json
def ask_insights(question, memory=None):
"""
Send a question to Claude and return SQL + explanation.
Simple memory caches the last query.
"""
# Add memory to prompt if present
mem_str = ""
if memory:
mem_str = f"\nPrevious question: {memory['question']}\nPrevious answer: {memory['answer']}"
full_prompt = base_prompt + mem_str + "\nQuestion: " + question
resp = model.generate(full_prompt, **params)
text = resp.choices[0].text.strip()
data = json.loads(text)
# Update memory
memory = {"question": question, "answer": data}
return data, memory
# Demo
memory = None
for q in [
"How many trips were recorded each day?",
"What was the average tip amount last Monday?"
]:
answer, memory = ask_insights(q, memory)
print("SQL:", answer["query"])
print("Note:", answer["explanation"])
This function keeps track of the last question and answer, so Claude can stay “in the loop.”
6. Serving the Agent as a REST Endpoint
You might want to let other apps call your agent. FastAPI makes this easy.
- Install FastAPI and Uvicorn:
pip install fastapi uvicorn
2. Write the API (app.py
):
from fastapi import FastAPI
from pydantic import BaseModel
app = FastAPI()
class Query(BaseModel):
question: str
memory = None
@app.post("/query")
async def query_agent(q: Query):
global memory
data, memory = ask_insights(q.question, memory)
return data
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
3. Run the API:
uvicorn app:app --reload
Now, POST JSON like {"question": "How many trips on Jan 15?"}
to http://localhost:8000/query
and you’ll get back the SQL and explanation.

To deploy on Databricks, you can use:
- Databricks Repos: push this code into a repo, attach to a cluster, and run as a job.
- Databricks Jobs API: schedule the FastAPI service as a long-running job.
7. Governance and Monitoring
Databricks makes it simple to keep an eye on usage.
- Model versions: In the UI, under Model Registry, you’ll see each version of Claude Opus 4 you’ve used.
- Usage logs: With the Python SDK, you can pull recent prompts:
from databricks_ai import UsageLog
logs = ai_client.get_usage_logs(
start_time="2025-05-23T00:00:00Z",
end_time="2025-05-24T00:00:00Z"
)
for entry in logs:
print(entry.prompt, entry.response_time_ms)

This helps you spot strange patterns, track cost, and ensure your agent is behaving as expected.
Wrap-Up and Next Steps
Congratulations you now have a live, data-aware AI agent running on Databricks with Claude Opus 4! Here are some ideas to take it further:
- Vector search: Add embeddings of past queries so your agent can recall older conversations.
- Dashboards: Plug the agent into a Streamlit or Dash app for a slick GUI.
- BI integration: Expose the
/query
endpoint to Power BI or Tableau via a web connector.
For more details, check out the official docs:
- Claude Opus 4 on Databricks AI
- Sonnet 4 capabilities
- Databricks Python SDK reference
I hope this tutorial gave you a clear, hands-on path to building and deploying your own AI assistants. Now go forth and let your data speak your team will thank you!