feat: improve chat template

Author: ulleo

backend/apps/chat/models/chat_model.py

@@ -176,15 +176,16 @@ def sql_sys_question(self):
         return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question,
                                                    lang=self.lang, terminologies=self.terminologies)
 
-    def sql_user_question(self):
+    def sql_user_question(self, current_time: str):
         return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question,
-                                                 rule=self.rule)
+                                                 rule=self.rule, current_time=current_time)
 
     def chart_sys_question(self):
         return get_chart_template()['system'].format(sql=self.sql, question=self.question, lang=self.lang)
 
-    def chart_user_question(self):
-        return get_chart_template()['user'].format(sql=self.sql, question=self.question, rule=self.rule)
+    def chart_user_question(self, chart_type: Optional[str] = None):
+        return get_chart_template()['user'].format(sql=self.sql, question=self.question, rule=self.rule,
+                                                   chart_type=chart_type)
 
     def analysis_sys_question(self):
         return get_analysis_template()['system'].format(lang=self.lang, terminologies=self.terminologies)

backend/apps/chat/task/llm.py

@@ -3,6 +3,7 @@
 import traceback
 import warnings
 from concurrent.futures import ThreadPoolExecutor, Future
+from datetime import datetime
 from typing import Any, List, Optional, Union, Dict
 
 import numpy as np
@@ -70,7 +71,8 @@ class LLMService:
     future: Future
 
     def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
-                 current_assistant: Optional[CurrentAssistant] = None, no_reasoning: bool = False, config: LLMConfig = None):
+                 current_assistant: Optional[CurrentAssistant] = None, no_reasoning: bool = False,
+                 config: LLMConfig = None):
         self.chunk_list = []
         engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))
         session_maker = sessionmaker(bind=engine)
@@ -126,7 +128,7 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
         self.llm = llm_instance.llm
 
         self.init_messages()
-    
+
     @classmethod
     async def create(cls, *args, **kwargs):
         config: LLMConfig = await get_default_config()
@@ -503,7 +505,8 @@ def select_datasource(self):
 
     def generate_sql(self):
         # append current question
-        self.sql_message.append(HumanMessage(self.chat_question.sql_user_question()))
+        self.sql_message.append(HumanMessage(
+            self.chat_question.sql_user_question(current_time=datetime.now().strftime('%Y-%m-%d %H:%M:%S'))))
 
         self.current_logs[OperationEnum.GENERATE_SQL] = start_log(session=self.session,
                                                                   ai_modal_id=self.chat_question.ai_modal_id,
@@ -670,9 +673,9 @@ def generate_assistant_filter(self, sql, tables: List):
             return None
         return self.build_table_filter(sql=sql, filters=filters)
 
-    def generate_chart(self):
+    def generate_chart(self, chart_type: Optional[str] = ''):
         # append current question
-        self.chart_message.append(HumanMessage(self.chat_question.chart_user_question()))
+        self.chart_message.append(HumanMessage(self.chat_question.chart_user_question(chart_type)))
 
         self.current_logs[OperationEnum.GENERATE_CHART] = start_log(session=self.session,
                                                                     ai_modal_id=self.chat_question.ai_modal_id,
@@ -714,7 +717,8 @@ def generate_chart(self):
                                                                   reasoning_content=full_thinking_text,
                                                                   token_usage=token_usage)
 
-    def check_sql(self, res: str) -> tuple[any]:
+    @staticmethod
+    def check_sql(res: str) -> tuple[str, Optional[list]]:
         json_str = extract_nested_json(res)
         if json_str is None:
             raise SingleMessageError(orjson.dumps({'message': 'Cannot parse sql from answer',
@@ -739,6 +743,26 @@ def check_sql(self, res: str) -> tuple[any]:
             raise SingleMessageError("SQL query is empty")
         return sql, data.get('tables')
 
+    @staticmethod
+    def get_chart_type_from_sql_answer(res: str) -> Optional[str]:
+        json_str = extract_nested_json(res)
+        if json_str is None:
+            return None
+
+        chart_type: Optional[str]
+        data: dict
+        try:
+            data = orjson.loads(json_str)
+
+            if data['success']:
+                chart_type = data['chart-type']
+            else:
+                return None
+        except Exception:
+            return None
+
+        return chart_type
+
     def check_save_sql(self, res: str) -> str:
         sql, *_ = self.check_sql(res=res)
         save_sql(session=self.session, sql=sql, record_id=self.record.id)
@@ -921,6 +945,9 @@ def run_task(self, in_chat: bool = True):
 
             # filter sql
             SQLBotLogUtil.info(full_sql_text)
+
+            chart_type = self.get_chart_type_from_sql_answer(full_sql_text)
+
             use_dynamic_ds: bool = self.current_assistant and self.current_assistant.type in dynamic_ds_types
 
             # todo row permission
@@ -962,7 +989,7 @@ def run_task(self, in_chat: bool = True):
                 yield 'data:' + orjson.dumps({'content': 'execute-success', 'type': 'sql-data'}).decode() + '\n\n'
 
             # generate chart
-            chart_res = self.generate_chart()
+            chart_res = self.generate_chart(chart_type)
             full_chart_text = ''
             for chunk in chart_res:
                 full_chart_text += chunk.get('content')