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"

Re-exported Unity Catalog Utilities

This module re-exports selected utilities from the Unity Catalog open source package.

Available aliases:

Refer to the Unity Catalog documentation for more information.

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"], filters=args.get("filters", None)
)  # 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
)

Note: Any additional keyword arguments passed to the constructor will be passed along to databricks.vector_search.client.VectorSearchIndex.similarity_search when executing the tool. See documentation to see the full set of supported keyword arguments, e.g. score_threshold. Also, see documentation for VectorSearchRetrieverToolMixin for additional supported constructor arguments not listed below, including query_type and num_results.

WorkspaceClient instances with auth types PAT, OAuth-M2M (client ID and client secret), or model serving credential strategy will be used to instantiate the underlying VectorSearchClient.

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: ChatCompletionFunctionToolParam = None

The tool input used in the OpenAI chat completion SDK

execute(query: str, filters: List[FilterItem] | None = None, openai_client: OpenAI = None, **kwargs: Any) → 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

class databricks_openai.UCFunctionToolkit

Bases: BaseModel

param client: BaseFunctionClient | None = None

The client for managing functions, must be an instance of BaseFunctionClient

param filter_accessible_functions: bool = False

When set to true, UCFunctionToolkit is initialized with functions that only the client has access to

param function_names: List[str] [Optional]

The list of function names in the form of ‘catalog.schema.function’

param tools_dict: Dict[str, ChatCompletionFunctionToolParam] [Optional]

The tools dictionary storing the function name and tool definition mapping, no need to provide this field

static uc_function_to_openai_function_definition(*, function_name: str, client: BaseFunctionClient | None = None, filter_accessible_functions: bool = False) → ChatCompletionFunctionToolParam | None

Convert a UC function to OpenAI function definition.

Parameters:

  • function_name – The full name of the function in the form of ‘catalog.schema.function’

  • client – The client for managing functions, must be an instance of BaseFunctionClient

property tools: List[ChatCompletionFunctionToolParam]
class databricks_openai.DatabricksFunctionClient(client: WorkspaceClient | None = None, *, profile: str | None = None, execution_mode: str = 'serverless', **kwargs: Any)

Bases: BaseFunctionClient

Databricks UC function calling client

set_spark_session()

Initialize the spark session with serverless compute if not already active.

stop_spark_session()
initialize_spark_session()

Initialize the spark session with serverless compute. This method is called when the spark session is not active.

refresh_client_and_session()

Refreshes the databricks client and spark session if the session_id has been invalidated due to expiration of temporary credentials. If the client is running within an interactive Databricks notebook environment, the spark session is not terminated.

create_function(*, sql_function_body: str | None = None) → FunctionInfo

Create a UC function with the given sql body or function info.

Note: databricks-connect is required to use this function, make sure its version is 15.1.0 or above to use

serverless compute.

Parameters:

sql_function_body – The sql body of the function. Defaults to None. It should follow the syntax of CREATE FUNCTION statement in Databricks. Ref: https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-sql-function.html#syntax

Returns:

The created function info.

Return type:

FunctionInfo

create_python_function(*, func: Callable[[...], Any], catalog: str, schema: str, replace: bool = False, dependencies: list[str] | None = None, environment_version: str = 'None') → FunctionInfo

Create a Unity Catalog (UC) function directly from a Python function.

This API allows you to convert a Python function into a Unity Catalog User-Defined Function (UDF). It automates the creation of UC functions while ensuring that the Python function meets certain criteria and adheres to best practices.

