feat(RAG): Introduce SentenceTransformer Reranker (#1810)

private_gpt/server/chat/chat_service.py

@@ -9,6 +9,7 @@ from llama_index.core.indices import VectorStoreIndex
 from llama_index.core.indices.postprocessor import MetadataReplacementPostProcessor
 from llama_index.core.llms import ChatMessage, MessageRole
 from llama_index.core.postprocessor import (
+    SentenceTransformerRerank,
     SimilarityPostprocessor,
 )
 from llama_index.core.storage import StorageContext
@@ -113,16 +114,24 @@ class ChatService:
                 context_filter=context_filter,
                 similarity_top_k=self.settings.rag.similarity_top_k,
             )
+            node_postprocessors = [
+                MetadataReplacementPostProcessor(target_metadata_key="window"),
+                SimilarityPostprocessor(
+                    similarity_cutoff=settings.rag.similarity_value
+                ),
+            ]
+
+            if settings.rag.rerank.enabled:
+                rerank_postprocessor = SentenceTransformerRerank(
+                    model=settings.rag.rerank.model, top_n=settings.rag.rerank.top_n
+                )
+                node_postprocessors.append(rerank_postprocessor)
+
             return ContextChatEngine.from_defaults(
                 system_prompt=system_prompt,
                 retriever=vector_index_retriever,
                 llm=self.llm_component.llm,  # Takes no effect at the moment
-                node_postprocessors=[
-                    MetadataReplacementPostProcessor(target_metadata_key="window"),
-                    SimilarityPostprocessor(
-                        similarity_cutoff=settings.rag.similarity_value
-                    ),
-                ],
+                node_postprocessors=node_postprocessors,
             )
         else:
             return SimpleChatEngine.from_defaults(

private_gpt/settings/settings.py

@@ -284,15 +284,31 @@ class UISettings(BaseModel):
     )
 
 
+class RerankSettings(BaseModel):
+    enabled: bool = Field(
+        False,
+        description="This value controls whether a reranker should be included in the RAG pipeline.",
+    )
+    model: str = Field(
+        "cross-encoder/ms-marco-MiniLM-L-2-v2",
+        description="Rerank model to use. Limited to SentenceTransformer cross-encoder models.",
+    )
+    top_n: int = Field(
+        2,
+        description="This value controls the number of documents returned by the RAG pipeline.",
+    )
+
+
 class RagSettings(BaseModel):
     similarity_top_k: int = Field(
         2,
-        description="This value controls the number of documents returned by the RAG pipeline",
+        description="This value controls the number of documents returned by the RAG pipeline or considered for reranking if enabled.",
     )
     similarity_value: float = Field(
         None,
         description="If set, any documents retrieved from the RAG must meet a certain match score. Acceptable values are between 0 and 1.",
     )
+    rerank: RerankSettings
 
 
 class PostgresSettings(BaseModel):