datahub-agent-context 1.3.1.8__py3-none-any.whl

datahub_agent_context/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Acryl Data, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""DataHub Agent Context - MCP Tools for AI Agents."""
+
+from datahub_agent_context._version import __version__
+from datahub_agent_context.context import (
+    DataHubContext,
+    get_graph,
+    reset_graph,
+    set_graph,
+)
+
+__all__ = ["__version__", "DataHubContext", "get_graph", "set_graph", "reset_graph"]

datahub_agent_context/_version.py

@@ -0,0 +1,16 @@
+# Copyright 2025 Acryl Data, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+__package_name__ = "datahub-agent-context"
+__version__ = "1.3.1.8"

datahub_agent_context/context.py

@@ -0,0 +1,97 @@
+"""Context management for DataHub tools.
+
+This module provides a context manager pattern for managing DataHubGraph instances
+across tool calls without explicit parameter passing.
+"""
+
+import contextvars
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from datahub.ingestion.graph.client import DataHubGraph
+
+# Context variable to store the current DataHubGraph instance
+_graph_context: contextvars.ContextVar[Optional["DataHubGraph"]] = (
+    contextvars.ContextVar("datahub_graph", default=None)
+)
+
+
+def get_graph() -> "DataHubGraph":
+    """Get the current DataHubGraph from context.
+
+    Returns:
+        DataHubGraph instance from context
+
+    Raises:
+        RuntimeError: If no graph is set in context
+    """
+    graph = _graph_context.get()
+    if graph is None:
+        raise RuntimeError(
+            "No DataHubGraph in context. "
+            "Make sure to use DataHubContext context manager or set_graph() before calling tools."
+        )
+    return graph
+
+
+def set_graph(graph: "DataHubGraph") -> contextvars.Token:
+    """Set the DataHubGraph in context.
+
+    Args:
+        graph: DataHubGraph instance to set
+
+    Returns:
+        Token that can be used to reset the context
+    """
+    return _graph_context.set(graph)
+
+
+def reset_graph(token: contextvars.Token) -> None:
+    """Reset the DataHubGraph context to its previous value.
+
+    Args:
+        token: Token returned by set_graph()
+    """
+    _graph_context.reset(token)
+
+
+class DataHubContext:
+    """Context manager for DataHub tool execution.
+
+    This context manager sets the DataHubGraph in context for the duration
+    of the with block, allowing tools to access it without explicit parameter passing.
+
+    Example:
+        from datahub.sdk.main_client import DataHubClient
+        from datahub_agent_context.context import DataHubContext
+        from datahub_agent_context.mcp_tools import search
+
+        client = DataHubClient(...)
+
+        with DataHubContext(client.graph):
+            results = search(query="users")  # No graph parameter needed!
+    """
+
+    def __init__(self, graph: "DataHubGraph"):
+        """Initialize the context manager.
+
+        Args:
+            graph: DataHubGraph instance to use in this context
+        """
+        self.graph = graph
+        self._token: Optional[contextvars.Token] = None
+
+    def __enter__(self) -> "DataHubGraph":
+        """Enter the context and set the graph.
+
+        Returns:
+            The DataHubGraph instance
+        """
+        self._token = set_graph(self.graph)
+        return self.graph
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the context and reset the graph."""
+        if self._token is not None:
+            reset_graph(self._token)
+            self._token = None

datahub_agent_context/langchain_tools/__init__.py

@@ -0,0 +1,8 @@
+"""LangChain integration for DataHub Agent Context.
+
+This module provides utilities for building LangChain tools from DataHub MCP tools.
+"""
+
+from datahub_agent_context.langchain_tools.builder import build_langchain_tools
+
+__all__ = ["build_langchain_tools"]

datahub_agent_context/langchain_tools/builder.py

