"""In-memory graph store for Headroom's knowledge graph memory system. Provides thread-safe in-memory storage for entities and relationships with efficient lookup via multiple indexes and BFS-based traversal. """ from __future__ import annotations from collections import deque from threading import RLock from typing import TYPE_CHECKING from .graph_models import Entity, Relationship, RelationshipDirection, Subgraph if TYPE_CHECKING: from ..tracker import ComponentStats class InMemoryGraphStore: """Thread-safe in-memory graph store implementing the GraphStore protocol. Provides storage for entities and relationships with efficient lookup via multiple indexes. All operations are thread-safe using RLock. Features: - O(1) entity and relationship lookup by ID - O(1) entity lookup by name (per user, case-insensitive) - O(1) lookup of relationships by source or target entity - BFS-based subgraph traversal with configurable hop limit - BFS-based shortest path finding between entities Indexes: - _entities: dict[str, Entity] - entity_id -> Entity - _relationships: dict[str, Relationship] - relationship_id -> Relationship - _entities_by_user: dict[str, set[str]] - user_id -> entity_ids - _entities_by_name: dict[str, dict[str, str]] - user_id -> {name_lower: entity_id} - _relationships_by_source: dict[str, set[str]] - entity_id -> relationship_ids (outgoing) - _relationships_by_target: dict[str, set[str]] - entity_id -> relationship_ids (incoming) Usage: store = InMemoryGraphStore() await store.add_entity(Entity(user_id="alice", name="Project X", entity_type="project")) entity = await store.get_entity_by_name("alice", "project x") # Case-insensitive subgraph = await store.query_subgraph(["entity-id"], max_hops=2) """ def __init__(self) -> None: """Initialize the in-memory graph store with empty indexes.""" # Primary storage self._entities: dict[str, Entity] = {} self._relationships: dict[str, Relationship] = {} # Secondary indexes for efficient lookup self._entities_by_user: dict[str, set[str]] = {} self._entities_by_name: dict[str, dict[str, str]] = {} # user_id -> {name_lower: entity_id} self._relationships_by_source: dict[str, set[str]] = {} # entity_id -> relationship_ids self._relationships_by_target: dict[str, set[str]] = {} # entity_id -> relationship_ids # Thread safety - use RLock to allow re-entrant locking self._lock = RLock() # ========================================================================= # Entity Operations # ========================================================================= async def add_entity(self, entity: Entity) -> None: """Add an entity to the graph store. If an entity with the same ID already exists, it will be replaced and indexes will be updated accordingly. Args: entity: The entity to add. """ with self._lock: # If entity already exists, remove from old indexes first if entity.id in self._entities: old_entity = self._entities[entity.id] self._remove_entity_from_indexes(old_entity) # Store the entity self._entities[entity.id] = entity # Update user index if entity.user_id not in self._entities_by_user: self._entities_by_user[entity.user_id] = set() self._entities_by_user[entity.user_id].add(entity.id) # Update name index (case-insensitive) if entity.user_id not in self._entities_by_name: self._entities_by_name[entity.user_id] = {} name_lower = entity.name.lower() self._entities_by_name[entity.user_id][name_lower] = entity.id async def get_entity(self, entity_id: str) -> Entity | None: """Retrieve an entity by ID. Args: entity_id: The unique identifier of the entity. Returns: The entity if found, None otherwise. """ with self._lock: return self._entities.get(entity_id) async def get_entity_by_name(self, user_id: str, name: str) -> Entity | None: """Retrieve an entity by name (case-insensitive). Args: user_id: The user who owns the entity. name: The name of the entity (case-insensitive). Returns: The entity if found, None otherwise. """ with self._lock: user_names = self._entities_by_name.get(user_id) if user_names is None: return None entity_id = user_names.get(name.lower()) if entity_id is None: return None return self._entities.get(entity_id) async def delete_entity(self, entity_id: str) -> bool: """Delete an entity and all its relationships. Removes the entity from storage and all indexes, and deletes all relationships where this entity is source or target. Args: entity_id: The unique identifier of the entity. Returns: True if the entity was deleted, False if not found. """ with self._lock: entity = self._entities.get(entity_id) if entity is None: return False # Collect all relationships to delete relationships_to_delete: set[str] = set() # Outgoing relationships outgoing = self._relationships_by_source.get(entity_id, set()) relationships_to_delete.update(outgoing) # Incoming relationships incoming = self._relationships_by_target.get(entity_id, set()) relationships_to_delete.update(incoming) # Delete relationships for rel_id in relationships_to_delete: self._remove_relationship_internal(rel_id) # Remove entity from indexes self._remove_entity_from_indexes(entity) # Remove from primary storage del self._entities[entity_id] return True def _remove_entity_from_indexes(self, entity: Entity) -> None: """Remove an entity from secondary indexes (internal, must hold lock).""" # Remove from user index user_entities = self._entities_by_user.get(entity.user_id) if user_entities is not None: user_entities.discard(entity.id) if not user_entities: del self._entities_by_user[entity.user_id] # Remove from name index user_names = self._entities_by_name.get(entity.user_id) if user_names is not None: name_lower = entity.name.lower() if user_names.get(name_lower) == entity.id: del user_names[name_lower] if not user_names: del self._entities_by_name[entity.user_id] # ========================================================================= # Relationship Operations # ========================================================================= async def add_relationship(self, relationship: Relationship) -> None: """Add a relationship to the graph store. If a relationship with the same ID already exists, it will be replaced and indexes will be updated accordingly. Args: relationship: The relationship to add. """ with self._lock: # If relationship already exists, remove from old indexes first if relationship.id in self._relationships: old_rel = self._relationships[relationship.id] self._remove_relationship_from_indexes(old_rel) # Store the relationship self._relationships[relationship.id] = relationship # Update source index if relationship.source_id not in self._relationships_by_source: self._relationships_by_source[relationship.source_id] = set() self._relationships_by_source[relationship.source_id].add(relationship.id) # Update target index if relationship.target_id not in self._relationships_by_target: self._relationships_by_target[relationship.target_id] = set() self._relationships_by_target[relationship.target_id].add(relationship.id) async def get_relationships( self, entity_id: str, direction: RelationshipDirection = RelationshipDirection.BOTH, relation_type: str | None = None, ) -> list[Relationship]: """Get relationships for an entity. Args: entity_id: The entity ID to get relationships for. direction: Whether to get outgoing, incoming, or both relationships. relation_type: Optional filter for relationship type. Returns: List of matching relationships. """ with self._lock: rel_ids: set[str] = set() if direction in (RelationshipDirection.OUTGOING, RelationshipDirection.BOTH): rel_ids.update(self._relationships_by_source.get(entity_id, set())) if direction in (RelationshipDirection.INCOMING, RelationshipDirection.BOTH): rel_ids.update(self._relationships_by_target.get(entity_id, set())) relationships = [] for rel_id in rel_ids: rel = self._relationships.get(rel_id) if rel is not None: if relation_type is None or rel.relation_type == relation_type: relationships.append(rel) return relationships async def delete_relationship(self, relationship_id: str) -> bool: """Delete a single relationship. Args: relationship_id: The unique identifier of the relationship. Returns: True if the relationship was deleted, False if not found. """ with self._lock: return self._remove_relationship_internal(relationship_id) def _remove_relationship_internal(self, relationship_id: str) -> bool: """Remove a relationship (internal, must hold lock). Returns: True if removed, False if not found. """ rel = self._relationships.get(relationship_id) if rel is None: return False # Remove from indexes self._remove_relationship_from_indexes(rel) # Remove from primary storage del self._relationships[relationship_id] return True def _remove_relationship_from_indexes(self, relationship: Relationship) -> None: """Remove a relationship from secondary indexes (internal, must hold lock).""" # Remove from source index source_rels = self._relationships_by_source.get(relationship.source_id) if source_rels is not None: source_rels.discard(relationship.id) if not source_rels: del self._relationships_by_source[relationship.source_id] # Remove from target index target_rels = self._relationships_by_target.get(relationship.target_id) if target_rels is not None: target_rels.discard(relationship.id) if not target_rels: del self._relationships_by_target[relationship.target_id] # ========================================================================= # Graph Traversal Operations # ========================================================================= async def query_subgraph( self, entity_ids: list[str], max_hops: int = 2, direction: RelationshipDirection = RelationshipDirection.BOTH, relation_types: list[str] | None = None, ) -> Subgraph: """Query a subgraph starting from given entities using BFS traversal. Performs a breadth-first traversal from the starting entities, collecting all entities and relationships within the specified number of hops. Args: entity_ids: Starting entity IDs for the traversal. max_hops: Maximum number of hops from starting entities (default 2). direction: Direction of relationship traversal. relation_types: Optional filter for relationship types. Returns: Subgraph containing all reachable entities and relationships. """ with self._lock: collected_entities: dict[str, Entity] = {} collected_relationships: dict[str, Relationship] = {} # BFS queue: (entity_id, current_depth) queue: deque[tuple[str, int]] = deque() visited: set[str] = set() # Initialize queue with starting entities for entity_id in entity_ids: if entity_id in self._entities: queue.append((entity_id, 0)) visited.add(entity_id) collected_entities[entity_id] = self._entities[entity_id] # BFS traversal while queue: current_id, depth = queue.popleft() if depth >= max_hops: continue # Get relationships based on direction rel_ids: set[str] = set() if direction in (RelationshipDirection.OUTGOING, RelationshipDirection.BOTH): rel_ids.update(self._relationships_by_source.get(current_id, set())) if direction in (RelationshipDirection.INCOMING, RelationshipDirection.BOTH): rel_ids.update(self._relationships_by_target.get(current_id, set())) # Process relationships for rel_id in rel_ids: rel = self._relationships.get(rel_id) if rel is None: continue # Filter by relation type if specified if relation_types is not None and rel.relation_type not in relation_types: continue # Add relationship collected_relationships[rel_id] = rel # Determine neighbor based on direction neighbor_id = None if direction == RelationshipDirection.OUTGOING: if rel.source_id == current_id: neighbor_id = rel.target_id elif direction == RelationshipDirection.INCOMING: if rel.target_id == current_id: neighbor_id = rel.source_id else: # BOTH if rel.source_id == current_id: neighbor_id = rel.target_id elif rel.target_id == current_id: neighbor_id = rel.source_id if neighbor_id is not None and neighbor_id not in visited: neighbor = self._entities.get(neighbor_id) if neighbor is not None: visited.add(neighbor_id) collected_entities[neighbor_id] = neighbor queue.append((neighbor_id, depth + 1)) return Subgraph( entities=list(collected_entities.values()), relationships=list(collected_relationships.values()), root_entity_ids=entity_ids, ) async def find_path( self, source_id: str, target_id: str, max_depth: int = 10, direction: RelationshipDirection = RelationshipDirection.BOTH, ) -> list[str] | None: """Find the shortest path between two entities using BFS. Args: source_id: Starting entity ID. target_id: Target entity ID. max_depth: Maximum path length to search (default 10). direction: Direction of relationship traversal. Returns: List of entity IDs representing the path (including source and target), or None if no path exists within max_depth. """ with self._lock: # Edge cases if source_id == target_id: return [source_id] if source_id in self._entities else None if source_id not in self._entities or target_id not in self._entities: return None # BFS with path tracking # Queue: (current_entity_id, path_so_far) queue: deque[tuple[str, list[str]]] = deque() visited: set[str] = set() queue.append((source_id, [source_id])) visited.add(source_id) while queue: current_id, path = queue.popleft() if len(path) > max_depth: continue # Get neighbors rel_ids: set[str] = set() if direction in (RelationshipDirection.OUTGOING, RelationshipDirection.BOTH): rel_ids.update(self._relationships_by_source.get(current_id, set())) if direction in (RelationshipDirection.INCOMING, RelationshipDirection.BOTH): rel_ids.update(self._relationships_by_target.get(current_id, set())) for rel_id in rel_ids: rel = self._relationships.get(rel_id) if rel is None: continue # Determine neighbor based on direction neighbor_id = None if direction == RelationshipDirection.OUTGOING: if rel.source_id == current_id: neighbor_id = rel.target_id elif direction == RelationshipDirection.INCOMING: if rel.target_id == current_id: neighbor_id = rel.source_id else: # BOTH if rel.source_id == current_id: neighbor_id = rel.target_id elif rel.target_id == current_id: neighbor_id = rel.source_id if neighbor_id is None or neighbor_id in visited: continue # Check if we found the target new_path = path + [neighbor_id] if neighbor_id == target_id: return new_path # Continue BFS if within depth limit if len(new_path) <= max_depth: visited.add(neighbor_id) queue.append((neighbor_id, new_path)) return None # ========================================================================= # User Management Operations # ========================================================================= async def clear_user(self, user_id: str) -> tuple[int, int]: """Clear all entities and relationships for a user. Args: user_id: The user ID to clear data for. Returns: Tuple of (entities_deleted, relationships_deleted). """ with self._lock: entities_deleted = 0 relationships_deleted = 0 # Get all entity IDs for this user entity_ids = list(self._entities_by_user.get(user_id, set())) # Delete each entity (which also deletes its relationships) for entity_id in entity_ids: entity = self._entities.get(entity_id) if entity is None: continue # Count and delete relationships outgoing = self._relationships_by_source.get(entity_id, set()) incoming = self._relationships_by_target.get(entity_id, set()) rel_ids_to_delete = outgoing | incoming for rel_id in list(rel_ids_to_delete): if self._remove_relationship_internal(rel_id): relationships_deleted += 1 # Remove entity from indexes self._remove_entity_from_indexes(entity) # Remove from primary storage del self._entities[entity_id] entities_deleted += 1 return entities_deleted, relationships_deleted # ========================================================================= # Utility Methods # ========================================================================= async def get_entities_for_user(self, user_id: str) -> list[Entity]: """Get all entities for a user. Args: user_id: The user ID to get entities for. Returns: List of all entities belonging to the user. """ with self._lock: entity_ids = self._entities_by_user.get(user_id, set()) return [self._entities[eid] for eid in entity_ids if eid in self._entities] async def clear(self) -> None: """Clear all data from the store.""" with self._lock: self._entities.clear() self._relationships.clear() self._entities_by_user.clear() self._entities_by_name.clear() self._relationships_by_source.clear() self._relationships_by_target.clear() @property def entity_count(self) -> int: """Get the total number of entities.""" with self._lock: return len(self._entities) @property def relationship_count(self) -> int: """Get the total number of relationships.""" with self._lock: return len(self._relationships) def stats(self) -> dict: """Get store statistics. Returns: Dict with counts and index sizes. """ with self._lock: return { "entity_count": len(self._entities), "relationship_count": len(self._relationships), "users_count": len(self._entities_by_user), "source_index_size": len(self._relationships_by_source), "target_index_size": len(self._relationships_by_target), } def get_memory_stats(self) -> ComponentStats: """Get memory statistics for the MemoryTracker. Returns: ComponentStats with current memory usage. """ import sys from ..tracker import ComponentStats with self._lock: # Calculate size of all data structures size_bytes = 0 # Entities size_bytes += sys.getsizeof(self._entities) for entity_id, entity in self._entities.items(): size_bytes += len(entity_id) size_bytes += sys.getsizeof(entity) size_bytes += len(entity.id) + len(entity.user_id) + len(entity.name) size_bytes += len(entity.entity_type) if entity.properties: size_bytes += sys.getsizeof(entity.properties) # Relationships size_bytes += sys.getsizeof(self._relationships) for rel_id, rel in self._relationships.items(): size_bytes += len(rel_id) size_bytes += sys.getsizeof(rel) size_bytes += len(rel.id) + len(rel.source_id) + len(rel.target_id) size_bytes += len(rel.relation_type) if rel.properties: size_bytes += sys.getsizeof(rel.properties) # Indexes size_bytes += sys.getsizeof(self._entities_by_user) for user_id, entity_ids in self._entities_by_user.items(): size_bytes += len(user_id) size_bytes += sys.getsizeof(entity_ids) size_bytes += sys.getsizeof(self._entities_by_name) for user_id, name_map in self._entities_by_name.items(): size_bytes += len(user_id) size_bytes += sys.getsizeof(name_map) size_bytes += sys.getsizeof(self._relationships_by_source) size_bytes += sys.getsizeof(self._relationships_by_target) entry_count = len(self._entities) + len(self._relationships) return ComponentStats( name="graph_store", entry_count=entry_count, size_bytes=size_bytes, budget_bytes=None, hits=0, misses=0, evictions=0, )