"""Navigation and organization layer for the knowledge plane. Provides topic clustering, tag extraction, and a navigation index for easy orientation in the knowledgebase. """ from __future__ import annotations from collections import defaultdict from dataclasses import dataclass, field from cyberbrain26.knowledge_plane.contracts import KnowledgeDocument @dataclass(frozen=True) class TopicCluster: """A cluster of related documents.""" cluster_id: str label: str document_ids: list[str] = field(default_factory=list) centroid_vector: list[float] | None = None tags: list[str] = field(default_factory=list) @dataclass(frozen=True) class NavigationEntry: """A navigation entry linking a topic to its documents.""" topic: str document_ids: list[str] chunk_count: int last_updated: float @dataclass(frozen=True) class EntityEdge: """An edge in the entity relationship graph. Represents a relationship between two entities in documents. Format: (source_entity, target_entity, relationship_type, document_id) """ source_entity: str target_entity: str relationship_type: str # "mentioned_with", "co_reference", "part_of", etc. document_id: str @dataclass(frozen=True) class KnowledgeGraph: """A knowledge graph built from document entity relationships. Entities: entity name -> set of related entity names Edges: (source_entity, target_entity) -> list of EntityEdge Entity documents: entity name -> set of document IDs """ entities: dict[str, set[str]] = field(default_factory=dict) edges: dict[str, list[EntityEdge]] = field(default_factory=dict) entity_documents: dict[str, set[str]] = field(default_factory=dict) class KnowledgeNavigator: """Organizes and navigates the knowledgebase. Provides: - Topic clustering based on document metadata - Tag-based navigation - Full-text search index - Temporal versioning view """ def __init__(self) -> None: self._topics: dict[str, list[str]] = defaultdict(list) self._tags: dict[str, list[str]] = defaultdict(list) self._clusters: dict[str, TopicCluster] = {} # Entity relationship graph (MemGraphRAG pattern) self._entity_graph: KnowledgeGraph = KnowledgeGraph() def add_document(self, document: KnowledgeDocument) -> None: """Index a document for navigation.""" # Extract topic from metadata topic = document.metadata.get("topic", document.metadata.get("classification_topic", "uncategorized")) self._topics[topic].append(document.document_id) # Extract tags from metadata tags = document.metadata.get("tags", document.metadata.get("classification_tags", [])) if isinstance(tags, str): tags = [t.strip() for t in tags.split(",")] for tag in tags: self._tags[tag].append(document.document_id) # Store entity information for graph-based retrieval (MemGraphRAG pattern) entities = document.metadata.get("classification_entities", document.metadata.get("entities", [])) if isinstance(entities, str): entities = [e.strip() for e in entities.split(",")] if isinstance(entities, list): for entity in entities: if entity not in self._tags: self._tags[entity] = [] if document.document_id not in self._tags[entity]: self._tags[entity].append(document.document_id) # Add entity edges for co-occurrence relationships (MemGraphRAG) # All entities in the same document share an edge with a default type if len(entities) > 1: for i, entity1 in enumerate(entities): for entity2 in entities[i + 1:]: self.add_entity_edge(entity1, entity2, "co_reference", document.document_id) def get_topics(self) -> list[str]: """Return all known topics.""" return sorted(self._topics.keys()) def get_documents_by_topic(self, topic: str) -> list[str]: """Return document IDs for a topic.""" return self._topics.get(topic, []) def get_documents_by_tag(self, tag: str) -> list[str]: """Return document IDs for a tag.""" return self._tags.get(tag, []) def get_all_tags(self) -> list[str]: """Return all known tags.""" return sorted(self._tags.keys()) def get_navigation_index(self) -> list[NavigationEntry]: """Build a navigation index of all topics.""" entries = [] for topic, doc_ids in self._topics.items(): entries.append( NavigationEntry( topic=topic, document_ids=doc_ids, chunk_count=len(doc_ids), last_updated=0.0, ) ) return sorted(entries, key=lambda e: e.topic) def cluster_by_source(self, documents: list[KnowledgeDocument]) -> dict[str, list[str]]: """Group documents by their source URI host.""" clusters: dict[str, list[str]] = defaultdict(list) for doc in documents: host = "unknown" if "://" in doc.source_uri: host = doc.source_uri.split("://")[1].split("/")[0] clusters[host].append(doc.document_id) return dict(clusters) def search_by_keyword(self, keyword: str, documents: list[KnowledgeDocument]) -> list[str]: """Simple keyword search over document metadata.""" keyword_lower = keyword.lower() results = [] for doc in documents: # Search in metadata values for value in doc.metadata.values(): if isinstance(value, str) and keyword_lower in value.lower(): results.append(doc.document_id) break # Search in source_uri if keyword_lower in doc.source_uri.lower() and doc.document_id not in results: results.append(doc.document_id) return results def check_temporal_validity( self, document: KnowledgeDocument, current_time: float | None = None ) -> bool: """MemStrata pattern: deterministic staleness check. Returns True if document is currently valid, False if superseded. Uses (subject, relation, object) supersedence rule — no similarity threshold, no LLM call for stale fact retirement. """ if not document.superseded_by: return True return False def get_valid_documents( self, documents: list[KnowledgeDocument] ) -> list[KnowledgeDocument]: """Filter out superseded documents using temporal validity.""" return [d for d in documents if self.check_temporal_validity(d)] def supersede_document( self, document: KnowledgeDocument, superseded_by: str, current_time: float | None = None, ) -> KnowledgeDocument: """Mark a document as superseded (MemStrata bi-temporal ledger). Sets valid_to to current time and records the superseding document ID. """ from dataclasses import replace if current_time is None: import time current_time = time.time() return replace( document, valid_to=current_time, superseded_by=superseded_by, ) # ------------------------------------------------------------------ # Entity relationship graph (MemGraphRAG pattern: memory-aware # hierarchical retrieval through entity co-occurrence) # ------------------------------------------------------------------ def add_entity_edge( self, source_entity: str, target_entity: str, relationship_type: str, document_id: str, ) -> None: """Add an edge to the entity relationship graph. Tracks co-occurrence and relationships between entities extracted from document classification metadata. """ # Initialize entity sets if not present if source_entity not in self._entity_graph.entities: self._entity_graph.entities[source_entity] = set() if target_entity not in self._entity_graph.entities: self._entity_graph.entities[target_entity] = set() if source_entity not in self._entity_graph.entity_documents: self._entity_graph.entity_documents[source_entity] = set() if target_entity not in self._entity_graph.entity_documents: self._entity_graph.entity_documents[target_entity] = set() self._entity_graph.entities[source_entity].add(target_entity) self._entity_graph.entities[target_entity].add(source_entity) edge = EntityEdge( source_entity=source_entity, target_entity=target_entity, relationship_type=relationship_type, document_id=document_id, ) edge_key = f"{source_entity}->{target_entity}" if edge_key not in self._entity_graph.edges: self._entity_graph.edges[edge_key] = [] self._entity_graph.edges[edge_key].append(edge) self._entity_graph.entity_documents[source_entity].add(document_id) self._entity_graph.entity_documents[target_entity].add(document_id) def get_entity_neighbors(self, entity: str) -> list[str]: """Return entities directly related to the given entity.""" return sorted(self._entity_graph.entities.get(entity, set())) def get_entity_documents(self, entity: str) -> list[str]: """Return document IDs containing the given entity.""" return sorted(self._entity_graph.entity_documents.get(entity, set())) def get_entity_relationships( self, entity: str, relationship_type: str | None = None ) -> list[EntityEdge]: """Return all relationship edges involving the given entity. Optionally filter by relationship type. """ edges: list[EntityEdge] = [] for key, edge_list in self._entity_graph.edges.items(): if key.startswith(f"{entity}->") or key.endswith(f"->{entity}"): for edge in edge_list: if relationship_type is None or edge.relationship_type == relationship_type: edges.append(edge) return edges def get_related_entities( self, entity: str, depth: int = 1, relationship_type: str | None = None, ) -> dict[int, list[str]]: """Return entities related to the given entity up to a traversal depth. Performs breadth-first traversal over the entity relationship graph. Returns a dict mapping distance -> list of entity names. depth=1 returns immediate neighbors (distance 1). depth=2 returns neighbors + their neighbors (distance 1 and 2). """ if depth < 1: return {} visited: set[str] = {entity} frontier: list[str] = [entity] result: dict[int, list[str]] = {} current_depth = 1 while frontier and current_depth <= depth: next_frontier: list[str] = [] for current_entity in frontier: neighbors = self.get_entity_neighbors(current_entity) for neighbor in neighbors: if neighbor not in visited: visited.add(neighbor) next_frontier.append(neighbor) result.setdefault(current_depth, []).append(neighbor) frontier = next_frontier current_depth += 1 # If no neighbors found, return empty dict (entity has no relationships) return result __all__ = ["EntityEdge", "KnowledgeGraph", "KnowledgeNavigator", "NavigationEntry", "TopicCluster"]