feat: 增加 Agent Excel 批量测试
This commit is contained in:
157
ai_knowledge_base_v2/apps/backend/app/api/admin_agent_batch.py
Normal file
157
ai_knowledge_base_v2/apps/backend/app/api/admin_agent_batch.py
Normal file
@@ -0,0 +1,157 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import UTC, datetime
|
||||
from typing import Annotated
|
||||
from urllib.parse import quote
|
||||
|
||||
from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile, status
|
||||
from fastapi.responses import StreamingResponse
|
||||
from sqlalchemy import func, select, update
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.api.pagination import page_result
|
||||
from app.core.database import get_db
|
||||
from app.core.dependencies import get_current_admin
|
||||
from app.core.responses import api_success
|
||||
from app.models.admin import Admin
|
||||
from app.models.agent_batch import AgentBatchTest, AgentBatchTestItem
|
||||
from app.services.admin_service import OperationLogService
|
||||
from app.services.agent_batch_test_service import AgentBatchTestService, item_dict, job_dict
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/agent/batch/template")
|
||||
def download_batch_template(
|
||||
current_admin: Admin = Depends(get_current_admin),
|
||||
) -> StreamingResponse:
|
||||
return StreamingResponse(
|
||||
AgentBatchTestService.template_workbook(),
|
||||
media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
headers={"Content-Disposition": "attachment; filename*=UTF-8''Agent_batch_test_template.xlsx"},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/agent/batch/jobs")
|
||||
def create_batch_job(
|
||||
file: Annotated[UploadFile, File(...)],
|
||||
configJson: Annotated[str, Form(...)],
|
||||
db: Session = Depends(get_db),
|
||||
current_admin: Admin = Depends(get_current_admin),
|
||||
) -> dict:
|
||||
try:
|
||||
config = json.loads(configJson)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="批量测试配置格式不正确") from exc
|
||||
if not isinstance(config, dict):
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="批量测试配置格式不正确")
|
||||
rows = AgentBatchTestService.parse_questions(file)
|
||||
job = AgentBatchTestService.create_job(
|
||||
db,
|
||||
admin=current_admin,
|
||||
filename=file.filename or "Agent批量测试.xlsx",
|
||||
rows=rows,
|
||||
config=config,
|
||||
)
|
||||
OperationLogService.write(db, admin_id=current_admin.id, module="agent", action="batch_create", target_id=job.id)
|
||||
db.commit()
|
||||
return api_success(job_dict(job))
|
||||
|
||||
|
||||
@router.get("/agent/batch/jobs")
|
||||
def list_batch_jobs(
|
||||
page: int = Query(default=1, ge=1),
|
||||
pageSize: int = Query(default=10, ge=10, le=100),
|
||||
db: Session = Depends(get_db),
|
||||
current_admin: Admin = Depends(get_current_admin),
|
||||
) -> dict:
|
||||
total = db.scalar(select(func.count(AgentBatchTest.id))) or 0
|
||||
rows = list(
|
||||
db.scalars(
|
||||
select(AgentBatchTest)
|
||||
.order_by(AgentBatchTest.created_at.desc(), AgentBatchTest.id.desc())
|
||||
.offset((page - 1) * pageSize)
|
||||
.limit(pageSize)
|
||||
)
|
||||
)
|
||||
return api_success(page_result([job_dict(row) for row in rows], total=total, page=page, page_size=pageSize))
|
||||
|
||||
|
||||
@router.get("/agent/batch/jobs/{job_id}")
|
||||
def get_batch_job(
|
||||
job_id: int,
|
||||
page: int = Query(default=1, ge=1),
|
||||
pageSize: int = Query(default=20, ge=10, le=100),
|
||||
db: Session = Depends(get_db),
|
||||
current_admin: Admin = Depends(get_current_admin),
|
||||
) -> dict:
|
||||
job = _job_or_404(db, job_id)
|
||||
total = db.scalar(select(func.count(AgentBatchTestItem.id)).where(AgentBatchTestItem.batch_id == job.id)) or 0
|
||||
items = list(
|
||||
db.scalars(
|
||||
select(AgentBatchTestItem)
|
||||
.where(AgentBatchTestItem.batch_id == job.id)
|
||||
.order_by(AgentBatchTestItem.row_number.asc(), AgentBatchTestItem.id.asc())
|
||||
.offset((page - 1) * pageSize)
|
||||
.limit(pageSize)
|
||||
)
|
||||
)
|
||||
return api_success({**job_dict(job), "items": page_result([item_dict(item) for item in items], total=total, page=page, page_size=pageSize)})
|
||||
|
||||
|
||||
@router.get("/agent/batch/jobs/{job_id}/export")
|
||||
def export_batch_job(
|
||||
job_id: int,
|
||||
db: Session = Depends(get_db),
|
||||
current_admin: Admin = Depends(get_current_admin),
|
||||
) -> StreamingResponse:
|
||||
job = _job_or_404(db, job_id)
|
||||
items = list(
|
||||
db.scalars(
|
||||
select(AgentBatchTestItem)
|
||||
.where(AgentBatchTestItem.batch_id == job.id)
|
||||
.order_by(AgentBatchTestItem.row_number.asc(), AgentBatchTestItem.id.asc())
|
||||
)
|
||||
)
|
||||
filename = quote(f"Agent批量测试结果_{job.id}.xlsx")
|
||||
OperationLogService.write(db, admin_id=current_admin.id, module="agent", action="batch_export", target_id=job.id)
|
||||
db.commit()
|
||||
return StreamingResponse(
|
||||
AgentBatchTestService.export_workbook(job, items),
|
||||
media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
headers={"Content-Disposition": f"attachment; filename*=UTF-8''{filename}"},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/agent/batch/jobs/{job_id}/cancel")
|
||||
def cancel_batch_job(
|
||||
job_id: int,
|
||||
db: Session = Depends(get_db),
|
||||
current_admin: Admin = Depends(get_current_admin),
|
||||
) -> dict:
|
||||
job = _job_or_404(db, job_id)
|
||||
if job.status not in {"pending", "running"}:
|
||||
raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="当前任务已结束，不能取消")
|
||||
now = datetime.now(UTC).replace(tzinfo=None)
|
||||
job.status = "cancelled"
|
||||
job.cancelled_at = now
|
||||
job.finished_at = now
|
||||
db.add(job)
|
||||
db.execute(
|
||||
update(AgentBatchTestItem)
|
||||
.where(AgentBatchTestItem.batch_id == job.id, AgentBatchTestItem.status == "pending")
|
||||
.values(status="cancelled", finished_at=now, next_run_at=None)
|
||||
)
|
||||
OperationLogService.write(db, admin_id=current_admin.id, module="agent", action="batch_cancel", target_id=job.id)
|
||||
db.commit()
|
||||
db.refresh(job)
|
||||
return api_success(job_dict(job))
|
||||
|
||||
|
||||
def _job_or_404(db: Session, job_id: int) -> AgentBatchTest:
|
||||
job = db.get(AgentBatchTest, job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="批量测试任务不存在")
|
||||
return job
|
||||
@@ -7,6 +7,7 @@ from app.api import (
|
||||
admin_auth,
|
||||
admin_content_generation,
|
||||
admin_agent_records,
|
||||
admin_agent_batch,
|
||||
admin_dashboard,
|
||||
admin_entitlements,
|
||||
admin_knowledge,
|
||||
@@ -38,6 +39,7 @@ api_router.include_router(admin_management.router, prefix="/admin", tags=["admin
|
||||
guard = [Depends(enforce_admin_access)]
|
||||
api_router.include_router(admin_content_generation.router, prefix="/admin", tags=["admin-content-generation"], dependencies=guard)
|
||||
api_router.include_router(admin_agent_records.router, prefix="/admin", tags=["admin-agent-records"], dependencies=guard)
|
||||
api_router.include_router(admin_agent_batch.router, prefix="/admin", tags=["admin-agent-batch"], dependencies=guard)
|
||||
api_router.include_router(admin_dashboard.router, prefix="/admin", tags=["admin"], dependencies=guard)
|
||||
api_router.include_router(admin_entitlements.router, prefix="/admin", tags=["admin-entitlements"], dependencies=guard)
|
||||
api_router.include_router(admin_users.router, prefix="/admin", tags=["admin"], dependencies=guard)
|
||||
|
||||
Reference in New Issue
Block a user