Merge branch 'main' of github.com:taylorwilsdon/google_workspace_mcp into integration-fresh

This commit is contained in:
Taylor Wilsdon
2025-08-17 14:01:00 -04:00
6 changed files with 1385 additions and 422 deletions

96
core/tool_registry.py Normal file
View File

@@ -0,0 +1,96 @@
"""
Tool Registry for Conditional Tool Registration
This module provides a registry system that allows tools to be conditionally registered
based on tier configuration, replacing direct @server.tool() decorators.
"""
import logging
from typing import Set, Optional, Callable
logger = logging.getLogger(__name__)
# Global registry of enabled tools
_enabled_tools: Optional[Set[str]] = None
def set_enabled_tools(tool_names: Optional[Set[str]]):
"""Set the globally enabled tools."""
global _enabled_tools
_enabled_tools = tool_names
def get_enabled_tools() -> Optional[Set[str]]:
"""Get the set of enabled tools, or None if all tools are enabled."""
return _enabled_tools
def is_tool_enabled(tool_name: str) -> bool:
"""Check if a specific tool is enabled."""
if _enabled_tools is None:
return True # All tools enabled by default
return tool_name in _enabled_tools
def conditional_tool(server, tool_name: str):
"""
Decorator that conditionally registers a tool based on the enabled tools set.
Args:
server: The FastMCP server instance
tool_name: The name of the tool to register
Returns:
Either the registered tool decorator or a no-op decorator
"""
def decorator(func: Callable) -> Callable:
if is_tool_enabled(tool_name):
logger.debug(f"Registering tool: {tool_name}")
return server.tool()(func)
else:
logger.debug(f"Skipping tool registration: {tool_name}")
return func
return decorator
def wrap_server_tool_method(server):
"""
Track tool registrations and filter them post-registration.
"""
original_tool = server.tool
server._tracked_tools = []
def tracking_tool(*args, **kwargs):
original_decorator = original_tool(*args, **kwargs)
def wrapper_decorator(func: Callable) -> Callable:
tool_name = func.__name__
server._tracked_tools.append(tool_name)
# Always apply the original decorator to register the tool
return original_decorator(func)
return wrapper_decorator
server.tool = tracking_tool
def filter_server_tools(server):
"""Remove disabled tools from the server after registration."""
enabled_tools = get_enabled_tools()
if enabled_tools is None:
return
tools_removed = 0
# Access FastMCP's tool registry via _tool_manager._tools
if hasattr(server, '_tool_manager'):
tool_manager = server._tool_manager
if hasattr(tool_manager, '_tools'):
tool_registry = tool_manager._tools
tools_to_remove = []
for tool_name in list(tool_registry.keys()):
if not is_tool_enabled(tool_name):
tools_to_remove.append(tool_name)
for tool_name in tools_to_remove:
del tool_registry[tool_name]
tools_removed += 1
if tools_removed > 0:
logger.info(f"🔧 Tool tier filtering: removed {tools_removed} tools, {len(enabled_tools)} enabled")

181
core/tool_tier_loader.py Normal file
View File

