FinRAG-Agent/FinRAG-Pro

1"""2金融投研RAG系统 - RAG检索与生成模块（终极防弹合体版）3工作流：双线检索 → 智能清洗 → 上下文融合 → 终极生成4已修复初始化死锁问题！5"""6 7import os8import json9import pickle10import numpy as np11import datetime12 13BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))14INDEX_DIR = os.path.join(BASE_DIR, "data", "index")15PROCESSED_DIR = os.path.join(BASE_DIR, "data", "processed")16TOP_K = 417 18class RAGPipeline:19    def __init__(self):20        self.index = None21        self.chunks = None22        self.metadata = None23        self.encoder = None24        self.embedding_method = "unknown"25        self.initialized = False26 27    def initialize(self):28        """初始化：强力重试机制"""29        try:30            import faiss31            index_path = os.path.join(INDEX_DIR, "faiss_index.bin")32            if not os.path.exists(index_path):33                raise FileNotFoundError(f"FAISS索引文件不存在: {index_path}")34 35            self.index = faiss.read_index(index_path)36            37            chunks_path = os.path.join(PROCESSED_DIR, "chunks.json")38            with open(chunks_path, "r", encoding="utf-8") as f:39                self.chunks = json.load(f)40 41            meta_path = os.path.join(PROCESSED_DIR, "chunk_metadata.json")42            with open(meta_path, "r", encoding="utf-8") as f:43                self.metadata = json.load(f)44 45            self._load_encoder()46            47            self.initialized = True48            print(f"🚀 RAG引擎初始化成功！已加载 {self.index.ntotal} 条向量。")49            return True50        except Exception as e:51            self.initialized = False52            print(f"❌ RAG初始化中止: {e}")53            raise e54 55    def _load_encoder(self):56        tfidf_path = os.path.join(INDEX_DIR, "tfidf_vectorizer.pkl")57        if os.path.exists(tfidf_path):58            from sklearn.feature_extraction.text import TfidfVectorizer59            with open(tfidf_path, "rb") as f:60                self.encoder = pickle.load(f)61            self.embedding_method = "tfidf"62        else:63            from sentence_transformers import SentenceTransformer64            self.encoder = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")65            self.embedding_method = "sentence-transformers"66 67    def _encode(self, texts: list[str]) -> np.ndarray:68        if self.embedding_method == "sentence-transformers":69            return self.encoder.encode(texts, show_progress_bar=False).astype(np.float32)70        else:71            return self.encoder.transform(texts).toarray().astype(np.float32)72 73    # ==================== 第一步：本地检索 ====================74    def retrieve(self, query: str, top_k: int = TOP_K) -> list[dict]:75        if not self.initialized:76            print("⚠️ 检测到引擎未就绪，尝试紧急初始化...")77            self.initialize()78 79        import faiss80        query_vec = self._encode([query])81        faiss.normalize_L2(query_vec)82        scores, indices = self.index.search(query_vec, top_k)83 84        results = []85        for i, (score, idx) in enumerate(zip(scores[0], indices[0])):86            if idx < 0 or idx >= len(self.chunks): continue87            meta = self.metadata[idx] if self.metadata else {}88            results.append({89                "rank": i + 1,90                "score": float(score),91                "chunk_id": meta.get("chunk_id", f"chunk_{idx}"),92                "doc_title": meta.get("doc_title", "未知"),93                "source": meta.get("source", "未知"),94                "source_type": meta.get("source_type", "未知来源"),95                "source_dir": meta.get("source_dir", ""),96                "content": self.chunks[idx],97            })98        return results99 100    # ==================== 第二步：Tavily全网搜索 ====================101    def search_web_tavily(self, query: str, tavily_api_key: str) -> dict:102        if not tavily_api_key or tavily_api_key.startswith("tvly-你的"):103            return {"success": False, "raw_text": "", "results": [], "error": "Tavily API Key未配置"}104 105        try:106            from tavily import TavilyClient107            client = TavilyClient(api_key=tavily_api_key)108            response = client.search(query=query, search_depth="advanced", max_results=5)109 110            results = response.get("results", [])111            if not results:112                return {"success": True, "raw_text": "全网未检索到相关内容", "results": []}113 114            raw_text = ""115            formatted = []116            for idx, r in enumerate(results):117                title = r.get("title", "无标题")118                content = r.get("content", "")119                url = r.get("url", "")120                raw_text += f"【线索{idx+1}】标题: {title}\n内容: {content}\n来源: {url}\n\n"121                formatted.append({"title": title, "content": content[:500], "url": url})122 123            return {"success": True, "raw_text": raw_text, "results": formatted}124        except Exception as e:125            return {"success": False, "raw_text": "", "results": [], "error": f"Tavily搜索失败: {str(e)}"}126 127    # ==================== 第三步：噪音过滤判别器 ====================128    def impact_filter(self, raw_news: str, deepseek_api_key: str, api_type: str = "deepseek", model: str = "", base_url: str = "") -> dict:129        if not raw_news.strip(): return {"success": True, "filtered_text": "", "kept_count": 0, "total_count": 0}130        if not deepseek_api_key: return {"success": False, "filtered_text": raw_news, "kept_count": 0, "total_count": 0, "error": "DeepSeek API Key未配置"}131 132        today = datetime.datetime.now().strftime("%Y年%m月%d日")133        system_prompt = f"""当前真实时间是 {today}。你是一个冷酷的金融噪音过滤机器。134请你严格评估用户提供的实时搜索资讯的'市场情绪影响因子'（满分 10 分）。135任务要求：逐条评估，直接剔除低于 6 分的资讯！只输出保留下来的高价值信息。"""136 137        if not base_url: base_url = "https://api.deepseek.com" if api_type == "deepseek" else "https://api.siliconflow.cn/v1"138        if not model: model = "deepseek-chat" if api_type == "deepseek" else "deepseek-ai/DeepSeek-V2.5"139 140        try:141            import requests142            url = f"{base_url.rstrip('/')}/chat/completions"143            headers = {"Authorization": f"Bearer {deepseek_api_key}", "Content-Type": "application/json"}144            payload = {"model": model, "messages": [{"role": "system", "content": system_prompt}, {"role": "user", "content": raw_news}], "temperature": 0.1, "max_tokens": 2048}145            response = requests.post(url, headers=headers, json=payload, timeout=60)146 147            if response.status_code == 200:148                filtered = response.json()["choices"][0]["message"]["content"]149                kept = filtered.count("【线索") if "【线索" in filtered else (1 if len(filtered) > 50 else 0)150                total = raw_news.count("【线索")151                return {"success": True, "filtered_text": filtered, "kept_count": kept, "total_count": total}152            else:153                return {"success": False, "filtered_text": raw_news, "kept_count": 0, "total_count": 0, "error": f"HTTP {response.status_code}"}154        except Exception as e:155            return {"success": False, "filtered_text": raw_news, "kept_count": 0, "total_count": 0, "error": str(e)}156 157    # ==================== Prompt构建（支持双源上下文） ====================158    def build_fusion_prompt(self, query: str, local_chunks: list[dict], web_filtered: str, web_formatted: list[dict] = None) -> str:159        local_context = ""160        if local_chunks:161            for chunk in local_chunks:162                local_context += f"[来源类型: {chunk.get('source_type', '未知')} | 文件: {chunk.get('source', '未知')}]\n内容: {chunk['content']}\n\n"163 164        web_context = web_filtered if web_filtered else "（本次未获取到全网高价值资讯）"165        166        seen_local = set()167        local_sources_list = []168        for c in local_chunks:169            key = (c.get("source_type", "未知"), c.get("source", ""))170            if key not in seen_local:171                local_sources_list.append(f"📄 [{key[0]}] {key[1]}")172                seen_local.add(key)173 174        web_sources_list = [f"🌐 {w.get('title','未知')} - {w.get('url','')}" for w in (web_formatted or [])]175        all_sources = "\n".join(local_sources_list + web_sources_list)176        today = datetime.datetime.now().strftime("%Y年%m月%d日")177 178        prompt = f"""# Role179【最高指令：当前的真实世界时间是 {today}！你的所有分析、估值和预测都必须基于这个时间点！如果检索到的研报或资讯严重滞后（例如2025年的旧闻），你必须在回答中严厉指出“数据滞后”，并拒绝将其作为最新行情的判断依据！】180你是一位拥有10年以上经验的资深金融分析师，精通企业财报分析、宏观经济研究与行业研报解读。181 182# 参考资料说明183一、【本地深度研报】（来自上市公司年报/行业研究报告）184{local_context}185二、【全网实时高价值资讯】（经过AI噪音过滤筛选后的最新动态）186{web_context}187 188# 用户问题189{query}190 191# 回答要求1921. 绝对忠于原文：财务数据必须严格基于知识库片段，不得捏造。1932. 标注来源：回答中必须使用格式「[来源类型: xxx | 文件: xxx]」标注每个数据的出处。1943. 交叉验证：发现官方年报预期与第三方客观分析冲突时，重点指出【认知预期差】。1954. 强制风险提示：在所有回答的最后，必须单列一段明确的投资风险提示。196 197# 回答格式198- **核心结论**199- **数据与分析**200- **认知预期差**201- **风险提示**202 203# 参考来源总清单204{all_sources}205"""206        return prompt207 208    # ==================== LLM调用 ====================209    def query_llm(self, prompt: str, api_type: str, api_key: str, model: str = "", base_url: str = "") -> str:210        if not api_key: return "[错误] 未配置API密钥"211        if not base_url: base_url = {"deepseek": "https://api.deepseek.com", "siliconflow": "https://api.siliconflow.cn/v1", "openai": "https://api.openai.com/v1"}.get(api_type, "https://api.deepseek.com")212        if not model: model = {"deepseek": "deepseek-chat", "siliconflow": "deepseek-ai/DeepSeek-V2.5", "openai": "gpt-3.5-turbo"}.get(api_type, "deepseek-chat")213 214        try:215            import requests216            url = f"{base_url.rstrip('/')}/chat/completions"217            headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}218            payload = {"model": model, "messages": [{"role": "user", "content": prompt}], "temperature": 0.1, "max_tokens": 4096}219            response = requests.post(url, headers=headers, json=payload, timeout=90)220            if response.status_code == 200:221                return response.json()["choices"][0]["message"]["content"]222            return f"[错误] API请求失败 (HTTP {response.status_code}): {response.text[:300]}"223        except Exception as e:224            return f"[错误] 请求异常: {str(e)}"225 226# ==================== 单例模式修复 ====================227_rag_pipeline = None228 229def get_rag_pipeline() -> "RAGPipeline":230    global _rag_pipeline231    if _rag_pipeline is None:232        p = RAGPipeline()233        p.initialize()234        _rag_pipeline = p235    elif not _rag_pipeline.initialized:236        _rag_pipeline.initialize()237    return _rag_pipeline