Change the retrieval and reranking into two steps search with two different indexes - which is supposed to make the latency much lower (faster)

rag_agent.py

@@ -5,6 +5,7 @@ import os
 
 from retriever import Retriever
 from reranker import Reranker
+from text_embedder_encoder import TextEmbedder, encoder_model_name
 
 
 retriever = Retriever()
@@ -27,15 +28,16 @@ class RAGAgent:
         self.model_name = model_name
         self.max_tokens = max_tokens
         self.temperature = temperature
+        self.text_embedder = TextEmbedder()
         self.conversation_summary = ""
         self.messages = []
 
     def get_context(self, query: str) -> List[str]:
         # Get initial candidates from retriever
-        retrieved_docs = self.retriever.search_similar(query)
-
+        query_vector = self.text_embedder.encode(query)
+        retrieved_answers_ids = self.retriever.search_similar(query_vector)
         # Rerank the candidates
-        context = self.reranker.rerank(query, retrieved_docs)
+        context = self.reranker.rerank(query_vector, retrieved_answers_ids)
 
         return context
 

reranker.py

@@ -1,22 +1,37 @@
+from pinecone import Pinecone
 from sklearn.metrics.pairwise import cosine_similarity
+import os
 
 
-from text_embedder_encoder import TextEmbedder
+from text_embedder_encoder import encoder_model_name
 
 
 class Reranker:
-    def __init__(self):
-        self.text_embedder = TextEmbedder()
+    def __init__(self,
+                 pinecone_api_key=os.environ["pinecone_api_key"],
+                 answer_index_name=f"hebrew-dentist-answers-{encoder_model_name.replace('/', '-')}".lower()):
+        self.pc = Pinecone(api_key=pinecone_api_key)
+        self.answer_index_name = answer_index_name
 
-    def rerank(self, query, retrieved_docs, top_n=5):
+    def rerank(self, query_vector, retrieved_answers_ids, top_n=5):
         # Encode query and documents
-        query_embedding = self.text_embedder.encode(query)
-        doc_embeddings = self.text_embedder.encode_many(retrieved_docs)
-        similarity_scores = cosine_similarity([query_embedding], doc_embeddings)[0]
+        try:
+            index = self.pc.Index(self.answer_index_name)
+            fetch_response = index.fetch(ids=retrieved_answers_ids)
 
-        similarity_scores_with_idxes = list(zip(similarity_scores, range(len(similarity_scores))))
-        similarity_scores_with_idxes.sort(reverse=True)
-        similarity_scores_with_idxes_final = similarity_scores_with_idxes[:top_n]
-        reranked_docs = [retrieved_docs[idx] for score, idx in similarity_scores_with_idxes_final if score >= 0.7]
+            doc_embeddings = []
+            answers = []
+            for i in range(len(retrieved_answers_ids)):
+                doc_embeddings.append(fetch_response['vectors'][retrieved_answers_ids[i]]['values'])
+                answers.append(fetch_response['vectors'][retrieved_answers_ids[i]]['metadata']['answer'])
 
-        return reranked_docs
+            similarity_scores = cosine_similarity([query_vector], doc_embeddings)[0]
+            similarity_scores_with_idxes = list(zip(similarity_scores, range(len(similarity_scores))))
+            similarity_scores_with_idxes.sort(reverse=True)
+            similarity_scores_with_idxes_final = similarity_scores_with_idxes[:top_n]
+            reranked_answers = [answers[idx] for score, idx in similarity_scores_with_idxes_final if score >= 0.7]
+
+            return reranked_answers
+        except Exception as e:
+            print(f"Error performing rerank: {e}")
+            return []

retriever.py

@@ -1,29 +1,25 @@
 from pinecone import Pinecone
 import os
 
-from text_embedder_encoder import TextEmbedder, encoder_model_name
+from text_embedder_encoder import encoder_model_name
 
 
 class Retriever:
     def __init__(self,
                  pinecone_api_key=os.environ["pinecone_api_key"],
-                 index_name=f"hebrew-dentist-qa-{encoder_model_name.replace('/', '-')}".lower()):
+                 question_index_name=f"hebrew-dentist-questions-{encoder_model_name.replace('/', '-')}".lower()):
         # Initialize Pinecone connection
         self.pc = Pinecone(api_key=pinecone_api_key)
-        self.index_name = index_name
-        self.text_embedder = TextEmbedder()
-        self.vector_dim = 768
+        self.question_index_name = question_index_name
 
-    def search_similar(self, query_text, top_k=50):
+    def search_similar(self, query_vector, top_k=50):
         """
         Search for similar content using vector similarity in Pinecone
         """
         try:
-            # Generate embedding for query
-            query_vector = self.text_embedder.encode(query_text)
 
             # Get Pinecone index
-            index = self.pc.Index(self.index_name)
+            index = self.pc.Index(self.question_index_name)
 
             # Execute search
             results = index.query(
@@ -32,12 +28,12 @@ class Retriever:
                 include_metadata=True,
             )
 
-            answers = []
+            answers_records_ids = []
             for match in results['matches']:
-                answer = match['metadata']['answer']
-                answers.append(answer)
+                answers_records_ids.append(
+                    ':'.join(match['id'].split(':')[:-1]) + ":" + str(int(match['metadata']['answer_id'])))
 
-            return answers
+            return answers_records_ids
         except Exception as e:
-            print(f"Error performing similarity search: {e}")
+            print(f"Error performing retriever: {e}")
             return []

text_embedder_encoder.py

@@ -1,7 +1,6 @@
 import torch
 import numpy as np
 from sentence_transformers import SentenceTransformer
-from typing import List
 
 
 encoder_model_name = 'MPA/sambert'
@@ -36,21 +35,21 @@ class TextEmbedder:
 
         return embeddings
 
-    def encode_many(self, texts: List[str]) -> np.ndarray:
-        """
-            Encode Hebrew text using LaBSE model with handling for texts longer than max_seq_length.
-
-            Args:
-                text (str): Hebrew text to encode
-                model_name (str): Name of the model to use
-                # max_seq_length (int): Maximum sequence length for the model
-                strategy (str): Strategy for combining sentence embeddings ('mean' or 'concat')
-
-            Returns:
-                numpy.ndarray: Text embedding
-            """
-        # Get embeddings for the text
-        embeddings = self.model.encode(texts)
-        embeddings = [[float(x) for x in embedding] for embedding in embeddings]
-
-        return embeddings
+    # def encode_many(self, texts: List[str]) -> np.ndarray:
+    #     """
+    #         Encode Hebrew text using LaBSE model with handling for texts longer than max_seq_length.
+    #
+    #         Args:
+    #             text (str): Hebrew text to encode
+    #             model_name (str): Name of the model to use
+    #             # max_seq_length (int): Maximum sequence length for the model
+    #             strategy (str): Strategy for combining sentence embeddings ('mean' or 'concat')
+    #
+    #         Returns:
+    #             numpy.ndarray: Text embedding
+    #         """
+    #     # Get embeddings for the text
+    #     embeddings = self.model.encode(texts)
+    #     embeddings = [[float(x) for x in embedding] for embedding in embeddings]
+    #
+    #     return embeddings