Spaces:
Sleeping
Sleeping
File size: 1,467 Bytes
2db8ee1 8a660ff 1d2f776 8a660ff 1d2f776 8a660ff 1d2f776 2db8ee1 1d2f776 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | class IndexManager:
def __init__(self, doc_index, chunk_index):
self.doc_index = doc_index
self.chunk_index = chunk_index
def retrieve(self, query, k=10, source_filter=None):
"""
Retrieve relevant chunks for a query.
Args:
query: The search query
k: Number of results to return
source_filter: Optional set of filenames to restrict search to.
If provided, only chunks from these sources are returned.
If None, all chunks are searched (backward compatible).
"""
# Get chunks from the global chunk index (fetch extra for filtering)
all_chunks = self.chunk_index.search(query, k=k * 3)
if source_filter:
# Filter to only chunks from the specified documents
filtered = [c for c in all_chunks if c["source"] in source_filter]
return filtered[:k]
# No filter — use two-stage doc-aware retrieval
relevant_docs = self.doc_index.search(query, k=3)
relevant_sources = {doc["source"] for doc in relevant_docs}
boosted = [c for c in all_chunks if c["source"] in relevant_sources]
others = [c for c in all_chunks if c["source"] not in relevant_sources]
result = boosted[:k]
remaining = k - len(result)
if remaining > 0:
result.extend(others[:remaining])
return result
|