""" Knowledge Graph for Breakpilot Agents Provides entity and relationship management: - Entity storage with properties - Relationship definitions - Graph traversal - Optional Qdrant integration for semantic search """ from typing import Dict, Any, List, Optional, Set, Tuple from datetime import datetime, timezone import json import logging from agent_core.brain.knowledge_models import ( EntityType, RelationshipType, Entity, Relationship, ) logger = logging.getLogger(__name__) class KnowledgeGraph: """ Knowledge graph for managing entity relationships. Provides: - Entity CRUD operations - Relationship management - Graph traversal (neighbors, paths) - Optional vector search via Qdrant """ def __init__( self, db_pool=None, qdrant_client=None, namespace: str = "breakpilot" ): """ Initialize the knowledge graph. Args: db_pool: Async PostgreSQL connection pool qdrant_client: Optional Qdrant client for vector search namespace: Namespace for isolation """ self.db_pool = db_pool self.qdrant = qdrant_client self.namespace = namespace self._entities: Dict[str, Entity] = {} self._relationships: Dict[str, Relationship] = {} self._adjacency: Dict[str, Set[str]] = {} # entity_id -> set of relationship_ids # Entity Operations def add_entity( self, entity_id: str, entity_type: EntityType, name: str, properties: Optional[Dict[str, Any]] = None ) -> Entity: """ Adds an entity to the graph. Args: entity_id: Unique entity identifier entity_type: Type of entity name: Human-readable name properties: Entity properties Returns: The created Entity """ entity = Entity( id=entity_id, entity_type=entity_type, name=name, properties=properties or {} ) self._entities[entity_id] = entity self._adjacency[entity_id] = set() logger.debug(f"Added entity: {entity_type.value}/{entity_id}") return entity def get_entity(self, entity_id: str) -> Optional[Entity]: """Gets an entity by ID""" return self._entities.get(entity_id) def update_entity( self, entity_id: str, name: Optional[str] = None, properties: Optional[Dict[str, Any]] = None ) -> Optional[Entity]: """ Updates an entity. Args: entity_id: Entity to update name: New name (optional) properties: Properties to update (merged) Returns: Updated entity or None if not found """ entity = self._entities.get(entity_id) if not entity: return None if name: entity.name = name if properties: entity.properties.update(properties) entity.updated_at = datetime.now(timezone.utc) return entity def delete_entity(self, entity_id: str) -> bool: """ Deletes an entity and its relationships. Args: entity_id: Entity to delete Returns: True if deleted """ if entity_id not in self._entities: return False # Delete all relationships involving this entity rel_ids = list(self._adjacency.get(entity_id, set())) for rel_id in rel_ids: self._delete_relationship_internal(rel_id) del self._entities[entity_id] del self._adjacency[entity_id] return True def get_entities_by_type( self, entity_type: EntityType ) -> List[Entity]: """Gets all entities of a specific type""" return [ e for e in self._entities.values() if e.entity_type == entity_type ] def search_entities( self, query: str, entity_type: Optional[EntityType] = None, limit: int = 10 ) -> List[Entity]: """ Searches entities by name. Args: query: Search query (case-insensitive substring) entity_type: Optional type filter limit: Maximum results Returns: Matching entities """ query_lower = query.lower() results = [] for entity in self._entities.values(): if entity_type and entity.entity_type != entity_type: continue if query_lower in entity.name.lower(): results.append(entity) if len(results) >= limit: break return results # Relationship Operations def add_relationship( self, relationship_id: str, source_id: str, target_id: str, relationship_type: RelationshipType, properties: Optional[Dict[str, Any]] = None, weight: float = 1.0 ) -> Optional[Relationship]: """ Adds a relationship between two entities. Args: relationship_id: Unique relationship identifier source_id: Source entity ID target_id: Target entity ID relationship_type: Type of relationship properties: Relationship properties weight: Relationship weight/strength Returns: The created Relationship or None if entities don't exist """ if source_id not in self._entities or target_id not in self._entities: logger.warning( f"Cannot create relationship: entity not found " f"(source={source_id}, target={target_id})" ) return None relationship = Relationship( id=relationship_id, source_id=source_id, target_id=target_id, relationship_type=relationship_type, properties=properties or {}, weight=weight ) self._relationships[relationship_id] = relationship self._adjacency[source_id].add(relationship_id) self._adjacency[target_id].add(relationship_id) logger.debug( f"Added relationship: {source_id} -[{relationship_type.value}]-> {target_id}" ) return relationship def get_relationship(self, relationship_id: str) -> Optional[Relationship]: """Gets a relationship by ID""" return self._relationships.get(relationship_id) def delete_relationship(self, relationship_id: str) -> bool: """Deletes a relationship""" return self._delete_relationship_internal(relationship_id) def _delete_relationship_internal(self, relationship_id: str) -> bool: """Internal relationship deletion""" relationship = self._relationships.get(relationship_id) if not relationship: return False # Remove from adjacency lists if relationship.source_id in self._adjacency: self._adjacency[relationship.source_id].discard(relationship_id) if relationship.target_id in self._adjacency: self._adjacency[relationship.target_id].discard(relationship_id) del self._relationships[relationship_id] return True # Graph Traversal def get_neighbors( self, entity_id: str, relationship_type: Optional[RelationshipType] = None, direction: str = "both" # "outgoing", "incoming", "both" ) -> List[Tuple[Entity, Relationship]]: """ Gets neighboring entities. Args: entity_id: Starting entity relationship_type: Optional filter by relationship type direction: Direction to traverse Returns: List of (entity, relationship) tuples """ if entity_id not in self._entities: return [] results = [] rel_ids = self._adjacency.get(entity_id, set()) for rel_id in rel_ids: rel = self._relationships.get(rel_id) if not rel: continue # Filter by relationship type if relationship_type and rel.relationship_type != relationship_type: continue # Determine neighbor based on direction neighbor_id = None if direction == "outgoing" and rel.source_id == entity_id: neighbor_id = rel.target_id elif direction == "incoming" and rel.target_id == entity_id: neighbor_id = rel.source_id elif direction == "both": neighbor_id = rel.target_id if rel.source_id == entity_id else rel.source_id if neighbor_id: neighbor = self._entities.get(neighbor_id) if neighbor: results.append((neighbor, rel)) return results def get_path( self, source_id: str, target_id: str, max_depth: int = 5 ) -> Optional[List[Tuple[Entity, Optional[Relationship]]]]: """ Finds a path between two entities using BFS. Args: source_id: Starting entity target_id: Target entity max_depth: Maximum path length Returns: Path as list of (entity, relationship) tuples, or None if no path """ if source_id not in self._entities or target_id not in self._entities: return None if source_id == target_id: return [(self._entities[source_id], None)] # BFS from collections import deque visited = {source_id} # Queue items: (entity_id, path so far) queue = deque([(source_id, [(self._entities[source_id], None)])]) while queue: current_id, path = queue.popleft() if len(path) > max_depth: continue for neighbor, rel in self.get_neighbors(current_id): if neighbor.id == target_id: return path + [(neighbor, rel)] if neighbor.id not in visited: visited.add(neighbor.id) queue.append((neighbor.id, path + [(neighbor, rel)])) return None def get_subgraph( self, entity_id: str, depth: int = 2 ) -> Tuple[List[Entity], List[Relationship]]: """ Gets a subgraph around an entity. Args: entity_id: Center entity depth: How many hops to include Returns: Tuple of (entities, relationships) """ if entity_id not in self._entities: return [], [] entities_set: Set[str] = {entity_id} relationships_set: Set[str] = set() frontier: Set[str] = {entity_id} for _ in range(depth): next_frontier: Set[str] = set() for e_id in frontier: for neighbor, rel in self.get_neighbors(e_id): if neighbor.id not in entities_set: entities_set.add(neighbor.id) next_frontier.add(neighbor.id) relationships_set.add(rel.id) frontier = next_frontier entities = [self._entities[e_id] for e_id in entities_set] relationships = [self._relationships[r_id] for r_id in relationships_set] return entities, relationships # Serialization def to_dict(self) -> Dict[str, Any]: """Serializes the graph to a dictionary""" return { "entities": [e.to_dict() for e in self._entities.values()], "relationships": [r.to_dict() for r in self._relationships.values()] } @classmethod def from_dict(cls, data: Dict[str, Any], **kwargs) -> "KnowledgeGraph": """Deserializes a graph from a dictionary""" graph = cls(**kwargs) # Load entities first for e_data in data.get("entities", []): entity = Entity.from_dict(e_data) graph._entities[entity.id] = entity graph._adjacency[entity.id] = set() # Load relationships for r_data in data.get("relationships", []): rel = Relationship.from_dict(r_data) graph._relationships[rel.id] = rel if rel.source_id in graph._adjacency: graph._adjacency[rel.source_id].add(rel.id) if rel.target_id in graph._adjacency: graph._adjacency[rel.target_id].add(rel.id) return graph def export_json(self) -> str: """Exports graph to JSON string""" return json.dumps(self.to_dict(), indent=2) @classmethod def import_json(cls, json_str: str, **kwargs) -> "KnowledgeGraph": """Imports graph from JSON string""" return cls.from_dict(json.loads(json_str), **kwargs) # Statistics def get_statistics(self) -> Dict[str, Any]: """Gets graph statistics""" entity_types: Dict[str, int] = {} for entity in self._entities.values(): entity_types[entity.entity_type.value] = \ entity_types.get(entity.entity_type.value, 0) + 1 rel_types: Dict[str, int] = {} for rel in self._relationships.values(): rel_types[rel.relationship_type.value] = \ rel_types.get(rel.relationship_type.value, 0) + 1 return { "total_entities": len(self._entities), "total_relationships": len(self._relationships), "entity_types": entity_types, "relationship_types": rel_types, "avg_connections": ( sum(len(adj) for adj in self._adjacency.values()) / max(len(self._adjacency), 1) ) } @property def entity_count(self) -> int: """Returns number of entities""" return len(self._entities) @property def relationship_count(self) -> int: """Returns number of relationships""" return len(self._relationships)