klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-10 18:26:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: execute sql & generate chart with history

Browse Source
This commit is contained in:
ulleo
2025-06-04 14:19:05 +08:00
parent cc19a8430b
commit b06891e710
10 changed files with 286 additions and 549 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/api/chat.py
+43 -39
View File
@@ -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'
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'
# 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)
# )
# 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 EventSourceResponse(event_generator(), headers={"Content-Type": "text/event-stream"})
return StreamingResponse(run_task(), media_type="text/event-stream")