server/classes/RAGRetriever.py

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 []