Databricks AI Bridge Python API

class databricks_ai_bridge.genie.GenieResponse(result: str | pandas.core.frame.DataFrame, query: str | None = '', description: str | None = '')

Bases: object

result: str | DataFrame

query: str | None = ''

description: str | None = ''

class databricks_ai_bridge.genie.Genie(space_id, client: WorkspaceClient | None = None)

Bases: object

start_conversation(content)

create_message(conversation_id, content)

poll_for_result(conversation_id, message_id)

ask_question(question)

databricks_ai_bridge.vector_search_retriever_tool.vector_search_retriever_tool_trace(func)

Decorator factory to trace VectorSearchRetrieverTool with the tool name

class databricks_ai_bridge.vector_search_retriever_tool.FilterItem

Bases: BaseModel

param key: str [Required]

The filter key, which includes the column name and can include operators like ‘NOT’, ‘<’, ‘>=’, ‘LIKE’, ‘OR’

param value: Any [Required]

The filter value, which can be a single value or an array of values

class databricks_ai_bridge.vector_search_retriever_tool.VectorSearchRetrieverToolInput

Bases: BaseModel

param filters: List[FilterItem] | None = None

Optional filters to refine vector search results as an array of key-value pairs. Supports the following operators:

• Inclusion: [{“key”: “column”, “value”: value}] or [{“key”: “column”, “value”: [value1, value2]}] (matches if the column equals any of the provided values)

• Exclusion: [{“key”: “column NOT”, “value”: value}]

• Comparisons: [{“key”: “column <”, “value”: value}], [{“key”: “column >=”, “value”: value}], etc.

• Pattern match: [{“key”: “column LIKE”, “value”: “word”}] (matches full tokens separated by whitespace)

• OR logic: [{“key”: “column1 OR column2”, “value”: [value1, value2]}] (matches if column1 equals value1 or column2 equals value2; matches are position-specific)

param query: str [Required]

The string used to query the index with and identify the most similar vectors and return the associated documents.

class databricks_ai_bridge.vector_search_retriever_tool.VectorSearchRetrieverToolMixin

Bases: BaseModel

Mixin class for Databricks Vector Search retrieval tools. This class provides the common structure and interface that framework-specific implementations should follow.

param columns: List[str] | None = None

Columns to return when doing the search.

param doc_uri: str | None = None

The URI for the document, used for rendering a link in the UI.

param filters: Dict[str, Any] | None = None

Filters to apply to the search.

param include_score: bool | None = False

When true, will return the similarity score with the metadata.

param index_name: str [Required]

The name of the index to use, format: ‘catalog.schema.index’.

param num_results: int = 5

The number of results to return.

param primary_key: str | None = None

Identifies the chunk that the document is a part of. This is used by some evaluation metrics.

param query_type: str = 'ANN'

The type of this query. Supported values are ‘ANN’ and ‘HYBRID’.

param resources: List[dict] | None = None

Resources required to log a model that uses this tool.

param tool_description: str | None = None

A description of the tool.

param tool_name: str | None = None

The name of the retrieval tool.

param workspace_client: WorkspaceClient | None = None

When specified, will use workspace client credential strategy to instantiate VectorSearchClient

classmethod validate_tool_name(tool_name)