feat: 增加 Agent Excel 批量测试

This commit is contained in:
2026-08-12 17:00:50 +08:00
parent c4231f268b
commit 3faecab6ac
26 changed files with 1327 additions and 6 deletions

View 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

View File

@@ -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)

View File

@@ -80,6 +80,9 @@ class Settings(BaseSettings):
topic_settlement_poll_seconds: int = 2
topic_settlement_stale_minutes: int = 30
topic_settlement_max_attempts: int = 3
agent_batch_worker_enabled: bool = True
agent_batch_poll_seconds: int = 2
agent_batch_stale_minutes: int = 30
bootstrap_admin_username: str = ""
bootstrap_admin_password: str = ""
bootstrap_admin_name: str = "系统管理员"

View File

@@ -100,6 +100,8 @@ def enforce_admin_access(
permission = "users.delete" if method == "DELETE" else ("users.view" if method == "GET" else ("users.create" if method == "POST" and (path in {"user", "user/import", "user/import/excel"}) else "users.edit"))
elif path.startswith("knowledge"):
permission = "knowledge.delete" if method == "DELETE" else ("knowledge.view" if method == "GET" else ("knowledge.publish" if path.endswith("open-status") or path.endswith("lifecycle") else "knowledge.edit"))
elif path.startswith("agent/batch"):
permission = "prompt.batch"
elif path.startswith("prompt") or path.startswith("agent/"):
permission = "prompt.view" if method == "GET" else "prompt.edit"
elif path.startswith("model"):

View File

@@ -15,6 +15,7 @@ from app.services.secret_service import SecretService
from app.services.maintenance_service import MaintenanceService
from app.services.periodic_report_worker import PeriodicReportWorker
from app.services.topic_settlement_worker import TopicSettlementWorker
from app.services.agent_batch_test_worker import AgentBatchTestWorker
@asynccontextmanager
@@ -26,12 +27,13 @@ async def lifespan(app: FastAPI):
maintenance_task = asyncio.create_task(MaintenanceService.run_forever())
periodic_report_task = asyncio.create_task(PeriodicReportWorker.run_forever())
topic_settlement_task = asyncio.create_task(TopicSettlementWorker.run_forever())
agent_batch_task = asyncio.create_task(AgentBatchTestWorker.run_forever())
try:
yield
finally:
for task in (maintenance_task, periodic_report_task, topic_settlement_task):
for task in (maintenance_task, periodic_report_task, topic_settlement_task, agent_batch_task):
task.cancel()
for task in (maintenance_task, periodic_report_task, topic_settlement_task):
for task in (maintenance_task, periodic_report_task, topic_settlement_task, agent_batch_task):
try:
await task
except asyncio.CancelledError:

View File

@@ -1,4 +1,5 @@
from app.models.admin import Admin, Role
from app.models.agent_batch import AgentBatchTest, AgentBatchTestItem
from app.models.ai_config import ContentGenerationConfig, ModelConfig, Prompt, SystemConfig
from app.models.base import Base
from app.models.chat import ChatMessage, ChatSession, TopicSession
@@ -29,6 +30,8 @@ from app.models.user import User
__all__ = [
"Admin",
"AgentBatchTest",
"AgentBatchTestItem",
"AiRequestLog",
"Base",
"ChatMessage",

View File

@@ -0,0 +1,71 @@
from __future__ import annotations
from datetime import datetime
from sqlalchemy import BigInteger, DateTime, ForeignKey, Index, Integer, String, Text, func
from sqlalchemy.orm import Mapped, mapped_column
from app.models.base import Base
PRIMARY_KEY_TYPE = BigInteger().with_variant(Integer, "sqlite")
class AgentBatchTest(Base):
__tablename__ = "sys_agent_batch_test"
__table_args__ = (
Index("ix_agent_batch_test_creator_created", "created_by", "created_at"),
Index("ix_agent_batch_test_status_created", "status", "created_at"),
)
id: Mapped[int] = mapped_column(PRIMARY_KEY_TYPE, primary_key=True, autoincrement=True)
name: Mapped[str] = mapped_column(String(160), nullable=False)
original_filename: Mapped[str] = mapped_column(String(255), nullable=False)
status: Mapped[str] = mapped_column(String(20), default="pending", nullable=False)
total_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
success_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
failed_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
processed_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
prompt_content: Mapped[str] = mapped_column(Text, nullable=False)
model_id: Mapped[int] = mapped_column(ForeignKey("sys_model.id"), nullable=False)
model_name: Mapped[str] = mapped_column(String(100), nullable=False)
knowledge_ids: Mapped[str] = mapped_column(Text, nullable=False)
knowledge_names: Mapped[str] = mapped_column(Text, nullable=False)
generation_config: Mapped[str] = mapped_column(Text, nullable=False)
created_by: Mapped[int] = mapped_column(ForeignKey("sys_admin.id"), nullable=False)
started_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
finished_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
cancelled_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), nullable=False)
updated_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), onupdate=func.now(), nullable=False)
class AgentBatchTestItem(Base):
__tablename__ = "sys_agent_batch_test_item"
__table_args__ = (
Index("ix_agent_batch_item_job_order", "batch_id", "row_number"),
Index("ix_agent_batch_item_claim", "status", "next_run_at", "id"),
Index("ix_agent_batch_item_stale", "status", "locked_at"),
)
id: Mapped[int] = mapped_column(PRIMARY_KEY_TYPE, primary_key=True, autoincrement=True)
batch_id: Mapped[int] = mapped_column(ForeignKey("sys_agent_batch_test.id", ondelete="CASCADE"), nullable=False)
row_number: Mapped[int] = mapped_column(Integer, nullable=False)
external_no: Mapped[str | None] = mapped_column(String(100), nullable=True)
question: Mapped[str] = mapped_column(Text, nullable=False)
answer: Mapped[str | None] = mapped_column(Text, nullable=True)
status: Mapped[str] = mapped_column(String(20), default="pending", nullable=False)
error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
model_name: Mapped[str | None] = mapped_column(String(100), nullable=True)
knowledge_ids: Mapped[str | None] = mapped_column(Text, nullable=True)
retrieve_count: Mapped[int | None] = mapped_column(Integer, nullable=True)
duration_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
attempt_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
max_attempts: Mapped[int] = mapped_column(Integer, default=2, nullable=False)
next_run_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
locked_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
locked_by: Mapped[str | None] = mapped_column(String(120), nullable=True)
started_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
finished_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), nullable=False)
updated_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), onupdate=func.now(), nullable=False)

