feat: add chat recommend questions

Author: ulleo

backend/apps/chat/api/chat.py

@@ -1,13 +1,12 @@
 import traceback
 
-import orjson
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
 
 from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, rename_chat, \
     delete_chat
 from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion
-from apps.chat.task.llm import LLMService, run_task
+from apps.chat.task.llm import LLMService, run_task, run_analysis_or_predict_task, run_recommend_questions_task
 from common.core.deps import SessionDep, CurrentUser
 
 router = APIRouter(tags=["Data Q&A"], prefix="/chat")
@@ -62,6 +61,26 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat
         )
 
 
+@router.get("/recommend_questions/{chat_record_id}")
+async def recommend_questions(session: SessionDep, current_user: CurrentUser, chat_record_id: int):
+    try:
+        record = session.query(ChatRecord).get(chat_record_id)
+        if not record:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Chat record with id {chat_record_id} not found"
+            )
+        request_question = ChatQuestion(chat_id=record.chat_id, question=record.question if record.question else '')
+
+        llm_service = LLMService(session, current_user, request_question)
+        llm_service.set_record(record)
+
+        return run_recommend_questions_task(llm_service)
+    except Exception as e:
+        traceback.print_exc()
+        return '[]'
+
+
 @router.post("/question")
 async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion):
     """Stream SQL analysis results
@@ -88,61 +107,38 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
     return StreamingResponse(run_task(llm_service, session), media_type="text/event-stream")
 
 
-@router.post("/record/{chart_record_id}/{action_type}")
-async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chart_record_id: int, action_type: str):
+@router.post("/record/{chat_record_id}/{action_type}")
+async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str):
     if action_type != 'analysis' and action_type != 'predict':
         raise HTTPException(
             status_code=404,
             detail="Not Found"
         )
 
-    record = session.query(ChatRecord).get(chart_record_id)
+    record = session.query(ChatRecord).get(chat_record_id)
     if not record:
         raise HTTPException(
             status_code=400,
-            detail=f"Chat record with id {chart_record_id} not found"
+            detail=f"Chat record with id {chat_record_id} not found"
         )
 
     if not record.chart:
         raise HTTPException(
             status_code=500,
-            detail=f"Chat record with id {chart_record_id} has not generated chart, do not support to analyze it"
+            detail=f"Chat record with id {chat_record_id} has not generated chart, do not support to analyze it"
         )
 
     request_question = ChatQuestion(chat_id=record.chat_id, question='')
 
-    llm_service = LLMService(session, current_user, request_question)
-    llm_service.set_record(record)
-
-    def run_task():
-        try:
-            if action_type == 'analysis':
-                # generate analysis
-                analysis_res = llm_service.generate_analysis(session=session)
-                for chunk in analysis_res:
-                    yield orjson.dumps({'content': chunk, 'type': 'analysis-result'}).decode() + '\n\n'
-                yield orjson.dumps({'type': 'info', 'msg': 'analysis generated'}).decode() + '\n\n'
-
-                yield orjson.dumps({'type': 'analysis_finish'}).decode() + '\n\n'
-
-            elif action_type == 'predict':
-                # generate predict
-                analysis_res = llm_service.generate_predict()
-                full_text = ''
-                for chunk in analysis_res:
-                    yield orjson.dumps({'content': chunk, 'type': 'predict-result'}).decode() + '\n\n'
-                    full_text += chunk
-                yield orjson.dumps({'type': 'info', 'msg': 'predict generated'}).decode() + '\n\n'
-
-                _data = llm_service.check_save_predict_data(res=full_text)
-                yield orjson.dumps({'type': 'predict', 'content': _data}).decode() + '\n\n'
-
-                yield orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n'
-
-
-        except Exception as e:
-            traceback.print_exc()
-            # llm_service.save_error(session=session, message=str(e))
-            yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
+    try:
+        llm_service = LLMService(session, current_user, request_question)
+        llm_service.set_record(record)
+    except Exception as e:
+        traceback.print_exc()
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )
 
-    return StreamingResponse(run_task(), media_type="text/event-stream")
+    return StreamingResponse(run_analysis_or_predict_task(llm_service, action_type),
+                             media_type="text/event-stream")

backend/apps/chat/curd/chat.py

@@ -1,12 +1,13 @@
 import datetime
 from typing import List
 
-from sqlalchemy import and_
+from sqlalchemy import and_, distinct
 from sqlalchemy.orm import load_only
 
 from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion
 from apps.datasource.models.datasource import CoreDatasource
 from common.core.deps import SessionDep, CurrentUser
+from common.utils.utils import extract_nested_json
 
 
 def list_chats(session: SessionDep, current_user: CurrentUser) -> List[Chat]:
@@ -130,8 +131,6 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
         _record.id = record.id
         session.commit()
 
-        # todo suggest questions
-
         chat_info.records.append(_record)
 
     return chat_info
@@ -252,6 +251,36 @@ def save_full_select_datasource_message_and_answer(session: SessionDep, record_i
     return result
 
 
+def save_full_recommend_question_message_and_answer(session: SessionDep, record_id: int, answer: str,
+                                                    full_message: str) -> ChatRecord:
+    if not record_id:
+        raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.full_recommended_question_message = full_message
+    record.recommended_question_answer = answer
+
+    json_str = '[]'
+    if answer and answer != '':
+        try:
+            json_str = extract_nested_json(answer)
+
+            if not json_str:
+                json_str = '[]'
+        except Exception as e:
+            pass
+    record.recommended_question = json_str
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+
 def save_sql(session: SessionDep, record_id: int, sql: str) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")
@@ -379,3 +408,12 @@ def finish_record(session: SessionDep, record_id: int) -> ChatRecord:
     session.commit()
 
     return result
+
+
+def get_old_questions(session: SessionDep, datasource: int):
+    if not datasource:
+        return []
+    records = session.query(ChatRecord.question, ChatRecord.create_time).filter(ChatRecord.datasource == datasource,
+                                                                  ChatRecord.question != None).order_by(
+        ChatRecord.create_time.desc()).limit(20).all()
+    return records

backend/apps/chat/models/chat_model.py

@@ -130,10 +130,10 @@ def datasource_sys_question(self):
     def datasource_user_question(self, datasource_list: str = "[]"):
         return get_datasource_template()['user'].format(question=self.question, data=datasource_list, lang=self.lang)
 
-    def datasource_guess_sys_question(self):
+    def guess_sys_question(self):
         return get_guess_question_template()['system']
 
-    def datasource_guess_user_question(self, old_questions: str = "[]"):
+    def guess_user_question(self, old_questions: str = "[]"):
         return get_guess_question_template()['user'].format(question=self.question, schema=self.db_schema,
                                                             old_questions=old_questions, lang=self.lang)
 

backend/apps/chat/task/llm.py

@@ -17,13 +17,15 @@
 from apps.chat.curd.chat import save_question, save_full_sql_message, save_full_sql_message_and_answer, save_sql, \
     save_error_message, save_sql_exec_data, save_full_chart_message, save_full_chart_message_and_answer, save_chart, \
     finish_record, save_full_analysis_message_and_answer, save_full_predict_message_and_answer, save_predict_data, \
-    save_full_select_datasource_message_and_answer, list_records
+    save_full_select_datasource_message_and_answer, list_records, save_full_recommend_question_message_and_answer, \
+    get_old_questions
 from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat
 from apps.datasource.crud.datasource import get_table_schema
 from apps.datasource.models.datasource import CoreDatasource
 from apps.db.db import exec_sql
 from common.core.config import settings
 from common.core.deps import SessionDep, CurrentUser
+from common.utils.utils import extract_nested_json
 
 warnings.filterwarnings("ignore")
 
@@ -59,7 +61,6 @@ def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question
 
             chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL'
 
-
         history_records: List[ChatRecord] = list(
             map(lambda x: ChatRecord(**x.model_dump()), filter(lambda r: True if r.first_chat != True else False,
                                                                list_records(session=self.session,
@@ -75,7 +76,6 @@ def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question
         self.chat_question = chat_question
         self.config = get_default_config()
         self.chat_question.ai_modal_id = self.config.model_id
-        
 
         # Create LLM instance through factory
         llm_instance = LLMFactory.create_llm(self.config)
@@ -176,7 +176,7 @@ def get_fields_from_chart(self):
                     fields.append(column_str)
         return fields
 
-    def generate_analysis(self, session: SessionDep):
+    def generate_analysis(self):
         fields = self.get_fields_from_chart()
 
         self.chat_question.fields = orjson.dumps(fields).decode()
@@ -189,7 +189,7 @@ def generate_analysis(self, session: SessionDep):
         if self.record.full_analysis_message and self.record.full_analysis_message.strip() != '':
             history_msg = orjson.loads(self.record.full_analysis_message)
 
-        self.record = save_full_analysis_message_and_answer(session=session, record_id=self.record.id, answer='',
+        self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id, answer='',
                                                             full_message=orjson.dumps(history_msg +
                                                                                       [{'type': msg.type,
                                                                                         'content': msg.content} for msg
@@ -210,7 +210,7 @@ def generate_analysis(self, session: SessionDep):
                 continue
 
         analysis_msg.append(AIMessage(full_analysis_text))
-        self.record = save_full_analysis_message_and_answer(session=session, record_id=self.record.id,
+        self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id,
                                                             answer=full_analysis_text,
                                                             full_message=orjson.dumps(history_msg +
                                                                                       [{'type': msg.type,
@@ -261,6 +261,47 @@ def generate_predict(self):
                                                                                       in
                                                                                       predict_msg]).decode())
 
+    def generate_recommend_questions_task(self):
+
+        # get schema
+        if self.ds and not self.chat_question.db_schema:
+            self.chat_question.db_schema = get_table_schema(session=self.session, ds=self.ds)
+
+        guess_msg: List[Union[BaseMessage, dict[str, Any]]] = []
+        guess_msg.append(SystemMessage(content=self.chat_question.guess_sys_question()))
+        #  todo old questions
+        old_questions = list(map(lambda q: q[0].strip(), get_old_questions(self.session, self.record.datasource)))
+        guess_msg.append(HumanMessage(content=self.chat_question.guess_user_question(orjson.dumps(old_questions).decode())))
+
+        self.record = save_full_recommend_question_message_and_answer(session=self.session, record_id=self.record.id,
+                                                                      answer='',
+                                                                      full_message=orjson.dumps([{'type': msg.type,
+                                                                                                  'content': msg.content}
+                                                                                                 for msg
+                                                                                                 in
+                                                                                                 guess_msg]).decode())
+
+        full_guess_text = ''
+        res = self.llm.stream(guess_msg)
+        for chunk in res:
+            print(chunk)
+            if isinstance(chunk, dict):
+                full_guess_text += chunk['content']
+                continue
+            if isinstance(chunk, AIMessageChunk):
+                full_guess_text += chunk.content
+                continue
+
+        guess_msg.append(AIMessage(full_guess_text))
+        self.record = save_full_recommend_question_message_and_answer(session=self.session, record_id=self.record.id,
+                                                                      answer=full_guess_text,
+                                                                      full_message=orjson.dumps([{'type': msg.type,
+                                                                                                  'content': msg.content}
+                                                                                                 for msg
+                                                                                                 in
+                                                                                                 guess_msg]).decode())
+        return self.record.recommended_question
+
     def select_datasource(self):
         datasource_msg: List[Union[BaseMessage, dict[str, Any]]] = []
         datasource_msg.append(SystemMessage(self.chat_question.datasource_sys_question()))
@@ -486,33 +527,6 @@ def execute_sql(self, sql: str):
         return exec_sql(self.ds, sql)
 
 
-def extract_nested_json(text):
-    stack = []
-    start_index = -1
-    results = []
-
-    for i, char in enumerate(text):
-        if char in '{[':
-            if not stack:  # 记录起始位置
-                start_index = i
-            stack.append(char)
-        elif char in '}]':
-            if stack and ((char == '}' and stack[-1] == '{') or (char == ']' and stack[-1] == '[')):
-                stack.pop()
-                if not stack:  # 栈空时截取完整JSON
-                    json_str = text[start_index:i + 1]
-                    try:
-                        orjson.loads(json_str)  # 验证有效性
-                        results.append(json_str)
-                    except:
-                        pass
-            else:
-                stack = []  # 括号不匹配则重置
-    if len(results) > 0 and results[0]:
-        return results[0]
-    return None
-
-
 def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
     """Execute SQL query using SQLDatabase
 