@@ -0,0 +1,181 @@
"""
Tool Tier Loader Module
This module provides functionality to load and resolve tool tiers from the YAML configuration.
It integrates with the existing tool enablement workflow to support tiered tool loading.
"""
import logging
from pathlib import Path
from typing import Dict, List, Set, Literal, Optional
import yaml
logger = logging.getLogger(__name__)
TierLevel = Literal["core", "extended", "complete"]
class ToolTierLoader:
"""Loads and manages tool tiers from configuration."""
def __init__(self, config_path: Optional[str] = None):
"""
Initialize the tool tier loader.
Args:
config_path: Path to the tool_tiers.yaml file. If None, uses default location.
"""
if config_path is None:
# Default to core/tool_tiers.yaml relative to this file
config_path = Path(__file__).parent / "tool_tiers.yaml"
self.config_path = Path(config_path)
self._tiers_config: Optional[Dict] = None
def _load_config(self) -> Dict:
"""Load the tool tiers configuration from YAML file."""
if self._tiers_config is not None:
return self._tiers_config
if not self.config_path.exists():
raise FileNotFoundError(f"Tool tiers configuration not found: {self.config_path}")
try:
with open(self.config_path, 'r', encoding='utf-8') as f:
self._tiers_config = yaml.safe_load(f)
logger.info(f"Loaded tool tiers configuration from {self.config_path}")
return self._tiers_config
except yaml.YAMLError as e:
raise ValueError(f"Invalid YAML in tool tiers configuration: {e}")
except Exception as e:
raise RuntimeError(f"Failed to load tool tiers configuration: {e}")
def get_available_services(self) -> List[str]:
"""Get list of all available services defined in the configuration."""
config = self._load_config()
return list(config.keys())
def get_tools_for_tier(self, tier: TierLevel, services: Optional[List[str]] = None) -> List[str]:
"""
Get all tools for a specific tier level.
Args:
tier: The tier level (core, extended, complete)
services: Optional list of services to filter by. If None, includes all services.
Returns:
List of tool names for the specified tier level
"""
config = self._load_config()
tools = []
# If no services specified, use all available services
if services is None:
services = self.get_available_services()
for service in services:
if service not in config:
logger.warning(f"Service '{service}' not found in tool tiers configuration")
continue
service_config = config[service]
if tier not in service_config:
logger.debug(f"Tier '{tier}' not defined for service '{service}'")
continue
tier_tools = service_config[tier]
if tier_tools: # Handle empty lists
tools.extend(tier_tools)
return tools
def get_tools_up_to_tier(self, tier: TierLevel, services: Optional[List[str]] = None) -> List[str]:
"""
Get all tools up to and including the specified tier level.
Args:
tier: The maximum tier level to include
services: Optional list of services to filter by. If None, includes all services.
Returns:
List of tool names up to the specified tier level
"""
tier_order = ["core", "extended", "complete"]
max_tier_index = tier_order.index(tier)
tools = []
for i in range(max_tier_index + 1):
current_tier = tier_order[i]
tools.extend(self.get_tools_for_tier(current_tier, services))
# Remove duplicates while preserving order
seen = set()
unique_tools = []
for tool in tools:
if tool not in seen:
seen.add(tool)
unique_tools.append(tool)
return unique_tools
def get_services_for_tools(self, tool_names: List[str]) -> Set[str]:
"""
Get the service names that provide the specified tools.
Args:
tool_names: List of tool names to lookup
Returns:
Set of service names that provide any of the specified tools
"""
config = self._load_config()
services = set()
for service, service_config in config.items():
for tier_name, tier_tools in service_config.items():
if tier_tools and any(tool in tier_tools for tool in tool_names):
services.add(service)
break
return services
def get_tools_for_tier(tier: TierLevel, services: Optional[List[str]] = None) -> List[str]:
"""
Convenience function to get tools for a specific tier.
Args:
tier: The tier level (core, extended, complete)
services: Optional list of services to filter by
Returns:
List of tool names for the specified tier level
"""
loader = ToolTierLoader()
return loader.get_tools_up_to_tier(tier, services)
def resolve_tools_from_tier(tier: TierLevel, services: Optional[List[str]] = None) -> tuple[List[str], List[str]]:
"""
Resolve tool names and service names for the specified tier.
Args:
tier: The tier level (core, extended, complete)
services: Optional list of services to filter by
Returns:
Tuple of (tool_names, service_names) where:
- tool_names: List of specific tool names for the tier
- service_names: List of service names that should be imported
"""
loader = ToolTierLoader()
# Get all tools for the tier
tools = loader.get_tools_up_to_tier(tier, services)
# Map back to service names
service_names = loader.get_services_for_tools(tools)
logger.info(f"Tier '{tier}' resolved to {len(tools)} tools across {len(service_names)} services: {sorted(service_names)}")
return tools, sorted(service_names)

132
core/tool_tiers.yaml Normal file
View File

@@ -0,0 +1,132 @@
gmail:
core:
- search_gmail_messages
- get_gmail_message_content
- get_gmail_messages_content_batch
- send_gmail_message
extended:
- get_gmail_thread_content
- modify_gmail_message_labels
- list_gmail_labels
- manage_gmail_label
- draft_gmail_message
complete:
- get_gmail_threads_content_batch
- batch_modify_gmail_message_labels
- start_google_auth
drive:
core:
- search_drive_files
- get_drive_file_content
- create_drive_file
extended:
- list_drive_items
complete: []
calendar:
core:
- list_calendars
- get_events
- create_event
- modify_event
extended:
- delete_event
complete: []
docs:
core:
- get_doc_content
- create_doc
- modify_doc_text
extended:
- search_docs
- find_and_replace_doc
- list_docs_in_folder
- insert_doc_elements
complete:
- insert_doc_image
- update_doc_headers_footers
- batch_update_doc
- inspect_doc_structure
- create_table_with_data
- debug_table_structure
- read_document_comments
- create_document_comment
- reply_to_document_comment
- resolve_document_comment
sheets:
core:
- create_spreadsheet
- read_sheet_values
- modify_sheet_values
extended:
- list_spreadsheets
- get_spreadsheet_info
complete:
- create_sheet
- read_spreadsheet_comments
- create_spreadsheet_comment
- reply_to_spreadsheet_comment
- resolve_spreadsheet_comment
chat:
core:
- send_message
- get_messages
- search_messages
extended:
- list_spaces
complete: []
forms:
core:
- create_form
- get_form
extended:
- list_form_responses
complete:
- set_publish_settings
- get_form_response
slides:
core:
- create_presentation
- get_presentation
extended:
- batch_update_presentation
- get_page
- get_page_thumbnail
complete:
- read_presentation_comments
- create_presentation_comment
- reply_to_presentation_comment
- resolve_presentation_comment
tasks:
core:
- get_task
- list_tasks
- create_task
- update_task
extended:
- delete_task
complete:
- list_task_lists
- get_task_list
- create_task_list
- update_task_list
- delete_task_list
- move_task
- clear_completed_tasks
search:
core:
- search_custom
extended:
- search_custom_siterestrict
complete:
- get_search_engine_info