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 按钮以获取 CurlPython 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.