Spaces:
Sleeping
Sleeping
| from collections import Counter | |
| from collections.abc import Sequence | |
| from techdocs_graphrag.models import RetrievalContext, RetrievalMode, RetrievedChunk | |
| from techdocs_graphrag.ports import KnowledgeGraph, LanguageModel | |
| MAX_LEXICAL_SEEDS = 5 | |
| MAX_MENTION_SEEDS = 8 | |
| class HybridRetriever: | |
| def __init__( | |
| self, | |
| *, | |
| llm: LanguageModel, | |
| store: KnowledgeGraph, | |
| top_k: int, | |
| max_hops: int, | |
| max_facts: int, | |
| ) -> None: | |
| self._llm = llm | |
| self._store = store | |
| self._top_k = top_k | |
| self._max_hops = max_hops | |
| self._max_facts = max_facts | |
| def retrieve(self, question: str, mode: RetrievalMode = RetrievalMode.GRAPH) -> RetrievalContext: | |
| embedding = self._llm.embed([question])[0] | |
| chunks = self._store.vector_search(embedding, self._top_k) | |
| if mode is RetrievalMode.VECTOR: | |
| return RetrievalContext(chunks=chunks) | |
| seeds = self._seed_entities(question, chunks) | |
| facts = self._store.expand(seeds, self._max_hops, self._max_facts) | |
| return RetrievalContext(chunks=chunks, facts=facts, seed_entities=seeds) | |
| def _seed_entities(self, question: str, chunks: Sequence[RetrievedChunk]) -> list[str]: | |
| lexical_matches = self._store.search_entities(question, MAX_LEXICAL_SEEDS) | |
| mentions = Counter(key for chunk in chunks for key in chunk.entity_keys) | |
| most_mentioned = [key for key, _ in mentions.most_common(MAX_MENTION_SEEDS)] | |
| return list(dict.fromkeys([*lexical_matches, *most_mentioned])) | |