83e0caae7d
实现RAG评估系统，包括基准管理、评估指标计算和结果展示功能 - 添加评估路由和API接口 - 实现评估指标计算工具类 - 新增基准上传和自动生成功能 - 添加评估结果展示组件 - 完善相关文档和测试用例
330 lines
12 KiB
Python
330 lines
12 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import uuid
|
|
from dataclasses import asdict, dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from collections.abc import Awaitable, Callable
|
|
from collections import Counter
|
|
|
|
from src.config import config
|
|
from src.utils.logging_config import logger
|
|
from src.utils.datetime_utils import utc_isoformat
|
|
|
|
TaskCoroutine = Callable[["TaskContext"], Awaitable[Any]]
|
|
TERMINAL_STATUSES = {"success", "failed", "cancelled"}
|
|
|
|
|
|
def _utc_timestamp() -> str:
|
|
return utc_isoformat()
|
|
|
|
|
|
@dataclass
|
|
class Task:
|
|
id: str
|
|
name: str
|
|
type: str
|
|
status: str = "pending"
|
|
progress: float = 0.0
|
|
message: str = ""
|
|
created_at: str = field(default_factory=_utc_timestamp)
|
|
updated_at: str = field(default_factory=_utc_timestamp)
|
|
started_at: str | None = None
|
|
completed_at: str | None = None
|
|
payload: dict[str, Any] = field(default_factory=dict)
|
|
result: Any | None = None
|
|
error: str | None = None
|
|
cancel_requested: bool = False
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
data = asdict(self)
|
|
return data
|
|
|
|
def to_summary_dict(self) -> dict[str, Any]:
|
|
data = asdict(self)
|
|
data.pop("payload", None)
|
|
data.pop("result", None)
|
|
return data
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any]) -> "Task":
|
|
return cls(
|
|
id=data["id"],
|
|
name=data.get("name", "Unnamed Task"),
|
|
type=data.get("type", "general"),
|
|
status=data.get("status", "pending"),
|
|
progress=data.get("progress", 0.0),
|
|
message=data.get("message", ""),
|
|
created_at=data.get("created_at", _utc_timestamp()),
|
|
updated_at=data.get("updated_at", _utc_timestamp()),
|
|
started_at=data.get("started_at"),
|
|
completed_at=data.get("completed_at"),
|
|
payload=data.get("payload", {}),
|
|
result=data.get("result"),
|
|
error=data.get("error"),
|
|
cancel_requested=data.get("cancel_requested", False),
|
|
)
|
|
|
|
|
|
class TaskContext:
|
|
def __init__(self, tasker: "Tasker", task_id: str):
|
|
self._tasker = tasker
|
|
self.task_id = task_id
|
|
|
|
async def set_progress(self, progress: float, message: str | None = None) -> None:
|
|
await self._tasker._update_task(
|
|
self.task_id,
|
|
progress=max(0.0, min(progress, 100.0)),
|
|
message=message,
|
|
)
|
|
|
|
async def set_message(self, message: str) -> None:
|
|
await self._tasker._update_task(self.task_id, message=message)
|
|
|
|
async def set_result(self, result: Any) -> None:
|
|
await self._tasker._update_task(self.task_id, result=result)
|
|
|
|
def is_cancel_requested(self) -> bool:
|
|
return self._tasker._is_cancel_requested(self.task_id)
|
|
|
|
async def raise_if_cancelled(self) -> None:
|
|
if self.is_cancel_requested():
|
|
raise asyncio.CancelledError("Task was cancelled")
|
|
|
|
|
|
class Tasker:
|
|
def __init__(self, worker_count: int = 2):
|
|
self.worker_count = max(1, worker_count)
|
|
self._queue: asyncio.Queue[tuple[str, TaskCoroutine]] = asyncio.Queue()
|
|
self._tasks: dict[str, Task] = {}
|
|
self._lock = asyncio.Lock()
|
|
self._workers: list[asyncio.Task[Any]] = []
|
|
self._storage_path = Path(config.save_dir) / "tasks" / "tasks.json"
|
|
os.makedirs(self._storage_path.parent, exist_ok=True)
|
|
self._started = False
|
|
|
|
async def start(self) -> None:
|
|
async with self._lock:
|
|
if self._started:
|
|
return
|
|
await self._load_state()
|
|
for _ in range(self.worker_count):
|
|
worker = asyncio.create_task(self._worker_loop(), name="tasker-worker")
|
|
self._workers.append(worker)
|
|
self._started = True
|
|
logger.info("Tasker started with {} workers", self.worker_count)
|
|
|
|
async def shutdown(self) -> None:
|
|
async with self._lock:
|
|
if not self._started:
|
|
return
|
|
for worker in self._workers:
|
|
worker.cancel()
|
|
await asyncio.gather(*self._workers, return_exceptions=True)
|
|
self._workers.clear()
|
|
await self._persist_state()
|
|
self._started = False
|
|
logger.info("Tasker shutdown complete")
|
|
|
|
async def enqueue(
|
|
self,
|
|
*,
|
|
name: str,
|
|
task_type: str,
|
|
payload: dict[str, Any] | None = None,
|
|
coroutine: TaskCoroutine,
|
|
) -> Task:
|
|
task_id = uuid.uuid4().hex
|
|
task = Task(id=task_id, name=name, type=task_type, payload=payload or {})
|
|
async with self._lock:
|
|
self._tasks[task_id] = task
|
|
await self._persist_state()
|
|
await self._queue.put((task_id, coroutine))
|
|
logger.info("Enqueued task {} ({})", task_id, name)
|
|
return task
|
|
|
|
async def list_tasks(self, status: str | None = None, limit: int = 100) -> dict[str, Any]:
|
|
async with self._lock:
|
|
all_tasks = list(self._tasks.values())
|
|
|
|
status_counter = Counter(task.status for task in all_tasks)
|
|
type_counter = Counter(task.type for task in all_tasks)
|
|
all_tasks.sort(key=lambda item: item.created_at or utc_isoformat(), reverse=True)
|
|
|
|
tasks = all_tasks
|
|
if status:
|
|
tasks = [task for task in tasks if task.status == status]
|
|
|
|
limited_tasks = tasks[: max(limit, 0)]
|
|
|
|
summary: dict[str, Any] = {
|
|
"total": len(all_tasks),
|
|
"filtered_total": len(tasks),
|
|
"status_counts": dict(status_counter),
|
|
"type_counts": dict(type_counter),
|
|
}
|
|
|
|
return {
|
|
"tasks": [task.to_summary_dict() for task in limited_tasks],
|
|
"summary": summary,
|
|
}
|
|
|
|
async def get_task(self, task_id: str) -> dict[str, Any] | None:
|
|
async with self._lock:
|
|
task = self._tasks.get(task_id)
|
|
return task.to_dict() if task else None
|
|
|
|
async def cancel_task(self, task_id: str) -> bool:
|
|
async with self._lock:
|
|
task = self._tasks.get(task_id)
|
|
if not task:
|
|
return False
|
|
if task.status in {"success", "failed", "cancelled"}:
|
|
return False
|
|
task.cancel_requested = True
|
|
task.updated_at = _utc_timestamp()
|
|
await self._persist_state()
|
|
logger.info("Cancellation requested for task {}", task_id)
|
|
return True
|
|
|
|
async def _worker_loop(self) -> None:
|
|
while True:
|
|
try:
|
|
task_id, coroutine = await self._queue.get()
|
|
try:
|
|
task = await self._get_task_instance(task_id)
|
|
if not task:
|
|
continue
|
|
if task.cancel_requested:
|
|
await self._mark_cancelled(task_id, "Task was cancelled before execution")
|
|
continue
|
|
await self._update_task(
|
|
task_id, status="running", progress=0.0, message="任务开始执行", started_at=_utc_timestamp()
|
|
)
|
|
context = TaskContext(self, task_id)
|
|
try:
|
|
result = await coroutine(context)
|
|
if task.cancel_requested:
|
|
await self._mark_cancelled(task_id, "Task cancelled during execution")
|
|
continue
|
|
await self._update_task(
|
|
task_id,
|
|
status="success",
|
|
progress=100.0,
|
|
message="任务已完成",
|
|
result=result,
|
|
completed_at=_utc_timestamp(),
|
|
)
|
|
except asyncio.CancelledError:
|
|
await self._mark_cancelled(task_id, "任务被取消")
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.exception("Task {} failed: {}", task_id, exc)
|
|
await self._update_task(
|
|
task_id,
|
|
status="failed",
|
|
progress=100.0,
|
|
message="任务执行失败",
|
|
error=str(exc),
|
|
completed_at=_utc_timestamp(),
|
|
)
|
|
finally:
|
|
self._queue.task_done()
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.exception("Tasker worker error: {}", exc)
|
|
|
|
async def _get_task_instance(self, task_id: str) -> Task | None:
|
|
async with self._lock:
|
|
return self._tasks.get(task_id)
|
|
|
|
async def _mark_cancelled(self, task_id: str, message: str) -> None:
|
|
await self._update_task(
|
|
task_id,
|
|
status="cancelled",
|
|
progress=100.0,
|
|
message=message,
|
|
completed_at=_utc_timestamp(),
|
|
)
|
|
|
|
async def _update_task(
|
|
self,
|
|
task_id: str,
|
|
*,
|
|
status: str | None = None,
|
|
progress: float | None = None,
|
|
message: str | None = None,
|
|
result: Any = None,
|
|
error: str | None = None,
|
|
started_at: str | None = None,
|
|
completed_at: str | None = None,
|
|
) -> None:
|
|
async with self._lock:
|
|
task = self._tasks.get(task_id)
|
|
if not task:
|
|
return
|
|
if status:
|
|
task.status = status
|
|
if progress is not None:
|
|
task.progress = max(0.0, min(progress, 100.0))
|
|
if message is not None:
|
|
task.message = message
|
|
if result is not None:
|
|
task.result = result
|
|
if error is not None:
|
|
task.error = error
|
|
if started_at is not None:
|
|
task.started_at = started_at
|
|
if completed_at is not None:
|
|
task.completed_at = completed_at
|
|
task.updated_at = _utc_timestamp()
|
|
await self._persist_state()
|
|
|
|
def _is_cancel_requested(self, task_id: str) -> bool:
|
|
task = self._tasks.get(task_id)
|
|
return bool(task and task.cancel_requested)
|
|
|
|
async def _load_state(self) -> None:
|
|
if not self._storage_path.exists():
|
|
return
|
|
try:
|
|
content = await asyncio.to_thread(self._storage_path.read_text, encoding="utf-8")
|
|
if not content.strip():
|
|
return
|
|
data = json.loads(content)
|
|
tasks = data.get("tasks", [])
|
|
for item in tasks:
|
|
task = Task.from_dict(item)
|
|
if task.status == "running":
|
|
task.status = "failed"
|
|
task.message = "服务重启时任务中断"
|
|
task.updated_at = _utc_timestamp()
|
|
elif task.status not in TERMINAL_STATUSES:
|
|
task.status = "failed"
|
|
task.message = "服务重启时任务未继续执行"
|
|
task.updated_at = _utc_timestamp()
|
|
self._tasks[task.id] = task
|
|
logger.info("Loaded {} task records from storage", len(tasks))
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.exception("Failed to load task state: {}", exc)
|
|
|
|
async def _persist_state(self) -> None:
|
|
tasks = [task.to_dict() for task in self._tasks.values()]
|
|
payload = {"tasks": tasks, "updated_at": _utc_timestamp()}
|
|
|
|
def _write() -> None:
|
|
self._storage_path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp_path = self._storage_path.with_suffix(".tmp")
|
|
with open(tmp_path, "w", encoding="utf-8") as fh:
|
|
json.dump(payload, fh, ensure_ascii=False, indent=2)
|
|
os.replace(tmp_path, self._storage_path)
|
|
|
|
await asyncio.to_thread(_write)
|
|
|
|
|
|
tasker = Tasker()
|
|
|
|
|
|
__all__ = ["tasker", "TaskContext", "Tasker"]
|