add: support for custom_metadata filterds in search_assets tool

Author: SatabrataPaul-GitAc

modelcontextprotocol/server.py

@@ -27,6 +27,7 @@
 @mcp.tool()
 def search_assets_tool(
     conditions=None,
+    custom_metadata_conditions=None,
     negative_conditions=None,
     some_conditions=None,
     min_somes=1,
@@ -50,6 +51,8 @@ def search_assets_tool(
     Args:
         conditions (Dict[str, Any], optional): Dictionary of attribute conditions to match.
             Format: {"attribute_name": value} or {"attribute_name": {"operator": operator, "value": value}}
+        custom_metadata_conditions (List[Dict[str, Any]], optional): List of custom metadata conditions to match.
+            Format: [{"custom_metadata": value}] or [{"custom_metadata": {"operator": operator, "value": value}}]
         negative_conditions (Dict[str, Any], optional): Dictionary of attribute conditions to exclude.
             Format: {"attribute_name": value} or {"attribute_name": {"operator": operator, "value": value}}
         some_conditions (Dict[str, Any], optional): Conditions for where_some() queries that require min_somes of them to match.
@@ -95,6 +98,15 @@ def search_assets_tool(
             include_attributes=["owner_users", "owner_groups"]
         )
 
+        # Search for assets with custom metadata
+        asset_list_1 = search_assets(
+            custom_metadata_conditions=[{"custom_metadata_filter": {"display_name": "test-mcp", "property_filters": [{"property_name": "mcp_allow_status", "property_value": "yes"}]}}]
+        )
+
+        asset_list_2 = search_assets(
+            custom_metadata_conditions=[{"custom_metadata_filter": {"display_name": "test-mcp", "property_filters": [{"property_name": "mcp_allow_status", "property_value": "yes", "operator": "eq"}]}}]
+        )
+
         # Search for columns with specific certificate status
         columns = search_assets(
             asset_type="Column",
@@ -219,6 +231,7 @@ def search_assets_tool(
     try:
         # Parse JSON string parameters if needed
         conditions = parse_json_parameter(conditions)
+        custom_metadata_conditions = parse_json_parameter(custom_metadata_conditions)
         negative_conditions = parse_json_parameter(negative_conditions)
         some_conditions = parse_json_parameter(some_conditions)
         date_range = parse_json_parameter(date_range)
@@ -229,6 +242,7 @@ def search_assets_tool(
 
         return search_assets(
             conditions,
+            custom_metadata_conditions,
             negative_conditions,
             some_conditions,
             min_somes,
@@ -683,42 +697,42 @@ def create_glossary_categories(categories) -> List[Dict[str, Any]]:
 def detect_custom_metadata_from_query(query_text: str) -> Dict[str, Any]:
     """
     Detect custom metadata triggers from natural language queries.
-    
-    This tool analyzes natural language text to identify when users are referencing 
-    custom metadata (business metadata) and automatically provides context about 
+
+    This tool analyzes natural language text to identify when users are referencing
+    custom metadata (business metadata) and automatically provides context about
     available custom metadata definitions. Use this tool when you receive natural
     language queries that might involve custom metadata concepts.
-    
+
     Args:
         query_text (str): Natural language query text to analyze for custom metadata references
-        
+
     Returns:
         Dict[str, Any]: Dictionary containing:
             - detected: Boolean indicating if custom metadata was detected
             - context: Custom metadata context if detected (list of metadata definitions)
             - detection_reasons: List of reasons why custom metadata was detected
             - suggested_attributes: List of suggested custom metadata attributes
-            
+
     Detection Triggers:
         The tool detects custom metadata usage when the query contains:
         - Business metadata keywords (e.g., "business metadata", "data classification")
         - Data governance terms (e.g., "PII", "GDPR", "compliance", "data quality")
         - Attribute patterns (e.g., "sensitivity level", "business owner", "data steward")
         - Quality and classification terms (e.g., "quality score", "classification level")
-        
+
     Examples:
         # Query mentioning data classification
         result = detect_custom_metadata_from_query("Find all tables with sensitive data classification")
-        
+
         # Query about data quality
         result = detect_custom_metadata_from_query("Show me assets with poor data quality scores")
-        
+
         # Query about business ownership
         result = detect_custom_metadata_from_query("Which datasets have John as the business owner?")
-        
+
         # Query about compliance
         result = detect_custom_metadata_from_query("Find all PII data that needs GDPR compliance review")
-        
+
     Use Cases:
         - Analyze user queries before executing searches to provide custom metadata context
         - Understand when users are asking about business metadata attributes
@@ -733,7 +747,7 @@ def detect_custom_metadata_from_query(query_text: str) -> Dict[str, Any]:
             "context": None,
             "detection_reasons": [],
             "suggested_attributes": [],
-            "error": f"Failed to detect custom metadata: {str(e)}"
+            "error": f"Failed to detect custom metadata: {str(e)}",
         }
 
 

modelcontextprotocol/tools/search.py

@@ -14,6 +14,7 @@
 
 def search_assets(
     conditions: Optional[Union[Dict[str, Any], str]] = None,
+    custom_metadata_conditions: Optional[List[Dict[str, Any]]] = None,
     negative_conditions: Optional[Dict[str, Any]] = None,
     some_conditions: Optional[Dict[str, Any]] = None,
     min_somes: int = 1,
@@ -40,6 +41,8 @@ def search_assets(
     Args:
         conditions (Dict[str, Any], optional): Dictionary of attribute conditions to match.
             Format: {"attribute_name": value} or {"attribute_name": {"operator": operator, "value": value}}
+        custom_metadata_conditions (List[Dict[str, Any]], optional): List of custom metadata conditions to match.
+            Format: [{"custom_metadata": value}] or [{"custom_metadata": {"operator": operator, "value": value}}]
         negative_conditions (Dict[str, Any], optional): Dictionary of attribute conditions to exclude.
             Format: {"attribute_name": value} or {"attribute_name": {"operator": operator, "value": value}}
         some_conditions (Dict[str, Any], optional): Conditions for where_some() queries that require min_somes of them to match.
@@ -187,6 +190,19 @@ def search_assets(
                 )
             search = search.min_somes(min_somes)
 
+        if custom_metadata_conditions:
+            logger.debug(
+                f"Applying custom metadata conditions: {custom_metadata_conditions}"
+            )
+            for custom_metadata_filter_onject in custom_metadata_conditions:
+                if isinstance(custom_metadata_filter_onject, dict):
+                    _, condition = next(iter(custom_metadata_filter_onject.items()))
+                else:
+                    condition = custom_metadata_filter_onject
+                search = SearchUtils._process_custom_metadata_condition(
+                    search, condition, "where"
+                )
+
         # Apply date range filters
         if date_range:
             logger.debug(f"Applying date range filters: {date_range}")

modelcontextprotocol/utils/search.py

@@ -1,6 +1,8 @@
-from typing import Dict, Any
 import logging
+from typing import Dict, Any
 from pyatlan.model.assets import Asset
+from pyatlan.model.fields.atlan_fields import CustomMetadataField
+from pyatlan.model.fluent_search import FluentSearch
 
 logger = logging.getLogger(__name__)
 
@@ -170,3 +172,73 @@ def _process_condition(
             )
             search = search_method(attr.eq(condition))
             return search
+
+    @staticmethod
+    def _process_custom_metadata_condition(
+        search: FluentSearch, condition: Dict[str, Any], search_method_name: str
+    ):
+        """
+        Process a single custom metadata condition and apply it to the search using the specified method.
+
+        Args:
+            search: The FluentSearch object
+            condition: Dictionary containing display_name, property_name, property_value, and optional operator
+            search_method_name: The search method to use ('where', 'where_not', 'where_some')
+
+        Returns:
+            FluentSearch: The updated search object
+        """
+        if not isinstance(condition, dict):
+            logger.warning("Custom metadata condition must be a dictionary")
+            return search
+
+        # Validate required fields
+        required_fields = ["display_name", "property_filters"]
+        if not all(field in condition for field in required_fields):
+            logger.warning(
+                f"Custom metadata condition missing required fields: {required_fields}"
+            )
+            return search
+
+        search_method = getattr(search, search_method_name)
+
+        # Get operator, default to "eq"
+        for property_filter in condition["property_filters"]:
+            operator = property_filter.get("operator", "eq")
+            property_name = property_filter.get("property_name")
+            property_value = property_filter.get("property_value")
+
+            try:
+                # Create the custom metadata field
+                custom_metadata_field = CustomMetadataField(
+                    set_name=condition["display_name"], attribute_name=property_name
+                )
+
+                # Apply the appropriate operator
+                if property_value == "any":
+                    # For "any" value, use has_any_value() method
+                    query_condition = custom_metadata_field.has_any_value()
+                else:
+                    # Get the operator method dynamically
+                    if hasattr(custom_metadata_field, operator):
+                        operator_method = getattr(custom_metadata_field, operator)
+                        query_condition = operator_method(property_value)
+                    else:
+                        # Fallback to eq if operator not found
+                        logger.warning(
+                            f"Operator '{operator}' not found, falling back to 'eq'"
+                        )
+                        query_condition = custom_metadata_field.eq(property_value)
+
+                # Apply the condition to the search
+                search = search_method(query_condition)
+                logger.info(search)
+                logger.debug(
+                    f"Applied custom metadata condition: {condition['display_name']}.{condition['property_name']} {operator} {condition['property_value']}"
+                )
+
+            except Exception as e:
+                logger.error(f"Error processing custom metadata condition: {e}")
+                logger.exception("Exception details:")
+
+        return search