Databricks Agent Bricks — UC 函数
简介
本指南演示如何直接在 Databricks Unity Catalog (UC) Functions 中使用 Neo4j 数据库。
此设置使您能够定义 MCP 工具,在远程 Neo4j 实例上执行查询,直接在 Databricks 笔记本或代理中完成,无需专用的 Neo4j MCP 服务器,且操作迅速。
按照本指南操作,您可以将 Neo4j 查询公开为 UC 函数,可在 Databricks 的 Python 中调用(作为工具),从而实现与大型语言模型 (LLM) 代理或其他工作流的集成。
示例展示了一个简单的代理实现,可返回给定公司名称的竞争对手。
架构概述
→ Databricks 代理 / Playground
→ UC 函数(Python)作为工具
→ Neo4j REST API
→ Neo4j 数据库(例如 demo.neo4jlabs.com / companies 数据集)
关键点
-
UC 函数封装查询逻辑并处理对 Neo4j 的 HTTP 调用。
-
不需要 MCP 服务器;LLM 或代理将 UC 函数视为工具进行交互。
-
Neo4j 连接使用基本身份验证或其他凭证进行安全保护。
限制
-
仅限 Python/SQL(UC 函数中不直接使用 Cypher)
-
必须使用 HTTP 将 Neo4j 驱动调用封装在函数中,服务器无状态函数无法在笔记本之外使用 Python 依赖。
-
连接凭证必须硬编码在 UC 函数中,UC 函数内部无法访问机密——这是架构限制。Unity Catalog 函数无权访问
-
dbutils(包括机密)
-
笔记本会话变量
-
Databricks SDK 认证上下文
-
实施
步骤 1 - 创建一个或多个 UC 函数
在您的 Workspace 中创建一个新的 Notebook 并定义如下单元格,这里我们定义一个 UC 函数来检索给定公司的竞争对手,该函数将作为 LLM 的工具。
%sql
/* Create the catalog and schema where the UC function will be available */
CREATE CATALOG IF NOT EXISTS knowledge_graph;
CREATE SCHEMA IF NOT EXISTS company_data;
CREATE SCHEMA IF NOT EXISTS knowledge_graph.company_data;
%sql
/* Define the UC function and implements its behavior, in this example, the function finds the competitors of a given company by performing a cypher query through HTTP */
CREATE OR REPLACE FUNCTION knowledge_graph.company_data.find_competitors(
company_name STRING,
max_results INT
)
RETURNS STRING
COMMENT 'Finds competitors using Neo4j graph traversal and returns structured competitor data'
LANGUAGE PYTHON
AS $$
import requests
import base64
import json
url = "https://demo.neo4jlabs.com:7473/db/companies/tx/commit"
# Databricks secrets (dbutils.secrets) are not available outside of the notebook, therefore, they cannot be used in a serverless UC Function
auth = base64.b64encode(b"companies:companies").decode()
headers = {
"Authorization": f"Basic {auth}",
"Content-Type": "application/json",
"Accept": "application/json"
}
cypher_query = f"""
MATCH (c:Organization {{name: '{company_name}'}})-[:HAS_COMPETITOR]->(competitor:Organization)
RETURN competitor.name as name, competitor.revenue as revenue
LIMIT {max_results}
"""
payload = {
"statements": [
{ "statement": cypher_query }
]
}
r = requests.post(url, json=payload, headers=headers)
r.raise_for_status()
return r.text
$$;
%sql
/* Test the function */
SELECT knowledge_graph.company_data.find_competitors(
'BigFix',
5
);
运行 Notebook 单元格并确保测试如预期通过后,UC 函数将出现在 Catalog 中,此时我们即可将其作为代理的工具使用。
步骤 2 - Playground 测试和部署
在 Playground 中,从 Tools → Add Tool 选择我们新建的 UC 函数,添加如下系统提示,然后开始提问第一个问题:What are the competitors of BigFix?
Purpose: Assist users in getting companies/organizations info.
Limitations:
- Focus on companies.
- Be conversational but do not answer any unrelated queries that are not related to companies.
- Handle queries for multiple companies.
- If there is no company information, do not attempt to retrieve otherwise – inform the user with an appropriate error message.
Parameters:
- Company name
- Max results
Data Sources:
- Use the find_competitors API tool when requested with questions about company's competitors.
Actions:
1. Retrieve company info
2. Retrieve competitors
Error Handling:
- Provide clear error messages if Neo4j Connection calls fail.
Sample Questions:
- "What are the competitors of 'BigFix'?"
LLM 将调用 UC 函数从 Neo4j 检索信息,并生成自然语言响应。
注意:可以同时使用多个工具,从而创建更复杂的代理。
现在我们已经测试了代理功能,准备部署它。
在 Get Code 按钮中,点击 Create Agent Notebook,然后按照步骤操作(替换所有 TODO)。
运行所有单元格后,模型将在 Serving 选项卡下可用。
最后,可以点击 Try in Playground 再次测试。
测试完成后,点击 Get Code 按钮以获取 Curl 和 Python APIs(使用 API 必须设置 Databricks 令牌)。随后即可使用提供的 API 调用代理!
Curl API 示例
curl https://dbc-...cloud.databricks.com/serving-endpoints/agents_knowledge_graph-company_data-find_competitors_model/invocations \
-X POST \
-H "Authorization: Bearer $DATABRICKS_TOKEN" \
-H "Content-Type: application/json" \
-d '{"input":[{"role":"user","content":"What is an LLM agent?"}],"databricks_options":{"conversation_id":"...","return_trace":true},"context":{"conversation_id":"...","user_id":"..."}}'