feat: chat view

backend/apps/chat/models/chat_model.py

@@ -70,18 +70,25 @@ class AiModelQuestion(BaseModel):
     engine: str = ""
     db_schema: str = ""
     sql: str = ""
+    rule: str = """
+    请逐步推理后给出答案：
+    推理过程中不需要输出JSON，仅在最终结果内输出符合要求的JSON
+    步骤1: [思考内容]
+    步骤2: [思考内容]
+    最终答案: [结果]
+    """
 
     def sql_sys_question(self):
         return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question)
 
     def sql_user_question(self):
-        return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question)
+        return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question, rule=self.rule)
 
     def chart_sys_question(self):
         return get_chart_template()['system'].format(sql=self.sql, question=self.question)
 
     def chart_user_question(self):
-        return get_chart_template()['user'].format(sql=self.sql, question=self.question)
+        return get_chart_template()['user'].format(sql=self.sql, question=self.question, rule=self.rule)
 
 
 class ChatQuestion(AiModelQuestion):

backend/apps/chat/task/llm.py

@@ -1,9 +1,7 @@
-import asyncio
 import json
 import logging
-import re
 import warnings
-from typing import Any, List, Union, AsyncGenerator, Dict
+from typing import Any, List, Union, Dict
 
 from langchain_community.utilities import SQLDatabase
 from langchain_core.language_models import BaseLLM
@@ -21,6 +19,8 @@ from common.core.deps import SessionDep, CurrentUser
 
 warnings.filterwarnings("ignore")
 
+base_message_count_limit = 5
+
 
 class LLMService:
     ds: CoreDatasource
@@ -55,34 +55,51 @@ class LLMService:
             last_chart_messages[-1].full_chart_message
 
         last_sql_messages: List[dict[str, Any]] = json.loads(last_sql_message_str)
+
+        # todo maybe can configure
+        count_limit = 0 - base_message_count_limit
+
+        self.sql_message = []
         if last_sql_messages is None or len(last_sql_messages) == 0:
             # add sys prompt
             self.sql_message.append(SystemMessage(content=self.chat_question.sql_sys_question()))
         else:
+            # limit count
             for last_sql_message in last_sql_messages:
+                if last_sql_message['type'] == 'system':
+                    _msg = SystemMessage(content=last_sql_message['content'])
+                    self.sql_message.append(_msg)
+                    break
+            for last_sql_message in last_sql_messages[count_limit:]:
                 _msg: BaseMessage
                 if last_sql_message['type'] == 'human':
                     _msg = HumanMessage(content=last_sql_message['content'])
+                    self.sql_message.append(_msg)
                 elif last_sql_message['type'] == 'ai':
                     _msg = AIMessage(content=last_sql_message['content'])
-                else:
-                    _msg = SystemMessage(content=last_sql_message['content'])
-                self.sql_message.append(_msg)
+                    self.sql_message.append(_msg)
 
         last_chart_messages: List[dict[str, Any]] = json.loads(last_chart_message_str)
+
+        self.chart_message = []
         if last_chart_messages is None or len(last_chart_messages) == 0:
             # add sys prompt
             self.chart_message.append(SystemMessage(content=self.chat_question.chart_sys_question()))
         else:
+            # limit count
+            for last_chart_message in last_chart_messages:
+                if last_chart_message['type'] == 'system':
+                    _msg = SystemMessage(content=last_chart_message['content'])
+                    self.chart_message.append(_msg)
+                    break
             for last_chart_message in last_chart_messages:
                 _msg: BaseMessage
                 if last_chart_message['type'] == 'human':
                     _msg = HumanMessage(content=last_chart_message['content'])
+                    self.chart_message.append(_msg)
                 elif last_chart_message['type'] == 'ai':
                     _msg = AIMessage(content=last_chart_message['content'])
-                else:
-                    _msg = SystemMessage(content=last_chart_message['content'])
-                self.chart_message.append(_msg)
+                    self.chart_message.append(_msg)
 
     def init_record(self, session: SessionDep, current_user: CurrentUser) -> ChatRecord:
         self.record = save_question(session=session, current_user=current_user, question=self.chat_question)
@@ -98,6 +115,7 @@ class LLMService:
         full_sql_text = ''
         res = self.llm.stream(self.sql_message)
         for chunk in res:
+            print(chunk)
             if isinstance(chunk, dict):
                 full_sql_text += chunk['content']
                 yield chunk['content']
@@ -196,9 +214,10 @@ class LLMService:
         return save_error_message(session=session, record_id=self.record.id, message=message)
 
     def save_sql_data(self, session: SessionDep, data_obj: Dict[str, Any]):
-        return save_sql_exec_data(session=session, record_id=self.record.id, data=json.dumps(data_obj, ensure_ascii=False))
+        return save_sql_exec_data(session=session, record_id=self.record.id,
+                                  data=json.dumps(data_obj, ensure_ascii=False))
 
-    def finish(self,session: SessionDep):
+    def finish(self, session: SessionDep):
         return finish_record(session=session, record_id=self.record.id)
 
     def execute_sql(self, sql: str):
@@ -237,7 +256,7 @@ def extract_nested_json(text):
                         pass
             else:
                 stack = []  # 括号不匹配则重置
-    if results[0]:
+    if len(results) > 0 and results[0]:
         return results[0]
     return None
 

backend/template.yaml

@@ -15,17 +15,23 @@ template:
       {{"success":false,"message":"抱歉，我无法回答您的问题。"}}
       - 如果根据提供的表结构不能生成符合问题与条件的SQL，回答：
       {{"success":false,"message":"无法生成SQL的原因"}}
-    user: |  
-      表结构:
+      - 如果问题是图表展示相关且与生成SQL查询无关时，请参考上一次回答的SQL来生成SQL
+      - 如果问题是图表展示相关，可参考的图表类型为表格(table)、条形图(bar)、折线图(line)或饼图(pie)，返回的JSON：
+       {{"success":true,"sql":"生成的SQL语句","chart-type":"选择的图表类型(table/bar/line/pie)"}}
+    user: |
+      ### 表结构:
       {schema}
       
-      问题:
+      ### 问题:
       {question}
+      
+      ### 其他规则:
+      {rule}
   chart:
     system: |
       ### 说明：
       您的任务是通过给定的问题和SQL生成 JSON 以进行数据可视化。
-      请遵守以下规则：
+      请遵守以下规则:
       - 如果需要表格，则生成的 JSON 格式应为：
       {{"type":"table", "title": "标题", "columns": [{{"name":"中文字段名1", "value": "SQL 查询列 1(有别名用别名)"}}, {{"name": "中文字段名 2", "value": "SQL 查询列 2(有别名用别名)"}}]}}
       必须从 SQL 查询列中提取“columns”。
@@ -36,22 +42,25 @@ template:
       {{"type":"line", "title": "标题", "axis": {{"x": {{"name":"x轴的中文名称","value": "x轴的SQL查询列(有别名用别名)"}}, "y": {{"name":"y轴的中文名称","value": "y轴的SQL查询列(有别名用别名)"}}}}
       其中“x”和“y”必须从SQL查询列中提取。
       - 如果需要饼图，则生成的 JSON 格式应为：
-      {{"type":"pie", "title": "标题", "column": {{"name":"中文字段名","value":"SQL查询列1(有别名用别名)"}}}}
+      {{"type":"pie", "title": "标题", "column": [{{"name":"中文字段名1","value":"SQL查询列1(有别名用别名)"}},{{"name":"中文字段名2","value":"SQL查询列2(有别名用别名)"}}]}}
       其中“column”必须从SQL查询列中提取。
       - 如果答案未知或者与生成JSON无关，则生成的 JSON 格式应为：
       {{"type":"error", "reason": "抱歉，我无法回答您的问题。"}}
       
-      ### 示例
+      ### 示例:
       如果 SQL 为： SELECT products_sales_data.category, AVG(products_sales_data.price) AS average_price FROM products_sales_data GROUP BY products_sales_data.category;
       问题是：每个商品分类的平均价格
       则生成的 JSON 可以是： {{"type":"table", "title": "每个商品分类的平均价格", "columns": [{{"name":"商品分类","value":"category"}}, {{"name":"平均价格","value":"average_price"}}]}}
       
-      ### 响应：
+      ### 响应:
       根据您的指示，这是我生成的与 问题 和 sql 匹配的 JSON：
       ```json
     user: |    
       ### SQL:
       {sql}
       
-      ### 问题：
-      {question}
+      ### 问题:
+      {question}
+      
+      ### 其他规则:
+      {rule}