1 line
6.0 KiB
JSON
1 line
6.0 KiB
JSON
{"repo": "3434912374/aiagent", "files": {"20260316/main.py": {"51064ef9d27c38c08b00ab49b0dcda644266a021": {"t": 1785746434.8043754, "n": 4245, "h": "11bf6e57ea624884", "c": "import yaml\nimport time\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.messages import HumanMessage, SystemMessage\n\nfrom rag_engine.embeddings import DenseEmbedder\nfrom rag_engine.retrievers import DualRetriever\nfrom rag_engine.fusion import reciprocal_rank_fusion\nfrom rag_engine.reranker import DeepReranker\nfrom rl_scheduler.ppo_agent import PPOScheduler\nfrom workflow.telemetry import setup_telemetry\nfrom workflow.dag_engine import DAGWorkflow\n\n# --- 新增：查询重写（Query Rewrite）引擎 ---\ndef rewrite_query(original_query, api_key):\n \"\"\"\n 检索前置拦截器（收紧版）：严格限制输出长度，防止大模型词汇爆炸。\n \"\"\"\n print(f\" [重写引擎] 正在请 DeepSeek 提取核心检索词...\")\n llm = ChatOpenAI(\n model=\"deepseek-chat\", \n temperature=0, # 必须设为 0，让它绝对冷静，不发散\n api_key=api_key,\n base_url=\"https://api.deepseek.com/v1\" \n )\n \n # 【核心修改点】：给 Prompt 加上极度严格的限制\n prompt = f\"\"\"你是一个搜索引擎优化专家。\n请从以下用户的口语化提问中，提取最重要的 3 个技术关键词。如果涉及到寻找特定工具，你可以直接补充 1 个最著名的开源工具名。\n\n用户的原始提问：“{original_query}”\n\n【严格要求】：\n1. 你的输出只能包含 3-4 个词语。\n2. 词语之间用空格隔开。\n3. 绝对不要输出任何标点符号、解释性文字或其他废话。\n\"\"\"\n messages = [\n SystemMessage(content=\"你是一个严格的关键词提取器。输出不能超过 4 个词。\"),\n HumanMessage(content=prompt)\n ]\n \n expanded_query = llm.invoke(messages).content.strip()\n print(f\" [重写引擎] 原始问题: {original_query}\")\n print(f\" [重写引擎] 精炼检索词: {expanded_query}\")\n return expanded_query\n# ----------------------------------------\n\ndef generate_answer(query, retrieved_docs, api_key):\n context = \"\\n\".join([f\"片段 {i+1}: {doc['text']}\" for i, doc in enumerate(retrieved_docs)])\n llm = ChatOpenAI(\n model=\"deepseek-chat\", \n temperature=0, \n api_key=api_key,\n base_url=\"https://api.deepseek.com/v1\" \n )\n prompt = f\"\"\"你是一个专业的企业级 AI 助手。请严格基于以下【检索到的参考知识】来回答用户的问题。\n如果知识库中没有相关内容，请直接回答“根据提供的知识库，我无法准确回答此问题”，不要胡编乱造。\n\n【检索到的参考知识】：\n{context}\n\n【用户问题】：\n{query}\n\"\"\"\n messages = [\n SystemMessage(content=\"你是一个精准、严谨的知识库问答助手。请始终用中文回答。\"),\n HumanMessage(content=prompt)\n ]\n print(\"🧠 [大模型] 正在呼叫 DeepSeek 阅读检索结果并生成最终回答...\")\n return llm.invoke(messages).content\n\n\ndef main():\n with open(\"config.yaml\", \"r\", encoding=\"utf-8\") as f:\n cfg = yaml.safe_load(f)\n\n print(\"====== 正在初始化 企业级 Deep RAG 智能系统 ======\")\n \n embedder = DenseEmbedder(cfg['rag']['dense_model'])\n retriever = DualRetriever(embedder)\n reranker = DeepReranker(cfg['rag']['rerank_model'])\n \n corpus = [\n \"PPO (近端策略优化) 是一种非常稳定的强化学习算法，适合用于智能体的训练。\",\n \"FAISS 是由 Facebook 开发的，专门用于高效的密集向量相似度搜索。\",\n \"BM25 是一种传统的稀疏检索算法，它主要基于词频(TF)和逆文档频率(IDF)的统计。\",\n \"RRF (倒数秩融合) 算法可以有效地把不同搜索引擎（比如向量搜索和关键词搜索）的结果混合在一起。\",\n \"Cross-Encoders（交叉编码器）通过自注意力机制同时处理用户问题和文档，打分极其精准，但速度较慢。\"\n ]\n retriever.build_index(corpus)\n \n scheduler = PPOScheduler(cfg['scheduler']['num_nodes'], cfg['scheduler']['ppo_timesteps'])\n tracer = setup_telemetry()\n workflow = DAGWorkflow(tracer, scheduler, retriever, reciprocal_rank_fusion, reranker)\n\n print(\"\\n====== 系统就绪，开始执行智能问答 ======\")\n DEEPSEEK_API_KEY = \"sk-c45174ee28f847f8bf59669c79cc7383\"\n \n queries = [\n \"如果我想做高效的向量搜索，有什么工具推荐吗？是哪个公司开发的？\"\n ]\n\n for q in queries:\n print(f\"\\n======================================\")\n print(f\"🙋♂️ 收到用户原始提问: '{q}'\")\n print(f\"======================================\")\n \n start_time = time.time()\n \n # 【关键改动】：在送入 DAG 工作流之前，先重写 Query\n with tracer.start_as_current_span(\"Query_Rewrite\"):\n optimized_query = rewrite_query(q, DEEPSEEK_API_KEY)\n \n # 使用优化后的超级检索词去本地知识库“找答案”\n results = workflow.execute(optimized_query, corpus, cfg['rag']['top_k'], cfg['rag']['rrf_k'])\n \n print(\"\\n🎯 [内部检索结果] 找到了以下最相关的本地知识：\")\n for rank, res in enumerate(results):\n print(f\" [{rank+1}] 相关度得分: {res['score']:.4f} | 内容: {res['text']}\")\n \n with tracer.start_as_current_span(\"LLM_Generation\"):\n # 注意：生成回答时，依然传入用户原始问题 (q)，以保证回答语气自然\n final_answer = generate_answer(q, results, DEEPSEEK_API_KEY)\n \n latency = (time.time() - start_time)\n \n print(f\"\\n🤖 DeepSeek 最终回复 (耗时 {latency:.2f} 秒):\")\n print(f\"\\n{final_answer}\\n\")\n\nif __name__ == \"__main__\":\n main()"}}}} |