From e7bd571ac589357b7a687024a642c46419183af4 Mon Sep 17 00:00:00 2001 From: Stan Girard Date: Thu, 25 Jan 2024 20:19:56 -0800 Subject: [PATCH] =?UTF-8?q?feat:=20=F0=9F=8E=B8=20rag?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit now works with 30 chunks --- backend/llm/knowledge_brain_qa.py | 2 +- backend/llm/rags/quivr_rag.py | 4 +++- backend/models/databases/repository.py | 4 +--- backend/models/databases/supabase/vectors.py | 4 ++-- .../repository/brain/get_question_context_from_brain.py | 1 + backend/vectorstore/supabase.py | 9 ++++++--- 6 files changed, 14 insertions(+), 10 deletions(-) diff --git a/backend/llm/knowledge_brain_qa.py b/backend/llm/knowledge_brain_qa.py index bc33681d3219..27adc99480e2 100644 --- a/backend/llm/knowledge_brain_qa.py +++ b/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 diff --git a/backend/llm/rags/quivr_rag.py b/backend/llm/rags/quivr_rag.py index dd8a945a06c8..d0eb23918bab 100644 --- a/backend/llm/rags/quivr_rag.py +++ b/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( diff --git a/backend/models/databases/repository.py b/backend/models/databases/repository.py index dc62ff06b4d5..cfc1c41f4f32 100644 --- a/backend/models/databases/repository.py +++ b/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 diff --git a/backend/models/databases/supabase/vectors.py b/backend/models/databases/supabase/vectors.py index 8c52d8eb2764..5b16d49cf581 100644 --- a/backend/models/databases/supabase/vectors.py +++ b/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() diff --git a/backend/repository/brain/get_question_context_from_brain.py b/backend/repository/brain/get_question_context_from_brain.py index 7aaca453a308..9f6fde6391e6 100644 --- a/backend/repository/brain/get_question_context_from_brain.py +++ b/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) diff --git a/backend/vectorstore/supabase.py b/backend/vectorstore/supabase.py index e303cbc2a788..1e2e1e07944a 100644 --- a/backend/vectorstore/supabase.py +++ b/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()