init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
56
scripts/analyze_chunks.py
Normal file
56
scripts/analyze_chunks.py
Normal file
@@ -0,0 +1,56 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""分析 exported_chunks_v2 中的切片质量"""
|
||||
import re
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
BASE = r"c:\Users\qq318\Desktop\rag-agent\exported_chunks_v2\public_kb"
|
||||
|
||||
print("=" * 60)
|
||||
print("切片质量分析报告")
|
||||
print("=" * 60)
|
||||
|
||||
for fname in os.listdir(BASE):
|
||||
if not fname.endswith('.md') or fname.startswith('_'):
|
||||
continue
|
||||
|
||||
fpath = os.path.join(BASE, fname)
|
||||
with open(fpath, encoding='utf-8') as f:
|
||||
data = f.read()
|
||||
|
||||
# 提取所有 Length
|
||||
lengths = [int(m) for m in re.findall(r'\*\*Length\*\*: (\d+) chars', data)]
|
||||
|
||||
if not lengths:
|
||||
print(f"\n{fname}: 未找到 Length 字段")
|
||||
continue
|
||||
|
||||
print(f"\n--- {fname} ---")
|
||||
print(f" 总切片数: {len(lengths)}")
|
||||
print(f" 最大: {max(lengths)}, 最小: {min(lengths)}, 平均: {sum(lengths)/len(lengths):.0f}")
|
||||
print(f" 空(0字符): {sum(1 for l in lengths if l == 0)}")
|
||||
print(f" 微短(<10字符): {sum(1 for l in lengths if l < 10)}")
|
||||
print(f" 短(<20字符): {sum(1 for l in lengths if l < 20)}")
|
||||
print(f" 短(<50字符): {sum(1 for l in lengths if l < 50)}")
|
||||
print(f" 超长(>1000字符): {sum(1 for l in lengths if l > 1000)}")
|
||||
print(f" 巨型(>2000字符): {sum(1 for l in lengths if l > 2000)}")
|
||||
print(f" 巨型(>3000字符): {sum(1 for l in lengths if l > 3000)}")
|
||||
|
||||
# 小型表格统计
|
||||
bad_table = data.count('小型表格：表格')
|
||||
empty_table_count = len(re.findall(r'小型表格：\s*```', data))
|
||||
good_table = len(re.findall(r'小型表格：\S', data)) - bad_table
|
||||
|
||||
if bad_table or empty_table_count or good_table:
|
||||
print(f" [表格摘要] 有意义: {good_table}, 无意义(仅'表格'): {bad_table}, 空标题: {empty_table_count}")
|
||||
|
||||
# 列出 top 5 最大切片
|
||||
if max(lengths) > 1000:
|
||||
sorted_l = sorted(enumerate(lengths), key=lambda x: x[1], reverse=True)
|
||||
print(f" Top 5 最大切片:")
|
||||
for idx, size in sorted_l[:5]:
|
||||
print(f" Chunk index {idx}: {size} chars")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
57
scripts/analyze_content_list.py
Normal file
57
scripts/analyze_content_list.py
Normal file
@@ -0,0 +1,57 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
分析 MinerU content_list.json，找出图片数量不一致的原因
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
# 找到 content_list.json
|
||||
base = Path(r".data/mineru_output/a2569e0bfa76/三峡公报_1-15页/auto")
|
||||
cl_path = base / "三峡公报_1-15页_content_list.json"
|
||||
cl_v2_path = base / "三峡公报_1-15页_content_list_v2.json"
|
||||
|
||||
for path, label in [(cl_path, "content_list.json"), (cl_v2_path, "content_list_v2.json")]:
|
||||
if not path.exists():
|
||||
print(f"{label} 不存在")
|
||||
continue
|
||||
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"文件: {label}, 共 {len(data)} 条")
|
||||
|
||||
types = Counter(item.get("type", "text") for item in data)
|
||||
print("\n按类型统计:")
|
||||
for t, c in sorted(types.items()):
|
||||
print(f" {t}: {c}")
|
||||
|
||||
# 所有 image/chart 条目
|
||||
img_items = [item for item in data if item.get("type") in ("image", "chart")]
|
||||
print(f"\n独立图片/图表条目 (image/chart): {len(img_items)}")
|
||||
for i, item in enumerate(img_items):
|
||||
img_path = item.get("img_path", "")
|
||||
caption = str(item.get("caption", ""))[:60]
|
||||
print(f" [{i}] type={item.get('type')}, img_path={img_path}, caption={caption}")
|
||||
|
||||
# 所有 table 条目且带 img_path
|
||||
table_img_items = [item for item in data if item.get("type") == "table" and item.get("img_path")]
|
||||
print(f"\n表格条目带 img_path: {len(table_img_items)}")
|
||||
for item in table_img_items:
|
||||
print(f" img_path={item.get('img_path')}, caption={str(item.get('table_caption',''))[:60]}")
|
||||
|
||||
# 所有带 img_path 的条目（任意类型）
|
||||
all_with_img = [item for item in data if item.get("img_path")]
|
||||
print(f"\n所有带 img_path 的条目: {len(all_with_img)}")
|
||||
for item in all_with_img:
|
||||
print(f" type={item.get('type')}, img_path={item.get('img_path')}")
|
||||
|
||||
# 实际移动到 .data/files/images 的图片
|
||||
images_dir = Path(".data/files/images")
|
||||
actual_images = list(images_dir.glob("*.*")) if images_dir.exists() else []
|
||||
print(f"\n{'='*60}")
|
||||
print(f".data/files/images 实际文件数: {len(actual_images)}")
|
||||
for f in actual_images:
|
||||
print(f" {f.name} ({f.stat().st_size} bytes)")
|
||||
71
scripts/check_tables.py
Normal file
71
scripts/check_tables.py
Normal file
@@ -0,0 +1,71 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""检查 MinerU 输出中表格的实际内容"""
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
BASE = r"c:\Users\qq318\Desktop\rag-agent\.data\mineru_output"
|
||||
|
||||
# hash -> 文件名映射
|
||||
hash_map = {}
|
||||
for h in os.listdir(BASE):
|
||||
subdir = os.path.join(BASE, h)
|
||||
for name in os.listdir(subdir):
|
||||
hash_map[h] = name
|
||||
|
||||
print("=" * 60)
|
||||
print("MinerU 表格内容分析")
|
||||
print("=" * 60)
|
||||
|
||||
for file_hash, doc_name in hash_map.items():
|
||||
# 查找 content_list
|
||||
for subpath in ["office", "auto"]:
|
||||
cl_path = os.path.join(BASE, file_hash, doc_name, subpath, f"{doc_name}_content_list.json")
|
||||
if os.path.exists(cl_path):
|
||||
break
|
||||
else:
|
||||
print(f"\n{doc_name}: content_list not found")
|
||||
continue
|
||||
|
||||
with open(cl_path, 'r', encoding='utf-8') as f:
|
||||
content_list = json.load(f)
|
||||
|
||||
tables = [(i, item) for i, item in enumerate(content_list) if item.get('type') == 'table']
|
||||
|
||||
if not tables:
|
||||
continue
|
||||
|
||||
print(f"\n--- {doc_name} ({len(tables)} tables) ---")
|
||||
|
||||
no_caption = 0
|
||||
no_body = 0
|
||||
|
||||
for idx, item in tables:
|
||||
caption = item.get('table_caption', '')
|
||||
if isinstance(caption, list):
|
||||
caption = ' '.join(str(c) for c in caption)
|
||||
body = item.get('table_body', '')
|
||||
if isinstance(body, list):
|
||||
body = ' '.join(str(b) for b in body)
|
||||
page = item.get('page_idx', '?')
|
||||
|
||||
has_caption = bool(caption and str(caption).strip() and str(caption).strip() != '表格')
|
||||
has_body = bool(body and str(body).strip())
|
||||
|
||||
if not has_caption:
|
||||
no_caption += 1
|
||||
if not has_body:
|
||||
no_body += 1
|
||||
|
||||
# 只打印前5个无caption的
|
||||
if not has_caption and no_caption <= 5:
|
||||
body_preview = body[:120].replace('\n', '\\n') if body else '(empty)'
|
||||
print(f" [idx={idx}, page={page}] caption={repr(caption)[:40]}")
|
||||
print(f" body: {body_preview}")
|
||||
|
||||
print(f" 总表格数: {len(tables)}")
|
||||
print(f" 无caption: {no_caption}")
|
||||
print(f" 无body: {no_body}")
|
||||
print(f" 有caption有body: {len(tables) - max(no_caption, no_body)}")
|
||||
245
scripts/compare_embedding_models.py
Normal file
245
scripts/compare_embedding_models.py
Normal file
@@ -0,0 +1,245 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
向量模型对比评估脚本
|
||||
|
||||
比较不同嵌入模型的检索效果：
|
||||
- bge-base-zh-v1.5 (当前使用)
|
||||
- bge-large-zh-v1.5 (可选升级)
|
||||
- bge-m3 (多语言支持)
|
||||
|
||||
用法:
|
||||
python scripts/compare_embedding_models.py --models bge-base-zh-v1.5 bge-large-zh-v1.5
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
PROJECT_ROOT = Path(__file__).parent.parent
|
||||
sys.path.insert(0, str(PROJECT_ROOT))
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# 模型配置
|
||||
MODEL_CONFIGS = {
|
||||
"bge-base-zh-v1.5": {
|
||||
"path": "BAAI/bge-base-zh-v1.5",
|
||||
"dimension": 768,
|
||||
"description": "中文基础模型，768维",
|
||||
"size_mb": 390
|
||||
},
|
||||
"bge-large-zh-v1.5": {
|
||||
"path": "BAAI/bge-large-zh-v1.5",
|
||||
"dimension": 1024,
|
||||
"description": "中文大模型，1024维，精度更高",
|
||||
"size_mb": 1300
|
||||
},
|
||||
"bge-m3": {
|
||||
"path": "BAAI/bge-m3",
|
||||
"dimension": 1024,
|
||||
"description": "多语言模型，支持100+语言",
|
||||
"size_mb": 2200
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def evaluate_model(model_name: str, eval_dataset_path: str, sample_size: int = None) -> dict:
|
||||
"""
|
||||
评估单个模型
|
||||
|
||||
Args:
|
||||
model_name: 模型名称
|
||||
eval_dataset_path: 评测数据集路径
|
||||
sample_size: 采样数量
|
||||
|
||||
Returns:
|
||||
评测结果
|
||||
"""
|
||||
from sentence_transformers import SentenceTransformer
|
||||
import numpy as np
|
||||
|
||||
config = MODEL_CONFIGS.get(model_name)
|
||||
if not config:
|
||||
raise ValueError(f"未知模型: {model_name}")
|
||||
|
||||
logger.info(f"加载模型: {model_name} ({config['description']})")
|
||||
|
||||
# 加载模型
|
||||
start_time = time.time()
|
||||
model = SentenceTransformer(config["path"])
|
||||
load_time = time.time() - start_time
|
||||
logger.info(f"模型加载耗时: {load_time:.2f}秒")
|
||||
|
||||
# 加载评测数据集
|
||||
with open(eval_dataset_path, 'r', encoding='utf-8') as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
queries = dataset.get('queries', [])
|
||||
if sample_size and sample_size < len(queries):
|
||||
import random
|
||||
queries = random.sample(queries, sample_size)
|
||||
|
||||
# 简单评测：计算查询与相关文档的相似度
|
||||
results = {
|
||||
"model": model_name,
|
||||
"dimension": config["dimension"],
|
||||
"load_time": load_time,
|
||||
"queries_evaluated": len(queries),
|
||||
"avg_similarity": 0.0,
|
||||
"embedding_times": []
|
||||
}
|
||||
|
||||
similarities = []
|
||||
for q in queries:
|
||||
query_text = q['query']
|
||||
reference_answer = q.get('reference_answer', '')
|
||||
|
||||
# 计算查询和参考答案的相似度
|
||||
start = time.time()
|
||||
embeddings = model.encode([query_text, reference_answer])
|
||||
embed_time = time.time() - start
|
||||
results["embedding_times"].append(embed_time)
|
||||
|
||||
similarity = np.dot(embeddings[0], embeddings[1]) / (
|
||||
np.linalg.norm(embeddings[0]) * np.linalg.norm(embeddings[1])
|
||||
)
|
||||
similarities.append(similarity)
|
||||
|
||||
results["avg_similarity"] = float(np.mean(similarities))
|
||||
results["avg_embed_time"] = float(np.mean(results["embedding_times"]))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def compare_models(models: list, eval_dataset_path: str, sample_size: int = None) -> dict:
|
||||
"""
|
||||
对比多个模型
|
||||
|
||||
Args:
|
||||
models: 模型列表
|
||||
eval_dataset_path: 评测数据集路径
|
||||
sample_size: 采样数量
|
||||
|
||||
Returns:
|
||||
对比结果
|
||||
"""
|
||||
results = {}
|
||||
|
||||
for model_name in models:
|
||||
try:
|
||||
result = evaluate_model(model_name, eval_dataset_path, sample_size)
|
||||
results[model_name] = result
|
||||
|
||||
logger.info(f"\n{'='*50}")
|
||||
logger.info(f"模型: {model_name}")
|
||||
logger.info(f"维度: {result['dimension']}")
|
||||
logger.info(f"加载时间: {result['load_time']:.2f}秒")
|
||||
logger.info(f"平均相似度: {result['avg_similarity']:.4f}")
|
||||
logger.info(f"平均编码时间: {result['avg_embed_time']*1000:.2f}ms")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"评估模型 {model_name} 失败: {e}")
|
||||
results[model_name] = {"error": str(e)}
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def print_comparison(results: dict):
|
||||
"""打印对比结果"""
|
||||
print("\n" + "=" * 70)
|
||||
print(" 向量模型对比结果")
|
||||
print("=" * 70)
|
||||
|
||||
# 表头
|
||||
print(f"\n{'模型':<25} {'维度':<8} {'加载时间':<12} {'平均相似度':<12} {'编码时间':<12}")
|
||||
print("-" * 70)
|
||||
|
||||
for model_name, result in results.items():
|
||||
if "error" in result:
|
||||
print(f"{model_name:<25} 错误: {result['error']}")
|
||||
else:
|
||||
print(f"{model_name:<25} {result['dimension']:<8} {result['load_time']:.2f}秒{'':<4} "
|
||||
f"{result['avg_similarity']:.4f}{'':<4} {result['avg_embed_time']*1000:.2f}ms")
|
||||
|
||||
# 推荐
|
||||
print("\n" + "-" * 70)
|
||||
print("推荐建议:")
|
||||
print("-" * 70)
|
||||
|
||||
# 找出最佳模型
|
||||
valid_results = {k: v for k, v in results.items() if "error" not in v}
|
||||
if valid_results:
|
||||
best_sim = max(valid_results.items(), key=lambda x: x[1].get('avg_similarity', 0))
|
||||
fastest = min(valid_results.items(), key=lambda x: x[1].get('avg_embed_time', float('inf')))
|
||||
|
||||
print(f"• 最高相似度: {best_sim[0]} ({best_sim[1]['avg_similarity']:.4f})")
|
||||
print(f"• 最快编码: {fastest[0]} ({fastest[1]['avg_embed_time']*1000:.2f}ms)")
|
||||
|
||||
# 给出建议
|
||||
current = "bge-base-zh-v1.5"
|
||||
if current in valid_results:
|
||||
current_sim = valid_results[current]['avg_similarity']
|
||||
improvement = best_sim[1]['avg_similarity'] - current_sim
|
||||
if improvement > 0.05:
|
||||
print(f"\n💡 建议升级到 {best_sim[0]}，预期相似度提升: {improvement:.4f}")
|
||||
else:
|
||||
print(f"\n✓ 当前模型 {current} 表现良好，升级收益有限")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='向量模型对比评估')
|
||||
parser.add_argument(
|
||||
'--models',
|
||||
nargs='+',
|
||||
default=['bge-base-zh-v1.5'],
|
||||
choices=list(MODEL_CONFIGS.keys()),
|
||||
help='要评估的模型列表'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--eval_dataset',
|
||||
type=str,
|
||||
default='data/eval_dataset.json',
|
||||
help='评测数据集路径'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--sample',
|
||||
type=int,
|
||||
default=10,
|
||||
help='采样数量（快速测试）'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--output',
|
||||
type=str,
|
||||
default=None,
|
||||
help='结果输出路径'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
eval_path = PROJECT_ROOT / args.eval_dataset
|
||||
if not eval_path.exists():
|
||||
logger.error(f"评测数据集不存在: {eval_path}")
|
||||
sys.exit(1)
|
||||
|
||||
# 运行对比
|
||||
results = compare_models(args.models, str(eval_path), args.sample)
|
||||
|
||||
# 打印结果
|
||||
print_comparison(results)
|
||||
|
||||
# 保存结果
|
||||
if args.output:
|
||||
output_path = PROJECT_ROOT / args.output
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||||
logger.info(f"结果已保存到: {output_path}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
670
scripts/eval_e2e.py
Normal file
670
scripts/eval_e2e.py
Normal file
@@ -0,0 +1,670 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
RAG 端到端评测脚本 (eval_e2e.py)
|
||||
|
||||
通过 HTTP API 调用 /rag 接口，解析 SSE 流式响应，
|
||||
对每个问题从关键词覆盖、LLM 质量评分两个维度进行评测，
|
||||
支持基线录制与阶段间对比。
|
||||
|
||||
用法:
|
||||
# 录制基线
|
||||
python scripts/eval_e2e.py --baseline
|
||||
|
||||
# 阶段评测并与基线对比
|
||||
python scripts/eval_e2e.py --phase 1
|
||||
|
||||
# 仅评测不对比
|
||||
python scripts/eval_e2e.py --phase test
|
||||
|
||||
# 指定数据集和输出路径
|
||||
python scripts/eval_e2e.py --dataset data/eval/eval_dataset.json --phase 0 --output data/eval_results/phase0.json
|
||||
|
||||
# 跳过 LLM 评分（快速模式）
|
||||
python scripts/eval_e2e.py --phase test --no_llm
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
# 添加项目根目录到路径
|
||||
PROJECT_ROOT = Path(__file__).parent.parent
|
||||
sys.path.insert(0, str(PROJECT_ROOT))
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s [%(levelname)s] %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def setup_file_logging(log_path: str):
|
||||
"""将日志同时输出到文件，避免 Windows 控制台编码问题"""
|
||||
fh = logging.FileHandler(log_path, encoding='utf-8')
|
||||
fh.setLevel(logging.INFO)
|
||||
fh.setFormatter(logging.Formatter('%(asctime)s [%(levelname)s] %(message)s'))
|
||||
logger.addHandler(fh)
|
||||
|
||||
# ───────────── 配置 ─────────────
|
||||
RAG_API_URL = "http://127.0.0.1:5001/rag"
|
||||
DEFAULT_DATASET = "data/eval/eval_dataset.json"
|
||||
RESULTS_DIR = "data/eval_results"
|
||||
REQUEST_TIMEOUT = 120 # 秒
|
||||
REQUEST_INTERVAL = 1.5 # 请求间隔（秒），避免过载
|
||||
|
||||
|
||||
# ───────────── SSE 解析 ─────────────
|
||||
def call_rag_sse(question: str, collections: list = None, api_url: str = None) -> dict:
|
||||
"""
|
||||
调用 /rag 接口，解析 SSE 流式响应，提取 finish 事件中的完整回答。
|
||||
参考 docs/curl测试手册.md 中的调用规范。
|
||||
|
||||
Returns:
|
||||
{
|
||||
"answer": str,
|
||||
"sources": list,
|
||||
"citations": list,
|
||||
"duration_ms": int,
|
||||
"error": str | None
|
||||
}
|
||||
"""
|
||||
import requests
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "text/event-stream"
|
||||
}
|
||||
# chat_history 在生产模式下是必填字段
|
||||
payload = {"message": question, "chat_history": []}
|
||||
if collections:
|
||||
payload["collections"] = collections
|
||||
|
||||
result = {
|
||||
"answer": "",
|
||||
"sources": [],
|
||||
"citations": [],
|
||||
"duration_ms": 0,
|
||||
"error": None
|
||||
}
|
||||
|
||||
try:
|
||||
resp = requests.post(
|
||||
api_url or RAG_API_URL, json=payload, headers=headers,
|
||||
stream=True, timeout=REQUEST_TIMEOUT
|
||||
)
|
||||
resp.raise_for_status()
|
||||
|
||||
# 逐行解析 SSE 事件
|
||||
for raw_line in resp.iter_lines(decode_unicode=True):
|
||||
if not raw_line or not raw_line.startswith("data:"):
|
||||
continue
|
||||
data_str = raw_line[len("data:"):].strip()
|
||||
if not data_str:
|
||||
continue
|
||||
try:
|
||||
event = json.loads(data_str)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
evt_type = event.get("type")
|
||||
if evt_type == "finish":
|
||||
result["answer"] = event.get("answer", "")
|
||||
result["sources"] = event.get("sources", [])
|
||||
result["citations"] = event.get("citations", [])
|
||||
result["duration_ms"] = event.get("duration_ms", 0)
|
||||
break
|
||||
elif evt_type == "error":
|
||||
result["error"] = event.get("message", "未知错误")
|
||||
break
|
||||
|
||||
except requests.exceptions.ConnectionError:
|
||||
result["error"] = "无法连接到 RAG 服务，请确认服务已启动 (localhost:5001)"
|
||||
except requests.exceptions.Timeout:
|
||||
result["error"] = f"请求超时 ({REQUEST_TIMEOUT}s)"
|
||||
except Exception as e:
|
||||
result["error"] = str(e)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ───────────── 评分模块 ─────────────
|
||||
def keyword_coverage(answer: str, expected_keywords: list) -> dict:
|
||||
"""
|
||||
关键词覆盖率评分。
|
||||
检查 answer 中包含 expected_keywords 的百分比。
|
||||
|
||||
Returns:
|
||||
{"score": 0~1, "matched": [...], "missing": [...]}
|
||||
"""
|
||||
if not expected_keywords or not answer:
|
||||
return {"score": 0.0, "matched": [], "missing": expected_keywords or []}
|
||||
|
||||
matched = [kw for kw in expected_keywords if kw in answer]
|
||||
missing = [kw for kw in expected_keywords if kw not in answer]
|
||||
score = len(matched) / len(expected_keywords)
|
||||
return {"score": round(score, 4), "matched": matched, "missing": missing}
|
||||
|
||||
|
||||
def llm_quality_score(query: str, answer: str, reference: str) -> dict:
|
||||
"""
|
||||
使用 LLM 对回答质量进行多维度评分 (0-10)。
|
||||
|
||||
Returns:
|
||||
{"accuracy": X, "completeness": X, "relevance": X, "fluency": X, "overall": X}
|
||||
"""
|
||||
try:
|
||||
from openai import OpenAI
|
||||
from config import DASHSCOPE_API_KEY, DASHSCOPE_BASE_URL, DASHSCOPE_MODEL
|
||||
except ImportError:
|
||||
logger.warning("无法导入 openai 或 config 模块，跳过 LLM 评分")
|
||||
return {"overall": 0.0, "error": "import_failed"}
|
||||
|
||||
client = OpenAI(api_key=DASHSCOPE_API_KEY, base_url=DASHSCOPE_BASE_URL)
|
||||
|
||||
prompt = f"""你是一个专业的 RAG 系统评测员。请对以下回答进行质量评分。
|
||||
|
||||
【用户问题】
|
||||
{query}
|
||||
|
||||
【参考答案（人工编写的高质量答案）】
|
||||
{reference}
|
||||
|
||||
【系统实际回答】
|
||||
{answer}
|
||||
|
||||
【评分维度】（每项 0-10 分）
|
||||
1. 准确性(accuracy)：系统回答中的事实信息是否与参考答案一致，有无错误信息
|
||||
2. 完整性(completeness)：是否覆盖了参考答案中的关键要点
|
||||
3. 相关性(relevance)：回答是否直接针对问题，没有跑题或冗余
|
||||
4. 流畅性(fluency)：回答是否通顺、结构清晰、易于理解
|
||||
|
||||
请严格按以下 JSON 格式返回，不要包含其他内容：
|
||||
{{"accuracy": <分数>, "completeness": <分数>, "relevance": <分数>, "fluency": <分数>, "overall": <总分>}}"""
|
||||
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model=DASHSCOPE_MODEL,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=0.1,
|
||||
max_tokens=300
|
||||
)
|
||||
text = response.choices[0].message.content.strip()
|
||||
|
||||
# 提取 JSON
|
||||
json_match = re.search(r'\{[^}]+\}', text)
|
||||
if json_match:
|
||||
scores = json.loads(json_match.group())
|
||||
# 归一化到 0-1
|
||||
for k in scores:
|
||||
scores[k] = round(min(10, max(0, scores[k])) / 10.0, 4)
|
||||
return scores
|
||||
else:
|
||||
logger.warning(f"LLM 返回内容无法解析 JSON: {text[:100]}")
|
||||
return {"overall": 0.0, "error": "parse_failed"}
|
||||
except Exception as e:
|
||||
logger.warning(f"LLM 评分异常: {e}")
|
||||
return {"overall": 0.0, "error": str(e)}
|
||||
|
||||
|
||||
# ───────────── 评测主流程 ─────────────
|
||||
def evaluate_dataset(dataset_path: str, use_llm: bool = True, api_url: str = None) -> dict:
|
||||
"""
|
||||
对数据集中的所有问题进行评测。
|
||||
|
||||
Returns:
|
||||
{
|
||||
"meta": {...},
|
||||
"questions": [ {id, query, answer, scores, ...}, ... ],
|
||||
"summary": { avg_keyword, avg_llm, by_type, by_difficulty }
|
||||
}
|
||||
"""
|
||||
with open(dataset_path, 'r', encoding='utf-8') as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
questions = dataset.get("questions", [])
|
||||
total = len(questions)
|
||||
if total == 0:
|
||||
logger.error("数据集中没有问题")
|
||||
return {}
|
||||
|
||||
logger.info(f"开始评测，共 {total} 个问题...")
|
||||
|
||||
results = []
|
||||
kw_scores = []
|
||||
llm_scores = []
|
||||
by_type = {}
|
||||
by_difficulty = {}
|
||||
|
||||
for i, q in enumerate(questions, 1):
|
||||
qid = q["id"]
|
||||
query = q["query"]
|
||||
qtype = q.get("query_type", "unknown")
|
||||
difficulty = q.get("difficulty", "medium")
|
||||
expected_keywords = q.get("expected_keywords", [])
|
||||
reference_answer = q.get("reference_answer", "")
|
||||
|
||||
logger.info(f"[{i}/{total}] {qid}: {query[:40]}...")
|
||||
|
||||
# 调用 RAG
|
||||
t0 = time.time()
|
||||
rag_result = call_rag_sse(query, api_url=api_url)
|
||||
elapsed = round(time.time() - t0, 2)
|
||||
answer = rag_result.get("answer", "")
|
||||
error = rag_result.get("error")
|
||||
|
||||
if error:
|
||||
logger.warning(f" !! API 错误: {error}")
|
||||
results.append({
|
||||
"id": qid, "query": query, "query_type": qtype,
|
||||
"difficulty": difficulty, "answer": "", "error": error,
|
||||
"elapsed_s": elapsed,
|
||||
"keyword_score": 0.0, "llm_scores": {"overall": 0.0}
|
||||
})
|
||||
continue
|
||||
|
||||
# 关键词覆盖评分
|
||||
kw_result = keyword_coverage(answer, expected_keywords)
|
||||
kw_score = kw_result["score"]
|
||||
kw_scores.append(kw_score)
|
||||
|
||||
# LLM 质量评分
|
||||
llm_result = {}
|
||||
if use_llm and reference_answer:
|
||||
llm_result = llm_quality_score(query, answer, reference_answer)
|
||||
llm_scores.append(llm_result.get("overall", 0.0))
|
||||
|
||||
entry = {
|
||||
"id": qid,
|
||||
"query": query,
|
||||
"query_type": qtype,
|
||||
"difficulty": difficulty,
|
||||
"answer": answer,
|
||||
"sources": rag_result.get("sources", []),
|
||||
"citations": rag_result.get("citations", []),
|
||||
"duration_ms": rag_result.get("duration_ms", 0),
|
||||
"elapsed_s": elapsed,
|
||||
"keyword_score": kw_score,
|
||||
"keyword_matched": kw_result["matched"],
|
||||
"keyword_missing": kw_result["missing"],
|
||||
"llm_scores": llm_result,
|
||||
"reference_answer": reference_answer[:300]
|
||||
}
|
||||
results.append(entry)
|
||||
|
||||
# 按类型/难度分组
|
||||
for group_dict, key in [(by_type, qtype), (by_difficulty, difficulty)]:
|
||||
if key not in group_dict:
|
||||
group_dict[key] = {"kw": [], "llm": []}
|
||||
group_dict[key]["kw"].append(kw_score)
|
||||
if "overall" in llm_result:
|
||||
group_dict[key]["llm"].append(llm_result["overall"])
|
||||
|
||||
# 打印进度
|
||||
llm_str = f", LLM={llm_result.get('overall', 'N/A')}" if use_llm else ""
|
||||
logger.info(f" 关键词={kw_score:.0%}{llm_str}, 耗时={elapsed}s")
|
||||
|
||||
# 请求间隔
|
||||
if i < total:
|
||||
time.sleep(REQUEST_INTERVAL)
|
||||
|
||||
# 汇总
|
||||
avg_kw = sum(kw_scores) / len(kw_scores) if kw_scores else 0.0
|
||||
avg_llm = sum(llm_scores) / len(llm_scores) if llm_scores else 0.0
|
||||
|
||||
summary = {
|
||||
"avg_keyword_coverage": round(avg_kw, 4),
|
||||
"avg_llm_overall": round(avg_llm, 4),
|
||||
"total_questions": total,
|
||||
"successful": len(kw_scores),
|
||||
"by_type": {
|
||||
k: {
|
||||
"avg_keyword": round(sum(v["kw"]) / len(v["kw"]), 4) if v["kw"] else 0,
|
||||
"avg_llm": round(sum(v["llm"]) / len(v["llm"]), 4) if v["llm"] else 0,
|
||||
"count": len(v["kw"])
|
||||
}
|
||||
for k, v in by_type.items()
|
||||
},
|
||||
"by_difficulty": {
|
||||
k: {
|
||||
"avg_keyword": round(sum(v["kw"]) / len(v["kw"]), 4) if v["kw"] else 0,
|
||||
"avg_llm": round(sum(v["llm"]) / len(v["llm"]), 4) if v["llm"] else 0,
|
||||
"count": len(v["kw"])
|
||||
}
|
||||
for k, v in by_difficulty.items()
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
"meta": {
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"dataset": dataset_path,
|
||||
"use_llm": use_llm
|
||||
},
|
||||
"questions": results,
|
||||
"summary": summary
|
||||
}
|
||||
|
||||
|
||||
# ───────────── 对比模块 ─────────────
|
||||
def compare_results(baseline_path: str, current_path: str) -> dict:
|
||||
"""
|
||||
对比基线和当前阶段的评测结果。
|
||||
|
||||
Returns:
|
||||
{
|
||||
"per_question": [{id, query, baseline_kw, current_kw, delta_kw, ...}],
|
||||
"summary_delta": { keyword_delta, llm_delta },
|
||||
"regressions": [{id, query, delta_kw, delta_llm}]
|
||||
}
|
||||
"""
|
||||
with open(baseline_path, 'r', encoding='utf-8') as f:
|
||||
baseline = json.load(f)
|
||||
with open(current_path, 'r', encoding='utf-8') as f:
|
||||
current = json.load(f)
|
||||
|
||||
# 按 ID 索引
|
||||
b_map = {q["id"]: q for q in baseline.get("questions", [])}
|
||||
c_map = {q["id"]: q for q in current.get("questions", [])}
|
||||
|
||||
per_question = []
|
||||
regressions = []
|
||||
|
||||
for qid in sorted(b_map.keys()):
|
||||
bq = b_map[qid]
|
||||
cq = c_map.get(qid)
|
||||
if not cq:
|
||||
continue
|
||||
|
||||
b_kw = bq.get("keyword_score", 0)
|
||||
c_kw = cq.get("keyword_score", 0)
|
||||
b_llm = bq.get("llm_scores", {}).get("overall", 0)
|
||||
c_llm = cq.get("llm_scores", {}).get("overall", 0)
|
||||
|
||||
delta_kw = round(c_kw - b_kw, 4)
|
||||
delta_llm = round(c_llm - b_llm, 4)
|
||||
|
||||
entry = {
|
||||
"id": qid,
|
||||
"query": bq.get("query", ""),
|
||||
"query_type": bq.get("query_type", ""),
|
||||
"difficulty": bq.get("difficulty", ""),
|
||||
"baseline_kw": b_kw,
|
||||
"current_kw": c_kw,
|
||||
"delta_kw": delta_kw,
|
||||
"baseline_llm": b_llm,
|
||||
"current_llm": c_llm,
|
||||
"delta_llm": delta_llm
|
||||
}
|
||||
per_question.append(entry)
|
||||
|
||||
# 回归判定：关键词覆盖或 LLM 评分下降超过 10%
|
||||
if delta_kw < -0.1 or delta_llm < -0.1:
|
||||
regressions.append(entry)
|
||||
|
||||
bs = baseline.get("summary", {})
|
||||
cs = current.get("summary", {})
|
||||
summary_delta = {
|
||||
"keyword_delta": round(
|
||||
cs.get("avg_keyword_coverage", 0) - bs.get("avg_keyword_coverage", 0), 4
|
||||
),
|
||||
"llm_delta": round(
|
||||
cs.get("avg_llm_overall", 0) - bs.get("avg_llm_overall", 0), 4
|
||||
),
|
||||
"baseline_avg_kw": bs.get("avg_keyword_coverage", 0),
|
||||
"current_avg_kw": cs.get("avg_keyword_coverage", 0),
|
||||
"baseline_avg_llm": bs.get("avg_llm_overall", 0),
|
||||
"current_avg_llm": cs.get("avg_llm_overall", 0)
|
||||
}
|
||||
|
||||
return {
|
||||
"per_question": per_question,
|
||||
"summary_delta": summary_delta,
|
||||
"regressions": regressions
|
||||
}
|
||||
|
||||
|
||||
# ───────────── 报告生成 ─────────────
|
||||
def generate_report(eval_result: dict, comparison: dict = None, phase_name: str = "") -> str:
|
||||
"""生成 Markdown 格式的评测报告"""
|
||||
lines = []
|
||||
meta = eval_result.get("meta", {})
|
||||
summary = eval_result.get("summary", {})
|
||||
questions = eval_result.get("questions", [])
|
||||
|
||||
lines.append(f"## RAG 端到端评测报告 — {phase_name or '未命名'}")
|
||||
lines.append("")
|
||||
lines.append(f"**评测时间**: {meta.get('timestamp', 'N/A')}")
|
||||
lines.append(f"**数据集**: {meta.get('dataset', 'N/A')}")
|
||||
lines.append(f"**LLM 评分**: {'启用' if meta.get('use_llm') else '跳过'}")
|
||||
lines.append("")
|
||||
|
||||
# 汇总
|
||||
lines.append("### 汇总指标")
|
||||
lines.append("")
|
||||
lines.append(f"| 指标 | 值 |")
|
||||
lines.append(f"|------|-----|")
|
||||
lines.append(f"| 平均关键词覆盖率 | {summary.get('avg_keyword_coverage', 0):.2%} |")
|
||||
lines.append(f"| 平均 LLM 总分 | {summary.get('avg_llm_overall', 0):.2f} |")
|
||||
lines.append(f"| 总问题数 | {summary.get('total_questions', 0)} |")
|
||||
lines.append(f"| 成功评测数 | {summary.get('successful', 0)} |")
|
||||
lines.append("")
|
||||
|
||||
# 按类型
|
||||
by_type = summary.get("by_type", {})
|
||||
if by_type:
|
||||
lines.append("### 按查询类型")
|
||||
lines.append("")
|
||||
lines.append("| 类型 | 数量 | 平均关键词 | 平均 LLM |")
|
||||
lines.append("|------|------|-----------|---------|")
|
||||
for t, v in sorted(by_type.items()):
|
||||
lines.append(f"| {t} | {v['count']} | {v['avg_keyword']:.2%} | {v['avg_llm']:.2f} |")
|
||||
lines.append("")
|
||||
|
||||
# 按难度
|
||||
by_diff = summary.get("by_difficulty", {})
|
||||
if by_diff:
|
||||
lines.append("### 按难度")
|
||||
lines.append("")
|
||||
lines.append("| 难度 | 数量 | 平均关键词 | 平均 LLM |")
|
||||
lines.append("|------|------|-----------|---------|")
|
||||
for d, v in sorted(by_diff.items()):
|
||||
lines.append(f"| {d} | {v['count']} | {v['avg_keyword']:.2%} | {v['avg_llm']:.2f} |")
|
||||
lines.append("")
|
||||
|
||||
# 对比结果
|
||||
if comparison:
|
||||
sd = comparison.get("summary_delta", {})
|
||||
lines.append("### 与基线对比")
|
||||
lines.append("")
|
||||
lines.append(f"| 指标 | 基线 | 当前 | 变化 |")
|
||||
lines.append(f"|------|------|------|------|")
|
||||
lines.append(
|
||||
f"| 关键词覆盖率 | {sd.get('baseline_avg_kw', 0):.2%} "
|
||||
f"| {sd.get('current_avg_kw', 0):.2%} "
|
||||
f"| {sd.get('keyword_delta', 0):+.2%} |"
|
||||
)
|
||||
lines.append(
|
||||
f"| LLM 总分 | {sd.get('baseline_avg_llm', 0):.2f} "
|
||||
f"| {sd.get('current_avg_llm', 0):.2f} "
|
||||
f"| {sd.get('llm_delta', 0):+.2f} |"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
regressions = comparison.get("regressions", [])
|
||||
if regressions:
|
||||
lines.append(f"### 回归问题（{len(regressions)} 个）")
|
||||
lines.append("")
|
||||
for r in regressions:
|
||||
lines.append(
|
||||
f"- **{r['id']}** {r['query'][:50]}... "
|
||||
f"(关键词 {r['delta_kw']:+.0%}, LLM {r['delta_llm']:+.2f})"
|
||||
)
|
||||
lines.append("")
|
||||
else:
|
||||
lines.append("### 无回归问题")
|
||||
lines.append("")
|
||||
|
||||
# 逐题详情
|
||||
lines.append("### 逐题结果")
|
||||
lines.append("")
|
||||
for q in questions:
|
||||
status = "ERROR" if q.get("error") else "OK"
|
||||
llm_overall = q.get("llm_scores", {}).get("overall", "N/A")
|
||||
if isinstance(llm_overall, float):
|
||||
llm_overall = f"{llm_overall:.2f}"
|
||||
lines.append(f"**{q['id']}** [{status}] {q['query']}")
|
||||
lines.append(f"- 关键词覆盖: {q.get('keyword_score', 0):.0%}")
|
||||
lines.append(f"- LLM 总分: {llm_overall}")
|
||||
if q.get("keyword_missing"):
|
||||
lines.append(f"- 缺失关键词: {', '.join(q['keyword_missing'])}")
|
||||
if q.get("error"):
|
||||
lines.append(f"- 错误: {q['error']}")
|
||||
# 截取回答前 150 字
|
||||
ans = q.get("answer", "")
|
||||
if ans:
|
||||
lines.append(f"- 回答摘要: {ans[:150]}...")
|
||||
lines.append("")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# ───────────── 主入口 ─────────────
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="RAG 端到端评测脚本")
|
||||
parser.add_argument(
|
||||
"--dataset", type=str, default=DEFAULT_DATASET,
|
||||
help=f"评测数据集路径 (默认: {DEFAULT_DATASET})"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--phase", type=str, default=None,
|
||||
help="当前阶段名称，如 baseline / 0 / 1 / 2 ..."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--baseline", action="store_true",
|
||||
help="录制基线（等价于 --phase baseline）"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", type=str, default=None,
|
||||
help="结果 JSON 输出路径（默认自动生成）"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no_llm", action="store_true",
|
||||
help="跳过 LLM 评分（快速模式）"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api_url", type=str, default=None,
|
||||
help=f"RAG API 地址 (默认: {RAG_API_URL})"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Windows 编码
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8', errors='replace')
|
||||
sys.stderr.reconfigure(encoding='utf-8', errors='replace')
|
||||
|
||||
# 确定阶段名
|
||||
if args.baseline:
|
||||
phase_name = "baseline"
|
||||
elif args.phase:
|
||||
phase_name = args.phase
|
||||
else:
|
||||
phase_name = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
# API 地址覆盖
|
||||
api_url = args.api_url or RAG_API_URL
|
||||
|
||||
# 路径处理
|
||||
dataset_path = PROJECT_ROOT / args.dataset
|
||||
if not dataset_path.exists():
|
||||
logger.error(f"dataset not found: {dataset_path}")
|
||||
sys.exit(1)
|
||||
|
||||
results_dir = PROJECT_ROOT / RESULTS_DIR
|
||||
results_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 设置文件日志
|
||||
log_path = results_dir / f"eval_{phase_name}.log"
|
||||
setup_file_logging(str(log_path))
|
||||
|
||||
# 输出路径
|
||||
if args.output:
|
||||
output_path = PROJECT_ROOT / args.output
|
||||
else:
|
||||
output_path = results_dir / f"phase_{phase_name}.json"
|
||||
|
||||
report_path = output_path.with_suffix(".md")
|
||||
|
||||
# 执行评测
|
||||
logger.info(f"=== RAG E2E Eval [{phase_name}] ===")
|
||||
logger.info(f"dataset: {dataset_path}")
|
||||
logger.info(f"API: {api_url}")
|
||||
logger.info(f"LLM scoring: {'on' if not args.no_llm else 'off'}")
|
||||
|
||||
eval_result = evaluate_dataset(
|
||||
str(dataset_path),
|
||||
use_llm=not args.no_llm,
|
||||
api_url=api_url
|
||||
)
|
||||
|
||||
if not eval_result:
|
||||
logger.error("eval failed, no results")
|
||||
sys.exit(1)
|
||||
|
||||
# 保存 JSON
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(eval_result, f, ensure_ascii=False, indent=2)
|
||||
logger.info(f"results saved: {output_path}")
|
||||
|
||||
# 基线对比
|
||||
comparison = None
|
||||
baseline_path = results_dir / "phase_baseline.json"
|
||||
if phase_name != "baseline" and baseline_path.exists():
|
||||
logger.info("comparing with baseline...")
|
||||
comparison = compare_results(str(baseline_path), str(output_path))
|
||||
regressions = comparison.get("regressions", [])
|
||||
if regressions:
|
||||
logger.warning(f"found {len(regressions)} regression(s)!")
|
||||
else:
|
||||
logger.info("no regressions, optimization is safe.")
|
||||
|
||||
# 生成报告
|
||||
report = generate_report(eval_result, comparison, phase_name)
|
||||
with open(report_path, 'w', encoding='utf-8') as f:
|
||||
f.write(report)
|
||||
logger.info(f"report saved: {report_path}")
|
||||
|
||||
# 控制台摘要（纯 ASCII，避免 Windows 编码问题）
|
||||
summary = eval_result.get("summary", {})
|
||||
print(f"\n{'='*60}")
|
||||
print(f" RAG E2E Eval [{phase_name}] DONE")
|
||||
print(f"{'='*60}")
|
||||
print(f" Avg Keyword Coverage: {summary.get('avg_keyword_coverage', 0):.2%}")
|
||||
print(f" Avg LLM Score: {summary.get('avg_llm_overall', 0):.2f}")
|
||||
print(f" Questions: {summary.get('total_questions', 0)}")
|
||||
|
||||
if comparison:
|
||||
sd = comparison["summary_delta"]
|
||||
print(f" --- vs Baseline ---")
|
||||
print(f" Keyword delta: {sd['keyword_delta']:+.2%}")
|
||||
print(f" LLM delta: {sd['llm_delta']:+.2f}")
|
||||
regressions = comparison.get("regressions", [])
|
||||
if regressions:
|
||||
print(f" !! Regressions: {len(regressions)}")
|
||||
for r in regressions:
|
||||
print(f" - {r['id']}: kw {r['delta_kw']:+.0%}")
|
||||
|
||||
print(f" JSON: {output_path}")
|
||||
print(f" Report: {report_path}")
|
||||
print(f" Log: {log_path}")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
553
scripts/evaluate_answer.py
Normal file
553
scripts/evaluate_answer.py
Normal file
@@ -0,0 +1,553 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
RAG 答案层评测脚本
|
||||
|
||||
评测指标：
|
||||
- LLM Score: 使用LLM对答案质量打分 (0-1)
|
||||
- ROUGE-L: 最长公共子序列相似度
|
||||
- Semantic Similarity: 语义相似度（使用embedding）
|
||||
|
||||
用法:
|
||||
python scripts/evaluate_answer.py --eval_dataset data/eval_dataset.json
|
||||
python scripts/evaluate_answer.py --topk 5 --output results/answer_results.json
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
# 添加项目根目录到路径
|
||||
PROJECT_ROOT = Path(__file__).parent.parent
|
||||
sys.path.insert(0, str(PROJECT_ROOT))
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AnswerEvaluator:
|
||||
"""答案层评测器"""
|
||||
|
||||
def __init__(self, engine=None):
|
||||
"""
|
||||
初始化评测器
|
||||
|
||||
Args:
|
||||
engine: RAG引擎实例，如果为None则自动创建
|
||||
"""
|
||||
self.engine = engine
|
||||
self._llm_client = None
|
||||
self._embedding_model = None
|
||||
|
||||
def _get_engine(self):
|
||||
"""延迟加载引擎"""
|
||||
if self.engine is None:
|
||||
from core.engine import get_engine
|
||||
self.engine = get_engine()
|
||||
return self.engine
|
||||
|
||||
def _get_llm_client(self):
|
||||
"""延迟加载LLM客户端"""
|
||||
if self._llm_client is None:
|
||||
try:
|
||||
from openai import OpenAI
|
||||
from config import API_KEY, BASE_URL, MODEL
|
||||
self._llm_client = OpenAI(
|
||||
api_key=API_KEY,
|
||||
base_url=BASE_URL
|
||||
)
|
||||
self._llm_model = MODEL
|
||||
except Exception as e:
|
||||
logger.warning(f"LLM客户端初始化失败: {e}")
|
||||
self._llm_client = None
|
||||
return self._llm_client
|
||||
|
||||
def _get_embedding_model(self):
|
||||
"""延迟加载embedding模型"""
|
||||
if self._embedding_model is None:
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from config import EMBEDDING_MODEL_PATH
|
||||
self._embedding_model = SentenceTransformer(EMBEDDING_MODEL_PATH)
|
||||
except Exception as e:
|
||||
logger.warning(f"Embedding模型初始化失败: {e}")
|
||||
self._embedding_model = None
|
||||
return self._embedding_model
|
||||
|
||||
def rouge_l(self, generated: str, reference: str) -> float:
|
||||
"""
|
||||
计算 ROUGE-L 分数
|
||||
|
||||
Args:
|
||||
generated: 生成的答案
|
||||
reference: 参考答案
|
||||
|
||||
Returns:
|
||||
ROUGE-L F1 分数 (0-1)
|
||||
"""
|
||||
if not generated or not reference:
|
||||
return 0.0
|
||||
|
||||
# 简单分词（按字符）
|
||||
gen_tokens = list(generated)
|
||||
ref_tokens = list(reference)
|
||||
|
||||
# 计算最长公共子序列长度
|
||||
m, n = len(gen_tokens), len(ref_tokens)
|
||||
if m == 0 or n == 0:
|
||||
return 0.0
|
||||
|
||||
# DP计算LCS
|
||||
dp = [[0] * (n + 1) for _ in range(m + 1)]
|
||||
for i in range(1, m + 1):
|
||||
for j in range(1, n + 1):
|
||||
if gen_tokens[i - 1] == ref_tokens[j - 1]:
|
||||
dp[i][j] = dp[i - 1][j - 1] + 1
|
||||
else:
|
||||
dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])
|
||||
|
||||
lcs_len = dp[m][n]
|
||||
|
||||
# 计算Precision和Recall
|
||||
precision = lcs_len / m if m > 0 else 0.0
|
||||
recall = lcs_len / n if n > 0 else 0.0
|
||||
|
||||
# F1分数
|
||||
if precision + recall == 0:
|
||||
return 0.0
|
||||
f1 = 2 * precision * recall / (precision + recall)
|
||||
|
||||
return f1
|
||||
|
||||
def semantic_similarity(self, generated: str, reference: str) -> float:
|
||||
"""
|
||||
计算语义相似度
|
||||
|
||||
Args:
|
||||
generated: 生成的答案
|
||||
reference: 参考答案
|
||||
|
||||
Returns:
|
||||
语义相似度 (0-1)
|
||||
"""
|
||||
model = self._get_embedding_model()
|
||||
if model is None:
|
||||
return 0.0
|
||||
|
||||
try:
|
||||
embeddings = model.encode([generated, reference])
|
||||
# 余弦相似度
|
||||
similarity = np.dot(embeddings[0], embeddings[1]) / (
|
||||
np.linalg.norm(embeddings[0]) * np.linalg.norm(embeddings[1])
|
||||
)
|
||||
return float(max(0, min(1, similarity))) # 确保在0-1范围内
|
||||
except Exception as e:
|
||||
logger.warning(f"计算语义相似度失败: {e}")
|
||||
return 0.0
|
||||
|
||||
def llm_score(self, query: str, generated: str, reference: str) -> float:
|
||||
"""
|
||||
使用LLM对答案质量打分
|
||||
|
||||
Args:
|
||||
query: 用户问题
|
||||
generated: 生成的答案
|
||||
reference: 参考答案
|
||||
|
||||
Returns:
|
||||
LLM评分 (0-1)
|
||||
"""
|
||||
client = self._get_llm_client()
|
||||
if client is None:
|
||||
logger.warning("LLM客户端不可用，跳过LLM评分")
|
||||
return 0.0
|
||||
|
||||
prompt = f"""请作为专业评测员，对RAG系统的回答质量进行评分。
|
||||
|
||||
【用户问题】
|
||||
{query}
|
||||
|
||||
【参考答案】
|
||||
{reference}
|
||||
|
||||
【系统回答】
|
||||
{generated}
|
||||
|
||||
【评分标准】
|
||||
请从以下维度评分（每项0-10分）：
|
||||
1. 准确性：回答是否与参考答案的核心信息一致
|
||||
2. 完整性：回答是否覆盖了参考答案的关键要点
|
||||
3. 相关性：回答是否直接回答了用户问题
|
||||
4. 流畅性：回答是否通顺、易于理解
|
||||
|
||||
请以JSON格式返回评分：
|
||||
{{"accuracy": X, "completeness": X, "relevance": X, "fluency": X, "overall": X}}
|
||||
|
||||
只返回JSON，不要其他内容。"""
|
||||
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model=self._llm_model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=0.1,
|
||||
max_tokens=200
|
||||
)
|
||||
|
||||
result_text = response.choices[0].message.content.strip()
|
||||
|
||||
# 尝试解析JSON
|
||||
import re
|
||||
json_match = re.search(r'\{[^}]+\}', result_text)
|
||||
if json_match:
|
||||
scores = json.loads(json_match.group())
|
||||
overall = scores.get('overall', 0)
|
||||
return min(1.0, max(0.0, overall / 10.0))
|
||||
else:
|
||||
# 尝试从文本中提取数字
|
||||
numbers = re.findall(r'\d+', result_text)
|
||||
if numbers:
|
||||
return min(1.0, max(0.0, int(numbers[-1]) / 10.0))
|
||||
return 0.5
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"LLM评分失败: {e}")
|
||||
return 0.0
|
||||
|
||||
def generate_answer(self, query: str, top_k: int = 5, collections: list = None) -> str:
|
||||
"""
|
||||
使用RAG引擎生成答案
|
||||
|
||||
Args:
|
||||
query: 用户问题
|
||||
top_k: 检索数量
|
||||
collections: 目标向量库列表
|
||||
|
||||
Returns:
|
||||
生成的答案
|
||||
"""
|
||||
engine = self._get_engine()
|
||||
|
||||
try:
|
||||
# 先检索
|
||||
results = engine.search_knowledge(
|
||||
query=query,
|
||||
top_k=top_k,
|
||||
collections=collections
|
||||
)
|
||||
|
||||
# 提取文档内容
|
||||
if isinstance(results, dict):
|
||||
documents = results.get('documents', [])
|
||||
if documents and isinstance(documents[0], list):
|
||||
documents = documents[0]
|
||||
else:
|
||||
documents = []
|
||||
|
||||
if not documents:
|
||||
return "抱歉，未找到相关信息。"
|
||||
|
||||
# 简单拼接作为答案（实际应该用LLM生成）
|
||||
# 这里为了评测，我们用检索到的内容拼接
|
||||
context = "\n\n".join(documents[:3])
|
||||
|
||||
# 使用LLM生成答案
|
||||
client = self._get_llm_client()
|
||||
if client:
|
||||
prompt = f"""基于以下检索到的信息回答用户问题。请简洁准确地回答。
|
||||
|
||||
【用户问题】
|
||||
{query}
|
||||
|
||||
【检索到的信息】
|
||||
{context}
|
||||
|
||||
请直接回答问题，不要重复问题本身："""
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self._llm_model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=0.3,
|
||||
max_tokens=500
|
||||
)
|
||||
return response.choices[0].message.content.strip()
|
||||
else:
|
||||
# 没有LLM时，直接返回最相关的文档片段
|
||||
return documents[0] if documents else "无法生成答案"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"生成答案失败: {e}")
|
||||
return f"生成答案时出错: {str(e)}"
|
||||
|
||||
def evaluate_query(
|
||||
self,
|
||||
query: str,
|
||||
reference_answer: str,
|
||||
top_k: int = 5,
|
||||
collections: list = None,
|
||||
use_llm: bool = True
|
||||
) -> dict:
|
||||
"""
|
||||
评估单个查询的答案质量
|
||||
|
||||
Args:
|
||||
query: 用户问题
|
||||
reference_answer: 参考答案
|
||||
top_k: 检索数量
|
||||
collections: 目标向量库列表
|
||||
use_llm: 是否使用LLM评分
|
||||
|
||||
Returns:
|
||||
评测结果字典
|
||||
"""
|
||||
# 生成答案
|
||||
generated_answer = self.generate_answer(query, top_k=top_k, collections=collections)
|
||||
|
||||
# 计算各项指标
|
||||
metrics = {
|
||||
'rouge_l': self.rouge_l(generated_answer, reference_answer),
|
||||
'semantic_similarity': self.semantic_similarity(generated_answer, reference_answer),
|
||||
'generated_answer': generated_answer[:500], # 截断保存
|
||||
'reference_answer': reference_answer[:500]
|
||||
}
|
||||
|
||||
# LLM评分（可选，因为较慢）
|
||||
if use_llm:
|
||||
metrics['llm_score'] = self.llm_score(query, generated_answer, reference_answer)
|
||||
|
||||
return metrics
|
||||
|
||||
def evaluate_dataset(
|
||||
self,
|
||||
eval_dataset_path: str,
|
||||
top_k: int = 5,
|
||||
collections: list = None,
|
||||
use_llm: bool = True,
|
||||
sample_size: int = None
|
||||
) -> dict:
|
||||
"""
|
||||
评估整个数据集
|
||||
|
||||
Args:
|
||||
eval_dataset_path: 评测数据集路径
|
||||
top_k: 检索数量
|
||||
collections: 目标向量库列表
|
||||
use_llm: 是否使用LLM评分
|
||||
sample_size: 采样数量（用于快速测试）
|
||||
|
||||
Returns:
|
||||
汇总评测结果
|
||||
"""
|
||||
# 加载数据集
|
||||
with open(eval_dataset_path, 'r', encoding='utf-8') as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
queries = dataset.get('queries', [])
|
||||
|
||||
# 采样（如果指定）
|
||||
if sample_size and sample_size < len(queries):
|
||||
import random
|
||||
queries = random.sample(queries, sample_size)
|
||||
|
||||
# 存储每个查询的结果
|
||||
all_metrics = {
|
||||
'rouge_l': [],
|
||||
'semantic_similarity': [],
|
||||
'llm_score': []
|
||||
}
|
||||
|
||||
# 按查询类型分组
|
||||
by_type = {}
|
||||
# 按难度分组
|
||||
by_difficulty = {}
|
||||
|
||||
logger.info(f"开始答案层评测，共 {len(queries)} 条查询...")
|
||||
|
||||
for i, q in enumerate(queries, 1):
|
||||
query_text = q['query']
|
||||
query_type = q.get('query_type', 'unknown')
|
||||
difficulty = q.get('difficulty', 'medium')
|
||||
reference_answer = q.get('reference_answer', '')
|
||||
|
||||
if not reference_answer:
|
||||
logger.warning(f"查询 {q['id']} 没有参考答案，跳过")
|
||||
continue
|
||||
|
||||
# 执行评测
|
||||
metrics = self.evaluate_query(
|
||||
query=query_text,
|
||||
reference_answer=reference_answer,
|
||||
top_k=top_k,
|
||||
collections=collections,
|
||||
use_llm=use_llm
|
||||
)
|
||||
|
||||
# 收集结果
|
||||
all_metrics['rouge_l'].append(metrics['rouge_l'])
|
||||
all_metrics['semantic_similarity'].append(metrics['semantic_similarity'])
|
||||
if 'llm_score' in metrics:
|
||||
all_metrics['llm_score'].append(metrics['llm_score'])
|
||||
|
||||
# 按类型分组
|
||||
if query_type not in by_type:
|
||||
by_type[query_type] = {k: [] for k in all_metrics}
|
||||
by_type[query_type]['rouge_l'].append(metrics['rouge_l'])
|
||||
by_type[query_type]['semantic_similarity'].append(metrics['semantic_similarity'])
|
||||
if 'llm_score' in metrics:
|
||||
by_type[query_type]['llm_score'].append(metrics['llm_score'])
|
||||
|
||||
# 按难度分组
|
||||
if difficulty not in by_difficulty:
|
||||
by_difficulty[difficulty] = {k: [] for k in all_metrics}
|
||||
by_difficulty[difficulty]['rouge_l'].append(metrics['rouge_l'])
|
||||
by_difficulty[difficulty]['semantic_similarity'].append(metrics['semantic_similarity'])
|
||||
if 'llm_score' in metrics:
|
||||
by_difficulty[difficulty]['llm_score'].append(metrics['llm_score'])
|
||||
|
||||
# 打印进度
|
||||
llm_str = f", LLM={metrics.get('llm_score', 0):.2f}" if 'llm_score' in metrics else ""
|
||||
logger.info(f" [{i}/{len(queries)}] {q['id']}: ROUGE-L={metrics['rouge_l']:.2f}, SemSim={metrics['semantic_similarity']:.2f}{llm_str}")
|
||||
|
||||
# 计算平均值
|
||||
results = {
|
||||
'overall': {
|
||||
key: np.mean(values) if values else 0.0
|
||||
for key, values in all_metrics.items()
|
||||
},
|
||||
'by_type': {
|
||||
qtype: {key: np.mean(vals) if vals else 0.0 for key, vals in metrics.items()}
|
||||
for qtype, metrics in by_type.items()
|
||||
},
|
||||
'by_difficulty': {
|
||||
diff: {key: np.mean(vals) if vals else 0.0 for key, vals in metrics.items()}
|
||||
for diff, metrics in by_difficulty.items()
|
||||
},
|
||||
'total_queries': len(queries),
|
||||
'evaluated_queries': len(all_metrics['rouge_l'])
|
||||
}
|
||||
|
||||
return results
|
||||
|
||||
def print_results(self, results: dict):
|
||||
"""打印评测结果"""
|
||||
print("\n" + "=" * 60)
|
||||
print(" RAG 答案层评测结果")
|
||||
print("=" * 60)
|
||||
|
||||
# 整体结果
|
||||
print("\n【整体指标】")
|
||||
overall = results.get('overall', {})
|
||||
print(f" ROUGE-L: {overall.get('rouge_l', 0):.4f}")
|
||||
print(f" Semantic Sim: {overall.get('semantic_similarity', 0):.4f}")
|
||||
if overall.get('llm_score'):
|
||||
print(f" LLM Score: {overall.get('llm_score', 0):.4f}")
|
||||
|
||||
# 按查询类型
|
||||
print("\n【按查询类型】")
|
||||
by_type = results.get('by_type', {})
|
||||
for qtype, metrics in sorted(by_type.items()):
|
||||
print(f" {qtype}:")
|
||||
print(f" ROUGE-L: {metrics.get('rouge_l', 0):.4f}")
|
||||
print(f" SemSim: {metrics.get('semantic_similarity', 0):.4f}")
|
||||
|
||||
# 按难度
|
||||
print("\n【按难度】")
|
||||
by_difficulty = results.get('by_difficulty', {})
|
||||
for diff, metrics in sorted(by_difficulty.items()):
|
||||
print(f" {diff}:")
|
||||
print(f" ROUGE-L: {metrics.get('rouge_l', 0):.4f}")
|
||||
print(f" SemSim: {metrics.get('semantic_similarity', 0):.4f}")
|
||||
|
||||
# 统计信息
|
||||
print("\n【统计信息】")
|
||||
print(f" 总查询数: {results.get('total_queries', 0)}")
|
||||
print(f" 已评测数: {results.get('evaluated_queries', 0)}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='RAG 答案层评测脚本')
|
||||
parser.add_argument(
|
||||
'--eval_dataset',
|
||||
type=str,
|
||||
default='data/eval_dataset.json',
|
||||
help='评测数据集路径'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--topk',
|
||||
type=int,
|
||||
default=5,
|
||||
help='检索返回数量 (默认: 5)'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--collections',
|
||||
type=str,
|
||||
nargs='+',
|
||||
default=None,
|
||||
help='目标向量库列表'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--no_llm',
|
||||
action='store_true',
|
||||
help='跳过LLM评分（更快但指标较少）'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--sample',
|
||||
type=int,
|
||||
default=None,
|
||||
help='采样数量（用于快速测试）'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--output',
|
||||
type=str,
|
||||
default=None,
|
||||
help='结果输出文件路径 (JSON格式)'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 检查数据集文件
|
||||
eval_path = PROJECT_ROOT / args.eval_dataset
|
||||
if not eval_path.exists():
|
||||
logger.error(f"评测数据集不存在: {eval_path}")
|
||||
sys.exit(1)
|
||||
|
||||
# 创建评测器
|
||||
logger.info("初始化答案层评测器...")
|
||||
evaluator = AnswerEvaluator()
|
||||
|
||||
# 执行评测
|
||||
start_time = time.time()
|
||||
results = evaluator.evaluate_dataset(
|
||||
eval_dataset_path=str(eval_path),
|
||||
top_k=args.topk,
|
||||
collections=args.collections,
|
||||
use_llm=not args.no_llm,
|
||||
sample_size=args.sample
|
||||
)
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
# 打印结果
|
||||
evaluator.print_results(results)
|
||||
print(f"\n评测耗时: {elapsed_time:.2f} 秒")
|
||||
|
||||
# 保存结果
|
||||
if args.output:
|
||||
output_path = PROJECT_ROOT / args.output
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||||
logger.info(f"结果已保存到: {output_path}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
459
scripts/evaluate_rag.py
Normal file
459
scripts/evaluate_rag.py
Normal file
@@ -0,0 +1,459 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
RAG 检索层评测脚本
|
||||
|
||||
评测指标：
|
||||
- Recall@k: 召回率 - 命中相关切片数 / 相关切片总数
|
||||
- MRR: 平均倒数排名 - 1 / 第一命中的排名
|
||||
- Hit Rate: 命中率 - 命中查询数 / 总查询数
|
||||
- nDCG: 归一化折损累积增益 - 考虑位置的加权得分
|
||||
|
||||
用法:
|
||||
python scripts/evaluate_rag.py --topk 5 --embedding bge-base --rerank on
|
||||
python scripts/evaluate_rag.py --eval_dataset data/eval_dataset.json
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
# 添加项目根目录到路径
|
||||
PROJECT_ROOT = Path(__file__).parent.parent
|
||||
sys.path.insert(0, str(PROJECT_ROOT))
|
||||
|
||||
from core.engine import get_engine, RAGEngine
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RetrievalEvaluator:
|
||||
"""检索层评测器"""
|
||||
|
||||
def __init__(self, engine: Optional[RAGEngine] = None):
|
||||
"""
|
||||
初始化评测器
|
||||
|
||||
Args:
|
||||
engine: RAG引擎实例，如果为None则自动创建
|
||||
"""
|
||||
self.engine = engine or get_engine()
|
||||
|
||||
def recall_at_k(self, retrieved_ids: list, relevant_ids: list, k: int) -> float:
|
||||
"""
|
||||
计算 Recall@k
|
||||
|
||||
Args:
|
||||
retrieved_ids: 检索返回的切片ID列表
|
||||
relevant_ids: 相关切片ID列表
|
||||
k: 截断位置
|
||||
|
||||
Returns:
|
||||
Recall@k 值 (0-1)
|
||||
"""
|
||||
if not relevant_ids:
|
||||
return 0.0
|
||||
|
||||
retrieved_set = set(retrieved_ids[:k])
|
||||
relevant_set = set(relevant_ids)
|
||||
|
||||
hits = len(retrieved_set & relevant_set)
|
||||
return hits / len(relevant_set)
|
||||
|
||||
def mrr(self, retrieved_ids: list, relevant_ids: list) -> float:
|
||||
"""
|
||||
计算 MRR (Mean Reciprocal Rank)
|
||||
|
||||
Args:
|
||||
retrieved_ids: 检索返回的切片ID列表
|
||||
relevant_ids: 相关切片ID列表
|
||||
|
||||
Returns:
|
||||
MRR 值 (0-1)
|
||||
"""
|
||||
if not relevant_ids:
|
||||
return 0.0
|
||||
|
||||
relevant_set = set(relevant_ids)
|
||||
|
||||
for rank, rid in enumerate(retrieved_ids, start=1):
|
||||
if rid in relevant_set:
|
||||
return 1.0 / rank
|
||||
|
||||
return 0.0
|
||||
|
||||
def hit_rate(self, retrieved_ids: list, relevant_ids: list, k: int) -> float:
|
||||
"""
|
||||
计算 Hit Rate@k
|
||||
|
||||
Args:
|
||||
retrieved_ids: 检索返回的切片ID列表
|
||||
relevant_ids: 相关切片ID列表
|
||||
k: 截断位置
|
||||
|
||||
Returns:
|
||||
Hit Rate 值 (0-1)
|
||||
"""
|
||||
if not relevant_ids:
|
||||
return 0.0
|
||||
|
||||
retrieved_set = set(retrieved_ids[:k])
|
||||
relevant_set = set(relevant_ids)
|
||||
|
||||
return 1.0 if retrieved_set & relevant_set else 0.0
|
||||
|
||||
def ndcg_at_k(self, retrieved_ids: list, relevant_ids: list, k: int) -> float:
|
||||
"""
|
||||
计算 nDCG@k (Normalized Discounted Cumulative Gain)
|
||||
|
||||
Args:
|
||||
retrieved_ids: 检索返回的切片ID列表
|
||||
relevant_ids: 相关切片ID列表
|
||||
k: 截断位置
|
||||
|
||||
Returns:
|
||||
nDCG@k 值 (0-1)
|
||||
"""
|
||||
if not relevant_ids:
|
||||
return 0.0
|
||||
|
||||
relevant_set = set(relevant_ids)
|
||||
|
||||
# 计算 DCG
|
||||
dcg = 0.0
|
||||
for i, rid in enumerate(retrieved_ids[:k], start=1):
|
||||
if rid in relevant_set:
|
||||
dcg += 1.0 / np.log2(i + 1)
|
||||
|
||||
# 计算 IDCG (理想情况)
|
||||
idcg = 0.0
|
||||
for i in range(1, min(len(relevant_ids), k) + 1):
|
||||
idcg += 1.0 / np.log2(i + 1)
|
||||
|
||||
return dcg / idcg if idcg > 0 else 0.0
|
||||
|
||||
def retrieve(self, query: str, top_k: int = 5, collections: list = None) -> list:
|
||||
"""
|
||||
执行检索
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
top_k: 返回数量
|
||||
collections: 目标向量库列表
|
||||
|
||||
Returns:
|
||||
检索结果列表，每个元素包含 id, score, content 等
|
||||
"""
|
||||
try:
|
||||
results = self.engine.search_knowledge(
|
||||
query=query,
|
||||
top_k=top_k,
|
||||
collections=collections
|
||||
)
|
||||
return results
|
||||
except Exception as e:
|
||||
logger.error(f"检索失败: {e}")
|
||||
return []
|
||||
|
||||
def evaluate_query(
|
||||
self,
|
||||
query: str,
|
||||
relevant_ids: list,
|
||||
top_k: int = 5,
|
||||
collections: list = None
|
||||
) -> dict:
|
||||
"""
|
||||
评估单个查询
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
relevant_ids: 相关切片ID列表
|
||||
top_k: 检索数量
|
||||
collections: 目标向量库列表
|
||||
|
||||
Returns:
|
||||
评测结果字典
|
||||
"""
|
||||
# 执行检索
|
||||
results = self.retrieve(query, top_k=top_k, collections=collections)
|
||||
|
||||
# 提取检索到的ID
|
||||
# search_knowledge 返回格式: {'ids': [[...]], 'documents': [[...]], 'metadatas': [[...]], 'distances': [[...]]}
|
||||
# 注意：ids 是嵌套列表，需要展平
|
||||
if isinstance(results, dict):
|
||||
ids = results.get('ids', [])
|
||||
# 展平嵌套列表
|
||||
if ids and isinstance(ids[0], list):
|
||||
retrieved_ids = ids[0] if ids else []
|
||||
else:
|
||||
retrieved_ids = ids
|
||||
elif isinstance(results, list):
|
||||
retrieved_ids = [r.get('id', r.get('chunk_id', '')) if isinstance(r, dict) else str(r) for r in results]
|
||||
else:
|
||||
retrieved_ids = []
|
||||
|
||||
# 计算各项指标
|
||||
metrics = {
|
||||
'recall@k': self.recall_at_k(retrieved_ids, relevant_ids, top_k),
|
||||
'mrr': self.mrr(retrieved_ids, relevant_ids),
|
||||
f'hit_rate@{top_k}': self.hit_rate(retrieved_ids, relevant_ids, top_k),
|
||||
f'ndcg@{top_k}': self.ndcg_at_k(retrieved_ids, relevant_ids, top_k),
|
||||
'retrieved_count': len(retrieved_ids),
|
||||
'relevant_count': len(relevant_ids)
|
||||
}
|
||||
|
||||
return metrics
|
||||
|
||||
def get_chunk_ids_by_doc(self, doc_name: str, collection: str = "public_kb") -> list:
|
||||
"""
|
||||
获取指定文档的所有切片ID
|
||||
|
||||
Args:
|
||||
doc_name: 文档名称
|
||||
collection: 向量库名称
|
||||
|
||||
Returns:
|
||||
切片ID列表
|
||||
"""
|
||||
try:
|
||||
# 使用向量库管理器查询
|
||||
from knowledge.manager import get_kb_manager
|
||||
manager = get_kb_manager()
|
||||
|
||||
# 获取collection对象
|
||||
coll = manager.get_collection(collection)
|
||||
if not coll:
|
||||
logger.warning(f"向量库不存在: {collection}")
|
||||
return []
|
||||
|
||||
# 查询指定source的所有切片
|
||||
results = coll.get(
|
||||
where={"source": doc_name},
|
||||
include=['metadatas']
|
||||
)
|
||||
|
||||
return results.get('ids', [])
|
||||
except Exception as e:
|
||||
logger.warning(f"获取文档切片ID失败: {e}")
|
||||
return []
|
||||
|
||||
def evaluate_dataset(
|
||||
self,
|
||||
eval_dataset_path: str,
|
||||
top_k: int = 5,
|
||||
collections: list = None
|
||||
) -> dict:
|
||||
"""
|
||||
评估整个数据集
|
||||
|
||||
Args:
|
||||
eval_dataset_path: 评测数据集路径
|
||||
top_k: 检索数量
|
||||
collections: 目标向量库列表
|
||||
|
||||
Returns:
|
||||
汇总评测结果
|
||||
"""
|
||||
# 加载数据集
|
||||
with open(eval_dataset_path, 'r', encoding='utf-8') as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
queries = dataset.get('queries', [])
|
||||
|
||||
# 存储每个查询的结果
|
||||
all_metrics = {
|
||||
'recall@k': [],
|
||||
'mrr': [],
|
||||
f'hit_rate@{top_k}': [],
|
||||
f'ndcg@{top_k}': []
|
||||
}
|
||||
|
||||
# 按查询类型分组
|
||||
by_type = {}
|
||||
# 按难度分组
|
||||
by_difficulty = {}
|
||||
|
||||
logger.info(f"开始评测，共 {len(queries)} 条查询...")
|
||||
|
||||
for i, q in enumerate(queries, 1):
|
||||
query_text = q['query']
|
||||
query_type = q.get('query_type', 'unknown')
|
||||
difficulty = q.get('difficulty', 'medium')
|
||||
relevant_docs = q.get('relevant_docs', [])
|
||||
|
||||
# 获取相关文档的切片ID作为 ground truth
|
||||
# 注意：这里简化处理，实际应该根据具体业务逻辑
|
||||
# 如果数据集有 relevant_chunks 字段则直接使用
|
||||
relevant_ids = q.get('relevant_chunks', [])
|
||||
|
||||
# 如果没有 relevant_chunks，尝试根据 relevant_docs 获取
|
||||
if not relevant_ids and relevant_docs:
|
||||
for doc in relevant_docs:
|
||||
ids = self.get_chunk_ids_by_doc(doc)
|
||||
relevant_ids.extend(ids)
|
||||
|
||||
if not relevant_ids:
|
||||
logger.warning(f"查询 {q['id']} 没有相关切片ID，跳过")
|
||||
continue
|
||||
|
||||
# 执行评测
|
||||
metrics = self.evaluate_query(
|
||||
query=query_text,
|
||||
relevant_ids=relevant_ids,
|
||||
top_k=top_k,
|
||||
collections=collections
|
||||
)
|
||||
|
||||
# 收集结果
|
||||
for key in all_metrics:
|
||||
if key in metrics:
|
||||
all_metrics[key].append(metrics[key])
|
||||
|
||||
# 按类型分组
|
||||
if query_type not in by_type:
|
||||
by_type[query_type] = {k: [] for k in all_metrics}
|
||||
for key in all_metrics:
|
||||
if key in metrics:
|
||||
by_type[query_type][key].append(metrics[key])
|
||||
|
||||
# 按难度分组
|
||||
if difficulty not in by_difficulty:
|
||||
by_difficulty[difficulty] = {k: [] for k in all_metrics}
|
||||
for key in all_metrics:
|
||||
if key in metrics:
|
||||
by_difficulty[difficulty][key].append(metrics[key])
|
||||
|
||||
logger.info(f" [{i}/{len(queries)}] {q['id']}: Recall@{top_k}={metrics['recall@k']:.2f}, Hit={metrics[f'hit_rate@{top_k}']:.2f}")
|
||||
|
||||
# 计算平均值
|
||||
results = {
|
||||
'overall': {
|
||||
key: np.mean(values) if values else 0.0
|
||||
for key, values in all_metrics.items()
|
||||
},
|
||||
'by_type': {
|
||||
qtype: {key: np.mean(vals) if vals else 0.0 for key, vals in metrics.items()}
|
||||
for qtype, metrics in by_type.items()
|
||||
},
|
||||
'by_difficulty': {
|
||||
diff: {key: np.mean(vals) if vals else 0.0 for key, vals in metrics.items()}
|
||||
for diff, metrics in by_difficulty.items()
|
||||
},
|
||||
'total_queries': len(queries),
|
||||
'evaluated_queries': len(all_metrics['recall@k'])
|
||||
}
|
||||
|
||||
return results
|
||||
|
||||
def print_results(self, results: dict):
|
||||
"""打印评测结果"""
|
||||
print("\n" + "=" * 60)
|
||||
print(" RAG 检索层评测结果")
|
||||
print("=" * 60)
|
||||
|
||||
# 整体结果
|
||||
print("\n【整体指标】")
|
||||
overall = results.get('overall', {})
|
||||
print(f" Recall@k: {overall.get('recall@k', 0):.4f}")
|
||||
print(f" MRR: {overall.get('mrr', 0):.4f}")
|
||||
print(f" Hit Rate: {overall.get('hit_rate@5', overall.get('hit_rate@k', 0)):.4f}")
|
||||
print(f" nDCG: {overall.get('ndcg@5', overall.get('ndcg@k', 0)):.4f}")
|
||||
|
||||
# 按查询类型
|
||||
print("\n【按查询类型】")
|
||||
by_type = results.get('by_type', {})
|
||||
for qtype, metrics in sorted(by_type.items()):
|
||||
print(f" {qtype}:")
|
||||
print(f" Recall@k: {metrics.get('recall@k', 0):.4f}")
|
||||
print(f" Hit Rate: {metrics.get('hit_rate@5', metrics.get('hit_rate@k', 0)):.4f}")
|
||||
|
||||
# 按难度
|
||||
print("\n【按难度】")
|
||||
by_difficulty = results.get('by_difficulty', {})
|
||||
for diff, metrics in sorted(by_difficulty.items()):
|
||||
print(f" {diff}:")
|
||||
print(f" Recall@k: {metrics.get('recall@k', 0):.4f}")
|
||||
print(f" Hit Rate: {metrics.get('hit_rate@5', metrics.get('hit_rate@k', 0)):.4f}")
|
||||
|
||||
# 统计信息
|
||||
print("\n【统计信息】")
|
||||
print(f" 总查询数: {results.get('total_queries', 0)}")
|
||||
print(f" 已评测数: {results.get('evaluated_queries', 0)}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='RAG 检索层评测脚本')
|
||||
parser.add_argument(
|
||||
'--eval_dataset',
|
||||
type=str,
|
||||
default='data/eval_dataset.json',
|
||||
help='评测数据集路径'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--topk',
|
||||
type=int,
|
||||
default=5,
|
||||
help='检索返回数量 (默认: 5)'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--collections',
|
||||
type=str,
|
||||
nargs='+',
|
||||
default=None,
|
||||
help='目标向量库列表'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--output',
|
||||
type=str,
|
||||
default=None,
|
||||
help='结果输出文件路径 (JSON格式)'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 检查数据集文件
|
||||
eval_path = PROJECT_ROOT / args.eval_dataset
|
||||
if not eval_path.exists():
|
||||
logger.error(f"评测数据集不存在: {eval_path}")
|
||||
sys.exit(1)
|
||||
|
||||
# 创建评测器
|
||||
logger.info("初始化评测器...")
|
||||
evaluator = RetrievalEvaluator()
|
||||
|
||||
# 执行评测
|
||||
start_time = time.time()
|
||||
results = evaluator.evaluate_dataset(
|
||||
eval_dataset_path=str(eval_path),
|
||||
top_k=args.topk,
|
||||
collections=args.collections
|
||||
)
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
# 打印结果
|
||||
evaluator.print_results(results)
|
||||
print(f"\n评测耗时: {elapsed_time:.2f} 秒")
|
||||
|
||||
# 保存结果
|
||||
if args.output:
|
||||
output_path = PROJECT_ROOT / args.output
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||||
logger.info(f"结果已保存到: {output_path}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
116
scripts/fix_image_paths.py
Normal file
116
scripts/fix_image_paths.py
Normal file
@@ -0,0 +1,116 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
修复向量库中的图片路径格式
|
||||
|
||||
问题：历史数据中 image_path 存储为 "images/xxx.jpg" 格式
|
||||
修复：统一改为只存储文件名 "xxx.jpg"
|
||||
|
||||
使用方式：
|
||||
python scripts/fix_image_paths.py [--dry-run]
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
# 添加项目根目录到路径
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
import chromadb
|
||||
|
||||
|
||||
def fix_image_paths(dry_run: bool = False):
|
||||
"""
|
||||
修复向量库中的图片路径格式
|
||||
|
||||
Args:
|
||||
dry_run: 仅预览，不实际修改
|
||||
"""
|
||||
chroma_path = "knowledge/vector_store/chroma"
|
||||
client = chromadb.PersistentClient(path=chroma_path)
|
||||
|
||||
# 获取所有集合
|
||||
collections = client.list_collections()
|
||||
print(f"找到 {len(collections)} 个集合")
|
||||
|
||||
for collection in collections:
|
||||
print(f"\n处理集合: {collection.name}")
|
||||
|
||||
# 获取所有图片类型的切片
|
||||
try:
|
||||
results = collection.get(
|
||||
where={"chunk_type": "image"},
|
||||
include=["metadatas", "documents", "embeddings"]
|
||||
)
|
||||
|
||||
if not results["ids"]:
|
||||
print(f" 无图片类型切片")
|
||||
continue
|
||||
|
||||
print(f" 找到 {len(results['ids'])} 个图片切片")
|
||||
|
||||
# 检查需要修复的路径
|
||||
needs_fix = []
|
||||
for i, (id, meta) in enumerate(zip(results["ids"], results["metadatas"])):
|
||||
image_path = meta.get("image_path", "")
|
||||
if image_path and ("/" in image_path or "\\" in image_path):
|
||||
# 路径包含目录，需要修复
|
||||
new_path = os.path.basename(image_path)
|
||||
needs_fix.append({
|
||||
"id": id,
|
||||
"old_path": image_path,
|
||||
"new_path": new_path,
|
||||
"index": i
|
||||
})
|
||||
|
||||
if not needs_fix:
|
||||
print(f" 所有路径格式正确，无需修复")
|
||||
continue
|
||||
|
||||
print(f" 需要修复 {len(needs_fix)} 条记录:")
|
||||
for item in needs_fix[:5]: # 只显示前5条
|
||||
print(f" {item['old_path']} -> {item['new_path']}")
|
||||
if len(needs_fix) > 5:
|
||||
print(f" ... 还有 {len(needs_fix) - 5} 条")
|
||||
|
||||
if dry_run:
|
||||
print(f" [DRY-RUN] 跳过实际修改")
|
||||
continue
|
||||
|
||||
# 执行修复
|
||||
ids_to_update = []
|
||||
metadatas_to_update = []
|
||||
|
||||
for item in needs_fix:
|
||||
# 更新元数据
|
||||
new_meta = results["metadatas"][item["index"]].copy()
|
||||
new_meta["image_path"] = item["new_path"]
|
||||
ids_to_update.append(item["id"])
|
||||
metadatas_to_update.append(new_meta)
|
||||
|
||||
# 批量更新
|
||||
collection.update(
|
||||
ids=ids_to_update,
|
||||
metadatas=metadatas_to_update
|
||||
)
|
||||
print(f" ✓ 已修复 {len(ids_to_update)} 条记录")
|
||||
|
||||
except Exception as e:
|
||||
print(f" 处理集合 {collection.name} 时出错: {e}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="修复向量库中的图片路径格式")
|
||||
parser.add_argument("--dry-run", action="store_true", help="仅预览，不实际修改")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.dry_run:
|
||||
print("=== DRY-RUN 模式：仅预览，不修改数据 ===\n")
|
||||
|
||||
fix_image_paths(dry_run=args.dry_run)
|
||||
|
||||
print("\n完成!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
44
scripts/migrate_add_metadata.py
Normal file
44
scripts/migrate_add_metadata.py
Normal file
@@ -0,0 +1,44 @@
|
||||
"""
|
||||
数据库迁移脚本：为 messages 表添加 metadata 列
|
||||
|
||||
运行方式：
|
||||
python scripts/migrate_add_metadata.py
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
import os
|
||||
|
||||
DB_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "rag_core.db")
|
||||
|
||||
def migrate():
|
||||
"""添加 metadata 列到 messages 表"""
|
||||
if not os.path.exists(DB_PATH):
|
||||
print(f"数据库不存在: {DB_PATH}")
|
||||
return
|
||||
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
# 检查列是否已存在
|
||||
cursor.execute("PRAGMA table_info(messages)")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if 'metadata' in columns:
|
||||
print("✓ metadata 列已存在，无需迁移")
|
||||
return
|
||||
|
||||
# 添加 metadata 列
|
||||
print("正在添加 metadata 列...")
|
||||
cursor.execute("ALTER TABLE messages ADD COLUMN metadata TEXT")
|
||||
conn.commit()
|
||||
print("✓ 迁移完成！")
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ 迁移失败: {e}")
|
||||
conn.rollback()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
migrate()
|
||||
44
scripts/migrate_mineru_models.bat
Normal file
44
scripts/migrate_mineru_models.bat
Normal file
@@ -0,0 +1,44 @@
|
||||
@echo off
|
||||
REM MinerU 模型迁移脚本 - Windows 版本
|
||||
REM
|
||||
REM 功能：将 HuggingFace 缓存中的 MinerU 模型迁移到项目 models/ 目录
|
||||
|
||||
echo ========================================
|
||||
echo MinerU 模型迁移工具
|
||||
echo ========================================
|
||||
echo.
|
||||
|
||||
REM 检查虚拟环境
|
||||
if not exist "venv\Scripts\activate.bat" (
|
||||
echo [错误] 虚拟环境不存在，请先创建虚拟环境
|
||||
echo 运行: python -m venv venv
|
||||
pause
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
REM 激活虚拟环境
|
||||
call venv\Scripts\activate.bat
|
||||
|
||||
REM 运行迁移脚本
|
||||
python scripts\migrate_mineru_models.py
|
||||
|
||||
REM 检查返回码
|
||||
if %ERRORLEVEL% EQU 0 (
|
||||
echo.
|
||||
echo ========================================
|
||||
echo 迁移成功！
|
||||
echo ========================================
|
||||
echo.
|
||||
echo 下一步：
|
||||
echo 1. 测试解析: python parsers\mineru_parser.py documents\test.pdf
|
||||
echo 2. 删除缓存: rmdir /s "%%USERPROFILE%%\.cache\huggingface\hub\models--opendatalab--*"
|
||||
echo.
|
||||
) else (
|
||||
echo.
|
||||
echo ========================================
|
||||
echo 迁移失败，请检查错误信息
|
||||
echo ========================================
|
||||
echo.
|
||||
)
|
||||
|
||||
pause
|
||||
182
scripts/migrate_mineru_models.py
Normal file
182
scripts/migrate_mineru_models.py
Normal file
@@ -0,0 +1,182 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
MinerU 模型迁移脚本
|
||||
|
||||
将 HuggingFace 缓存中的 MinerU 模型迁移到项目 models/ 目录
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
# 项目根目录（脚本在 scripts/ 目录下，需要回到上级）
|
||||
PROJECT_ROOT = Path(__file__).parent.parent
|
||||
MODELS_DIR = PROJECT_ROOT / "models" / "mineru"
|
||||
|
||||
# 当前配置文件路径
|
||||
USER_CONFIG = Path.home() / "mineru.json"
|
||||
|
||||
def read_current_config():
|
||||
"""读取当前配置"""
|
||||
if not USER_CONFIG.exists():
|
||||
print(f"❌ 配置文件不存在: {USER_CONFIG}")
|
||||
return None
|
||||
|
||||
with open(USER_CONFIG, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
|
||||
def migrate_models():
|
||||
"""迁移模型文件"""
|
||||
print("=" * 60)
|
||||
print("MinerU 模型迁移工具")
|
||||
print("=" * 60)
|
||||
|
||||
# 读取当前配置
|
||||
config = read_current_config()
|
||||
if not config:
|
||||
return False
|
||||
|
||||
models_dir_config = config.get('models-dir', {})
|
||||
if not models_dir_config:
|
||||
print("❌ 配置文件中没有 models-dir 配置")
|
||||
return False
|
||||
|
||||
print(f"\n📂 当前模型路径:")
|
||||
for model_type, model_path in models_dir_config.items():
|
||||
print(f" {model_type}: {model_path}")
|
||||
|
||||
# 创建目标目录
|
||||
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 迁移每个模型
|
||||
new_config = {}
|
||||
for model_type, src_path in models_dir_config.items():
|
||||
src_path = Path(src_path)
|
||||
|
||||
if not src_path.exists():
|
||||
print(f"\n⚠️ 源路径不存在，跳过: {src_path}")
|
||||
continue
|
||||
|
||||
# 目标路径
|
||||
dst_path = MODELS_DIR / model_type
|
||||
|
||||
print(f"\n📦 迁移 {model_type} 模型...")
|
||||
print(f" 源: {src_path}")
|
||||
print(f" 目标: {dst_path}")
|
||||
|
||||
# 检查目标是否已存在
|
||||
if dst_path.exists():
|
||||
print(f" ⚠️ 目标已存在，是否覆盖？(y/n): ", end='')
|
||||
choice = input().strip().lower()
|
||||
if choice != 'y':
|
||||
print(f" ⏭️ 跳过")
|
||||
new_config[model_type] = str(dst_path.absolute())
|
||||
continue
|
||||
else:
|
||||
shutil.rmtree(dst_path)
|
||||
|
||||
# 复制模型文件
|
||||
try:
|
||||
shutil.copytree(src_path, dst_path)
|
||||
print(f" ✅ 迁移成功")
|
||||
new_config[model_type] = str(dst_path.absolute())
|
||||
except Exception as e:
|
||||
print(f" ❌ 迁移失败: {e}")
|
||||
return False
|
||||
|
||||
if not new_config:
|
||||
print("\n❌ 没有成功迁移任何模型")
|
||||
return False
|
||||
|
||||
# 更新配置文件
|
||||
print(f"\n📝 更新配置文件...")
|
||||
|
||||
# 更新用户目录配置
|
||||
config['models-dir'] = new_config
|
||||
with open(USER_CONFIG, 'w', encoding='utf-8') as f:
|
||||
json.dump(config, f, indent=4, ensure_ascii=False)
|
||||
print(f" ✅ 已更新: {USER_CONFIG}")
|
||||
|
||||
# 创建项目配置文件（使用相对路径）
|
||||
project_config_path = PROJECT_ROOT / "mineru.json"
|
||||
project_config = {
|
||||
"models-dir": {
|
||||
model_type: f"models/mineru/{model_type}"
|
||||
for model_type in new_config.keys()
|
||||
},
|
||||
"config_version": config.get("config_version", "1.3.1")
|
||||
}
|
||||
|
||||
with open(project_config_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(project_config, f, indent=4, ensure_ascii=False)
|
||||
print(f" ✅ 已创建: {project_config_path}")
|
||||
|
||||
# 显示新配置
|
||||
print(f"\n✅ 迁移完成！新的模型路径:")
|
||||
for model_type, model_path in new_config.items():
|
||||
print(f" {model_type}: {model_path}")
|
||||
|
||||
# 计算模型大小
|
||||
total_size = 0
|
||||
for model_type in new_config.keys():
|
||||
model_path = MODELS_DIR / model_type
|
||||
if model_path.exists():
|
||||
size = sum(f.stat().st_size for f in model_path.rglob('*') if f.is_file())
|
||||
total_size += size
|
||||
print(f" {model_type} 大小: {size / 1024 / 1024:.1f} MB")
|
||||
|
||||
print(f"\n📊 总大小: {total_size / 1024 / 1024:.1f} MB ({total_size / 1024 / 1024 / 1024:.2f} GB)")
|
||||
|
||||
return True
|
||||
|
||||
def verify_migration():
|
||||
"""验证迁移结果"""
|
||||
print("\n" + "=" * 60)
|
||||
print("验证迁移结果")
|
||||
print("=" * 60)
|
||||
|
||||
# 检查项目配置文件
|
||||
project_config_path = PROJECT_ROOT / "mineru.json"
|
||||
if not project_config_path.exists():
|
||||
print("❌ 项目配置文件不存在")
|
||||
return False
|
||||
|
||||
with open(project_config_path, 'r', encoding='utf-8') as f:
|
||||
config = json.load(f)
|
||||
|
||||
models_dir_config = config.get('models-dir', {})
|
||||
|
||||
all_ok = True
|
||||
for model_type, model_path in models_dir_config.items():
|
||||
model_path = Path(model_path)
|
||||
if model_path.exists():
|
||||
print(f"✅ {model_type}: {model_path}")
|
||||
else:
|
||||
print(f"❌ {model_type}: {model_path} (不存在)")
|
||||
all_ok = False
|
||||
|
||||
return all_ok
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
# Windows 控制台编码
|
||||
if sys.platform == 'win32':
|
||||
import locale
|
||||
if sys.stdout.encoding != 'utf-8':
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
success = migrate_models()
|
||||
|
||||
if success:
|
||||
verify_migration()
|
||||
print("\n🎉 迁移完成！现在可以删除 HuggingFace 缓存中的模型以节省空间。")
|
||||
print(f"\n💡 提示：")
|
||||
print(f" 1. 项目配置文件: {PROJECT_ROOT / 'mineru.json'}")
|
||||
print(f" 2. 模型目录: {MODELS_DIR}")
|
||||
print(f" 3. 环境变量: MINERU_MODEL_SOURCE=local")
|
||||
else:
|
||||
print("\n❌ 迁移失败")
|
||||
sys.exit(1)
|
||||
151
scripts/migrate_version_status.py
Normal file
151
scripts/migrate_version_status.py
Normal file
@@ -0,0 +1,151 @@
|
||||
"""
|
||||
数据迁移脚本 - 为现有chunks添加版本状态字段
|
||||
|
||||
运行此脚本为现有向量库数据添加status、version等字段
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Windows 编码设置
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
# 项目路径
|
||||
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
os.chdir(PROJECT_ROOT)
|
||||
sys.path.insert(0, PROJECT_ROOT)
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
def migrate_chunks():
|
||||
"""为现有chunks添加status字段"""
|
||||
from knowledge_base_manager import get_kb_manager
|
||||
|
||||
kb_manager = get_kb_manager()
|
||||
collections = kb_manager.list_collections()
|
||||
|
||||
print("=" * 60)
|
||||
print("开始迁移 - 为chunks添加版本状态字段")
|
||||
print("=" * 60)
|
||||
|
||||
total_migrated = 0
|
||||
|
||||
for coll_info in collections:
|
||||
coll_name = coll_info.name
|
||||
print(f"\n处理向量库: {coll_name}")
|
||||
|
||||
collection = kb_manager.get_collection(coll_name)
|
||||
if not collection:
|
||||
continue
|
||||
|
||||
result = collection.get()
|
||||
ids = result['ids']
|
||||
metadatas = result['metadatas']
|
||||
|
||||
if not ids:
|
||||
print(f" 向量库为空，跳过")
|
||||
continue
|
||||
|
||||
migrated = 0
|
||||
updated_metadatas = []
|
||||
|
||||
for i, (chunk_id, meta) in enumerate(zip(ids, metadatas)):
|
||||
# 检查是否已有status字段
|
||||
if 'status' not in meta:
|
||||
# 添加默认字段
|
||||
updated_meta = {
|
||||
**meta,
|
||||
'status': 'active',
|
||||
'version': 'v1',
|
||||
'effective_date': datetime.now().strftime('%Y-%m-%d')
|
||||
}
|
||||
updated_metadatas.append(updated_meta)
|
||||
migrated += 1
|
||||
else:
|
||||
updated_metadatas.append(meta)
|
||||
|
||||
# 批量更新
|
||||
if migrated > 0:
|
||||
collection.update(
|
||||
ids=ids,
|
||||
metadatas=updated_metadatas
|
||||
)
|
||||
print(f" 迁移了 {migrated} 个chunks")
|
||||
total_migrated += migrated
|
||||
else:
|
||||
print(f" 所有chunks已有status字段，无需迁移")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(f"迁移完成！总计迁移 {total_migrated} 个chunks")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
def init_version_tables():
|
||||
"""初始化版本管理数据库表"""
|
||||
import sqlite3
|
||||
|
||||
db_path = "./data/exam_analysis.db"
|
||||
|
||||
print("\n初始化版本管理表...")
|
||||
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 文档版本表
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS document_versions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
document_id TEXT NOT NULL,
|
||||
collection TEXT NOT NULL,
|
||||
version TEXT NOT NULL DEFAULT 'v1',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
effective_date DATE,
|
||||
expiry_date DATE,
|
||||
deprecated_date DATETIME,
|
||||
deprecated_reason TEXT,
|
||||
deprecated_by TEXT,
|
||||
change_summary TEXT,
|
||||
supersedes TEXT,
|
||||
chunk_count INTEGER DEFAULT 0,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
created_by TEXT,
|
||||
UNIQUE(document_id, collection, version)
|
||||
)
|
||||
''')
|
||||
|
||||
# 版本变更日志表
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS version_change_logs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
document_id TEXT NOT NULL,
|
||||
collection TEXT NOT NULL,
|
||||
old_version TEXT,
|
||||
new_version TEXT,
|
||||
old_status TEXT,
|
||||
new_status TEXT,
|
||||
change_type TEXT NOT NULL,
|
||||
reason TEXT,
|
||||
changed_by TEXT,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
''')
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
print("版本管理表初始化完成")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("向量知识库版本管理数据迁移")
|
||||
print()
|
||||
|
||||
# 1. 初始化数据库表
|
||||
init_version_tables()
|
||||
|
||||
# 2. 迁移chunks元数据
|
||||
migrate_chunks()
|
||||
|
||||
print("\n迁移脚本执行完毕！")
|
||||
266
scripts/rebuild_multi_kb.py
Normal file
266
scripts/rebuild_multi_kb.py
Normal file
@@ -0,0 +1,266 @@
|
||||
"""
|
||||
重建多向量知识库脚本（v5 统一解析版）
|
||||
|
||||
将现有文档按部门/类别分配到不同的向量库中
|
||||
使用统一的 parse_document() 入口
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Windows 编码设置
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
# 项目路径
|
||||
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
os.chdir(PROJECT_ROOT)
|
||||
sys.path.insert(0, PROJECT_ROOT)
|
||||
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
from config import DOCUMENTS_PATH, EMBEDDING_MODEL_PATH
|
||||
from parsers import parse_document, convert_to_rag_format, SUPPORTED_FORMATS
|
||||
from knowledge.manager import KnowledgeBaseManager, PUBLIC_KB_NAME
|
||||
|
||||
|
||||
def get_target_kb(filepath: str) -> str:
|
||||
"""
|
||||
判断文档应归属的向量库
|
||||
|
||||
根据文件所在目录判断：
|
||||
- public/ -> public_kb
|
||||
- dept_finance/ -> dept_finance
|
||||
- dept_hr/ -> dept_hr
|
||||
等
|
||||
"""
|
||||
# 标准化路径
|
||||
filepath_lower = filepath.replace('\\', '/').lower()
|
||||
|
||||
# 1. 根据目录路径判断（优先级最高）
|
||||
if 'public/' in filepath_lower or filepath_lower.startswith('public'):
|
||||
return PUBLIC_KB_NAME
|
||||
|
||||
# 检查是否在部门目录下
|
||||
for dept in ['finance', 'hr', 'tech', 'admin', 'operation', 'legal', 'strategy', 'marketing']:
|
||||
if f'dept_{dept}/' in filepath_lower or filepath_lower.startswith(f'dept_{dept}'):
|
||||
return f'dept_{dept}'
|
||||
|
||||
# 默认放入 public
|
||||
return PUBLIC_KB_NAME
|
||||
|
||||
|
||||
def scan_documents(documents_path: str) -> list:
|
||||
"""扫描文档目录"""
|
||||
documents = []
|
||||
|
||||
for root, dirs, files in os.walk(documents_path):
|
||||
for filename in files:
|
||||
ext = os.path.splitext(filename)[1].lower()
|
||||
if ext in SUPPORTED_FORMATS:
|
||||
filepath = os.path.join(root, filename)
|
||||
relpath = os.path.relpath(filepath, documents_path)
|
||||
documents.append({
|
||||
'filepath': filepath,
|
||||
'relpath': relpath,
|
||||
'filename': filename,
|
||||
'ext': ext
|
||||
})
|
||||
|
||||
return documents
|
||||
|
||||
|
||||
def process_document(doc_info: dict, embedding_model) -> tuple:
|
||||
"""
|
||||
处理单个文档，返回 (target_kb, chunks)
|
||||
|
||||
使用统一的 parse_document() 入口
|
||||
|
||||
Returns:
|
||||
(目标向量库名, [(text, metadata), ...])
|
||||
"""
|
||||
filepath = doc_info['filepath']
|
||||
filename = doc_info['filename']
|
||||
relpath = doc_info['relpath']
|
||||
|
||||
# 确定目标向量库
|
||||
target_kb = get_target_kb(relpath)
|
||||
|
||||
chunks = []
|
||||
|
||||
try:
|
||||
# 使用统一解析入口
|
||||
parse_result = parse_document(
|
||||
filepath,
|
||||
output_base=".data/mineru_temp",
|
||||
images_output=".data/images",
|
||||
cleanup_after_image_move=True
|
||||
)
|
||||
|
||||
# 转换为 RAG 格式
|
||||
pages_content = convert_to_rag_format(parse_result)
|
||||
raw_chunks = parse_result.get('chunks', [])
|
||||
|
||||
for i, (page_info, chunk) in enumerate(zip(pages_content, raw_chunks)):
|
||||
text = chunk.content if hasattr(chunk, 'content') else page_info.get('text', '')
|
||||
if not text.strip():
|
||||
continue
|
||||
|
||||
# 构建元数据
|
||||
metadata = {
|
||||
'source': filename,
|
||||
'page': page_info.get('page', 0),
|
||||
'chunk_index': i,
|
||||
'chunk_type': page_info.get('chunk_type', 'text'),
|
||||
'has_table': page_info.get('chunk_type') == 'table',
|
||||
'section': page_info.get('section_path', '') or page_info.get('section', ''),
|
||||
'collection': target_kb,
|
||||
'status': 'active'
|
||||
}
|
||||
|
||||
# 图片信息
|
||||
if hasattr(chunk, 'image_path') and chunk.image_path:
|
||||
import json
|
||||
metadata['images_json'] = json.dumps([{'id': chunk.image_path}], ensure_ascii=False)
|
||||
|
||||
chunks.append((text, metadata))
|
||||
|
||||
except Exception as e:
|
||||
print(f" 解析错误 {filename}: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
return target_kb, chunks
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("重建多向量知识库（v5 统一解析版）")
|
||||
print("=" * 60)
|
||||
|
||||
# 0. 清理现有向量库（避免重复数据）
|
||||
print("\n[0/5] 清理现有向量库...")
|
||||
import shutil
|
||||
vector_store_path = os.path.join(PROJECT_ROOT, "knowledge", "vector_store")
|
||||
if os.path.exists(vector_store_path):
|
||||
shutil.rmtree(vector_store_path)
|
||||
print(" [OK] 已删除旧向量库")
|
||||
else:
|
||||
print(" [-] 无旧向量库需要清理")
|
||||
|
||||
# 1. 加载向量模型
|
||||
print("\n[1/6] 加载向量模型...")
|
||||
embedding_model = SentenceTransformer(EMBEDDING_MODEL_PATH)
|
||||
print(" [OK] 向量模型加载完成")
|
||||
|
||||
# 2. 初始化知识库管理器
|
||||
print("\n[2/6] 初始化知识库管理器...")
|
||||
kb_manager = KnowledgeBaseManager()
|
||||
print(" [OK] 知识库管理器初始化完成")
|
||||
|
||||
# 3. 创建向量库
|
||||
print("\n[3/6] 创建向量库...")
|
||||
collections_to_create = [
|
||||
(PUBLIC_KB_NAME, '公开知识库', '全公司公开文档'),
|
||||
('dept_finance', '财务部知识库', '财务部专属文档'),
|
||||
('dept_hr', '人事部知识库', '人事部专属文档'),
|
||||
('dept_tech', '技术部知识库', '技术部专属文档'),
|
||||
('dept_admin', '行政部知识库', '行政部专属文档'),
|
||||
('dept_operation', '运营部知识库', '运营部专属文档'),
|
||||
('dept_legal', '法务部知识库', '法务部专属文档'),
|
||||
('dept_strategy', '战略部知识库', '战略部专属文档'),
|
||||
('dept_marketing', '市场部知识库', '市场部专属文档'),
|
||||
]
|
||||
|
||||
for name, display_name, desc in collections_to_create:
|
||||
success, msg = kb_manager.create_collection(name, display_name=display_name, description=desc)
|
||||
if success:
|
||||
print(f" [OK] 创建: {name}")
|
||||
else:
|
||||
print(f" [-] {name}: {msg}")
|
||||
|
||||
# 4. 扫描文档
|
||||
print("\n[4/6] 扫描文档...")
|
||||
documents = scan_documents(DOCUMENTS_PATH)
|
||||
print(f" 共发现 {len(documents)} 个文档")
|
||||
print(f" 支持的格式: {', '.join(SUPPORTED_FORMATS.keys())}")
|
||||
|
||||
# 5. 向量化并写入
|
||||
print("\n[5/6] 向量化并写入向量库...")
|
||||
stats = {}
|
||||
total_chunks = 0
|
||||
BATCH_SIZE = 100 # 每批写入数量
|
||||
|
||||
for i, doc in enumerate(documents):
|
||||
target_kb, chunks = process_document(doc, embedding_model)
|
||||
|
||||
if not chunks:
|
||||
continue
|
||||
|
||||
# 生成向量
|
||||
texts = [c[0] for c in chunks]
|
||||
metadatas = [c[1] for c in chunks]
|
||||
vectors = embedding_model.encode(texts).tolist()
|
||||
|
||||
# 生成 ID
|
||||
ids = [f'{doc["filename"]}_{j}' for j in range(len(texts))]
|
||||
|
||||
# 分批写入
|
||||
try:
|
||||
collection = kb_manager.get_collection(target_kb)
|
||||
if collection:
|
||||
for batch_start in range(0, len(ids), BATCH_SIZE):
|
||||
batch_end = min(batch_start + BATCH_SIZE, len(ids))
|
||||
collection.add(
|
||||
ids=ids[batch_start:batch_end],
|
||||
documents=texts[batch_start:batch_end],
|
||||
embeddings=vectors[batch_start:batch_end],
|
||||
metadatas=metadatas[batch_start:batch_end]
|
||||
)
|
||||
|
||||
stats[target_kb] = stats.get(target_kb, 0) + len(chunks)
|
||||
total_chunks += len(chunks)
|
||||
|
||||
if (i + 1) % 10 == 0:
|
||||
print(f" 已处理 {i + 1}/{len(documents)} 文档, 累计 {total_chunks} chunks")
|
||||
except Exception as e:
|
||||
print(f" [X] {doc['filename'][:30]}... 写入失败: {e}")
|
||||
|
||||
# 重建 BM25 索引
|
||||
print("\n重建 BM25 索引...")
|
||||
for kb_name in stats.keys():
|
||||
try:
|
||||
kb_manager._rebuild_bm25_index(kb_name)
|
||||
print(f" [OK] {kb_name} BM25 索引完成")
|
||||
except Exception as e:
|
||||
print(f" [!] {kb_name} BM25 索引失败: {e}")
|
||||
|
||||
# 汇总
|
||||
print("\n" + "=" * 60)
|
||||
print("重建完成")
|
||||
print("=" * 60)
|
||||
print(f"总文档数: {len(documents)}")
|
||||
print(f"总 chunks: {total_chunks}")
|
||||
print("\n各向量库统计:")
|
||||
for kb, count in sorted(stats.items()):
|
||||
print(f" {kb}: {count} chunks")
|
||||
|
||||
# 验证
|
||||
print("\n向量库列表:")
|
||||
for coll in kb_manager.list_collections():
|
||||
doc_count = coll.document_count
|
||||
print(f" {coll.name}: {doc_count} documents")
|
||||
|
||||
# 显式强制回收
|
||||
import gc
|
||||
import time
|
||||
print("\n等待底层向量引擎写入...")
|
||||
time.sleep(10)
|
||||
kb_manager._clients.clear()
|
||||
del kb_manager
|
||||
gc.collect()
|
||||
print("完成！")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
194
scripts/test_all_52_endpoints.sh
Normal file
194
scripts/test_all_52_endpoints.sh
Normal file
@@ -0,0 +1,194 @@
|
||||
#!/bin/bash
|
||||
# RAG API 完整测试脚本 - 覆盖 curl测试手册.md 所有 52 个接口
|
||||
|
||||
set -e
|
||||
BASE_URL="http://localhost:5001"
|
||||
PASS=0
|
||||
FAIL=0
|
||||
RESULTS=""
|
||||
|
||||
log_result() {
|
||||
local name="$1"
|
||||
local status="$2"
|
||||
if [ "$status" = "PASS" ]; then
|
||||
PASS=$((PASS + 1))
|
||||
echo "✅ $name"
|
||||
else
|
||||
FAIL=$((FAIL + 1))
|
||||
echo "❌ $name"
|
||||
fi
|
||||
}
|
||||
|
||||
test_endpoint() {
|
||||
local name="$1"
|
||||
local method="$2"
|
||||
local path="$3"
|
||||
local data="$4"
|
||||
local check="$5"
|
||||
|
||||
if [ "$method" = "GET" ]; then
|
||||
result=$(curl -s "$BASE_URL$path" 2>&1)
|
||||
else
|
||||
result=$(curl -s -X $method "$BASE_URL$path" -H "Content-Type: application/json" -d "$data" 2>&1)
|
||||
fi
|
||||
|
||||
if echo "$result" | grep -q "$check"; then
|
||||
log_result "$name" "PASS"
|
||||
else
|
||||
log_result "$name" "FAIL"
|
||||
fi
|
||||
}
|
||||
|
||||
echo "=========================================="
|
||||
echo "RAG API 完整测试 (52 接口)"
|
||||
echo "=========================================="
|
||||
|
||||
# ===== 2. 健康检查 (2个) =====
|
||||
echo -e "\n--- 2. 健康检查 ---"
|
||||
test_endpoint "GET /health" "GET" "/health" "" '"status"'
|
||||
curl -s "$BASE_URL/health" > /dev/null && log_result "GET /health (HTTP)" "PASS" || log_result "GET /health (HTTP)" "FAIL"
|
||||
|
||||
# ===== 3. 问答接口 (2个) =====
|
||||
echo -e "\n--- 3. 问答接口 ---"
|
||||
echo '{"message":"三峡工程","chat_history":[]}' > /tmp/rag.json
|
||||
curl -s -X POST "$BASE_URL/rag" -H "Content-Type: application/json" -d @/tmp/rag.json 2>&1 | grep -q '"type":"finish"' && log_result "POST /rag" "PASS" || log_result "POST /rag" "FAIL"
|
||||
|
||||
echo '{"message":"你好","chat_history":[]}' > /tmp/chat.json
|
||||
test_endpoint "POST /chat" "POST" "/chat" '{"message":"你好","chat_history":[]}' '"answer"'
|
||||
|
||||
# ===== 4. 检索接口 (1个) =====
|
||||
echo -e "\n--- 4. 检索接口 ---"
|
||||
test_endpoint "POST /search" "POST" "/search" '{"query":"三峡工程","top_k":3}' '"contexts"'
|
||||
|
||||
# ===== 5. 向量库管理 (8个) =====
|
||||
echo -e "\n--- 5. 向量库管理 ---"
|
||||
test_endpoint "GET /collections" "GET" "/collections" "" '"collections"'
|
||||
test_endpoint "POST /collections" "POST" "/collections" '{"name":"test_api_kb","display_name":"测试库"}' '"success"'
|
||||
test_endpoint "PUT /collections/test_api_kb" "PUT" "/collections/test_api_kb" '{"display_name":"测试库-修改"}' '"success"'
|
||||
test_endpoint "GET /collections/public_kb/documents" "GET" "/collections/public_kb/documents" "" '"documents"'
|
||||
test_endpoint "GET /collections/public_kb/chunks" "GET" "/collections/public_kb/chunks?limit=2" "" '"chunks"'
|
||||
test_endpoint "GET /collections/public_kb/documents/versions" "GET" "/collections/public_kb/documents/%E4%B8%89%E5%B3%A1%E5%85%AC%E6%8A%A5_1-15%E9%A1%B5.pdf/versions" "" '"versions"'
|
||||
test_endpoint "POST /collections/public_kb/update-image-descriptions" "POST" "/collections/public_kb/update-image-descriptions" "" '"success"'
|
||||
curl -s -X DELETE "$BASE_URL/collections/test_api_kb" | grep -q '"success"' && log_result "DELETE /collections/test_api_kb" "PASS" || log_result "DELETE /collections/test_api_kb" "FAIL"
|
||||
|
||||
# ===== 6. 文档管理 (10个) =====
|
||||
echo -e "\n--- 6. 文档管理 ---"
|
||||
test_endpoint "GET /documents/list" "GET" "/documents/list?collection=public_kb&page=1&page_size=5" "" '"documents"'
|
||||
test_endpoint "GET /documents/status" "GET" "/documents/public_kb%2Ftest.txt/status" "" '"status"'
|
||||
|
||||
# 创建测试文件上传
|
||||
echo "测试文档内容" > /tmp/test_upload.txt
|
||||
result=$(curl -s -X POST "$BASE_URL/documents/upload" -F "file=@/tmp/test_upload.txt" -F "collection=public_kb")
|
||||
echo "$result" | grep -q '"success"' && log_result "POST /documents/upload" "PASS" || log_result "POST /documents/upload" "FAIL"
|
||||
|
||||
# 批量上传
|
||||
echo "文件1" > /tmp/file1.txt
|
||||
echo "文件2" > /tmp/file2.txt
|
||||
result=$(curl -s -X POST "$BASE_URL/documents/batch-upload" -F "files=@/tmp/file1.txt" -F "files=@/tmp/file2.txt" -F "collection=public_kb")
|
||||
echo "$result" | grep -q '"success_count"' && log_result "POST /documents/batch-upload" "PASS" || log_result "POST /documents/batch-upload" "FAIL"
|
||||
|
||||
# 更新文档
|
||||
echo "更新内容" > /tmp/update.txt
|
||||
result=$(curl -s -X PUT "$BASE_URL/documents/public_kb%2Ftest_upload.txt" -F "file=@/tmp/update.txt")
|
||||
echo "$result" | grep -q '"success"' && log_result "PUT /documents/<path>" "PASS" || log_result "PUT /documents/<path>" "FAIL"
|
||||
|
||||
# 删除文档
|
||||
curl -s -X DELETE "$BASE_URL/documents/public_kb%2Ftest_upload.txt" | grep -q '"success"' && log_result "DELETE /documents/<path>" "PASS" || log_result "DELETE /documents/<path>" "FAIL"
|
||||
|
||||
# 废弃/恢复
|
||||
test_endpoint "POST /documents/deprecate" "POST" "/collections/public_kb/documents/test.txt/deprecate" '' '"success"'
|
||||
test_endpoint "POST /documents/restore" "POST" "/collections/public_kb/documents/test.txt/restore" '' '"success"'
|
||||
|
||||
# ===== 7. 切片管理 (4个) =====
|
||||
echo -e "\n--- 7. 切片管理 ---"
|
||||
test_endpoint "GET /documents/chunks" "GET" "/documents/public_kb%2F%E4%B8%89%E5%B3%A1%E5%85%AC%E6%8A%A5_1-15%E9%A1%B5.pdf/chunks?page=1&page_size=2" "" '"chunks"'
|
||||
test_endpoint "POST /chunks" "POST" "/chunks" '{"collection":"public_kb","content":"测试切片内容"}' '"success"'
|
||||
|
||||
# 获取一个真实切片ID
|
||||
chunk_id=$(curl -s "$BASE_URL/collections/public_kb/chunks?limit=1" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
|
||||
if [ -n "$chunk_id" ]; then
|
||||
curl -s -X PUT "$BASE_URL/chunks/$chunk_id" -H "Content-Type: application/json" -d '{"collection":"public_kb","content":"更新内容"}' | grep -q '"success"' && log_result "PUT /chunks/<id>" "PASS" || log_result "PUT /chunks/<id>" "FAIL"
|
||||
else
|
||||
log_result "PUT /chunks/<id>" "FAIL (无切片)"
|
||||
fi
|
||||
|
||||
# ===== 8. 同步服务 (6个) =====
|
||||
echo -e "\n--- 8. 同步服务 ---"
|
||||
test_endpoint "POST /sync" "POST" "/sync" '' '"success"'
|
||||
test_endpoint "GET /sync/status" "GET" "/sync/status" "" '"enabled"'
|
||||
test_endpoint "GET /sync/history" "GET" "/sync/history" "" '"history"'
|
||||
test_endpoint "GET /sync/changes" "GET" "/sync/changes" "" '"changes"'
|
||||
test_endpoint "POST /sync/start" "POST" "/sync/start" '' '"success"'
|
||||
test_endpoint "POST /sync/stop" "POST" "/sync/stop" '' '"success"'
|
||||
|
||||
# ===== 9. 反馈系统 (5个) =====
|
||||
echo -e "\n--- 9. 反馈系统 ---"
|
||||
test_endpoint "POST /feedback" "POST" "/feedback" '{"session_id":"test","query":"测试问题","answer":"测试回答","rating":1}' '"success"'
|
||||
test_endpoint "GET /feedback/list" "GET" "/feedback/list?page=1&page_size=5" "" '"feedbacks"'
|
||||
test_endpoint "GET /feedback/stats" "GET" "/feedback/stats" "" '"stats"'
|
||||
test_endpoint "GET /feedback/bad-cases" "GET" "/feedback/bad-cases" "" '"bad_cases"'
|
||||
test_endpoint "GET /feedback/blacklist" "GET" "/feedback/blacklist" "" '"blacklist"'
|
||||
|
||||
# ===== 10. FAQ 管理 (7个) =====
|
||||
echo -e "\n--- 10. FAQ 管理 ---"
|
||||
test_endpoint "GET /faq" "GET" "/faq?page=1&page_size=5" "" '"faqs"'
|
||||
result=$(curl -s -X POST "$BASE_URL/faq" -H "Content-Type: application/json" -d '{"question":"测试FAQ问题","answer":"测试FAQ答案"}')
|
||||
echo "$result" | grep -q '"faq_id"' && log_result "POST /faq" "PASS" || log_result "POST /faq" "FAIL"
|
||||
faq_id=$(echo "$result" | grep -o '"faq_id":[0-9]*' | cut -d':' -f2)
|
||||
|
||||
if [ -n "$faq_id" ]; then
|
||||
curl -s -X PUT "$BASE_URL/faq/$faq_id" -H "Content-Type: application/json" -d '{"question":"更新问题","answer":"更新答案"}' | grep -q '"success"' && log_result "PUT /faq/<id>" "PASS" || log_result "PUT /faq/<id>" "FAIL"
|
||||
fi
|
||||
|
||||
test_endpoint "GET /faq/suggestions" "GET" "/faq/suggestions?page=1&page_size=5" "" '"suggestions"'
|
||||
test_endpoint "POST /faq/suggestions/approve" "POST" "/faq/suggestions/1/approve" '{}' '"success"'
|
||||
test_endpoint "POST /faq/suggestions/reject" "POST" "/faq/suggestions/1/reject" '{}' '"success"'
|
||||
|
||||
if [ -n "$faq_id" ]; then
|
||||
curl -s -X DELETE "$BASE_URL/faq/$faq_id" | grep -q '"success"' && log_result "DELETE /faq/<id>" "PASS" || log_result "DELETE /faq/<id>" "FAIL"
|
||||
fi
|
||||
|
||||
# ===== 11. 出题系统 (3个) =====
|
||||
echo -e "\n--- 11. 出题系统 ---"
|
||||
test_endpoint "GET /exam/health" "GET" "/exam/health" "" '"status"'
|
||||
test_endpoint "POST /exam/generate" "POST" "/exam/generate" '{"file_path":"public_kb/三峡公报_1-15页.pdf","collection":"public_kb","question_types":{"single_choice":1}}' '"questions"'
|
||||
test_endpoint "POST /exam/grade" "POST" "/exam/grade" '{"answers":[{"question_id":"q1","question_type":"single_choice","question_content":{"answer":"A"},"student_answer":"A","max_score":2}]}' '"results"'
|
||||
|
||||
# ===== 12. 图片服务 (4个) =====
|
||||
echo -e "\n--- 12. 图片服务 ---"
|
||||
test_endpoint "GET /images/list" "GET" "/images/list?limit=5" "" '"images"'
|
||||
|
||||
# 获取图片ID测试
|
||||
image_id=$(curl -s "$BASE_URL/images/list?limit=1" | grep -o '"image_id":"[^"]*"' | head -1 | cut -d'"' -f4)
|
||||
if [ -n "$image_id" ]; then
|
||||
http_code=$(curl -s -o /dev/null -w "%{http_code}" "$BASE_URL/images/$image_id.jpg")
|
||||
[ "$http_code" = "200" ] && log_result "GET /images/<id>" "PASS" || log_result "GET /images/<id>" "FAIL"
|
||||
|
||||
curl -s "$BASE_URL/images/$image_id/info" | grep -q '"image_id"' && log_result "GET /images/<id>/info" "PASS" || log_result "GET /images/<id>/info" "FAIL"
|
||||
else
|
||||
log_result "GET /images/<id>" "FAIL (无图片)"
|
||||
log_result "GET /images/<id>/info" "FAIL (无图片)"
|
||||
fi
|
||||
|
||||
test_endpoint "GET /images/stats" "GET" "/images/stats" "" '"total_images"'
|
||||
|
||||
# ===== 13. 报告服务 (2个) =====
|
||||
echo -e "\n--- 13. 报告服务 ---"
|
||||
test_endpoint "GET /reports/weekly" "GET" "/reports/weekly" "" '"report"'
|
||||
test_endpoint "GET /reports/monthly" "GET" "/reports/monthly" "" '"report"'
|
||||
|
||||
# ===== 14. 知识库路由 (1个) =====
|
||||
echo -e "\n--- 14. 知识库路由 ---"
|
||||
test_endpoint "POST /kb/route" "POST" "/kb/route" '{"query":"三峡工程"}' '"target_collections"'
|
||||
|
||||
# 输出结果
|
||||
echo -e "\n=========================================="
|
||||
echo "测试结果汇总"
|
||||
echo "=========================================="
|
||||
echo "总计: $PASS 通过, $FAIL 失败"
|
||||
echo "=========================================="
|
||||
|
||||
# 清理
|
||||
rm -f /tmp/rag.json /tmp/chat.json /tmp/test_upload.txt /tmp/file1.txt /tmp/file2.txt /tmp/update.txt
|
||||
|
||||
exit $FAIL
|
||||
210
scripts/test_all_endpoints.sh
Normal file
210
scripts/test_all_endpoints.sh
Normal file
@@ -0,0 +1,210 @@
|
||||
#!/bin/bash
|
||||
# RAG API 完整测试脚本
|
||||
# 基于 docs/curl测试手册.md
|
||||
|
||||
set -e
|
||||
BASE_URL="http://localhost:5001"
|
||||
PASS=0
|
||||
FAIL=0
|
||||
RESULTS=""
|
||||
|
||||
log_result() {
|
||||
local name="$1"
|
||||
local status="$2"
|
||||
if [ "$status" = "PASS" ]; then
|
||||
PASS=$((PASS + 1))
|
||||
RESULTS="${RESULTS}✅ ${name}\n"
|
||||
else
|
||||
FAIL=$((FAIL + 1))
|
||||
RESULTS="${RESULTS}❌ ${name}\n"
|
||||
fi
|
||||
}
|
||||
|
||||
echo "=========================================="
|
||||
echo "RAG API 完整测试"
|
||||
echo "=========================================="
|
||||
|
||||
# 1. 健康检查
|
||||
echo -e "\n--- 1. 健康检查 ---"
|
||||
result=$(curl -s "$BASE_URL/health")
|
||||
if echo "$result" | grep -q '"status":"ok"'; then
|
||||
log_result "GET /health" "PASS"
|
||||
else
|
||||
log_result "GET /health" "FAIL"
|
||||
fi
|
||||
|
||||
# 2. 问答接口
|
||||
echo -e "\n--- 2. 问答接口 ---"
|
||||
|
||||
# POST /rag (简化测试)
|
||||
echo '{"message":"三峡工程","chat_history":[]}' > /tmp/rag.json
|
||||
result=$(curl -s -X POST "$BASE_URL/rag" -H "Content-Type: application/json" -d @/tmp/rag.json 2>&1)
|
||||
if echo "$result" | grep -q '"type":"finish"'; then
|
||||
log_result "POST /rag" "PASS"
|
||||
else
|
||||
log_result "POST /rag" "FAIL"
|
||||
fi
|
||||
|
||||
# POST /chat
|
||||
echo '{"message":"你好","chat_history":[]}' > /tmp/chat.json
|
||||
result=$(curl -s -X POST "$BASE_URL/chat" -H "Content-Type: application/json" -d @/tmp/chat.json)
|
||||
if echo "$result" | grep -q '"answer"'; then
|
||||
log_result "POST /chat" "PASS"
|
||||
else
|
||||
log_result "POST /chat" "FAIL"
|
||||
fi
|
||||
|
||||
# 3. 检索接口
|
||||
echo -e "\n--- 3. 检索接口 ---"
|
||||
echo '{"query":"三峡工程","top_k":3}' > /tmp/search.json
|
||||
result=$(curl -s -X POST "$BASE_URL/search" -H "Content-Type: application/json" -d @/tmp/search.json)
|
||||
if echo "$result" | grep -q '"contexts"'; then
|
||||
log_result "POST /search" "PASS"
|
||||
else
|
||||
log_result "POST /search" "FAIL"
|
||||
fi
|
||||
|
||||
# 4. 向量库管理
|
||||
echo -e "\n--- 4. 向量库管理 ---"
|
||||
|
||||
# GET /collections
|
||||
result=$(curl -s "$BASE_URL/collections")
|
||||
if echo "$result" | grep -q '"collections"'; then
|
||||
log_result "GET /collections" "PASS"
|
||||
else
|
||||
log_result "GET /collections" "FAIL"
|
||||
fi
|
||||
|
||||
# 创建测试向量库
|
||||
echo '{"name":"test_kb_auto","display_name":"自动测试库"}' > /tmp/create_kb.json
|
||||
result=$(curl -s -X POST "$BASE_URL/collections" -H "Content-Type: application/json" -d @/tmp/create_kb.json)
|
||||
if echo "$result" | grep -q '"success":true'; then
|
||||
log_result "POST /collections (创建)" "PASS"
|
||||
else
|
||||
log_result "POST /collections (创建)" "FAIL"
|
||||
fi
|
||||
|
||||
# 修改向量库
|
||||
echo '{"display_name":"自动测试库-已修改"}' > /tmp/update_kb.json
|
||||
result=$(curl -s -X PUT "$BASE_URL/collections/test_kb_auto" -H "Content-Type: application/json" -d @/tmp/update_kb.json)
|
||||
if echo "$result" | grep -q '"success"'; then
|
||||
log_result "PUT /collections/test_kb_auto" "PASS"
|
||||
else
|
||||
log_result "PUT /collections/test_kb_auto" "FAIL"
|
||||
fi
|
||||
|
||||
# 删除测试向量库
|
||||
result=$(curl -s -X DELETE "$BASE_URL/collections/test_kb_auto")
|
||||
if echo "$result" | grep -q '"success"'; then
|
||||
log_result "DELETE /collections/test_kb_auto" "PASS"
|
||||
else
|
||||
log_result "DELETE /collections/test_kb_auto" "FAIL"
|
||||
fi
|
||||
|
||||
# GET /collections/<name>/documents
|
||||
result=$(curl -s "$BASE_URL/collections/public_kb/documents")
|
||||
if echo "$result" | grep -q '"documents"'; then
|
||||
log_result "GET /collections/public_kb/documents" "PASS"
|
||||
else
|
||||
log_result "GET /collections/public_kb/documents" "FAIL"
|
||||
fi
|
||||
|
||||
# 5. 文档管理
|
||||
echo -e "\n--- 5. 文档管理 ---"
|
||||
result=$(curl -s "$BASE_URL/documents/list?collection=public_kb&page=1&page_size=5")
|
||||
if echo "$result" | grep -q '"documents"'; then
|
||||
log_result "GET /documents/list" "PASS"
|
||||
else
|
||||
log_result "GET /documents/list" "FAIL"
|
||||
fi
|
||||
|
||||
# 6. 同步服务
|
||||
echo -e "\n--- 6. 同步服务 ---"
|
||||
result=$(curl -s "$BASE_URL/sync/status")
|
||||
if echo "$result" | grep -q '"enabled"'; then
|
||||
log_result "GET /sync/status" "PASS"
|
||||
else
|
||||
log_result "GET /sync/status" "FAIL"
|
||||
fi
|
||||
|
||||
# 7. 反馈系统
|
||||
echo -e "\n--- 7. 反馈系统 ---"
|
||||
result=$(curl -s "$BASE_URL/feedback/stats")
|
||||
if echo "$result" | grep -q '"stats"'; then
|
||||
log_result "GET /feedback/stats" "PASS"
|
||||
else
|
||||
log_result "GET /feedback/stats" "FAIL"
|
||||
fi
|
||||
|
||||
# 8. FAQ 管理
|
||||
echo -e "\n--- 8. FAQ 管理 ---"
|
||||
result=$(curl -s "$BASE_URL/faq?page=1&page_size=5")
|
||||
if echo "$result" | grep -q '"faqs"'; then
|
||||
log_result "GET /faq" "PASS"
|
||||
else
|
||||
log_result "GET /faq" "FAIL"
|
||||
fi
|
||||
|
||||
# 9. 出题系统
|
||||
echo -e "\n--- 9. 出题系统 ---"
|
||||
result=$(curl -s "$BASE_URL/exam/health")
|
||||
if echo "$result" | grep -q '"status":"ok"'; then
|
||||
log_result "GET /exam/health" "PASS"
|
||||
else
|
||||
log_result "GET /exam/health" "FAIL"
|
||||
fi
|
||||
|
||||
# 10. 图片服务
|
||||
echo -e "\n--- 10. 图片服务 ---"
|
||||
result=$(curl -s "$BASE_URL/images/list?limit=5")
|
||||
if echo "$result" | grep -q '"images"'; then
|
||||
log_result "GET /images/list" "PASS"
|
||||
else
|
||||
log_result "GET /images/list" "FAIL"
|
||||
fi
|
||||
|
||||
# 获取第一张图片ID并测试
|
||||
image_id=$(curl -s "$BASE_URL/images/list?limit=1" | grep -o '"image_id":"[^"]*"' | head -1 | cut -d'"' -f4)
|
||||
if [ -n "$image_id" ]; then
|
||||
http_code=$(curl -s -o /dev/null -w "%{http_code}" "$BASE_URL/images/$image_id")
|
||||
if [ "$http_code" = "200" ]; then
|
||||
log_result "GET /images/<id>" "PASS"
|
||||
else
|
||||
log_result "GET /images/<id>" "FAIL"
|
||||
fi
|
||||
else
|
||||
log_result "GET /images/<id>" "FAIL (无图片)"
|
||||
fi
|
||||
|
||||
# 11. 报告服务
|
||||
echo -e "\n--- 11. 报告服务 ---"
|
||||
result=$(curl -s "$BASE_URL/reports/weekly")
|
||||
if echo "$result" | grep -q '"report"'; then
|
||||
log_result "GET /reports/weekly" "PASS"
|
||||
else
|
||||
log_result "GET /reports/weekly" "FAIL"
|
||||
fi
|
||||
|
||||
# 12. 知识库路由
|
||||
echo -e "\n--- 12. 知识库路由 ---"
|
||||
echo '{"query":"三峡工程"}' > /tmp/route.json
|
||||
result=$(curl -s -X POST "$BASE_URL/kb/route" -H "Content-Type: application/json" -d @/tmp/route.json)
|
||||
if echo "$result" | grep -q '"target_collections"'; then
|
||||
log_result "POST /kb/route" "PASS"
|
||||
else
|
||||
log_result "POST /kb/route" "FAIL"
|
||||
fi
|
||||
|
||||
# 输出结果
|
||||
echo -e "\n=========================================="
|
||||
echo "测试结果汇总"
|
||||
echo "=========================================="
|
||||
echo -e "$RESULTS"
|
||||
echo "------------------------------------------"
|
||||
echo "总计: $PASS 通过, $FAIL 失败"
|
||||
echo "=========================================="
|
||||
|
||||
# 清理临时文件
|
||||
rm -f /tmp/rag.json /tmp/chat.json /tmp/search.json /tmp/create_kb.json /tmp/update_kb.json /tmp/route.json
|
||||
|
||||
exit $FAIL
|
||||
116
scripts/test_chunk_fixes.py
Normal file
116
scripts/test_chunk_fixes.py
Normal file
@@ -0,0 +1,116 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
验证切片质量修复效果
|
||||
|
||||
测试内容：
|
||||
1. _post_process_chunks: 空切片过滤、碎片合并、超长拆分
|
||||
2. _extract_table_title: 从表格内容提取标题
|
||||
3. _generate_table_summary: 不再产生 "小型表格：表格"
|
||||
"""
|
||||
import sys
|
||||
sys.path.insert(0, '.')
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
from parsers.mineru_parser import MinerUChunk, _post_process_chunks
|
||||
|
||||
print("=" * 60)
|
||||
print("测试 1: _post_process_chunks")
|
||||
print("=" * 60)
|
||||
|
||||
# 模拟实际碎片化数据
|
||||
test_chunks = [
|
||||
MinerUChunk(content="**货源投放工作规范", chunk_type="text", text_level=1, title="货源投放工作规范", section_path="货源投放工作规范", page_start=1, page_end=1),
|
||||
MinerUChunk(content="（2023版）", chunk_type="text", text_level=2, title="（2023版）", section_path="货源投放工作规范 > （2023版）", page_start=1, page_end=1),
|
||||
MinerUChunk(content="按照行业**营销市场化取向改革和全省系统**营销高质量发展的总体要求，结合全省**营销工作实际，特制定本规范。", chunk_type="text", text_level=0, page_start=1, page_end=1),
|
||||
MinerUChunk(content="**一、适用范围**", chunk_type="text", text_level=2, title="一、适用范围", section_path="一、适用范围", page_start=1, page_end=1),
|
||||
MinerUChunk(content="适用于全省地市公司**货源投放工作所涉及的基础工作、投放准备、投放规则、投放策略制定。", chunk_type="text", text_level=0, page_start=1, page_end=1),
|
||||
MinerUChunk(content="**二、总体要求**", chunk_type="text", text_level=2, title="二、总体要求", section_path="二、总体要求", page_start=1, page_end=1),
|
||||
MinerUChunk(content="货源投放是**营销的核心业务，总体要求是：", chunk_type="text", text_level=0, page_start=1, page_end=1),
|
||||
MinerUChunk(content="1.坚持市场导向、供需匹配；", chunk_type="text", text_level=0, page_start=1, page_end=1),
|
||||
MinerUChunk(content="2.坚持总量控制、稍紧平衡；", chunk_type="text", text_level=0, page_start=1, page_end=1),
|
||||
MinerUChunk(content="3.坚持增速合理、贵在持续；", chunk_type="text", text_level=0, page_start=1, page_end=1),
|
||||
MinerUChunk(content="4.坚持公平公正、严格规范；", chunk_type="text", text_level=0, page_start=1, page_end=1),
|
||||
MinerUChunk(content="5.坚持状态优先、科学投放；", chunk_type="text", text_level=0, page_start=1, page_end=1),
|
||||
MinerUChunk(content="6.坚持区域协同、高效运作。", chunk_type="text", text_level=0, page_start=1, page_end=1),
|
||||
# 空切片
|
||||
MinerUChunk(content="", chunk_type="text", text_level=0, page_start=2, page_end=2),
|
||||
MinerUChunk(content=" \n ", chunk_type="text", text_level=0, page_start=2, page_end=2),
|
||||
# 表格（不应被合并）
|
||||
MinerUChunk(content="表格内容", chunk_type="table", page_start=2, page_end=2, title="表格", table_html="<table>...</table>"),
|
||||
]
|
||||
|
||||
print(f"输入: {len(test_chunks)} 个切片")
|
||||
result = _post_process_chunks(test_chunks)
|
||||
print(f"输出: {len(result)} 个切片")
|
||||
print()
|
||||
|
||||
for i, chunk in enumerate(result):
|
||||
content_preview = chunk.content[:80].replace('\n', '\\n')
|
||||
print(f" [{i}] type={chunk.chunk_type}, level={chunk.text_level}, len={len(chunk.content)}")
|
||||
print(f" title={chunk.title}")
|
||||
print(f" content: {content_preview}...")
|
||||
print()
|
||||
|
||||
# 断言检查
|
||||
assert len(result) < len(test_chunks), f"合并后应减少切片数: {len(result)} >= {len(test_chunks)}"
|
||||
empty_chunks = [c for c in result if not c.content.strip()]
|
||||
assert len(empty_chunks) == 0, f"不应有空切片: 发现 {len(empty_chunks)} 个"
|
||||
table_chunks = [c for c in result if c.chunk_type == 'table']
|
||||
assert len(table_chunks) == 1, f"表格应保持独立: 发现 {len(table_chunks)} 个"
|
||||
print("✅ 碎片合并 + 空切片过滤 测试通过")
|
||||
|
||||
# 测试超长切片拆分
|
||||
print("\n" + "=" * 60)
|
||||
print("测试 2: 超长切片拆分")
|
||||
print("=" * 60)
|
||||
|
||||
long_text = "这是一段很长的测试文本。" * 200 # ~1200 chars
|
||||
long_chunks = [
|
||||
MinerUChunk(content=long_text, chunk_type="text", page_start=1, page_end=1),
|
||||
]
|
||||
|
||||
result2 = _post_process_chunks(long_chunks, max_chunk_size=500)
|
||||
print(f"输入: 1 个切片 ({len(long_text)} chars)")
|
||||
print(f"输出: {len(result2)} 个切片")
|
||||
for i, c in enumerate(result2):
|
||||
print(f" [{i}] len={len(c.content)}")
|
||||
assert all(len(c.content) <= 500 for c in result2), "拆分后不应超过 max_chunk_size"
|
||||
print("✅ 超长切片拆分 测试通过")
|
||||
|
||||
# 测试表格标题提取
|
||||
print("\n" + "=" * 60)
|
||||
print("测试 3: _extract_table_title")
|
||||
print("=" * 60)
|
||||
|
||||
sys.path.insert(0, '.')
|
||||
# 需要 import manager 但不初始化，直接测试静态方法
|
||||
from knowledge.manager import KnowledgeBaseManager
|
||||
|
||||
# Markdown 表头
|
||||
md_table = "【表格】表格\n\n| 序号 | 设备名称 | 数量 | 单价 |\n| --- | --- | --- | --- |\n| 1 | 电脑 | 10 | 5000 |"
|
||||
title = KnowledgeBaseManager._extract_table_title(md_table)
|
||||
print(f" Markdown 表头: '{title}'")
|
||||
assert '序号' in title, f"应包含列名: {title}"
|
||||
|
||||
# HTML <strong> 标签
|
||||
html_table = "<table><tr><td><strong>检查环节</strong></td><td><strong>检查项目</strong></td></tr></table>"
|
||||
title2 = KnowledgeBaseManager._extract_table_title(html_table)
|
||||
print(f" HTML strong: '{title2}'")
|
||||
assert '检查环节' in title2, f"应从 HTML 提取: {title2}"
|
||||
|
||||
# 【表格】带标题
|
||||
titled_table = "【表格】三峡上游主要水文站\n\n| 河流 | 站名 |\n| --- | --- |"
|
||||
title3 = KnowledgeBaseManager._extract_table_title(titled_table)
|
||||
print(f" 【表格】标题: '{title3}'")
|
||||
assert '三峡' in title3, f"应提取【表格】后标题: {title3}"
|
||||
|
||||
# 纯文本 "表格" 回退
|
||||
title4 = KnowledgeBaseManager._extract_table_title("表格")
|
||||
print(f" 纯'表格'回退: '{title4}'")
|
||||
assert title4 == "数据表格", f"应回退为'数据表格': {title4}"
|
||||
|
||||
print("✅ 表格标题提取 测试通过")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("所有测试通过! ✅")
|
||||
print("=" * 60)
|
||||
97
scripts/test_image_metadata.py
Normal file
97
scripts/test_image_metadata.py
Normal file
@@ -0,0 +1,97 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
阶段 1 验证脚本：检查向量库中的 metadata 是否包含图片信息
|
||||
|
||||
运行方式：
|
||||
.\venv\Scripts\python.exe scripts/test_image_metadata.py
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
# 添加项目根目录到路径
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from knowledge.manager import get_kb_manager
|
||||
|
||||
|
||||
def test_image_metadata():
|
||||
"""检查向量库中的 metadata 是否包含图片信息"""
|
||||
print("=" * 60)
|
||||
print("阶段 1 验证：检查 metadata 中的图片信息")
|
||||
print("=" * 60)
|
||||
|
||||
kb = get_kb_manager()
|
||||
coll = kb.get_collection('public_kb')
|
||||
|
||||
if not coll:
|
||||
print("❌ public_kb 不存在")
|
||||
return
|
||||
|
||||
# 获取所有切片
|
||||
result = coll.get(include=['metadatas', 'documents'])
|
||||
|
||||
if not result or not result.get('ids'):
|
||||
print("❌ 向量库为空，请先上传一个包含图片的 PDF")
|
||||
return
|
||||
|
||||
total = len(result['ids'])
|
||||
print(f"\n📊 向量库总切片数: {total}")
|
||||
|
||||
# 统计图片相关字段
|
||||
has_images_json = 0
|
||||
has_image_path = 0
|
||||
image_chunks = 0
|
||||
table_chunks = 0
|
||||
|
||||
print("\n📋 切片详情（前 20 个）:")
|
||||
print("-" * 60)
|
||||
|
||||
for i, (chunk_id, meta, doc) in enumerate(zip(result['ids'], result['metadatas'], result['documents'])):
|
||||
chunk_type = meta.get('chunk_type', 'text')
|
||||
images_json = meta.get('images_json')
|
||||
image_path = meta.get('image_path')
|
||||
|
||||
if images_json:
|
||||
has_images_json += 1
|
||||
if image_path:
|
||||
has_image_path += 1
|
||||
if chunk_type in ('image', 'chart'):
|
||||
image_chunks += 1
|
||||
if chunk_type == 'table':
|
||||
table_chunks += 1
|
||||
|
||||
# 打印前 20 个切片的详情
|
||||
if i < 20:
|
||||
print(f"[{i+1}] chunk_type: {chunk_type}")
|
||||
if images_json:
|
||||
print(f" images_json: {images_json[:100]}...")
|
||||
if image_path:
|
||||
print(f" image_path: {image_path}")
|
||||
print()
|
||||
|
||||
print("-" * 60)
|
||||
print("\n📈 统计结果:")
|
||||
print(f" - 包含 images_json 的切片: {has_images_json}")
|
||||
print(f" - 包含 image_path 的切片: {has_image_path}")
|
||||
print(f" - 图片类型切片 (image/chart): {image_chunks}")
|
||||
print(f" - 表格类型切片: {table_chunks}")
|
||||
|
||||
# 验证结果
|
||||
print("\n✅ 验证结果:")
|
||||
if has_images_json > 0:
|
||||
print(f" ✅ images_json 字段已正确写入 ({has_images_json} 个切片)")
|
||||
else:
|
||||
print(" ⚠️ images_json 字段未找到，可能需要重新上传 PDF")
|
||||
|
||||
if has_image_path > 0:
|
||||
print(f" ✅ image_path 字段已正确写入 ({has_image_path} 个切片)")
|
||||
else:
|
||||
print(" ⚠️ image_path 字段未找到，可能没有独立图片切片")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
test_image_metadata()
|
||||
255
scripts/test_image_pipeline.py
Normal file
255
scripts/test_image_pipeline.py
Normal file
@@ -0,0 +1,255 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
图片处理完整链路测试
|
||||
|
||||
测试文件: documents/public_kb/三峡公报_16-30页.pdf
|
||||
MinerU 临时输出: .data/mineru_temp
|
||||
|
||||
测试步骤:
|
||||
1. MinerU 解析 → 检查 chunk.images 和 chunk.image_path
|
||||
2. 同步到向量库 → 检查 metadata.images_json 和 metadata.image_path
|
||||
3. RAG 检索 → 检查图片召回结果
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
# 测试文件路径
|
||||
TEST_PDF = "documents/public_kb/三峡公报_16-30页.pdf"
|
||||
TEMP_DIR = ".data/mineru_temp"
|
||||
KB_NAME = "my_kb"
|
||||
|
||||
|
||||
def step1_mineru_parse():
|
||||
"""步骤1: MinerU 解析，检查图片信息"""
|
||||
print("\n" + "=" * 60)
|
||||
print("步骤1: MinerU 解析")
|
||||
print("=" * 60)
|
||||
|
||||
if not os.path.exists(TEST_PDF):
|
||||
print(f"❌ 测试文件不存在: {TEST_PDF}")
|
||||
return None
|
||||
|
||||
print(f"📄 解析文件: {TEST_PDF}")
|
||||
|
||||
# 清理临时目录
|
||||
if os.path.exists(TEMP_DIR):
|
||||
shutil.rmtree(TEMP_DIR)
|
||||
|
||||
# 调用 MinerU 解析
|
||||
from parsers.mineru_parser import parse_with_mineru_persistent
|
||||
|
||||
result = parse_with_mineru_persistent(
|
||||
TEST_PDF,
|
||||
output_base=TEMP_DIR,
|
||||
cleanup_after_image_move=False # 保留临时文件用于检查
|
||||
)
|
||||
|
||||
if not result:
|
||||
print("❌ MinerU 解析失败")
|
||||
return None
|
||||
|
||||
chunks = result.get('chunks', [])
|
||||
images = result.get('images', [])
|
||||
|
||||
print(f"\n📊 解析结果:")
|
||||
print(f" - 切片数量: {len(chunks)}")
|
||||
print(f" - 图片数量: {len(images)}")
|
||||
|
||||
# 检查 chunk.images 和 chunk.image_path
|
||||
chunks_with_images = 0
|
||||
chunks_with_image_path = 0
|
||||
image_chunks = 0
|
||||
|
||||
print(f"\n📋 切片图片信息（前 20 个）:")
|
||||
print("-" * 40)
|
||||
|
||||
for i, chunk in enumerate(chunks[:20]):
|
||||
chunk_type = getattr(chunk, 'chunk_type', 'text')
|
||||
has_images = hasattr(chunk, 'images') and chunk.images
|
||||
has_image_path = hasattr(chunk, 'image_path') and chunk.image_path
|
||||
|
||||
if has_images:
|
||||
chunks_with_images += 1
|
||||
if has_image_path:
|
||||
chunks_with_image_path += 1
|
||||
if chunk_type in ('image', 'chart'):
|
||||
image_chunks += 1
|
||||
|
||||
print(f"[{i+1}] type={chunk_type}")
|
||||
if has_images:
|
||||
print(f" images: {chunk.images}")
|
||||
if has_image_path:
|
||||
print(f" image_path: {chunk.image_path}")
|
||||
|
||||
print("-" * 40)
|
||||
print(f"\n📈 统计:")
|
||||
print(f" - 有 images 字段的切片: {chunks_with_images}")
|
||||
print(f" - 有 image_path 字段的切片: {chunks_with_image_path}")
|
||||
print(f" - 图片类型切片: {image_chunks}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def step2_sync_to_kb(result):
|
||||
"""步骤2: 同步到向量库，检查 metadata"""
|
||||
print("\n" + "=" * 60)
|
||||
print("步骤2: 同步到向量库")
|
||||
print("=" * 60)
|
||||
|
||||
from knowledge.manager import get_kb_manager
|
||||
|
||||
kb = get_kb_manager()
|
||||
|
||||
# 检查/创建知识库
|
||||
if KB_NAME not in kb.list_collections():
|
||||
kb.create_collection(KB_NAME, display_name="测试知识库")
|
||||
print(f"✅ 创建知识库: {KB_NAME}")
|
||||
else:
|
||||
# 清空现有数据
|
||||
coll = kb.get_collection(KB_NAME)
|
||||
if coll:
|
||||
existing = coll.get()
|
||||
if existing['ids']:
|
||||
coll.delete(ids=existing['ids'])
|
||||
print(f"🗑️ 清空知识库: {len(existing['ids'])} 条记录")
|
||||
|
||||
# 同步文件
|
||||
print(f"📤 同步文件: {TEST_PDF}")
|
||||
count = kb.add_file_to_kb(KB_NAME, TEST_PDF)
|
||||
print(f"✅ 同步完成: {count} 个切片")
|
||||
|
||||
# 检查 metadata
|
||||
coll = kb.get_collection(KB_NAME)
|
||||
data = coll.get(include=['metadatas', 'documents'])
|
||||
|
||||
has_images_json = 0
|
||||
has_image_path = 0
|
||||
|
||||
print(f"\n📋 Metadata 检查（前 20 个）:")
|
||||
print("-" * 40)
|
||||
|
||||
for i, meta in enumerate(data['metadatas'][:20]):
|
||||
images_json = meta.get('images_json')
|
||||
image_path = meta.get('image_path')
|
||||
|
||||
if images_json:
|
||||
has_images_json += 1
|
||||
if image_path:
|
||||
has_image_path += 1
|
||||
|
||||
print(f"[{i+1}] type={meta.get('chunk_type')}")
|
||||
if images_json:
|
||||
print(f" images_json: {images_json[:80]}...")
|
||||
if image_path:
|
||||
print(f" image_path: {image_path}")
|
||||
|
||||
print("-" * 40)
|
||||
print(f"\n📈 统计:")
|
||||
print(f" - 有 images_json 的切片: {has_images_json}")
|
||||
print(f" - 有 image_path 的切片: {has_image_path}")
|
||||
|
||||
# 验证结果
|
||||
print(f"\n✅ 验证结果:")
|
||||
if has_images_json > 0:
|
||||
print(f" ✅ images_json 字段正确写入")
|
||||
else:
|
||||
print(f" ❌ images_json 字段未写入")
|
||||
|
||||
if has_image_path > 0:
|
||||
print(f" ✅ image_path 字段正确写入")
|
||||
else:
|
||||
print(f" ⚠️ image_path 字段未写入（可能没有独立图片切片）")
|
||||
|
||||
return has_images_json > 0 or has_image_path > 0
|
||||
|
||||
|
||||
def step3_rag_recall():
|
||||
"""步骤3: RAG 检索，检查图片召回"""
|
||||
print("\n" + "=" * 60)
|
||||
print("步骤3: RAG 检索图片召回")
|
||||
print("=" * 60)
|
||||
|
||||
from core.engine import get_engine
|
||||
|
||||
engine = get_engine()
|
||||
if not engine._initialized:
|
||||
engine.initialize()
|
||||
|
||||
# 测试查询
|
||||
queries = [
|
||||
"三峡公报中有哪些图表？",
|
||||
"文档中的图片",
|
||||
"有哪些数据图表",
|
||||
]
|
||||
|
||||
for query in queries:
|
||||
print(f"\n🔍 查询: {query}")
|
||||
print("-" * 40)
|
||||
|
||||
result = engine.search_knowledge(
|
||||
query=query,
|
||||
top_k=10,
|
||||
collections=[KB_NAME]
|
||||
)
|
||||
|
||||
if not result or not result.get('ids'):
|
||||
print(" ❌ 无检索结果")
|
||||
continue
|
||||
|
||||
print(f" 📊 检索到 {len(result['ids'])} 个切片")
|
||||
|
||||
# 统计图片
|
||||
images_found = []
|
||||
for meta in result.get('metadatas', []):
|
||||
images_json = meta.get('images_json')
|
||||
image_path = meta.get('image_path')
|
||||
|
||||
if images_json:
|
||||
try:
|
||||
imgs = json.loads(images_json)
|
||||
images_found.extend(imgs)
|
||||
except:
|
||||
pass
|
||||
if image_path:
|
||||
images_found.append({'id': image_path})
|
||||
|
||||
if images_found:
|
||||
print(f" ✅ 召回 {len(images_found)} 张图片:")
|
||||
for img in images_found:
|
||||
print(f" 📷 {img.get('id', img)}")
|
||||
else:
|
||||
print(" ⚠️ 未召回图片")
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("图片处理完整链路测试")
|
||||
print("=" * 60)
|
||||
|
||||
# 步骤1: MinerU 解析
|
||||
result = step1_mineru_parse()
|
||||
if not result:
|
||||
return
|
||||
|
||||
# 步骤2: 同步到向量库
|
||||
success = step2_sync_to_kb(result)
|
||||
|
||||
# 步骤3: RAG 检索
|
||||
step3_rag_recall()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("测试完成")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
main()
|
||||
139
scripts/test_rag_image_recall.py
Normal file
139
scripts/test_rag_image_recall.py
Normal file
@@ -0,0 +1,139 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
RAG 接口测试：验证图片召回功能
|
||||
|
||||
运行方式：
|
||||
.\venv\Scripts\python.exe scripts/test_rag_image_recall.py
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from knowledge.manager import get_kb_manager
|
||||
from core.engine import get_engine
|
||||
|
||||
|
||||
def test_rag_image_recall():
|
||||
"""测试 RAG 检索的图片召回"""
|
||||
print("=" * 60)
|
||||
print("RAG 图片召回测试")
|
||||
print("=" * 60)
|
||||
|
||||
# 初始化引擎
|
||||
engine = get_engine()
|
||||
if not engine._initialized:
|
||||
engine.initialize()
|
||||
|
||||
# 测试查询
|
||||
test_queries = [
|
||||
"文档中有哪些图片？",
|
||||
"有哪些图表？",
|
||||
"表格内容是什么？",
|
||||
]
|
||||
|
||||
for query in test_queries:
|
||||
print(f"\n🔍 查询: {query}")
|
||||
print("-" * 40)
|
||||
|
||||
# 执行检索
|
||||
result = engine.search_knowledge(
|
||||
query=query,
|
||||
collection_name='public_kb',
|
||||
top_k=10
|
||||
)
|
||||
|
||||
if not result or not result.get('ids'):
|
||||
print(" ❌ 无检索结果")
|
||||
continue
|
||||
|
||||
print(f" 📊 检索到 {len(result['ids'])} 个切片")
|
||||
|
||||
# 检查图片信息
|
||||
images_found = []
|
||||
for i, meta in enumerate(result.get('metadatas', [])):
|
||||
images_json = meta.get('images_json')
|
||||
image_path = meta.get('image_path')
|
||||
chunk_type = meta.get('chunk_type')
|
||||
|
||||
if images_json:
|
||||
try:
|
||||
imgs = json.loads(images_json)
|
||||
images_found.extend(imgs)
|
||||
print(f" [{i+1}] chunk_type={chunk_type}, images_json={len(imgs)}张图片")
|
||||
except:
|
||||
pass
|
||||
|
||||
if image_path:
|
||||
images_found.append({'id': image_path, 'type': 'image_path'})
|
||||
print(f" [{i+1}] chunk_type={chunk_type}, image_path={image_path}")
|
||||
|
||||
if images_found:
|
||||
print(f"\n ✅ 共召回 {len(images_found)} 张图片")
|
||||
else:
|
||||
print("\n ⚠️ 未召回任何图片")
|
||||
|
||||
|
||||
def test_extract_rich_media():
|
||||
"""测试 _extract_rich_media 函数"""
|
||||
print("\n" + "=" * 60)
|
||||
print("_extract_rich_media 函数测试")
|
||||
print("=" * 60)
|
||||
|
||||
from core.agentic import AgenticRAG
|
||||
|
||||
# 创建 AgenticRAG 实例
|
||||
rag = AgenticRAG()
|
||||
|
||||
# 模拟检索结果
|
||||
mock_contexts = [
|
||||
{
|
||||
'doc': '这是一段文本',
|
||||
'meta': {
|
||||
'chunk_type': 'text',
|
||||
'source': 'test.pdf',
|
||||
'page': 1,
|
||||
}
|
||||
},
|
||||
{
|
||||
'doc': '这是一张图片',
|
||||
'meta': {
|
||||
'chunk_type': 'image',
|
||||
'source': 'test.pdf',
|
||||
'page': 2,
|
||||
'image_path': 'abc123.jpg',
|
||||
}
|
||||
},
|
||||
{
|
||||
'doc': '这是带关联图片的文本',
|
||||
'meta': {
|
||||
'chunk_type': 'text',
|
||||
'source': 'test.pdf',
|
||||
'page': 3,
|
||||
'images_json': '[{"id": "img1.jpg", "order": 1}, {"id": "img2.jpg", "order": 2}]',
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
# 调用 _extract_rich_media
|
||||
rich_media = rag._extract_rich_media(mock_contexts)
|
||||
|
||||
print(f"\n📊 提取结果:")
|
||||
print(f" - 图片数量: {len(rich_media.get('images', []))}")
|
||||
print(f" - 表格数量: {len(rich_media.get('tables', []))}")
|
||||
|
||||
for img in rich_media.get('images', []):
|
||||
print(f" 📷 {img}")
|
||||
|
||||
if not rich_media.get('images'):
|
||||
print(" ⚠️ 未提取到图片，检查 _extract_rich_media 逻辑")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
test_rag_image_recall()
|
||||
test_extract_rich_media()
|
||||
354
scripts/test_rag_questions.py
Normal file
354
scripts/test_rag_questions.py
Normal file
@@ -0,0 +1,354 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
RAG 系统测试脚本
|
||||
|
||||
自动执行 40 个测试问题并记录结果
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
|
||||
# 添加项目根目录到路径
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
# 测试问题列表
|
||||
TEST_QUESTIONS = [
|
||||
# 一、精确匹配测试（10题）
|
||||
{"id": 1, "category": "精确匹配", "question": "智启科技成立于哪一年？", "expected": "2015年3月", "source": "公司简介.txt / 员工手册.txt"},
|
||||
{"id": 2, "category": "精确匹配", "question": "公司的客服热线是多少？", "expected": "400-888-8888", "source": "公司简介.txt"},
|
||||
{"id": 3, "category": "精确匹配", "question": "公司总部位于哪里？", "expected": "北京市海淀区中关村大街1号科技大厦15层", "source": "公司简介.txt"},
|
||||
{"id": 4, "category": "精确匹配", "question": "ZDAP平台标准版支持多少并发用户？", "expected": "50人", "source": "产品手册.txt"},
|
||||
{"id": 5, "category": "精确匹配", "question": "企业版单表最大支持多少行数据？", "expected": "1亿行", "source": "产品手册.txt"},
|
||||
{"id": 6, "category": "精确匹配", "question": "年假满10年不满20年可以休多少天？", "expected": "10天", "source": "员工手册.txt / 常见问题.txt"},
|
||||
{"id": 7, "category": "精确匹配", "question": "产假可以休多少天？", "expected": "158天", "source": "员工手册.txt / 常见问题.txt"},
|
||||
{"id": 8, "category": "精确匹配", "question": "技术研发中心的负责人是谁？", "expected": "张明远（技术总监）", "source": "公司简介.txt / 组织架构说明.txt"},
|
||||
{"id": 9, "category": "精确匹配", "question": "市场营销部有多少人？", "expected": "120人", "source": "公司简介.txt / 组织架构说明.txt"},
|
||||
{"id": 10, "category": "精确匹配", "question": "上海分公司的负责人是谁？", "expected": "华东区总经理 张华", "source": "组织架构说明.txt"},
|
||||
|
||||
# 二、语义理解测试（6题）
|
||||
{"id": 11, "category": "语义理解", "question": "公司的愿景是什么？", "expected": "成为全球领先的智能数据服务提供商", "source": "公司简介.txt / 员工手册.txt"},
|
||||
{"id": 12, "category": "语义理解", "question": "ZDAP的智能预警功能有哪些通知方式？", "expected": "邮件、短信、企业微信、钉钉", "source": "产品手册.txt"},
|
||||
{"id": 13, "category": "语义理解", "question": "什么是直连数据集？", "expected": "直接查询源数据库，实时性强", "source": "产品手册.txt"},
|
||||
{"id": 14, "category": "语义理解", "question": "入职当天需要做什么？", "expected": "签订劳动合同、领取工牌、开通账号、参观、培训等", "source": "常见问题.txt"},
|
||||
{"id": 15, "category": "语义理解", "question": "请假4天需要谁审批？", "expected": "部门负责人 + 人力资源部审批", "source": "常见问题.txt"},
|
||||
{"id": 16, "category": "语义理解", "question": "如何申请外部培训？", "expected": "OA系统提交申请→填写信息→部门负责人审批→人力资源部审批→超5000元签服务协议", "source": "常见问题.txt"},
|
||||
|
||||
# 三、跨文档关联测试（4题）
|
||||
{"id": 17, "category": "跨文档关联", "question": "公司有哪些分公司，分别在哪些城市？", "expected": "上海、深圳、成都、武汉四家分公司", "source": "公司简介.txt + 组织架构说明.txt"},
|
||||
{"id": 18, "category": "跨文档关联", "question": "技术研发中心下设哪些团队，各自的职责是什么？", "expected": "AI算法组、数据工程组、平台开发组、前端开发组、测试组", "source": "公司简介.txt + 组织架构说明.txt"},
|
||||
{"id": 19, "category": "跨文档关联", "question": "公司的薪酬由哪些部分组成？绩效奖金的范围是多少？", "expected": "基本工资+绩效奖金+年终奖金+津贴补贴，绩效奖金0-30%基本工资", "source": "员工手册.txt + 常见问题.txt"},
|
||||
{"id": 20, "category": "跨文档关联", "question": "病假工资怎么算？不同工龄有什么区别？", "expected": "按工龄60%-100%发放", "source": "常见问题.txt"},
|
||||
|
||||
# 四、复杂推理测试（3题）
|
||||
{"id": 21, "category": "复杂推理", "question": "一个入职3年的员工，累计病假1个月，能拿到多少病假工资？", "expected": "工龄2-4年按基本工资70%发放", "source": "常见问题.txt"},
|
||||
{"id": 22, "category": "复杂推理", "question": "如果我绩效考核连续两个月是D级会怎样？", "expected": "进入绩效改进期(PIP)，1-3个月改进期，仍不达标可调岗或解除合同", "source": "常见问题.txt"},
|
||||
{"id": 23, "category": "复杂推理", "question": "ZDAP企业版的简单查询响应时间要求是多少？", "expected": "<1秒", "source": "产品手册.txt"},
|
||||
|
||||
# 五、表格数据测试（3题）
|
||||
{"id": 24, "category": "表格数据", "question": "P4级工程师的年薪范围是多少？", "expected": "25-35万元", "source": "组织架构说明.txt"},
|
||||
{"id": 25, "category": "表格数据", "question": "M3级经理对应的职称是什么？", "expected": "高级经理", "source": "组织架构说明.txt"},
|
||||
{"id": 26, "category": "表格数据", "question": "数据工程组有多少人？负责人是谁？", "expected": "80人，负责人王工", "source": "组织架构说明.txt"},
|
||||
|
||||
# 六、否定性测试（3题）
|
||||
{"id": 27, "category": "否定性测试", "question": "公司有员工宿舍吗？", "expected": "没有员工宿舍，但为外地新员工提供15天免费过渡住宿", "source": "常见问题.txt"},
|
||||
{"id": 28, "category": "否定性测试", "question": "公司的股票代码是什么？", "expected": "文档中未提及", "source": "无"},
|
||||
{"id": 29, "category": "否定性测试", "question": "公司有食堂吗？", "expected": "没有食堂，但提供早餐和午餐", "source": "常见问题.txt"},
|
||||
|
||||
# 七、关键词检索测试（3题）
|
||||
{"id": 30, "category": "关键词检索", "question": "五险一金缴纳比例是多少？", "expected": "养老保险个人8%公司16%，医疗保险个人2%公司10%等", "source": "常见问题.txt"},
|
||||
{"id": 31, "category": "关键词检索", "question": "Kong Gateway在系统架构中的作用是什么？", "expected": "API网关层，JWT认证、限流熔断", "source": "产品手册.txt"},
|
||||
{"id": 32, "category": "关键词检索", "question": "ClickHouse用于什么用途？", "expected": "分析引擎", "source": "产品手册.txt"},
|
||||
|
||||
# 八、长文档/分块测试（2题）
|
||||
{"id": 33, "category": "长文档测试", "question": "三峡工程2024年发电量是多少？", "expected": "需从三峡公报PDF中检索", "source": "三峡公报_*.pdf"},
|
||||
{"id": 34, "category": "长文档测试", "question": "三峡水库的水位调节范围是多少？", "expected": "需从三峡公报PDF中检索", "source": "三峡公报_*.pdf"},
|
||||
|
||||
# 九、学术论文测试（2题）
|
||||
{"id": 35, "category": "学术论文", "question": "这篇论文的主要贡献是什么？", "expected": "需从论文PDF中提取", "source": "2604.09205v1.pdf"},
|
||||
{"id": 36, "category": "学术论文", "question": "论文使用了什么方法或模型？", "expected": "需从论文PDF中提取", "source": "2604.09205v1.pdf"},
|
||||
|
||||
# 十、边缘案例测试（4题）
|
||||
{"id": 37, "category": "边缘案例", "question": "如果我要报销5000元以内的费用，需要谁审批？", "expected": "主管审批，备案即可", "source": "组织架构说明.txt"},
|
||||
{"id": 38, "category": "边缘案例", "question": "公司的核心工作时间是什么时候？", "expected": "10:00-16:00（必须到岗）", "source": "常见问题.txt"},
|
||||
{"id": 39, "category": "边缘案例", "question": "ZDAP支持哪些图表类型？", "expected": "柱状图、折线图、饼图、散点图、热力图、地图、雷达图、漏斗图", "source": "产品手册.txt"},
|
||||
{"id": 40, "category": "边缘案例", "question": "技术支持热线的服务时间是怎样的？", "expected": "工作日9:00-21:00，周末及节假日10:00-18:00", "source": "产品手册.txt"},
|
||||
]
|
||||
|
||||
|
||||
def call_rag_api(question: str, kb_name: str = "public", top_k: int = 5) -> dict:
|
||||
"""
|
||||
调用 RAG API 进行问答
|
||||
|
||||
Args:
|
||||
question: 问题内容
|
||||
kb_name: 知识库名称
|
||||
top_k: 返回结果数量
|
||||
|
||||
Returns:
|
||||
包含答案和检索上下文的字典
|
||||
"""
|
||||
import requests
|
||||
|
||||
url = "http://localhost:5001/rag"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": "Bearer mock-token-admin"
|
||||
}
|
||||
data = {
|
||||
"message": question # 使用 message 字段
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(url, json=data, headers=headers, timeout=120)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
return {
|
||||
"answer": result.get("answer"),
|
||||
"sources": result.get("sources", []),
|
||||
"contexts": [] # /rag 接口不返回 contexts，但有 sources
|
||||
}
|
||||
except requests.exceptions.RequestException as e:
|
||||
return {"error": str(e), "answer": None, "sources": [], "contexts": []}
|
||||
|
||||
|
||||
def evaluate_answer(actual: str, expected: str) -> dict:
|
||||
"""
|
||||
评估答案准确性
|
||||
|
||||
Args:
|
||||
actual: 实际答案
|
||||
expected: 预期答案
|
||||
|
||||
Returns:
|
||||
评估结果字典
|
||||
"""
|
||||
if not actual:
|
||||
return {"score": 0, "grade": "不合格", "reason": "未获取到答案", "matched_keywords": []}
|
||||
|
||||
actual_lower = actual.lower()
|
||||
expected_lower = expected.lower()
|
||||
|
||||
# 使用 jieba 分词（中文友好）
|
||||
try:
|
||||
import jieba
|
||||
expected_keywords = set(jieba.cut(expected_lower))
|
||||
actual_keywords = set(jieba.cut(actual_lower))
|
||||
except ImportError:
|
||||
# jieba 未安装时的降级方案
|
||||
expected_keywords = set(expected_lower.replace("，", " ").replace("、", " ").replace("。", " ").split())
|
||||
actual_keywords = set(actual_lower.replace("，", " ").replace("、", " ").replace("。", " ").split())
|
||||
|
||||
# 过滤停用词和短词
|
||||
stop_words = {"的", "是", "有", "在", "和", "与", "或", "等", "了", "到", "为", "按", "以", "可以", "需要", "应该"}
|
||||
expected_keywords = {w for w in expected_keywords if len(w) >= 2 and w not in stop_words}
|
||||
actual_keywords = {w for w in actual_keywords if len(w) >= 2 and w not in stop_words}
|
||||
|
||||
matched_keywords = expected_keywords & actual_keywords
|
||||
|
||||
if len(expected_keywords) == 0:
|
||||
keyword_score = 50
|
||||
else:
|
||||
keyword_score = int(len(matched_keywords) / len(expected_keywords) * 100)
|
||||
|
||||
# 完全匹配加分
|
||||
if expected_lower in actual_lower:
|
||||
keyword_score = min(100, keyword_score + 20)
|
||||
|
||||
# 评分等级
|
||||
if keyword_score >= 80:
|
||||
grade = "优秀"
|
||||
elif keyword_score >= 60:
|
||||
grade = "良好"
|
||||
elif keyword_score >= 40:
|
||||
grade = "合格"
|
||||
else:
|
||||
grade = "不合格"
|
||||
|
||||
return {
|
||||
"score": keyword_score,
|
||||
"grade": grade,
|
||||
"matched_keywords": list(matched_keywords),
|
||||
"reason": f"关键词匹配度 {keyword_score}%"
|
||||
}
|
||||
|
||||
|
||||
def run_tests():
|
||||
"""执行所有测试"""
|
||||
print("=" * 80)
|
||||
print("RAG 系统测试开始")
|
||||
print(f"测试时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
print(f"测试问题数量: {len(TEST_QUESTIONS)}")
|
||||
print("=" * 80)
|
||||
print()
|
||||
|
||||
results = []
|
||||
|
||||
for i, test in enumerate(TEST_QUESTIONS):
|
||||
print(f"[{i+1}/{len(TEST_QUESTIONS)}] 测试类别: {test['category']}")
|
||||
print(f"问题: {test['question']}")
|
||||
|
||||
# 调用 API
|
||||
start_time = time.time()
|
||||
response = call_rag_api(test['question'])
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
# 提取答案和上下文
|
||||
answer = response.get("answer", "")
|
||||
sources = response.get("sources", [])
|
||||
|
||||
# 评估答案
|
||||
evaluation = evaluate_answer(answer, test['expected'])
|
||||
|
||||
result = {
|
||||
"id": test['id'],
|
||||
"category": test['category'],
|
||||
"question": test['question'],
|
||||
"expected": test['expected'],
|
||||
"expected_source": test['source'],
|
||||
"actual_answer": answer,
|
||||
"actual_sources": sources,
|
||||
"response_time": round(elapsed_time, 2),
|
||||
"evaluation": evaluation
|
||||
}
|
||||
results.append(result)
|
||||
|
||||
# 打印结果
|
||||
print(f"预期答案: {test['expected']}")
|
||||
# 过滤 emoji 和特殊字符，避免编码问题
|
||||
import re
|
||||
def sanitize_text(text):
|
||||
if not text:
|
||||
return "无答案"
|
||||
# 移除 emoji 和特殊 Unicode 字符
|
||||
text = re.sub(r'[\U00010000-\U0010ffff]', '', text)
|
||||
return text[:200] + "..." if len(text) > 200 else text
|
||||
answer_display = sanitize_text(answer)
|
||||
print(f"实际答案: {answer_display}")
|
||||
print(f"检索来源: {sources[:3]}..." if len(sources) > 3 else f"检索来源: {sources}")
|
||||
print(f"评分: {evaluation['score']} ({evaluation['grade']}) - {evaluation['reason']}")
|
||||
print(f"响应时间: {elapsed_time:.2f}s")
|
||||
print("-" * 80)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def generate_report(results: list) -> str:
|
||||
"""生成测试报告"""
|
||||
# 统计各类别得分
|
||||
category_stats = {}
|
||||
for r in results:
|
||||
cat = r['category']
|
||||
if cat not in category_stats:
|
||||
category_stats[cat] = {"total": 0, "scores": [], "times": []}
|
||||
category_stats[cat]['total'] += 1
|
||||
category_stats[cat]['scores'].append(r['evaluation']['score'])
|
||||
category_stats[cat]['times'].append(r['response_time'])
|
||||
|
||||
# 计算总体统计
|
||||
all_scores = [r['evaluation']['score'] for r in results]
|
||||
all_times = [r['response_time'] for r in results]
|
||||
|
||||
# 生成报告
|
||||
report = []
|
||||
report.append("# RAG 系统测试报告")
|
||||
report.append("")
|
||||
report.append(f"**测试时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
report.append(f"**测试问题数量**: {len(results)}")
|
||||
report.append("")
|
||||
|
||||
# 总体统计
|
||||
report.append("## 一、总体统计")
|
||||
report.append("")
|
||||
report.append(f"- **平均得分**: {sum(all_scores)/len(all_scores):.1f} 分")
|
||||
report.append(f"- **最高得分**: {max(all_scores)} 分")
|
||||
report.append(f"- **最低得分**: {min(all_scores)} 分")
|
||||
report.append(f"- **平均响应时间**: {sum(all_times)/len(all_times):.2f} 秒")
|
||||
report.append(f"- **优秀率 (≥80分)**: {len([s for s in all_scores if s >= 80])/len(all_scores)*100:.1f}%")
|
||||
report.append(f"- **合格率 (≥40分)**: {len([s for s in all_scores if s >= 40])/len(all_scores)*100:.1f}%")
|
||||
report.append("")
|
||||
|
||||
# 分类统计
|
||||
report.append("## 二、分类统计")
|
||||
report.append("")
|
||||
report.append("| 测试类别 | 题数 | 平均分 | 最高分 | 最低分 | 平均响应时间 |")
|
||||
report.append("|----------|------|--------|--------|--------|--------------|")
|
||||
for cat, stats in category_stats.items():
|
||||
avg_score = sum(stats['scores']) / len(stats['scores'])
|
||||
avg_time = sum(stats['times']) / len(stats['times'])
|
||||
report.append(f"| {cat} | {stats['total']} | {avg_score:.1f} | {max(stats['scores'])} | {min(stats['scores'])} | {avg_time:.2f}s |")
|
||||
report.append("")
|
||||
|
||||
# 详细结果
|
||||
report.append("## 三、详细测试结果")
|
||||
report.append("")
|
||||
|
||||
for r in results:
|
||||
report.append(f"### Q{r['id']}: {r['question']}")
|
||||
report.append("")
|
||||
report.append(f"- **测试类别**: {r['category']}")
|
||||
report.append(f"- **预期答案**: {r['expected']}")
|
||||
report.append(f"- **预期来源**: {r['expected_source']}")
|
||||
report.append(f"- **实际答案**: {r['actual_answer']}")
|
||||
# 处理 sources 格式（可能是 dict 列表或 str 列表）
|
||||
sources = r['actual_sources']
|
||||
if sources and isinstance(sources[0], dict):
|
||||
sources_str = ', '.join([s.get('source', str(s)) for s in sources[:5]])
|
||||
else:
|
||||
sources_str = ', '.join([str(s) for s in sources[:5]]) if sources else '无'
|
||||
report.append(f"- **检索来源**: {sources_str}")
|
||||
report.append(f"- **评分**: {r['evaluation']['score']} ({r['evaluation']['grade']})")
|
||||
report.append(f"- **响应时间**: {r['response_time']}s")
|
||||
report.append("")
|
||||
|
||||
# 问题与建议
|
||||
report.append("## 四、问题分析")
|
||||
report.append("")
|
||||
|
||||
# 低分问题
|
||||
low_score_questions = [r for r in results if r['evaluation']['score'] < 40]
|
||||
if low_score_questions:
|
||||
report.append("### 低分问题 (< 40分)")
|
||||
report.append("")
|
||||
for r in low_score_questions:
|
||||
report.append(f"- Q{r['id']}: {r['question']} (得分: {r['evaluation']['score']})")
|
||||
report.append("")
|
||||
|
||||
return "\n".join(report)
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
# Windows 控制台编码处理
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8', errors='replace')
|
||||
sys.stderr.reconfigure(encoding='utf-8', errors='replace')
|
||||
|
||||
# 执行测试
|
||||
results = run_tests()
|
||||
|
||||
# 生成报告
|
||||
report = generate_report(results)
|
||||
|
||||
# 保存报告
|
||||
report_file = "rag_test_report.md"
|
||||
with open(report_file, "w", encoding="utf-8") as f:
|
||||
f.write(report)
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("测试完成！")
|
||||
print(f"报告已保存到: {report_file}")
|
||||
print("=" * 80)
|
||||
|
||||
# 保存详细结果 JSON
|
||||
json_file = "rag_test_results.json"
|
||||
with open(json_file, "w", encoding="utf-8") as f:
|
||||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"详细结果已保存到: {json_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user