feat: 🎸 rag

now works with 30 chunks

backend/llm/knowledge_brain_qa.py

@@ -60,7 +60,7 @@ class Config:
     temperature: float = 0.1
     chat_id: str = None  # pyright: ignore reportPrivateUsage=none
     brain_id: str  # pyright: ignore reportPrivateUsage=none
-    max_tokens: int = 256
+    max_tokens: int = 2000
     streaming: bool = False
     knowledge_qa: Optional[RAGInterface]
     metadata: Optional[dict] = None

backend/llm/rags/quivr_rag.py

@@ -60,7 +60,7 @@ class Config:
     temperature: float = 0.1
     chat_id: str = None  # pyright: ignore reportPrivateUsage=none
     brain_id: str = None  # pyright: ignore reportPrivateUsage=none
-    max_tokens: int = 256
+    max_tokens: int = 2000
     streaming: bool = False
 
     @property
@@ -91,6 +91,7 @@ def __init__(
         chat_id: str,
         streaming: bool = False,
         prompt_id: Optional[UUID] = None,
+        max_tokens: int = 2000,
         **kwargs,
     ):
         super().__init__(
@@ -103,6 +104,7 @@ def __init__(
         self.supabase_client = self._create_supabase_client()
         self.vector_store = self._create_vector_store()
         self.prompt_id = prompt_id
+        self.max_tokens = max_tokens
 
     def _create_supabase_client(self) -> Client:
         return create_client(

backend/models/databases/repository.py

@@ -59,9 +59,7 @@ def get_vectors_by_file_name(self, file_name: str):
         pass
 
     @abstractmethod
-    def similarity_search(
-        self, query_embedding, table: str, top_k: int, threshold: float
-    ):
+    def similarity_search(self, query_embedding, table: str, k: int, threshold: float):
         pass
 
     @abstractmethod

backend/models/databases/supabase/vectors.py

@@ -30,12 +30,12 @@ def get_vectors_by_file_sha1(self, file_sha1):
         return response
 
     # TODO: remove duplicate similarity_search in supabase vector store
-    def similarity_search(self, query_embedding, table, top_k, threshold):
+    def similarity_search(self, query_embedding, table, k, threshold):
         response = self.db.rpc(
             table,
             {
                 "query_embedding": query_embedding,
-                "match_count": top_k,
+                "match_count": k,
                 "match_threshold": threshold,
             },
         ).execute()

backend/repository/brain/get_question_context_from_brain.py

@@ -29,6 +29,7 @@ def get_question_context_from_brain(brain_id: UUID, question: str) -> str:
         embeddings,
         table_name="vectors",
         brain_id=str(brain_id),
+        number_docs=20,
     )
     documents = vector_store.similarity_search(question, k=20, threshold=0.8)
 

backend/vectorstore/supabase.py

@@ -14,6 +14,7 @@ class CustomSupabaseVectorStore(SupabaseVectorStore):
 
     brain_id: str = "none"
     user_id: str = "none"
+    number_docs: int = 35
 
     def __init__(
         self,
@@ -22,10 +23,12 @@ def __init__(
         table_name: str,
         brain_id: str = "none",
         user_id: str = "none",
+        number_docs: int = 35,
     ):
         super().__init__(client, embedding, table_name)
         self.brain_id = brain_id
         self.user_id = user_id
+        self.number_docs = number_docs
 
     def find_brain_closest_query(
         self,
@@ -42,7 +45,7 @@ def find_brain_closest_query(
             table,
             {
                 "query_embedding": query_embedding,
-                "match_count": k,
+                "match_count": self.number_docs,
                 "p_user_id": str(self.user_id),
             },
         ).execute()
@@ -62,7 +65,7 @@ def find_brain_closest_query(
     def similarity_search(
         self,
         query: str,
-        k: int = 6,
+        k: int = 35,
         table: str = "match_vectors",
         threshold: float = 0.5,
         **kwargs: Any,
@@ -73,7 +76,7 @@ def similarity_search(
             table,
             {
                 "query_embedding": query_embedding,
-                "match_count": k,
+                "match_count": self.number_docs,
                 "p_brain_id": str(self.brain_id),
             },
         ).execute()