How to invoke Databricks AI Assistant from a notebook cell?

[shashankB]

Hello Community,

I am exploring the Databricks AI Assistant and wondering if there is a way to invoke or interact with it directly from a notebook cell instead of using the workspace sidebar UI.

Is there any built-in command (like %assistant) to open or trigger the Assistant inside notebooks?

If not, what is the recommended way to mimic the Assistant’s functionality programmatically? For example, can we send prompts and get responses inside notebooks?

I often upload notebooks from my local system into Databricks, and I would like the AI Assistant to review or analyze the uploaded code before I actually run it.

Any official documentation or examples would also be very helpful.

Thanks in advance!

[szymon_dybczak · Accepted solution]

Hi @shashankB ,

Here's a good documentation of all possible interaction with an assistant:

What is Databricks Assistant? | Databricks on AWS

Basically, is supports various modes:

• Chat mode: Get answers to your questions by chatting with the Assistant. It responds with relevant information, including citations from Databricks documentation.
• Edit mode: Allow the Assistant to make suggestions across multiple cells in your notebook from a single prompt.
• Agent mode: Use the Data Science Agent (Beta) to automate entire multi-step data science workflows in notebooks and the SQL editor.

And you can use shortcut commands to interact with agent:

Get coding help from Databricks Assistant | Databricks on AWS

So, to answer your question in title. Here what you can do:

[nayan_wylde · Accepted solution]

@shashankB There are no command like %assistant exists today to interact with Databricks Assistant. As @szymon_dybczak mentioned in the reply  the exiting modes that you can interact with Assistant today.

Also there is no published Assistant‑specific REST API in the Databricks today. But you can built your own custom assistant in notebooks using Model Serving / Foundation Models endpoints.

Here are few examples.

1. Using Model Serving:

# Databricks notebook (Python)
from databricks.sdk import WorkspaceClient

# Get OpenAI-compatible client for your workspace serving endpoints
w = WorkspaceClient()
client = w.serving_endpoints.get_open_ai_client()

code_to_review = """[Put your code here]"""

prompt = f"""You are a senior Databricks reviewer.
Identify bugs, performance issues, and style problems in this code.
Return a short bullet list with actionable fixes, and a corrected snippet.
Code:{code_to_review}"""

system_msg = (
        "You are a senior Databricks code reviewer. "
        "Focus on correctness, Spark/Delta best practices, performance, security/secrets handling, "
        "cluster/UC implications, and style. "
        "Return Markdown with three sections: "
        "1) Findings (bullets), 2) Proposed fixes (bullets), 3) Revised code (single fenced block)."
    )


resp = client.chat.completions.create(
        model="databricks-meta-llama-3-3-70b-instruct",
        messages=[
            {"role": "system", "content": system_msg},
            {"role": "user", "content": prompt},
        ],
        temperature=0.1,
        max_tokens=900,
    )
print(resp.choices[0].message.content)

2. Use SQL AI functions:

%sql

SELECT ai_query(
  "databricks-meta-llama-3-3-70b-instruct",
  "Review this PySpark code for correctness and performance. " ||
  "Suggest changes if needed:\n\n" ||
  :code_text  -- pass as SQL variable that holds the code text
) AS review;