num/knowledge/scripts/test_rag_retrieval.py

"""
Dify RAG Retrieval Test Script (Dual Mode)
===========================================
Tests knowledge base retrieval across 5 query types.

Modes:
  --mode local   Local retriever (no Dify needed)
  --mode dify    Dify API (requires API key)

Usage:
  python test_rag_retrieval.py --mode local
  python test_rag_retrieval.py --mode dify --api-key <KEY>
"""
import os
import sys
import json
import argparse


# (type, chinese_query, expected_file_substring, doc_hint)
TEST_QUERIES = [
    # Type 1: Position queries -> should retrieve A1 (命盘24位置详解)
    ("position", "O位置代表什么含义", "命盘", "A1"),
    ("position", "P位置如何计算", "命盘", "A1"),
    ("position", "I和J位置代表什么能量", "命盘", "A1"),

    # Type 2: Combo pair queries -> should retrieve B2 (三角命盘组合对判读规则)
    ("combo", "M和N组合出现天医星怎么解读", "判读", "B2"),
    ("combo", "横向组合对和纵向组合对的区", "判读", "B2"),
    ("combo", "组合对分析的优先顺序是什么", "判读", "B2"),

    # Type 3: Dimension queries -> should retrieve C1 (数字能量分析维度手册)
    ("dimension", "天医星财富维度如何分析", "维度", "C1"),
    ("dimension", "延年星在事业维度的含义", "维度", "C1"),
    ("dimension", "五鬼星对健康维度的影响", "维度", "C1"),

    # Type 4: Zone/group queries -> should retrieve C2 (五区三组分析指南)
    ("zone", "父源区包含哪些位置", "五区", "C2"),
    ("zone", "左侧组年龄段和人生课题", "五区", "C2"),
    ("zone", "五区和三组的交叉分析", "五区", "C2"),

    # Type 5: Comprehensive
    ("comprehensive", "主性格6号人的感情事业特点", "维度", "C1"),
]


def _dify_retrieve(query_text, api_key, base_url, dataset_id, top_k=6):
    """Dify API retrieval"""
    import requests
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    url = f"{base_url}/v1/datasets/{dataset_id}/retrieve"
    payload = {
        "query": query_text,
        "retrieval_model": {
            "search_method": "hybrid_search",
            "reranking_enable": True,
            "top_k": top_k,
            "score_threshold_enabled": False
        }
    }
    try:
        resp = requests.post(url, headers=headers, json=payload, timeout=30)
        if resp.status_code == 200:
            records = resp.json().get("records", [])
            return records, None
        else:
            return None, f"HTTP {resp.status_code} {resp.text[:200]}"
    except Exception as e:
        return None, str(e)


def _dify_run_all(api_key, base_url, dataset_id):
    """Execute all Dify retrieval tests"""
    import requests

    print("=" * 70)
    print("Dify KB RAG Retrieval Test")
    print("=" * 70)
    print(f"Dataset: {dataset_id}")
    print(f"URL:     {base_url}")

    # Show config
    try:
        resp = requests.get(f"{base_url}/v1/datasets/{dataset_id}",
                          headers={"Authorization": f"Bearer {api_key}"}, timeout=15)
        if resp.status_code == 200:
            info = resp.json()
            rm = info.get("retrieval_model_dict", {})
            print(f"KB: {info.get('name', '?')}")
            print(f"Docs: {info.get('document_count')}")
            print(f"Config: method={rm.get('search_method')} top_k={rm.get('top_k')}")
    except Exception:
        pass

    print(f"Queries: {len(TEST_QUERIES)}")
    print()

    stats = {"pass": 0, "fail": 0, "error": 0, "total": len(TEST_QUERIES), "mode": "dify"}

    for qtype, query, expected, hint in TEST_QUERIES:
        print(f"[{qtype}] {query[:40]}...", end=" ")

        results, error = _dify_retrieve(query, api_key, base_url, dataset_id)

        if error:
            print(f"ERROR: {error}")
            stats["error"] += 1
            continue

        if not results:
            print("NO RESULTS")
            stats["fail"] += 1
            continue

        doc_names = []
        for r in results:
            seg = r.get("segment", {})
            doc = seg.get("document", {})
            doc_name = doc.get("name", "")
            if doc_name:
                doc_names.append(doc_name)

        hit = any(expected in name for name in doc_names)

        if hit:
            print(f"HIT  [{hint}]")
            stats["pass"] += 1
        else:
            sources = list(doc_names[:5])
            print(f"MISS expected='{expected}' got={sources}")
            stats["fail"] += 1

        if results:
            top = results[0]
            seg = top.get("segment", {})
            snippet = (seg.get("content", "") or "")[:120].replace("\n", " ")
            score = top.get("score", "?")
            print(f"  score={score:.4f} | doc: {doc_names[0] if doc_names else '?'}")
            print(f"  {snippet}...")
        print()

    return stats


