tog_02.py

class KnowledgeGraph:
    def __init__(self, model=None):
        self.entities = {}
        self.relations = {}
        self.triples = []
        self.documents = {}
        self.model = model
        self.llm = None
    
    def get_relation_candidates(self, entity_id: str) -> List[Dict]:
        """Get all possible relations for an entity"""
        candidates = []
        
        # Find all triples where entity is either head or tail
        for triple in self.triples:
            head_id, relation_id, tail_id = triple
            
            if head_id == entity_id:
                # Entity is head
                candidates.append({
                    'entity_id': entity_id,
                    'relation': relation_id,
                    'relation_name': self.relations[relation_id]['name'],
                    'target_id': tail_id,
                    'target_name': self.entities[tail_id]['name'],
                    'head': True
                })
            
            if tail_id == entity_id:
                # Entity is tail
                candidates.append({
                    'entity_id': entity_id,
                    'relation': relation_id,
                    'relation_name': self.relations[relation_id]['name'],
                    'target_id': head_id,
                    'target_name': self.entities[head_id]['name'],
                    'head': False
                })
        
        return candidates