datahub-agent-context 1.4.0rc1__py3-none-any.whl → 1.4.0rc2__py3-none-any.whl

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/_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/context.py

@@ -1,6 +1,6 @@
 """Context management for DataHub tools.
 
-This module provides a context manager pattern for managing DataHubGraph instances
+This module provides a context manager pattern for managing DataHubClient instances
 across tool calls without explicit parameter passing.
 """
 
@@ -9,56 +9,69 @@ 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 DataHubGraph instance
-_graph_context: contextvars.ContextVar[Optional["DataHubGraph"]] = (
-    contextvars.ContextVar("datahub_graph", default=None)
+# Context variable to store the current DataHubClient instance
+_client_context: contextvars.ContextVar[Optional["DataHubClient"]] = (
+    contextvars.ContextVar("datahub_client", default=None)
 )
 
 
-def get_graph() -> "DataHubGraph":
-    """Get the current DataHubGraph from context.
+def get_datahub_client() -> "DataHubClient":
+    """Get the current DataHubClient from context.
 
     Returns:
-        DataHubGraph instance from context
+        DataHubClient instance from context
 
     Raises:
-        RuntimeError: If no graph is set in context
+        RuntimeError: If no client is set in context
     """
-    graph = _graph_context.get()
-    if graph is None:
+    client = _client_context.get()
+    if client is None:
         raise RuntimeError(
-            "No DataHubGraph in context. "
-            "Make sure to use DataHubContext context manager or set_graph() before calling tools."
+            "No DataHubClient in context. "
+            "Make sure to use DataHubContext context manager or set_client() before calling tools."
         )
-    return graph
+    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_graph(graph: "DataHubGraph") -> contextvars.Token:
-    """Set the DataHubGraph in context.
+def set_client(client: "DataHubClient") -> contextvars.Token:
+    """Set the DataHubClient in context.
 
     Args:
-        graph: DataHubGraph instance to set
+        client: DataHubClient instance to set
 
     Returns:
         Token that can be used to reset the context
     """
-    return _graph_context.set(graph)
+    return _client_context.set(client)
 
 
-def reset_graph(token: contextvars.Token) -> None:
-    """Reset the DataHubGraph context to its previous value.
+def reset_client(token: contextvars.Token) -> None:
+    """Reset the DataHubClient context to its previous value.
 
     Args:
-        token: Token returned by set_graph()
+        token: Token returned by set_client()
     """
-    _graph_context.reset(token)
+    _client_context.reset(token)
 
 
 class DataHubContext:
     """Context manager for DataHub tool execution.
 
-    This context manager sets the DataHubGraph in context for the duration
+    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:
@@ -68,30 +81,30 @@ class DataHubContext:
 
         client = DataHubClient(...)
 
-        with DataHubContext(client.graph):
-            results = search(query="users")  # No graph parameter needed!
+        with DataHubContext(client):
+            results = search(query="users")  # No client parameter needed!
     """
 
-    def __init__(self, graph: "DataHubGraph"):
+    def __init__(self, client: "DataHubClient"):
         """Initialize the context manager.
 
         Args:
-            graph: DataHubGraph instance to use in this context
+            client: DataHubClient instance to use in this context
         """
-        self.graph = graph
+        self.client = client
         self._token: Optional[contextvars.Token] = None
 
-    def __enter__(self) -> "DataHubGraph":
-        """Enter the context and set the graph.
+    def __enter__(self) -> "DataHubClient":
+        """Enter the context and set the client.
 
         Returns:
-            The DataHubGraph instance
+            The DataHubClient instance
         """
-        self._token = set_graph(self.graph)
-        return self.graph
+        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 graph."""
+        """Exit the context and reset the client."""
         if self._token is not None:
-            reset_graph(self._token)
+            reset_client(self._token)
             self._token = None

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

datahub_agent_context/mcp_tools/save_document.py

@@ -0,0 +1,634 @@
+"""Document saving tool for DataHub MCP server.
+
+This tool enables AI agents to save documents to DataHub's knowledge base.
+Documents are organized under a configurable parent folder (default: "Shared"),
+optionally with per-user subfolders for organization.
+
+Configuration via environment variables:
+- SAVE_DOCUMENT_TOOL_ENABLED: Set to "false" to disable this tool (default: enabled). Also requires TOOLS_IS_MUTATION_ENABLED enabled.
+- SAVE_DOCUMENT_PARENT_TITLE: Custom title for the parent folder (default: "Shared")
+- SAVE_DOCUMENT_ORGANIZE_BY_USER: Set to "true" to enable per-user organization (default: false)
+- SAVE_DOCUMENT_RESTRICT_UPDATES: Set to "false" to allow updating any document (default: true - only agent-created docs can be updated)
+"""
+
+import logging
+import os
+import re
+import uuid
+from datetime import datetime
+from typing import Dict, List, Literal, Optional, Tuple
+
+from datahub.metadata import schema_classes as models
+from datahub.sdk import Document
+from datahub_agent_context.context import get_datahub_client
+from datahub_agent_context.mcp_tools.base import execute_graphql
+
+logger = logging.getLogger(__name__)
+
+# Fixed root parent document ID - independent of title for future flexibility
+ROOT_PARENT_DOC_ID = "__system_shared_documents"
+
+
+def _get_parent_title() -> str:
+    """Get the configurable parent document title from environment."""
+    return os.environ.get("SAVE_DOCUMENT_PARENT_TITLE", "Shared")
+
+
+def _is_organize_by_user_enabled() -> bool:
+    """Check if per-user organization is enabled (default: False)."""
+    value = os.environ.get("SAVE_DOCUMENT_ORGANIZE_BY_USER", "false")
+    return value.lower() in ("true", "1", "yes")
+
+
+def _restrict_updates_to_shared_folder() -> bool:
+    """Check if updates should be restricted to the shared folder (default: True).
+
+    When enabled, only documents inside the shared folder can be updated.
+    This prevents accidental modification of user-created or imported documents.
+    """
+    value = os.environ.get("SAVE_DOCUMENT_RESTRICT_UPDATES", "true")
+    return value.lower() in ("true", "1", "yes")
+
+
+def _make_safe_id(text: str, max_length: int = 30) -> str:
+    """Convert text to a safe ID string."""
+    safe_id = "".join(c if c.isalnum() else "-" for c in text.lower())[:max_length]
+    safe_id = re.sub(r"-+", "-", safe_id)  # Collapse multiple dashes
+    return safe_id.strip("-")
+
+
+def _get_root_parent_id() -> str:
+    """Get the root parent document ID.
+
+    Uses a fixed ID independent of the display title to allow changing
+    the title without requiring data migration.
+    """
+    return ROOT_PARENT_DOC_ID
+
+
+def _get_root_parent_urn() -> str:
+    """Get the root parent document URN."""
+    return f"urn:li:document:{_get_root_parent_id()}"
+
+
+# Supported document types (subtypes)
+DocumentType = Literal[
+    "Insight",
+    "Decision",
+    "FAQ",
+    "Analysis",
+    "Summary",
+    "Recommendation",
+    "Note",
+    "Context",
+]
+
+
+def _get_current_user_info() -> Optional[Dict]:
+    """Fetch the current authenticated user's information."""
+
+    client = get_datahub_client()
+
+    query = """
+        query getMe {
+            me {
+                corpUser {
+                    urn
+                    username
+                    info {
+                        displayName
+                        fullName
+                        firstName
+                        lastName
+                    }
+                    editableProperties {
+                        displayName
+                    }
+                }
+            }
+        }
+    """
+
+    try:
+        result = execute_graphql(
+            client._graph,
+            query=query,
+            variables={},
+            operation_name="getMe",
+        )
+        me_data = result.get("me", {})
+        return me_data.get("corpUser") if me_data else None
+    except Exception as e:
+        logger.warning(f"Failed to get current user info: {e}")
+        return None
+
+
+def _get_user_display_name(user_info: Optional[Dict]) -> str:
+    """Extract the best display name from user info."""
+    if not user_info:
+        return "Unknown User"
+
+    # Try editable displayName first, then info fields
+    editable = user_info.get("editableProperties") or {}
+    info = user_info.get("info") or {}
+
+    return (
+        editable.get("displayName")
+        or info.get("displayName")
+        or info.get("fullName")
+        or f"{info.get('firstName', '')} {info.get('lastName', '')}".strip()
+        or user_info.get("username")
+        or "Unknown User"
+    )
+
+
+def _generate_document_id() -> str:
+    """Generate a unique document ID using UUID.
+
+    Each save creates a new document with a unique ID.
+    Format: shared-<uuid>
+    """
+    unique_id = str(uuid.uuid4())
+    return f"shared-{unique_id}"
+
+
+def _is_document_in_shared_folder(document_urn: str) -> Tuple[bool, Optional[str]]:
+    """Check if a document is within the shared documents folder.
+
+    Simple validation: document must have the shared folder as a parent/ancestor.
+
+    Returns:
+        Tuple of (is_valid, error_message)
+        - (True, None) if document is in the shared folder
+        - (False, error_message) if document is outside the folder
+    """
+
+    client = get_datahub_client()
+    root_parent_urn = _get_root_parent_urn()
+
+    # Can't update the root folder itself
+    if document_urn == root_parent_urn:
+        return False, (
+            "Cannot update the root shared documents folder. "
+            "Only documents within this folder can be updated."
+        )
+
+    try:
+        # Fetch the document
+        doc = client.entities.get(document_urn)
+        logger.debug(
+            f"Validating document {document_urn} for update, fetched: {doc is not None}"
+        )
+        if doc is None:
+            # Document doesn't exist yet - allow (will be created)
+            logger.debug(f"Document {document_urn} does not exist, allowing update")
+            return True, None
+
+        # Get documentInfo aspect
+        aspects = getattr(doc, "aspects", None) or getattr(doc, "_aspects", {})
+        logger.debug(
+            f"Document aspects type: {type(aspects)}, keys: {list(aspects.keys()) if isinstance(aspects, dict) else 'N/A'}"
+        )
+        doc_info = aspects.get("documentInfo") if isinstance(aspects, dict) else None
+
+        if doc_info is None:
+            logger.debug(f"Document {document_urn} has no documentInfo aspect")
+            return False, (
+                f"Document '{document_urn}' has no document info. "
+                "Cannot verify it's in the shared folder."
+            )
+
+        # Walk up the parent chain looking for the shared folder
+        # parentDocument is a ParentDocumentClass with a .document field containing the URN
+        parent_doc_obj = getattr(doc_info, "parentDocument", None)
+        logger.debug(f"Document parentDocument object: {parent_doc_obj}")
+        current_parent_urn = (
+            getattr(parent_doc_obj, "document", None) if parent_doc_obj else None
+        )
+        logger.debug(
+            f"Document parent URN: {current_parent_urn}, looking for root: {root_parent_urn}"
+        )
+        visited = set()
+
+        while current_parent_urn:
+            if current_parent_urn in visited:
+                break
+            visited.add(current_parent_urn)
+
+            # Found the shared folder - document is valid
+            if current_parent_urn == root_parent_urn:
+                return True, None
+
+            # Fetch parent and continue walking up
+            try:
+                parent_doc = client.entities.get(current_parent_urn)
+                if parent_doc is None:
+                    break
+                parent_aspects = getattr(parent_doc, "aspects", None) or getattr(
+                    parent_doc, "_aspects", {}
+                )
+                parent_info = (
+                    parent_aspects.get("documentInfo")
+                    if isinstance(parent_aspects, dict)
+                    else None
+                )
+                if parent_info is None:
+                    break
+                # Get next parent - again, it's a ParentDocumentClass object
+                next_parent_obj = getattr(parent_info, "parentDocument", None)
+                current_parent_urn = (
+                    getattr(next_parent_obj, "document", None)
+                    if next_parent_obj
+                    else None
+                )
+            except Exception:
+                break
+
+        return False, (
+            f"Document '{document_urn}' is not in the shared documents folder. "
+            "Only documents in this folder can be updated."
+        )
+
+    except Exception as e:
+        logger.error(f"Failed to validate document hierarchy: {e}", exc_info=True)
+        # Fail closed - if we can't validate, don't allow the update
+        return False, (
+            f"Failed to validate document hierarchy for '{document_urn}': {str(e)}. "
+            "Cannot update document without verifying it's in the shared folder."
+        )
+
+
+def _ensure_document_exists(
+    doc_id: str,
+    title: str,
+    description: str,
+    parent_urn: Optional[str] = None,
+) -> str:
+    """Ensure a document exists, creating it if necessary. Returns the URN."""
+
+    client = get_datahub_client()
+    doc_urn = f"urn:li:document:{doc_id}"
+
+    try:
+        existing = client.entities.get(doc_urn)
+        if existing is not None:
+            return doc_urn
+    except Exception:
+        pass
+
+    # Create the document
+    doc = Document.create_document(
+        id=doc_id,
+        title=title,
+        text=description,
+        subtype="Folder",
+        parent_document=parent_urn,
+        show_in_global_context=True,
+    )
+
+    try:
+        client.entities.upsert(doc)
+        logger.info(f"Created folder document: {doc_urn}")
+    except Exception as e:
+        logger.warning(f"Failed to create folder document (may already exist): {e}")
+
+    return doc_urn
+
+
+def _ensure_parent_hierarchy(user_info: Optional[Dict]) -> Tuple[str, Optional[str]]:
+    """Ensure the parent document hierarchy exists.
+
+    Returns:
+        Tuple of (parent_urn_for_document, user_urn_if_available)
+    """
+    root_title = _get_parent_title()
+    root_id = _get_root_parent_id()
+
+    # Always create the root parent
+    root_urn = _ensure_document_exists(
+        doc_id=root_id,
+        title=root_title,
+        description="Contains shared documents authored through AI agents like Ask DataHub.",
+        parent_urn=None,
+    )
+
+    # If per-user organization is disabled, return root as parent
+    if not _is_organize_by_user_enabled():
+        return root_urn, user_info.get("urn") if user_info else None
+
+    # Create user-specific folder if we have user info
+    if user_info:
+        user_urn = user_info.get("urn")
+        username = user_info.get("username", "unknown")
+        display_name = _get_user_display_name(user_info)
+
+        # Create user folder under root
+        user_folder_id = f"agent-docs-user-{_make_safe_id(username, max_length=30)}"
+        user_folder_urn = _ensure_document_exists(
+            doc_id=user_folder_id,
+            title=display_name,
+            description=f"Contains documents authored in sessions for {display_name}.",
+            parent_urn=root_urn,
+        )
+        return user_folder_urn, user_urn
+
+    # No user info available, use root
+    return root_urn, None
+
+
+def save_document(
+    document_type: DocumentType,
+    title: str,
+    content: str,
+    urn: Optional[str] = None,
+    topics: Optional[List[str]] = None,
+    related_documents: Optional[List[str]] = None,
+    related_assets: Optional[List[str]] = None,
+) -> dict:
+    """Save or update a STANDALONE document in DataHub's knowledge base. Once saved,
+    a document will be visible to all users of DataHub and to Ask DataHub AI assistant.
+
+    NOTE: This tool is for creating standalone documents (insights, FAQs, notes, etc.),
+    NOT for updating descriptions on data assets like datasets or dashboards.
+    Use update_description for asset descriptions.
+
+    WHEN TO USE THIS TOOL:
+
+    Use this tool when the user explicitly requests to save information:
+    - "Save this for later..."
+    - "Bookmark this.."
+    - "Document this insight.."
+    - "Remember this.."
+    - "Add this to our knowledge base.."
+    - "Create a document about this.."
+
+    Also SUGGEST using this tool when the user provides valuable information such as:
+    - Useful SQL queries they want to reuse
+    - Decisions about data modeling or architecture
+    - FAQs or common questions about data
+    - Analysis results worth sharing with the team
+    - Corrections or clarifications about data, service, business definitions, etc.
+
+    ⚠️  IMPORTANT: Before calling this tool, you SHOULD confirm with the user that
+    they want to save this document. Present the title, content summary,
+    and any related assets, and ask for their approval before proceeding. Do not attempt to save
+    information that would be private or user-specific.
+
+    This tool persists insights, decisions, FAQs, and other contextual information
+    as documents in DataHub. Documents are organized hierarchically:
+    - Under a configurable parent folder (default: "Shared" for global context)
+    - Optionally grouped by the user who authored them
+
+    UPSERT BEHAVIOR:
+    - If `urn` is NOT provided: Creates a NEW document with a unique URN
+    - If `urn` IS provided: Updates the EXISTING document with that URN
+
+    IMPORTANT USAGE GUIDELINES:
+    - Always confirm with the user before saving
+    - Provide a clear summary of what will be saved
+    - Ask if the user wants to proceed with creating/updating the document
+
+    REQUIRED PARAMETERS:
+
+    document_type: The type of document being saved. For example:
+        - "Insight": Data insights or discoveries
+        - "Decision": Documented decisions with rationale
+        - "FAQ": Frequently asked questions and answers
+        - "Analysis": Data analysis findings
+        - "Summary": Summaries of complex information
+        - "Recommendation": Suggested actions or improvements
+        - "Note": General notes or observations
+
+    title: A descriptive title for the document.
+        - Example: "Sales Data Quality Issues - Q4 2024"
+        - Example: "Decision: Deprecating Legacy Customer Table"
+
+    content: The full content of the document (supports markdown formatting).
+        - Can include headers, lists, code blocks, tables, etc.
+        - Example: "## Summary\\n\\nThe orders table shows 15% null values..."
+
+    OPTIONAL PARAMETERS:
+
+    urn: The URN of an existing document to update.
+        - ONLY use after a search_documents or get_entity call returns a document URN
+        - Example: "urn:li:document:agent-insight-abc123"
+        - If not provided, a new document is created with a unique URN
+        - If provided, the existing document is updated (upsert operation)
+
+    topics: List of topic tags for categorization and discovery (like a word cloud).
+        - These become searchable tags in DataHub that users can click to find related documents
+        - Example: ["data-quality", "customer-data", "Q4-2024"]
+        - Example: ["high-priority", "sales", "email", "null-values"]
+
+    related_documents: URNs of related documents.
+        - Example: ["urn:li:document:agent-insight-sales-abc123"]
+        - Creates links between related knowledge
+
+    related_assets: URNs of related data assets (tables, dashboards, etc).
+        - Example: ["urn:li:dataset:(urn:li:dataPlatform:snowflake,db.orders,PROD)"]
+        - Links the document to specific data assets in the catalog
+        - Users can then see this document when viewing those assets
+
+    Returns:
+        Dictionary with:
+        - success: Boolean indicating if the operation succeeded
+        - urn: The URN of the created/updated document
+        - message: Success or error message
+        - author: The user who authored the document (if available)
+
+    RECOMMENDED WORKFLOW:
+
+    1. Gather information you want to save publicly
+    2. Present a summary to the user:
+       "I'd like to save the following insight to DataHub:
+        - Title: High Null Rate in Customer Emails
+        - Type: Insight
+        - Related to: customers table
+        Would you like me to save this?"
+    3. Only call save_document after user confirms
+
+    EXAMPLE USAGE:
+
+    1. Create a new insight (after user confirmation):
+        save_document(
+            document_type="Insight",
+            title="High Null Rate in Customer Emails",
+            content="## Finding\\n\\n23% of customer records have null email...",
+            topics=["data-quality", "customer-data", "email", "high-severity"],
+            related_assets=["urn:li:dataset:(urn:li:dataPlatform:snowflake,customers,PROD)"]
+        )
+
+    2. Update an existing document (after finding it via search_documents):
+        save_document(
+            urn="urn:li:document:agent-insight-abc123",  # From search_documents result
+            document_type="Insight",
+            title="High Null Rate in Customer Emails (Updated)",
+            content="## Finding\\n\\nUpdated: Now 18% of customer records have null email...",
+            topics=["data-quality", "customer-data", "email", "resolved"]
+        )
+
+    3. Document a decision:
+        save_document(
+            document_type="Decision",
+            title="Migrating to New Production Database",
+            content="## Decision\\n\\nWe will migrate to v2 schema...\\n\\n## Rationale\\n...",
+            topics=["architecture", "data-model", "migration", "approved"]
+        )
+    """
+
+    client = get_datahub_client()
+
+    # Validate inputs
+    if not title or not title.strip():
+        return {
+            "success": False,
+            "urn": None,
+            "message": "title cannot be empty",
+            "author": None,
+        }
+
+    if not content or not content.strip():
+        return {
+            "success": False,
+            "urn": None,
+            "message": "content cannot be empty",
+            "author": None,
+        }
+
+    valid_document_types = [
+        "Insight",
+        "Decision",
+        "FAQ",
+        "Analysis",
+        "Summary",
+        "Recommendation",
+        "Note",
+        "Context",
+    ]
+    if document_type not in valid_document_types:
+        return {
+            "success": False,
+            "urn": None,
+            "message": f"Invalid document_type '{document_type}'. Must be one of: {', '.join(valid_document_types)}",
+            "author": None,
+        }
+
+    # Validate URN format if provided
+    if urn is not None:
+        if not urn.startswith("urn:li:document:"):
+            return {
+                "success": False,
+                "urn": None,
+                "message": f"Invalid urn format '{urn}'. Must start with 'urn:li:document:'",
+                "author": None,
+            }
+
+        # Validate that the document is within the agent-authored hierarchy
+        # This prevents accidental modification of user-created or imported documents
+        if _restrict_updates_to_shared_folder():
+            is_valid, error_message = _is_document_in_shared_folder(urn)
+            if not is_valid:
+                return {
+                    "success": False,
+                    "urn": None,
+                    "message": error_message,
+                    "author": None,
+                }
+
+        is_update = True
+        document_urn = urn
+        # Extract document ID from URN
+        document_id = urn.replace("urn:li:document:", "")
+    else:
+        is_update = False
+        # Generate new document ID
+        document_id = _generate_document_id()
+        document_urn = f"urn:li:document:{document_id}"
+
+    try:
+        # Get current user info for attribution and organization
+        user_info = _get_current_user_info()
+        user_display_name = _get_user_display_name(user_info) if user_info else None
+        user_urn = user_info.get("urn") if user_info else None
+
+        # Ensure parent hierarchy exists and get the parent URN for our document
+        # For updates, we still want to maintain proper parent hierarchy
+        parent_urn, _ = _ensure_parent_hierarchy(user_info)
+
+        # Set current user as owner only for NEW documents (captures authorship)
+        # For updates, preserve existing ownership by not setting owners
+        # The SDK expects owner URNs as strings in a list
+        if is_update:
+            owners = None  # Don't overwrite existing ownership on updates
+            logger.info("Updating existing document - preserving existing ownership")
+        else:
+            owners = [user_urn] if user_urn else None
+            logger.info(f"Creating new document - setting owners: {owners}")
+
+        # Convert topics to tag URNs (DataHub expects full URNs)
+        # TODO: Decide whether tags are the right abstraction here.
+        # Alternative: Use Structured Properties, custom properties, or custom label.
+        # Just needs to be searchable later on.
+        tag_urns = None
+        if topics:
+            tag_urns = [f"urn:li:tag:{topic}" for topic in topics]
+
+        # Create the document
+        doc = Document.create_document(
+            id=document_id,
+            title=title,
+            text=content,
+            subtype=document_type,
+            parent_document=parent_urn,
+            related_documents=related_documents,
+            related_assets=related_assets,
+            owners=owners,
+            tags=tag_urns,  # Use topics as tags for searchability
+            show_in_global_context=True,
+        )
+
+        # Manually add DocumentSettings to ensure showInGlobalContext=True is set
+        # This is a workaround until the SDK is updated to always emit this aspect
+        # TODO: Remove this once SDK properly emits DocumentSettings for show_in_global_context=True
+        actor_urn = user_urn or "urn:li:corpuser:datahub"
+        settings_audit = models.AuditStampClass(
+            time=int(datetime.now().timestamp() * 1000),
+            actor=actor_urn,
+        )
+        document_settings = models.DocumentSettingsClass(
+            showInGlobalContext=True,
+            lastModified=settings_audit,
+        )
+        doc._set_aspect(document_settings)
+
+        # Log document details before upsert
+        logger.info(
+            f"Document to upsert: URN={document_urn}, owners={owners}, parent={parent_urn}"
+        )
+
+        # Upsert the document
+        try:
+            client.entities.upsert(doc)
+            logger.info("Upsert completed successfully")
+        except Exception as upsert_error:
+            logger.error(f"Failed to upsert document: {upsert_error}", exc_info=True)
+            raise
+
+        action = "updated" if is_update else "created"
+        logger.info(f"Successfully {action} document: {document_urn}")
+
+        return {
+            "success": True,
+            "urn": document_urn,
+            "message": f"Successfully {action} document: {title}",
+            "author": user_display_name,
+        }
+
+    except Exception as e:
+        logger.error(f"Failed to save document: {e}")
+        return {
+            "success": False,
+            "urn": None,
+            "message": f"Error saving document: {str(e)}",
+            "author": None,
+        }

{datahub_agent_context-1.4.0rc1.dist-info → datahub_agent_context-1.4.0rc2.dist-info}/METADATA

@@ -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: cachetools<7.0.0,>=5.0.0
-Requires-Dist: httpcore<2.0,>=1.0.9
-Requires-Dist: json-repair<1.0.0,>=0.25.0
 Requires-Dist: h11<1.0,>=0.16
-Requires-Dist: jmespath<2.0.0,>=1.0.0
-Requires-Dist: acryl-datahub==1.4.0rc1
-Requires-Dist: pydantic<3.0.0,>=2.0.0
 Requires-Dist: google-re2<2.0,>=1.0
+Requires-Dist: acryl-datahub==1.4.0rc2
+Requires-Dist: pydantic<3.0.0,>=2.0.0
+Requires-Dist: httpcore<2.0,>=1.0.9
+Requires-Dist: jmespath<2.0.0,>=1.0.0
+Requires-Dist: json-repair<1.0.0,>=0.25.0
+Requires-Dist: cachetools<7.0.0,>=5.0.0
 Provides-Extra: dev
-Requires-Dist: tox<5.0.0,>=4.0.0; extra == "dev"
-Requires-Dist: types-PyYAML<7.0.0,>=6.0.0; extra == "dev"
-Requires-Dist: mypy==1.17.1; 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: click<9.0.0,>=8.0.0; extra == "dev"
-Requires-Dist: types-toml<1.0.0,>=0.10.0; extra == "dev"
-Requires-Dist: pytest-cov<7.0.0,>=2.8.0; extra == "dev"
 Requires-Dist: ruff==0.11.7; extra == "dev"
-Requires-Dist: pytest<9.0.0,>=8.3.4; extra == "dev"
-Requires-Dist: types-requests<3.0.0,>=2.0.0; extra == "dev"
 Requires-Dist: langchain-core<2.0.0,>=1.2.7; extra == "dev"
+Requires-Dist: snowflake-connector-python<4.0.0,>=3.0.0; extra == "dev"
+Requires-Dist: pytest-cov<7.0.0,>=2.8.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: tox<5.0.0,>=4.0.0; extra == "dev"
+Requires-Dist: click<9.0.0,>=8.0.0; extra == "dev"
+Requires-Dist: types-requests<3.0.0,>=2.0.0; extra == "dev"
+Requires-Dist: mypy==1.17.1; extra == "dev"
+Requires-Dist: types-PyYAML<7.0.0,>=6.0.0; extra == "dev"
+Requires-Dist: pytest<9.0.0,>=8.3.4; extra == "dev"
+Requires-Dist: types-jmespath<2.0.0,>=1.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.dist-info → datahub_agent_context-1.4.0rc2.dist-info}/RECORD

