wehub-resource-sync
723 lines
30 KiB
Python
723 lines
30 KiB
Python
import json
|
|
import re
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from yuxi.knowledge import knowledge_base
|
|
from yuxi.knowledge.eval.benchmark_generation import (
|
|
dump_benchmark_item,
|
|
iter_generated_benchmark_items,
|
|
normalize_generation_concurrency_count,
|
|
)
|
|
from yuxi.knowledge.eval.evaluator import aggregate_metrics, evaluate_question
|
|
from yuxi.models import select_model
|
|
from yuxi.repositories.evaluation_repository import EvaluationRepository
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
|
from yuxi.repositories.knowledge_chunk_repository import KnowledgeChunkRepository
|
|
from yuxi.repositories.task_repository import TaskRepository
|
|
from yuxi.services.task_service import TaskContext, tasker
|
|
from yuxi.utils import logger
|
|
from yuxi.utils.datetime_utils import format_utc_datetime, utc_now_naive
|
|
|
|
|
|
def build_evaluation_run_name(started_at=None, hash_value: str | None = None) -> str:
|
|
date_part = (started_at or utc_now_naive()).strftime("%Y%m%d")
|
|
hash_part = re.sub(r"[^a-fA-F0-9]", "", hash_value or uuid.uuid4().hex).lower()[:6]
|
|
if len(hash_part) < 6:
|
|
hash_part = (hash_part + uuid.uuid4().hex)[:6]
|
|
return f"eval-{date_part}-{hash_part}"
|
|
|
|
|
|
class EvaluationService:
|
|
"""RAG评估服务"""
|
|
|
|
def __init__(self):
|
|
self.eval_repo = EvaluationRepository()
|
|
self.kb_repo = KnowledgeBaseRepository()
|
|
self.chunk_repo = KnowledgeChunkRepository()
|
|
self.task_repo = TaskRepository()
|
|
|
|
def _dataset_to_dict(self, row) -> dict[str, Any]:
|
|
return {
|
|
"id": row.dataset_id,
|
|
"dataset_id": row.dataset_id,
|
|
"name": row.name,
|
|
"description": row.description,
|
|
"kb_id": row.kb_id,
|
|
"item_count": row.item_count,
|
|
"has_gold_chunks": row.has_gold_chunks,
|
|
"has_gold_answers": row.has_gold_answers,
|
|
"build_metadata": row.build_metadata or {},
|
|
"created_by": row.created_by,
|
|
"created_at": format_utc_datetime(row.created_at),
|
|
"updated_at": format_utc_datetime(row.updated_at),
|
|
}
|
|
|
|
def _dataset_item_to_dict(self, item) -> dict[str, Any]:
|
|
return {
|
|
"item_id": item.item_id,
|
|
"item_index": item.item_index,
|
|
"query": item.query_text,
|
|
"gold_chunk_ids": item.gold_chunk_ids or [],
|
|
"gold_answer": item.gold_answer,
|
|
}
|
|
|
|
def _run_item_to_dict(self, item) -> dict[str, Any]:
|
|
return {
|
|
"query": item.query_text,
|
|
"gold_chunk_ids": item.gold_chunk_ids,
|
|
"gold_answer": item.gold_answer,
|
|
"generated_answer": item.generated_answer,
|
|
"retrieved_chunks": item.retrieved_chunks,
|
|
"metrics": item.metrics or {},
|
|
}
|
|
|
|
def _is_error_run_item(self, item) -> bool:
|
|
metrics = item.metrics or {}
|
|
return metrics.get("score", 1.0) <= 0.5 or any(
|
|
metrics.get(key, 1.0) < 0.3 for key in metrics if key.startswith("recall@")
|
|
)
|
|
|
|
def _normalize_run_name(self, name: str | None, run_id: str) -> str:
|
|
run_name = (name or "").strip()
|
|
if run_name:
|
|
return run_name
|
|
return build_evaluation_run_name(hash_value=run_id.removeprefix("run_"))
|
|
|
|
def _run_name_from_row(self, row) -> str:
|
|
name = (getattr(row, "name", None) or "").strip()
|
|
if name:
|
|
return name
|
|
return build_evaluation_run_name(row.started_at, hash_value=row.run_id.removeprefix("run_"))
|
|
|
|
async def _sync_dataset_build_metadata(self, row) -> None:
|
|
metadata = dict(row.build_metadata or {})
|
|
if metadata.get("source") != "generated" or metadata.get("status") not in {"pending", "running"}:
|
|
return
|
|
|
|
task_id = metadata.get("task_id")
|
|
task = await self.task_repo.get_by_id(task_id) if task_id else None
|
|
if task is None:
|
|
metadata.pop("progress", None)
|
|
metadata.update(status="failed", message="生成任务不存在")
|
|
elif task.status == "success":
|
|
metadata.update(status="completed", progress=100, message=task.message or "完成")
|
|
elif task.status in {"failed", "cancelled"}:
|
|
metadata.pop("progress", None)
|
|
metadata.update(status="failed", message=task.error or task.message or "生成任务失败")
|
|
else:
|
|
metadata.update(status=task.status, progress=task.progress, message=task.message)
|
|
|
|
if metadata != (row.build_metadata or {}):
|
|
await self.eval_repo.update_dataset(row.dataset_id, {"build_metadata": metadata})
|
|
row.build_metadata = metadata
|
|
|
|
def _build_dataset_items(
|
|
self, dataset_id: str, kb_id: str, questions: list[dict[str, Any]]
|
|
) -> list[dict[str, Any]]:
|
|
return [
|
|
{
|
|
"item_id": f"dataset_item_{uuid.uuid4().hex[:12]}",
|
|
"dataset_id": dataset_id,
|
|
"kb_id": kb_id,
|
|
"item_index": index,
|
|
"query_text": item["query"],
|
|
"gold_chunk_ids": item.get("gold_chunk_ids") or [],
|
|
"gold_answer": item.get("gold_answer"),
|
|
}
|
|
for index, item in enumerate(questions)
|
|
]
|
|
|
|
def _build_jsonl_content(self, items: list[Any]) -> str:
|
|
lines = []
|
|
for item in items:
|
|
payload = {"query": item.query_text}
|
|
if item.gold_chunk_ids:
|
|
payload["gold_chunk_ids"] = item.gold_chunk_ids
|
|
if item.gold_answer:
|
|
payload["gold_answer"] = item.gold_answer
|
|
lines.append(dump_benchmark_item(payload).rstrip("\n"))
|
|
return "\n".join(lines) + ("\n" if lines else "")
|
|
|
|
def _safe_jsonl_filename(self, name: str | None, fallback: str) -> str:
|
|
filename = (name or "").strip() or fallback
|
|
filename = re.sub(r"[\\/:*?\"<>|]+", "_", filename).strip()
|
|
if not filename or filename in {".", ".."}:
|
|
filename = fallback
|
|
return filename if filename.endswith(".jsonl") else f"{filename}.jsonl"
|
|
|
|
def _parse_jsonl_questions(self, file_content: bytes) -> tuple[list[dict[str, Any]], bool, bool]:
|
|
questions = []
|
|
has_gold_chunks = False
|
|
has_gold_answers = False
|
|
content = file_content.decode("utf-8")
|
|
|
|
for line_num, line in enumerate(content.strip().split("\n"), 1):
|
|
if not line.strip():
|
|
continue
|
|
try:
|
|
item = json.loads(line)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"第{line_num}行JSON格式错误: {str(e)}")
|
|
if "query" not in item:
|
|
raise ValueError(f"第{line_num}行缺少必需的'query'字段")
|
|
if item.get("gold_chunk_ids"):
|
|
has_gold_chunks = True
|
|
if item.get("gold_answer"):
|
|
has_gold_answers = True
|
|
questions.append(item)
|
|
|
|
if not questions:
|
|
raise ValueError("文件中没有有效的问题数据")
|
|
return questions, has_gold_chunks, has_gold_answers
|
|
|
|
async def upload_dataset(
|
|
self, kb_id: str, file_content: bytes, filename: str, name: str, description: str, created_by: str
|
|
) -> dict[str, Any]:
|
|
try:
|
|
questions, has_gold_chunks, has_gold_answers = self._parse_jsonl_questions(file_content)
|
|
dataset_id = f"dataset_{uuid.uuid4().hex[:8]}"
|
|
dataset_name = name.strip() or filename or dataset_id
|
|
|
|
row = await self.eval_repo.create_dataset_with_items(
|
|
{
|
|
"dataset_id": dataset_id,
|
|
"kb_id": kb_id,
|
|
"name": dataset_name,
|
|
"description": description,
|
|
"item_count": len(questions),
|
|
"has_gold_chunks": has_gold_chunks,
|
|
"has_gold_answers": has_gold_answers,
|
|
"build_metadata": {
|
|
"source": "upload",
|
|
"status": "completed",
|
|
"progress": 100,
|
|
"filename": filename,
|
|
},
|
|
"created_by": created_by,
|
|
},
|
|
self._build_dataset_items(dataset_id, kb_id, questions),
|
|
)
|
|
return self._dataset_to_dict(row)
|
|
except Exception as e:
|
|
logger.error(f"上传评估数据集失败: {e}")
|
|
raise
|
|
|
|
async def list_datasets(self, kb_id: str) -> list[dict[str, Any]]:
|
|
try:
|
|
rows = await self.eval_repo.list_datasets(kb_id)
|
|
for row in rows:
|
|
await self._sync_dataset_build_metadata(row)
|
|
return [self._dataset_to_dict(row) for row in rows]
|
|
except Exception as e:
|
|
logger.error(f"获取评估数据集列表失败: {e}")
|
|
raise
|
|
|
|
async def get_dataset_detail(
|
|
self, kb_id: str, dataset_id: str, page: int = 1, page_size: int = 10
|
|
) -> dict[str, Any]:
|
|
try:
|
|
row = await self.eval_repo.get_dataset(dataset_id)
|
|
if row is None or row.kb_id != kb_id:
|
|
raise ValueError("Dataset not found")
|
|
if (row.build_metadata or {}).get("status", "completed") != "completed":
|
|
raise ValueError("Dataset is not ready")
|
|
|
|
total_items = await self.eval_repo.count_dataset_items(dataset_id)
|
|
items = await self.eval_repo.list_dataset_items(dataset_id, (page - 1) * page_size, page_size)
|
|
total_pages = (total_items + page_size - 1) // page_size
|
|
data = self._dataset_to_dict(row)
|
|
data.update(
|
|
{
|
|
"items": [self._dataset_item_to_dict(item) for item in items],
|
|
"pagination": {
|
|
"current_page": page,
|
|
"page_size": page_size,
|
|
"total_items": total_items,
|
|
"total_pages": total_pages,
|
|
"has_next": page < total_pages,
|
|
"has_prev": page > 1,
|
|
},
|
|
}
|
|
)
|
|
return data
|
|
except Exception as e:
|
|
logger.error(f"获取评估数据集详情失败: {e}")
|
|
raise
|
|
|
|
async def export_dataset_jsonl(self, dataset_id: str) -> dict[str, str]:
|
|
row = await self.eval_repo.get_dataset(dataset_id)
|
|
if row is None:
|
|
raise ValueError("Dataset not found")
|
|
if (row.build_metadata or {}).get("status", "completed") != "completed":
|
|
raise ValueError("Dataset is not ready")
|
|
items = await self.eval_repo.list_all_dataset_items(dataset_id)
|
|
return {
|
|
"filename": self._safe_jsonl_filename(row.name, row.dataset_id),
|
|
"content": self._build_jsonl_content(items),
|
|
}
|
|
|
|
async def delete_dataset(self, dataset_id: str) -> None:
|
|
try:
|
|
row = await self.eval_repo.get_dataset(dataset_id)
|
|
if row is None:
|
|
raise ValueError("Dataset not found")
|
|
await self.eval_repo.delete_dataset(dataset_id)
|
|
logger.info(f"成功删除评估数据集: {dataset_id}")
|
|
except Exception as e:
|
|
logger.error(f"删除评估数据集失败: {e}")
|
|
raise
|
|
|
|
async def generate_dataset(
|
|
self,
|
|
kb_id: str,
|
|
name: str,
|
|
description: str,
|
|
count: int,
|
|
neighbors_count: int,
|
|
concurrency_count: int,
|
|
llm_model_spec: str,
|
|
generation_mode: str = "vector",
|
|
graph_expand_top_k: int = 1,
|
|
created_by: str = "system",
|
|
) -> dict[str, Any]:
|
|
dataset_id = f"dataset_{uuid.uuid4().hex[:8]}"
|
|
count = int(count)
|
|
neighbors_count = int(neighbors_count)
|
|
concurrency_count = normalize_generation_concurrency_count(concurrency_count)
|
|
graph_expand_top_k = min(max(1, int(graph_expand_top_k)), 3)
|
|
if generation_mode not in {"vector", "graph_enhanced"}:
|
|
raise ValueError("不支持的评估基准生成方式")
|
|
if generation_mode == "graph_enhanced":
|
|
indexed_count = await self.chunk_repo.count_graph_indexed_by_kb_id(kb_id)
|
|
if indexed_count <= 0:
|
|
raise ValueError("当前知识库尚未完成图索引，无法使用图增强构建")
|
|
build_metadata = {
|
|
"source": "generated",
|
|
"status": "pending",
|
|
"progress": 0,
|
|
"params": {
|
|
"count": count,
|
|
"neighbors_count": neighbors_count,
|
|
"concurrency_count": concurrency_count,
|
|
"llm_model_spec": llm_model_spec,
|
|
"generation_mode": generation_mode,
|
|
"graph_expand_top_k": graph_expand_top_k,
|
|
},
|
|
}
|
|
await self.eval_repo.create_dataset(
|
|
{
|
|
"dataset_id": dataset_id,
|
|
"kb_id": kb_id,
|
|
"name": name,
|
|
"description": description,
|
|
"item_count": 0,
|
|
"has_gold_chunks": True,
|
|
"has_gold_answers": True,
|
|
"build_metadata": build_metadata,
|
|
"created_by": created_by,
|
|
}
|
|
)
|
|
task = await tasker.enqueue(
|
|
name="生成评估数据集",
|
|
task_type="dataset_generation",
|
|
payload={
|
|
"dataset_id": dataset_id,
|
|
"kb_id": kb_id,
|
|
"created_by": created_by,
|
|
"name": name,
|
|
"description": description,
|
|
"count": count,
|
|
"neighbors_count": neighbors_count,
|
|
"concurrency_count": concurrency_count,
|
|
"llm_model_spec": llm_model_spec,
|
|
"generation_mode": generation_mode,
|
|
"graph_expand_top_k": graph_expand_top_k,
|
|
},
|
|
coroutine=self._generate_dataset_task,
|
|
)
|
|
build_metadata["task_id"] = task.id
|
|
await self.eval_repo.update_dataset(dataset_id, {"build_metadata": build_metadata})
|
|
return {"dataset_id": dataset_id, "task_id": task.id, "message": "评估数据集生成任务已提交"}
|
|
|
|
async def _update_dataset_build_metadata(
|
|
self, dataset_id: str, metadata: dict[str, Any], **updates
|
|
) -> dict[str, Any]:
|
|
metadata.update(updates)
|
|
await self.eval_repo.update_dataset(dataset_id, {"build_metadata": metadata})
|
|
return metadata
|
|
|
|
async def _generate_dataset_task(self, context: TaskContext):
|
|
await context.set_progress(0, "初始化")
|
|
payload = context.payload
|
|
|
|
dataset_id = payload.get("dataset_id")
|
|
kb_id = payload.get("kb_id")
|
|
count = int(payload.get("count", 10))
|
|
neighbors_count = int(payload.get("neighbors_count", 1))
|
|
concurrency_count = normalize_generation_concurrency_count(payload.get("concurrency_count"))
|
|
llm_model_spec = payload.get("llm_model_spec")
|
|
generation_mode = payload.get("generation_mode") or "vector"
|
|
graph_expand_top_k = min(max(1, int(payload.get("graph_expand_top_k", 1))), 3)
|
|
build_metadata = {
|
|
"source": "generated",
|
|
"status": "running",
|
|
"progress": 0,
|
|
"task_id": context.task_id,
|
|
"params": {
|
|
"count": count,
|
|
"neighbors_count": neighbors_count,
|
|
"concurrency_count": concurrency_count,
|
|
"llm_model_spec": llm_model_spec,
|
|
"generation_mode": generation_mode,
|
|
"graph_expand_top_k": graph_expand_top_k,
|
|
},
|
|
}
|
|
await self._update_dataset_build_metadata(dataset_id, build_metadata)
|
|
|
|
async def report_progress(progress: float, message: str | None = None) -> None:
|
|
await context.set_progress(progress, message)
|
|
await self._update_dataset_build_metadata(
|
|
dataset_id,
|
|
build_metadata,
|
|
progress=max(0, min(round(progress), 100)),
|
|
message=message or build_metadata.get("message", ""),
|
|
)
|
|
|
|
try:
|
|
kb_instance = await knowledge_base.aget_kb(kb_id)
|
|
if not kb_instance:
|
|
await report_progress(100, "知识库不存在")
|
|
raise ValueError("Knowledge Base not found")
|
|
if kb_instance.kb_type != "milvus":
|
|
await report_progress(100, "仅支持 commonrag/Milvus 类型知识库生成评估数据集")
|
|
raise ValueError("Unsupported KB type for dataset generation")
|
|
|
|
questions = []
|
|
try:
|
|
async for item in iter_generated_benchmark_items(
|
|
kb_instance=kb_instance,
|
|
kb_id=kb_id,
|
|
count=count,
|
|
neighbors_count=neighbors_count,
|
|
llm_model_spec=llm_model_spec,
|
|
concurrency_count=concurrency_count,
|
|
generation_mode=generation_mode,
|
|
graph_expand_top_k=graph_expand_top_k,
|
|
progress_cb=report_progress,
|
|
cancel_cb=context.raise_if_cancelled,
|
|
):
|
|
questions.append(item)
|
|
except ValueError as e:
|
|
if str(e) == "No chunks found in knowledge base":
|
|
await report_progress(100, "知识库为空或未解析到chunks")
|
|
raise
|
|
|
|
if not questions:
|
|
raise ValueError("未生成有效评估题目")
|
|
|
|
await self.eval_repo.add_dataset_items(self._build_dataset_items(dataset_id, kb_id, questions))
|
|
await self.eval_repo.update_dataset(dataset_id, {"item_count": len(questions)})
|
|
await self._update_dataset_build_metadata(
|
|
dataset_id,
|
|
build_metadata,
|
|
status="completed",
|
|
progress=100,
|
|
message="完成",
|
|
)
|
|
await context.set_progress(100, "完成")
|
|
except Exception as e:
|
|
await self._update_dataset_build_metadata(
|
|
dataset_id,
|
|
build_metadata,
|
|
status="failed",
|
|
progress=100,
|
|
error_message=str(e),
|
|
message=str(e),
|
|
)
|
|
raise
|
|
|
|
async def run_evaluation(
|
|
self,
|
|
kb_id: str,
|
|
dataset_id: str,
|
|
name: str | None = None,
|
|
model_config: dict[str, Any] = None,
|
|
created_by: str = "system",
|
|
) -> str:
|
|
try:
|
|
run_id = f"run_{uuid.uuid4().hex[:8]}"
|
|
run_name = self._normalize_run_name(name, run_id)
|
|
dataset_row = await self.eval_repo.get_dataset(dataset_id)
|
|
if dataset_row is None or dataset_row.kb_id != kb_id:
|
|
raise ValueError("Dataset not found")
|
|
if (dataset_row.build_metadata or {}).get("status", "completed") != "completed":
|
|
raise ValueError("Dataset is not ready")
|
|
|
|
retrieval_config = {}
|
|
try:
|
|
kb_row = await self.kb_repo.get_by_kb_id(kb_id)
|
|
query_params = (kb_row.query_params if kb_row else None) or {}
|
|
retrieval_config = query_params.get("options", {}) if isinstance(query_params, dict) else {}
|
|
if not retrieval_config:
|
|
kb_instance = await knowledge_base.aget_kb(kb_id)
|
|
if kb_instance:
|
|
retrieval_config = kb_instance._get_default_query_params(kb_id).get("options", {})
|
|
logger.info(f"从知识库 {kb_id} 加载检索配置: {list(retrieval_config.keys())}")
|
|
except Exception as e:
|
|
logger.error(f"获取知识库检索配置失败: {e}")
|
|
|
|
if model_config:
|
|
retrieval_config.update(model_config)
|
|
|
|
await self.eval_repo.create_run(
|
|
{
|
|
"run_id": run_id,
|
|
"name": run_name,
|
|
"kb_id": kb_id,
|
|
"dataset_id": dataset_id,
|
|
"status": "running",
|
|
"retrieval_config": retrieval_config,
|
|
"metrics": {},
|
|
"overall_score": None,
|
|
"total_items": dataset_row.item_count or 0,
|
|
"completed_items": 0,
|
|
"started_at": utc_now_naive(),
|
|
"completed_at": None,
|
|
"created_by": created_by,
|
|
}
|
|
)
|
|
|
|
await tasker.enqueue(
|
|
name=f"RAG评估({run_name})",
|
|
task_type="rag_evaluation",
|
|
payload={
|
|
"run_id": run_id,
|
|
"name": run_name,
|
|
"kb_id": kb_id,
|
|
"dataset_id": dataset_id,
|
|
"retrieval_config": retrieval_config,
|
|
"created_by": created_by,
|
|
},
|
|
coroutine=self._run_evaluation_task,
|
|
)
|
|
return run_id
|
|
except Exception as e:
|
|
logger.error(f"启动评估失败: {e}")
|
|
raise
|
|
|
|
async def _run_evaluation_task(self, context: TaskContext):
|
|
try:
|
|
payload = context.payload
|
|
|
|
run_id = payload["run_id"]
|
|
kb_id = payload["kb_id"]
|
|
dataset_id = payload["dataset_id"]
|
|
retrieval_config = payload["retrieval_config"]
|
|
|
|
await context.set_progress(5, "加载评估数据集")
|
|
dataset_row = await self.eval_repo.get_dataset(dataset_id)
|
|
if dataset_row is None or dataset_row.kb_id != kb_id:
|
|
raise ValueError("Dataset not found")
|
|
dataset_items = await self.eval_repo.list_all_dataset_items(dataset_id)
|
|
if not dataset_items:
|
|
raise ValueError("Dataset has no items")
|
|
|
|
kb_instance = await knowledge_base.aget_kb(kb_id)
|
|
if not kb_instance:
|
|
raise ValueError(f"Knowledge Base {kb_id} not found")
|
|
|
|
judge_llm = None
|
|
if dataset_row.has_gold_answers:
|
|
judge_model_spec = retrieval_config.get("judge_llm") or retrieval_config.get("answer_llm")
|
|
if judge_model_spec:
|
|
try:
|
|
logger.debug(f"Initializing Judge LLM: {judge_model_spec}")
|
|
judge_llm = select_model(model_spec=judge_model_spec)
|
|
except Exception as e:
|
|
logger.error(f"Failed to load judge LLM: {e}")
|
|
|
|
all_retrieval_metrics = []
|
|
all_answer_metrics = []
|
|
total_items = len(dataset_items)
|
|
|
|
async def update_run_db(status=None, completed=None, metrics=None, final_score=None):
|
|
data = {}
|
|
if status is not None:
|
|
data["status"] = status
|
|
if status in ["completed", "failed"]:
|
|
data["completed_at"] = utc_now_naive()
|
|
if completed is not None:
|
|
data["completed_items"] = completed
|
|
if metrics is not None:
|
|
data["metrics"] = metrics
|
|
if final_score is not None:
|
|
data["overall_score"] = final_score
|
|
if data:
|
|
await self.eval_repo.update_run(run_id, data)
|
|
|
|
for index, item in enumerate(dataset_items):
|
|
await context.raise_if_cancelled()
|
|
progress = 10 + (index / total_items) * 80
|
|
await context.set_progress(progress, f"评估 {index + 1}/{total_items}")
|
|
|
|
question_data = {
|
|
"query": item.query_text,
|
|
"gold_chunk_ids": item.gold_chunk_ids or [],
|
|
"gold_answer": item.gold_answer,
|
|
}
|
|
question_result = await evaluate_question(
|
|
kb_instance=kb_instance,
|
|
kb_id=kb_id,
|
|
question_data=question_data,
|
|
retrieval_config=retrieval_config,
|
|
has_gold_chunks=dataset_row.has_gold_chunks,
|
|
has_gold_answers=dataset_row.has_gold_answers,
|
|
judge_llm=judge_llm,
|
|
select_model_fn=select_model,
|
|
)
|
|
|
|
if dataset_row.has_gold_chunks and question_data.get("gold_chunk_ids"):
|
|
all_retrieval_metrics.append(question_result["retrieval_scores"])
|
|
if dataset_row.has_gold_answers and question_data.get("gold_answer") and judge_llm:
|
|
all_answer_metrics.append(question_result["answer_scores"])
|
|
|
|
await self.eval_repo.upsert_run_item(
|
|
run_id=run_id,
|
|
item_index=index,
|
|
data={"dataset_item_id": item.item_id, **question_result["detail"]},
|
|
)
|
|
|
|
if (index + 1) % 5 == 0 or (index + 1) == total_items:
|
|
current_metrics, _ = aggregate_metrics(all_retrieval_metrics, all_answer_metrics)
|
|
await context.set_result(
|
|
{"current_metrics": current_metrics, "completed_items": index + 1, "total_items": total_items}
|
|
)
|
|
await update_run_db(completed=index + 1)
|
|
|
|
await context.set_progress(95, "计算最终指标")
|
|
overall_metrics, overall_score = aggregate_metrics(
|
|
all_retrieval_metrics, all_answer_metrics, include_overall_score=True
|
|
)
|
|
await update_run_db(
|
|
status="completed",
|
|
completed=total_items,
|
|
metrics=overall_metrics,
|
|
final_score=overall_score,
|
|
)
|
|
await context.set_progress(100, "完成")
|
|
except Exception as e:
|
|
logger.error(f"Task failed: {e}")
|
|
try:
|
|
if "payload" in locals():
|
|
await self.eval_repo.update_run(
|
|
payload["run_id"],
|
|
{"status": "failed", "metrics": {"error": str(e)}, "completed_at": utc_now_naive()},
|
|
)
|
|
except Exception as exc:
|
|
logger.error(f"Error updating run record: {exc}")
|
|
await context.set_message(f"Error: {str(e)}")
|
|
raise
|
|
|
|
async def list_runs(self, kb_id: str) -> list[dict[str, Any]]:
|
|
try:
|
|
rows = await self.eval_repo.list_runs(kb_id)
|
|
running_run_ids = {row.run_id for row in rows if row.status == "running"}
|
|
task_by_run_id = {}
|
|
if running_run_ids:
|
|
tasks = await self.task_repo.list_all()
|
|
task_by_run_id = {
|
|
(task.payload or {}).get("run_id"): task
|
|
for task in tasks
|
|
if task.type == "rag_evaluation"
|
|
and task.status in {"pending", "running"}
|
|
and (task.payload or {}).get("run_id") in running_run_ids
|
|
}
|
|
|
|
runs = []
|
|
for row in rows:
|
|
run = {
|
|
"run_id": row.run_id,
|
|
"name": self._run_name_from_row(row),
|
|
"dataset_id": row.dataset_id,
|
|
"status": row.status,
|
|
"started_at": format_utc_datetime(row.started_at),
|
|
"completed_at": format_utc_datetime(row.completed_at),
|
|
"total_items": row.total_items,
|
|
"completed_items": row.completed_items,
|
|
"overall_score": row.overall_score,
|
|
"retrieval_config": row.retrieval_config or {},
|
|
"metrics": row.metrics or {},
|
|
}
|
|
if row.status == "running":
|
|
task = task_by_run_id.get(row.run_id)
|
|
if task:
|
|
run.update(progress=task.progress, message=task.message)
|
|
runs.append(run)
|
|
return runs
|
|
except Exception as e:
|
|
logger.error(f"获取评估运行历史失败: {e}")
|
|
raise
|
|
|
|
async def get_run_results(
|
|
self, kb_id: str, run_id: str, page: int = 1, page_size: int = 20, error_only: bool = False
|
|
) -> dict[str, Any]:
|
|
if not re.match(r"^run_[a-f0-9]{8}$", run_id):
|
|
raise ValueError("Invalid run_id format")
|
|
row = await self.eval_repo.get_run(run_id)
|
|
if row is None or row.kb_id != kb_id:
|
|
task = await tasker.get_task(run_id)
|
|
if task:
|
|
return {"run_id": run_id, "status": task.status, "progress": task.progress, "message": task.message}
|
|
raise ValueError(f"Run not found for {run_id}")
|
|
|
|
start_idx = (page - 1) * page_size
|
|
if error_only:
|
|
total = 0
|
|
paged_items = []
|
|
offset = 0
|
|
batch_size = 200
|
|
while True:
|
|
batch = await self.eval_repo.list_run_items(run_id, offset, batch_size)
|
|
if not batch:
|
|
break
|
|
for item in batch:
|
|
if not self._is_error_run_item(item):
|
|
continue
|
|
if start_idx <= total < start_idx + page_size:
|
|
paged_items.append(self._run_item_to_dict(item))
|
|
total += 1
|
|
offset += batch_size
|
|
else:
|
|
total = await self.eval_repo.count_run_items(run_id)
|
|
details = await self.eval_repo.list_run_items(run_id, start_idx, page_size)
|
|
paged_items = [self._run_item_to_dict(item) for item in details]
|
|
return {
|
|
"run_id": row.run_id,
|
|
"name": self._run_name_from_row(row),
|
|
"status": row.status,
|
|
"started_at": format_utc_datetime(row.started_at),
|
|
"completed_at": format_utc_datetime(row.completed_at),
|
|
"total_items": row.total_items or 0,
|
|
"completed_items": row.completed_items or 0,
|
|
"overall_score": row.overall_score,
|
|
"retrieval_config": row.retrieval_config or {},
|
|
"items": paged_items,
|
|
"pagination": {
|
|
"current_page": page,
|
|
"page_size": page_size,
|
|
"total": total,
|
|
"total_pages": (total + page_size - 1) // page_size,
|
|
"error_only": error_only,
|
|
},
|
|
}
|
|
|
|
async def delete_run(self, kb_id: str, run_id: str) -> None:
|
|
if not re.match(r"^run_[a-f0-9]{8}$", run_id):
|
|
raise ValueError("Invalid run_id format")
|
|
row = await self.eval_repo.get_run(run_id)
|
|
if row is None or row.kb_id != kb_id:
|
|
raise ValueError("Run not found")
|
|
await self.eval_repo.delete_run(run_id)
|
|
logger.info(f"成功删除评估运行: {run_id}")
|