@lzhzzzzwill/cofos 1.0.1 → 1.0.2

package/runtime/src/retrieval/retrieve_evidence.py

@@ -0,0 +1,859 @@
+#!/usr/bin/env python3
+"""
+COF-ROS-Reasoner: 检索证据脚本
+
+目标：
+- 使用 FAISS + BM25 进行混合检索
+- 合并结果并重排
+- 返回证据列表
+"""
+
+import os
+import json
+import yaml
+import logging
+import pickle
+import re
+from pathlib import Path
+from typing import Dict, List, Any, Optional, Tuple
+from datetime import datetime
+from collections import defaultdict
+
+import numpy as np
+
+# 添加 src 路径
+
+try:
+    from sentence_transformers import SentenceTransformer
+except ImportError:
+    SentenceTransformer = None
+
+try:
+    import faiss
+except ImportError:
+    faiss = None
+
+try:
+    from rank_bm25 import BM25Okapi
+except ImportError:
+    BM25Okapi = None
+
+from data_processing.utils import setup_logging, ensure_directory
+from retrieval.text_search import tokenize_search_text
+
+
+class EvidenceRetriever:
+    """检索证据"""
+
+    def __init__(self, config: Dict[str, Any]):
+        self.config = config
+        self.faiss_dir = Path(config['paths']['retrieval_store'])
+        self.bm25_dir = Path(config['paths']['bm25_store'])
+
+        # 设置日志
+        setup_logging(config)
+
+        # 检索配置
+        self.vector_top_k = config.get('retrieval', {}).get('vector_top_k', 20)
+        self.bm25_top_k = config.get('retrieval', {}).get('bm25_top_k', 20)
+        self.rerank_top_k = config.get('retrieval', {}).get('rerank_top_k', 5)
+        self.enable_faiss = config.get('retrieval', {}).get('enable_faiss', False)
+        self.enable_bm25 = config.get('retrieval', {}).get('enable_bm25', True)
+        self.min_bm25_score = float(
+            config.get('retrieval', {}).get('min_bm25_score', 7.0)
+        )
+        self.min_vector_score = float(
+            config.get('retrieval', {}).get('min_vector_score', 0.35)
+        )
+        self.min_kg_confidence = float(
+            config.get('retrieval', {}).get('min_kg_confidence', 2.0)
+        )
+
+        self.source_priority = config.get('retrieval', {}).get('source_priority', {
+            'gold_json': 4,
+            'silver_json': 3,
+            'pdf_fulltext': 2,
+            'wos_abstract': 1
+        })
+
+        # ROS 标准化
+        self.ros_normalization = config.get('data_processing', {}).get('ros_normalization', {})
+
+        # 加载索引与 KG 三元组
+        self._load_indices()
+        self._load_kg_triples()
+
+    def _load_indices(self):
+        """加载 FAISS 和 BM25 索引"""
+        # 加载 FAISS 索引
+        self.faiss_index = None
+        self.faiss_metadata = []
+
+        if self.enable_faiss and self.faiss_dir.exists():
+            faiss_path = self.faiss_dir / 'faiss_index.bin'
+            metadata_path = self.faiss_dir / 'faiss_metadata.jsonl'
+
+            if faiss is not None and faiss_path.exists() and metadata_path.exists():
+                self.faiss_index = faiss.read_index(str(faiss_path))
+
+                with open(metadata_path, 'r', encoding='utf-8') as f:
+                    for line in f:
+                        if line.strip():
+                            self.faiss_metadata.append(json.loads(line))
+
+                logging.info(f"Loaded FAISS index with {len(self.faiss_metadata)} vectors")
+
+        # 加载 BM25 索引
+        self.bm25_index = None
+        self.bm25_metadata = []
+
+        if self.enable_bm25 and self.bm25_dir.exists():
+            metadata_path = self.bm25_dir / 'bm25_metadata.jsonl'
+            index_path = self.bm25_dir / 'bm25_index.pkl'
+
+            if metadata_path.exists():
+                with open(metadata_path, 'r', encoding='utf-8') as f:
+                    for line in f:
+                        if line.strip():
+                            self.bm25_metadata.append(json.loads(line))
+
+                logging.info(f"Loaded BM25 index with {len(self.bm25_metadata)} documents")
+
+            if not index_path.exists():
+                raise FileNotFoundError(
+                    f"BM25 index not found: {index_path}. Run scripts/build_retrieval.py."
+                )
+            with open(index_path, 'rb') as f:
+                self.bm25_index = pickle.load(f)
+            logging.info(f"Loaded persisted BM25 index from {index_path}")
+
+
+    def _load_kg_triples(self) -> None:
+        """Load real KG triples and build lightweight in-memory indexes."""
+        self.kg_triples = []
+        self.kg_by_head = defaultdict(list)
+        self.kg_by_tail = defaultdict(list)
+
+        kg_path = Path(self.config['paths']['kg_dir']) / 'cofros_kg_triples.jsonl'
+        if not kg_path.exists():
+            logging.warning("KG triples not found: %s", kg_path)
+            return
+
+        try:
+            with open(kg_path, 'r', encoding='utf-8') as handle:
+                for line in handle:
+                    if not line.strip():
+                        continue
+                    triple = json.loads(line)
+                    head = triple.get('head', '')
+                    tail = triple.get('tail', '')
+                    if not head or not tail:
+                        continue
+                    self.kg_triples.append(triple)
+                    self.kg_by_head[head].append(triple)
+                    self.kg_by_tail[tail].append(triple)
+            logging.info("Loaded %s KG triples for graph retrieval", len(self.kg_triples))
+        except Exception as exc:
+            logging.error("Failed to load KG triples from %s: %s", kg_path, exc)
+            self.kg_triples = []
+            self.kg_by_head = defaultdict(list)
+            self.kg_by_tail = defaultdict(list)
+
+    @staticmethod
+    def _kg_label(node_id: str) -> str:
+        """Convert a KG node id into a readable label."""
+        label = str(node_id or '')
+        if ':' in label:
+            label = label.split(':', 1)[1]
+        return label.replace('_', ' ').strip()
+
+    @staticmethod
+    def _kg_search_text(text: str) -> str:
+        """Normalize labels for fuzzy query matching without changing displayed chemistry."""
+        text = str(text or '').lower().replace('_', ' ')
+        text = text.replace('−', '-').replace('⋅', '·')
+        text = re.sub(r"[^a-z0-9\u4e00-\u9fff·*+\-\s]", " ", text)
+        return re.sub(r"\s+", " ", text).strip()
+
+    @staticmethod
+    def _kg_tokens(text: str) -> set:
+        """Extract informative tokens for KG node matching.
+
+        Only generic English stopwords are removed.  Domain terms like *cof*,
+        *ros*, *photocatalysis* are deliberately KEPT — they are the very
+        terms users query with and must contribute to KG node scoring.
+        """
+        tokens = set(re.findall(r"[a-z0-9]+", text.lower()))
+        return tokens - {
+            'the', 'and', 'with', 'under', 'what', 'which',
+            'does', 'do', 'is', 'are', 'a', 'an', 'of', 'in',
+            'on', 'or', 'to', 'for', 'from', 'at', 'by', 'as',
+            'it', 'be', 'how', 'not', 'no', 'but', 'its',
+        }
+
+    def _kg_label_score(self, label: str, query: str) -> float:
+        label_text = self._kg_search_text(label)
+        query_text = self._kg_search_text(query)
+        if not label_text:
+            return 0.0
+        if label_text in query_text:
+            return 100.0 + len(label_text) / 100.0
+
+        label_tokens = self._kg_tokens(label_text)
+        query_tokens = self._kg_tokens(query_text)
+        if not label_tokens or not query_tokens:
+            return 0.0
+
+        overlap = label_tokens & query_tokens
+        if not overlap:
+            return 0.0
+
+        generic_tokens = {
+            'cof', 'cofs', 'ros', 'light', 'visible', 'photocatalysis',
+            'photocatalytic', 'phenol', 'degradation', 'generate', 'generates',
+            'dominate', 'dominates', 'dominant', 'which', 'what', 'under',
+        }
+        specific_overlap = overlap - generic_tokens
+        if not specific_overlap:
+            return 0.0
+
+        # Specific material or condition tokens such as fe3o4, tpma, tt, tempo,
+        # h2o2, and pms should outrank short generic matches such as phenol/light.
+        digit_or_formula_hits = sum(
+            1 for token in specific_overlap if any(ch.isdigit() for ch in token)
+        )
+        score = (len(specific_overlap) * 3.0) + (len(overlap) * 0.5) + (digit_or_formula_hits * 2.0)
+        return score / max(len(label_tokens) ** 0.25, 1.0)
+
+    def _find_kg_nodes(self, query: str, node_prefix: str) -> List[str]:
+        return [node for _, node in self._find_scored_kg_nodes(query, node_prefix)]
+
+    def _find_scored_kg_nodes(self, query: str, node_prefix: str) -> List[Tuple[float, str]]:
+        nodes = set()
+        for triple in self.kg_triples:
+            for field in ('head', 'tail'):
+                node = triple.get(field, '')
+                if node.startswith(node_prefix):
+                    nodes.add(node)
+
+        scored = []
+        for node in nodes:
+            score = self._kg_label_score(self._kg_label(node), query)
+            if score > 0:
+                scored.append((score, node))
+        scored.sort(key=lambda item: item[0], reverse=True)
+        return scored
+
+    @staticmethod
+    def is_domain_query(query: str) -> bool:
+        """Return whether a query is specific enough for COF/ROS retrieval."""
+        return bool(re.search(
+            r"\b(cof|cofs|ros|h2o2|o2|orr|oxygen|radical|peroxide|hydroxyl|"
+            r"superoxide|singlet|photocatal|photocatalysis|photocatalytic|"
+            r"redox|oxidation|reduction|electron|radicals?)\b",
+            query,
+            flags=re.IGNORECASE,
+        ))
+
+    @staticmethod
+    def _looks_like_kg_query(query: str) -> bool:
+        return EvidenceRetriever.is_domain_query(query)
+
+    @staticmethod
+    def _kg_fact_priority(fact: Dict[str, Any]) -> float:
+        relation_weight = {
+            'dominantly_generates': 5.0,
+            'secondarily_generates': 4.0,
+            'has_intermediate': 3.0,
+            'under_condition': 2.0,
+            'reported_in': 0.5,
+        }.get(fact.get('relation', ''), 1.0)
+        evidence = fact.get('evidence') or {}
+        if not isinstance(evidence, dict):
+            evidence = {'evidence_level': str(evidence)}
+        support_weight = {
+            'supported': 3.0,
+            'partially_supported': 2.0,
+            'unsupported': -2.0,
+        }.get(evidence.get('support_status') or fact.get('support_status', ''), 0.0)
+        level_weight = {
+            'high': 2.0,
+            'medium': 1.0,
+            'low': 0.0,
+        }.get(evidence.get('evidence_level') or fact.get('evidence_level', ''), 0.0)
+        return relation_weight + support_weight + level_weight
+
+    @staticmethod
+    def _kg_fact_support_status(fact: Dict[str, Any]) -> str:
+        evidence = fact.get('evidence') or {}
+        if isinstance(evidence, dict):
+            return str(evidence.get('support_status') or fact.get('support_status') or '').lower()
+        return str(fact.get('support_status') or '').lower()
+
+    def is_confident_kg_fact(
+        self,
+        fact: Dict[str, Any],
+        min_confidence: Optional[float] = None,
+    ) -> bool:
+        """Return whether a KG fact is suitable for RAG context injection."""
+        if self._kg_fact_support_status(fact) == 'unsupported':
+            return False
+        threshold = getattr(self, 'min_kg_confidence', 2.0) if min_confidence is None else float(min_confidence)
+        return self._kg_fact_priority(fact) >= threshold
+
+    def _filter_confident_kg_facts(
+        self,
+        facts: List[Dict[str, Any]],
+        min_confidence: Optional[float] = None,
+    ) -> List[Dict[str, Any]]:
+        if min_confidence is None:
+            return facts
+        return [
+            fact for fact in facts
+            if self.is_confident_kg_fact(fact, min_confidence=min_confidence)
+        ]
+
+    def _format_kg_fact(self, triple: Dict[str, Any]) -> Dict[str, Any]:
+        fact = dict(triple)
+        fact['head'] = self._kg_label(triple.get('head', ''))
+        fact['tail'] = self._kg_label(triple.get('tail', ''))
+        fact['source'] = 'kg_triples'
+        return fact
+
+    def _dedupe_kg_facts(self, facts: List[Dict[str, Any]], top_k: int) -> List[Dict[str, Any]]:
+        seen = set()
+        unique = []
+        for fact in sorted(facts, key=self._kg_fact_priority, reverse=True):
+            key = (
+                fact.get('head', ''),
+                fact.get('relation', ''),
+                fact.get('tail', ''),
+                json.dumps(fact.get('condition', {}), sort_keys=True, ensure_ascii=False),
+            )
+            if key in seen:
+                continue
+            seen.add(key)
+            unique.append(fact)
+            if len(unique) >= top_k:
+                break
+        return unique
+
+    def _retrieve_graph_kg_facts(
+        self,
+        query: str,
+        top_k: int,
+        min_confidence: Optional[float] = None,
+    ) -> List[Dict[str, Any]]:
+        if not self.kg_triples or not self._looks_like_kg_query(query):
+            return []
+
+        facts = []
+        generation_relations = {
+            'dominantly_generates',
+            'secondarily_generates',
+            'has_intermediate',
+            'under_condition',
+        }
+
+        reverse_species_query = bool(re.search(
+            r"\bwhich\b.*\b(?:cof|cofs|materials?)\b.*\b(?:generate|generates|produce|produces|yield|yields)\b",
+            query,
+            flags=re.IGNORECASE,
+        ))
+
+        if not reverse_species_query:
+            scored_material_nodes = self._find_scored_kg_nodes(query, 'Material:')
+            if scored_material_nodes:
+                top_score = scored_material_nodes[0][0]
+                # If a query clearly names a material, use that material neighborhood
+                # instead of mixing in weaker generic material matches.
+                selected_nodes = [scored_material_nodes[0][1]] if top_score >= 2.5 else [
+                    node for _, node in scored_material_nodes[:5]
+                ]
+                for material_node in selected_nodes:
+                    for triple in self.kg_by_head.get(material_node, []):
+                        if triple.get('relation') in generation_relations:
+                            facts.append(self._format_kg_fact(triple))
+
+        # A concrete material mention should drive the graph neighborhood. Only use
+        # reverse species lookup when no material node was found.
+        if facts:
+            facts = self._filter_confident_kg_facts(facts, min_confidence=min_confidence)
+            return self._dedupe_kg_facts(facts, top_k)
+
+        for species_node in self._find_kg_nodes(query, 'ROS_or_Product:')[:5]:
+            for triple in self.kg_by_tail.get(species_node, []):
+                if triple.get('relation') in generation_relations:
+                    facts.append(self._format_kg_fact(triple))
+
+        for species_node in self._find_kg_nodes(query, 'Intermediate:')[:5]:
+            for triple in self.kg_by_tail.get(species_node, []):
+                if triple.get('relation') in generation_relations:
+                    facts.append(self._format_kg_fact(triple))
+
+        facts = self._filter_confident_kg_facts(facts, min_confidence=min_confidence)
+        return self._dedupe_kg_facts(facts, top_k)
+
+    def get_schema_record(self, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Return Gold schema record or validated Silver extracted payload."""
+        extracted = record.get('extracted_json')
+        if isinstance(extracted, dict):
+            if extracted.get('status') == 'insufficient_evidence' or extracted.get('error'):
+                return None
+            payload = dict(extracted)
+            payload['source_chunk_id'] = record.get('chunk_id', '')
+            payload['evidence_level'] = record.get('evidence_level', 'low')
+            payload['support_status'] = record.get('support_status', 'unsupported')
+            return payload
+
+        if 'material_context' in record or 'dominant_ros_or_product' in record:
+            return record
+
+        return None
+
+    def _load_text_metadata(self) -> Dict[str, Dict[str, Any]]:
+        """加载所有元数据以便检索文本"""
+        metadata = {}
+
+        # 加载 Gold/Silver JSON 元数据
+        processed_dir = Path(self.config['paths']['processed_dir'])
+        jsonl_paths = [
+            processed_dir / 'json_records_cleaned.jsonl',
+            processed_dir / 'silver_json_validated.jsonl',
+        ]
+        for jsonl_path in jsonl_paths:
+            if not jsonl_path.exists():
+                continue
+
+            with open(jsonl_path, 'r', encoding='utf-8') as f:
+                for line in f:
+                    if line.strip():
+                        raw_record = json.loads(line)
+                        record = self.get_schema_record(raw_record)
+                        if not record:
+                            continue
+                        doc_id = record.get('source_chunk_id', raw_record.get('chunk_id', ''))
+                        if doc_id:
+                            metadata[doc_id] = {
+                                'source_type': 'gold_silver_json',
+                                'text': self._construct_json_text(record),
+                                'evidence_level': record.get('evidence_level', 'low'),
+                                **record
+                            }
+
+        # 加载 PDF 元数据
+        pdf_path = processed_dir / 'pdf_chunks.jsonl'
+        if pdf_path.exists():
+            with open(pdf_path, 'r', encoding='utf-8') as f:
+                for line in f:
+                    if line.strip():
+                        record = json.loads(line)
+                        doc_id = record.get('chunk_id', '')
+                        if doc_id:
+                            metadata[doc_id] = {
+                                'source_type': 'pdf_fulltext',
+                                'text': record.get('text', ''),
+                                **record
+                            }
+
+        # 加载 WOS 元数据
+        wos_path = processed_dir / 'wos_abstract_chunks.jsonl'
+        if wos_path.exists():
+            with open(wos_path, 'r', encoding='utf-8') as f:
+                for line in f:
+                    if line.strip():
+                        record = json.loads(line)
+                        doc_id = record.get('doc_id', '')
+                        if doc_id:
+                            metadata[doc_id] = {
+                                'source_type': 'wos_abstract',
+                                'text': record.get('text', ''),
+                                **record
+                            }
+
+        return metadata
+
+    def _construct_json_text(self, record: Dict[str, Any]) -> str:
+        """构造 JSON 记录的文本表示"""
+        parts = []
+
+        material = record.get('material_context', {})
+        if material.get('material_type'):
+            parts.append(f"Material: {material['material_type']}")
+        if material.get('active_site'):
+            parts.append(f"Active site: {material['active_site']}")
+
+        env = record.get('reaction_environment', {})
+        if env.get('phase'):
+            parts.append(f"Phase: {env['phase']}")
+        if env.get('oxidant_source'):
+            parts.append(f"Oxidant: {env['oxidant_source']}")
+        if env.get('driving_force'):
+            parts.append(f"Driving: {env['driving_force']}")
+
+        for item in record.get('dominant_ros_or_product', []):
+            if item.get('species'):
+                parts.append(f"ROS: {item['species']}")
+
+        for step in record.get('mechanism', []):
+            if step:
+                parts.append(f"Mechanism: {step}")
+
+        return ' '.join(parts)
+
+    def _load_metadata_dict(self) -> Dict[str, Dict[str, Any]]:
+        """加载所有元数据为字典"""
+        if hasattr(self, '_metadata_dict_cache'):
+            return self._metadata_dict_cache
+
+        metadata = self._load_text_metadata()
+
+        self._metadata_dict_cache = metadata
+        return metadata
+
+    @staticmethod
+    def _truncate_at_sentence(text: str, max_chars: int = 500) -> str:
+        """Truncate *text* at the last sentence boundary before *max_chars*."""
+        if len(text) <= max_chars:
+            return text
+        # Find the last sentence terminator within the limit.
+        snippet = text[:max_chars]
+        for sep in ('. ', '.\n', '? ', '?\n', '! ', '!\n', '。', '；'):
+            idx = snippet.rfind(sep)
+            if idx > max_chars * 0.5:  # don't cut too early
+                return snippet[:idx + len(sep.rstrip())]
+        # Fallback: truncate at the last space.
+        last_space = snippet.rfind(' ')
+        if last_space > max_chars * 0.5:
+            return snippet[:last_space]
+        return snippet
+
+    def normalize_ros(self, name: str) -> str:
+        """标准化 ROS 名称"""
+        if not name:
+            return ""
+        name = name.strip()
+        return self.ros_normalization.get(name, self.ros_normalization.get(name.lower(), name))
+
+    def tokenize(self, text: str) -> List[str]:
+        """Use the same tokenizer as index construction."""
+        return tokenize_search_text(text)
+
+    def is_confident_result(self, result: Dict[str, Any]) -> bool:
+        """Return whether a text result clears at least one retrieval confidence gate."""
+        return (
+            result.get('bm25_score', 0.0) >= self.min_bm25_score
+            or result.get('vector_score', 0.0) >= self.min_vector_score
+        )
+
+    def _is_confident_result(self, result: Dict[str, Any]) -> bool:
+        """Backward-compatible alias for older callers/tests."""
+        return self.is_confident_result(result)
+
+    def hybrid_retrieve(self, query: str, top_k: int = 10) -> List[Dict[str, Any]]:
+        """混合检索"""
+        if not self.is_domain_query(query):
+            return []
+
+        results = []
+
+        # 1. 向量检索
+        vector_results = self._vector_search(query, self.vector_top_k)
+
+        # 2. BM25 检索
+        bm25_results = self._bm25_search(query, self.bm25_top_k)
+
+        # 3. 合并结果
+        combined = self._merge_results(vector_results, bm25_results)
+
+        # 4. 重排
+        reranked = self._rerank_results(combined, query)
+
+        # 5. 丢弃低置信度结果，避免无关问题也被强制注入 RAG 上下文
+        confident = [result for result in reranked if self.is_confident_result(result)]
+        return confident[:top_k]
+
+    def _vector_search(self, query: str, top_k: int) -> List[Tuple[int, float]]:
+        """向量搜索"""
+        if self.faiss_index is None or SentenceTransformer is None:
+            return []
+
+        try:
+            # 嵌入查询
+            embedding_model = self.config.get('models', {}).get('embedding_model', '')
+            if not embedding_model:
+                return []
+
+            model = SentenceTransformer(embedding_model)
+            query_vec = model.encode([query], normalize_embeddings=True).astype('float32')
+
+            # 搜索
+            D, I = self.faiss_index.search(query_vec, top_k)
+
+            results = []
+            for i, idx in enumerate(I[0]):
+                if idx >= 0 and idx < len(self.faiss_metadata):
+                    results.append((idx, float(D[0][i])))
+
+            return results
+
+        except Exception as e:
+            logging.error(f"Vector search error: {e}")
+            return []
+
+    def _bm25_search(self, query: str, top_k: int) -> List[Tuple[int, float]]:
+        """BM25 搜索"""
+        if BM25Okapi is None or not self.bm25_metadata:
+            return []
+
+        try:
+            # 分词查询
+            query_tokens = self.tokenize(query)
+
+            # 计算 BM25 得分
+            scores = self.bm25_index.get_scores(query_tokens)
+
+            # 排序
+            results = [(i, float(score)) for i, score in enumerate(scores)]
+            results.sort(key=lambda x: x[1], reverse=True)
+
+            return results[:top_k]
+
+        except Exception as e:
+            logging.error(f"BM25 search error: {e}")
+            return []
+
+    def _merge_results(self, vector_results: List[Tuple[int, float]],
+                       bm25_results: List[Tuple[int, float]]) -> List[Dict[str, Any]]:
+        """合并检索结果，对 BM25 分数做归一化后再组合。"""
+        merged = {}
+        metadata_dict = self._load_metadata_dict()
+
+        # Normalise BM25 scores to [0, 1] so they combine fairly with vector scores.
+        bm25_max = max((s for _, s in bm25_results), default=1.0)
+        bm25_max = max(bm25_max, 1.0)  # avoid division by zero
+
+        # 处理向量结果
+        for idx, score in vector_results:
+            if idx < len(self.faiss_metadata):
+                meta = self.faiss_metadata[idx]
+                doc_id = meta.get('doc_id') or f"faiss_{idx}"
+
+                if doc_id not in merged:
+                    merged[doc_id] = {
+                        'doc_id': doc_id,
+                        'vector_score': score,
+                        'bm25_score': 0.0,
+                        'combined_score': score,
+                        'metadata': meta
+                    }
+                else:
+                    merged[doc_id]['vector_score'] = score
+                    merged[doc_id]['combined_score'] += score
+
+        # 处理 BM25 结果
+        for idx, score in bm25_results:
+            if idx < len(self.bm25_metadata):
+                meta = self.bm25_metadata[idx]
+                doc_id = meta.get('doc_id') or f"bm25_{idx}"
+                norm_score = score / bm25_max
+
+                if doc_id not in merged:
+                    merged[doc_id] = {
+                        'doc_id': doc_id,
+                        'vector_score': 0.0,
+                        'bm25_score': score,
+                        'combined_score': norm_score,
+                        'metadata': meta
+                    }
+                else:
+                    merged[doc_id]['bm25_score'] = score
+                    merged[doc_id]['combined_score'] += norm_score
+
+        # 添加来自 metadata_dict 的完整数据
+        for doc_id, result in merged.items():
+            if doc_id in metadata_dict:
+                result['full_data'] = metadata_dict[doc_id]
+
+        return list(merged.values())
+
+    def _rerank_results(self, results: List[Dict[str, Any]], query: str) -> List[Dict[str, Any]]:
+        """重排结果"""
+        query_lower = query.lower()
+        query_tokens = set(self.tokenize(query))
+
+        # 提取查询中的关键词
+        ros_keywords = ['·oh', 'o2·-', '1o2', 'h2o2', 'so4·-', 'hydroxyl', 'superoxide', 'singlet oxygen', 'hydrogen peroxide', 'sulfate radical']
+        condition_keywords = [' photocatalysis', 'light', 'pms', 'pds', '2e-', 'orr', 'electrochemical']
+
+        # 重排逻辑
+        def score_result(result: Dict[str, Any]) -> float:
+            score = result['combined_score']
+
+            # 来源优先级
+            metadata = result.get('metadata', {})
+            source_type = metadata.get('source_type', '')
+
+            source_bonus = self.source_priority.get(source_type, 1)
+            score += 0.1 * source_bonus
+
+            # 材料匹配
+            if 'full_data' in result:
+                full_data = result['full_data']
+
+                # 检查 ROS 匹配
+                requested_ros = [ros for ros in ros_keywords if ros in query_lower]
+                if requested_ros:
+                    for item in full_data.get('dominant_ros_or_product', []):
+                        species = item.get('species', '').lower()
+                        if any(ros in species or species in ros for ros in requested_ros):
+                            score += 2.0
+                            break
+
+                # 检查条件匹配
+                env = full_data.get('reaction_environment', {})
+                if any(kw in query_lower for kw in condition_keywords):
+                    if env.get('phase') or env.get('driving_force'):
+                        score += 1.0
+
+                # 检查材料匹配
+                material = full_data.get('material_context', {})
+                material_text = material.get('material_type', '').lower()
+                material_tokens = set(self.tokenize(material_text))
+                specific_query_tokens = query_tokens - {'cof', 'cofs'}
+                if specific_query_tokens & material_tokens:
+                    score += 1.5
+
+            result['rerank_score'] = score
+            return score
+
+        # 排序
+        results.sort(key=score_result, reverse=True)
+
+        return results
+
+    def retrieve_with_sources(self, query: str, top_k: int = 10) -> List[Dict[str, Any]]:
+        """检索并返回带完整来源信息的结果"""
+        results = self.hybrid_retrieve(query, top_k * 2)  # 检索更多结果用于重排
+
+        # 构建返回结果
+        output = []
+        for result in results:
+            full_data = result.get('full_data') or result.get('metadata', {})
+
+            output.append({
+                'doc_id': result['doc_id'],
+                'score': result.get('rerank_score', result['combined_score']),
+                'vector_score': result['vector_score'],
+                'bm25_score': result['bm25_score'],
+                'source_type': result['metadata'].get('source_type', 'unknown'),
+                'evidence_level': full_data.get('evidence_level', 'low'),
+                'material_context': full_data.get('material_context', {}),
+                'reaction_environment': full_data.get('reaction_environment', {}),
+                'dominant_ros_or_product': full_data.get('dominant_ros_or_product', []),
+                'secondary_or_intermediate_ros': full_data.get('secondary_or_intermediate_ros', []),
+                'mechanism': full_data.get('mechanism', []),
+                'text': full_data.get('text', ''),
+                'source_chunk_id': full_data.get('source_chunk_id', '')
+            })
+
+        return output[:top_k]
+
+    def _retrieve_bm25_kg_facts(self, query: str, top_k: int = 10) -> List[Dict[str, Any]]:
+        """Fallback: derive compact facts from confident text retrieval results."""
+        results = self.hybrid_retrieve(query, top_k)
+
+        kg_facts = []
+
+        for result in results:
+            full_data = result.get('full_data', {})
+            if str(full_data.get('support_status', '')).lower() == 'unsupported':
+                continue
+
+            material = full_data.get('material_context', {})
+            if material.get('material_type'):
+                kg_facts.append({
+                    'head': material.get('material_type', ''),
+                    'relation': 'has_condition',
+                    'tail': ', '.join(filter(None, [
+                        full_data.get('reaction_environment', {}).get('phase', ''),
+                        full_data.get('reaction_environment', {}).get('driving_force', '')
+                    ])),
+                    'evidence': full_data.get('evidence_level', 'low'),
+                    'source': 'bm25_fallback',
+                })
+
+            for item in full_data.get('dominant_ros_or_product', []):
+                if item.get('species'):
+                    kg_facts.append({
+                        'head': material.get('material_type', ''),
+                        'relation': 'generates',
+                        'tail': item.get('species', ''),
+                        'category': item.get('category', ''),
+                        'evidence': full_data.get('evidence_level', 'low'),
+                        'source': 'bm25_fallback',
+                    })
+
+        return kg_facts
+
+    def retrieve_kg_facts(
+        self,
+        query: str,
+        top_k: int = 10,
+        min_confidence: Optional[float] = None,
+    ) -> List[Dict[str, Any]]:
+        """Retrieve confidence-gated KG facts from graph first, then BM25 fallback."""
+        threshold = getattr(self, 'min_kg_confidence', 2.0) if min_confidence is None else float(min_confidence)
+        graph_facts = self._retrieve_graph_kg_facts(
+            query,
+            top_k,
+            min_confidence=threshold,
+        )
+        if graph_facts:
+            return graph_facts
+        return self._retrieve_bm25_kg_facts(query, top_k)
+
+    def run(self):
+        """运行检索测试"""
+        logging.info("Evidence Retriever initialized")
+        logging.info(f"FAISS vectors: {len(self.faiss_metadata) if self.faiss_metadata else 0}")
+        logging.info(f"BM25 documents: {len(self.bm25_metadata) if self.bm25_metadata else 0}")
+
+
+def main():
+    """主函数"""
+    import argparse
+
+    parser = argparse.ArgumentParser(description='Retrieve evidence for COF-ROS queries')
+    parser.add_argument('--config', type=str, default='config/config.yaml',
+                        help='Path to config file')
+    parser.add_argument('--query', type=str, default=None,
+                        help='Query to retrieve (if not provided, run in test mode)')
+    parser.add_argument('--top-k', type=int, default=5,
+                        help='Number of results to return')
+    args = parser.parse_args()
+
+    # 加载配置
+    config_path = Path(args.config)
+    with open(config_path, 'r', encoding='utf-8') as f:
+        config = yaml.safe_load(f)
+
+    # 初始化检索器
+    retriever = EvidenceRetriever(config)
+
+    if args.query:
+        # 检索查询
+        results = retriever.retrieve_with_sources(args.query, args.top_k)
+
+        for i, result in enumerate(results, 1):
+            print(f"\n--- Result {i} (score: {result['score']:.3f}) ---")
+            print(f"Source: {result['source_type']}")
+            print(f"Material: {result['material_context'].get('material_type', 'N/A')}")
+            print(f"Dominant ROS: {[item['species'] for item in result['dominant_ros_or_product']]}")
+            print(f"Text preview: {result['text'][:200]}...")
+    else:
+        retriever.run()
+
+
+if __name__ == '__main__':
+    main()