diff --git a/src/agent/chatbot_agent.py b/src/agent/chatbot_agent.py index 380e29d..a24329e 100644 --- a/src/agent/chatbot_agent.py +++ b/src/agent/chatbot_agent.py @@ -1,7 +1,7 @@ # agent/chatbot_agent.py -from typing import Any, Dict +from typing import Any -from langchain_core.messages import HumanMessage, SystemMessage, AIMessage +from langchain_core.messages import HumanMessage, AIMessage from langchain.agents import create_agent from agent.llm import Qwen36_35B_A3B_FP8_LLM diff --git a/src/agent/tools/tool_rag.py b/src/agent/tools/tool_rag.py index b34701e..dce1e2a 100644 --- a/src/agent/tools/tool_rag.py +++ b/src/agent/tools/tool_rag.py @@ -4,10 +4,8 @@ from __future__ import annotations import hashlib import json import os -import tempfile -import time from datetime import datetime -from typing import List, Optional, Union, Literal +from typing import List, Optional, Literal, Union import chromadb import requests @@ -17,13 +15,14 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter from utils.env_utils import RAG_BASE_URL, RAG_CHUNK_SIZE, RAG_CHUNK_OVERLAP -# ---------- 新增 RAG 工具 ---------- # Chroma 持久化路径：agent/memory/chroma_rag_db CHROMA_PERSIST_DIR = os.path.join(os.path.dirname(__file__), '..', 'memory', 'chroma_rag_db') -# 全局 Chroma 客户端和集合缓存 +# 全局 Chroma 客户端和统一集合 _chroma_client = None -_collection_cache = {} +_global_collection = None + +GLOBAL_COLLECTION_NAME = "rag_global" def get_chroma_client(): global _chroma_client @@ -31,21 +30,21 @@ def get_chroma_client(): _chroma_client = chromadb.PersistentClient(path=CHROMA_PERSIST_DIR) return _chroma_client -def get_collection(agent_name: str): - """获取或创建某个 agent 的 Collection""" - client = get_chroma_client() - collection_name = f"rag_{agent_name}" - if collection_name not in _collection_cache: - # 若 Collection 不存在则创建，使用余弦相似度 +def get_global_collection(): + """获取或创建全局知识库 Collection""" + global _global_collection + if _global_collection is None: + client = get_chroma_client() try: - collection = client.get_collection(name=collection_name, embedding_function=None) + _global_collection = client.get_collection( + name=GLOBAL_COLLECTION_NAME, embedding_function=None + ) except Exception: - collection = client.create_collection( - name=collection_name, + _global_collection = client.create_collection( + name=GLOBAL_COLLECTION_NAME, metadata={"hnsw:space": "cosine"} ) - _collection_cache[collection_name] = collection - return _collection_cache[collection_name] + return _global_collection # ---------- Embedding 接口封装 ---------- def embed_texts(queries: List[str], documents: List[str] = None) -> tuple: @@ -165,21 +164,19 @@ class RAGTool(BaseTool): description: str = ( "通用知识库工具，支持增删查操作。\n" "参数说明：\n" - "- agent_name: 所属 Agent 名称或名称列表，必填。\n" "- action_type: 操作类型，必填。可选：'query'(查询), 'add_document'(添加文档), 'clear'(清除)。\n" "- content: 查询内容(用于query) 或 直接文本内容(用于add_document，与source二选一)。\n" "- source: 文件路径(用于add_document，与content二选一)。\n" "- filter: 过滤条件(用于query或clear)。\n" "使用示例：\n" - "- 查询：action_type='query', content='用户问题', agent_name='main'\n" - "- 添加文本：action_type='add_document', content='知识文本', agent_name='main'\n" - "- 添加文件：action_type='add_document', source='/path/to/file.pdf', agent_name=['agent1','agent2']\n" - "- 清除：action_type='clear', filter='某个来源', agent_name='main'" + "- 查询：action_type='query', content='用户问题'\n" + "- 添加文本：action_type='add_document', content='知识文本'\n" + "- 添加文件：action_type='add_document', source='/path/to/file.pdf'\n" + "- 清除：action_type='clear', filter='某个来源'" ) args_schema: type[BaseModel] = create_model( "RAGInput", - agent_name=(Union[str, List[str]], Field(..., description="所属 Agent 名称或名称列表")), action_type=(Literal["query", "add_document", "clear"], Field(..., description="操作类型")), content=(Optional[str], Field(None, description="查询内容或直接添加的文本内容")), source=(Optional[str], Field(None, description="文件路径，用于添加文档")), @@ -188,29 +185,25 @@ class RAGTool(BaseTool): def _run( self, - agent_name: Union[str, List[str]], action_type: str, content: Optional[str] = None, source: Optional[str] = None, filter: Optional[Union[str, List[str]]] = None ) -> str: try: - # 标准化 agent_name 为列表 - agents = [agent_name] if isinstance(agent_name, str) else agent_name - if action_type == "query": - return self._handle_query(agents, content, filter) + return self._handle_query(content, filter) elif action_type == "add_document": - return self._handle_add(agents, content, source) + return self._handle_add(content, source) elif action_type == "clear": - return self._handle_clear(agents, filter) + return self._handle_clear(filter) else: return f"未知操作: {action_type}" except Exception as e: return f"RAG 工具执行出错: {str(e)}" # ---------- 查询 ---------- - def _handle_query(self, agents: List[str], query: str, filter: Optional[Union[str, List[str]]]) -> str: + def _handle_query(self, query: str, filter: Optional[Union[str, List[str]]]) -> str: if not query: return "查询内容不能为空。" @@ -218,50 +211,42 @@ class RAGTool(BaseTool): q_embs, _ = embed_texts(queries=[query]) query_vector = q_embs[0] - # 2. 从相关 Collection 中检索 - all_results = [] - for agent in agents: - coll = get_collection(agent) - # 构建 where 条件（基础过滤可叠加元数据过滤，这里简单处理） - where = None - if filter: - if isinstance(filter, str): - # 简单实现：按 source 字段包含 filter 字符串 - where = {"source": {"$contains": filter}} - else: - # 列表情况，暂不支持，返回空 - pass + # 2. 从全局 Collection 中检索 + coll = get_global_collection() + where = None + if filter: + if isinstance(filter, str): + where = {"source": {"$contains": filter}} + # 列表暂简化为只支持单个字符串，但保留接口 + elif isinstance(filter, list) and len(filter) > 0: + where = {"source": {"$contains": filter[0]}} - res = coll.query( - query_embeddings=[query_vector], - n_results=10, - where=where - ) - # 整理结果：[(id, document, metadata), ...] - if res["ids"] and res["ids"][0]: - ids = res["ids"][0] - docs = res["documents"][0] - metas = res["metadatas"][0] - for i, d, m in zip(ids, docs, metas): - all_results.append({"id": i, "document": d, "metadata": m}) + res = coll.query( + query_embeddings=[query_vector], + n_results=10, + where=where + ) - if not all_results: + if not res["ids"] or not res["ids"][0]: return "未在知识库中找到相关内容。" - # 去重（同一个文档可能被多个 agent 同时索引） + ids = res["ids"][0] + docs = res["documents"][0] + metas = res["metadatas"][0] + + all_results = [] seen = set() - unique_results = [] - for r in all_results: - if r["document"] not in seen: - seen.add(r["document"]) - unique_results.append(r) - unique_results = unique_results[:10] # 确保最多10条 + for i, d, m in zip(ids, docs, metas): + if d not in seen: + seen.add(d) + all_results.append({"id": i, "document": d, "metadata": m}) + + unique_results = all_results[:10] # 3. Reranker 精排 docs_for_rerank = [r["document"] for r in unique_results] scores = rerank_documents(query, docs_for_rerank) - # 按分数降序取前3 sorted_items = sorted( zip(unique_results, scores), key=lambda x: x[1], @@ -271,7 +256,6 @@ class RAGTool(BaseTool): if not sorted_items: return "Reranker 未返回有效结果。" - # 拼接输出 outputs = [] for item, score in sorted_items: meta = item["metadata"] @@ -281,7 +265,7 @@ class RAGTool(BaseTool): return "\n\n".join(outputs) # ---------- 添加文档 ---------- - def _handle_add(self, agents: List[str], content: Optional[str], source: Optional[str]) -> str: + def _handle_add(self, content: Optional[str], source: Optional[str]) -> str: if not content and not source: return "必须提供 content 或 source。" if content and source: @@ -292,12 +276,9 @@ class RAGTool(BaseTool): timestamp = datetime.now().isoformat() if content: - # 直接使用 content 分块 texts = split_into_chunks(content) chunks_meta = [(t, "直接输入") for t in texts] - elif source: - # 解析文件 if not os.path.isfile(source): return f"文件不存在: {source}" chunks_meta = [(c["content"], c["source"]) for c in parse_file_to_chunks(source)] @@ -309,64 +290,57 @@ class RAGTool(BaseTool): doc_texts = [cm[0] for cm in chunks_meta] _, doc_embs = embed_texts(queries=["dummy"], documents=doc_texts) - # 3. 为每个 agent 写入 - total_added = 0 - for agent in agents: - coll = get_collection(agent) - ids = [] - embeddings = [] - metadatas = [] - documents = [] + # 3. 写入全局 Collection + coll = get_global_collection() + ids = [] + embeddings = [] + metadatas = [] + documents = [] - for idx, (text, src_desc) in enumerate(chunks_meta): - # 使用 hash 作为唯一 ID - uid = hashlib.md5(f"{agent}_{timestamp}_{text[:80]}".encode()).hexdigest() - ids.append(uid) - embeddings.append(doc_embs[idx]) - metadatas.append({ - "source": src_desc, - "agent_name": agent, - "timestamp": timestamp, - }) - documents.append(text) + for idx, (text, src_desc) in enumerate(chunks_meta): + uid = hashlib.md5(f"{timestamp}_{text[:80]}".encode()).hexdigest() + ids.append(uid) + embeddings.append(doc_embs[idx]) + metadatas.append({ + "source": src_desc, + "timestamp": timestamp, + }) + documents.append(text) - coll.add( - ids=ids, - embeddings=embeddings, - metadatas=metadatas, - documents=documents - ) - total_added += len(ids) + coll.add( + ids=ids, + embeddings=embeddings, + metadatas=metadatas, + documents=documents + ) - return f"成功添加 {total_added} 条知识块到 {', '.join(agents)}。" + return f"成功添加 {len(ids)} 条知识块到全局知识库。" # ---------- 清除 ---------- - def _handle_clear(self, agents: List[str], filter: Optional[Union[str, List[str]]]) -> str: + def _handle_clear(self, filter: Optional[Union[str, List[str]]]) -> str: if not filter: return "清除操作必须提供 filter 条件。" - total_deleted = 0 - for agent in agents: - coll = get_collection(agent) - # 根据 filter 构建 where，当前简化：假设 filter 是 source 的包含条件 - where = None - if isinstance(filter, str): - where = {"source": {"$contains": filter}} - elif isinstance(filter, list): - # 多个条件用 $or - where = {"$or": [{"source": {"$contains": f}} for f in filter]} + coll = get_global_collection() + where = None + if isinstance(filter, str): + where = {"source": {"$contains": filter}} + elif isinstance(filter, list) and len(filter) > 0: + # 简化为使用第一个条件 + where = {"source": {"$contains": filter[0]}} - try: - # 查询要删除的 IDs - res = coll.get(where=where) - if res["ids"]: - coll.delete(ids=res["ids"]) - total_deleted += len(res["ids"]) - except Exception as e: - return f"清除时出错 (agent={agent}): {e}" + if not where: + return "无法构建有效的过滤条件。" - return f"已从 {', '.join(agents)} 中删除 {total_deleted} 条记录。" + try: + res = coll.get(where=where) + if res["ids"]: + coll.delete(ids=res["ids"]) + return f"已从全局知识库中删除 {len(res['ids'])} 条记录。" + else: + return "未找到匹配的记录。" + except Exception as e: + return f"清除时出错: {e}" async def _arun(self, **kwargs): - # 直接调用同步版本 - return self._run(**kwargs) + return self._run(**kwargs) \ No newline at end of file