RAG 的下一站：能自己改代码、跑实验的知识底座

用户提问 → 向量检索 topK 段 → 拼接 prompt → LLM 生成回答

"""
最小自评估 RAG demo —— 检索后自动评分，低于阈值则调整 topK 重试
依赖: pip install openai chromadb sentence-transformers
运行前设置: export OPENAI_API_KEY=sk-xxx
"""

import os
import chromadb
from sentence_transformers import SentenceTransformer
from openai import OpenAI

# ── 1. 初始化向量库和 embedding ──
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
client = chromadb.Client()
collection = client.get_or_create_collection("demo_kb")

# 写入示例文档
docs = [
    "Python 的 GIL 使得多线程无法真正并行执行 CPU 密集任务，应改用 multiprocessing。",
    "FastAPI 的依赖注入系统支持 async generator，可在请求结束时自动清理资源。",
    "ChromaDB 默认使用 cosine similarity，也可切换为 L2 或 IP 距离函数。",
    "RAG 的 chunk 大小建议 256-512 字符，过大会稀释语义，过小会丢失上下文。",
    "topK 参数不是越大越好，超过 10 段后噪声增加，生成质量反而下降。",
]
embeddings = embed_model.encode(docs).tolist()
collection.add(ids=[f"d{i}" for i in range(len(docs))],
               embeddings=embeddings, documents=docs)

# ── 2. 检索 + 评估函数 ──
def retrieve_and_score(query: str, top_k: int) -> dict:
    q_emb = embed_model.encode([query]).tolist()
    results = collection.query(query_embeddings=q_emb, n_results=top_k)
    passages = results["documents"][0]

    # 用 LLM 评估检索质量（1-5 分）
    eval_prompt = f"""你是一个检索质量评估专家。
用户问题: {query}
检索到的段落:
{chr(10).join(f'[{i+1}] {p}' for i, p in enumerate(passages)}

请评分(1-5): 段落是否包含回答问题所需的关键信息？
只输出一个数字。"""
    llm = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
    score_text = llm.chat.completions.create(
        model="gpt-4o-mini", messages=[{"role": "user", "content": eval_prompt}],
        max_tokens=5
    ).choices[0].message.content.strip()
    try:
        score = int(score_text)
    except ValueError:
        score = 3  # 解析失败给中间分

    return {"passages": passages, "score": score, "top_k": top_k}

# ── 3. 自调整循环 ──
def self_evolving_rag(query: str, initial_k: int = 3, max_retries: int = 3) -> dict:
    k = initial_k
    history = []
    for attempt in range(max_retries):
        result = retrieve_and_score(query, k)
        history.append(result)
        print(f"尝试 {attempt+1}: topK={k}, 评分={result['score']}")

        if result["score"] >= 4:
            print("✓ 评分达标，采用当前参数")
            return result

        # 评分不足，策略：逐步扩大 topK
        k = min(k + 2, len(docs))
        print(f"→ 评分不足，调整 topK 为 {k}")

    print("⚠ 达到最大重试次数，采用最佳历史结果")
    best = max(history, key=lambda r: r["score"])
    return best

# ── 4. 运行 ──
if __name__ == "__main__":
    answer = self_evolving_rag("RAG 里 topK 怎么选？")
    print(f"\n最终 topK={answer['top_k']}, 评分={answer['score']}")
    print("检索段落:")
    for p in answer["passages"]:
        print(f"  - {p}")

pip install openai chromadb sentence-transformers
export OPENAI_API_KEY=sk-your-key
python self_eval_rag.py

# eval_benchmark.yaml —— 评估基准示例
benchmark:
  name: rag_v2_qa_set
  pairs:
    - question: "RAG 里 topK 怎么选？"
      expected_facts: ["topK 不是越大越好", "超过 10 段噪声增加"]
    - question: "Python 多线程为什么不能并行 CPU 任务？"
      expected_facts: ["GIL 限制", "应改用 multiprocessing"]
  metrics:
    - fact_coverage    # 期望事实被检索到的比例
    - answer_accuracy  # 最终答案与期望的一致性

# 假设商涛开源项目提供了 pipeline 版本管理 CLI
pit version commit --config current_config.yaml --message "topK 动态化: 从固定5改为 min(3+2*retry, 15)"
pit version list
# v3  | topK=dynamic(chunk_count)  | fact_coverage=0.87 | 2025-06-01
# v2  | topK=10                     | fact_coverage=0.72 | 2025-05-28
# v1  | topK=5                      | fact_coverage=0.65 | 2025-05-20
pit version rollback --to v2   # 如果 v3 在新文档集上表现变差