266 lines
10 KiB
Python
266 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""Export all loop-test Q&A batches for remote dagent from SQLite (fast path)."""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import sqlite3
|
||
import sys
|
||
from collections import defaultdict
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
ROOT = Path(__file__).resolve().parent.parent.parent
|
||
DB_PATH = ROOT / "server" / "data" / "rag_eval.db"
|
||
PLAN_PATH = ROOT / "docs" / "task_groups_plan.json"
|
||
EXPORT_DIR = ROOT / "docs" / "exports"
|
||
sys.path.insert(0, str(ROOT / "server"))
|
||
sys.path.insert(0, str(ROOT / "server" / "service"))
|
||
from loop_recall_md import DEFAULT_LLM_NOTE, append_recall_md_section # noqa: E402
|
||
|
||
|
||
def get_task_questions_fast(conn: sqlite3.Connection, task_id: str) -> list[dict]:
|
||
"""Approved Q&A from qa_gen_question; fallback to single_jump_result for legacy tasks."""
|
||
cursor = conn.cursor()
|
||
cursor.execute(
|
||
"""SELECT COUNT(*) as cnt FROM loop_round
|
||
WHERE loop_task_id=? AND qa_gen_task_id IS NOT NULL""",
|
||
(task_id,),
|
||
)
|
||
if cursor.fetchone()["cnt"] > 0:
|
||
cursor.execute(
|
||
"""SELECT
|
||
q.id as qa_question_id,
|
||
q.section_path, q.file_name, q.question, q.reference_answer,
|
||
q.source_chunk, q.quality_score, q.status,
|
||
q.dup_of, q.dup_similarity,
|
||
q.chunk_headers, q.chunk_id, q.file_id,
|
||
lr.round_number
|
||
FROM qa_gen_question q
|
||
JOIN loop_round lr ON q.task_id = lr.qa_gen_task_id
|
||
WHERE lr.loop_task_id = ? AND q.status = 'approved'
|
||
ORDER BY lr.round_number, q.chunk_headers, q.created_at""",
|
||
(task_id,),
|
||
)
|
||
return [dict(row) for row in cursor.fetchall()]
|
||
|
||
cursor.execute(
|
||
"""SELECT
|
||
r.section_path, r.file_name, r.question, r.reference_answer,
|
||
COALESCE(r.raw_chunk_headers, r.section_path) as chunk_headers,
|
||
r.expected_chunk_id as chunk_id,
|
||
lr.round_number
|
||
FROM single_jump_result r
|
||
JOIN loop_round lr ON r.task_id = lr.single_jump_task_id
|
||
WHERE lr.loop_task_id = ?
|
||
ORDER BY lr.round_number, r.section_path""",
|
||
(task_id,),
|
||
)
|
||
rows = []
|
||
for row in cursor.fetchall():
|
||
d = dict(row)
|
||
d.setdefault("quality_score", 1.0)
|
||
d.setdefault("status", "approved")
|
||
rows.append(d)
|
||
return rows
|
||
|
||
|
||
def rows_to_md(rows: list[dict]) -> str:
|
||
if not rows:
|
||
return ""
|
||
sections: dict[str, list] = defaultdict(list)
|
||
for row in rows:
|
||
key = row.get("chunk_headers") or row.get("section_path") or row.get("file_name") or "default"
|
||
sections[key].append(row)
|
||
|
||
lines: list[str] = []
|
||
for section_index, (section_key, items) in enumerate(sections.items(), 1):
|
||
file_name = (items[0].get("file_name") or "").strip()
|
||
slice_title = (items[0].get("chunk_headers") or "").strip() or section_key
|
||
meta = [f"> 代表轮次: {items[0]['round_number']}", DEFAULT_LLM_NOTE]
|
||
qa_items = [
|
||
{
|
||
"question": it["question"],
|
||
"reference_answer": it["reference_answer"],
|
||
"chunk_id": (it.get("chunk_id") or ""),
|
||
}
|
||
for it in items
|
||
]
|
||
append_recall_md_section(
|
||
lines, section_index,
|
||
file_name=file_name, slice_title=slice_title,
|
||
qa_items=qa_items, meta_lines=meta,
|
||
)
|
||
return "\n".join(lines)
|
||
|
||
|
||
def rows_to_json_questions(rows: list[dict]) -> list[dict]:
|
||
return [
|
||
{
|
||
"section_path": r.get("section_path"),
|
||
"file_name": r.get("file_name"),
|
||
"file_id": r.get("file_id"),
|
||
"chunk_headers": r.get("chunk_headers"),
|
||
"chunk_id": r.get("chunk_id"),
|
||
"round": r.get("round_number"),
|
||
"question": r["question"],
|
||
"reference_answer": r["reference_answer"],
|
||
"source_chunk": r.get("source_chunk"),
|
||
"quality_score": r.get("quality_score"),
|
||
"status": r.get("status"),
|
||
"is_duplicate": bool(r.get("dup_of")),
|
||
"dup_similarity": r.get("dup_similarity"),
|
||
"qa_question_id": r.get("qa_question_id"),
|
||
}
|
||
for r in rows
|
||
]
|
||
|
||
|
||
def resolve_task_id_from_db(conn: sqlite3.Connection, group_id: int, batch_id: int) -> dict | None:
|
||
"""Pick the loop_task with most approved questions when duplicates exist."""
|
||
name = f"循环测试_组{group_id}_批次{batch_id}"
|
||
cur = conn.cursor()
|
||
cur.execute(
|
||
"""SELECT id, name, status, total_approved, env_url, created_at
|
||
FROM loop_task
|
||
WHERE name=? AND env_url LIKE '%dagent%'
|
||
ORDER BY total_approved DESC, created_at DESC
|
||
LIMIT 1""",
|
||
(name,),
|
||
)
|
||
row = cur.fetchone()
|
||
return dict(row) if row else None
|
||
|
||
|
||
def build_export_plan(conn: sqlite3.Connection, plan: dict) -> list[dict]:
|
||
"""Merge task_groups_plan with DB tasks for pending groups missing task_ids."""
|
||
groups_by_id = {g["task_group_id"]: dict(g) for g in plan.get("task_groups") or []}
|
||
for gid in range(1, 15):
|
||
if gid not in groups_by_id:
|
||
groups_by_id[gid] = {"task_group_id": gid, "batch_ids": [], "status": "unknown", "task_ids": []}
|
||
|
||
for gid, group in groups_by_id.items():
|
||
batch_ids = list(group.get("batch_ids") or [])
|
||
plan_tasks = {t["batch_id"]: t for t in (group.get("task_ids") or [])}
|
||
|
||
# Infer batch ids from DB when plan only has pending stub
|
||
if not batch_ids:
|
||
cur = conn.cursor()
|
||
cur.execute(
|
||
"""SELECT DISTINCT CAST(substr(name, instr(name, '批次') + 2) AS INTEGER) AS bid
|
||
FROM loop_task
|
||
WHERE name LIKE ? AND env_url LIKE '%dagent%'
|
||
ORDER BY bid""",
|
||
(f"循环测试_组{gid}_批次%",),
|
||
)
|
||
batch_ids = [r["bid"] for r in cur.fetchall() if r["bid"]]
|
||
|
||
merged_tasks = []
|
||
for bid in sorted(batch_ids):
|
||
if bid in plan_tasks and plan_tasks[bid].get("task_id"):
|
||
merged_tasks.append(plan_tasks[bid])
|
||
continue
|
||
db_task = resolve_task_id_from_db(conn, gid, bid)
|
||
if db_task:
|
||
merged_tasks.append({
|
||
"batch_id": bid,
|
||
"task_id": db_task["id"],
|
||
"task_name": db_task["name"],
|
||
"db_status": db_task["status"],
|
||
"total_approved": db_task["total_approved"],
|
||
})
|
||
group["task_ids"] = merged_tasks
|
||
group["batch_ids"] = batch_ids
|
||
|
||
return [groups_by_id[i] for i in sorted(groups_by_id)]
|
||
|
||
|
||
def main():
|
||
if not DB_PATH.exists():
|
||
print(f"Database not found: {DB_PATH}")
|
||
sys.exit(1)
|
||
|
||
plan = json.loads(PLAN_PATH.read_text(encoding="utf-8")) if PLAN_PATH.exists() else {}
|
||
exported_at = datetime.now().isoformat()
|
||
env = plan.get("environment", "")
|
||
org_id = plan.get("org_id", "")
|
||
|
||
EXPORT_DIR.mkdir(exist_ok=True)
|
||
conn = sqlite3.connect(DB_PATH)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.execute("PRAGMA query_only=ON")
|
||
export_groups = build_export_plan(conn, plan)
|
||
|
||
md_parts = [
|
||
"# 远程 dagent 循环测试 — 全部组别批次问答集汇总\n",
|
||
f"\n> 导出时间: {exported_at}\n> 环境: {env}\n> 组织ID: {org_id}\n> 说明: 已批准问答（qa_gen_question.status=approved）\n\n---\n",
|
||
]
|
||
json_export = {
|
||
"exported_at": exported_at,
|
||
"environment": env,
|
||
"org_id": org_id,
|
||
"source_db": str(DB_PATH),
|
||
"task_groups": [],
|
||
"summary": {"groups": 0, "batches": 0, "batches_with_data": 0, "total_questions": 0},
|
||
}
|
||
|
||
total_batches = batches_with_data = total_questions = 0
|
||
|
||
for group in export_groups:
|
||
gid = group.get("task_group_id")
|
||
gstatus = group.get("status", "unknown")
|
||
group_entry = {
|
||
"task_group_id": gid,
|
||
"status": gstatus,
|
||
"batch_ids": group.get("batch_ids", []),
|
||
"total_chunks": group.get("total_chunks"),
|
||
"total_files": group.get("total_files"),
|
||
"completed_at": group.get("completed_at"),
|
||
"batches": [],
|
||
}
|
||
json_export["task_groups"].append(group_entry)
|
||
json_export["summary"]["groups"] += 1
|
||
|
||
md_parts.append(f"\n# 任务组 {gid}（{gstatus}）\n批次: {group.get('batch_ids', [])}\n")
|
||
|
||
for ti in group.get("task_ids") or []:
|
||
task_id, task_name, batch_id = ti["task_id"], ti.get("task_name"), ti.get("batch_id")
|
||
total_batches += 1
|
||
print(f"组{gid} 批次{batch_id} {task_name}", flush=True)
|
||
rows = get_task_questions_fast(conn, task_id)
|
||
n = len(rows)
|
||
total_questions += n
|
||
group_entry["batches"].append({
|
||
"batch_id": batch_id,
|
||
"task_id": task_id,
|
||
"task_name": task_name,
|
||
"chunk_count": ti.get("chunk_count"),
|
||
"question_count": n,
|
||
"questions": rows_to_json_questions(rows),
|
||
})
|
||
if n:
|
||
batches_with_data += 1
|
||
md_parts.append(f"\n## 批次 {batch_id}: {task_name}（{n} 题）\n\n{rows_to_md(rows)}\n\n---\n")
|
||
else:
|
||
md_parts.append(f"\n## 批次 {batch_id}: {task_name}（无数据）\n\n---\n")
|
||
|
||
conn.close()
|
||
json_export["summary"].update(
|
||
batches=total_batches,
|
||
batches_with_data=batches_with_data,
|
||
total_questions=total_questions,
|
||
)
|
||
|
||
md_path = EXPORT_DIR / "loop_dagent_全部组别批次_问答集汇总.md"
|
||
json_path = EXPORT_DIR / "loop_dagent_全部组别批次_问答集汇总.json"
|
||
md_path.write_text("".join(md_parts), encoding="utf-8")
|
||
json_path.write_text(json.dumps(json_export, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
print("=" * 60)
|
||
print(f"完成: {json_export['summary']['groups']} 组, {batches_with_data}/{total_batches} 批有数据, {total_questions} 题")
|
||
print(md_path)
|
||
print(json_path)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|