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)
|
||||
|
||||
@@ -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 = "系统管理员"
|
||||
|
||||
@@ -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"):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
|
||||
71
ai_knowledge_base_v2/apps/backend/app/models/agent_batch.py
Normal file
71
ai_knowledge_base_v2/apps/backend/app/models/agent_batch.py
Normal 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)
|
||||
@@ -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": "修改配置"}]},
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user