""" 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 """ 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: Main Character Personality -> should retrieve A2 (主性格深度解读) ("personality", "7号人深度解读：性格特征的全面分析", "A2", "A2"), ("personality", "卓越数11的直觉力和人生课题", "A2", "A2"), ("personality", "3号人的适合职业和情感模式", "A2", "A2"), # Type 6: Supplementary knowledge -> should retrieve E (天赋数空缺数) ("supplement", "天赋数的含义速查表和计算方法", "E天赋", "E"), ("supplement", "空缺数代表什么挑战领域", "E天赋", "E"), ("supplement", "天赋数的含义和空缺数的挑战领域", "E天赋", "E"), # Type 7: Progressive energy -> should retrieve D1 (生命数1组合递进能量) ("progression", "生命数1的28-10-1组合递进三阶段", "D1", "D1"), ("progression", "46-10-1组合的务实关怀特性", "D1", "D1"), # Type 8: Comprehensive (now targets A2 specifically for main character) ("comprehensive", "主性格6号人的感情和事业特点", "A2", "A2"), ] 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()