feat (RAG): Implement configurable search mode in RAGQueryConfig

Signed-off-by: Varsha Prasad Narsing <varshaprasad96@gmail.com>
This commit is contained in:
Varsha Prasad Narsing 2025-04-14 16:53:17 -07:00
parent 85b5f3172b
commit e2a7022d3c
14 changed files with 210 additions and 43 deletions

View file

@ -99,9 +99,15 @@ class FaissIndex(EmbeddingIndex):
# Save updated index
await self._save_index()
async def query(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
async def query(
self,
embedding: NDArray,
query_string: Optional[str],
k: int,
score_threshold: float,
mode: Optional[str],
) -> QueryChunksResponse:
distances, indices = await asyncio.to_thread(self.index.search, embedding.reshape(1, -1).astype(np.float32), k)
chunks = []
scores = []
for d, i in zip(distances[0], indices[0], strict=False):