def _local_run_all():
    """Execute all local retrieval tests"""
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from local_retriever import LocalRetriever

    retriever = LocalRetriever()
    idx_stats = retriever.index_documents()

    print("=" * 70)
    print("Local KB RAG Retrieval Test")
    print("=" * 70)
    print(f"Docs: {idx_stats['documents']}, Chunks: {idx_stats['chunks']}")
    print(f"Vocab: {idx_stats['vocab_size']}")
    print(f"Queries: {len(TEST_QUERIES)}")
    print()

    stats = {"pass": 0, "fail": 0, "error": 0, "total": len(TEST_QUERIES), "mode": "local"}

    for qtype, query, expected, hint in TEST_QUERIES:
        print(f"[{qtype}] {query[:40]}...", end=" ")

        results = retriever.query(query, top_k=6)

        if not results:
            print("NO RESULTS")
            stats["error"] += 1
            continue

        hit = any(expected in r["source"] for r in results)

        if hit:
            print(f"HIT  [{hint}]")
            stats["pass"] += 1
        else:
            top_src = [r["source"] for r in results[:5]]
            print(f"MISS expected='{expected}' got={top_src}")
            stats["fail"] += 1

        top = results[0]
        print(f"  score={top['score']:.4f} | {top['source']} > {top['heading']}")
        print(f"  {top['content_preview'][:120].replace(chr(10), ' ')}")
        print()

    return stats


def print_report(stats):
    total = stats["total"]
    passed = stats["pass"]
    failed = stats["fail"]
    errored = stats["error"]

    print("=" * 70)
    print("TEST REPORT")
    print("=" * 70)
    print(f"Mode:    {stats.get('mode', '?')}")
    print(f"Total:   {total}")
    print(f"Pass:    {passed}")
    print(f"Fail:    {failed}")
    print(f"Error:   {errored}")

    effective = total - errored
    if effective > 0:
        rate = passed * 100 // effective
        print(f"Rate:    {rate}% ({passed}/{effective})")

    if failed == 0 and errored == 0:
        print(f"\n  ALL TESTS PASSED!")
    elif failed > 0:
        print(f"\n  {failed} queries missed expected docs")
    if errored > 0:
        print(f"\n  {errored} queries errored")

    return stats


def parse_args():
    p = argparse.ArgumentParser(description="RAG Retrieval Test")
    p.add_argument("--mode", choices=["local", "dify"], default="local")
    p.add_argument("--api-key")
    p.add_argument("--base-url", default="http://dify.bianwoyou.cn")
    p.add_argument("--dataset-id", default="3ff939b3-8686-44f6-8ef5-65b1e53b55d3")
    return p.parse_args()


def main():
    args = parse_args()
    if args.mode == "local":
        stats = _local_run_all()
    elif args.mode == "dify":
        api_key = args.api_key or os.environ.get("DIFY_API_KEY")
        if not api_key:
            print("Error: DIFY_API_KEY required for dify mode")
            sys.exit(1)
        stats = _dify_run_all(api_key, args.base_url, args.dataset_id)
    else:
        print(f"Unknown mode: {args.mode}")
        sys.exit(1)

    stats = print_report(stats)

    if stats.get("error", 0) > 0:
        sys.exit(2)
    if stats.get("fail", 0) > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()