diff --git a/backend/modules/chat/controller/chat/brainful_chat.py b/backend/modules/chat/controller/chat/brainful_chat.py index ef8f48600..4605803c1 100644 --- a/backend/modules/chat/controller/chat/brainful_chat.py +++ b/backend/modules/chat/controller/chat/brainful_chat.py @@ -64,7 +64,7 @@ class BrainfulChat(ChatInterface): else: embeddings = OpenAIEmbeddings() vector_store = CustomSupabaseVectorStore( - supabase_client, embeddings, table_name="vectors" + supabase_client, embeddings, table_name="vectors", user_id=user_id ) # Get the first question from the chat_question @@ -72,7 +72,7 @@ class BrainfulChat(ChatInterface): history = chat_service.get_chat_history(chat_id) if history: question = history[0].user_message - list_brains = vector_store.find_brain_closest_query(question) + list_brains = vector_store.find_brain_closest_query(user_id, question) if list_brains: brain_id_to_use = list_brains[0]["id"] else: diff --git a/backend/vectorstore/supabase.py b/backend/vectorstore/supabase.py index e520fdcb9..d77230a6e 100644 --- a/backend/vectorstore/supabase.py +++ b/backend/vectorstore/supabase.py @@ -13,6 +13,7 @@ class CustomSupabaseVectorStore(SupabaseVectorStore): """A custom vector store that uses the match_vectors table instead of the vectors table.""" brain_id: str = "none" + user_id: str = "none" def __init__( self, @@ -20,12 +21,15 @@ class CustomSupabaseVectorStore(SupabaseVectorStore): embedding: Embeddings, table_name: str, brain_id: str = "none", + user_id: str = "none", ): super().__init__(client, embedding, table_name) self.brain_id = brain_id + self.user_id = user_id def find_brain_closest_query( self, + user_id: str, query: str, k: int = 6, table: str = "match_brain", @@ -33,11 +37,14 @@ class CustomSupabaseVectorStore(SupabaseVectorStore): ) -> [dict]: vectors = self._embedding.embed_documents([query]) query_embedding = vectors[0] + logger.info("🤯🤯") + res = self._client.rpc( table, { "query_embedding": query_embedding, "match_count": k, + "p_user_id": str(self.user_id), }, ).execute() diff --git a/supabase/migrations/20240121195523_fix-public.sql b/supabase/migrations/20240121195523_fix-public.sql new file mode 100644 index 000000000..1f2454e81 --- /dev/null +++ b/supabase/migrations/20240121195523_fix-public.sql @@ -0,0 +1,34 @@ +create type "public"."thumbs" as enum ('up', 'down'); + +drop function if exists "public"."match_brain"(query_embedding vector, match_count integer); + +alter table "public"."chat_history" add column "user_feedback" thumbs; + +set check_function_bodies = off; + +CREATE OR REPLACE FUNCTION public.match_brain(query_embedding vector, match_count integer, p_user_id uuid) + RETURNS TABLE(id uuid, name text, similarity double precision) + LANGUAGE plpgsql +AS $function$ +BEGIN + RETURN QUERY + SELECT + b.brain_id, + b.name, + 1 - (b.meaning <=> query_embedding) as similarity + FROM + brains b + LEFT JOIN + brains_users bu ON b.brain_id = bu.brain_id + WHERE + (b.status = 'public') OR + (bu.user_id = p_user_id AND bu.rights IN ('Owner', 'Editor', 'Viewer')) + ORDER BY + b.meaning <=> query_embedding + LIMIT + match_count; +END; +$function$ +; + +