|
|
"""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: |
|
|
|
|
|
label = e.split("/entity/", 1)[-1].replace("_", " ") |
|
|
chunks = self.kg.find_chunks_for_entity(label) |
|
|
chunk_ids.extend(chunks) |
|
|
summaries.append(label) |
|
|
|
|
|
chunk_ids = list(dict.fromkeys(chunk_ids)) |
|
|
return chunk_ids, summaries |
|
|
|
|
|
|
