Merge pull request #40 from ls1intum/optimize-reranking

Optimize reranking

Author: ninori9

rag/app/api/question_router.py

@@ -17,17 +17,21 @@ async def ask(request: UserRequest):
 
     if not question or not classification:
         raise HTTPException(status_code=400, detail="No question or classification provided")
+    
+    answer = request_handler.handle_question(question, classification, language, org_id=org_id)
+    return {"answer": answer}
 
-    if config.TEST_MODE == "true":
-        answer, used_tokens, general_context, specific_context, sq_context = request_handler.handle_question_test_mode(question,
-                                                                                                           classification,
-                                                                                                           language,
-                                                                                                           org_id=org_id)
-        return {"answer": answer, "used_tokens": used_tokens, "general_context": general_context,
-                "specific_context": specific_context, "sq_context": sq_context}
-    else:
-        answer = request_handler.handle_question(question, classification, language, org_id=org_id)
-        return {"answer": answer}
+    # Uncomment to use test mode and calculate RAG metrics
+    # if config.TEST_MODE == "true":
+    #     answer, used_tokens, general_context, specific_context, sq_context = request_handler.handle_question_test_mode(question,
+    #                                                                                                        classification,
+    #                                                                                                        language,
+    #                                                                                                        org_id=org_id)
+    #     return {"answer": answer, "used_tokens": used_tokens, "general_context": general_context,
+    #             "specific_context": specific_context, "sq_context": sq_context}
+    # else:
+    #     answer = request_handler.handle_question(question, classification, language, org_id=org_id)
+    #     return {"answer": answer}
 
 
 @question_router.post("/chat", tags=["chatbot"], dependencies=[Depends(auth_handler.verify_api_key)])

rag/app/managers/request_handler.py

@@ -1,6 +1,6 @@
 import logging
 from enum import Enum
-from typing import List, Union, Tuple, Optional
+from typing import List, Union, Tuple, Optional, Dict
 
 import weaviate
 import weaviate.classes as wvc
@@ -239,28 +239,23 @@ def get_question_embedding(self, question: str) -> List[float]:
         question_embedding = self.model.embed(question)
         return question_embedding
 
-    def get_relevant_context(self, question: str, question_embedding: str, study_program: str, language: str, org_id: Optional[int],
-                             test_mode: bool = False,
-                             limit=10, top_n=5, filter_by_org: bool = True) -> Union[str, Tuple[str, List[str]]]:
+
+    def get_relevant_context(self, question_embedding: List[float], study_program: str, org_id: Optional[int],
+                             limit=10, filter_by_org: bool = True) -> List[Dict]:
         """
-        Retrieve relevant documents based on the question embedding and study program.
-        Optionally returns both the concatenated context and the sorted context list for testing purposes.
+        Retrieves relevant context documents based on the given question embedding and study program.
 
         Args:
-            question (str): The student's question.
-            study_program (str): The study program of the student.
-            keywords (str, optional): Extracted keywords for boosting. Defaults to None.
-            test_mode (bool, optional): If True, returns both context and sorted_context. Defaults to False.
+            question_embedding (List[float]): The vector embedding representing the student's question.
+            study_program (str): The name of the study program to filter documents.
+            org_id (Optional[int]): The organization ID to filter documents (if applicable).
+            limit (int, optional): The maximum number of documents to retrieve. Defaults to 10.
+            filter_by_org (bool, optional): Whether to filter results by organization ID. Defaults to True.
 
         Returns:
-            Union[str, Tuple[str, List[str]]]:
-                - If test_mode is False: Returns the concatenated context string.
-                - If test_mode is True: Returns a tuple of (context, sorted_context list).
+            List[Dict]: A list of document dictionaries relevant to the query.
         """
         try:
-            # Define the number of documents to retrieve
-            min_relevance_score = 0.35
-            
             # Normalize the study program name
             study_program = WeaviateManager.normalize_study_program_name(study_program)
 
@@ -273,11 +268,6 @@ def get_relevant_context(self, question: str, question_embedding: str, study_pro
             else:
                 filters = Filter.by_property(DocumentSchema.STUDY_PROGRAMS.value).contains_any([study_program])
 
-            # If getting general context, adjust the parameters
-            if study_program.lower() != "general":
-                limit = 10
-                min_relevance_score = 0.25
-
 
             # Perform the vector-based query with filters
             query_result = self.documents.query.near_vector(
@@ -295,51 +285,28 @@ def get_relevant_context(self, question: str, question_embedding: str, study_pro
                 }
                 for result in query_result.objects
             ]
-            content_content_list: List[str] = [doc['content'] for doc in context_list]
-
-            # Remove exact duplicates from context_list
-            content_content_list = WeaviateManager.remove_exact_duplicates(content_content_list)
-
-            # Rerank the unique contexts using Cohere
-            sorted_context = self.reranker.rerank_with_cohere(context_list=content_content_list, query=question,
-                                                              language=language,
-                                                              min_relevance_score=min_relevance_score, top_n=top_n)
-            # Integrate links
-            sorted_context_with_links = []
-            for sorted_content in sorted_context:
-                for doc in context_list:
-                    if doc['content'] == sorted_content:
-                        if doc['link']:
-                            sorted_context_with_links.append(f'Link: {doc["link"]}\nContent: {doc["content"]}')
-                        else:
-                            sorted_context_with_links.append(f'Link: -\nContent: {doc["content"]}')
-                        break
-
-            context = "\n-----\n".join(sorted_context_with_links)
-
-            # Return based on test_mode
-            if test_mode:
-                return context, sorted_context_with_links
-            else:
-                return context
+            
+            return context_list
 
         except Exception as e:
             logging.error(f"Error retrieving relevant context: {e}")
             # tb = traceback.format_exc()
             # logging.error("Traceback:\n%s", tb)
-            return "" if not test_mode else ("", [])
+            return []
 
-    def get_relevant_sample_questions(self, question: str, question_embedding: str, language: str, org_id: int) -> List[SampleQuestion]:
+
+    def get_relevant_sample_questions(self, question: str, question_embedding: List[float], language: str, org_id: int) -> List[SampleQuestion]:
         """
-        Retrieve relevant sample questions and answers based on the question embedding.
+        Retrieves relevant sample questions and their answers based on the provided question and its embedding.
 
         Args:
-            question (str): The student's question.
+            question (str): The original student question.
+            question_embedding (List[float]): The vector embedding of the question.
             language (str): The language of the question.
-            top_k (int): The number of top relevant sample questions to return.
+            org_id (int): The organization ID to filter sample questions.
 
         Returns:
-            List[SampleQuestion]: A list of SampleQuestion objects, sorted based on reranking results.
+            List[SampleQuestion]: A list of SampleQuestion objects, sorted by relevance.
         """
         try:
             limit = 5
@@ -364,19 +331,22 @@ def get_relevant_sample_questions(self, question: str, question_embedding: str,
                                                        study_programs=study_programs))
 
             # Rerank the sample questions using the reranker
