使用 DataHub Agent Context 构建 Google ADK 自主数据智能体-开发者社区

使用 DataHub Agent Context 构建 Google ADK 自主数据智能体

【免费下载链接】datahubThe Context Platform for your Data and AI Stack项目地址: https://gitcode.com/GitHub_Trending/da/datahub

DataHub 的 Agent Context Kit 提供了将企业数据上下文（数据所有权、血缘、文档、质量信号等）直接注入 AI 智能体的能力。本文聚焦 Google ADK（Agent Development Kit）这一官方集成路径，讲解如何在 Google ADK 中通过 Python 工具直嵌或 MCP Server 两种方式接入 DataHub，读完即可基于当前仓库中的示例代码搭建一个能检索数据集、追踪血缘、查询文档乃至写入元数据的自主数据智能体。

集成总览：两种接入方式

在 Google ADK 中消费 DataHub 上下文，官方提供两条互补路径（见 google-adk.md）：

Python 工具直嵌：通过datahub-agent-context包的build_google_adk_tools()把 DataHub 能力包装成普通 Python 函数，直接传给 ADKAgent的tools参数，零 MCP 基础设施依赖；
MCP Server 连接：让 ADK 通过内置的McpToolset连接 DataHub 的 MCP server（DataHub Cloud 托管端点或自建端点），由 MCP 协议统一发现工具。

两种方式共享同一套底层 MCP 工具实现（位于 datahub-agent-context/src/datahub_agent_context/mcp_tools/），因此能力集一致，区别仅在于工具如何被 ADK 发现与调用。

前置条件

开始之前需要准备：

Python 3.10 及以上版本；
Google ADK：pip install google-adk；
一个可访问的 DataHub 实例，以及个人访问令牌；
一个 Google API Key（Gemini Developer API）或Google Cloud 凭据（Vertex AI）。

安装

pip install datahub-agent-context[google-adk]

[google-adk]是可选依赖组。从当前仓库的 requirements.txt 可以看到该示例锁定的版本约束为google-adk>=1.0.0,<2.0.0，并因 CVE-2025-43859 对httpcore>=1.0.9与h11>=0.16做了下限约束。本地开发可改用pip install -e "datahub-agent-context[google-adk]"方式安装。

快速开始：Python 工具直嵌

第一步：创建 DataHub 客户端

from datahub.sdk.main_client import DataHubClient client = DataHubClient.from_env()

DataHubClient.from_env()会从环境变量读取连接配置（GMS 地址与令牌）。若需显式指定，可参考 simple_search.py 的写法：

import os from datahub.sdk.main_client import DataHubClient datahub_gms_url = os.getenv("DATAHUB_GMS_URL") if datahub_gms_url is None: client = DataHubClient.from_env() else: client = DataHubClient(server=datahub_gms_url, token=os.getenv("DATAHUB_GMS_TOKEN"))

对应环境变量为DATAHUB_GMS_URL（默认http://localhost:8080）与DATAHUB_GMS_TOKEN。

第二步：构建工具集

from datahub_agent_context.google_adk_tools import build_google_adk_tools # 默认只读；设置 include_mutations=True 可启用写操作 tools = build_google_adk_tools(client, include_mutations=False)

从 builder.py 的源码可以确认默认只读工具清单：

工具	作用
`search`	按关键词搜索数据集、仪表盘等实体
`get_entities`	获取实体完整详情（schema、所有权、文档、标签）
`list_schema_fields`	列出数据集的字段（列）
`get_lineage`	追踪上游/下游血缘
`get_lineage_paths_between`	查询两个实体之间的血缘路径
`get_dataset_queries`	获取与数据集关联的 SQL 查询
`get_dataset_assertions`	获取数据集的质量断言（assertions）
`list_incidents`	列出事件（incidents）
`search_documents`/`grep_documents`	搜索知识库文章与文档
`get_me`	获取当前用户信息

当include_mutations=True时追加写操作工具：update_description（更新描述）、set_domains/remove_domains（管理域）、add_owners/remove_owners（管理所有者）、add_tags/remove_tags（管理标签）、add_glossary_terms/remove_glossary_terms（关联术语表术语）、add_structured_properties/remove_structured_properties（结构化属性）、save_document（保存文档）、raise_incident/resolve_incident（创建/解决事件）。

这些包装函数的实现原理见 utils.py：create_context_wrapper利用contextvars在执行函数前将DataHubClient注入上下文，执行后重置；同时把工具抛出的ItemNotFoundError转换为结构化字典返回，保证 LLM 拿到的是可读消息而非未处理异常。

另有一个 Cloud 专属的 build_google_adk_cloud_tools 函数，可启用Ask DataHubAI 助手工具（ask_datahub_chat/get_datahub_chat），仅适用于 DataHub Cloud 实例，OSS 实例请勿启用。

第三步：注入 ADK Agent

from google.adk.agents import Agent agent = Agent( model="gemini-2.5-flash", name="datahub_agent", description="A data discovery assistant with access to DataHub.", instruction="Use the available tools to search for datasets, get entity details, and trace lineage. Always include URNs in your answers.", tools=tools, )

