Spaces:
Sleeping
Sleeping
| 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 | |