@@ -1,10 +1,10 @@
-datahub_agent_context/__init__.py,sha256=VGBOuNztxuwUi5Ofnrpe7tw8EmUrQD-i-eMbSKwvMtU,890
-datahub_agent_context/_version.py,sha256=oeJ65E7WFebrM-WKg-BHPZ37x3hoyefmc19QzVAGYYs,648
+datahub_agent_context/__init__.py,sha256=WgJFMZaA5ae_9ntP686UXd0TvZpbGwQRdSISi0JHsvU,967
+datahub_agent_context/_version.py,sha256=sjXQOjG_dBO3UkwODDbFi7aLjoKMlkbubWQwtnl0qh0,648
 datahub_agent_context/cli.py,sha256=ND0KLT3cFb6KnQl6kEb7B74tAOu6yfS4dO6mJjZW1x4,4441
-datahub_agent_context/context.py,sha256=qRc44o38Y-LoDQH1oFm38hIatOWRKnRxxytQNIR93kU,2771
+datahub_agent_context/context.py,sha256=wj9q9hGf72q6oarnfEFHzqgS-vwtMO79hhjz8GNC0QQ,3163
 datahub_agent_context/py.typed,sha256=kO13kg6OXApIRwKRcPpEOL09GZHx2Pk8Rp2KZpxv0lw,63
 datahub_agent_context/langchain_tools/__init__.py,sha256=M0tn6fD9qY5Wc1XdptQuIf_7MSKLX8OSBaBxcPo5wmw,259