@@ -0,0 +1,127 @@
+"""Builder for LangChain tools from DataHub MCP tools."""
+
+import functools
+from typing import TYPE_CHECKING, Callable
+
+from datahub_agent_context.context import set_graph
+from datahub_agent_context.mcp_tools import get_me
+from datahub_agent_context.mcp_tools.documents import grep_documents, search_documents
+from datahub_agent_context.mcp_tools.domains import remove_domains, set_domains
+from datahub_agent_context.mcp_tools.structured_properties import (
+    add_structured_properties,
+    remove_structured_properties,
+)
+
+if TYPE_CHECKING:
+    from datahub.sdk.main_client import DataHubClient
+
+try:
+    from langchain_core.tools import tool  # type: ignore[import-not-found]
+    from langchain_core.tools.base import BaseTool  # type: ignore[import-not-found]
+except ImportError as e:
+    raise ImportError(
+        "langchain-core is required for LangChain tools. "
+        "Install with: pip install 'datahub-agent-context[langchain]'"
+    ) from e
+
+from datahub_agent_context.mcp_tools.descriptions import update_description
+from datahub_agent_context.mcp_tools.entities import get_entities, list_schema_fields
+from datahub_agent_context.mcp_tools.lineage import (
+    get_lineage,
+    get_lineage_paths_between,
+)
+from datahub_agent_context.mcp_tools.owners import add_owners, remove_owners
+from datahub_agent_context.mcp_tools.queries import get_dataset_queries
+from datahub_agent_context.mcp_tools.search import search
+from datahub_agent_context.mcp_tools.tags import add_tags, remove_tags
+from datahub_agent_context.mcp_tools.terms import (
+    add_glossary_terms,
+    remove_glossary_terms,
+)
+
+
+def create_context_wrapper(func: Callable, client: "DataHubClient") -> Callable:
+    """Create a wrapper that sets DataHubGraph context before calling the function.
+
+    This wrapper uses contextvars to set the graph in context for the duration
+    of the function call, allowing the tool to retrieve it using get_graph().
+
+    Args:
+        func: The tool function that retrieves graph from context
+        client: DataHubClient instance whose graph will be set in context
+
+    Returns:
+        Wrapped function that sets context before execution
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        # Set graph in context for this function call
+        token = set_graph(client._graph)
+        try:
+            return func(*args, **kwargs)
+        finally:
+            # Always reset context, even if function raises
+            from datahub_agent_context.context import reset_graph
+
+            reset_graph(token)
+
+    return wrapper
+
+
+def build_langchain_tools(
+    client: "DataHubClient",
+    include_mutations: bool = False,
+) -> list[BaseTool]:
+    """Build LangChain tools with automatic context management.
+
+    Each tool is wrapped to automatically set the DataHubGraph in context
+    before execution, allowing tools to retrieve it using get_graph().
+
+    Args:
+        client: DataHubClient instance
+        include_mutations: Whether to include mutation tools (default: False)
+
+    Returns:
+        List of LangChain BaseTool instances
+
+    Example:
+        from datahub.sdk.main_client import DataHubClient
+        from datahub_agent_context.langchain_tools import build_langchain_tools
+
+        client = DataHubClient(...)
+        tools = build_langchain_tools(client, include_mutations=True)
+
+        # Use with LangChain agents - context is managed automatically
+        agent = create_react_agent(llm, tools, prompt)
+        result = agent.invoke({"input": "search for datasets"})
+    """
+    tools = []
+
+    # Wrap each tool to manage context automatically
+    tools.append(tool(create_context_wrapper(search_documents, client)))
+    tools.append(tool(create_context_wrapper(grep_documents, client)))
+
+    tools.append(tool(create_context_wrapper(get_entities, client)))
+    tools.append(tool(create_context_wrapper(list_schema_fields, client)))
+    tools.append(tool(create_context_wrapper(get_me, client)))
+    tools.append(tool(create_context_wrapper(get_lineage, client)))
+    tools.append(tool(create_context_wrapper(get_lineage_paths_between, client)))
+
+    tools.append(tool(create_context_wrapper(get_dataset_queries, client)))
+    tools.append(tool(create_context_wrapper(search, client)))
+
+    if include_mutations:
+        tools.append(tool(create_context_wrapper(update_description, client)))
+        tools.append(tool(create_context_wrapper(set_domains, client)))
+        tools.append(tool(create_context_wrapper(remove_domains, client)))
+        tools.append(tool(create_context_wrapper(add_owners, client)))
+        tools.append(tool(create_context_wrapper(remove_owners, client)))
+        tools.append(tool(create_context_wrapper(add_structured_properties, client)))
+        tools.append(tool(create_context_wrapper(remove_structured_properties, client)))
+        tools.append(tool(create_context_wrapper(add_tags, client)))
+        tools.append(tool(create_context_wrapper(remove_tags, client)))
+        tools.append(tool(create_context_wrapper(add_glossary_terms, client)))
+        tools.append(tool(create_context_wrapper(remove_glossary_terms, client)))
+
+    return tools

datahub_agent_context/mcp_tools/__init__.py

@@ -0,0 +1,46 @@
+"""MCP tools for interacting with DataHub metadata."""
+
+from datahub_agent_context.mcp_tools.descriptions import update_description
+from datahub_agent_context.mcp_tools.documents import grep_documents, search_documents
+from datahub_agent_context.mcp_tools.domains import remove_domains, set_domains
+from datahub_agent_context.mcp_tools.entities import get_entities, list_schema_fields
+from datahub_agent_context.mcp_tools.get_me import get_me
+from datahub_agent_context.mcp_tools.lineage import (
+    get_lineage,
+    get_lineage_paths_between,
+)
+from datahub_agent_context.mcp_tools.owners import add_owners, remove_owners
+from datahub_agent_context.mcp_tools.queries import get_dataset_queries
+from datahub_agent_context.mcp_tools.search import search
+from datahub_agent_context.mcp_tools.structured_properties import (
+    add_structured_properties,
+    remove_structured_properties,
+)
+from datahub_agent_context.mcp_tools.tags import add_tags, remove_tags
+from datahub_agent_context.mcp_tools.terms import (
+    add_glossary_terms,
+    remove_glossary_terms,
+)
+
+__all__ = [
+    "search",
+    "get_entities",
+    "list_schema_fields",
+    "get_lineage",
+    "get_lineage_paths_between",
+    "get_dataset_queries",
+    "search_documents",
+    "grep_documents",
+    "add_tags",
+    "remove_tags",
+    "update_description",
+    "set_domains",
+    "remove_domains",
+    "add_owners",
+    "remove_owners",
+    "add_glossary_terms",
+    "remove_glossary_terms",
+    "add_structured_properties",
+    "remove_structured_properties",
+    "get_me",
+]

datahub_agent_context/mcp_tools/_token_estimator.py

@@ -0,0 +1,71 @@
+"""Token count estimation utilities for MCP responses."""
+
+import logging
+from typing import Union
+
+logger = logging.getLogger(__name__)
+
+
+class TokenCountEstimator:
+    """Fast token estimation for MCP response budget management.
+
+    Uses character-based heuristics instead of actual tokenization for performance.
+    Accuracy is sufficient given the 90% budget buffer used in practice.
+    """
+
+    @staticmethod
+    def estimate_dict_tokens(
+        obj: Union[dict, list, str, int, float, bool, None],
+    ) -> int:
+        """Fast approximation of token count for dict/list structures.
+
+        Recursively walks structure counting characters. Much faster than json.dumps + estimate_tokens.
+
+        IMPORTANT: Assumes no circular references in the structure.
+        Protected against infinite recursion with MAX_DEPTH=100.
+
+        Args:
+            obj: Dict, list, or primitive value (must not contain circular references)
+
+        Returns:
+            Approximate token count
+        """
+        MAX_DEPTH = 100
+
+        def _count_chars(item, depth: int = 0) -> int:
+            if depth > MAX_DEPTH:
+                logger.error(
+                    f"Max depth {MAX_DEPTH} exceeded in structure, stopping recursion"
+                )
+                return 0
+
+            if item is None:
+                return 4  # "null"
+            elif isinstance(item, bool):
+                return 5  # "true" or "false"
+            elif isinstance(item, str):
+                # Account for:
+                # - Quotes around string values: "value" → +6
+                # - Escape characters (\n, \", \\, etc.) → +10% of length
+                # Structural chars weighted heavier as they often tokenize separately
+                base_length = len(item)
+                escape_overhead = int(base_length * 0.1)
+                return base_length + 6 + escape_overhead
+            elif isinstance(item, (int, float)):
+                return 6  # Average number length
+            elif isinstance(item, list):
+                return sum(_count_chars(elem, depth + 1) for elem in item) + len(item)
+            elif isinstance(item, dict):
+                total = 0
+                for key, value in item.items():
+                    # Account for: "key": value, → 2 quotes + colon + space + comma
+                    # Structural chars weighted heavier (often separate tokens)
+                    total += len(str(key)) + 9
+                    total += _count_chars(value, depth + 1)
+                return total + len(item)  # Additional padding for structure
+            else:
+                return 10  # Fallback for other types
+
+        chars = _count_chars(obj, depth=0)
+        # Use same formula as estimate_tokens for consistency
+        return int(1.3 * chars / 4)