feat: add export Sample SQL

This commit is contained in:
ulleo
2025-11-14 17:53:02 +08:00
committed by ulleo
parent 58f43271ea
commit 07f1337305
10 changed files with 204 additions and 72 deletions
@@ -1,9 +1,15 @@
import asyncio
import io
from typing import Optional
import pandas as pd
from fastapi import APIRouter, Query
from fastapi.responses import StreamingResponse
from apps.chat.models.chat_model import AxisObj
from apps.chat.task.llm import LLMService
from apps.data_training.curd.data_training import page_data_training, create_training, update_training, delete_training, \
enable_training
enable_training, get_all_data_training
from apps.data_training.models.data_training_model import DataTrainingInfo
from common.core.deps import SessionDep, CurrentUser, Trans
@@ -43,3 +49,44 @@ async def delete(session: SessionDep, id_list: list[int]):
@router.get("/{id}/enable/{enabled}")
async def enable(session: SessionDep, id: int, enabled: bool, trans: Trans):
enable_training(session, id, enabled, trans)
@router.get("/export")
async def export_excel(session: SessionDep, trans: Trans, current_user: CurrentUser,
word: Optional[str] = Query(None, description="搜索术语(可选)")):
def inner():
_list = get_all_data_training(session, word, oid=current_user.oid)
data_list = []
for obj in _list:
_data = {
"question": obj.question,
"description": obj.description,
"datasource_name": obj.datasource_name,
"advanced_application_name": obj.advanced_application_name,
}
data_list.append(_data)
fields = []
fields.append(AxisObj(name=trans('i18n_data_training.data_training'), value='question'))
fields.append(AxisObj(name=trans('i18n_data_training.problem_description'), value='description'))
fields.append(AxisObj(name=trans('i18n_data_training.effective_data_sources'), value='datasource_name'))
if current_user.oid == 1:
fields.append(
AxisObj(name=trans('i18n_data_training.advanced_application'), value='advanced_application_name'))
md_data, _fields_list = LLMService.convert_object_array_for_pandas(fields, data_list)
df = pd.DataFrame(md_data, columns=_fields_list)
buffer = io.BytesIO()
with pd.ExcelWriter(buffer, engine='xlsxwriter',
engine_kwargs={'options': {'strings_to_numbers': False}}) as writer:
df.to_excel(writer, sheet_name='Sheet1', index=False)
buffer.seek(0)
return io.BytesIO(buffer.getvalue())
result = await asyncio.to_thread(inner)
return StreamingResponse(result, media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
@@ -18,42 +18,61 @@ from common.core.deps import SessionDep, Trans
from common.utils.embedding_threads import run_save_data_training_embeddings
def page_data_training(session: SessionDep, current_page: int = 1, page_size: int = 10, name: Optional[str] = None,
oid: Optional[int] = 1):
_list: List[DataTrainingInfoResult] = []
current_page = max(1, current_page)
page_size = max(10, page_size)
total_count = 0
total_pages = 0
def get_data_training_base_query(oid: int, name: Optional[str] = None):
"""
获取数据训练查询的基础查询结构
"""
if name and name.strip() != "":
keyword_pattern = f"%{name.strip()}%"
parent_ids_subquery = (
select(DataTraining.id)
.where(and_(DataTraining.question.ilike(keyword_pattern), DataTraining.oid == oid)) # LIKE查询条件
.where(and_(DataTraining.question.ilike(keyword_pattern), DataTraining.oid == oid))
)
else:
parent_ids_subquery = (
select(DataTraining.id).where(and_(DataTraining.oid == oid))
)
return parent_ids_subquery
def build_data_training_query(session: SessionDep, oid: int, name: Optional[str] = None,
paginate: bool = True, current_page: int = 1, page_size: int = 10):
"""
构建数据训练查询的通用方法
"""
parent_ids_subquery = get_data_training_base_query(oid, name)
# 计算总数
count_stmt = select(func.count()).select_from(parent_ids_subquery.subquery())
total_count = session.execute(count_stmt).scalar()
total_pages = (total_count + page_size - 1) // page_size
if current_page > total_pages:
if paginate:
# 分页处理
page_size = max(10, page_size)
total_pages = (total_count + page_size - 1) // page_size
current_page = max(1, min(current_page, total_pages)) if total_pages > 0 else 1
paginated_parent_ids = (
parent_ids_subquery
.order_by(DataTraining.create_time.desc())
.offset((current_page - 1) * page_size)
.limit(page_size)
.subquery()
)
else:
# 不分页,获取所有数据
total_pages = 1
current_page = 1
page_size = total_count if total_count > 0 else 1
paginated_parent_ids = (
parent_ids_subquery
.order_by(DataTraining.create_time.desc())
.offset((current_page - 1) * page_size)
.limit(page_size)
.subquery()
)
paginated_parent_ids = (
parent_ids_subquery
.order_by(DataTraining.create_time.desc())
.subquery()
)
# 构建主查询
stmt = (
select(
DataTraining.id,
@@ -74,6 +93,14 @@ def page_data_training(session: SessionDep, current_page: int = 1, page_size: in
.order_by(DataTraining.create_time.desc())
)
return stmt, total_count, total_pages, current_page, page_size
def execute_data_training_query(session: SessionDep, stmt) -> List[DataTrainingInfoResult]:
"""
执行查询并返回数据训练信息列表
"""
_list = []
result = session.execute(stmt)
for row in result:
@@ -90,9 +117,34 @@ def page_data_training(session: SessionDep, current_page: int = 1, page_size: in
advanced_application_name=row.advanced_application_name,
))
return _list
def page_data_training(session: SessionDep, current_page: int = 1, page_size: int = 10,
name: Optional[str] = None, oid: Optional[int] = 1):
"""
分页查询数据训练(原方法保持不变)
"""
stmt, total_count, total_pages, current_page, page_size = build_data_training_query(
session, oid, name, True, current_page, page_size
)
_list = execute_data_training_query(session, stmt)
return current_page, page_size, total_count, total_pages, _list
def get_all_data_training(session: SessionDep, name: Optional[str] = None, oid: Optional[int] = 1):
"""
获取所有数据训练(不分页)
"""
stmt, total_count, total_pages, current_page, page_size = build_data_training_query(
session, oid, name, False
)
_list = execute_data_training_query(session, stmt)
return _list
def create_training(session: SessionDep, info: DataTrainingInfo, oid: int, trans: Trans):
create_time = datetime.datetime.now()
if info.datasource is None and info.advanced_application is None:
+6 -1
View File
@@ -53,7 +53,12 @@
"datasource_cannot_be_none": "Datasource cannot be empty",
"datasource_assistant_cannot_be_none": "Datasource or advanced application cannot both be empty",
"data_training_not_exists": "This example does not exist",
"exists_in_db": "This question already exists"
"exists_in_db": "This question already exists",
"data_training": "SQL Example Library",
"problem_description": "Problem Description",
"sample_sql": "Sample SQL",
"effective_data_sources": "Effective Data Sources",
"advanced_application": "Advanced Application"
},
"i18n_custom_prompt": {
"exists_in_db": "Template name already exists",
+6 -1
View File
@@ -53,7 +53,12 @@
"datasource_cannot_be_none": "데이터 소스는 비울 수 없습니다",
"datasource_assistant_cannot_be_none": "데이터 소스와 고급 애플리케이션을 모두 비울 수 없습니다",
"data_training_not_exists": "이 예시가 존재하지 않습니다",
"exists_in_db": "이 질문이 이미 존재합니다"
"exists_in_db": "이 질문이 이미 존재합니다",
"data_training": "SQL 예시 라이브러리",
"problem_description": "문제 설명",
"sample_sql": "예시 SQL",
"effective_data_sources": "유효 데이터 소스",
"advanced_application": "고급 애플리케이션"
},
"i18n_custom_prompt": {
"exists_in_db": "템플릿 이름이 이미 존재합니다",
+6 -1
View File
@@ -53,7 +53,12 @@
"datasource_cannot_be_none": "数据源不能为空",
"datasource_assistant_cannot_be_none": "数据源或高级应用不能都为空",
"data_training_not_exists": "该示例不存在",
"exists_in_db": "该问题已存在"
"exists_in_db": "该问题已存在",
"data_training": "SQL 示例库",
"problem_description": "问题描述",
"sample_sql": "示例 SQL",
"effective_data_sources": "生效数据源",
"advanced_application": "高级应用"
},
"i18n_custom_prompt": {
"exists_in_db": "模版名称已存在",
+6
View File
@@ -9,4 +9,10 @@ export const trainingApi = {
deleteEmbedded: (params: any) => request.delete('/system/data-training', { data: params }),
getOne: (id: any) => request.get(`/system/data-training/${id}`),
enable: (id: any, enabled: any) => request.get(`/system/data-training/${id}/enable/${enabled}`),
export2Excel: (params: any) =>
request.get(`/system/data-training/export`, {
params,
responseType: 'blob',
requestOptions: { customError: true },
}),
}
+1
View File
@@ -40,6 +40,7 @@
"training_data_items": "Do you want to delete the {msg} selected SQL Sample items?",
"sql_statement": "SQL Statement",
"edit_training_data": "Edit SQL Sample",
"all_236_terms": "Export all {msg} sample SQL records?",
"sales_this_year": "Do you want to delete the SQL Sample: {msg}?"
},
"professional": {
+1
View File
@@ -40,6 +40,7 @@
"training_data_items": "선택된 {msg}개의 예제 SQL을 삭제하시겠습니까?",
"sql_statement": "SQL 문",
"edit_training_data": "예제 SQL 편집",
"all_236_terms": "모든 {msg}개의 예시 SQL 기록을 내보내시겠습니까?",
"sales_this_year": "예제 SQL을 삭제하시겠습니까: {msg}?"
},
"professional": {
+1
View File
@@ -40,6 +40,7 @@
"training_data_items": "是否删除选中的 {msg} 条示例 SQL?",
"sql_statement": "SQL 语句",
"edit_training_data": "编辑示例 SQL",
"all_236_terms": "是否导出全部 {msg} 条示例 SQL?",
"sales_this_year": "是否删除示例 SQL:{msg}?"
},
"professional": {
+57 -48
View File
@@ -85,46 +85,60 @@ const cancelDelete = () => {
checkAll.value = false
isIndeterminate.value = false
}
const exportBatchUser = () => {
ElMessageBox.confirm(
t('professional.selected_2_terms_de', { msg: multipleSelectionAll.value.length }),
{
confirmButtonType: 'primary',
confirmButtonText: t('professional.export'),
cancelButtonText: t('common.cancel'),
customClass: 'confirm-no_icon',
autofocus: false,
}
).then(() => {
trainingApi.deleteEmbedded(multipleSelectionAll.value.map((ele) => ele.id)).then(() => {
ElMessage({
type: 'success',
message: t('dashboard.delete_success'),
})
multipleSelectionAll.value = []
search()
})
})
}
const exportAllUser = () => {
ElMessageBox.confirm(t('professional.all_236_terms', { msg: pageInfo.total }), {
const exportExcel = () => {
ElMessageBox.confirm(t('training.all_236_terms', { msg: pageInfo.total }), {
confirmButtonType: 'primary',
confirmButtonText: t('professional.export'),
cancelButtonText: t('common.cancel'),
customClass: 'confirm-no_icon',
autofocus: false,
}).then(() => {
trainingApi.deleteEmbedded(multipleSelectionAll.value.map((ele) => ele.id)).then(() => {
ElMessage({
type: 'success',
message: t('dashboard.delete_success'),
searchLoading.value = true
trainingApi
.export2Excel(keywords.value ? { question: keywords.value } : {})
.then((res) => {
const blob = new Blob([res], {
type: 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
})
const link = document.createElement('a')
link.href = URL.createObjectURL(blob)
link.download = `${t('training.data_training')}.xlsx`
document.body.appendChild(link)
link.click()
document.body.removeChild(link)
})
.catch(async (error) => {
if (error.response) {
try {
let text = await error.response.data.text()
try {
text = JSON.parse(text)
} finally {
ElMessage({
message: text,
type: 'error',
showClose: true,
})
}
} catch (e) {
console.error('Error processing error response:', e)
}
} else {
console.error('Other error:', error)
ElMessage({
message: error,
type: 'error',
showClose: true,
})
}
})
.finally(() => {
searchLoading.value = false
})
multipleSelectionAll.value = []
search()
})
})
}
const deleteBatchUser = () => {
ElMessageBox.confirm(
t('training.training_data_items', { msg: multipleSelectionAll.value.length }),
@@ -363,20 +377,18 @@ const onRowFormClose = () => {
</el-icon>
</template>
</el-input>
<template v-if="false">
<el-button secondary @click="exportAllUser">
<template #icon>
<icon_export_outlined />
</template>
{{ $t('professional.export_all') }}
</el-button>
<el-button secondary @click="editHandler(null)">
<template #icon>
<ccmUpload></ccmUpload>
</template>
{{ $t('user.batch_import') }}
</el-button>
</template>
<el-button secondary @click="exportExcel">
<template #icon>
<icon_export_outlined />
</template>
{{ $t('professional.export_all') }}
</el-button>
<el-button v-if="false" secondary @click="editHandler(null)">
<template #icon>
<ccmUpload></ccmUpload>
</template>
{{ $t('user.batch_import') }}
</el-button>
<el-button type="primary" @click="editHandler(null)">
<template #icon>
<icon_add_outlined></icon_add_outlined>
@@ -501,9 +513,6 @@ const onRowFormClose = () => {
>
{{ $t('datasource.select_all') }}
</el-checkbox>
<button v-if="false" class="primary-button" @click="exportBatchUser">
{{ $t('professional.export') }}
</button>
<button class="danger-button" @click="deleteBatchUser">{{ $t('dashboard.delete') }}</button>
@@ -646,7 +655,7 @@ const onRowFormClose = () => {
position: relative;
:deep(.ed-table__empty-text) {
padding-top: 160px;
padding-top: 160px;
}
.datasource-yet {