AI Agents & Agent Context Kit
Relevant source files
The following files were used as context for generating this wiki page:
DataHub serves as an open-source context platform for the data and AI stack, providing a unified context graph that powers discovery, governance, and observability. The Agent Context Kit is a specialized framework designed to bring DataHub's rich metadata context to various AI agents and assistants, enabling them to better understand data structures, lineage, and business definitions. Sources: docs-website/static/llms.txt:3-17
The kit facilitates connections between AI assistants (such as Claude, Cursor, and Copilot) and DataHub metadata via the Model Context Protocol (MCP) and dedicated integrations for platforms like Snowflake Cortex, Databricks Genie, and LangChain. Sources: docs-website/static/llms.txt:19-25, docs-website/sidebars.js:88-150
Architecture and Integrations
The Agent Context Kit is organized as a collection of tools and generators that allow AI agents to interface directly with the DataHub GMS (Generalized Metadata Service). This enables agents to perform natural language search, lineage analysis, and data quality checks. Sources: datahub-agent-context/setup.py:84-118
Component Relationship
The following diagram illustrates how the Agent Context Kit bridges the gap between AI platforms and DataHub's metadata repository.
The diagram shows the flow of requests from various AI platforms through the Agent Context Kit to DataHub's backend services. Sources: docs-website/static/llms.txt:19-25, datahub-agent-context/setup.py:90-110
Supported Integrations
The Agent Context Kit includes specific support for the following ecosystems:
- LangChain: Adapters for MCP and core LangChain frameworks.
- Snowflake: Cortex Agents and UDF-based metadata tools.
- Google: Gemini CLI, Vertex AI, and Agent Development Kit (ADK).
- Databricks: Genie Code and Agent Bricks.
- IDE Tools: Integrations for Cursor and Claude.
Sources: docs-website/sidebars.js:88-150, datahub-agent-context/setup.py:61-75
DataHub Skills
DataHub Skills is an AI-assisted framework within the ingestion system that allows users to generate production-ready connectors from simple configurations. It serves as a higher-level abstraction for building integrations compared to manual Python development. Sources: metadata-ingestion/adding-source.md:7-11
| Feature | Description |
|---|---|
| Speed | Designed for building connectors in minutes instead of weeks. |
| Automation | Generates connector code based on simple configuration. |
| Skill Docs | Included within connectors (e.g., Pinecone) to provide PR reviews and implementation plans. |
| Sources: metadata-ingestion/adding-source.md:7-11, metadata-ingestion/src/datahub/ingestion/source/pinecone/skill_docs/PINECONE_datahub-connector-pr-review-2026-03-08.md:1-20 |
Snowflake Cortex Agent Implementation
The Agent Context Kit provides a SQL generator for Snowflake Cortex Agents. This generator creates a Snowflake Agent specification that incorporates DataHub metadata tools as functions or procedures within the Snowflake environment. Sources: datahub-agent-context/src/datahub_agent_context/snowflake/generators/cortex_agent.py:10-25
Core Agent Capabilities
Snowflake Cortex Agents configured via the kit have access to several capability groups:
- Search and Discovery:
search_datahub,search_documents. - Schema Exploration:
get_entities,list_schema_fields. - Lineage Analysis:
get_lineage,get_lineage_paths_between. - Query Patterns:
get_dataset_queries. - Data Quality:
get_dataset_assertions. - Metadata Management (Mutations): Adding/removing tags, owners, domains, and glossary terms.
Sources: datahub-agent-context/src/datahub_agent_context/snowflake/generators/cortex_agent.py:32-45, datahub-agent-context/src/datahub_agent_context/snowflake/generators/cortex_agent.py:440-520
Agent Interaction Flow
The following sequence diagram demonstrates how a Snowflake Cortex Agent utilizes DataHub tools to answer a user query.
This diagram illustrates the multi-step process where the agent first validates context via DataHub before executing SQL on the data platform. Sources: datahub-agent-context/src/datahub_agent_context/snowflake/generators/cortex_agent.py:55-100, datahub-agent-context/src/datahub_agent_context/snowflake/generators/cortex_agent.py:640-660
Tooling and Configuration
The kit relies on a specific Python environment and set of dependencies to bridge DataHub with AI frameworks.
Python Requirements
The datahub-agent-context package requires Python 3.9+ and includes several extras for specific integrations. Sources: datahub-agent-context/setup.py:53-118
| Extra | Dependencies |
|---|---|
| base | acryl-datahub[datahub-rest], pydantic>=2.0.0, json-repair, google-re2 |
| langchain | langchain, langchain-core, langchain-mcp-adapters |
| google-adk | google-adk, opentelemetry-resourcedetector-gcp |
| snowflake | snowflake-connector-python>=4.0.0, click |
| Sources: datahub-agent-context/setup.py:53-81 |
Tool Definitions for Agents
Agents are provided with tool_spec objects that define the interface for interacting with DataHub.
- tool_spec:
type: "generic"
name: "get_lineage"
description: "Get upstream or downstream lineage for entities or columns."
input_schema:
type: "object"
properties:
urn:
type: "string"
description: "Entity URN"
upstream:
type: "number"
description: "1 for upstream, 0 for downstream"
required: [urn, upstream]
Sources: datahub-agent-context/src/datahub_agent_context/snowflake/generators/cortex_agent.py:465-485
Summary
The AI Agents & Agent Context Kit provides the necessary infrastructure to transform DataHub from a passive catalog into an active context provider for LLMs. By providing standardized tools for search, lineage, and metadata management, it enables agents to operate with higher accuracy and reliability across complex data ecosystems. Sources: docs-website/static/llms.txt:3-10, datahub-agent-context/src/datahub_agent_context/snowflake/generators/cortex_agent.py:50-70