datahub-agent-context 1.4.0rc1__tar.gz → 1.4.0rc2__tar.gz

{datahub_agent_context-1.4.0rc1 → datahub_agent_context-1.4.0rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datahub-agent-context
-Version: 1.4.0rc1
+Version: 1.4.0rc2
 Summary: DataHub Agent Context - MCP Tools for AI Agents
 Home-page: https://datahub.io/
 License: Apache License 2.0
@@ -28,28 +28,28 @@ Classifier: Environment :: MacOS X
 Classifier: Topic :: Software Development
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: json-repair<1.0.0,>=0.25.0
 Requires-Dist: jmespath<2.0.0,>=1.0.0
+Requires-Dist: acryl-datahub==1.4.0rc2
+Requires-Dist: cachetools<7.0.0,>=5.0.0
 Requires-Dist: httpcore<2.0,>=1.0.9
-Requires-Dist: h11<1.0,>=0.16
-Requires-Dist: acryl-datahub==1.4.0rc1
 Requires-Dist: pydantic<3.0.0,>=2.0.0
-Requires-Dist: cachetools<7.0.0,>=5.0.0
+Requires-Dist: h11<1.0,>=0.16
 Requires-Dist: google-re2<2.0,>=1.0
+Requires-Dist: json-repair<1.0.0,>=0.25.0
 Provides-Extra: dev
-Requires-Dist: snowflake-connector-python<4.0.0,>=3.0.0; extra == "dev"
-Requires-Dist: types-cachetools<7.0.0,>=5.0.0; extra == "dev"
-Requires-Dist: types-toml<1.0.0,>=0.10.0; extra == "dev"
 Requires-Dist: types-requests<3.0.0,>=2.0.0; extra == "dev"
 Requires-Dist: click<9.0.0,>=8.0.0; extra == "dev"
-Requires-Dist: tox<5.0.0,>=4.0.0; extra == "dev"
-Requires-Dist: types-jmespath<2.0.0,>=1.0.0; extra == "dev"
-Requires-Dist: ruff==0.11.7; extra == "dev"
+Requires-Dist: types-PyYAML<7.0.0,>=6.0.0; extra == "dev"
+Requires-Dist: types-toml<1.0.0,>=0.10.0; extra == "dev"
 Requires-Dist: langchain-core<2.0.0,>=1.2.7; extra == "dev"
+Requires-Dist: types-cachetools<7.0.0,>=5.0.0; extra == "dev"
+Requires-Dist: pytest-cov<7.0.0,>=2.8.0; extra == "dev"
 Requires-Dist: mypy==1.17.1; extra == "dev"
+Requires-Dist: ruff==0.11.7; extra == "dev"
+Requires-Dist: types-jmespath<2.0.0,>=1.0.0; extra == "dev"
+Requires-Dist: snowflake-connector-python<4.0.0,>=3.0.0; extra == "dev"
 Requires-Dist: pytest<9.0.0,>=8.3.4; extra == "dev"
-Requires-Dist: types-PyYAML<7.0.0,>=6.0.0; extra == "dev"
-Requires-Dist: pytest-cov<7.0.0,>=2.8.0; extra == "dev"
+Requires-Dist: tox<5.0.0,>=4.0.0; extra == "dev"
 Provides-Extra: langchain
 Requires-Dist: langchain-core<2.0.0,>=1.2.7; extra == "langchain"
 Provides-Extra: snowflake
@@ -111,7 +111,7 @@ from datahub_agent_context.mcp_tools.entities import get_entities
 client = DataHubClient.from_env()
 
 # Search for datasets
-with client.graph as graph:
+with client as client:
     results = search(
         query="user_data",
         filters={"entity_type": ["dataset"]},
@@ -119,7 +119,7 @@ with client.graph as graph:
     )
 
 # Get detailed entity information
-with client.graph as graph:
+with client as client:
     entities = get_entities(
         urns=[result["entity"]["urn"] for result in results["searchResults"]]
     )
@@ -181,6 +181,7 @@ agent = create_agent(model, tools=tools, system_prompt="...")
 - `add_owners()`, `remove_owners()` - Manage owners
 - `add_glossary_terms()`, `remove_glossary_terms()` - Manage glossary terms
 - `add_structured_properties()`, `remove_structured_properties()` - Manage structured properties
+- `save_document()` - Save or update a Document.
 
 #### User Tools
 

{datahub_agent_context-1.4.0rc1 → datahub_agent_context-1.4.0rc2}/README.md

@@ -43,7 +43,7 @@ from datahub_agent_context.mcp_tools.entities import get_entities
 client = DataHubClient.from_env()
 
 # Search for datasets
-with client.graph as graph:
+with client as client:
     results = search(
         query="user_data",
         filters={"entity_type": ["dataset"]},
@@ -51,7 +51,7 @@ with client.graph as graph:
     )
 
 # Get detailed entity information
-with client.graph as graph:
+with client as client:
     entities = get_entities(
         urns=[result["entity"]["urn"] for result in results["searchResults"]]
     )
@@ -113,6 +113,7 @@ agent = create_agent(model, tools=tools, system_prompt="...")
 - `add_owners()`, `remove_owners()` - Manage owners
 - `add_glossary_terms()`, `remove_glossary_terms()` - Manage glossary terms
 - `add_structured_properties()`, `remove_structured_properties()` - Manage structured properties
+- `save_document()` - Save or update a Document.
 
 #### User Tools
 

{datahub_agent_context-1.4.0rc1 → datahub_agent_context-1.4.0rc2}/src/datahub_agent_context/__init__.py

@@ -17,9 +17,17 @@
 from datahub_agent_context._version import __version__
 from datahub_agent_context.context import (
     DataHubContext,
+    get_datahub_client,
     get_graph,
-    reset_graph,
-    set_graph,
+    reset_client,
+    set_client,
 )
 
-__all__ = ["__version__", "DataHubContext", "get_graph", "set_graph", "reset_graph"]
+__all__ = [
+    "__version__",
+    "DataHubContext",
+    "get_datahub_client",
+    "get_graph",
+    "set_client",
+    "reset_client",
+]

{datahub_agent_context-1.4.0rc1 → datahub_agent_context-1.4.0rc2}/src/datahub_agent_context/_version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 __package_name__ = "datahub-agent-context"
-__version__ = "1.4.0rc1"
+__version__ = "1.4.0rc2"

datahub_agent_context-1.4.0rc2/src/datahub_agent_context/context.py

@@ -0,0 +1,110 @@
+"""Context management for DataHub tools.
+
+This module provides a context manager pattern for managing DataHubClient 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
+    from datahub.sdk.main_client import DataHubClient
+
+# Context variable to store the current DataHubClient instance
+_client_context: contextvars.ContextVar[Optional["DataHubClient"]] = (
+    contextvars.ContextVar("datahub_client", default=None)
+)
+
+
+def get_datahub_client() -> "DataHubClient":
+    """Get the current DataHubClient from context.
+
+    Returns:
+        DataHubClient instance from context
+
+    Raises:
+        RuntimeError: If no client is set in context
+    """
+    client = _client_context.get()
+    if client is None:
+        raise RuntimeError(
+            "No DataHubClient in context. "
+            "Make sure to use DataHubContext context manager or set_client() before calling tools."
+        )
+    return client
+
+
+def get_graph() -> "DataHubGraph":
+    """Get the current DataHubGraph from context (convenience method).
+
+    Returns:
+        DataHubGraph instance from the client in context
+
+    Raises:
+        RuntimeError: If no client is set in context
+    """
+    return get_datahub_client()._graph
+
+
+def set_client(client: "DataHubClient") -> contextvars.Token:
+    """Set the DataHubClient in context.
+
+    Args:
+        client: DataHubClient instance to set
+
+    Returns:
+        Token that can be used to reset the context
+    """
+    return _client_context.set(client)
+
+
+def reset_client(token: contextvars.Token) -> None:
+    """Reset the DataHubClient context to its previous value.
+
+    Args:
+        token: Token returned by set_client()
+    """
+    _client_context.reset(token)
+
+
+class DataHubContext:
+    """Context manager for DataHub tool execution.
+
+    This context manager sets the DataHubClient 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):
+            results = search(query="users")  # No client parameter needed!
+    """
+
+    def __init__(self, client: "DataHubClient"):
+        """Initialize the context manager.
+
+        Args:
+            client: DataHubClient instance to use in this context
+        """
+        self.client = client
+        self._token: Optional[contextvars.Token] = None
+
+    def __enter__(self) -> "DataHubClient":
+        """Enter the context and set the client.
+
+        Returns:
+            The DataHubClient instance
+        """
+        self._token = set_client(self.client)
+        return self.client
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the context and reset the client."""
+        if self._token is not None:
+            reset_client(self._token)
+            self._token = None

{datahub_agent_context-1.4.0rc1 → datahub_agent_context-1.4.0rc2}/src/datahub_agent_context/langchain_tools/builder.py

@@ -3,7 +3,7 @@
 import functools
 from typing import TYPE_CHECKING, Callable
 
-from datahub_agent_context.context import set_graph
+from datahub_agent_context.context import set_client
 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
@@ -32,6 +32,7 @@ from datahub_agent_context.mcp_tools.lineage import (
 )
 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.save_document import save_document
 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 (
@@ -57,14 +58,14 @@ def create_context_wrapper(func: Callable, client: "DataHubClient") -> Callable:
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         # Set graph in context for this function call
-        token = set_graph(client._graph)
+        token = set_client(client)
         try:
             return func(*args, **kwargs)
         finally:
             # Always reset context, even if function raises
-            from datahub_agent_context.context import reset_graph
+            from datahub_agent_context.context import reset_client
 
-            reset_graph(token)
+            reset_client(token)
 
     return wrapper
 
@@ -123,5 +124,6 @@ def build_langchain_tools(
         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)))
+        tools.append(tool(create_context_wrapper(save_document, client)))
 
     return tools