dtSpark 1.0.4__py3-none-any.whl

dtSpark/mcp_integration/tool_selector.py

@@ -0,0 +1,225 @@
+"""
+Tool selection module for optimising token usage.
+
+This module implements intelligent tool selection to reduce the number
+of tools sent with each API request, significantly reducing token consumption.
+"""
+
+import logging
+from typing import List, Dict, Any, Set
+
+
+class ToolSelector:
+    """Selects relevant tools based on conversation context."""
+
+    # Define tool categories and their associated tool name patterns
+    TOOL_CATEGORIES = {
+        'aws_security': ['prowler', 'scan', 'findings', 'security', 'compliance', 'vulnerability'],
+        'docker': ['container', 'docker', 'image', 'compose', 'network', 'volume'],
+        'documentation': ['note', 'obsidian', 'create_note', 'update_note', 'search_notes', 'vault'],
+        'threat_intelligence': ['opencti', 'indicator', 'threat', 'malware', 'stix', 'observable'],
+        'aws_infrastructure': ['ec2', 's3', 'lambda', 'cloudwatch', 'iam', 'vpc', 'rds', 'dynamodb', 'diagram'],
+        'elasticsearch': ['elasticsearch', 'search', 'index', 'query', 'aggregation'],
+        'ragstore': ['ragstore', 'rag', 'embedding', 'vector', 'semantic'],
+    }
+
+    # Keywords in user messages that trigger specific categories
+    CATEGORY_KEYWORDS = {
+        'aws_security': ['security', 'prowler', 'scan', 'findings', 'vulnerabilities', 'compliance',
+                        'threat', 'risk', 'audit', 'cis', 'benchmark'],
+        'docker': ['container', 'docker', 'image', 'compose', 'containerised', 'containerized'],
+        'documentation': ['note', 'obsidian', 'document', 'report', 'write', 'create note',
+                         'update note', 'markdown', 'vault'],
+        'threat_intelligence': ['threat', 'indicator', 'malware', 'opencti', 'ioc', 'attack',
+                               'campaign', 'actor', 'ttp'],
+        'aws_infrastructure': ['ec2', 's3', 'lambda', 'resource', 'aws', 'cloud', 'infrastructure',
+                              'vpc', 'subnet', 'instance', 'bucket', 'function', 'diagram'],
+        'elasticsearch': ['elasticsearch', 'search', 'query', 'index', 'log', 'aggregate'],
+        'ragstore': ['ragstore', 'rag', 'embedding', 'semantic', 'vector', 'similarity'],
+    }
+
+    def __init__(self, max_tools_per_request: int = 30):
+        """
+        Initialise the tool selector.
+
+        Args:
+            max_tools_per_request: Maximum number of tools to include in a single request
+        """
+        self.max_tools_per_request = max_tools_per_request
+        logging.info(f"ToolSelector initialised with max {max_tools_per_request} tools per request")
+
+    def select_tools(self, all_tools: List[Dict[str, Any]], user_message: str,
+                    conversation_history: List[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
+        """
+        Select relevant tools based on user message and conversation context.
+
+        Args:
+            all_tools: List of all available tools
+            user_message: The current user message
+            conversation_history: Recent conversation messages
+
+        Returns:
+            List of selected tools
+        """
+        if not all_tools:
+            logging.debug("No tools available for selection")
+            return []
+
+        # Always include built-in tools (they're small and always useful)
+        selected_tools = [t for t in all_tools if t.get('server') == 'builtin']
+        logging.debug(f"Included {len(selected_tools)} built-in tools")
+
+        # Detect relevant categories based on user message and history
+        relevant_categories = self._detect_categories(user_message, conversation_history)
+
+        if not relevant_categories:
+            # If no specific categories detected, include a diverse sample
+            logging.info("No specific tool categories detected, selecting diverse sample")
+            return self._select_diverse_sample(all_tools, selected_tools)
+
+        logging.info(f"Detected relevant categories: {relevant_categories}")
+
+        # Build set of relevant tool name patterns
+        relevant_patterns = set()
+        for category in relevant_categories:
+            relevant_patterns.update(self.TOOL_CATEGORIES.get(category, []))
+
+        # Track selected tool names to avoid duplicates
+        selected_tool_names = {t.get('name') for t in selected_tools}
+
+        # Select tools that match the relevant patterns
+        for tool in all_tools:
+            if len(selected_tools) >= self.max_tools_per_request:
+                break
+
+            tool_name = tool.get('name', '')
+            if tool_name in selected_tool_names:
+                continue
+
+            tool_name_lower = tool_name.lower()
+            tool_desc = tool.get('description', '').lower()
+
+            # Check if tool name or description matches any relevant pattern
+            if any(pattern in tool_name_lower or pattern in tool_desc for pattern in relevant_patterns):
+                selected_tools.append(tool)
+                selected_tool_names.add(tool_name)
+
+        # If still below limit and we have room, add some general-purpose tools
+        if len(selected_tools) < self.max_tools_per_request:
+            remaining = self.max_tools_per_request - len(selected_tools)
+            logging.debug(f"Adding up to {remaining} additional tools to reach limit")
+
+            for tool in all_tools:
+                if len(selected_tools) >= self.max_tools_per_request:
+                    break
+
+                tool_name = tool.get('name', '')
+                if tool_name not in selected_tool_names:
+                    selected_tools.append(tool)
+                    selected_tool_names.add(tool_name)
+
+        logging.info(f"Selected {len(selected_tools)} tools from {len(all_tools)} available " +
+                    f"(categories: {', '.join(relevant_categories)})")
+
+        # Log which tools were selected for debugging
+        tool_names = [t.get('name') for t in selected_tools]
+        logging.debug(f"Selected tools: {', '.join(tool_names[:10])}{'...' if len(tool_names) > 10 else ''}")
+
+        return selected_tools
+
+    def _detect_categories(self, user_message: str,
+                          conversation_history: List[Dict[str, Any]] = None) -> Set[str]:
+        """
+        Detect relevant tool categories from user message and history.
+
+        Args:
+            user_message: Current user message
+            conversation_history: Recent conversation messages
+
+        Returns:
+            Set of relevant category names
+        """
+        categories = set()
+
+        # Analyse user message
+        message_lower = user_message.lower()
+        for category, keywords in self.CATEGORY_KEYWORDS.items():
+            if any(keyword in message_lower for keyword in keywords):
+                categories.add(category)
+                logging.debug(f"Category '{category}' detected from user message")
+
+        # Analyse recent conversation history (last 5 messages)
+        if conversation_history:
+            recent_messages = conversation_history[-5:]
+            for msg in recent_messages:
+                # Handle both string content and dict content
+                if isinstance(msg, dict):
+                    content = str(msg.get('content', '')).lower()
+                else:
+                    content = str(msg).lower()
+
+                for category, keywords in self.CATEGORY_KEYWORDS.items():
+                    if any(keyword in content for keyword in keywords):
+                        if category not in categories:
+                            categories.add(category)
+                            logging.debug(f"Category '{category}' detected from conversation history")
+
+        return categories
+
+    def _select_diverse_sample(self, all_tools: List[Dict[str, Any]],
+                               already_selected: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Select a diverse sample of tools when no specific categories are detected.
+
+        Attempts to get tools from each category to provide broad coverage.
+
+        Args:
+            all_tools: All available tools
+            already_selected: Tools already selected (e.g., built-in tools)
+
+        Returns:
+            List of selected tools
+        """
+        selected = list(already_selected)
+        selected_tool_names = {t.get('name') for t in selected}
+
+        # Calculate how many tools to get from each category
+        remaining_slots = self.max_tools_per_request - len(selected)
+        tools_per_category = max(1, remaining_slots // len(self.TOOL_CATEGORIES))
+
+        logging.debug(f"Selecting ~{tools_per_category} tools from each category for diversity")
+
+        for category, category_patterns in self.TOOL_CATEGORIES.items():
+            if len(selected) >= self.max_tools_per_request:
+                break
+
+            added = 0
+            for tool in all_tools:
+                tool_name = tool.get('name', '')
+                if tool_name in selected_tool_names:
+                    continue
+
+                tool_name_lower = tool_name.lower()
+                tool_desc = tool.get('description', '').lower()
+
+                # Check if tool matches this category
+                if any(pattern in tool_name_lower or pattern in tool_desc for pattern in category_patterns):
+                    selected.append(tool)
+                    selected_tool_names.add(tool_name)
+                    added += 1
+                    if added >= tools_per_category or len(selected) >= self.max_tools_per_request:
+                        break
+
+        # If still below limit, add remaining tools
+        if len(selected) < self.max_tools_per_request:
+            for tool in all_tools:
+                if len(selected) >= self.max_tools_per_request:
+                    break
+
+                tool_name = tool.get('name', '')
+                if tool_name not in selected_tool_names:
+                    selected.append(tool)
+                    selected_tool_names.add(tool_name)
+
+        logging.info(f"Selected {len(selected)} diverse tools (no specific category detected)")
+        return selected