DataHub Agent Context
DataHub Agent Context provides a collection of tools and utilities for building AI agents that interact with DataHub metadata. This package contains MCP (Model Context Protocol) tools that enable AI agents to search, retrieve, and manipulate metadata in DataHub. These can be used directly to create an agent, or be included in an MCP server such as Datahub's open source MCP server.
Features
Installation
Base Installation
python3 -m pip install --upgrade pip wheel setuptools
python3 -m pip install --upgrade datahub-agent-context
With LangChain Support
For building LangChain agents with pre-built tools:
python3 -m pip install --upgrade "datahub-agent-context[langchain]"
Prerequisites
This package requires:
- Python 3.9 or higher
acryl-datahubpackage
Quick Start
Basic Example
These tools are designed to be used with an AI agent and have the responses passed directly to an LLM, so the return schema is a simple dict, but they can be used independently if desired.
from datahub.ingestion.graph.client import DataHubGraph
from datahub_agent_context.mcp_tools.search import search
from datahub_agent_context.mcp_tools.entities import get_entities
# Initialize DataHub graph client
client = DataHubClient.from_env()
# Search for datasets
with client as client:
results = search(
query="user_data",
filters={"entity_type": ["dataset"]},
num_results=10
)
# Get detailed entity information
with client as client:
entities = get_entities(
urns=[result["entity"]["urn"] for result in results["searchResults"]]
)
LangChain Integration
Build AI agents with pre-built LangChain tools:
from datahub.sdk.main_client import DataHubClient
from datahub_agent_context.langchain_tools import build_langchain_tools
# Initialize DataHub client
client = DataHubClient.from_env()
# Build all tools (read-only by default)
tools = build_langchain_tools(client, include_mutations=False)
# Or include mutation tools for tagging, descriptions, etc.
tools = build_langchain_tools(client, include_mutations=True)
# Create agent
agent = create_agent(model, tools=tools, system_prompt="...")
DataHub Cloud Tools
If you're connected to a DataHub Cloud instance, you can add Cloud-only tools like Ask DataHub (AI-powered data assistant):
from datahub_agent_context.langchain_tools import build_langchain_tools, build_langchain_cloud_tools
client = DataHubClient.from_env()
# Base tools (works on any DataHub instance)
tools = build_langchain_tools(client, include_mutations=True)
# Add Cloud-only tools (requires DataHub Cloud)
tools += build_langchain_cloud_tools(client, ask_datahub=True)
The same pattern works for Google ADK:
from datahub_agent_context.google_adk_tools import build_google_adk_tools, build_google_adk_cloud_tools
tools = build_google_adk_tools(client, include_mutations=True)
tools += build_google_adk_cloud_tools(client, ask_datahub=True)
See examples/langchain/ for complete LangChain agent examples including:
- simple_search.py - Minimal example with AWS Bedrock
Available Tools
Search Tools
search()- Search across all entity types with filters and sortingsearch_documents()- Search specifically for Document entitiesgrep_documents()- Grep for patterns in document content
Entity Tools
get_entities()- Get detailed information about entities by URNlist_schema_fields()- List and filter schema fields for datasets
Lineage Tools
get_lineage()- Get upstream or downstream lineageget_lineage_paths_between()- Get detailed paths between two entities
Query Tools
get_dataset_queries()- Get SQL queries for datasets or columns
Mutation Tools
add_tags(),remove_tags()- Manage tagsupdate_description()- Update entity descriptionsset_domains(),remove_domains()- Manage domainsadd_owners(),remove_owners()- Manage ownersadd_glossary_terms(),remove_glossary_terms()- Manage glossary termsadd_structured_properties(),remove_structured_properties()- Manage structured propertiessave_document()- Save or update a Document.
User Tools
get_me()- Get information about the authenticated user
Cloud-Only Tools (DataHub Cloud)
ask_datahub_chat()- Ask the DataHub AI assistant a question about your data catalogget_datahub_chat()- Retrieve messages and status from an Ask DataHub conversation
Architecture
The package is organized into the following modules:
mcp_tools/- Core MCP tool implementationsbase.py- Base GraphQL execution and response cleaningsearch.py- Search functionalitydocuments.py- Document search and grepentities.py- Entity retrievallineage.py- Lineage queryingqueries.py- Query retrievaltags.py,descriptions.py,domains.py, etc. - Mutation toolshelpers.py- Shared utility functionsgql/- GraphQL query definitions
Development
Setup
# Clone the repository
git clone https://github.com/datahub-project/datahub.git
cd datahub/datahub-agent-context
# Set up development environment
./gradlew :datahub-agent-context:installDev
# Run tests
./gradlew :datahub-agent-context:testFull
# Run linting
./gradlew :datahub-agent-context:lintFix
Testing
The package includes comprehensive unit tests for all tools:
# Run full test suite
./gradlew :datahub-agent-context:testFull
Support
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file datahub_agent_context-1.6.0.9.post1.tar.gz.
File metadata
- Download URL: datahub_agent_context-1.6.0.9.post1.tar.gz
- Upload date:
- Size: 100.9 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.10.20
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
7c580638e6bec0a085963a86cb4a99658f309be8e2f176d672b2c82dd7f63750
|
|
| MD5 |
3e4395411f0ead8dac180d625e822ea2
|
|
| BLAKE2b-256 |
187e0cbd454ceb5d4e0fbb2463746bde5834cc9f53308d5a8cc566a6aef72fc2
|
File details
Details for the file datahub_agent_context-1.6.0.9.post1-py3-none-any.whl.
File metadata
- Download URL: datahub_agent_context-1.6.0.9.post1-py3-none-any.whl
- Upload date:
- Size: 136.2 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.10.20
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
0dc86f59b1171b14b177ec0de83a093fd23806dcedba3c5e671e2d492c8bb5d9
|
|
| MD5 |
9160e1671316f0a1755705b747e9c326
|
|
| BLAKE2b-256 |
650c4a43c2b515db1284f866ca1dfe1fcb3cbb3b62a0863ebba51af5e3acb390
|
