imporve code for memory and prompt

Author: lwj2015

kubechat/chat/websocket/common_consumer.py

@@ -31,6 +31,7 @@ async def predict(self, query, **kwargs):
     async def receive(self, text_data=None, bytes_data=None):
         message = ""
         message_id = f"{now_unix_milliseconds()}"
+        related_question = []
 
         # 先在text_data里传file_name，再在bytes_data里传file_content，最后问问题
         if bytes_data:
@@ -95,4 +96,4 @@ 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, [], related_question, self.pipeline.memory_count))
+            await self.send(text_data=stop_response(message_id, [], related_question, self.related_question_prompt, self.pipeline.memory_count))

kubechat/pipeline/common_pipeline.py

@@ -17,11 +17,10 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
         # self.prompt_template=None
-        if not self.prompt_template:
-            if self.memory:
-                self.prompt_template = COMMON_MEMORY_TEMPLATE
-            else:
-                self.prompt_template = COMMON_TEMPLATE
+        if self.memory:
+            self.prompt_template = COMMON_MEMORY_TEMPLATE
+        elif not self.prompt_template:
+            self.prompt_template = COMMON_TEMPLATE
         self.prompt = PromptTemplate(template=self.prompt_template, input_variables=["query"])
         self.file_prompt = PromptTemplate(template=COMMON_FILE_TEMPLATE,
                                           input_variables=["query", "context"])
@@ -51,6 +50,9 @@ async def run(self, message, gen_references=False, message_id="", file=None):
         related_questions = []
         response = ""
 
+        need_generate_answer = True
+        need_related_question = True
+
         if self.oops != "":
             response = self.oops
             yield self.oops

kubechat/views/main.py

@@ -712,7 +712,7 @@ async def update_bot(request, bot_id, bot_in: BotIn):
     new_config = json.loads(bot_in.config)
     model = new_config.get("model")
     llm_config = new_config.get("llm")
-    valid, msg = validate_bot_config(model, llm_config)
+    valid, msg = validate_bot_config(model, llm_config, bot)
     if not valid:
         return fail(HTTPStatus.BAD_REQUEST, msg)
     old_config = json.loads(bot.config)

kubechat/views/utils.py

@@ -12,7 +12,7 @@
 from config import settings
 from kubechat.chat.history.redis import RedisChatMessageHistory
 from kubechat.llm.base import Predictor, PredictorType
-from kubechat.db.models import ssl_file_path, ssl_temp_file_path
+from kubechat.db.models import ssl_file_path, ssl_temp_file_path, BotType
 from kubechat.db.ops import query_chat_feedbacks, logger, PagedResult
 from kubechat.source.base import get_source, CustomSourceInitializationError
 from kubechat.utils.utils import AVAILABLE_SOURCE
@@ -82,7 +82,7 @@ def validate_source_connect_config(config: Dict) -> (bool, str):
     return True, ""
 
 
-def validate_bot_config(model, config: Dict) -> (bool, str):
+def validate_bot_config(model, config: Dict, bot) -> (bool, str):
     try:
         Predictor.from_model(model, PredictorType.CUSTOM_LLM, **config)
     except Exception as e:
@@ -91,14 +91,20 @@ def validate_bot_config(model, config: Dict) -> (bool, str):
     try:
         # validate the prompt
         prompt_template = config.get("prompt_template", None)
-        PromptTemplate(template=prompt_template, input_variables=["query", "context"])
+        if bot.type == BotType.KNOWLEDGE:
+            PromptTemplate(template=prompt_template, input_variables=["query", "context"])
+        elif bot.type == BotType.COMMON:
+            # PromptTemplate(template=prompt_template, input_variables=["query"])
+            pass
+        else:
+            return False, "Unsupported bot type"
     except ValidationError:
         return False, "Invalid prompt template"
 
     try:
         # validate the memory prompt
         prompt_template = config.get("memory_prompt_template", None)
-        if prompt_template:
+        if prompt_template and bot.type == BotType.KNOWLEDGE:
             PromptTemplate(template=prompt_template, input_variables=["query", "context"])
     except ValidationError:
         return False, "Invalid memory prompt template"