From e65d47fad5fe535ddb8b8b02f0dc6a0693806a44 Mon Sep 17 00:00:00 2001
From: Wesley Stewart <wstewart3@gmail.com>
Date: Fri, 15 Mar 2024 19:21:39 +0000
Subject: [PATCH] Added second variable for the rag

---
 private_gpt/components/vector_store/vector_store_component.py | 4 ++--
 private_gpt/settings/settings.py                              | 4 ++++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/private_gpt/components/vector_store/vector_store_component.py b/private_gpt/components/vector_store/vector_store_component.py
index 5641f25e..5bde9cec 100644
--- a/private_gpt/components/vector_store/vector_store_component.py
+++ b/private_gpt/components/vector_store/vector_store_component.py
@@ -130,12 +130,12 @@ class VectorStoreComponent:
         self,
         index: VectorStoreIndex,
         context_filter: ContextFilter | None = None,
-        similarity_top_k: int = 2,
+        sim_top_k = settings.rag.similarity_top_k,
     ) -> VectorIndexRetriever:
         # This way we support qdrant (using doc_ids) and the rest (using filters)
         return VectorIndexRetriever(
             index=index,
-            similarity_top_k=similarity_top_k,
+            similarity_top_k=sim_top_k,
             doc_ids=context_filter.docs_ids if context_filter else None,
             filters=(
                 _doc_id_metadata_filter(context_filter)
diff --git a/private_gpt/settings/settings.py b/private_gpt/settings/settings.py
index 05f24dca..80c8df82 100644
--- a/private_gpt/settings/settings.py
+++ b/private_gpt/settings/settings.py
@@ -108,6 +108,10 @@ class RagSettings(BaseModel):
         None,
         description="If set, any documents retrieved from the RAG must meet a certain score. Acceptable values are between 0 and 1.",
     )
+    similarity_top_k: int = Field(
+        2,
+        description="An integer that specifies the number of documents to use when searching the RAG database",
+    )
 
 
 class VectorstoreSettings(BaseModel):