agent/auto_agent
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修改了文件系统架构

Browse Source
This commit is contained in:
minorQAQ committed 2026-05-19 16:45:42 +08:00
1 parent 4a4bd5c365
commit 95fefe96dd
2 files changed
+95 -121
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/agent/chatbot_agent.py
+2 -2
View File
@@ -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
src/agent/tools/tool_rag.py
+93 -119
View File
@@ -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)