Databricks OpenAI Integrations Python API

Setup:

Install databricks-openai.

pip install -U databricks-openai

If you are outside Databricks, set the Databricks workspace hostname and personal access token to environment variables:

export DATABRICKS_HOSTNAME="https://your-databricks-workspace"
export DATABRICKS_TOKEN="your-personal-access-token"

class databricks_openai.VectorSearchRetrieverTool

Bases: VectorSearchRetrieverToolMixin

A utility class to create a vector search-based retrieval tool for querying indexed embeddings. This class integrates with Databricks Vector Search and provides a convenient interface for tool calling using the OpenAI SDK.

Example

Step 1: Call model with VectorSearchRetrieverTool defined

dbvs_tool = VectorSearchRetrieverTool(index_name="catalog.schema.my_index_name")
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {
        "role": "user",
        "content": "Using the Databricks documentation, answer what is Spark?",
    },
]
first_response = client.chat.completions.create(
    model="gpt-4o", messages=messages, tools=[dbvs_tool.tool]
)

Step 2: Execute function code – parse the model’s response and handle function calls.

tool_call = first_response.choices[0].message.tool_calls[0]
args = json.loads(tool_call.function.arguments)
result = dbvs_tool.execute(
    query=args["query"]
)  # For self-managed embeddings, optionally pass in openai_client=client

Step 3: Supply model with results – so it can incorporate them into its final response.

messages.append(first_response.choices[0].message)
messages.append(
    {"role": "tool", "tool_call_id": tool_call.id, "content": json.dumps(result)}
)
second_response = client.chat.completions.create(
    model="gpt-4o", messages=messages, tools=tools
)

param embedding_model_name: str | None = None

The name of the embedding model to use for embedding the query text.Required for direct-access index or delta-sync index with self-managed embeddings.

param text_column: str | None = None

The name of the text column to use for the embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.

param tool: ChatCompletionToolParam = None

The tool input used in the OpenAI chat completion SDK

execute(query: str, openai_client: OpenAI = None) → List[Dict]

Execute the VectorSearchIndex tool calls from the ChatCompletions response that correspond to the self.tool VectorSearchRetrieverToolInput and attach the retrieved documents into tool call messages.

Parameters:

• query – The query text to use for the retrieval.

• openai_client – The OpenAI client object used to generate embeddings for retrieval queries. If not provided, the default OpenAI client in the current environment will be used.

Returns:

A list of documents