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:

databricks_openai.UCFunctionToolkit
databricks_openai.DatabricksFunctionClient
databricks_openai.set_uc_function_client()

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
)

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, 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, ChatCompletionToolParam] [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) → ChatCompletionToolParam | 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[ChatCompletionToolParam]

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:

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.
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)
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