Requirements:

  1. Type Annotations:

    • The Python function must use argument and return type annotations. These annotations are used

    to generate the SQL signature of the UC function. - Supported Python types and their corresponding UC types are as follows:

    Python Type | Unity Catalog Type |

    |----------------------|————————–| | int | LONG | | float | DOUBLE | | str | STRING | | bool | Boolean Operations — and, or, not | | Decimal | DECIMAL | | datetime.date | DATE | | datetime.timedelta | INTERVAL DAY TO SECOND | | datetime.datetime | TIMESTAMP | | list | ARRAY | | tuple | ARRAY | | dict | MAP | | bytes | Binary arithmetic operations |

    • Example of a valid function:

    ```python def my_function(a: int, b: str) -> float:

    return a + len(b)

    ```

    • Invalid function (missing type annotations):

    ```python def my_function(a, b):

    return a + len(b)

    ``` Attempting to create a UC function from a function without type hints will raise an error, as the system relies on type hints to generate the UC function’s signature.

    • For container types like list, tuple and dict, the inner types must be specified and must be

    uniform (Union types are not permitted). For example:

    ```python def my_function(a: List[int], b: Dict[str, float]) -> List[str]:

    return [str(x) for x in a]

    ```

    • var args and kwargs are not supported. All arguments must be explicitly defined in the function signature.

  2. Google Docstring Guidelines:

    • It is required to include detailed Python docstrings in your function to provide additional context.

    The docstrings will be used to auto-generate parameter descriptions and a function-level comment.

    • A function description must be provided at the beginning of the docstring (within the triple quotes)

    to describe the function’s purpose. This description will be used as the function-level comment in the UC function. The description must be included in the first portion of the docstring prior to any argument descriptions.

    • Parameter descriptions are optional but recommended. If provided, they should be included in the

    Google-style docstring. The parameter descriptions will be used to auto-generate detailed descriptions for each parameter in the UC function. The additional context provided by these argument descriptions can be useful for agent applications to understand context of the arguments and their purpose.

    • Only Google-style docstrings are supported for this auto-generation. For example:

    ```python def my_function(a: int, b: str) -> float:

    “”” Adds the length of a string to an integer.

    Args:

    a (int): The integer to add to. b (str): The string whose length will be added.

    Returns:

    float: The sum of the integer and the string length.

    “”” return a + len(b)

    ``` - If docstrings do not conform to Google-style for specifying arguments descriptions, parameter descriptions

    will default to "Parameter <name>", and no further information will be provided in the function comment for the given parameter.

    For examples of Google docstring guidelines, see [this link](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html)

  3. External Dependencies:

    • Unity Catalog UDFs are limited to Python standard libraries and Databricks-provided libraries. If your

    function relies on unsupported external dependencies, the created UC function may fail at runtime. - It is strongly recommended to test the created function by executing it before integrating it into GenAI or other tools.

Function Metadata: - Docstrings (if provided and Google-style) will automatically be included as detailed descriptions for function parameters as well as for the function itself, enhancing the discoverability of the utility of your UC function.

