feat: execute sql & generate chart with history

Author: ulleo

backend/apps/chat/api/chat.py

@@ -1,12 +1,12 @@
-import asyncio
 import json
+import traceback
 from typing import List
 
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
 from sqlmodel import select
 
-from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, save_question, save_answer, rename_chat, \
+from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, rename_chat, \
     delete_chat, list_records
 from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat, ChatQuestion
 from apps.chat.task.llm import LLMService
@@ -79,7 +79,6 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
     Returns:
         Streaming response with analysis results
     """
-    question = request_question.question
 
     chat = session.query(Chat).filter(Chat.id == request_question.chat_id).first()
     if not chat:
@@ -113,46 +112,51 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
                                                      chart_id=request_question.chat_id)
     # get schema
     request_question.db_schema = get_table_schema(session=session, ds=ds)
-    llm_service = LLMService(request_question, history_records, ds, aimodel)
+    llm_service = LLMService(request_question, history_records, CoreDatasource(**ds.model_dump()), aimodel)
 
     llm_service.init_record(session=session, current_user=current_user)
 
     def run_task():
-        sql_res = llm_service.generate_sql(session=session)
-        for chunk in sql_res:
-            yield json.dumps({'content': chunk, 'type': 'sql'}) + '\n\n'
-        yield json.dumps({'type': 'info', 'msg': 'sql generated'}) + '\n\n'
-
-    # async def event_generator():
-    #     all_text = ''
-    #     try:
-    #         async for chunk in llm_service.async_generate(question, request_question.db_schema):
-    #             data = json.loads(chunk.replace('data: ', ''))
-    #
-    #             if data['type'] in ['final', 'tool_result']:
-    #                 content = data['content']
-    #                 print('--  ' + content)
-    #                 all_text += content
-    #                 for char in content:
-    #                     yield f"data: {json.dumps({'type': 'char', 'content': char})}\n\n"
-    #                     await asyncio.sleep(0.05)
-    #
-    #                 if 'html' in data:
-    #                     yield f"data: {json.dumps({'type': 'html', 'content': data['html']})}\n\n"
-    #             else:
-    #                 yield chunk
-    #
-    #     except Exception as e:
-    #         all_text = 'Exception:' + str(e)
-    #         yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n"
-    #
-    #     try:
-    #         save_answer(session=session, id=record.id, answer=all_text)
-    #     except Exception as e:
-    #         raise HTTPException(
-    #             status_code=500,
-    #             detail=str(e)
-    #         )
-
-    # return EventSourceResponse(event_generator(), headers={"Content-Type": "text/event-stream"})
+        try:
+            # generate sql
+            sql_res = llm_service.generate_sql(session=session)
+            full_sql_text = ''
+            for chunk in sql_res:
+                full_sql_text += chunk
+                yield json.dumps({'content': chunk, 'type': 'sql-result'}, ensure_ascii=False) + '\n\n'
+            yield json.dumps({'type': 'info', 'msg': 'sql generated'}) + '\n\n'
+
+            # filter sql
+            print(full_sql_text)
+            sql = llm_service.check_save_sql(session=session, res=full_sql_text)
+            print(sql)
+            yield json.dumps({'content': sql, 'type': 'sql'}) + '\n\n'
+
+            # execute sql
+            result = llm_service.execute_sql(sql=sql)
+            llm_service.save_sql_data(session=session, data_obj=result)
+            yield json.dumps({'content': result, 'type': 'sql-data'}, ensure_ascii=False) + '\n\n'
+
+            # generate chart
+            chart_res = llm_service.generate_chart(session=session)
+            full_chart_text = ''
+            for chunk in chart_res:
+                full_chart_text += chunk
+                yield json.dumps({'content': chunk, 'type': 'chart-result'}, ensure_ascii=False) + '\n\n'
+            yield json.dumps({'type': 'info', 'msg': 'chart generated'}) + '\n\n'
+
+            # filter chart
+            print(full_chart_text)
+            chart = llm_service.check_save_chart(session=session, res=full_chart_text)
+            print(chart)
+            yield json.dumps({'content': chart, 'type': 'chart'}, ensure_ascii=False) + '\n\n'
+
+            llm_service.finish(session=session)
+            yield json.dumps({'type': 'finish'})
+
+        except Exception as e:
+            traceback.print_exc()
+            llm_service.save_error(session=session, message=str(e))
+            yield json.dumps({'content': str(e), 'type': 'error'}, ensure_ascii=False) + '\n\n'
+
     return StreamingResponse(run_task(), media_type="text/event-stream")

backend/apps/chat/curd/chat.py

@@ -150,11 +150,35 @@ def save_full_sql_message_and_answer(session: SessionDep, record_id: int, answer
     return result
 
 
-def save_full_chart_message(session: SessionDep, id: int, full_message: str) -> ChatRecord:
-    if not id:
+def save_sql(session: SessionDep, record_id: int, sql: str) -> ChatRecord:
+    if not record_id:
         raise Exception("Record id cannot be None")
-    record = session.query(ChatRecord).filter(ChatRecord.id == id).first()
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.sql = sql
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+
+def save_full_chart_message(session: SessionDep, record_id: int, full_message: str) -> ChatRecord:
+    return save_full_chart_message_and_answer(session=session, record_id=record_id, full_message=full_message,
+                                              answer='')
+
+
+def save_full_chart_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_chart_message = full_message
+    record.chart_answer = answer
 
     result = ChatRecord(**record.model_dump())
 
@@ -167,12 +191,62 @@ def save_full_chart_message(session: SessionDep, id: int, full_message: str) ->
     return result
 
 
-def save_answer(session: SessionDep, id: int, answer: str) -> ChatRecord:
-    if not id:
+def save_chart(session: SessionDep, record_id: int, chart: str) -> ChatRecord:
+    if not record_id:
         raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.chart = chart
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
 
-    record = session.query(ChatRecord).filter(ChatRecord.id == id).first()
-    record.answer = answer
+
+def save_error_message(session: SessionDep, record_id: int, 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.error = message
+    record.finish = True
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+
+def save_sql_exec_data(session: SessionDep, record_id: int, data: str) -> ChatRecord:
+    if not record_id:
+        raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.data = data
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+def finish_record(session: SessionDep, record_id: int) -> ChatRecord:
+    if not record_id:
+        raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.finish = True
 
     result = ChatRecord(**record.model_dump())
 
@@ -183,3 +257,4 @@ def save_answer(session: SessionDep, id: int, answer: str) -> ChatRecord:
     session.commit()
 
     return result
+