Spaces:
Sleeping
Sleeping
File size: 2,479 Bytes
2e9afea | 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 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | from typing import List,Dict,Any,Tuple
from .EmbeddingManager import EmbeddingManager
from .VectorStore import VectorStore
import numpy as np
class RAGRetriever:
def __init__(self,vector_store: VectorStore, embedding_manager:EmbeddingManager):
self.vector_store= vector_store
self.embedding_manager= embedding_manager
def retrieve(self,query: str, top_k: int=10, score_threshold: float= 0.5) -> List[Dict[str,Any]]:
print(f"retrieving documents for query: {query}")
print(f"Top_k: {top_k} score_threshold: {score_threshold}")
query_embedding= self.embedding_manager.generate_embeddings([query])[0]
# 1D array representing just 1 query
# search in vector store
try:
results= self.vector_store.collection.query(
query_embeddings= [query_embedding.tolist()],
# this expects batch of queries
n_results= top_k
)
retrieved_docs= []
if results['documents'] and results['documents'][0]:
documents= results['documents'][0]
metadatas= results['metadatas'][0]
distances= results['distances'][0]
ids= results['ids'][0]
metadatas= results['metadatas'][0]
for i, (doc_id,document,metadata,distance) in enumerate(zip(ids,documents,metadatas,distances)):
# convert distance to similarity score (chromadb uses cosine distance)
print(distance)
similarity_score= float(1.0-distance)
source_file = metadata.get('source', metadata.get('source_file', 'Unknown Source'))
print(source_file)
if similarity_score>=score_threshold:
retrieved_docs.append({
'id': doc_id,
'content': document,
'metadata': metadata,
'similarity_score': similarity_score,
'distance': distance,
'rank': i+1
})
print(f"Retrieved {len(retrieved_docs)} document after filtering")
else:
print("No documents found")
return retrieved_docs
except Exception as e:
print(f"erorr in retrieving documents for query: {query}")
return []
|