Example: ```python def example_function(x: int, y: int) -> float:

“”” Multiplies an integer by the length of a string.

Args:

x (int): The number to be multiplied. y (int): A string whose length will be used for multiplication.

Returns:

float: The product of the integer and the string length.

“”” return x * len(y)

client.create_python_function(

func=example_function, catalog=”my_catalog”, schema=”my_schema”

)

Overwriting a function: - If a function with the same name already exists in the specified catalog and schema, the function will not be created by default. To overwrite the existing function, set the replace parameter to True.

param func:

The Python function to convert into a UDF.

param catalog:

The catalog name in which to create the function.

param schema:

The schema name in which to create the function.

param replace:

Whether to replace the function if it already exists. Defaults to False.

param dependencies:

A list of external dependencies required by the function. Defaults to an empty list. Note that the dependencies parameter is not supported in all runtimes. Ensure that you are using a runtime that supports environment and dependency declaration prior to creating a function that defines dependencies. Standard PyPI package declarations are supported (i.e., requests>=2.25.1).

param environment_version:

The version of the environment in which the function will be executed. Defaults to ‘None’. Note that the environment_version parameter is not supported in all runtimes. Ensure that you are using a runtime that supports environment and dependency declaration prior to creating a function that declares an environment verison.

returns:

Metadata about the created function, including its name and signature.

rtype:

FunctionInfo

create_wrapped_function(*, primary_func: Callable[[...], Any], functions: list[Callable[[...], Any]], catalog: str, schema: str, replace=False, dependencies: list[str] | None = None, environment_version: str = 'None') → FunctionInfo

Create a wrapped function comprised of a primary_func function and in-lined wrapped Functions within the primary_func body.

Note: databricks-connect is required to use this function, make sure its version is 15.1.0 or above to use

serverless compute.

Parameters:

  • primary_func – The primary function to be wrapped.

  • functions – A list of functions to be wrapped inline within the body of primary_func.

  • catalog – The catalog name.

  • schema – The schema name.

  • replace – Whether to replace the function if it already exists. Defaults to False.

  • dependencies – A list of external dependencies required by the function. Defaults to an empty list. Note that the dependencies parameter is not supported in all runtimes. Ensure that you are using a runtime that supports environment and dependency declaration prior to creating a function that defines dependencies. Standard PyPI package declarations are supported (i.e., requests>=2.25.1).

  • environment_version – The version of the environment in which the function will be executed. Defaults to ‘None’. Note that the environment_version parameter is not supported in all runtimes. Ensure that you are using a runtime that supports environment and dependency declaration prior to creating a function that declares an environment verison.

Returns:

Metadata about the created function, including its name and signature.

Return type:

FunctionInfo

get_function(function_name: str, **kwargs: Any) → FunctionInfo

Get a function by its name.

Parameters:

  • function_name – The name of the function to get.

  • kwargs – additional key-value pairs to include when getting the function.

  • are (Allowed keys for retrieving functions)

  • include_browse (-) – bool (default to None) Whether to include functions in the response for which the principal can only access selective metadata for.

Note

The function name shouldn’t be *, to get all functions in a catalog and schema, please use list_functions API instead.

Returns:

The function info.

Return type:

FunctionInfo

list_functions(catalog: str, schema: str, max_results: int | None = None, page_token: str | None = None, include_browse: bool | None = None) → PagedList[FunctionInfo]

List functions in a catalog and schema.

Parameters:

  • catalog – The catalog name.

  • schema – The schema name.

  • max_results – The maximum number of functions to return. Defaults to None.

  • page_token – The token for the next page. Defaults to None.

  • include_browse – Whether to include functions in the response for which the

  • None. (principal can only access selective metadata for. Defaults to)

Returns:

The paginated list of function infos.

Return type:

PageList[FunctionInfo]

execute_function(function_name: str, parameters: Dict[str, Any] | None = None, **kwargs: Any) → FunctionExecutionResult

Execute a UC function by name with the given parameters.

Parameters:

  • function_name – The name of the function to execute.

  • parameters – The parameters to pass to the function. Defaults to None.

  • kwargs –

    additional key-value pairs to include when executing the function. Allowed keys for retrieving functions are: - include_browse: bool (default to False)

    Whether to include functions in the response for which the principal can only access selective metadata for.

    Allowed keys for executing functions are: - wait_timeout: str (default to 30s)

    The time in seconds the call will wait for the statement’s result set as Ns, where N can be set to 0 or to a value between 5 and 50.

    When set to 0s, the statement will execute in asynchronous mode and the call will not wait for the execution to finish. In this case, the call returns directly with PENDING state and a statement ID which can be used for polling with :method:statementexecution/getStatement.

    When set between 5 and 50 seconds, the call will behave synchronously up to this timeout and wait for the statement execution to finish. If the execution finishes within this time, the call returns immediately with a manifest and result data (or a FAILED state in case of an execution error). If the statement takes longer to execute, on_wait_timeout determines what should happen after the timeout is reached.

    • row_limit: int (default to 100)

      Applies the given row limit to the statement’s result set, but unlike the LIMIT clause in SQL, it also sets the truncated field in the response to indicate whether the result was trimmed due to the limit or not.

    • byte_limit: int (default to 1048576 = 1MB)

      Applies the given byte limit to the statement’s result size. Byte counts are based on internal data representations and might not match the final size in the requested format. If the result was truncated due to the byte limit, then truncated in the response is set to true. When using EXTERNAL_LINKS disposition, a default byte_limit of 100 GiB is applied if byte_limit is not explcitly set.

Returns:

The result of executing the function.

Return type:

FunctionExecutionResult

delete_function(function_name: str, force: bool | None = None) → None

Delete a function by its full name.

Parameters:

  • function_name – The full name of the function to delete. It should be in the format of “catalog.schema.function_name”.

  • force – Force deletion even if the function is not empty. This parameter is used by underlying databricks workspace client when deleting a function. If it is None then the parameter is not included in the request. Defaults to None.

to_dict()

Store the client configuration in a dictionary. Sensitive information should be excluded.

classmethod from_dict(config: Dict[str, Any])
get_function_source(function_name: str) → str

Returns the Python callable definition as a string for an EXTERNAL Python function that is stored within Unity Catalog. This function can only parse and extract the full callable definition for Python functions and cannot be used on SQL or TABLE functions.

Parameters:

function_name – The name of the function to retrieve the Python callable definition for.

Returns:

The Python callable definition as a string.

Return type:

str

get_function_as_callable(function_name: str, register_function: bool = True, namespace: dict[str, Any] | None = None) → Callable[[...], Any]

Returns the Python callable for an EXTERNAL Python function that is stored within Unity Catalog. This function can only parse and extract the full callable definition for Python functions and cannot be used on SQL or TABLE functions.

Parameters:

  • function_name – The name of the function to retrieve the Python callable for.

  • register_function – Whether to register the function in the namespace. Defaults to True.

  • namespace – The namespace to register the function in. Defaults to None (global)

Returns:

The Python callable for the function.

Return type:

Callable[…, Any]

databricks_openai.set_uc_function_client(client: BaseFunctionClient) → None
class databricks_openai.DatabricksOpenAI(workspace_client: WorkspaceClient | None = None)

Bases: OpenAI

OpenAI client authenticated with Databricks to query LLMs and agents hosted on Databricks.

This client extends the standard OpenAI client with Databricks authentication, allowing you to interact with foundation models and AI agents deployed on Databricks using the familiar OpenAI SDK interface.

The client automatically handles authentication using your Databricks credentials.

Parameters:

workspace_client – Databricks WorkspaceClient to use for authentication. Pass a custom WorkspaceClient to set up your own authentication method. If not provided, a default WorkspaceClient will be created using standard Databricks authentication resolution.

Example

>>> # Use default Databricks authentication
>>> client = DatabricksOpenAI()
>>> response = client.chat.completions.create(
...     model="databricks-meta-llama-3-1-70b-instruct",
...     messages=[{"role": "user", "content": "Hello!"}],
... )
>>> # Use custom WorkspaceClient for authentication
>>> from databricks.sdk import WorkspaceClient
>>> ws = WorkspaceClient(host="https://my-workspace.cloud.databricks.com", token="...")
>>> client = DatabricksOpenAI(workspace_client=ws)
class databricks_openai.AsyncDatabricksOpenAI(workspace_client: WorkspaceClient | None = None)

Bases: AsyncOpenAI

Async OpenAI client authenticated with Databricks to query LLMs and agents hosted on Databricks.

This client extends the standard AsyncOpenAI client with Databricks authentication, allowing you to interact with foundation models and AI agents deployed on Databricks using the familiar OpenAI SDK interface with async/await support.

The client automatically handles authentication using your Databricks credentials.

Parameters:

workspace_client – Databricks WorkspaceClient to use for authentication. Pass a custom WorkspaceClient to set up your own authentication method. If not provided, a default WorkspaceClient will be created using standard Databricks authentication resolution.

Example

>>> # Use default Databricks authentication
>>> client = AsyncDatabricksOpenAI()
>>> response = await client.chat.completions.create(
...     model="databricks-meta-llama-3-1-70b-instruct",
...     messages=[{"role": "user", "content": "Hello!"}],
... )
>>> # Use custom WorkspaceClient for authentication
>>> from databricks.sdk import WorkspaceClient
>>> ws = WorkspaceClient(host="https://my-workspace.cloud.databricks.com", token="...")
>>> client = AsyncDatabricksOpenAI(workspace_client=ws)
class databricks_openai.McpServerToolkit(url: str, name: str | None = None, workspace_client: WorkspaceClient | None = None)

Bases: object

Toolkit for accessing MCP server tools with the OpenAI SDK.

This class provides a simplified interface to MCP (Model Context Protocol) servers, automatically converting MCP tools into tool specifications for the OpenAI SDK. It’s designed for easy integration with OpenAI clients and agents that use function calling.

The toolkit handles authentication with Databricks, fetches available tools from the MCP server, and provides execution functions for each tool.

Parameters:

  • url – The URL of the MCP server to connect to. (Required parameter)

  • name – A readable name for the MCP server. This name will be used as a prefix for tool names to avoid conflicts when using multiple MCP servers (e.g., “server_name__tool_name”). If not provided, tool names will not be prefixed.

  • workspace_client – Databricks WorkspaceClient to use for authentication. Pass a custom WorkspaceClient to set up your own authentication method. If not provided, a default WorkspaceClient will be created using standard Databricks authentication resolution.

Example

Step 1: Create toolkit and get tools from MCP server

from databricks_openai import McpServerToolkit
from openai import OpenAI

toolkit = McpServerToolkit(url="https://my-mcp-server.com/mcp", name="my_tools")
tools = toolkit.get_tools()
tool_specs = [tool.spec for tool in tools]

Step 2: Call model with MCP tools defined

client = OpenAI()
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Help me search for information about Databricks."},
]
first_response = client.chat.completions.create(
    model="gpt-4o", messages=messages, tools=tool_specs
)

Step 3: Execute function code – parse the model’s response and handle tool calls

import json

tool_call = first_response.choices[0].message.tool_calls[0]
args = json.loads(tool_call.function.arguments)

# Find and execute the appropriate tool
tool_to_execute = next(t for t in tools if t.name == tool_call.function.name)
result = tool_to_execute.execute(**args)

Step 4: 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": result})
second_response = client.chat.completions.create(
    model="gpt-4o", messages=messages, tools=tool_specs
)
classmethod from_uc_function(catalog: str, schema: str, function_name: str | None = None, name: str | None = None, workspace_client: WorkspaceClient | None = None) → McpServerToolkit

Create an MCP server toolkit from Unity Catalog function path.

Convenience method to create a toolkit for UC functions by specifying Unity Catalog components instead of constructing the full URL manually.

Parameters:

  • catalog – Unity Catalog catalog name.

  • schema – Schema name within the catalog.

  • function_name – Optional UC function name. If omitted, provides access to all functions in the schema.

  • name – Readable name for the MCP server. See __init__ for details.

  • workspace_client – WorkspaceClient for authentication. See __init__ for details.

Returns:

McpServerToolkit instance for the specified Unity Catalog function.

Example

from databricks_openai import McpServerToolkit
from openai import OpenAI

# Create toolkit from UC function
toolkit = McpServerToolkit.from_uc_function(
    catalog="main", schema="tools", function_name="send_email", name="email_tools"
)

# Get tools and use with OpenAI
tools = toolkit.get_tools()
tool_specs = [tool.spec for tool in tools]

Create an MCP server toolkit from Unity Catalog vector search index path.

Convenience method to create a toolkit for vector search by specifying Unity Catalog components instead of constructing the full URL manually.

Parameters:

  • catalog – Unity Catalog catalog name.

  • schema – Schema name within the catalog.

  • index_name – Optional vector search index name. If omitted, provides access to all indexes in the schema.

  • name – Readable name for the MCP server. See __init__ for details.

  • workspace_client – WorkspaceClient for authentication. See __init__ for details.

Returns:

McpServerToolkit instance for the specified Unity Catalog vector search index.

Example

from databricks_openai import McpServerToolkit
from openai import OpenAI

# Create toolkit from vector search index
toolkit = McpServerToolkit.from_vector_search(
    catalog="main", schema="embeddings", index_name="product_docs", name="search_tools"
)

# Get tools and use with OpenAI
tools = toolkit.get_tools()
tool_specs = [tool.spec for tool in tools]
get_tools() → List[ToolInfo]
async aget_tools() → List[ToolInfo]
class databricks_openai.ToolInfo

Bases: BaseModel

param execute: Callable [Required]
param name: str [Required]
param spec: ChatCompletionFunctionToolParam [Required]