diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 90ec8f2f..23d37996 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -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) diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index 6369536d..5d6497f5 100644 --- a/backend/apps/chat/curd/chat.py +++ b/backend/apps/chat/curd/chat.py @@ -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: