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