"""Knowledge Graph builder for Myanmar Ghost project. Represents conversational context as a knowledge graph for better understanding of complex social interactions. Example: (Speaker, Role, Customer) --[located_in]--> (Restaurant) """ import json from dataclasses import dataclass, field from enum import Enum from pathlib import Path from typing import Any, Dict, List, Optional, Set, Tuple import networkx as nx class NodeType(str, Enum): """Types of nodes in the knowledge graph.""" SPEAKER = "speaker" UTTERANCE = "utterance" LOCATION = "location" ORGANIZATION = "organization" EMOTION = "emotion" TOPIC = "topic" ACTION = "action" TIME = "time" class RelationType(str, Enum): """Types of relations between nodes.""" SPEAKS = "speaks" LOCATED_IN = "located_in" WORKS_AT = "works_at" VISITS = "visits" FEELS = "feels" ABOUT = "about" BEFORE = "before" AFTER = "after" IN_RESPONSE_TO = "in_response_to" CONTAINS = "contains" HAS_ROLE = "has_role" @dataclass class Entity: """Represents an entity in the knowledge graph.""" id: str type: NodeType properties: Dict[str, Any] = field(default_factory=dict) aliases: List[str] = field(default_factory=list) def to_dict(self) -> Dict[str, Any]: return { "id": self.id, "type": self.type.value, "properties": self.properties, "aliases": self.aliases, } @dataclass class Relation: """Represents a relation between entities.""" source: str # Entity ID target: str # Entity ID type: RelationType properties: Dict[str, Any] = field(default_factory=dict) confidence: float = 1.0 def to_dict(self) -> Dict[str, Any]: return { "source": self.source, "target": self.target, "type": self.type.value, "properties": self.properties, "confidence": self.confidence, } class MyanmarKnowledgeGraph: """Build and manage knowledge graph for Myanmar conversations.""" # Common Myanmar entities LOCATIONS = { "စားသောက်ဆိုင်": NodeType.LOCATION, "ဆေးရုံ": NodeType.LOCATION, "ဈေး": NodeType.LOCATION, "ရုံး": NodeType.LOCATION, "အိမ်": NodeType.LOCATION, } EMOTIONS = { "ပျော်": NodeType.EMOTION, "စိတ်ဓာတ်ကျ": NodeType.EMOTION, "ဒေါသ": NodeType.EMOTION, "ဝမ်းနည်း": NodeType.EMOTION, "ပိုးပါး": NodeType.EMOTION, } ROLES = { "ဖေါ်သည်": "customer", "ဝန်ထမ်း": "staff", "ဆရာဝန်": "doctor", "ပါးရှင်း": "patient", "အရာရှိ": "manager", } def __init__(self): self.graph = nx.MultiDiGraph() self.entity_index: Dict[str, Entity] = {} self.session_id = 0 def add_entity(self, entity: Entity) -> None: """Add an entity to the graph.""" self.entity_index[entity.id] = entity self.graph.add_node( entity.id, type=entity.type.value, **entity.properties, ) def add_relation(self, relation: Relation) -> None: """Add a relation between entities.""" self.graph.add_edge( relation.source, relation.target, type=relation.type.value, **relation.properties, ) def extract_speaker_entity( self, speaker_id: str, role: Optional[str] = None, ) -> Entity: """Create a speaker entity from utterance metadata.""" entity = Entity( id=f"speaker_{speaker_id}", type=NodeType.SPEAKER, properties={ "role": role or "unknown", "session": self.session_id, }, ) self.add_entity(entity) return entity def extract_utterance_entity( self, text: str, speaker_id: str, timestamp: float, prosody: Optional[Dict] = None, ) -> Tuple[Entity, List[Entity], List[Relation]]: """Extract utterance and related entities from text.""" utterance_id = f"utt_{speaker_id}_{int(timestamp * 1000)}" utterance = Entity( id=utterance_id, type=NodeType.UTTERANCE, properties={ "text": text, "timestamp": timestamp, "prosody": prosody or {}, }, ) self.add_entity(utterance) # Extract related entities related_entities = [] relations = [] # Extract location mentions for loc, _ in self.LOCATIONS.items(): if loc in text: loc_entity = Entity( id=f"loc_{loc}_{self.session_id}", type=NodeType.LOCATION, properties={"name": loc}, ) self.add_entity(loc_entity) related_entities.append(loc_entity) relation = Relation( source=utterance_id, target=loc_entity.id, type=RelationType.LOCATED_IN, ) self.add_relation(relation) relations.append(relation) # Extract emotion mentions for emotion, _ in self.EMOTIONS.items(): if emotion in text: emotion_entity = Entity( id=f"emotion_{emotion}_{self.session_id}", type=NodeType.EMOTION, properties={"name": emotion}, ) self.add_entity(emotion_entity) related_entities.append(emotion_entity) relation = Relation( source=utterance_id, target=emotion_entity.id, type=RelationType.FEELS, ) self.add_relation(relation) relations.append(relation) # Link to speaker speaker_entity = self.entity_index.get(f"speaker_{speaker_id}") if speaker_entity: relation = Relation( source=speaker_entity.id, target=utterance_id, type=RelationType.SPEAKS, ) self.add_relation(relation) relations.append(relation) return utterance, related_entities, relations def build_from_conversation( self, utterances: List[Dict], context: Optional[Dict] = None, ) -> nx.MultiDiGraph: """Build knowledge graph from conversation data.""" self.session_id += 1 # Set context entities if context: for key, value in context.items(): if key == "location" and value in self.LOCATIONS: loc_entity = Entity( id=f"context_location", type=NodeType.LOCATION, properties={"name": value}, ) self.add_entity(loc_entity) prev_utterance = None for i, utt_data in enumerate(utterances): speaker_id = utt_data.get("speaker_id", f"s_{i}") text = utt_data.get("text", "") timestamp = utt_data.get("timestamp", i) prosody = utt_data.get("prosody") role = utt_data.get("role") # Add speaker self.extract_speaker_entity(speaker_id, role) # Add utterance utterance, related, _ = self.extract_utterance_entity( text, speaker_id, timestamp, prosody ) # Link to previous utterance (temporal relation) if prev_utterance: relation = Relation( source=prev_utterance.id, target=utterance.id, type=RelationType.BEFORE, ) self.add_relation(relation) # In response relation response_relation = Relation( source=utterance.id, target=prev_utterance.id, type=RelationType.IN_RESPONSE_TO, ) self.add_relation(response_relation) prev_utterance = utterance return self.graph def query_path( self, source_type: NodeType, target_type: NodeType, relation_type: Optional[RelationType] = None, ) -> List[Tuple[Entity, Entity, Relation]]: """Query paths between entity types.""" results = [] for source_id in self.entity_index: source = self.entity_index[source_id] if source.type != source_type: continue for target_id in self.entity_index: target = self.entity_index[target_id] if target.type != target_type: continue # Find paths try: if relation_type: edges = self.graph.get_edge_data(source_id, target_id) if edges: for edge_data in edges.values(): if edge_data.get("type") == relation_type.value: relation = Relation( source=source_id, target=target_id, type=relation_type, properties=edge_data, ) results.append((source, target, relation)) else: if nx.has_path(self.graph, source_id, target_id): path = nx.shortest_path( self.graph, source_id, target_id ) if len(path) == 2: relation = Relation( source=source_id, target=target_id, type=RelationType.CONTAINS, ) results.append((source, target, relation)) except nx.NetworkXError: continue return results def get_utterance_context(self, utterance_id: str) -> Dict: """Get full context for an utterance.""" if utterance_id not in self.entity_index: return {} context = { "utterance": self.entity_index[utterance_id].to_dict(), "speaker": None, "previous": None, "next": None, "locations": [], "emotions": [], } # Get speaker for edge in self.graph.out_edges(utterance_id, data=True): if edge[2].get("type") == RelationType.FEELS.value: context["emotions"].append(self.entity_index[edge[1]].to_dict()) if edge[2].get("type") == RelationType.LOCATED_IN.value: context["locations"].append(self.entity_index[edge[1]].to_dict()) # Get predecessor/successor predecessors = list(self.graph.predecessors(utterance_id)) successors = list(self.graph.successors(utterance_id)) for pred_id in predecessors: pred = self.entity_index.get(pred_id) if pred and pred.type == NodeType.UTTERANCE: context["previous"] = pred.to_dict() break for succ_id in successors: succ = self.entity_index.get(succ_id) if succ and succ.type == NodeType.UTTERANCE: context["next"] = succ.to_dict() break return context def export_to_json(self, path: str) -> None: """Export graph to JSON format.""" entities = [e.to_dict() for e in self.entity_index.values()] relations = [] for source, target, data in self.graph.edges(data=True): relations.append({ "source": source, "target": target, "type": data.get("type"), **data, }) output = { "entities": entities, "relations": relations, "metadata": { "num_entities": len(entities), "num_relations": len(relations), "session_id": self.session_id, }, } with open(path, "w", encoding="utf-8") as f: json.dump(output, f, indent=2, ensure_ascii=False) def load_from_json(self, path: str) -> None: """Load graph from JSON format.""" with open(path, "r", encoding="utf-8") as f: data = json.load(f) self.entity_index = {} self.graph = nx.MultiDiGraph() for entity_data in data.get("entities", []): entity = Entity( id=entity_data["id"], type=NodeType(entity_data["type"]), properties=entity_data.get("properties", {}), aliases=entity_data.get("aliases", []), ) self.add_entity(entity) for rel_data in data.get("relations", []): relation = Relation( source=rel_data["source"], target=rel_data["target"], type=RelationType(rel_data["type"]), properties=rel_data, confidence=rel_data.get("confidence", 1.0), ) self.add_relation(relation) def visualize(self) -> nx.MultiDiGraph: """Return the graph for visualization.""" return self.graph def create_knowledge_graph() -> MyanmarKnowledgeGraph: """Factory function to create knowledge graph.""" return MyanmarKnowledgeGraph() if __name__ == "__main__": # Example usage kg = create_knowledge_graph() # Sample conversation utterances = [ { "speaker_id": "customer_1", "text": "ဆိုင်သို့ ကျွန်ုပ်လာပါပြီ", "timestamp": 0, "role": "customer", }, { "speaker_id": "staff_1", "text": "ကြိုဆိုပါတယ်", "timestamp": 1, "role": "staff", }, { "speaker_id": "customer_1", "text": "ကျေးဇူးပါ", "timestamp": 2, "prosody": {"mean_pitch": 150, "speaking_rate": 3}, "role": "customer", }, ] context = {"location": "စားသောက်ဆိုင်"} kg.build_from_conversation(utterances, context) # Export kg.export_to_json("data/graph/conversation_graph.json") print(f"Graph exported with {len(kg.entity_index)} entities")