feat: related question prompt (#560)

Author: ikkvix

kubechat/chat/utils.py

@@ -44,7 +44,7 @@ def start_response(message_id):
     )
 
 
-def stop_response(message_id, references, related_question = [], memory_count=0):
+def stop_response(message_id, references, related_question = [], related_question_prompt = '' ,memory_count=0):
     if references is None:
         references = []
     return json.dumps(
@@ -53,6 +53,7 @@ def stop_response(message_id, references, related_question = [], memory_count=0)
             "id": message_id,
             "data": references,
             "memoryCount": memory_count,
+            "related_question_prompt":related_question_prompt,
             "related_question":related_question,
             "timestamp": now_unix_milliseconds()
         }

kubechat/chat/websocket/base_consumer.py

@@ -4,6 +4,7 @@
 from abc import abstractmethod
 import ast
 from datetime import datetime
+import random
 
 import websockets
 from asgiref.sync import sync_to_async
@@ -80,12 +81,18 @@ async def connect(self):
 
         message_id = f"{now_unix_milliseconds()}"
         bot_config = json.loads(self.bot.config)
+        self.related_question_prompt = bot_config.get("related_question_prompt","")
+        
         welcome = bot_config.get("welcome",{})
         faq = welcome.get("faq",[])
         questions = []
         for qa in faq:
             questions.append(qa["question"])
-        welcome_message = {"hello":welcome.get("hello",""), "faq":questions}
+        if len(questions) < 3:
+            random_questions = questions
+        else:
+            random_questions = random.sample(questions, 3)
+        welcome_message = {"hello":welcome.get("hello",""), "faq":random_questions}
         await self.send(text_data=welcome_response(message_id, welcome_message))
 
     async def disconnect(self, close_code):
@@ -163,6 +170,6 @@ async def receive(self, text_data=None, bytes_data=None):
             if self.use_default_token and self.conversation_limit:
                 await self.manage_quota_usage()
             # send stop message
-            await self.send(text_data=stop_response(message_id, references, related_question, self.pipeline.memory_count))
+            await self.send(text_data=stop_response(message_id, references, related_question, self.related_question_prompt, self.pipeline.memory_count))
 
 

kubechat/pipeline/pipeline.py

@@ -323,6 +323,7 @@ async def run(self, message, gen_references=False, message_id=""):
         logger.info("[%s] start processing", log_prefix)
 
         references = []
+        related_questions = []
         response = ""
         vector = self.embedding_model.embed_query(message)
         logger.info("[%s] embedding query end", log_prefix)
@@ -366,8 +367,11 @@ async def run(self, message, gen_references=False, message_id=""):
                     yield self.oops
                     need_generate_answer = False
                 if self.welcome_question != []:
-                    yield KUBE_CHAT_RELATED_QUESTIONS + str(self.welcome_question)
-                    need_related_question = False
+                    if len(self.welcome_question)>=3:
+                        related_questions = random.sample(self.welcome_question, 3)
+                        need_related_question = False
+                    else:
+                        related_questions = self.welcome_question
 
             if self.use_related_question and need_related_question:
                 related_question_prompt = self.related_question_prompt.format(query=message, context=context)
@@ -405,18 +409,20 @@ async def run(self, message, gen_references=False, message_id=""):
         await self.add_ai_message(message, message_id, response, references)
         logger.info("[%s] add ai message end and the pipeline is succeed", log_prefix)
 
-        if self.use_related_question and need_related_question:
-            related_question = await related_question_task
-            related_question = re.sub(r'\n+', '\n', related_question).split('\n')
-            for i,question in enumerate(related_question):
-                match = re.match(r"\s*-\s*(.*)", question)
-                if match:
-                    question = match.group(1)
-                match = re.match(r"\s*\d+\.\s*(.*)", question)
-                if match:  
-                    question = match.group(1)
-                related_question[i] = question
-            yield KUBE_CHAT_RELATED_QUESTIONS + str(related_question[:3])
+        if self.use_related_question:
+            if need_related_question:
+                related_question_generate = await related_question_task
+                related_question = re.sub(r'\n+', '\n', related_question_generate).split('\n')
+                for i,question in enumerate(related_question):
+                    match = re.match(r"\s*-\s*(.*)", question)
+                    if match:
+                        question = match.group(1)
+                    match = re.match(r"\s*\d+\.\s*(.*)", question)
+                    if match:  
+                        question = match.group(1)
+                    related_question[i] = question
+                related_questions.extend(related_question)
+            yield KUBE_CHAT_RELATED_QUESTIONS + str(related_questions[:3])
 
         if gen_references:
             yield KUBE_CHAT_DOC_QA_REFERENCES + json.dumps(references)