-datahub_agent_context/langchain_tools/builder.py,sha256=X59zdmdUqltKiTo3HZrE4-JOd7CztGprW2O32jIYt2o,5145
+datahub_agent_context/langchain_tools/builder.py,sha256=-h8IuFWfIJGhQXnuK5ASCMwhoNmIhxLV2yu3ZOZPwpg,5288
 datahub_agent_context/mcp_tools/__init__.py,sha256=7iUoWuT-KvszOqnmL3_co2LVQdhZtkQKRLRE98Hn8WM,1544
 datahub_agent_context/mcp_tools/_token_estimator.py,sha256=U0kTqPZKBkKwxe7JZaLxIIFEobNSrEEHoM4NQbrmmAE,2782
 datahub_agent_context/mcp_tools/base.py,sha256=UFqMe9yS-YikgaKOnu2DkaMLaHRzwSUBOFXOWxBjULA,11054
@@ -17,6 +17,7 @@ datahub_agent_context/mcp_tools/helpers.py,sha256=NRIoVEB62vDWDg26UOFv-IhM8mEQd4
 datahub_agent_context/mcp_tools/lineage.py,sha256=sJVR2jJkbGU_KjjtqZ8IJVOKDaIjDdtQKtAIxYWq71Q,26753
 datahub_agent_context/mcp_tools/owners.py,sha256=LGZ5n5a3xRKSttay2NLf_rq97_Dl9pGIcVFi-l7uJK8,11798
 datahub_agent_context/mcp_tools/queries.py,sha256=V4-yFcCi3c8r4Xy7XVKfQ7s3SsIWXMAHRrI8Sqf2g20,6864
