search_unified.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722
  1. #!/usr/bin/env python3
  2. """
  3. LambdaRAG v3.1: Intent-Entity-Fusion-Slot Pipeline
  4. Changes from v3:
  5. P0-1: Fix graph keywords error (ensure all paths return keywords)
  6. P0-2: Scoring-based intent classification (not first-match)
  7. P0-3: Query logging to feedback/query_log.jsonl
  8. P0-4: Missing query logging to feedback/missing_queries.jsonl
  9. P1-1: Corrective RAG — retrieval quality check + reformulate
  10. P1-2: Compare skips entity enhancement
  11. P1-5: Temporal time expansion ("近五年" → absolute years)
  12. Architecture:
  13. λ = IntentEntityExtract(q)
  14. >> EntityEnhancedFusionSearch(q, entities, intent)
  15. >> [Optional] RetrievalQualityCheck → Reformulate+Retry
  16. >> SlotFillFromResults(intent, entities, results)
  17. >> BuildEnrichedContext(slots, results, relations)
  18. >> LogQuery
  19. """
  20. import json
  21. import os
  22. import re
  23. import sys
  24. import time
  25. import urllib.request
  26. from collections import defaultdict
  27. from datetime import datetime
  28. from pathlib import Path
  29. from concurrent.futures import ThreadPoolExecutor, as_completed
  30. # ── Imports ──
  31. try:
  32. from search_engine import search as _bm25_search_fn
  33. except ImportError:
  34. _bm25_search_fn = None
  35. try:
  36. from search_engine_v2 import _vector_search as _vector_search_fn, _embed_query
  37. except ImportError:
  38. _vector_search_fn = None
  39. _embed_query = None
  40. try:
  41. from graph_engine import get_graph
  42. except ImportError:
  43. get_graph = None
  44. try:
  45. from config import cfg
  46. WIKI_DIR = str(cfg.wiki_dir)
  47. VLLM_URL = cfg.vllm_url
  48. WEIGHT_PROFILES = cfg.weight_profiles
  49. DEFAULT_WEIGHTS = cfg.default_weights
  50. RRF_K = cfg.rrf_k
  51. PROMPTS = cfg.raw.get('prompts', {})
  52. FEEDBACK_DIR = cfg.base_dir / 'feedback'
  53. except ImportError:
  54. WIKI_DIR = os.environ.get('WIKI_DIR', str(Path(__file__).parent.parent / 'wiki'))
  55. VLLM_URL = os.environ.get('VLLM_URL', 'http://127.0.0.1:8000/v1/chat/completions')
  56. WEIGHT_PROFILES = {
  57. 'person': {'bm25': 0.8, 'vector': 0.5, 'graph': 0.3, 'wiki': 0.5},
  58. 'process': {'bm25': 0.7, 'vector': 0.5, 'graph': 0.3, 'wiki': 0.6},
  59. 'fact': {'bm25': 1.0, 'vector': 0.5, 'graph': 0.1, 'wiki': 0.1},
  60. 'compare': {'bm25': 0.5, 'vector': 0.8, 'graph': 0.4, 'wiki': 0.7},
  61. 'temporal': {'bm25': 1.0, 'vector': 0.2, 'graph': 0.05, 'wiki': 0.2},
  62. 'data': {'bm25': 1.0, 'vector': 0.3, 'graph': 0.05, 'wiki': 0.1},
  63. }
  64. DEFAULT_WEIGHTS = {'bm25': 0.7, 'vector': 0.5, 'graph': 0.3, 'wiki': 0.3}
  65. RRF_K = 60
  66. PROMPTS = {}
  67. FEEDBACK_DIR = Path(os.environ.get('KNOWLEDGE_DIR', '.')) / 'feedback'
  68. # Ensure weights for all types
  69. for t, w in [('person', {'bm25':0.8,'vector':0.5,'graph':0.3,'wiki':0.5}),
  70. ('process', {'bm25':0.7,'vector':0.5,'graph':0.3,'wiki':0.6}),
  71. ('data', {'bm25':1.0,'vector':0.3,'graph':0.05,'wiki':0.1})]:
  72. if t not in WEIGHT_PROFILES:
  73. WEIGHT_PROFILES[t] = w
  74. CURRENT_YEAR = datetime.now().year
  75. def _llm_call(prompt, max_tokens=500):
  76. body = json.dumps({
  77. "model": "qwen2.5-32b",
  78. "messages": [{"role": "user", "content": prompt}],
  79. "temperature": 0.0, "max_tokens": max_tokens,
  80. }).encode('utf-8')
  81. req = urllib.request.Request(VLLM_URL, data=body,
  82. headers={"Content-Type": "application/json"}, method="POST")
  83. with urllib.request.urlopen(req, timeout=30) as resp:
  84. data = json.loads(resp.read())
  85. return data["choices"][0]["message"]["content"].strip()
  86. # ══════════════════════════════════════════
  87. # P0-3 / P0-4: Feedback Logging
  88. # ══════════════════════════════════════════
  89. def _log_query(query, intent, strategy, elapsed, path_counts, has_result):
  90. """P0-3: Append query record to feedback/query_log.jsonl"""
  91. try:
  92. FEEDBACK_DIR.mkdir(parents=True, exist_ok=True)
  93. record = {
  94. "ts": datetime.now().strftime("%Y-%m-%dT%H:%M:%S"),
  95. "query": query[:200],
  96. "intent": intent,
  97. "strategy": strategy,
  98. "elapsed": elapsed,
  99. "path_counts": path_counts,
  100. "has_result": has_result,
  101. }
  102. with open(FEEDBACK_DIR / "query_log.jsonl", "a", encoding="utf-8") as f:
  103. f.write(json.dumps(record, ensure_ascii=False) + "\n")
  104. except:
  105. pass # Logging should never crash the pipeline
  106. def _log_missing(query, intent):
  107. """P0-4: Log queries with no useful results"""
  108. try:
  109. FEEDBACK_DIR.mkdir(parents=True, exist_ok=True)
  110. record = {
  111. "ts": datetime.now().strftime("%Y-%m-%dT%H:%M:%S"),
  112. "query": query[:200],
  113. "intent": intent,
  114. }
  115. with open(FEEDBACK_DIR / "missing_queries.jsonl", "a", encoding="utf-8") as f:
  116. f.write(json.dumps(record, ensure_ascii=False) + "\n")
  117. except:
  118. pass
  119. # ══════════════════════════════════════════
  120. # P0-2: Scoring-based Intent Classification
  121. # ══════════════════════════════════════════
  122. # Intent keyword rules with scores (higher = stronger signal)
  123. _INTENT_RULES = {
  124. 'person': [
  125. (r'谁负责|找谁|联系方式|打什么电话|哪位老师|联系人|哪个部门管辖|主管机关', 5),
  126. ],
  127. 'process': [
  128. (r'怎么办|操作步骤|怎么报销|审批流程|怎么预约|怎么退回|怎么申请|办理流程', 5),
  129. (r'能不能报销|可以报销吗|是否可以|能否|需要.*吗|要不要', 4),
  130. (r'报销|预约|退回|建账|扣税|签字|审批|发放|申请|备案|登记', 2),
  131. ],
  132. 'data': [
  133. (r'及格率|通过率|排名|统计数据|人次.*统计|吞吐量|收费.*标准.*表|费用.*表', 5),
  134. (r'清单.*内容|哪些.*机构|标准.*表', 3),
  135. ],
  136. 'fact': [
  137. (r'\d{4}年.*多少|数量.*多少|人数|统计|税率|额度|识别号|有效期', 4),
  138. (r'标准.*多少|是什么|是多少|包括哪些|有哪些', 2),
  139. ],
  140. 'compare': [
  141. (r'异同|比较|区别|对比|不同|相同|差异', 5),
  142. ],
  143. 'relation': [
  144. (r'隶属|上级|下级|依据|谁.*管|属于|归.*管|由谁制定|哪个.*认可', 4),
  145. ],
  146. 'synthesis': [
  147. (r'总结|综合|梳理|分析.*体系|框架|概述|全面|整体', 4),
  148. ],
  149. 'temporal': [
  150. (r'最新|修订|什么时候|截止|几月份|何时.*施行|变化.*趋势', 4),
  151. (r'近.*年|历史|演变', 3),
  152. ],
  153. }
  154. def _score_intent(query):
  155. """P0-2: Score-based classification — highest score wins"""
  156. q = query.lower()
  157. scores = defaultdict(float)
  158. for intent, rules in _INTENT_RULES.items():
  159. for pattern, weight in rules:
  160. if re.search(pattern, q):
  161. scores[intent] += weight
  162. if not scores:
  163. return None
  164. # Return highest scoring intent
  165. best = max(scores.items(), key=lambda x: x[1])
  166. return best[0] if best[1] >= 2 else None # Minimum threshold of 2
  167. # ══════════════════════════════════════════
  168. # Step 1: Intent + Entity Extraction
  169. # ══════════════════════════════════════════
  170. def intent_entity_extract(query):
  171. """
  172. Classify intent (scoring-based) + extract entities.
  173. Fast path for most intents, LLM only for person/process or ambiguous.
  174. """
  175. # P0-2: Score-based classification
  176. fast_intent = _score_intent(query)
  177. # Fast path: rule-classified non-person/process → skip LLM
  178. if fast_intent and fast_intent not in ('person', 'process'):
  179. entities = _rule_based_entities(query)
  180. return {
  181. 'intent': fast_intent,
  182. 'entities': entities,
  183. 'keywords': list(entities.values())[:5],
  184. 'original_query': query,
  185. }
  186. # LLM extraction for: unknown / person / process
  187. try:
  188. prompt = f"""分析以下问题,输出JSON(不要其他文字):
  189. {{
  190. "intent": "person/process/fact/data/compare/temporal 之一",
  191. "entities": {{"关键实体名": "值"}},
  192. "keywords": ["扩展关键词"]
  193. }}
  194. 问题:{query}"""
  195. result = _llm_call(prompt, max_tokens=200)
  196. match = re.search(r'\{[\s\S]*\}', result)
  197. if match:
  198. parsed = json.loads(match.group())
  199. intent = fast_intent or parsed.get('intent', 'fact')
  200. entities = parsed.get('entities', {})
  201. keywords = parsed.get('keywords', [])
  202. return {
  203. 'intent': intent,
  204. 'entities': entities,
  205. 'keywords': keywords,
  206. 'original_query': query,
  207. }
  208. except Exception as e:
  209. print(f"[intent_entity] LLM error: {e}")
  210. return {
  211. 'intent': fast_intent or 'fact',
  212. 'entities': _rule_based_entities(query),
  213. 'keywords': [],
  214. 'original_query': query,
  215. }
  216. def _rule_based_entities(query):
  217. """Extract entities from query using simple rules (no LLM)"""
  218. entities = {}
  219. quoted = re.findall(r'[《「](.+?)[》」]', query)
  220. for q in quoted:
  221. entities[q] = q
  222. nouns = re.findall(r'[\u4e00-\u9fff]{2,8}', query)
  223. seen = set()
  224. for n in nouns:
  225. if n not in seen and len(n) >= 2:
  226. seen.add(n)
  227. entities[n] = n
  228. if len(entities) >= 5:
  229. break
  230. return entities
  231. # ══════════════════════════════════════════
  232. # P1-5: Temporal Time Expansion
  233. # ══════════════════════════════════════════
  234. def _expand_temporal_query(query):
  235. """Expand vague time words to absolute years"""
  236. expanded = query
  237. year = CURRENT_YEAR
  238. expanded = re.sub(r'近五年|近5年', f'{year-4}年至{year}年', expanded)
  239. expanded = re.sub(r'近三年|近3年', f'{year-2}年至{year}年', expanded)
  240. expanded = re.sub(r'近年来|近年', f'{year-2}年至{year}年', expanded)
  241. expanded = re.sub(r'最新', f'{year}年最新', expanded)
  242. return expanded
  243. # ══════════════════════════════════════════
  244. # Step 2: Entity-Enhanced 4-Path Fusion Search
  245. # ══════════════════════════════════════════
  246. def entity_enhanced_search(query_info, top_k=5):
  247. query = query_info['original_query']
  248. intent = query_info['intent']
  249. entities = query_info.get('entities', {})
  250. keywords = query_info.get('keywords', [])
  251. weights = WEIGHT_PROFILES.get(intent, DEFAULT_WEIGHTS)
  252. # P1-2: Compare skips entity enhancement (entities dilute precise matching)
  253. if intent == 'compare':
  254. enhanced_query = query
  255. else:
  256. entity_terms = ' '.join(str(v) for v in entities.values() if v)
  257. keyword_terms = ' '.join(keywords[:5])
  258. enhanced_query = f"{query} {entity_terms} {keyword_terms}".strip()
  259. # P1-5: Temporal time expansion
  260. if intent == 'temporal':
  261. enhanced_query = _expand_temporal_query(enhanced_query)
  262. entity_names = [str(v) for v in entities.values() if v and len(str(v)) > 1]
  263. path_results = {}
  264. with ThreadPoolExecutor(max_workers=4) as pool:
  265. futures = {
  266. pool.submit(_path_bm25, enhanced_query, 10): 'bm25',
  267. pool.submit(_path_vector, query, 10): 'vector',
  268. pool.submit(_path_graph, query, entity_names, 10): 'graph',
  269. pool.submit(_path_wiki, query, entity_names, 5): 'wiki',
  270. }
  271. for future in as_completed(futures):
  272. name = futures[future]
  273. try:
  274. path_results[name] = future.result(timeout=15)
  275. except Exception as e:
  276. print(f"[{name}] Error: {e}")
  277. path_results[name] = []
  278. fused = _weighted_rrf_fuse(path_results, weights)
  279. if intent == 'temporal':
  280. fused = _time_filter(fused, query)
  281. if intent == 'data':
  282. fused = _boost_table_chunks(fused)
  283. return fused, path_results, weights
  284. # ══════════════════════════════════════════
  285. # P1-1: Corrective RAG — Retrieval Quality Check
  286. # ══════════════════════════════════════════
  287. def _check_retrieval_quality(query, fused_results, threshold=0.3):
  288. """
  289. Quick check: do top results actually contain query-relevant content?
  290. Returns True if quality is acceptable, False if reformulation needed.
  291. Uses simple keyword overlap (no LLM call for speed).
  292. """
  293. if not fused_results:
  294. return False
  295. # Check if top-3 results contain at least some query keywords
  296. query_chars = set(re.findall(r'[\u4e00-\u9fff]', query))
  297. if not query_chars:
  298. return True
  299. top_text = ' '.join(r.get('text', '')[:300] for r in fused_results[:3])
  300. overlap = sum(1 for c in query_chars if c in top_text)
  301. ratio = overlap / len(query_chars)
  302. return ratio >= threshold
  303. def _reformulate_query(query):
  304. """Simple query reformulation: extract core nouns"""
  305. nouns = re.findall(r'[\u4e00-\u9fff]{2,6}', query)
  306. if len(nouns) >= 2:
  307. return ' '.join(nouns[:4])
  308. return query
  309. # ══════════════════════════════════════════
  310. # 4-Path Retrieval Functions
  311. # ══════════════════════════════════════════
  312. def _path_bm25(query, top_k=10):
  313. if _bm25_search_fn:
  314. try:
  315. results = _bm25_search_fn(query, top_k=top_k)
  316. for r in results:
  317. r['_path'] = 'bm25'
  318. return results
  319. except Exception as e:
  320. print(f"[bm25] Error: {e}")
  321. return []
  322. def _path_vector(query, top_k=10):
  323. if _vector_search_fn:
  324. try:
  325. results = _vector_search_fn(query, top_k=top_k)
  326. for r in results:
  327. r['_path'] = 'vector'
  328. return results
  329. except Exception as e:
  330. print(f"[vector] Error: {e}")
  331. return []
  332. def _path_graph(query, entity_names=None, top_k=10):
  333. if not get_graph:
  334. return []
  335. try:
  336. graph = get_graph(wiki_dir=WIKI_DIR)
  337. seeds = entity_names or []
  338. if not seeds:
  339. seeds = graph.find_entities(query, max_results=5)
  340. if not seeds:
  341. return []
  342. expand = graph.expand_query(query, seed_entities=seeds, max_hops=2)
  343. expanded_query = expand.get('expanded_query', query)
  344. results = []
  345. if _bm25_search_fn:
  346. r = _bm25_search_fn(expanded_query, top_k=top_k // 2)
  347. if r:
  348. results.extend(r)
  349. if _vector_search_fn:
  350. r = _vector_search_fn(expanded_query, top_k=top_k // 2)
  351. if r:
  352. results.extend(r)
  353. for rel in expand.get('relations', [])[:5]:
  354. results.append({
  355. 'source': f"relation:{rel.get('subject','')}→{rel.get('object','')}",
  356. 'chunk_id': 0,
  357. 'text': f"{rel.get('subject','')} {rel.get('relation','')} {rel.get('object','')}。{rel.get('description', '')}",
  358. 'score': 0.5,
  359. })
  360. for r in results:
  361. r['_path'] = 'graph'
  362. if results:
  363. results[0]['_relation_context'] = expand.get('relation_context', '')
  364. return results
  365. except Exception as e:
  366. print(f"[graph] Error: {e}")
  367. return []
  368. def _path_wiki(query, entity_names=None, top_k=5):
  369. wiki_path = Path(WIKI_DIR)
  370. if not wiki_path.exists():
  371. return []
  372. pages = []
  373. for subdir in ['sources', 'entities', 'topics', 'analyses']:
  374. d = wiki_path / subdir
  375. if not d.exists():
  376. continue
  377. for f in d.glob('*.md'):
  378. try:
  379. content = f.read_text(encoding='utf-8')
  380. score = 0
  381. if entity_names:
  382. for e in entity_names:
  383. if e in f.stem or e in content[:500]:
  384. score += 10
  385. score += sum(1 for c in set(query) if c in content[:500])
  386. if score > len(query) * 0.2:
  387. pages.append({
  388. 'source': f'wiki:{subdir}/{f.stem}',
  389. 'chunk_id': 0,
  390. 'text': content[:800],
  391. 'score': score / 100.0,
  392. '_path': 'wiki',
  393. })
  394. except:
  395. pass
  396. pages.sort(key=lambda x: -x['score'])
  397. return pages[:top_k]
  398. def _weighted_rrf_fuse(path_results, weights, k=RRF_K):
  399. scores = defaultdict(float)
  400. chunk_map = {}
  401. path_tags = defaultdict(set)
  402. for path_name, results in path_results.items():
  403. w = weights.get(path_name, 0.5)
  404. if w == 0:
  405. continue
  406. for rank, r in enumerate(results):
  407. key = f"{r.get('source', '')}#{r.get('chunk_id', 0)}"
  408. scores[key] += w / (k + rank + 1)
  409. if key not in chunk_map:
  410. chunk_map[key] = r.copy()
  411. path_tags[key].add(path_name)
  412. sorted_keys = sorted(scores.items(), key=lambda x: -x[1])
  413. fused = []
  414. for key, score in sorted_keys:
  415. r = chunk_map[key]
  416. r['score'] = round(score, 4)
  417. r['_fusion_paths'] = list(path_tags[key])
  418. fused.append(r)
  419. return fused
  420. def _boost_table_chunks(results):
  421. table_results = []
  422. other_results = []
  423. for r in results:
  424. text = r.get('text', '')
  425. if 'TABLE_START' in text or (text.count('|') > 10 and '---' in text):
  426. r['score'] = r.get('score', 0) + 0.1
  427. table_results.append(r)
  428. else:
  429. other_results.append(r)
  430. return table_results + other_results
  431. def _time_filter(results, query):
  432. scored = []
  433. for r in results:
  434. boost = 0
  435. text = r.get('text', '') + r.get('source', '')
  436. if any(k in text for k in ['最新修订', '现行有效', '起施行']):
  437. boost += 2
  438. src_years = re.findall(r'(20\d{2})', r.get('source', ''))
  439. if src_years:
  440. boost += (max(int(y) for y in src_years) - 2000) * 0.1
  441. scored.append((r, r.get('score', 0) + boost))
  442. scored.sort(key=lambda x: -x[1])
  443. return [r for r, _ in scored]
  444. # ══════════════════════════════════════════
  445. # Step 3: Slot Fill from Search Results
  446. # ══════════════════════════════════════════
  447. def slot_fill(query_info, fused_results):
  448. intent = query_info['intent']
  449. entities = query_info.get('entities', {})
  450. query = query_info['original_query']
  451. # Fast path: fact/compare/temporal/synthesis skip slot fill
  452. if intent in ('fact', 'compare', 'temporal', 'synthesis'):
  453. return {}
  454. context = '\n'.join([r.get('text', '')[:500] for r in fused_results[:5]])
  455. if not context.strip():
  456. return {}
  457. if intent == 'person' and 'person_extract' in PROMPTS:
  458. prompt = PROMPTS['person_extract'].format(context=context[:2500], question=query)
  459. elif intent == 'process' and 'process_extract' in PROMPTS:
  460. prompt = PROMPTS['process_extract'].format(context=context[:2500], question=query)
  461. elif intent == 'fact' and 'fact_extract' in PROMPTS:
  462. prompt = PROMPTS['fact_extract'].format(context=context[:2500], question=query)
  463. elif intent in ('relation', 'synthesis') and 'fact_extract' in PROMPTS:
  464. prompt = PROMPTS['fact_extract'].format(context=context[:2500], question=query)
  465. elif intent == 'data':
  466. entity_desc = ', '.join(f'{k}: {v}' for k, v in entities.items()) if entities else ''
  467. prompt = f"""基于以下知识库内容(可能包含表格),精确提取问题要求的数据。
  468. 仅使用知识库信息,不编造。
  469. 知识库内容:
  470. {context[:2500]}
  471. 问题:{query}
  472. 查找目标:{entity_desc}
  473. 输出JSON:
  474. {{"answer_summary": "直接回答(含具体数字)",
  475. "data_values": {{"指标名": "数值"}},
  476. "table_source": "数据来源文件名",
  477. "source": "信息来源"}}"""
  478. else:
  479. entity_desc = ', '.join(f'{k}: {v}' for k, v in entities.items()) if entities else ''
  480. prompt = f"""基于以下知识库内容,提取关键信息。仅使用知识库信息,不编造。输出JSON。
  481. 知识库内容:
  482. {context[:2000]}
  483. 问题:{query}
  484. {{"answer_summary": "一句话回答",
  485. "key_values": {{"指标": "值"}},
  486. "conditions": "适用条件",
  487. "policy_basis": "政策依据",
  488. "source": "来源文件名"}}"""
  489. try:
  490. result = _llm_call(prompt, max_tokens=500)
  491. match = re.search(r'\{[\s\S]*\}', result)
  492. if match:
  493. return json.loads(match.group())
  494. except Exception as e:
  495. print(f"[slot_fill] Error: {e}")
  496. return {}
  497. # ══════════════════════════════════════════
  498. # Step 4: Build Enriched Context
  499. # ══════════════════════════════════════════
  500. def build_enriched_context(query_info, fused_results, slots, top_k=5):
  501. context_parts = []
  502. relation_context = ""
  503. wiki_summaries = []
  504. if slots:
  505. slots_text = "## 提取的关键信息\n"
  506. for k, v in slots.items():
  507. if v and k != 'source':
  508. if isinstance(v, dict):
  509. for sk, sv in v.items():
  510. if sv:
  511. slots_text += f"- {sk}: {sv}\n"
  512. else:
  513. slots_text += f"- {k}: {v}\n"
  514. context_parts.append(slots_text)
  515. for r in fused_results[:top_k]:
  516. source = r.get('source', '')
  517. text = r.get('text', '')
  518. if source.startswith('wiki:'):
  519. wiki_summaries.append(f"[Wiki: {source[5:]}]\n{text[:600]}")
  520. elif source.startswith('relation:'):
  521. if not relation_context:
  522. relation_context = "## 实体关系\n"
  523. relation_context += f"- {text}\n"
  524. else:
  525. context_parts.append(f"[doc: {source}]\n{text}")
  526. if '_relation_context' in r and r['_relation_context']:
  527. if not relation_context:
  528. relation_context = "## 实体关系\n"
  529. relation_context += r['_relation_context'] + "\n"
  530. enriched = ""
  531. if relation_context:
  532. enriched += relation_context + "\n"
  533. if wiki_summaries:
  534. enriched += "## Wiki 知识\n" + "\n".join(wiki_summaries) + "\n\n"
  535. enriched += "\n".join(context_parts)
  536. return enriched
  537. # ══════════════════════════════════════════
  538. # Main search function
  539. # ══════════════════════════════════════════
  540. def search(query, top_k=5):
  541. """
  542. LambdaRAG v3.1 Pipeline:
  543. 1. Intent + Entity extraction (scoring-based, fast)
  544. 2. Entity-enhanced 4-path fusion search
  545. 2b. [P1-1] Retrieval quality check → reformulate if needed
  546. 3. Slot filling (person/process/data only)
  547. 4. Build enriched context
  548. 5. Log query + log missing
  549. """
  550. t0 = time.time()
  551. # Step 1: Intent + Entity
  552. query_info = intent_entity_extract(query)
  553. intent = query_info['intent']
  554. # Step 2: Entity-enhanced fusion search
  555. fused, path_results, weights = entity_enhanced_search(query_info, top_k)
  556. # Step 2b: P1-1 Corrective RAG — check retrieval quality
  557. if not _check_retrieval_quality(query, fused):
  558. # Reformulate and retry once
  559. alt_query = _reformulate_query(query)
  560. if alt_query != query:
  561. alt_info = {**query_info, 'original_query': alt_query}
  562. alt_fused, alt_paths, _ = entity_enhanced_search(alt_info, top_k)
  563. if _check_retrieval_quality(query, alt_fused):
  564. fused = alt_fused
  565. path_results = alt_paths
  566. # Step 3: Slot fill
  567. slots = slot_fill(query_info, fused)
  568. # Step 4: Build enriched context
  569. enriched = build_enriched_context(query_info, fused, slots, top_k)
  570. elapsed = time.time() - t0
  571. # Metadata
  572. strategy = f"v3.1({intent}|e={len(query_info.get('entities',{}))}|s={len(slots)})"
  573. pc = {k: len(v) for k, v in path_results.items()}
  574. if fused:
  575. fused[0]['_query_type'] = intent
  576. fused[0]['_entities'] = query_info.get('entities', {})
  577. fused[0]['_slots'] = slots
  578. fused[0]['_weights'] = weights
  579. fused[0]['_strategy'] = strategy
  580. fused[0]['_total_time'] = round(elapsed, 3)
  581. fused[0]['_path_counts'] = pc
  582. fused[0]['_enriched_context'] = enriched[:3500]
  583. # Step 5: P0-3 + P0-4 Logging
  584. has_result = bool(fused)
  585. _log_query(query, intent, strategy, round(elapsed, 2), pc, has_result)
  586. if not has_result:
  587. _log_missing(query, intent)
  588. return fused[:top_k]
  589. # ══════════════════════════════════════════
  590. # CLI test
  591. # ══════════════════════════════════════════
  592. if __name__ == '__main__':
  593. query = sys.argv[1] if len(sys.argv) > 1 else '横向经费报销接待餐费的标准'
  594. print(f"Query: {query}")
  595. print("=" * 60)
  596. info = intent_entity_extract(query)
  597. print(f"Intent: {info['intent']}")
  598. print(f"Entities: {info['entities']}")
  599. print(f"Keywords: {info['keywords']}")
  600. print(f"\nSearching...")
  601. results = search(query, top_k=5)
  602. if results:
  603. meta = results[0]
  604. print(f"Time: {meta.get('_total_time', '?')}s")
  605. print(f"Strategy: {meta.get('_strategy', '?')}")
  606. print(f"Slots: {json.dumps(meta.get('_slots', {}), ensure_ascii=False)[:200]}")
  607. print(f"Path counts: {meta.get('_path_counts', {})}")
  608. for i, r in enumerate(results):
  609. paths = r.get('_fusion_paths', [])
  610. print(f"\n[{i+1}] score={r.get('score','?')} paths={paths} {r.get('source','?')[:60]}")
  611. # Show feedback log
  612. if FEEDBACK_DIR.exists():
  613. log = FEEDBACK_DIR / "query_log.jsonl"
  614. if log.exists():
  615. lines = log.read_text().strip().split('\n')
  616. print(f"\nFeedback log: {len(lines)} entries")