src/kg/retriever.py

"""KG retriever that returns chunk IDs and short node summaries for a question."""
from typing import List, Tuple
from .store import KGStore


class KGRetriever:
    def __init__(self, kg_store: KGStore):
        self.kg = kg_store

    def get_context_for_question(self, question: str, hops: int = 1) -> Tuple[List[str], List[str]]:
        """Return (chunk_ids, node_summaries).

        This simple implementation finds entities whose labels appear in the question
        and returns linked chunk ids. For hops >1 you could expand to related entities.
        """
        entity_uris = self.kg.query_entities(question)
        chunk_ids = []
        summaries = []
        for e in entity_uris:
            # uri like http://.../entity/<label>
            label = e.split("/entity/", 1)[-1].replace("_", " ")
            chunks = self.kg.find_chunks_for_entity(label)
            chunk_ids.extend(chunks)
            summaries.append(label)
        # dedupe
        chunk_ids = list(dict.fromkeys(chunk_ids))
        return chunk_ids, summaries