+datahub_agent_context/mcp_tools/save_document.py,sha256=r3nYUQSQjmgCybX7X39zu8Ef4FON284K_D1X9Y53cIE,23221
 datahub_agent_context/mcp_tools/search.py,sha256=z5Hy1jLV4uDO26nb_oFuP5w6GX0DYcYWRIWn3kDp7dY,9880
 datahub_agent_context/mcp_tools/structured_properties.py,sha256=amj7C-sbeAyctrXY_rpc2vCNTaJy2aTRx21TioeKEJk,15745
 datahub_agent_context/mcp_tools/tags.py,sha256=5_Wg1Jqf_FgPgYuUV5bDwQ6J8t_sECcSM5yVtwQruPs,10814
@@ -59,7 +60,7 @@ datahub_agent_context/snowflake/udfs/search_datahub.py,sha256=-El0JnpkClaaxX5tCZ
 datahub_agent_context/snowflake/udfs/search_documents.py,sha256=yrcLsSNyq_M-CsgwGrsgSXFVeusKAvuNqEI_DvHtC08,1944
 datahub_agent_context/snowflake/udfs/set_domains.py,sha256=nKFJ9ZMBCz9qWgJv-FkOkeEmRE-QYBUCKbHw1RFdJmc,1701
 datahub_agent_context/snowflake/udfs/update_description.py,sha256=laJjqqRARDy3VCo0xX_5lvTrZn2U3Dwm7XEliPdDRSA,2096
-datahub_agent_context-1.4.0rc1.dist-info/METADATA,sha256=ikdYmmpFuJB8wQ7Ybs4J-9VnMumP9xEV-qR6V4ScgT8,7756
-datahub_agent_context-1.4.0rc1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-datahub_agent_context-1.4.0rc1.dist-info/top_level.txt,sha256=Tv1bg7ZwDOKM9u9RHj5m1Zbx2LDf4lVBBRNHi_gBBTI,22
-datahub_agent_context-1.4.0rc1.dist-info/RECORD,,
+datahub_agent_context-1.4.0rc2.dist-info/METADATA,sha256=GzhPDkJ4HtHmsyCeu_Lm_z44amlbDRHevW7R9hSH0Y8,7795
+datahub_agent_context-1.4.0rc2.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+datahub_agent_context-1.4.0rc2.dist-info/top_level.txt,sha256=Tv1bg7ZwDOKM9u9RHj5m1Zbx2LDf4lVBBRNHi_gBBTI,22
+datahub_agent_context-1.4.0rc2.dist-info/RECORD,,