feat(brains): added now multiple brains close by (#2039)

# Description

Please include a summary of the changes and the related issue. Please
also include relevant motivation and context.

## Checklist before requesting a review

Please delete options that are not relevant.

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my code
- [ ] I have commented hard-to-understand areas
- [ ] I have ideally added tests that prove my fix is effective or that
my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged

## Screenshots (if appropriate):

backend/llm/knowledge_brain_qa.py

@@ -63,6 +63,7 @@ class KnowledgeBrainQA(BaseModel, QAInterface):
     max_tokens: int = 256
     streaming: bool = False
     knowledge_qa: Optional[RAGInterface]
+    metadata: Optional[dict] = None
 
     callbacks: List[
         AsyncIteratorCallbackHandler
@@ -77,6 +78,7 @@ class KnowledgeBrainQA(BaseModel, QAInterface):
         chat_id: str,
         streaming: bool = False,
         prompt_id: Optional[UUID] = None,
+        metadata: Optional[dict] = None,
         **kwargs,
     ):
         super().__init__(
@@ -94,6 +96,7 @@ class KnowledgeBrainQA(BaseModel, QAInterface):
             streaming=streaming,
             **kwargs,
         )
+        self.metadata = metadata
 
     @property
     def prompt_to_use(self):
@@ -260,7 +263,7 @@ class KnowledgeBrainQA(BaseModel, QAInterface):
                     if self.prompt_to_use
                     else None,
                     "brain_name": brain.name if brain else None,
-                    "sources": None,
+                    "metadata": self.metadata,
                 }
             )
         else:
@@ -275,6 +278,7 @@ class KnowledgeBrainQA(BaseModel, QAInterface):
                     if self.prompt_to_use
                     else None,
                     "brain_name": brain.name if brain else None,
+                    "metadata": self.metadata,
                 }
             )
 
@@ -302,7 +306,7 @@ class KnowledgeBrainQA(BaseModel, QAInterface):
                 # Create metadata if it doesn't exist
                 if not streamed_chat_history.metadata:
                     streamed_chat_history.metadata = {}
-                    streamed_chat_history.metadata["sources"] = sources_list
+                streamed_chat_history.metadata["sources"] = sources_list
                 yield f"data: {json.dumps(streamed_chat_history.dict())}"
             else:
                 logger.info(

backend/modules/chat/controller/chat/brainful_chat.py

@@ -52,6 +52,7 @@ class BrainfulChat(ChatInterface):
         chat_question,
     ):
         brain_id_to_use = brain_id
+        metadata = {}
         if not brain_id:
             brain_settings = BrainSettings()
             supabase_client = get_supabase_client()
@@ -71,7 +72,16 @@ class BrainfulChat(ChatInterface):
             history = chat_service.get_chat_history(chat_id)
             if history:
                 question = history[0].user_message
-            brain_id_to_use = vector_store.find_brain_closest_query(question)
+            list_brains = vector_store.find_brain_closest_query(question)
+            if list_brains:
+                brain_id_to_use = list_brains[0]["id"]
+            else:
+                brain_id_to_use = None
+            # Add to metadata close_brains and close_brains_similarity
+            metadata["close_brains"] = list_brains
+
+        follow_up_questions = chat_service.get_follow_up_question(chat_id)
+        metadata["follow_up_questions"] = follow_up_questions
 
         brain = brain_service.get_brain_by_id(brain_id_to_use)
         if (
@@ -87,6 +97,7 @@ class BrainfulChat(ChatInterface):
                 brain_id=str(brain.brain_id),
                 streaming=streaming,
                 prompt_id=prompt_id,
+                metadata=metadata,
             )
         if brain.brain_type == BrainType.COMPOSITE:
             return CompositeBrainQA(
@@ -98,6 +109,7 @@ class BrainfulChat(ChatInterface):
                 streaming=streaming,
                 prompt_id=prompt_id,
                 user_id=user_id,
+                metadata=metadata,
             )
 
         if brain.brain_type == BrainType.API:
@@ -110,4 +122,5 @@ class BrainfulChat(ChatInterface):
                 streaming=streaming,
                 prompt_id=prompt_id,
                 user_id=user_id,
+                metadata=metadata,
             )

backend/modules/chat/service/chat_service.py

@@ -1,3 +1,4 @@
+import random
 from typing import List, Optional
 from uuid import UUID
 
@@ -46,6 +47,25 @@ class ChatService:
 
         return insert_response.data[0]
 
+    def get_follow_up_question(
+        self, brain_id: UUID = None, question: str = None
+    ) -> [str]:
+        follow_up = [
+            "Summarize the conversation",
+            "Explain in more detail",
+            "Explain like I'm 5",
+            "Provide a list",
+            "Give examples",
+            "Use simpler language",
+            "Elaborate on a specific point",
+            "Provide pros and cons",
+            "Break down into steps",
+            "Illustrate with an image or diagram",
+        ]
+        # Return 3 random follow up questions amongs the list
+        random3 = random.sample(follow_up, 3)
+        return random3
+
     def add_question_and_answer(
         self, chat_id: UUID, question_and_answer: QuestionAndAnswer
     ) -> Optional[Chat]:

backend/vectorstore/supabase.py

@@ -1,5 +1,4 @@
 from typing import Any, List
-from uuid import UUID
 
 from langchain.docstore.document import Document
 from langchain.embeddings.base import Embeddings
@@ -31,7 +30,7 @@ class CustomSupabaseVectorStore(SupabaseVectorStore):
         k: int = 6,
         table: str = "match_brain",
         threshold: float = 0.5,
-    ) -> UUID | None:
+    ) -> [dict]:
         vectors = self._embedding.embed_documents([query])
         query_embedding = vectors[0]
         res = self._client.rpc(
@@ -43,10 +42,16 @@ class CustomSupabaseVectorStore(SupabaseVectorStore):
         ).execute()
 
         # Get the brain_id of the brain that is most similar to the query
-        brain_id = res.data[0].get("id", None)
-        if not brain_id:
-            return None
-        return str(brain_id)
+        # Get the brain_id and name of the brains that are most similar to the query
+        brain_details = [
+            {
+                "id": item.get("id", None),
+                "name": item.get("name", None),
+                "similarity": item.get("similarity", 0.0),
+            }
+            for item in res.data
+        ]
+        return brain_details
 
     def similarity_search(
         self,