@@ -647,6 +661,42 @@ def run_task(llm_service: LLMService, session: SessionDep, in_chat: bool = True)
             yield f'> ❌ **ERROR**\n\n> \n\n> {str(e)}。'
 
 
+def run_analysis_or_predict_task(llm_service: LLMService, action_type: str):
+    try:
+        if action_type == 'analysis':
+            # generate analysis
+            analysis_res = llm_service.generate_analysis()
+            for chunk in analysis_res:
+                yield orjson.dumps({'content': chunk, 'type': 'analysis-result'}).decode() + '\n\n'
+            yield orjson.dumps({'type': 'info', 'msg': 'analysis generated'}).decode() + '\n\n'
+
+            yield orjson.dumps({'type': 'analysis_finish'}).decode() + '\n\n'
+
+        elif action_type == 'predict':
+            # generate predict
+            analysis_res = llm_service.generate_predict()
+            full_text = ''
+            for chunk in analysis_res:
+                yield orjson.dumps({'content': chunk, 'type': 'predict-result'}).decode() + '\n\n'
+                full_text += chunk
+            yield orjson.dumps({'type': 'info', 'msg': 'predict generated'}).decode() + '\n\n'
+
+            _data = llm_service.check_save_predict_data(res=full_text)
+            yield orjson.dumps({'type': 'predict', 'content': _data}).decode() + '\n\n'
+
+            yield orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n'
+
+
+    except Exception as e:
+        traceback.print_exc()
+        # llm_service.save_error(session=session, message=str(e))
+        yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
+
+
+def run_recommend_questions_task(llm_service: LLMService):
+    return llm_service.generate_recommend_questions_task()
+
+
 def request_picture(chat_id: int, record_id: int, chart: dict, data: dict):
     file_name = f'c_{chat_id}_r_{record_id}'