topic Re: DatabricksVectorSearch seems to crash when served in Generative AI

DatabricksVectorSearch seems to crash when served

brahaman — Mon, 16 Jun 2025 20:54:15 GMT

Hey y'all !

So I'm experimenting with the Databricks' DatanircksVectorSearch class in Python to serve as a tool that can be used by an agent. When I run it on a notebook, I get the following error:
"[NOTICE] Using a notebook authentication token. Recommended for development only. For improved performance, please use Service Principal based authentication. To disable this message, pass disable_notice=True."
But I do correctly get an output in the notebook. But when I serve it, and test it via the Playground, the agent tells me that it has encountered an issue about authentification and I think it has to do with this notice but I'm not really sure and I don't know how to fix it...

Here's my code:

# ----------------------------- # Initialize the Vector Search tool vs_tool = DatabricksVectorSearch( index_name="dbw_genai_sbx_databricks.unstructured.pdf_docs_chunked_index", columns=["chunk_id", "doc_uri", "content_chunked"], ) # ----------------------------- # Function to retrieve information def retrieve_infs_for_databricks(query: str) -> list[dict[str, Any]]: docs = vs_tool.similarity_search(query=query) res = [] for doc in docs: res.append({ "content": doc.page_content, "source": doc.metadata["doc_uri"], "chunk_id": doc.metadata["chunk_id"], }) print(f"Res:\n{res}") return res # ----------------------------- # LangChain tool definition tool = Tool( name="databricks_docs_retriever", func=retrieve_infs_for_databricks, description=""" Searches for infos for Databricks products from the Databricks documentation. It accepts a string query which is the main keywords to search for. And it returns a list of dictionaries with the following keys: - content: contains the retrieved information - source: contains the URL of the document where the information was retrieved - chunk_id: contains the ID of the chunk where the information was When it returns a source, you should also cite the source in your final result. """ ) print("----------------------------------------------------------------------") # ----------------------------- # Initialize LangChain LLM with tools llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct") agent = create_react_agent( llm, tools=[tool], ) # ----------------------------- # Example usage agent.invoke({ "messages": [ {"role": "user", "content": "Based on the Databricks documentation, What is the best way to develop an AI agent? And give me the sources"} ] })

Re: DatabricksVectorSearch seems to crash when served

Louis_Frolio — Mon, 20 Oct 2025 19:00:25 GMT

Hello @brahaman ,

The notice you're seeing is exactly what's causing the authentication issue when your agent is deployed. When running in a notebook, the `DatabricksVectorSearch` class automatically uses your notebook authentication token (which works fine for development), but this authentication method fails when the agent is served because the deployed endpoint doesn't have access to your personal credentials.

The Solution: Use Automatic Authentication Passthrough

The recommended approach for production deployment is to use Databricks' automatic authentication passthrough feature. This eliminates authentication issues by automatically managing credentials for your agent.

Here's how to fix your code:

1. Update your agent logging to declare the Vector Search index as a resource:

```python
import mlflow
from mlflow.models.resources import DatabricksVectorSearchIndex, DatabricksServingEndpoint

# When you log your model, add the resources parameter
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
python_model="your_agent.py", # Your agent file
artifact_path="agent",
input_example=input_example,
example_no_conversion=True,
resources=[
DatabricksVectorSearchIndex(
index_name="dbw_genai_sbx_databricks.unstructured.pdf_docs_chunked_index"
),
DatabricksServingEndpoint(
endpoint_name="databricks-meta-llama-3-3-70b-instruct"
)
]
)
```

2. You don't need to modify your tool initialization code.

When the agent is deployed with resources declared, Databricks automatically provisions a service principal with the necessary permissions and handles all authentication.

How It Works

When you declare resources during logging, Databricks:
1. Verifies you (the deployer) have access to those resources
2. Automatically creates a service principal for your agent with least-privilege access
3. Manages short-lived OAuth tokens that are automatically rotated

This approach provides better security than manual authentication and eliminates the performance overhead mentioned in the notice (up to 100ms per query).

Requirements

- Ensure you're using MLflow 2.13.1 or above for Vector Search automatic authentication
- You must have `Can Use` permission on the Vector Search index when deploying

Hope this helps, Louis.

Re: DatabricksVectorSearch seems to crash when served

ZD — Wed, 25 Mar 2026 13:38:45 GMT

@Louis_Frolio Are there audit logs or visibility into which service principal is being used and what access it has?

Another question: With automatic authentication passthrough, do we still need to set DATABRICKS_HOST and DATABRICKS_TOKEN as environment variables in the serving endpoint configuration? or can those be omitted entirely?