from typing import List, Dict
from .vector_store import VectorStore
import time
class Retriever:
def __init__(self):
self.vector_store = VectorStore()
def retrieve(self, query: str, top_k: int = 5, filters: Dict = None) -> Dict:
"""Retrieve relevant documents with timing"""
start_time = time.time()
results = self.vector_store.search(query, top_k, filters)
query_time_ms = (time.time() - start_time) * 1000
return {
"query": query,
"results": results,
"query_time_ms": round(query_time_ms, 2),
"total_results": len(results)
}
def retrieve_with_context(self, query: str, top_k: int = 5) -> Dict:
"""Retrieve with expanded context (neighboring chunks)"""
results = self.retrieve(query, top_k)
# TODO: Add logic to fetch ±1 neighboring chunks
return results