View File

@@ -12,7 +12,7 @@ PERMISSION_TREE = [
{"code": "users", "name": "用户管理", "children": [{"code": "users.view", "name": "查看用户"}, {"code": "users.create", "name": "新增/导入"}, {"code": "users.edit", "name": "编辑/权益续期"}, {"code": "users.delete", "name": "删除用户"}]},
{"code": "entitlements", "name": "权益管理", "children": [{"code": "entitlements.view", "name": "查看权益"}, {"code": "entitlements.edit", "name": "编辑权益"}]},
{"code": "knowledge", "name": "知识库管理", "children": [{"code": "knowledge.view", "name": "查看知识库"}, {"code": "knowledge.edit", "name": "新增/编辑/同步"}, {"code": "knowledge.publish", "name": "开放/归档"}, {"code": "knowledge.delete", "name": "删除知识库"}]},
{"code": "prompt", "name": "Agent 管理", "children": [{"code": "prompt.view", "name": "查看 Agent"}, {"code": "prompt.edit", "name": "编辑/测试 Agent"}]},
{"code": "prompt", "name": "Agent 管理", "children": [{"code": "prompt.view", "name": "查看 Agent"}, {"code": "prompt.edit", "name": "编辑/单条测试 Agent"}, {"code": "prompt.batch", "name": "批量测试/导出"}]},
{"code": "models", "name": "模型管理", "children": [{"code": "models.view", "name": "查看模型"}, {"code": "models.edit", "name": "新增/编辑/测试"}, {"code": "models.delete", "name": "删除模型"}]},
{"code": "content-generation", "name": "内容生成", "children": [{"code": "content-generation.view", "name": "查看配置"}, {"code": "content-generation.edit", "name": "编辑/测试配置"}]},
{"code": "configs", "name": "系统配置", "children": [{"code": "configs.view", "name": "查看配置"}, {"code": "configs.edit", "name": "修改配置"}]},

View File

@@ -0,0 +1,274 @@
from __future__ import annotations
import json
from datetime import UTC, datetime
from io import BytesIO
from fastapi import HTTPException, UploadFile, status
from openpyxl import Workbook, load_workbook
from openpyxl.styles import Alignment, Font, PatternFill
from openpyxl.worksheet.table import Table, TableStyleInfo
from sqlalchemy import func, select
from sqlalchemy.orm import Session
from app.models.admin import Admin
from app.models.agent_batch import AgentBatchTest, AgentBatchTestItem
from app.models.ai_config import ModelConfig
from app.models.knowledge import Knowledge
MAX_BATCH_QUESTIONS = 200
MAX_QUESTION_LENGTH = 2000
TEMPLATE_HEADERS = ["序号", "问题"]
class AgentBatchTestService:
@staticmethod
def parse_questions(file: UploadFile) -> list[dict]:
filename = file.filename or ""
if not filename.lower().endswith(".xlsx"):
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="请上传 .xlsx 格式的 Excel 文件")
content = file.file.read()
if not content:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Excel 文件为空")
if len(content) > 5 * 1024 * 1024:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Excel 文件不能超过 5MB")
try:
workbook = load_workbook(BytesIO(content), read_only=True, data_only=False)
except Exception as exc:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Excel 文件解析失败,请重新下载模板填写") from exc
sheet = workbook.active
if sheet.max_row > 1000:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Excel 行数异常,请删除多余空行后重试")
headers = [str(value or "").strip() for value in next(sheet.iter_rows(min_row=1, max_row=1, values_only=True), ())]
header_map = {name: index for index, name in enumerate(headers) if name}
if "问题" not in header_map:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="模板缺少必填列:问题")
rows: list[dict] = []
seen: set[str] = set()
for row_number, values in enumerate(sheet.iter_rows(min_row=2, max_row=min(sheet.max_row, 1000), values_only=True), start=2):
raw_question = _value_at(values, header_map["问题"])
if isinstance(raw_question, str) and raw_question.startswith("="):
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"{row_number} 行问题不能使用 Excel 公式")
question = _cell_text(raw_question)
external_no = _cell_text(_value_at(values, header_map.get("序号"))) or str(row_number - 1)
if not question:
continue
if len(question) > MAX_QUESTION_LENGTH:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"{row_number} 行问题超过 {MAX_QUESTION_LENGTH}")
normalized = " ".join(question.split())
if normalized in seen:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"{row_number} 行问题与前面重复")
seen.add(normalized)
rows.append({"rowNumber": row_number, "externalNo": external_no[:100], "question": question})
if len(rows) > MAX_BATCH_QUESTIONS:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"单次最多导入 {MAX_BATCH_QUESTIONS} 个问题")
if not rows:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Excel 中没有可测试的问题")
return rows
@staticmethod
def create_job(
db: Session,
*,
admin: Admin,
filename: str,
rows: list[dict],
config: dict,
) -> AgentBatchTest:
prompt_content = str(config.get("promptContent") or "").strip()
if not prompt_content:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="主提示词不能为空")
try:
model_id = int(config.get("modelId"))
except (TypeError, ValueError) as exc:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="请选择批量测试模型") from exc
model = db.get(ModelConfig, model_id)
if model is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="批量测试模型不存在")
knowledge_ids = _positive_int_list(config.get("knowledgeIds"))
if knowledge_ids:
found = list(db.scalars(select(Knowledge).where(Knowledge.id.in_(knowledge_ids))))
if len(found) != len(knowledge_ids):
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="所选知识库中包含已删除的数据")
knowledge_names = [item.name for item in sorted(found, key=lambda item: knowledge_ids.index(item.id))]
else:
knowledge_names = ["全部正式开放知识库"]
generation_config = {
"temperature": config.get("temperature"),
"topP": config.get("topP"),
"topK": config.get("topK"),
"presencePenalty": config.get("presencePenalty"),
"frequencyPenalty": config.get("frequencyPenalty"),
"maxToken": config.get("maxToken") or 8192,
"streamEnabled": int(config.get("streamEnabled", 1)),
"reasoningVisible": 0,
"responseDepth": int(config.get("responseDepth", 35)),
}
job = AgentBatchTest(
name=str(config.get("name") or filename.rsplit(".", 1)[0] or "Agent 批量测试")[:160],
original_filename=filename[:255],
total_count=len(rows),
prompt_content=prompt_content,
model_id=model.id,
model_name=model.display_name or model.model_name,
knowledge_ids=json.dumps(knowledge_ids, ensure_ascii=False),
knowledge_names=json.dumps(knowledge_names, ensure_ascii=False),
generation_config=json.dumps(generation_config, ensure_ascii=False),
created_by=admin.id,
)
db.add(job)
db.flush()
db.add_all([
AgentBatchTestItem(
batch_id=job.id,
row_number=row["rowNumber"],
external_no=row["externalNo"],
question=row["question"],
)
for row in rows
])
db.flush()
return job
@staticmethod
def template_workbook() -> BytesIO:
workbook = Workbook()
sheet = workbook.active
sheet.title = "批量测试问题"
sheet.append(TEMPLATE_HEADERS)
sheet.append([1, "请填写需要测试的问题,保留一行一个问题。"])
sheet.append([2, "例如:遇到情绪焦虑时,可以先做什么?"])
_style_sheet(sheet, widths=(12, 72), table_ref="A1:B3", table_name="AgentBatchQuestionTable")
note = workbook.create_sheet("填写说明")
note.append(["填写规则", "说明"])
note.append(["必填列", "问题"])
note.append(["单次数量", f"最多 {MAX_BATCH_QUESTIONS} 个问题"])
note.append(["问题长度", f"每个问题最多 {MAX_QUESTION_LENGTH}"])
note.append(["重复问题", "同一文件内不允许重复"])
note.append(["不要修改", "不要修改工作表名称和表头;可以删除示例行后填写"])
_style_sheet(note, widths=(18, 72), table_ref="A1:B6", table_name="AgentBatchInstructions")
stream = BytesIO()
workbook.save(stream)
stream.seek(0)
return stream
@staticmethod
def export_workbook(job: AgentBatchTest, items: list[AgentBatchTestItem]) -> BytesIO:
workbook = Workbook()
sheet = workbook.active
sheet.title = "批量测试结果"
headers = ["序号", "问题", "答案", "状态", "失败原因", "实际模型", "召回数量", "耗时(秒)"]
sheet.append(headers)
status_labels = {"success": "成功", "failed": "失败", "cancelled": "已取消", "pending": "等待中", "running": "生成中"}
for item in items:
sheet.append([
item.external_no or item.row_number - 1,
_excel_safe(item.question),
_excel_safe(item.answer or ""),
status_labels.get(item.status, item.status),
_excel_safe(item.error_message or ""),
_excel_safe(item.model_name or job.model_name),
item.retrieve_count,
round(item.duration_ms / 1000, 2) if item.duration_ms is not None else None,
])
end_row = max(2, len(items) + 1)
_style_sheet(sheet, widths=(12, 42, 80, 12, 36, 22, 12, 12), table_ref=f"A1:H{end_row}", table_name="AgentBatchResultTable")
for row in sheet.iter_rows(min_row=2, max_row=end_row):
row[1].alignment = Alignment(vertical="top", wrap_text=True)
row[2].alignment = Alignment(vertical="top", wrap_text=True)
row[4].alignment = Alignment(vertical="top", wrap_text=True)
summary = workbook.create_sheet("任务信息")
summary.append(["项目", "内容"])
summary.append(["任务名称", job.name])
summary.append(["任务状态", job.status])
summary.append(["问题总数", job.total_count])
summary.append(["成功", job.success_count])
summary.append(["失败", job.failed_count])
summary.append(["测试模型", job.model_name])
summary.append(["知识库", "".join(json.loads(job.knowledge_names or "[]"))])
summary.append(["创建时间", job.created_at])
summary.append(["完成时间", job.finished_at])
_style_sheet(summary, widths=(20, 86), table_ref="A1:B10", table_name="AgentBatchSummary")
summary.column_dimensions["B"].width = 86
stream = BytesIO()
workbook.save(stream)
stream.seek(0)
return stream
def job_dict(job: AgentBatchTest) -> dict:
return {
"id": job.id,
"name": job.name,
"originalFilename": job.original_filename,
"status": job.status,
"totalCount": job.total_count,
"processedCount": job.processed_count,
"successCount": job.success_count,
"failedCount": job.failed_count,
"progress": round(job.processed_count * 100 / job.total_count) if job.total_count else 0,
"modelId": job.model_id,
"modelName": job.model_name,
"knowledgeIds": json.loads(job.knowledge_ids or "[]"),
"knowledgeNames": json.loads(job.knowledge_names or "[]"),
"createdAt": job.created_at,
"startedAt": job.started_at,
"finishedAt": job.finished_at,
}
def item_dict(item: AgentBatchTestItem) -> dict:
return {
"id": item.id,
"rowNumber": item.row_number,
"externalNo": item.external_no,
"question": item.question,
"answer": item.answer,
"status": item.status,
"errorMessage": item.error_message,
"modelName": item.model_name,
"retrieveCount": item.retrieve_count,
"durationMs": item.duration_ms,
}
def _style_sheet(sheet, *, widths: tuple[int, ...], table_ref: str, table_name: str) -> None:
sheet.freeze_panes = "A2"
sheet.sheet_view.showGridLines = False
header_fill = PatternFill("solid", fgColor="176B57")
for cell in sheet[1]:
cell.fill = header_fill
cell.font = Font(color="FFFFFF", bold=True)
cell.alignment = Alignment(vertical="center")
sheet.row_dimensions[1].height = 28
for index, width in enumerate(widths, start=1):
sheet.column_dimensions[chr(64 + index)].width = width
table = Table(displayName=table_name, ref=table_ref)
table.tableStyleInfo = TableStyleInfo(name="TableStyleMedium4", showRowStripes=True, showFirstColumn=False, showLastColumn=False)
sheet.add_table(table)
def _cell_text(value) -> str:
if value is None:
return ""
return str(value).strip()
def _value_at(row: tuple, index: int | None):
return row[index] if index is not None and index < len(row) else None
def _positive_int_list(value) -> list[int]:
if not isinstance(value, list):
return []
try:
result = list(dict.fromkeys(int(item) for item in value if int(item) > 0))
except (TypeError, ValueError) as exc:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="知识库参数格式不正确") from exc
return result
def _excel_safe(value: str) -> str:
text = str(value or "")
return f"'{text}" if text.lstrip().startswith(("=", "+", "-", "@")) else text