-            context_list = [sq.question for sq in sample_questions]
-            sorted_questions = self.reranker.rerank_with_cohere(
+            context_list = [
+                (f"Question: {sq.question}\nAnswer: {sq.answer}" if language == "English"
+                else f"Frage: {sq.question}\nAntwort: {sq.answer}")
+                for sq in sample_questions
+            ]
+            
+            rerank_results = self.reranker.rerank_with_cohere(
                 context_list=context_list, query=question, language=language, top_n=top_n,
-                min_relevance_score=min_relevance_score
             )
 
-            # Map the sorted questions back to SampleQuestion objects
             sorted_sample_questions: List[SampleQuestion] = []
-            for sorted_question in sorted_questions:
-                for sq in sample_questions:
-                    if sq.question == sorted_question:
-                        sorted_sample_questions.append(sq)
-                        break
+            for result in rerank_results:
+                idx = result['index']
+                score = result['relevance_score']
+                if score >= min_relevance_score and idx < len(sample_questions):
+                    sorted_sample_questions.append(sample_questions[idx])
 
             return sorted_sample_questions
 
@@ -673,3 +643,15 @@ def remove_exact_duplicates(context_list: List[str]) -> List[str]:
                 unique_context.append(context)
                 seen.add(context)
         return unique_context
+    
+    @staticmethod
+    def remove_exact_duplicates_from_dict(dicts: List[Dict], key: str = 'content') -> list:
+        """Remove dicts with duplicate values for given key, preserving order."""
+        seen = set()
+        deduped = []
+        for d in dicts:
+            val = d.get(key)
+            if val not in seen:
+                deduped.append(d)
+                seen.add(val)
+        return deduped

rag/app/managers/weaviate_manager.py

@@ -3,7 +3,7 @@
 import cohere
 from app.models.base_model import BaseModelClient
 from sklearn.metrics.pairwise import cosine_similarity
-from typing import List
+from typing import List, Dict
 import requests
 
 class DocumentWithEmbedding:
@@ -64,23 +64,22 @@ def rerank_with_embeddings(self, context_list: List[DocumentWithEmbedding], keyw
 
         return ranked_context_list
 
-    def rerank_with_cohere(self, context_list: List[str], query: str, language: str, min_relevance_score: float, top_n: int = 5) -> List[str]:
+    def rerank_with_cohere(self, context_list: List[str], query: str, language: str, top_n: int = 5) -> List[Dict]:
         """
         Re-ranks the context list using the Cohere reranking model deployed on Azure.
 
         Args:
             context_list (List[str]): List of document texts to be re-ranked.
             query (str): The query string to rerank the documents against.
             language (str): The language of the documents ('english' or other).
-            min_relevance_score (float): The minimum relevance score to consider.
             top_n (int): The number of top results to return after re-ranking.
 
         Returns:
-            List[str]: A list of the re-ranked document contents based on relevance.
+            List[Dict]: A list of the re-ranked document contents based on relevance.
         """
         try:
-            if len(context_list) == 0:
-                return context_list
+            if not context_list:
+                return []
 
             # Determine the correct endpoint URL and API key based on language
             if language.lower() == "english":
@@ -103,30 +102,13 @@ def rerank_with_cohere(self, context_list: List[str], query: str, language: str,
             }
 
             response = requests.post(rerank_url, headers=headers, json=payload)
-
             if response.status_code != 200:
                 logging.error(f"Error during Cohere re-ranking: {response.status_code} {response.text}")
-                return context_list[:top_n]
-
-            response_json = response.json()
-
-            # Log the full response from the API for debugging
-            results = response_json.get('results', [])
-
-            # Log the ranked documents that are in the top_n
-            ranked_indices = []
-            for i, result in enumerate(results):
-                index = result['index']
-                relevance_score = result.get('relevance_score')
-                # Filter results based on min_relevance_score
-                if relevance_score >= min_relevance_score:
-                    ranked_indices.append(index)
-
-            # Get the ranked documents based on the indices
-            ranked_context_list = [context_list[result['index']] for result in results]
+                return [{'index': i, 'relevance_score': 1.0} for i in range(min(top_n, len(context_list)))]
 
-            return ranked_context_list
+            results = response.json().get('results', [])
+            return results
 
         except Exception as e:
             logging.error(f"Error during Cohere re-ranking: {e}")
-            return context_list[:top_n]
+            return [{'index': i, 'relevance_score': 1.0} for i in range(min(top_n, len(context_list)))]