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

datahub_agent_context/snowflake/generators/cortex_agent.py

@@ -0,0 +1,725 @@
+"""Generate Snowflake Cortex Agent SQL."""
+
+
+def generate_cortex_agent_sql(
+    agent_name: str,
+    agent_display_name: str,
+    agent_color: str,
+    sf_warehouse: str | None,
+    sf_database: str | None,
+    sf_schema: str | None,
+    include_mutations: bool = True,
+) -> str:
+    """Generate Cortex Agent SQL that uses configuration variables with DataHub tools.
+
+    Args:
+        agent_name: Agent name
+        agent_display_name: Agent display name
+        agent_color: Agent color
+        sf_warehouse: Snowflake warehouse name (uses placeholder if None)
+        sf_database: Snowflake database name (uses placeholder if None)
+        sf_schema: Snowflake schema name (uses placeholder if None)
+        include_mutations: Whether to include mutation/write tools (default: True)
+    """
+    # Use placeholders for None values - these will be set via SQL variables at runtime
+    warehouse = sf_warehouse or "MY_WAREHOUSE"
+    database = sf_database or "MY_DATABASE"
+    schema = sf_schema or "MY_SCHEMA"
+
+    # Build instructions based on whether mutations are enabled
+    if include_mutations:
+        capabilities = """1. Find and query data (search, schema exploration, SQL generation)
+      2. Understand data lineage and relationships
+      3. Manage metadata (tags, descriptions, owners, domains, glossary terms)
+      4. Search documentation and runbooks"""
+
+        system_capabilities = """- 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)
+      - Metadata management (tags, descriptions, owners, domains, glossary terms)
+      - User information (get_me)"""
+
+        orchestration_guidance = """For data queries:
+      1. Use search_datahub to find relevant datasets
+      2. Use get_entities or list_schema_fields for schema details
+      3. Generate SQL based on actual schema
+      4. Execute using SqlExecutor
+
+      For lineage questions:
+      1. Use get_lineage to explore upstream/downstream dependencies
+      2. Use get_lineage_paths_between for detailed transformation chains
+
+      For metadata management:
+      1. Search for entities first to get URNs
+      2. Use appropriate tools (add_tags, update_description, etc.)
+      3. Confirm changes were successful"""
+
+        metadata_note = """Always use DataHub tools before generating SQL to ensure accuracy.
+      When managing metadata, confirm changes with the user first."""
+    else:
+        capabilities = """1. Find and query data (search, schema exploration, SQL generation)
+      2. Understand data lineage and relationships
+      3. Search documentation and runbooks"""
+
+        system_capabilities = """- 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)
+      - User information (get_me)"""
+
+        orchestration_guidance = """For data queries:
+      1. Use search_datahub to find relevant datasets
+      2. Use get_entities or list_schema_fields for schema details
+      3. Generate SQL based on actual schema
+      4. Execute using SqlExecutor
+
+      For lineage questions:
+      1. Use get_lineage to explore upstream/downstream dependencies
+      2. Use get_lineage_paths_between for detailed transformation chains"""
+
+        metadata_note = (
+            "Always use DataHub tools before generating SQL to ensure accuracy."
+        )
+
+    # Build mutation tools section if enabled
+    mutation_tools = (
+        """
+    # Tag Management Tools
+    - tool_spec:
+        type: "generic"
+        name: "add_tags"
+        description: "Add tags to entities or columns. Confirm with user first before making changes."
+        input_schema:
+          type: "object"
+          properties:
+            tag_urns:
+              type: "string"
+              description: "JSON array of tag URNs (e.g., '[\\\"urn:li:tag:PII\\\"]')"
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+            column_paths:
+              type: "string"
+              description: "JSON array of column names. Default: null (entity-level tagging)"
+          required: [tag_urns, entity_urns, column_paths]
+
+    - tool_spec:
+        type: "generic"
+        name: "remove_tags"
+        description: "Remove tags from entities or columns. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            tag_urns:
+              type: "string"
+              description: "JSON array of tag URNs to remove"
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+            column_paths:
+              type: "string"
+              description: "JSON array of column names. Default: null (entity-level tag removal)"
+          required: [tag_urns, entity_urns, column_paths]
+
+    # Description Management
+    - tool_spec:
+        type: "generic"
+        name: "update_description"
+        description: "Update entity/column descriptions. Operations: 'replace', 'append', 'remove'. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            entity_urn:
+              type: "string"
+              description: "Entity URN"
+            operation:
+              type: "string"
+              description: "'replace', 'append', or 'remove'"
+            description:
+              type: "string"
+              description: "Description text. Default: null (not needed for 'remove' operation)"
+            column_path:
+              type: "string"
+              description: "Column name. Default: null (entity-level description)"
+          required: [entity_urn, operation, description, column_path]
+
+    # Domain Management
+    - tool_spec:
+        type: "generic"
+        name: "set_domains"
+        description: "Assign a domain to entities. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            domain_urn:
+              type: "string"
+              description: "Domain URN (e.g., 'urn:li:domain:marketing')"
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+          required: [domain_urn, entity_urns]
+
+    - tool_spec:
+        type: "generic"
+        name: "remove_domains"
+        description: "Remove domain assignments from entities. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+          required: [entity_urns]
+
+    # Owner Management
+    - tool_spec:
+        type: "generic"
+        name: "add_owners"
+        description: "Add owners to entities. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            owner_urns:
+              type: "string"
+              description: "JSON array of owner URNs (CorpUser or CorpGroup)"
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+            ownership_type_urn:
+              type: "string"
+              description: "Ownership type URN. Default: null (uses default ownership type)"
+          required: [owner_urns, entity_urns, ownership_type_urn]
+
+    - tool_spec:
+        type: "generic"
+        name: "remove_owners"
+        description: "Remove owners from entities. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            owner_urns:
+              type: "string"
+              description: "JSON array of owner URNs"
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+            ownership_type_urn:
+              type: "string"
+              description: "Ownership type URN. Default: null (removes all ownership types)"
+          required: [owner_urns, entity_urns, ownership_type_urn]
+
+    # Glossary Term Management
+    - tool_spec:
+        type: "generic"
+        name: "add_glossary_terms"
+        description: "Add glossary terms to entities or columns. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            term_urns:
+              type: "string"
+              description: "JSON array of glossary term URNs"
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+            column_paths:
+              type: "string"
+              description: "JSON array of column names. Default: null (entity-level glossary terms)"
+          required: [term_urns, entity_urns, column_paths]
+
+    - tool_spec:
+        type: "generic"
+        name: "remove_glossary_terms"
+        description: "Remove glossary terms from entities or columns. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            term_urns:
+              type: "string"
+              description: "JSON array of glossary term URNs"
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+            column_paths:
+              type: "string"
+              description: "JSON array of column names. Default: null (entity-level glossary terms)"
+          required: [term_urns, entity_urns, column_paths]
+
+    # Structured Property Management
+    - tool_spec:
+        type: "generic"
+        name: "add_structured_properties"
+        description: "Add structured properties to entities or columns. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            property_values:
+              type: "string"
+              description: "JSON array of {{propertyUrn, value}} objects"
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+            column_paths:
+              type: "string"
+              description: "JSON array of column names. Default: null (entity-level structured properties)"
+          required: [property_values, entity_urns, column_paths]
+
+    - tool_spec:
+        type: "generic"
+        name: "remove_structured_properties"
+        description: "Remove structured properties from entities or columns. Confirm with user first."
+        input_schema:
+          type: "object"
+          properties:
+            property_urns:
+              type: "string"
+              description: "JSON array of property URNs to remove"
+            entity_urns:
+              type: "string"
+              description: "JSON array of entity URNs"
+            column_paths:
+              type: "string"
+              description: "JSON array of column names. Default: null (entity-level structured properties)"
+          required: [property_urns, entity_urns, column_paths]
+"""
+        if include_mutations
+        else ""
+    )
+
+    # Build mutation tool resources section if enabled
+    mutation_tool_resources = (
+        f"""
+    # Tags
+    add_tags:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.ADD_TAGS
+
+    remove_tags:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.REMOVE_TAGS
+
+    # Descriptions
+    update_description:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.UPDATE_DESCRIPTION
+
+    # Domains
+    set_domains:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.SET_DOMAINS
+
+    remove_domains:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.REMOVE_DOMAINS
+
+    # Owners
+    add_owners:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.ADD_OWNERS
+
+    remove_owners:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.REMOVE_OWNERS
+
+    # Glossary Terms
+    add_glossary_terms:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.ADD_GLOSSARY_TERMS
+
+    remove_glossary_terms:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.REMOVE_GLOSSARY_TERMS
+
+    # Structured Properties
+    add_structured_properties:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.ADD_STRUCTURED_PROPERTIES
+
+    remove_structured_properties:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.REMOVE_STRUCTURED_PROPERTIES
+"""
+        if include_mutations
+        else ""
+    )
+
+    tool_count_note = (
+        "20 tools (read + write)" if include_mutations else "9 tools (read-only)"
+    )
+    query_description = " and manage metadata" if include_mutations else ""
+    comment_suffix = " and metadata management" if include_mutations else ""
+
+    # Build sample questions based on whether mutations are enabled
+    sample_questions_with_mutations = '''
+      - question: "What tables contain customer data?"
+        answer: "I'll search DataHub for datasets related to customer data."
+      - question: "Show me the lineage for the sales_monthly table"
+        answer: "I'll retrieve the lineage information for the sales_monthly table."
+      - question: "Tag all PII datasets in the finance domain"
+        answer: "I'll search for datasets in the finance domain and add PII tags to them."
+      - question: "What queries use the users table?"
+        answer: "I'll retrieve the SQL queries that reference the users table."
+      - question: "Add a description to the revenue column"
+        answer: "I'll update the description for the revenue column."
+      - question: "Who owns the analytics datasets?"
+        answer: "I'll search for analytics datasets and show their ownership information."'''
+
+    sample_questions_readonly = '''
+      - question: "What tables contain customer data?"
+        answer: "I'll search DataHub for datasets related to customer data."
+      - question: "Show me the lineage for the sales_monthly table"
+        answer: "I'll retrieve the lineage information for the sales_monthly table."
+      - question: "What queries use the users table?"
+        answer: "I'll retrieve the SQL queries that reference the users table."
+      - question: "Who owns the analytics datasets?"
+        answer: "I'll search for analytics datasets and show their ownership information."'''
+
+    sample_questions = (
+        sample_questions_with_mutations
+        if include_mutations
+        else sample_questions_readonly
+    )
+    return f"""-- ============================================================================
+-- Step 4: Create Cortex Agent with DataHub Tools
+-- ============================================================================
+-- This creates a Snowflake Cortex Agent that uses DataHub metadata
+-- to generate accurate SQL queries{query_description}
+--
+-- Prerequisites:
+-- - Run 00_configuration.sql first to set variables
+-- - Run 01_network_rules.sql to set up network access
+-- - Run 02_datahub_udfs.sql to create DataHub UDFs ({tool_count_note})
+-- - Run 03_stored_procedure.sql to create EXECUTE_DYNAMIC_SQL
+-- ============================================================================
+
+USE DATABASE IDENTIFIER($SF_DATABASE);
+USE SCHEMA IDENTIFIER($SF_SCHEMA);
+USE WAREHOUSE IDENTIFIER($SF_WAREHOUSE);
+
+CREATE OR REPLACE AGENT {agent_name}
+  COMMENT = 'Agent that uses DataHub metadata for SQL generation{comment_suffix}'
+  PROFILE = '{{"display_name": "{agent_display_name}", "color": "{agent_color}"}}'
+  FROM SPECIFICATION
+  $$
+  models:
+    orchestration: auto
+
+  orchestration:
+    budget:
+      seconds: 60
+      tokens: 32000
+
+  instructions:
+    response: |
+      You are a comprehensive data assistant with access to DataHub metadata.
+      You can help users:
+      {capabilities}
+
+      {metadata_note}
+
+    orchestration: |
+      {orchestration_guidance}
+
+    system: |
+      You have comprehensive access to DataHub including:
+      {system_capabilities}
+
+    sample_questions:{sample_questions}
+
+  tools:
+    # Core Search & Discovery Tools
+    - tool_spec:
+        type: "generic"
+        name: "search_datahub"
+        description: "Search DataHub for entities (datasets, dashboards, etc.). Use /q prefix for structured queries. Returns URNs, names, descriptions, and metadata."
+        input_schema:
+          type: "object"
+          properties:
+            search_query:
+              type: "string"
+              description: "Search query (e.g., 'customer', '/q user+transaction')"
+            entity_type:
+              type: "string"
+              description: "Entity type filter (e.g., 'dataset', 'tag', etc.). Default: null (all entity types)"
+          required: [search_query, entity_type]
+
+    - tool_spec:
+        type: "generic"
+        name: "get_entities"
+        description: "Get detailed entity information including schema, tags, owners, lineage summary. Use URN from search results."
+        input_schema:
+          type: "object"
+          properties:
+            entity_urn:
+              type: "string"
+              description: "Entity URN from search results"
+          required: [entity_urn]
+
+    - tool_spec:
+        type: "generic"
+        name: "list_schema_fields"
+        description: "List schema fields with filtering and pagination. Useful for large schemas or finding specific columns."
+        input_schema:
+          type: "object"
+          properties:
+            dataset_urn:
+              type: "string"
+              description: "Dataset URN"
+            keywords:
+              type: "string"
+              description: "Keywords to filter fields (single string or JSON array). Default: null (no filtering)"
+            limit:
+              type: "number"
+              description: "Max fields to return. Default: 100"
+          required: [dataset_urn, keywords, limit]
+
+    # Lineage Tools
+    - tool_spec:
+        type: "generic"
+        name: "get_lineage"
+        description: "Get upstream or downstream lineage for entities or columns. Returns lineage graph with metadata."
+        input_schema:
+          type: "object"
+          properties:
+            urn:
+              type: "string"
+              description: "Entity URN"
+            column_name:
+              type: "string"
+              description: "Column name for column-level lineage. Default: null (entity-level lineage)"
+            upstream:
+              type: "number"
+              description: "1 for upstream, 0 for downstream. Default: 1"
+            max_hops:
+              type: "number"
+              description: "Max hops (1-3+). Default: 1"
+            max_results:
+              type: "number"
+              description: "Max results. Default: 30"
+          required: [urn, column_name, upstream, max_hops, max_results]
+
+    - tool_spec:
+        type: "generic"
+        name: "get_lineage_paths_between"
+        description: "Get detailed transformation paths between two entities/columns. Shows intermediate steps and queries."
+        input_schema:
+          type: "object"
+          properties:
+            source_urn:
+              type: "string"
+              description: "Source dataset URN"
+            target_urn:
+              type: "string"
+              description: "Target dataset URN"
+            source_column:
+              type: "string"
+              description: "Source column name. Default: null (dataset-level lineage)"
+            target_column:
+              type: "string"
+              description: "Target column name. Default: null (dataset-level lineage)"
+          required: [source_urn, target_urn, source_column, target_column]
+
+    # Query Analysis Tools
+    - tool_spec:
+        type: "generic"
+        name: "get_dataset_queries"
+        description: "Get SQL queries that use a dataset/column. Filter by MANUAL (user queries) or SYSTEM (BI tools)."
+        input_schema:
+          type: "object"
+          properties:
+            urn:
+              type: "string"
+              description: "Dataset URN"
+            column_name:
+              type: "string"
+              description: "Column name to filter queries. Default: null (queries for all columns)"
+            source:
+              type: "string"
+              description: "'MANUAL', 'SYSTEM', or null for both. Default: null"
+            count:
+              type: "number"
+              description: "Number of queries. Default: 10"
+          required: [urn, column_name, source, count]
+
+    # Document Search Tools
+    - tool_spec:
+        type: "generic"
+        name: "search_documents"
+        description: "Search organization documents (runbooks, FAQs, knowledge articles from Notion, Confluence, etc.)."
+        input_schema:
+          type: "object"
+          properties:
+            search_query:
+              type: "string"
+              description: "Search query"
+            num_results:
+              type: "number"
+              description: "Max results. Default: 10"
+          required: [search_query, num_results]
+
+    - tool_spec:
+        type: "generic"
+        name: "grep_documents"
+        description: "Search within document content using regex patterns. Use after search_documents to find specific content."
+        input_schema:
+          type: "object"
+          properties:
+            urns:
+              type: "string"
+              description: "JSON array of document URNs"
+            pattern:
+              type: "string"
+              description: "Regex pattern to search for"
+            context_chars:
+              type: "number"
+              description: "Context characters. Default: 200"
+            max_matches_per_doc:
+              type: "number"
+              description: "Max matches per document. Default: 5"
+          required: [urns, pattern, context_chars, max_matches_per_doc]
+{mutation_tools}
+    # User Info
+    - tool_spec:
+        type: "generic"
+        name: "get_me"
+        description: "Get information about the authenticated user (profile, groups, privileges). This tool takes no parameters."
+        input_schema:
+          type: "object"
+          properties: {{}}
+
+    # SQL Executor
+    - tool_spec:
+        type: "generic"
+        name: "SqlExecutor"
+        description: "Execute SELECT SQL queries and return results. Use after generating SQL from DataHub metadata."
+        input_schema:
+          type: "object"
+          properties:
+            SQL_TEXT:
+              type: "string"
+              description: "SELECT SQL query (must start with SELECT)"
+          required: [SQL_TEXT]
+
+  tool_resources:
+    # Search & Discovery
+    search_datahub:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.SEARCH_DATAHUB
+
+    get_entities:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.GET_ENTITIES
+
+    list_schema_fields:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.LIST_SCHEMA_FIELDS
+
+    # Lineage
+    get_lineage:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.GET_LINEAGE
+
+    get_lineage_paths_between:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.GET_LINEAGE_PATHS_BETWEEN
+
+    # Query Analysis
+    get_dataset_queries:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.GET_DATASET_QUERIES
+
+    # Documents
+    search_documents:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.SEARCH_DOCUMENTS
+
+    grep_documents:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.GREP_DOCUMENTS
+{mutation_tool_resources}
+    # User Info
+    get_me:
+      type: "function"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.GET_ME
+
+    # SQL Executor
+    SqlExecutor:
+      type: "procedure"
+      execution_environment:
+        type: "warehouse"
+        warehouse: {warehouse}
+      identifier: {database}.{schema}.EXECUTE_DYNAMIC_SQL
+  $$;
+
+-- Grant usage to the specified role
+GRANT USAGE ON AGENT {agent_name} TO ROLE IDENTIFIER($SF_ROLE);
+
+-- Verify the agent was created
+DESCRIBE AGENT {agent_name};
+
+SELECT
+    'Agent created successfully with {"20 DataHub tools (read + write)" if include_mutations else "9 DataHub tools (read-only)"}!' AS status,
+    '{agent_name}' AS agent_name,
+    'You can now use this agent in Snowflake Intelligence UI for {"SQL generation and metadata management" if include_mutations else "SQL generation and metadata exploration"}' AS next_steps;
+"""