instruction提示词是决定智能体是否真正调用工具的关键——建议在其中明确指定应使用的工具类别（如先搜索、再取详情、再追血缘），并约定回答格式（如"始终在回答中包含 URN"），这一经验同样体现在 basic_agent.py 的SYSTEM_PROMPT中。

完整可运行的最小示例

仓库中的 simple_search.py 是"绝对最小"可运行版本，完整展示了从建客户端、建工具、建 Agent 到用Runner流式执行一次查询的闭环：

import asyncio from google.adk.agents import Agent from google.adk.runners import Runner from google.adk.sessions import InMemorySessionService from google.genai import types from datahub.sdk.main_client import DataHubClient from datahub_agent_context.google_adk_tools import build_google_adk_tools client = DataHubClient.from_env() tools = build_google_adk_tools(client, include_mutations=False) agent = Agent( model="gemini-2.5-flash", name="datahub_agent", instruction="You help users find datasets in DataHub. Provide clear, concise answers.", tools=tools, ) async def main() -> None: session_service = InMemorySessionService() session = await session_service.create_session(app_name="datahub_simple_search", user_id="user") runner = Runner(agent=agent, app_name="datahub_simple_search", session_service=session_service) async for event in runner.run_async( user_id="user", session_id=session.id, new_message=types.Content(role="user", parts=[types.Part(text="Find datasets about users")]), ): if event.is_final_response() and event.content and event.content.parts: print(f"Agent: {event.content.parts[0].text}") asyncio.run(main())

进阶示例：DataHub + BigQuery 数据分析智能体

basic_agent.py 演示了更完整的"数据分析师"智能体：同时挂载 DataHub 工具（含写操作）与 Google ADK 自带的BigQueryToolset，形成"先到 DataHub 发现和确认表结构 → 再写 SQL 查真实数据 → 最后综合解释结果"的工作流。其中：

通过google.auth.default()探测 GCP 凭据，未找到时静默降级、只保留 DataHub 工具；
BigQuery 工具以WriteMode.BLOCKED配置为只读；
使用InMemorySessionService维持跨轮次会话上下文；
流式事件循环中实时打印工具调用名与参数，便于观察智能体行为。

通过 MCP Server 连接

不想直嵌 Python 工具时，可让 ADK 通过McpToolset连接 DataHub 的 MCP server：

from google.adk.tools.mcp_tool import McpToolset from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams toolset = McpToolset( connection_params=StreamableHTTPConnectionParams( url="https://<tenant>.acryl.io/integrations/ai/mcp" ), headers={"Authorization": f"Bearer {YOUR_TOKEN}"}, ) # 在当前任务内主动初始化，确保 AsyncExitStack 归属于本任务 await toolset.get_tools() agent = Agent( model="gemini-2.5-flash", name="datahub_agent", instruction="You help users find datasets in DataHub.", tools=[toolset], )

用完记得await toolset.close()。

自建 DataHub Core（OSS）时，MCP 端点默认为http://localhost:8080/mcp，示例 simple_mcp.py 展示了用DATAHUB_MCP_SERVER_URL环境变量覆盖默认值、并把DataHubClient的 token 直接放进Authorization头的写法，且把await toolset.get_tools()与toolset.close()放进try/finally保证资源释放。

关于 AsyncExitStack 的要点

ADK 的McpToolset通过异步上下文管理器持有 MCP 会话。若不在当前任务内主动调用await toolset.get_tools()，ADK 会在派生的任务中创建会话，随后调用close()可能抛出"Attempted to exit cancel scope in a different task"错误。正确做法是：在拥有 toolset 的同一 async 任务中初始化，并在finally块中关闭。

使用 Vertex AI 替代 Gemini Developer API

默认情况下 ADK 通过GOOGLE_API_KEY使用 Gemini Developer API。若改用 Vertex AI：

不要设置GOOGLE_API_KEY——ADK 会自动回退到 Application Default Credentials（ADC）；
确保已执行过gcloud auth application-default login。

工具上下文注入原理

无论走哪条接入路径，DataHub Agent Context 的 Python 工具都依赖contextvars实现客户端注入：create_context_wrapper在执行工具函数前set_client(client)、结束后reset_client(token)，工具内部通过get_datahub_client()获取客户端。这保证了对同一批 MCP 工具函数的复用——它们既能在 MCP 服务端按标准协议调用，也能被包装成普通函数直接交给 ADK，这是"两种接入方式、一套工具实现"的底层机制（见 utils.py 与 context.py）。

常见问题排查

工具执行报错？检查 DataHub 连接（client.config）与令牌权限，确认 token 是否具备对应实体的读取/写入范围。
智能体不调用工具？强化instruction提示词，明确列出应使用的工具与回答格式；或换用工具调用能力更强的模型（Gemini 2.0 及以上）。
出现AsyncExitStack/ 任务错误？在与 toolset 相同的 async 任务中调用await toolset.get_tools()，并在finally块中await toolset.close()。
导入报错？执行pip install datahub-agent-context[google-adk] google-adk，确保可选依赖组与 ADK 本体都已安装。