Dynamic Per-User Tool Generation MCP Server

""" Hierarchical Entity Cache System Provides in-memory tree caches for hierarchical entities (locations, departments) with fast lookups and path generation capabilities. Based on the OptimizedLocationCache pattern from treiTree.py but generalized to support multiple entity types. """ import sys import bisect import logging from dataclasses import dataclass from typing import Dict, List, Optional, Tuple, Union from collections import defaultdict logger = logging.getLogger(__name__) @dataclass class HierarchicalNode: """ Lightweight hierarchical node with minimal memory footprint. Represents a single entity in a hierarchy (location, department, etc.) """ name: str id: Union[str, int] parent_id: Optional[Union[str, int]] = None def __post_init__(self): # Intern strings to reduce memory usage for duplicate names if isinstance(self.name, str): self.name = sys.intern(self.name) if isinstance(self.id, str): self.id = sys.intern(self.id) if isinstance(self.parent_id, str): self.parent_id = sys.intern(self.parent_id) def __repr__(self): return f"HierarchicalNode(name='{self.name}', id='{self.id}')" class HierarchicalCache: """ High-performance hierarchical entity cache. Features: - O(1) ID lookups via HashMap - O(log N) name searches via sorted arrays + binary search - O(1) path lookups via pre-computed cache - Handles deep hierarchies efficiently - Thread-safe for read operations after construction """ def __init__(self, entity_type: str = "entity"): """ Initialize hierarchical cache. Args: entity_type: Type of entity being cached (for logging/debugging) """ self.entity_type = entity_type # Primary storage: O(1) ID lookups self.nodes_by_id: Dict[Union[str, int], HierarchicalNode] = {} # Sorted indices for binary search: O(log N) lookups self.names_sorted: List[Tuple[str, Union[str, int]]] = [] # (name, id) self.names_lower_sorted: List[Tuple[str, Union[str, int]]] = [] # (name.lower(), id) # Name-to-IDs mapping for handling duplicates self.name_to_ids: Dict[str, List[Union[str, int]]] = defaultdict(list) # Pre-computed paths cache: O(1) path lookups self.paths_cache: Dict[Union[str, int], str] = {} self.paths_name_cache: Dict[Union[str, int], str] = {} # Children mapping for hierarchy traversal self.children_by_parent: Dict[Union[str, int], List[Union[str, int]]] = defaultdict(list) # Build state self._indices_built = False self._paths_computed = False def add_node(self, name: str, id_: Union[str, int], parent_id: Optional[Union[str, int]] = None): """ Add a node to the cache. Args: name: Entity name id_: Unique entity identifier parent_id: Parent entity ID (None for root nodes) """ # Create and store node node = HierarchicalNode(name, id_, parent_id) self.nodes_by_id[id_] = node # Track children for hierarchy if parent_id is not None: self.children_by_parent[parent_id].append(id_) # Track name duplicates self.name_to_ids[name.lower()].append(id_) # Mark indices as needing rebuild self._indices_built = False self._paths_computed = False def update_node(self, id_: Union[str, int], name: str, parent_id: Optional[Union[str, int]] = None) -> bool: """ Update an existing node in the cache, or add it if it doesn't exist (upsert). This method replaces an existing node's name and parent_id while maintaining cache integrity. If the node doesn't exist, it will be added to the cache. All related indices and tracking structures are updated. Args: id_: Unique entity identifier of the node to update/add name: New entity name parent_id: New parent entity ID (None for root nodes) Returns: True if node was updated/added successfully Example: >>> cache.update_node(123, "New York Office", parent_id=100) True """ # Check if node exists - if not, add it if id_ not in self.nodes_by_id: logger.debug(f"Node {id_} not found in cache, adding it") self.add_node(name, id_, parent_id) return True old_node = self.nodes_by_id[id_] old_name = old_node.name old_parent_id = old_node.parent_id # Remove old name from name tracking if old_name.lower() in self.name_to_ids: try: self.name_to_ids[old_name.lower()].remove(id_) # Clean up empty lists if not self.name_to_ids[old_name.lower()]: del self.name_to_ids[old_name.lower()] except ValueError: logger.warning(f"Node {id_} not found in name_to_ids for '{old_name}'") # Remove from old parent's children list if old_parent_id is not None and old_parent_id in self.children_by_parent: try: self.children_by_parent[old_parent_id].remove(id_) # Clean up empty lists if not self.children_by_parent[old_parent_id]: del self.children_by_parent[old_parent_id] except ValueError: logger.warning(f"Node {id_} not found in children list of parent {old_parent_id}") # Create updated node updated_node = HierarchicalNode(name, id_, parent_id) self.nodes_by_id[id_] = updated_node # Add to new parent's children list if parent_id is not None: self.children_by_parent[parent_id].append(id_) # Track new name self.name_to_ids[name.lower()].append(id_) # Mark indices as needing rebuild self._indices_built = False self._paths_computed = False logger.debug( f"Updated {self.entity_type} node {id_}: " f"name '{old_name}' -> '{name}', " f"parent {old_parent_id} -> {parent_id}" ) return True def remove_node(self, id_: Union[str, int]) -> bool: """ Remove a node from the cache by its ID. This method removes a node and updates all related indices and tracking structures. If the node has children, they will be orphaned (parent_id set to None). Args: id_: Unique entity identifier of the node to remove Returns: True if node was removed successfully, False if node doesn't exist Example: >>> cache.remove_node(123) True """ # Check if node exists if id_ not in self.nodes_by_id: logger.warning(f"Cannot remove node {id_}: node not found in cache") return False node = self.nodes_by_id[id_] node_name = node.name node_parent_id = node.parent_id # Handle children - orphan them by setting their parent_id to None if id_ in self.children_by_parent: child_ids = self.children_by_parent[id_].copy() for child_id in child_ids: if child_id in self.nodes_by_id: child_node = self.nodes_by_id[child_id] # Create new node with parent_id = None orphaned_node = HierarchicalNode(child_node.name, child_node.id, None) self.nodes_by_id[child_id] = orphaned_node # Remove the children mapping del self.children_by_parent[id_] logger.warning(f"Orphaned {len(child_ids)} children of node {id_}") # Remove from parent's children list if node_parent_id is not None and node_parent_id in self.children_by_parent: try: self.children_by_parent[node_parent_id].remove(id_) # Clean up empty lists if not self.children_by_parent[node_parent_id]: del self.children_by_parent[node_parent_id] except ValueError: logger.warning(f"Node {id_} not found in children list of parent {node_parent_id}") # Remove from name tracking if node_name.lower() in self.name_to_ids: try: self.name_to_ids[node_name.lower()].remove(id_) # Clean up empty lists if not self.name_to_ids[node_name.lower()]: del self.name_to_ids[node_name.lower()] except ValueError: logger.warning(f"Node {id_} not found in name_to_ids for '{node_name}'") # Remove from paths caches self.paths_cache.pop(id_, None) self.paths_name_cache.pop(id_, None) # Remove from primary storage del self.nodes_by_id[id_] # Mark indices as needing rebuild self._indices_built = False self._paths_computed = False logger.debug(f"Removed {self.entity_type} node {id_} ('{node_name}')") return True def _build_indices(self): """Build sorted indices for fast searching. Call after all nodes added.""" if self._indices_built: return # Build sorted name indices for binary search name_id_pairs = [(node.name, node.id) for node in self.nodes_by_id.values()] self.names_sorted = sorted(name_id_pairs, key=lambda x: x[0]) # Build lowercase sorted index with lowercase names name_lower_id_pairs = [(node.name.lower(), node.id) for node in self.nodes_by_id.values()] self.names_lower_sorted = sorted(name_lower_id_pairs, key=lambda x: x[0]) self._indices_built = True def _compute_paths(self): """Pre-compute all entity paths for O(1) access.""" if self._paths_computed: return self._build_indices() # Ensure indices are built first # Compute paths for all nodes for node_id, node in self.nodes_by_id.items(): # Build ID path id_path = self._build_id_path(node_id) self.paths_cache[node_id] = " -> ".join(str(id_) for id_ in id_path) # Build name path name_path = self._build_name_path(node_id) self.paths_name_cache[node_id] = " -> ".join(name_path) self._paths_computed = True def _build_id_path(self, node_id: Union[str, int]) -> List[Union[str, int]]: """Build path of IDs from root to given node.""" path = [] current_id = node_id visited = set() # Prevent infinite loops while current_id is not None and current_id in self.nodes_by_id: if current_id in visited: logger.warning(f"Circular reference detected in {self.entity_type} hierarchy at node {current_id}") break # Circular reference detected visited.add(current_id) path.append(current_id) current_id = self.nodes_by_id[current_id].parent_id return list(reversed(path)) def _build_name_path(self, node_id: Union[str, int]) -> List[str]: """Build path of names from root to given node.""" id_path = self._build_id_path(node_id) return [self.nodes_by_id[id_].name for id_ in id_path] def search_by_name_exact(self, name: str) -> List[HierarchicalNode]: """ Find nodes whose name exactly matches. Args: name: Exact name to search for Returns: List of matching HierarchicalNode objects """ self._build_indices() # Handle case-insensitive search matching_ids = self.name_to_ids.get(name.lower(), []) return [self.nodes_by_id[id_] for id_ in matching_ids] def search_by_name_prefix(self, prefix: str) -> List[HierarchicalNode]: """ Find nodes whose name starts with given prefix. Args: prefix: Name prefix to search for Returns: List of matching HierarchicalNode objects """ self._build_indices() prefix_lower = prefix.lower() results = [] # Binary search for first match start_idx = bisect.bisect_left(self.names_lower_sorted, (prefix_lower, "")) # Collect all matches for i in range(start_idx, len(self.names_lower_sorted)): name_lower, node_id = self.names_lower_sorted[i] if not name_lower.startswith(prefix_lower): break results.append(self.nodes_by_id[node_id]) return results def get_full_path(self, node_id: Union[str, int]) -> str: """ Return the full path (IDs) of a given node. Args: node_id: Node identifier Returns: Path string with IDs separated by " -> " """ self._compute_paths() return self.paths_cache.get(node_id, "") def get_full_path_name(self, identifier: Union[str, int]) -> Union[str, List[str]]: """ Return full path (names) given either node_id or name. Args: identifier: Either a node ID or an entity name Returns: Path string with names, or list of paths if multiple matches """ self._compute_paths() # Try as node ID first if identifier in self.nodes_by_id: return self.paths_name_cache.get(identifier, "") # Try as name matching_nodes = self.search_by_name_exact(str(identifier)) if not matching_nodes: return "" if len(matching_nodes) == 1: return self.paths_name_cache.get(matching_nodes[0].id, "") # Multiple matches - return all paths return [self.paths_name_cache.get(node.id, "") for node in matching_nodes] def get_children(self, node_id: Union[str, int]) -> List[HierarchicalNode]: """ Get direct children of a node. Args: node_id: Parent node identifier Returns: List of child HierarchicalNode objects """ child_ids = self.children_by_parent.get(node_id, []) return [self.nodes_by_id[child_id] for child_id in child_ids] def get_stats(self) -> Dict[str, Union[int, float]]: """ Get cache statistics. Returns: Dictionary with cache statistics """ total_nodes = len(self.nodes_by_id) root_nodes = sum(1 for node in self.nodes_by_id.values() if node.parent_id is None) # Calculate depth statistics depths = [] for node_id in self.nodes_by_id: depth = len(self._build_id_path(node_id)) - 1 depths.append(depth) return { 'entity_type': self.entity_type, 'total_nodes': total_nodes, 'root_nodes': root_nodes, 'max_depth': max(depths) if depths else 0, 'avg_depth': sum(depths) / len(depths) if depths else 0, 'indices_built': self._indices_built, 'paths_computed': self._paths_computed } # ------------------------------- # Elasticsearch Data Loaders # ------------------------------- class BaseEntityLoader: """ Base class for loading hierarchical entities from Elasticsearch. Handles common Elasticsearch operations and provides template for entity-specific field mapping. """ def __init__(self, tenant_id: str, entity_type: str, index_suffix: str): """ Initialize entity loader. Args: tenant_id: Tenant identifier for index naming entity_type: Type of entity (for logging) index_suffix: Index suffix (e.g., "location", "department") """ self.tenant_id = tenant_id self.entity_type = entity_type self.index_name = f"{tenant_id}_{index_suffix}" self.es_client = None def connect(self) -> bool: """Connect to Elasticsearch using existing client infrastructure.""" try: from elasticsearch_client import get_elasticsearch_client es_wrapper = get_elasticsearch_client() if es_wrapper.connect(): self.es_client = es_wrapper.get_client() logger.info(f"Connected to Elasticsearch for {self.entity_type} loading") return True else: logger.error(f"Failed to connect to Elasticsearch for {self.entity_type}") return False except ImportError as e: logger.error(f"elasticsearch_client module not found: {e}") return False except Exception as e: logger.error(f"Failed to connect to Elasticsearch: {e}", exc_info=True) return False def load_entities(self, cache: HierarchicalCache) -> bool: """ Load all entities from Elasticsearch into the cache. Args: cache: HierarchicalCache instance to populate Returns: True if successful, False otherwise """ if not self.es_client: logger.error(f"Not connected to Elasticsearch for {self.entity_type}") return False try: # Check if index exists if not self.es_client.indices.exists(index=self.index_name): logger.warning(f"Index '{self.index_name}' does not exist") return False # Get all entities using scroll API for large datasets query = { "query": {"match_all": {}}, "size": 1000 # Process in batches } response = self.es_client.search( index=self.index_name, body=query, scroll='2m' ) scroll_id = response['_scroll_id'] hits = response['hits']['hits'] total_loaded = 0 # Process first batch total_loaded += self._process_batch(cache, hits) # Process remaining batches while len(hits) > 0: response = self.es_client.scroll(scroll_id=scroll_id, scroll='2m') hits = response['hits']['hits'] total_loaded += self._process_batch(cache, hits) # Clear scroll self.es_client.clear_scroll(scroll_id=scroll_id) logger.info(f"✅ Loaded {total_loaded} {self.entity_type} entities from Elasticsearch") return True except Exception as e: logger.error(f"Error loading {self.entity_type} entities: {e}", exc_info=True) return False def _process_batch(self, cache: HierarchicalCache, hits: list) -> int: """ Process a batch of Elasticsearch hits. Must be implemented by subclasses to handle entity-specific field mapping. Args: cache: HierarchicalCache to populate hits: List of Elasticsearch hit documents Returns: Number of entities successfully loaded """ raise NotImplementedError("Subclasses must implement _process_batch") class LocationCacheLoader(BaseEntityLoader): """Loads location entities from Elasticsearch into HierarchicalCache.""" def __init__(self, tenant_id: str = "apolo"): """ Initialize location loader. Args: tenant_id: Tenant identifier for index naming """ super().__init__(tenant_id, "location", "location") def _process_batch(self, cache: HierarchicalCache, hits: list) -> int: """Process a batch of location documents.""" loaded_count = 0 for hit in hits: try: source = hit['_source'] # Extract location data based on ES schema location_id = source.get('dbid') location_name = source.get('location_name', '') parent_id = source.get('location_parentid') # Validate required fields if location_id is None or not location_name: logger.warning(f"Skipping location with missing required fields: {hit.get('_id', 'unknown')}") continue # Convert parent_id = 0 to None (root nodes) if parent_id == 0: parent_id = None # Add to cache cache.add_node(location_name, location_id, parent_id) loaded_count += 1 except Exception as e: logger.warning(f"Error processing location {hit.get('_id', 'unknown')}: {e}") continue return loaded_count class DepartmentCacheLoader(BaseEntityLoader): """Loads department entities from Elasticsearch into HierarchicalCache.""" def __init__(self, tenant_id: str = "apolo"): """ Initialize department loader. Args: tenant_id: Tenant identifier for index naming """ super().__init__(tenant_id, "department", "department") def _process_batch(self, cache: HierarchicalCache, hits: list) -> int: """Process a batch of department documents.""" loaded_count = 0 for hit in hits: try: source = hit['_source'] # Extract department data based on ES schema department_id = source.get('dbid') department_name = source.get('department_name', '') parent_id = source.get('department_parentid') # Validate required fields if department_id is None or not department_name: logger.warning(f"Skipping department with missing required fields: {hit.get('_id', 'unknown')}") continue # Convert parent_id = 0 to None (root nodes) if parent_id == 0: parent_id = None # Add to cache cache.add_node(department_name, department_id, parent_id) loaded_count += 1 except Exception as e: logger.warning(f"Error processing department {hit.get('_id', 'unknown')}: {e}") continue return loaded_count # ------------------------------- # Hierarchy Cache Manager # ------------------------------- class HierarchyCacheManager: """ Manages multiple hierarchical entity caches (locations, departments). Provides a unified interface for accessing and managing separate tree structures for different entity types. Implements singleton pattern to ensure single initialization. """ _instance: Optional['HierarchyCacheManager'] = None _initialized: bool = False def __init__(self): """Initialize cache manager (use get_instance() instead).""" self.location_cache: Optional[HierarchicalCache] = None self.department_cache: Optional[HierarchicalCache] = None self._tenant_id: Optional[str] = None @classmethod def get_instance(cls) -> 'HierarchyCacheManager': """ Get singleton instance of cache manager. Returns: HierarchyCacheManager instance """ if cls._instance is None: cls._instance = cls() return cls._instance def initialize(self, tenant_id: str) -> bool: """ Initialize both location and department caches. Args: tenant_id: Tenant identifier for index naming Returns: True if at least one cache was successfully initialized """ if self._initialized: logger.info("Hierarchy caches already initialized") return True self._tenant_id = tenant_id logger.info(f"Initializing hierarchy caches for tenant '{tenant_id}'...") location_success = self._initialize_location_cache(tenant_id) department_success = self._initialize_department_cache(tenant_id) # Mark as initialized if at least one succeeded if location_success or department_success: self._initialized = True logger.info("✅ Hierarchy cache initialization complete") self._log_statistics() return True else: logger.error("❌ Failed to initialize any hierarchy caches") return False def _initialize_location_cache(self, tenant_id: str) -> bool: """Initialize location cache.""" try: logger.info("Loading location hierarchy...") # Create cache self.location_cache = HierarchicalCache(entity_type="location") # Load data from Elasticsearch loader = LocationCacheLoader(tenant_id) if not loader.connect(): logger.error("Failed to connect to Elasticsearch for locations") return False if not loader.load_entities(self.location_cache): logger.error("Failed to load location data") return False # Build indices and compute paths self.location_cache._build_indices() self.location_cache._compute_paths() logger.info("✅ Location cache initialized successfully") return True except Exception as e: logger.error(f"Error initializing location cache: {e}", exc_info=True) self.location_cache = None return False def _initialize_department_cache(self, tenant_id: str) -> bool: """Initialize department cache.""" try: logger.info("Loading department hierarchy...") # Create cache self.department_cache = HierarchicalCache(entity_type="department") # Load data from Elasticsearch loader = DepartmentCacheLoader(tenant_id) if not loader.connect(): logger.error("Failed to connect to Elasticsearch for departments") return False if not loader.load_entities(self.department_cache): logger.error("Failed to load department data") return False # Build indices and compute paths self.department_cache._build_indices() self.department_cache._compute_paths() logger.info("✅ Department cache initialized successfully") return True except Exception as e: logger.error(f"Error initializing department cache: {e}", exc_info=True) self.department_cache = None return False def _log_statistics(self): """Log statistics for all initialized caches.""" if self.location_cache: stats = self.location_cache.get_stats() logger.info( f"📊 Location Cache: {stats['total_nodes']} nodes, " f"{stats['root_nodes']} roots, max depth {stats['max_depth']}" ) if self.department_cache: stats = self.department_cache.get_stats() logger.info( f"📊 Department Cache: {stats['total_nodes']} nodes, " f"{stats['root_nodes']} roots, max depth {stats['max_depth']}" ) def get_location_cache(self) -> Optional[HierarchicalCache]: """ Get location cache instance. Returns: HierarchicalCache for locations or None if not initialized """ return self.location_cache def get_department_cache(self) -> Optional[HierarchicalCache]: """ Get department cache instance. Returns: HierarchicalCache for departments or None if not initialized """ return self.department_cache def is_initialized(self) -> bool: """Check if caches have been initialized.""" return self._initialized def get_statistics(self) -> Dict[str, any]: """ Get statistics for all caches. Returns: Dictionary with statistics for each cache """ stats = { 'tenant_id': self._tenant_id, 'initialized': self._initialized, 'location': self.location_cache.get_stats() if self.location_cache else None, 'department': self.department_cache.get_stats() if self.department_cache else None, } return stats # ------------------------------- # Public API # ------------------------------- def initialize_hierarchy_caches(tenant_id: str = "apolo") -> Optional[HierarchyCacheManager]: """ Initialize hierarchical entity caches for locations and departments. This is the main entry point for setting up the cache system. It should be called once during server startup. Args: tenant_id: Tenant identifier for index naming (default: "apolo") Returns: HierarchyCacheManager instance if successful, None if completely failed Example: >>> # During server startup >>> cache_manager = initialize_hierarchy_caches("apolo") >>> if cache_manager: ... location_cache = cache_manager.get_location_cache() ... dept_cache = cache_manager.get_department_cache() """ try: logger.info("=" * 60) logger.info("🚀 Initializing Hierarchy Cache System") logger.info("=" * 60) # Get singleton manager instance manager = HierarchyCacheManager.get_instance() # Initialize caches success = manager.initialize(tenant_id) if success: logger.info("=" * 60) logger.info("✅ Hierarchy Cache System Ready") logger.info("=" * 60) return manager else: logger.warning("⚠️ Hierarchy cache initialization failed, but server can continue") return manager # Return manager even if initialization failed except Exception as e: logger.error(f"❌ Critical error during hierarchy cache initialization: {e}", exc_info=True) logger.warning("Server will continue without hierarchy caches") return None def get_hierarchy_cache_manager() -> Optional[HierarchyCacheManager]: """ Get the singleton hierarchy cache manager instance. Returns: HierarchyCacheManager instance or None if not initialized Example: >>> manager = get_hierarchy_cache_manager() >>> if manager and manager.is_initialized(): ... location_cache = manager.get_location_cache() """ return HierarchyCacheManager.get_instance() if HierarchyCacheManager._instance else None