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 连接使用基本身份验证或其他凭证进行安全保护。

优势

  • 零基础设施(由 Databricks 无服务器管理)

  • 快速原型化

  • 自动权限继承

  • 模式级别曝光(多个函数 → 多个工具)

  • 在 Playground 中可立即使用

限制

  • 仅限 Python/SQL(UC 函数中不直接使用 Cypher)

  • 必须使用 HTTP 将 Neo4j 驱动调用封装在函数中,服务器无状态函数无法在笔记本之外使用 Python 依赖。

  • 连接凭证必须硬编码在 UC 函数中,UC 函数内部无法访问机密——这是架构限制。Unity Catalog 函数无权访问

    • dbutils(包括机密)

    • 笔记本会话变量

    • Databricks SDK 认证上下文

先决条件

  • 具备计算能力的 Databricks 订阅。

  • 在个人区域的 Developer → Access Tokens 中创建 Databricks 令牌。

实施

步骤 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 检索信息,并生成自然语言响应。

Playground Results

注意:可以同时使用多个工具,从而创建更复杂的代理。

现在我们已经测试了代理功能,准备部署它。

在 Get Code 按钮中,点击 Create Agent Notebook,然后按照步骤操作(替换所有 TODO)。

运行所有单元格后,模型将在 Serving 选项卡下可用。

Serving Page

最后,可以点击 Try in Playground 再次测试。

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":"..."}}'
© . This site is unofficial and not affiliated with Neo4j, Inc.