mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 09:17:06 +08:00
feat: add get_chat_with_data
This commit is contained in:
@@ -8,7 +8,7 @@ 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, get_chat_chart_data, get_chat_predict_data
|
||||
delete_chat, get_chat_chart_data, get_chat_predict_data, get_chat_with_records_with_data
|
||||
from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion, ExcelData
|
||||
from apps.chat.task.llm import LLMService
|
||||
from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
|
||||
@@ -33,6 +33,18 @@ async def get_chat(session: SessionDep, current_user: CurrentUser, chart_id: int
|
||||
)
|
||||
|
||||
|
||||
@router.get("/get/with_data/{chart_id}")
|
||||
async def get_chat_with_data(session: SessionDep, current_user: CurrentUser, chart_id: int, current_assistant: CurrentAssistant):
|
||||
try:
|
||||
return get_chat_with_records_with_data(chart_id=chart_id, session=session, current_user=current_user,
|
||||
current_assistant=current_assistant)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=str(e)
|
||||
)
|
||||
|
||||
|
||||
@router.get("/record/get/{chart_record_id}/data")
|
||||
async def chat_record_data(session: SessionDep, chart_record_id: int):
|
||||
try:
|
||||
@@ -205,7 +217,8 @@ async def export_excel(excel_data: ExcelData):
|
||||
|
||||
buffer = io.BytesIO()
|
||||
|
||||
with pd.ExcelWriter(buffer, engine='xlsxwriter', engine_kwargs={'options': {'strings_to_numbers': True}}) as writer:
|
||||
with pd.ExcelWriter(buffer, engine='xlsxwriter',
|
||||
engine_kwargs={'options': {'strings_to_numbers': True}}) as writer:
|
||||
df.to_excel(writer, sheet_name='Sheet1', index=False)
|
||||
|
||||
buffer.seek(0)
|
||||
|
||||
@@ -67,8 +67,13 @@ def get_chat_predict_data(session: SessionDep, chart_record_id: int):
|
||||
return {}
|
||||
|
||||
|
||||
def get_chat_with_records_with_data(session: SessionDep, chart_id: int, current_user: CurrentUser,
|
||||
current_assistant: CurrentAssistant) -> ChatInfo:
|
||||
return get_chat_with_records(session, chart_id, current_user, current_assistant, True)
|
||||
|
||||
|
||||
def get_chat_with_records(session: SessionDep, chart_id: int, current_user: CurrentUser,
|
||||
current_assistant: CurrentAssistant) -> ChatInfo:
|
||||
current_assistant: CurrentAssistant, with_data: bool = False) -> ChatInfo:
|
||||
chat = session.get(Chat, chart_id)
|
||||
if not chat:
|
||||
raise Exception(f"Chat with id {chart_id} not found")
|
||||
@@ -96,6 +101,16 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
|
||||
ChatRecord.recommended_question, ChatRecord.first_chat,
|
||||
ChatRecord.finish, ChatRecord.error).where(
|
||||
and_(ChatRecord.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time)
|
||||
if with_data:
|
||||
stmt = select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time,
|
||||
ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql,
|
||||
ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict,
|
||||
ChatRecord.datasource_select_answer, ChatRecord.analysis_record_id, ChatRecord.predict_record_id,
|
||||
ChatRecord.recommended_question, ChatRecord.first_chat,
|
||||
ChatRecord.finish, ChatRecord.error, ChatRecord.data, ChatRecord.predict_data).where(
|
||||
and_(ChatRecord.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(
|
||||
ChatRecord.create_time)
|
||||
|
||||
result = session.execute(stmt).all()
|
||||
record_list: list[ChatRecord] = []
|
||||
for row in result:
|
||||
|
||||
Reference in New Issue
Block a user