View File

@@ -0,0 +1,227 @@
from __future__ import annotations
import asyncio
import json
import logging
import os
import socket
import time
import uuid
from datetime import UTC, datetime, timedelta
from sqlalchemy import func, or_, select
from sqlalchemy.orm import Session
from app.core.config import get_settings
from app.core.database import SessionLocal
from app.models.admin import Admin
from app.models.agent_batch import AgentBatchTest, AgentBatchTestItem
from app.schemas.admin import AgentDebugRequest
from app.services.agent_debug_service import AgentDebugService
logger = logging.getLogger(__name__)
class AgentBatchTestWorker:
"""Persistent, restart-safe queue for Agent batch questions."""
@classmethod
async def run_forever(cls) -> None:
settings = get_settings()
if not settings.agent_batch_worker_enabled:
logger.info("agent batch test worker disabled")
return
worker_id = f"{socket.gethostname()}:{os.getpid()}:{uuid.uuid4().hex[:8]}"
while True:
try:
processed = await cls.run_once(worker_id)
except Exception:
processed = False
logger.exception("agent batch test worker iteration failed")
await asyncio.sleep(0 if processed else max(1, settings.agent_batch_poll_seconds))
@classmethod
async def run_once(cls, worker_id: str) -> bool:
now = _now()
with SessionLocal() as db:
recovered_batch_ids = cls.recover_stale_items(db, now=now)
db.commit()
for batch_id in recovered_batch_ids:
cls.refresh_job(db, batch_id)
with SessionLocal() as db:
item_id = cls.claim_next(db, worker_id=worker_id, now=now)
if item_id is None:
return False
await cls.execute_claimed(item_id=item_id, worker_id=worker_id)
return True
@staticmethod
def claim_next(db: Session, *, worker_id: str, now: datetime | None = None) -> int | None:
current = now or _now()
item = db.scalar(
select(AgentBatchTestItem)
.join(AgentBatchTest, AgentBatchTest.id == AgentBatchTestItem.batch_id)
.where(
AgentBatchTest.status.in_(("pending", "running")),
AgentBatchTestItem.status == "pending",
AgentBatchTestItem.attempt_count < AgentBatchTestItem.max_attempts,
or_(AgentBatchTestItem.next_run_at.is_(None), AgentBatchTestItem.next_run_at <= current),
)
.order_by(AgentBatchTestItem.batch_id.asc(), AgentBatchTestItem.row_number.asc())
.with_for_update(skip_locked=True)
.limit(1)
)
if item is None:
db.rollback()
return None
job = db.get(AgentBatchTest, item.batch_id)
item.status = "running"
item.attempt_count += 1
item.locked_at = current
item.locked_by = worker_id
item.started_at = current
if job is not None and job.status == "pending":
job.status = "running"
job.started_at = current
db.add(job)
db.add(item)
db.commit()
return item.id
@classmethod
async def execute_claimed(cls, *, item_id: int, worker_id: str) -> None:
with SessionLocal() as db:
item = db.get(AgentBatchTestItem, item_id)
if item is None or item.status != "running" or item.locked_by != worker_id:
return
job = db.get(AgentBatchTest, item.batch_id)
admin = db.get(Admin, job.created_by) if job is not None else None
if job is None or admin is None:
cls._finish_item(db, item, error="批量任务或创建管理员不存在")
return
if job.status == "cancelled":
item.status = "cancelled"
item.finished_at = _now()
item.locked_at = None
item.locked_by = None
db.add(item)
db.commit()
return
payload = _payload_from(job, item)
started = time.monotonic()
try:
result = await AgentDebugService.run_complete(payload, db, admin)
duration_ms = round((time.monotonic() - started) * 1000)
item.answer = result["answer"]
item.model_name = result.get("modelName") or job.model_name
item.knowledge_ids = json.dumps(result.get("knowledgeIds") or [], ensure_ascii=False)
item.retrieve_count = result.get("retrieveCount")
item.duration_ms = duration_ms
item.status = "success"
item.error_message = None
item.finished_at = _now()
item.next_run_at = None
item.locked_at = None
item.locked_by = None
db.add(item)
db.commit()
except Exception as exc:
db.rollback()
item = db.get(AgentBatchTestItem, item_id)
if item is None:
return
item.duration_ms = round((time.monotonic() - started) * 1000)
if item.attempt_count < item.max_attempts:
item.status = "pending"
item.next_run_at = _now() + timedelta(seconds=15 * item.attempt_count)
item.finished_at = None
else:
item.status = "failed"
item.next_run_at = None
item.finished_at = _now()
item.error_message = (str(exc) or "Agent 生成失败")[:2000]
item.locked_at = None
item.locked_by = None
db.add(item)
db.commit()
cls.refresh_job(db, job.id)
@staticmethod
def refresh_job(db: Session, batch_id: int) -> None:
job = db.get(AgentBatchTest, batch_id)
if job is None or job.status == "cancelled":
return
counts = dict(
db.execute(
select(AgentBatchTestItem.status, func.count(AgentBatchTestItem.id))
.where(AgentBatchTestItem.batch_id == batch_id)
.group_by(AgentBatchTestItem.status)
).all()
)
job.success_count = counts.get("success", 0)
job.failed_count = counts.get("failed", 0)
job.processed_count = job.success_count + job.failed_count
if job.processed_count >= job.total_count:
job.status = "completed" if job.failed_count == 0 else "completed_with_errors"
job.finished_at = _now()
else:
job.status = "running"
db.add(job)
db.commit()
@staticmethod
def recover_stale_items(db: Session, *, now: datetime | None = None) -> set[int]:
current = now or _now()
stale_before = current - timedelta(minutes=max(5, get_settings().agent_batch_stale_minutes))
items = list(
db.scalars(
select(AgentBatchTestItem)
.where(
AgentBatchTestItem.status == "running",
AgentBatchTestItem.locked_at.is_not(None),
AgentBatchTestItem.locked_at < stale_before,
)
.order_by(AgentBatchTestItem.id.asc())
.limit(100)
.with_for_update(skip_locked=True)
)
)
affected_batch_ids: set[int] = set()
for item in items:
affected_batch_ids.add(item.batch_id)
item.status = "pending" if item.attempt_count < item.max_attempts else "failed"
item.error_message = "任务执行进程中断,系统已自动恢复" if item.status == "pending" else "任务执行进程连续中断"
item.next_run_at = current if item.status == "pending" else None
item.finished_at = current if item.status == "failed" else None
item.locked_at = None
item.locked_by = None
db.add(item)
return affected_batch_ids
@classmethod
def _finish_item(cls, db: Session, item: AgentBatchTestItem, *, error: str) -> None:
item.status = "failed"
item.error_message = error
item.finished_at = _now()
item.locked_at = None
item.locked_by = None
db.add(item)
db.commit()
cls.refresh_job(db, item.batch_id)
def _payload_from(job: AgentBatchTest, item: AgentBatchTestItem) -> AgentDebugRequest:
generation = json.loads(job.generation_config or "{}")
return AgentDebugRequest(
promptContent=job.prompt_content,
modelId=job.model_id,
knowledgeIds=json.loads(job.knowledge_ids or "[]"),
question=item.question,
history=[],
**generation,
)
def _now() -> datetime:
return datetime.now(UTC).replace(tzinfo=None)

View File

@@ -283,3 +283,20 @@ class AgentDebugService:
except Exception as exc:
db.rollback()
yield {"type": "error", "message": str(exc) or "Agent 调试失败"}
@classmethod
async def run_complete(cls, payload: AgentDebugRequest, db: Session, current_admin: Admin) -> dict:
"""Consume the same routed debug stream used by the UI and return one complete answer."""
answer_parts: list[str] = []
completion: dict = {}
async for event in cls.stream(payload, db, current_admin):
if event.get("type") == "content":
answer_parts.append(str(event.get("content") or ""))
elif event.get("type") == "complete":
completion = event
elif event.get("type") == "error":
raise RuntimeError(str(event.get("message") or "Agent 生成失败"))
answer = "".join(answer_parts).strip()
if not answer:
raise RuntimeError("模型未返回正式回答")
return {"answer": answer, **completion}