记忆与检索 第八章 记忆与检索
框架现有的缺陷? 最显著的一点是:所有的历史上下文,都只在程序运行期间存在的(也就是内存里),程序运行完毕历史就没了,会话之间的上下文也没法共享。
另一个缺陷则是,我们太过于依赖 LLM 本身的能力了,我们必须承认不同供应商的 LLM 的知识库水平是不同的,比如 gemini 知识库很全,claude 在 coding 知识上很强。完全依靠 LLM 的知识库是不稳定的,我们需要构建自己的知识库。
所以我们的框架,当务之急是解决以上问题,我们最容易想到的是将历史上下文持久化,这样就可以解决上下文共享,也可以建立我们自己的知识库。
官方依然给出了完整的 rag 组件,使用 pip 安装:
1 2 3 pip install "hello-agents[all]==0.2.9" python -m spacy download zh_core_web_sm python -m spacy download en_core_web_sm
随后,需要三个组件:Qdrant、Neo4J、Embedding,这里我们使用 docker 进行本地部署前两个,Embedding 模型用云服务。
Qdrant Qdrant 是向量数据库,擅长存储实体之间的 语义相似度,主要用于 语义搜索 和 Rag增强。也就是说,可以将 Qdrant 作为文本知识库。
在 docker 部署:
1 2 3 4 5 6 7 8 9 10 docker pull qdrant/qdrant mkdir /mydata/qrantdocker run -d \ --name qdrant-helloagent \ -p 6333:6333 \ -p 6334:6334 \ -v /mydata/qdrant \ --restart always \ qdrant/qdrant
数据放在哪个文件夹,自定义。
创建完毕后可以尝试进入 http://localhost:6333/dashboard 控制台看看是否成功。此时,服务内一个 Collection 都没有,实际上也不需要我们手动创建。不过这里给出创建一个官方参数 Collection 的方法:
1 2 3 4 5 6 7 8 PUT collections/HelloAgent { "vectors" : { "size" : 384 , "distance" : "Cosine" , "timeout" : 30 } }
控制台会返回
1 2 3 4 5 { "result" : true , "status" : "ok" , "time" : 0.139126398 }
最后,在 .env 中配置:
1 2 3 4 5 6 7 # Qdrant QDRANT_URL = http://localhost:6333 # Qdrant集合配置 QDRANT_COLLECTION = HelloAgentRag QDRANT_VECTOR_SIZE=384 QDRANT_DISTANCE=cosine QDRANT_TIMEOUT=30
docker 部署 qdrant,一般情况下不需要 key。
Neo4j 相比起 Qdrant 专门用来存储相似程度, Neo4j 的重心则放在了如何存储 知识之间的确切关系 ,与 Qdrant 知识库一同为 LLM 匹配更多有关知识的上下文。在 HelloAgent 中,被用于存储 语义记忆。
docker 部署:
1 2 3 4 5 6 7 8 9 docker pull neo4j mkdir /mydata/neo4jdocker run -it -d -p 7474:7474 -p 7687:7687 \ -v /mydata/neo4j \ -e NEO4J_AUTH=neo4j/password \ --name neo4j \ neo4j
控制台的端口是 7474,服务本身的端口是 7687。username 和 password 可以随意配置,这里设置成 neo4j & password。
完成后,进入 http://localhost:7474/browser/ 可以查看是否完成。
随后,配置 .env:
1 2 3 4 5 6 7 8 9 10 # 或使用本地Neo4j (需要Docker) NEO4J_URI= bolt://localhost:7687 NEO4J_USERNAME=neo4j NEO4J_PASSWORD=password # Neo4j连接配置 NEO4J_DATABASE=neo4j NEO4J_MAX_CONNECTION_LIFETIME=3600 NEO4J_MAX_CONNECTION_POOL_SIZE=50 NEO4J_CONNECTION_TIMEOUT=60
注意 uri 不要填 http,要填 bolt
SQLite SQLite 是一种十分常见的轻量级数据库,官方将其作为文档数据库使用,用于存储有着格式化数据的存储与查询。
SQLite 不需要我们额外安装,python 中包含 sqlite3 这个库可以直接用,对本地的 .db 文件进行读写操作。
在 HelloAgent 中,被用于存储完整的记忆,包括 content 在内。
Embedding Embedding 是专门的模型,将文本拆分成 token 再转化为向量。
因为嵌入模型在本地部署压力太大了,官方推荐使用云端 api ,比如阿里云百炼,我这里用了 text-embedding-v4 模型。在 .env 中配置:
1 2 3 4 5 6 7 8 # ========================== # 嵌入(Embedding)配置示例 - 可从阿里云控制台获取:https://dashscope.aliyun.com/ # ========================== # - 若为空,dashscope 默认 text-embedding-v3;local 默认 sentence-transformers/all-MiniLM-L6-v2 EMBED_MODEL_TYPE=dashscope EMBED_MODEL_NAME=text-embedding-v4 EMBED_API_KEY= 自己填 EMBED_BASE_URL=https://ws-91vjyamuomyrjpla.cn-beijing.maas.aliyuncs.com/compatible-mode/v1
测试程序 我们使用官方的测试程序,但是 SimpleAgent、LLM、ToolRegistry 可以使用我们自己实现的版本:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 from agent.simpleAgent import SimpleAgentfrom core.llm import LLMfrom tool.toolRegistry import ToolRegistryfrom hello_agents.tools import MemoryTool, RAGToolllm = LLM() agent = SimpleAgent( name="智能助手" , llm=llm, system_prompt="你是一个有记忆和知识检索能力的AI助手" ) tool_registry = ToolRegistry() memory_tool = MemoryTool(user_id="user123" ) tool_registry.register_tool(memory_tool) rag_tool = RAGTool(knowledge_base_path="./knowledge_base" ) tool_registry.register_tool(rag_tool) agent.tool_registry = tool_registry response = agent.run("你好!请记住我叫张三,我是一名Python开发者" ) print (response)
结果:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 ✅ 智能助手 初始化完成,工具调用: 禁用 [OK] SQLite 数据库表和索引创建完成 [OK] SQLite 文档存储初始化完成: ./memory_data\memory.db INFO:hello_agents.memory.storage.qdrant_store:✅ 成功连接到Qdrant服务: http://localhost:6333 INFO:hello_agents.memory.storage.qdrant_store:✅ 使用现有Qdrant集合: helloAgent INFO:hello_agents.memory.types.semantic:✅ 嵌入模型就绪,维度: 1024 INFO:hello_agents.memory.types.semantic:✅ Qdrant向量数据库初始化完成 INFO:hello_agents.memory.storage.neo4j_store:✅ 成功连接到Neo4j服务: bolt://localhost:7687 INFO:hello_agents.memory.storage.neo4j_store:✅ Neo4j索引创建完成 INFO:hello_agents.memory.types.semantic:✅ Neo4j图数据库初始化完成 INFO:hello_agents.memory.types.semantic:🏥 数据库健康状态: Qdrant=✅, Neo4j=✅ INFO:hello_agents.memory.types.semantic:✅ 加载中文spaCy模型: zh_core_web_sm INFO:hello_agents.memory.types.semantic:✅ 加载英文spaCy模型: en_core_web_sm INFO:hello_agents.memory.types.semantic:🎯 主要使用中文spaCy模型 INFO:hello_agents.memory.types.semantic:📚 可用语言模型: 中文, 英文 INFO:hello_agents.memory.types.semantic:增强语义记忆初始化完成(使用Qdrant+Neo4j专业数据库) INFO:hello_agents.memory.manager:MemoryManager初始化完成,启用记忆类型: ['working', 'episodic', 'semantic'] ✅ 工具 'memory' 已注册。 INFO:hello_agents.memory.storage.qdrant_store:✅ 成功连接到Qdrant服务: http://localhost:6333 INFO:hello_agents.memory.storage.qdrant_store:✅ 创建Qdrant集合: rag_knowledge_base ✅ RAG工具初始化成功: namespace=default, collection=rag_knowledge_base ✅ 工具 'rag' 已注册。 🤖 智能助手 正在处理: 你好!请记住我叫张三,我是一名Python开发者 🧠 正在调用 deepseek-v4-pro 模型... ✅ 大语言模型响应成功: ✅ 智能助手 响应完成 你好,张三!我已经记住了,你是一名 Python 开发者。有什么需要帮忙的吗?
我们现在不清楚到底干了什么,但是可以验证所有服务部署成功了。
memory 我们创建一个 memory 包,专用于处理记忆。这个包与 core、agent、tool 平行,其结构为:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 HelloAgent/ │ ├── memory/ # 记忆 │ ├── memoryItem.py # 记忆实体类 │ ├── memoryConfig.py # 记忆统一配置中心 │ ├── memoryManager.py # 记忆管理器 │ ├── memory.py # Memory 基类 │ └── types # 四种不同类型的记忆 │ ├── workingMemory.py # 工作记忆(TTL管理,纯内存) │ ├── episodicMemory.py # 情景记忆(事件序列,SQLite+Qdrant) │ └── semanticMemory.py # 语义记忆(知识图谱,Qdrant+Neo4j) │ ├── tool/ │ └── builtin # 四种不同类型的记忆 │ └── memoryTool.py # 对外暴露的记忆存取工具接口
这一部分中,我们的核心目标就是构建三种不同的记忆,并完成对不同类型的记忆的统一管理。
官方给出了测试程序,用来测试 MemoryTool:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 from agent.simpleAgent import SimpleAgentfrom core.llm import LLMfrom tool.toolRegistry import ToolRegistryfrom hello_agents.tools import MemoryToolllm = LLM() agent = SimpleAgent(name="记忆助手" , llm=llm) memory_tool = MemoryTool(user_id="user123" ) tool_registry = ToolRegistry() tool_registry.register_tool(memory_tool) agent.tool_registry = tool_registry print ("=== 添加多个记忆 ===" )result1 = memory_tool.run({"action" : "add" , "content" : "用户张三是一名Python开发者,专注于机器学习和数据分析" , "memory_type" : "semantic" , "importance" : 0.8 }) print (f"记忆1: {result1} " )result2 = memory_tool.run({"action" : "add" , "content" : "李四是前端工程师,擅长React和Vue.js开发" , "memory_type" : "semantic" , "importance" : 0.7 }) print (f"记忆2: {result2} " )result3 = memory_tool.run({"action" : "add" , "content" : "王五是产品经理,负责用户体验设计和需求分析" , "memory_type" : "semantic" , "importance" : 0.6 }) print (f"记忆3: {result3} " )print ("\n=== 搜索特定记忆 ===" )print ("🔍 搜索 '前端工程师':" )result = memory_tool.run({"action" : "search" , "query" : "前端工程师" , "limit" : 3 }) print (result)print ("\n=== 记忆摘要 ===" )result = memory_tool.run({"action" : "summary" }) print (result)
测试结果:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 python -m test.__testMemoryTool ✅ 记忆助手 初始化完成,工具调用: 禁用 [OK] SQLite 数据库表和索引创建完成 [OK] SQLite 文档存储初始化完成: ./memory_data\memory.db INFO:hello_agents.memory.storage.qdrant_store:✅ 成功连接到Qdrant服务: http://localhost:6333 INFO:hello_agents.memory.storage.qdrant_store:✅ 使用现有Qdrant集合: helloAgent INFO:hello_agents.memory.types.semantic:✅ 嵌入模型就绪,维度: 1024 INFO:hello_agents.memory.types.semantic:✅ Qdrant向量数据库初始化完成 INFO:hello_agents.memory.storage.neo4j_store:✅ 成功连接到Neo4j服务: bolt://localhost:7687 INFO:hello_agents.memory.storage.neo4j_store:✅ Neo4j索引创建完成 INFO:hello_agents.memory.types.semantic:✅ Neo4j图数据库初始化完成 INFO:hello_agents.memory.types.semantic:🏥 数据库健康状态: Qdrant=✅, Neo4j=✅ INFO:hello_agents.memory.types.semantic:✅ 加载中文spaCy模型: zh_core_web_sm INFO:hello_agents.memory.types.semantic:✅ 加载英文spaCy模型: en_core_web_sm INFO:hello_agents.memory.types.semantic:🎯 主要使用中文spaCy模型 INFO:hello_agents.memory.types.semantic:📚 可用语言模型: 中文, 英文 INFO:hello_agents.memory.types.semantic:增强语义记忆初始化完成(使用Qdrant+Neo4j专业数据库) INFO:hello_agents.memory.manager:MemoryManager初始化完成,启用记忆类型: ['working', 'episodic', 'semantic'] ✅ 工具 'memory' 已注册。 === 添加多个记忆 === INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] add_vectors start: n_vectors=1 n_meta=1 collection=helloAgent INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert begin: points=1 INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert done INFO:hello_agents.memory.storage.qdrant_store:✅ 成功添加 1 个向量到Qdrant INFO:hello_agents.memory.types.semantic:✅ 添加语义记忆: 1个实体, 0个关系 记忆1: ✅ 记忆已添加 (ID: be8b1a9f...) INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] add_vectors start: n_vectors=1 n_meta=1 collection=helloAgent INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert begin: points=1 INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert done INFO:hello_agents.memory.storage.qdrant_store:✅ 成功添加 1 个向量到Qdrant INFO:hello_agents.memory.types.semantic:✅ 添加语义记忆: 2个实体, 1个关系 记忆2: ✅ 记忆已添加 (ID: dd982d4f...) INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] add_vectors start: n_vectors=1 n_meta=1 collection=helloAgent INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert begin: points=1 INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert done INFO:hello_agents.memory.storage.qdrant_store:✅ 成功添加 1 个向量到Qdrant INFO:hello_agents.memory.types.semantic:✅ 添加语义记忆: 0个实体, 0个关系 记忆3: ✅ 记忆已添加 (ID: 89173a66...) === 搜索特定记忆 === 🔍 搜索 '前端工程师': INFO:hello_agents.memory.types.semantic:✅ 检索到 1 条相关记忆 🔍 找到 1 条相关记忆: 1. [语义记忆] 李四是前端工程师,擅长React和Vue.js开发 (重要性: 0.70) === 记忆摘要 === ERROR:hello_agents.memory.storage.qdrant_store:❌ 获取集合信息失败: 'CollectionInfo' object has no attribute 'vectors_count' INFO:hello_agents.memory.types.semantic:✅ 检索到 3 条相关记忆 📊 记忆系统摘要 总记忆数: 3 当前会话: session_20260723_132000 对话轮次: 0 📋 记忆类型分布: • 工作记忆: 0 条 (平均重要性: 0.00) • 情景记忆: 0 条 (平均重要性: 0.00) • 语义记忆: 3 条 (平均重要性: 0.70) ⭐ 重要记忆 (前3条): 1. 用户张三是一名Python开发者,专注于机器学习和数据分析 (重要性: 0.80) 2. 李四是前端工程师,擅长React和Vue.js开发 (重要性: 0.70) 3. 王五是产品经理,负责用户体验设计和需求分析 (重要性: 0.60)
记忆实体类 MemoryItem 首先,我们要定义一个严密的 Memory 最小单位,且要继承数据类 BaseModel :
1 2 3 4 5 6 7 8 9 10 11 12 13 14 from datetime import datetimefrom pydantic import BaseModelfrom typing import Any , Dict class MemoryItem (BaseModel ): """记忆项数据结构""" id : str content: str memory_type: str user_id: str timestamp: datetime importance: float = 0.5 metadata: Dict [str , Any ] = {} arbitrary_types_allowed: bool = True
id:全局唯一标识符,采用 uuid;
content:核心数据,是一串自然语言;
userId:标识该记忆的用户;
timestamp:时间戳;
importance:重要程度,后续用于计算是否遗忘;
metadata:附加数据,根据记忆类型的不同,会携带不同的 kv 字段;例如所有记忆都要自带 type,以表明自身的类型,例如 [type: “working”]。
arbitraryTypesAllowed:允许 datetime 和 Dict 等非标准 JSON 类型。
记忆配置类 MemoryConfig 在定义了最小单位的记忆数据类后,我们再对记忆存储的参数进行统一配置,包括存在哪、怎么存、存什么:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 from pydantic import BaseModelfrom typing import List class MemoryConfig (BaseModel ): """记忆系统配置""" storage_path: str = "./memory_data" max_capacity: int = 100 importance_threshold: float = 0.1 decay_factor: float = 0.95 working_memory_capacity: int = 10 working_memory_tokens: int = 2000 working_memory_ttl_minutes: int = 120 perceptual_memory_modalities: List [str ] = ["text" , "image" , "audio" ]
perceptual_memory_modalities 表示我们的记忆系统支持哪些模态:目前硬编码为文本、图像、音频。
记忆行为抽象类 Memory 我们定义了记忆数据类,和记忆统一配置,接下来我们实现一个抽象类 Memory,对所有三种记忆实现统一规范。
初始化 1 2 3 4 5 6 7 8 9 10 11 12 13 from abc import ABC, abstractmethodfrom typing import List , Dict , Any from .memoryConfg import MemoryConfig from .memoryItem import MemoryItemclass Memory (ABC ): """记忆基类 定义所有记忆类型的通用接口和行为 """ def __init__ (self, config: MemoryConfig, storage_backend=None ): self.config = config self.storage = storage_backend self.memory_type = self.__class__.__name__.lower().replace("memory" , "" )
接受一个 Config 来配置这段记忆的参数;
而 memory_type 表示该记忆的类型,这里没有给输入传参,而是通过继承 Memory 的子类的类名来推断,具体而言:
1 2 3 WorkingMemory → "working" EpisodicMemory → "episodic" SemanticMemory → "semantic"
这是为了方便未来的记忆类型拓展,当然,我承认这种方式很诡异,直接传入参数不就行了。
抽象方法 这里规定所有 Memory 必须实现以下方法:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 @abstractmethod def add (self, memory_item: MemoryItem ) -> str : """添加记忆项 Args: memory_item: 记忆项对象 Returns: 记忆ID """ pass @abstractmethod def retrieve (self, query: str , limit: int = 5 , **kwargs ) -> List [MemoryItem]: """检索相关记忆 Args: query: 查询内容 limit: 返回数量限制 **kwargs: 其他检索参数 Returns: 相关记忆列表 """ pass @abstractmethod def update (self, memory_id: str , content: str = None , importance: float = None , metadata: Dict [str , Any ] = None ) -> bool : """更新记忆 Args: memory_id: 记忆ID content: 新内容 importance: 新重要性 metadata: 新元数据 Returns: 是否更新成功 """ pass @abstractmethod def remove (self, memory_id: str ) -> bool : """删除记忆 Args: memory_id: 记忆ID Returns: 是否删除成功 """ pass @abstractmethod def has_memory (self, memory_id: str ) -> bool : """检查记忆是否存在 Args: memory_id: 记忆ID Returns: 是否存在 """ pass @abstractmethod def clear (self ): """清空所有记忆""" pass @abstractmethod def get_stats (self ) -> Dict [str , Any ]: """获取记忆统计信息 Returns: 统计信息字典 """ pass
基本方法就是增删改查。但值得注意的是,Memory 抽象类中并没有强制要求实现类实现 forget 方法,因为可能存在过于重要以至于必须永远记住的记忆存在。
同时,我们再明确一下这里面没有实现的设计:有关 MemoryItem 的存储。记忆存储的最小单元就是 MemoryItem,每个实现了 Memory 的实现类,内部都会维护一个用于存储 MemoryItem 的数据结构。
3+1 种记忆 有了 Memory 基类后,我们就可以通过实现,构建四种 Agent 工作必不可少的记忆。
工作记忆 WorkingMemory:主要用于原封不动地存储当前对话的上下文信息,生命周期较短,其容量也不需要太大,更重要的是不需要持久化,存在内存中就行。比方说 时间:21:04:39;用户:"使用 openAI 规范",Agent:"好的,我会严格遵守 openAI 规范。";当前会话轮数:12、时间:21:07:39;用户:"/grill-with-docs 我要实现一个 Memory 基类,帮我拟定计划。",Agent:"好的,我先调用。";当前会话轮数:13 等,都适合存入工作记忆。
情景记忆 EpisodicMemory:适合存用户与 Agent 在特定场合下的重要 交互,需要长期存储以支持 /resume 或者 --continue,例如 用户在21:04:39告诉我需要使用 openAI 规范等。相当于抛弃不重要的,保留重要的,然后进行上下文压缩,这就是情景记忆。
语义记忆 SemanticMemory:用来存储必须遵守 的抽象规范,这些记忆一般是由 LLM 自行总结并判断存入的。比如 规范:openAI 或者 用户更加偏好于......,当然这些也需要长期存储。
有了以上三种记忆后,我们的 Agent 的记忆功能就基本成型了:用户使用中,工作记忆保留完整上下文,尽可能保证当前会话信息不失误;同时,随时将重要信息压缩本地化,实现长期记忆和过滤无关信息;最后将必须遵守的抽象规范单独存放,在读取时尽可能保证完全遵守。
此外,官方也提出了一个感知记忆 PerceptualMemory,主要用于处理图片、视频、音频等多模态内容,我们这里就不做实现了,有上面三种记忆支持就可以初步实现用户与 Agent 的完整交互了。
工作记忆 WorkingMemory 前面提到,工作记忆无需存放进 Qdrant 或者 Neo4j,而是存在内存中,也就是说在会话期间直接以一个 WorkingMemory 实例存在就足以满足了。有关 Rag 的功能还没有实现,其他两种记忆依赖于 Rag,所以我们先尝试实现无需 Rag 的 WorkingMemory,存放到 memory/types。
记忆储存机制 继承 Memory,输入 MemoryConfig 作为配置中心。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 class WorkingMemory (Memory ): def __init__ (self, config: MemoryConfig, storage_backend=None ): super ().__init__(config, storage_backend) self.max_capacity = self.config.working_memory_capacity self.max_tokens = self.config.working_memory_tokens self.max_age_minutes = getattr (self.config, 'working_memory_ttl_minutes' , 120 ) self.current_tokens = 0 self.session_start = datetime.now() self.memories: List [MemoryItem] = []
可以看到,我们采用的是 List 对 MemoryItem 进行存储。
不过在官方提供的 hello-agents 包内,官方使用 List 和 heapq 来存储 MemoryItem,官方的意图很明显:分别用作存储和查询,但问题是 heapq 除了在添加记忆和淘汰记忆时,会对 heapq 进行入堆和重建,在关键的查询上,反而没有使用堆,所以我们可以直接看成不存在,只使用 List 其实更便于理解,在后续的代码构成中,我们也会直接忽略 heapq。
新增 添加记忆的方法如下:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 def add (self, memory_item: MemoryItem ) -> str : """添加工作记忆""" self._expire_old_memories() priority = self._calculate_priority(memory_item) self.memories.append(memory_item) self.current_tokens += len (memory_item.content.split()) self._enforce_capacity_limits() return memory_item.id
在真正添加到 List 前,我们还有两件事要做:过期清理 _expire_old_memories 和计算优先级 _calculate_priority。
先来看 _expire_old_memories,大体基本上就是对 List 进行遍历,删除超时的数据。代码:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 def _expire_old_memories (self ): """按TTL清理过期记忆,并同步更新堆与token计数""" if not self.memories: return cutoff_time = datetime.now() - timedelta(minutes=self.max_age_minutes) kept: List [MemoryItem] = [] removed_token_sum = 0 for m in self.memories: if m.timestamp >= cutoff_time: kept.append(m) else : removed_token_sum += len (m.content.split()) if len (kept) == len (self.memories): return self.memories = kept self.current_tokens = max (0 , self.current_tokens - removed_token_sum)
看 _calculate_priority 函数:
1 2 3 4 5 6 7 8 9 10 def _calculate_priority (self, memory: MemoryItem ) -> float : """计算记忆优先级""" priority = memory.importance time_decay = self._calculate_time_decay(memory.timestamp) priority *= time_decay return priority
优先级等于重要程度乘时间衰减 _calculate_time_decay,记忆存在时间越长,衰减越大,这意味着越晚加入 List 的 MemoryItem,优先级也越小,越容易被清除:
1 2 3 4 5 6 7 8 def _calculate_time_decay (self, timestamp: datetime ) -> float : """计算时间衰减因子""" time_diff = datetime.now() - timestamp hours_passed = time_diff.total_seconds() / 3600 decay_factor = self.config.decay_factor ** (hours_passed / 6 ) return max (0.1 , decay_factor)
而在添加到 List 后,还要干两件事:更新token计数和检查记忆容量 _enforce_capacity_limits,这个函数会一直清除优先级最低的 MemoryItem,直到在容量与 token 数量上均不超过最大限制。
1 2 3 4 5 6 7 8 9 def _enforce_capacity_limits (self ): """强制执行容量限制""" while len (self.memories) > self.max_capacity: self._remove_lowest_priority_memory() while self.current_tokens > self.max_tokens: self._remove_lowest_priority_memory()
检索 检索是最复杂的功能,主要分为四个步骤:1.过滤用户 id;2.TF-IDF 向量计算;3.关键词检索;4.最终合成。
TF-IDF 算法就不做解释了,因为我也不懂。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 def retrieve (self, query: str , limit: int = 5 , user_id: str = None , **kwargs ) -> List [MemoryItem]: """检索工作记忆 - 混合语义向量检索和关键词匹配""" self._expire_old_memories() if not self.memories: return [] active_memories = self.memories filtered_memories = active_memories if user_id: filtered_memories = [m for m in active_memories if m.user_id == user_id] if not filtered_memories: return [] vector_scores = {} try : from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity import numpy as np documents = [query] + [m.content for m in filtered_memories] vectorizer = TfidfVectorizer(stop_words=None , lowercase=True ) tfidf_matrix = vectorizer.fit_transform(documents) query_vector = tfidf_matrix[0 :1 ] doc_vectors = tfidf_matrix[1 :] similarities = cosine_similarity(query_vector, doc_vectors).flatten() for i, memory in enumerate (filtered_memories): vector_scores[memory.id ] = similarities[i] except Exception as e: vector_scores = {} query_lower = query.lower() scored_memories = [] for memory in filtered_memories: content_lower = memory.content.lower() vector_score = vector_scores.get(memory.id , 0.0 ) keyword_score = 0.0 if query_lower in content_lower: keyword_score = len (query_lower) / len (content_lower) else : query_words = set (query_lower.split()) content_words = set (content_lower.split()) intersection = query_words.intersection(content_words) if intersection: keyword_score = len (intersection) / len (query_words.union(content_words)) * 0.8 if vector_score > 0 : base_relevance = vector_score * 0.7 + keyword_score * 0.3 else : base_relevance = keyword_score time_decay = self._calculate_time_decay(memory.timestamp) base_relevance *= time_decay importance_weight = 0.8 + (memory.importance * 0.4 ) final_score = base_relevance * importance_weight if final_score > 0 : scored_memories.append((final_score, memory)) scored_memories.sort(key=lambda x: x[0 ], reverse=True ) return [memory for _, memory in scored_memories[:limit]]
遗忘 Memory 抽象类中并没有要求实现 forget 函数,这是出于某些记忆不需要遗忘机制而做出的决定。但对于工作记忆来说,遗忘是必须的,所以我们有必要对其进行实现:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 def forget (self, strategy: str = "importance_based" , threshold: float = 0.1 , max_age_days: int = 1 ) -> int : """工作记忆遗忘机制""" forgotten_count = 0 current_time = datetime.now() to_remove = [] cutoff_ttl = current_time - timedelta(minutes=self.max_age_minutes) for memory in self.memories: if memory.timestamp < cutoff_ttl: to_remove.append(memory.id ) if strategy == "importance_based" : for memory in self.memories: if memory.importance < threshold: to_remove.append(memory.id ) elif strategy == "time_based" : cutoff_time = current_time - timedelta(hours=max_age_days * 24 ) for memory in self.memories: if memory.timestamp < cutoff_time: to_remove.append(memory.id ) elif strategy == "capacity_based" : if len (self.memories) > self.max_capacity: sorted_memories = sorted ( self.memories, key=lambda m: self._calculate_priority(m) ) excess_count = len (self.memories) - self.max_capacity for memory in sorted_memories[:excess_count]: to_remove.append(memory.id ) for memory_id in to_remove: if self.remove(memory_id): forgotten_count += 1 return forgotten_count
三种遗忘策略:importance_based, time_based, capacity_based。
MemoryManager 我们已经定义了一个具体的 Memory 实现类,那这个记忆是如何生效的?先让我们回想一下,官方给出的包演示是怎么做的:
1 2 3 4 5 6 7 8 9 10 11 12 13 tool_registry = ToolRegistry() memory_tool = MemoryTool(user_id="user123" ) tool_registry.register_tool(memory_tool) rag_tool = RAGTool(knowledge_base_path="./knowledge_base" ) tool_registry.register_tool(rag_tool) agent.tool_registry = tool_registry
Memory 被作为工具被 ToolRegistry 统一注册发现,然后被 LLM 主动调用。
不过有个问题是,我们现在只实现了底层的 Memory 类及其实现类,如果我们要让工具意图类调用底层记忆,我们最好再加一层统一管理类,这就是 MemoryManager 类的由来。官方对 MemoryTool 和 MemoryManager 的定位是:
这种分层设计体现了软件工程中的关注点分离原则,MemoryTool专注于用户接口和参数处理,而MemoryManager则负责核心的记忆管理逻辑。
初始化 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 def __init__ ( self, config: Optional [MemoryConfig] = None , user_id: str = "default_user" , enable_working: bool = True , enable_episodic: bool = False , enable_semantic: bool = False ): self.config = config or MemoryConfig() self.user_id = user_id self.memory_types = {} if enable_working: self.memory_types['working' ] = WorkingMemory(self.config) logger.info(f"MemoryManager初始化完成,启用记忆类型: {list (self.memory_types.keys())} " )
除了 config,其内部还有一个 kv 对成员 memory_types,用来存四种记忆的对象,对其进行管理。我们先熟悉一下其结构。
memory_types 在大多数情况下可以看作一个有着三个元素的 kv map,而每个 v 内部则有着一个 List 用来存 MemoryItem,整个 memory_types 如下:
1 2 3 4 5 6 7 memory_types: "working" -> WorkingMemory() -> List<MemoryItem> -> [memory1, memory2 ......] "episodic" -> EpisodicMemory() -> List<MemoryItem> -> [memory1, memory2 ......] "semantic" -> SemanticMemory() -> List<MemoryItem> -> [memory1, memory2 ......]
记忆分类 对于 MemoryTool 来说,它并不需要对记忆的类型进行区分,我们规定让 LLM 在 metadata 中存记忆的类型,便于 MemoryManager 直接存取:
1 2 3 4 5 6 7 8 9 10 11 12 def _classify_memory_type (self, content: str , metadata: Optional [Dict [str , Any ]] ) -> str : """自动分类记忆类型""" if metadata and metadata.get("type" ): return metadata["type" ] if self._is_episodic_content(content): return "episodic" elif self._is_semantic_content(content): return "semantic" else : return "working"
同时,官方又给出了另外两个函数,根据 content 关键词来判断是否为情景记忆和语义记忆:
1 2 3 4 5 6 7 8 9 def _is_episodic_content (self, content: str ) -> bool : """判断是否为情景记忆内容""" episodic_keywords = ["昨天" , "今天" , "明天" , "上次" , "记得" , "发生" , "经历" ] return any (keyword in content for keyword in episodic_keywords) def _is_semantic_content (self, content: str ) -> bool : """判断是否为语义记忆内容""" semantic_keywords = ["定义" , "概念" , "规则" , "知识" , "原理" , "方法" ] return any (keyword in content for keyword in semantic_keywords)
如果 content 中包含任意关键词,则可以直接判定记忆类型,主要是在 LLM 没有响应记忆类型的时候兜底,平时几乎不参与。
CRUD 增:新增一个记忆,就是将其加入到成员 memory_types 对应 type 的 Memory 类对象,这里调用了 Memory 类对象必须实现的 add 函数。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 def add_memory ( self, content: str , memory_type: str = "working" , importance: Optional [float ] = None , metadata: Optional [Dict [str , Any ]] = None , auto_classify: bool = True ) -> str : """添加记忆 Args: content: 记忆内容 memory_type: 记忆类型 importance: 重要性分数 (0-1) metadata: 元数据 auto_classify: 是否自动分类到合适的记忆类型 Returns: 记忆ID """ if auto_classify: memory_type = self._classify_memory_type(content, metadata) if importance is None : importance = self._calculate_importance(content, metadata) memory_item = MemoryItem( id =str (uuid.uuid4()), content=content, memory_type=memory_type, user_id=self.user_id, timestamp=datetime.now(), importance=importance, metadata=metadata or {} ) if memory_type in self.memory_types: memory_id = self.memory_types[memory_type].add(memory_item) logger.debug(f"添加记忆到 {memory_type} : {memory_id} " ) return memory_id else : raise ValueError(f"不支持的记忆类型: {memory_type} " )
查:因为我们是按照关键词查找内容的,事先不知道预查找记忆的类型,所以这里需要遍历 memory_types 中每个 key-value 中的 value,调用其内部实现的 retrieve 函数,加入到最终返回。查找完后,根据重要度排序。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 def retrieve_memories ( self, query: str , memory_types: Optional [List [str ]] = None , limit: int = 10 , min_importance: float = 0.0 , time_range: Optional [tuple ] = None ) -> List [MemoryItem]: """检索记忆 Args: query: 查询内容 memory_types: 要检索的记忆类型列表 limit: 返回数量限制 min_importance: 最小重要性阈值 time_range: 时间范围 (start_time, end_time) Returns: 检索到的记忆列表 """ if memory_types is None : memory_types = list (self.memory_types.keys()) all_results = [] per_type_limit = max (1 , limit // len (memory_types)) for memory_type in memory_types: if memory_type in self.memory_types: memory_instance = self.memory_types[memory_type] try : type_results = memory_instance.retrieve( query=query, limit=per_type_limit, min_importance=min_importance, user_id=self.user_id ) all_results.extend(type_results) except Exception as e: logger.warning(f"检索 {memory_type} 记忆时出错: {e} " ) continue all_results.sort(key=lambda x: x.importance, reverse=True ) return all_results[:limit]
改:指定了 memory_id,这个和 memoryItem 中的成员变量 id 一致。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 def update_memory ( self, memory_id: str , content: Optional [str ] = None , importance: Optional [float ] = None , metadata: Optional [Dict [str , Any ]] = None ) -> bool : """更新记忆 Args: memory_id: 记忆ID content: 新内容 importance: 新重要性 metadata: 新元数据 Returns: 是否更新成功 """ for memory_type, memory_instance in self.memory_types.items(): if memory_instance.has_memory(memory_id): return memory_instance.update(memory_id, content, importance, metadata) logger.warning(f"未找到记忆: {memory_id} " ) return False
删:遍历 memory_types 里的 key-value 中的 value,将其 List 中对应 id 的 memoryItem 移除就行。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 def remove_memory (self, memory_id: str ) -> bool : """删除记忆 Args: memory_id: 记忆ID Returns: 是否删除成功 """ for memory_type, memory_instance in self.memory_types.items(): if memory_instance.has_memory(memory_id): return memory_instance.remove(memory_id) logger.warning(f"未找到记忆: {memory_id} " ) return False
遗忘机制 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 def forget_memories ( self, strategy: str = "importance_based" , threshold: float = 0.1 , max_age_days: int = 30 ) -> int : """记忆遗忘机制 Args: strategy: 遗忘策略 ("importance_based", "time_based", "capacity_based") threshold: 遗忘阈值 max_age_days: 最大保存天数 Returns: 遗忘的记忆数量 """ total_forgotten = 0 for memory_type, memory_instance in self.memory_types.items(): if hasattr (memory_instance, 'forget' ): forgotten = memory_instance.forget(strategy, threshold, max_age_days) total_forgotten += forgotten logger.info(f"记忆遗忘完成: {total_forgotten} 条记忆" ) return total_forgotten
官方包内也给出了一些其他的功能,遗忘功能是其中之一,就是简单地调用 memory_types 中的 Memory 实现类对象的 froget 函数就行。
记忆整合 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 def consolidate_memories ( self, from_type: str = "working" , to_type: str = "episodic" , importance_threshold: float = 0.7 ) -> int : """记忆整合 - 将重要的短期记忆转换为长期记忆 Args: from_type: 源记忆类型 to_type: 目标记忆类型 importance_threshold: 重要性阈值 Returns: 整合的记忆数量 """ if from_type not in self.memory_types or to_type not in self.memory_types: logger.warning(f"记忆类型不存在: {from_type} -> {to_type} " ) return 0 source_memory = self.memory_types[from_type] target_memory = self.memory_types[to_type] all_memories = source_memory.get_all() candidates = [ m for m in all_memories if m.importance >= importance_threshold ] consolidated_count = 0 for memory in candidates: if source_memory.remove(memory.id ): memory.memory_type = to_type memory.importance *= 1.1 target_memory.add(memory) consolidated_count += 1 logger.info(f"记忆整合完成: {consolidated_count} 条记忆从 {from_type} 转移到 {to_type} " ) return consolidated_count
记忆整合,实际上就是记忆类型转换,默认是将工作记忆转化为情景记忆。
我们已经构建了一个基本能用的 MemoryManager,接下来就是实现 MemoryTool。官方的设计理念是:
MemoryTool作为记忆系统的统一接口,其设计遵循了”统一入口,分发处理”的架构模式。
我们在 /tool/builtin 下面构建 memoryTool.py
初始化 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 class MemoryTool (Tool ): """记忆工具 —— 让 LLM 能够存取和管理记忆 作为记忆系统的统一接口,遵循"统一入口,分发处理"的架构模式。 MemoryTool 专注于参数解析和结果格式化,核心逻辑委托给 MemoryManager。 """ def __init__ ( self, user_id: str = "default_user" , config: Optional [MemoryConfig] = None ): super ().__init__( name="memory" , description=( "记忆管理工具,用于存储、检索、更新和删除记忆。" "支持以下操作:\n" "- add: 添加一条新记忆。需要 content(记忆内容),可选 memory_type " "(working/episodic/semantic,默认自动推断)、importance(0-1,默认自动计算)\n" "- search: 按关键词检索记忆。需要 query(查询内容),可选 limit(返回数量,默认5)\n" "- update: 更新已有记忆。需要 memory_id,可选 content/importance\n" "- remove: 删除指定记忆。需要 memory_id\n" "- summary: 查看记忆系统的统计摘要(总数、各类型容量等)\n" "- forget: 触发记忆遗忘机制。可选 strategy(importance_based/time_based/capacity_based)、" "threshold(遗忘阈值,默认0.1)\n" "- consolidate: 将重要的短期记忆整合为长期记忆。" "可选 from_type(默认working)、to_type(默认episodic)、importance_threshold(默认0.7)\n\n" "记忆类型说明:\n" "- working: 工作记忆(当前会话上下文,容量有限,过期自动清理)\n" "- episodic: 情景记忆(重要交互记录,长期保留)\n" "- semantic: 语义记忆(抽象规则和用户偏好,永久保留)\n\n" "使用建议:\n" "- 用户说'记住XXX'时,根据内容类型选择合适的 memory_type\n" "- 用户偏好/规范类信息 → semantic,重要性 0.8+\n" "- 临时上下文/当前话题 → working,重要性 0.5\n" "- 重要事件/决定 → episodic,重要性 0.7+\n" "- 在回答用户问题前,先用 search 检索相关记忆" ) ) self.user_id = user_id self._manager = MemoryManager( config=config, user_id=user_id, enable_working=True , enable_episodic=False , enable_semantic=False )
显然,MemoryTool 需要继承 Tool,内部也定义了 description 以让 LLM 调用,且里面提供了详细的 action 字段的作用。
此外,其包含了一个 MemoryManager 成员,用来充当调用 Memory 实现类的中间层。
统一入口 run 函数 我们希望外部只需要通过一个函数就可以调用所有记忆相关的功能,如同官方做法一样,我们定义一个网关函数 run,通过输入 Dict 字典的参数来路由到不同的函数。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 def run (self, parameters: Dict [str , Any ] ) -> str : """统一入口:根据 action 分发给具体处理方法 Args: parameters: 包含 action 及其他所需参数的字典 Returns: 格式化后的结果字符串 """ action = parameters.get("action" , "search" ) dispatch = { "add" : self._handle_add, "search" : self._handle_search, "update" : self._handle_update, "remove" : self._handle_remove, "summary" : self._handle_summary, "forget" : self._handle_forget, "consolidate" : self._handle_consolidate, } handler = dispatch.get(action) if handler is None : return ( f"❌ 不支持的操作: '{action} '。" f"可用操作: {', ' .join(dispatch.keys())} " ) try : return handler(parameters) except ValueError as e: return f"❌ 参数错误: {e} " except Exception as e: return f"❌ 操作失败: {e} "
此外,run 函数返回字符串类型。
测试 我们只测试一下 Tool 的基本功能正不正常,所以我们把所有记忆类型设置为 working:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 from tool.builtin.memoryTool import MemoryToolmemory_tool = MemoryTool(user_id="user123" ) print ("=== 添加多个记忆 ===" )result1 = memory_tool.run({"action" : "add" , "content" : "用户张三是一名Python开发者,专注于机器学习和数据分析" , "memory_type" : "working" , "importance" : 0.8 }) print (f"记忆1: {result1} " )result2 = memory_tool.run({"action" : "add" , "content" : "李四是前端工程师,擅长React和Vue.js开发" , "memory_type" : "working" , "importance" : 0.7 }) print (f"记忆2: {result2} " )result3 = memory_tool.run({"action" : "add" , "content" : "王五是产品经理,负责用户体验设计和需求分析" , "memory_type" : "working" , "importance" : 0.6 }) print (f"记忆3: {result3} " )print ("\n=== 搜索特定记忆 ===" )print ("🔍 搜索 '前端工程师':" )result = memory_tool.run({"action" : "search" , "query" : "前端工程师" , "limit" : 3 }) print (result)print ("\n=== 记忆摘要 ===" )result = memory_tool.run({"action" : "summary" }) print (result)
结果:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 > python -m test.__testMemoryTool ✅ 记忆助手 初始化完成,工具调用: 禁用 INFO:memory.memoryManager:MemoryManager初始化完成,启用记忆类型: ['working' ] ✅ 工具 'memory' 已注册。 === 添加多个记忆 === 记忆1 : ✅ 记忆已添加 ID: 5c0ba6bf-bf42-47a9-83cc-32880e7297fc 类型: working 内容: 用户张三是一名Python开发者,专注于机器学习和数据分析 记忆2 : ✅ 记忆已添加 ID: e8c899f9-fa4c-4511 -a458-166f470658e3 类型: working 内容: 李四是前端工程师,擅长React和Vue.js开发 记忆3 : ✅ 记忆已添加 ID: 3ebdb60e-c9df-46f7-9b06-b50d80910080 类型: working 内容: 王五是产品经理,负责用户体验设计和需求分析 === 搜索特定记忆 === 🔍 搜索 '前端工程师' : 🔍 搜索 '前端工程师' 的结果(共 1 条): [1 ] [working] 李四是前端工程师,擅长React和Vue.js开发 (importance=0.70 , id =e8c899f9...) === 记忆摘要 === 📊 记忆系统摘要: 用户: user123 总记忆数: 3 启用的记忆类型: working 配置: 最大容量=100 , 重要性阈值=0.1 , 衰减因子=0.95 [working] 活跃=3 , 容量使用率=30.0 %
记忆持久化 数据库操作 我们直接大大方方使用官方包提供的四个包,新建一个 storage 包放到下面:
1 2 3 4 5 6 7 8 HelloAgent/ │ ├── storage/ # 记忆 │ ├── __init__.py # 包含 load_env() 函数 │ ├── document_store.py # 文档数据库 │ ├── embedding.py # 向量模型 │ ├── neo4j_store.py # neo4j │ └── qdrant_store.py # qdrant
不过 qdrant_store.py 里有点问题,版本不兼容,会出现错误 'CollectionInfo' object has no attribute 'vectors_count'。解决就是把里面的 get_collection_info 改成:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 """ 获取集合信息 Returns: Dict: 集合信息 """ try : collection_info = self.client.get_collection(self.collection_name) count=0 if hasattr (collection_info, "points_count" ): count=collection_info.points_count elif hasattr (collection_info, "vectors_count" ): count=collection_info.vectors_count info = { "name" : self.collection_name, "points_count" : collection_info.points_count, "segments_count" : collection_info.segments_count, "indexed_vectors_count" : collection_info.indexed_vectors_count, "config" : { "vector_size" : self.vector_size, "distance" : self.distance.value, } } return info except Exception as e: logger.error(f"❌ 获取集合信息失败: {e} " ) return {}
这样就不会报错了。
情景记忆 EpisodicMemory 官方对情景记忆的定义是:
情景记忆负责存储具体的事件和经历,它的设计重点在于保持事件的完整性和时间序列关系。
采用了两种数据库:
采用了SQLite+Qdrant的混合存储方案,SQLite负责结构化数据的存储和复杂查询,Qdrant负责高效的向量检索。
所以我们创建一个 EpisodeMemory。
情景实体类 Episodic 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 class Episode : """情景记忆中的单个情景""" def __init__ ( self, episode_id: str , user_id: str , session_id: str , timestamp: datetime, content: str , context: Dict [str , Any ], outcome: Optional [str ] = None , importance: float = 0.5 ): self.episode_id = episode_id self.user_id = user_id self.session_id = session_id self.timestamp = timestamp self.content = content self.context = context self.outcome = outcome self.importance = importance
初始化 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 class EpisodicMemory (Memory ): """情景记忆实现 特点: - 存储具体的交互事件 - 包含丰富的上下文信息 - 按时间序列组织 - 支持模式识别和回溯 """ def __init__ (self, config: MemoryConfig, storage_backend=None ): super ().__init__(config, storage_backend) self.episodes: List [Episode] = [] self.sessions: Dict [str , List [str ]] = {} self.patterns_cache = {} self.last_pattern_analysis = None db_dir = self.config.storage_path if hasattr (self.config, 'storage_path' ) else "./memory_data" os.makedirs(db_dir, exist_ok=True ) db_path = os.path.join(db_dir, "memory.db" ) self.doc_store = SQLiteDocumentStore(db_path=db_path) self.embedder = get_text_embedder() qdrant_url = os.getenv("QDRANT_URL" ) qdrant_api_key = os.getenv("QDRANT_API_KEY" ) self.vector_store = QdrantConnectionManager.get_instance( url=qdrant_url, api_key=qdrant_api_key, collection_name=os.getenv("QDRANT_COLLECTION" , "hello_agents_vectors" ), vector_size=get_dimension(getattr (self.embedder, 'dimension' , 384 )), distance=os.getenv("QDRANT_DISTANCE" , "cosine" ) )
重要的成员变量有以下:
episodes:存储当前 session 下的情景知识;
sessions:一个字典,将不同的 session 下的情景知识按照 session_id 分开来,便于用户 /resume;
doc_store:sqlite 连接器,可以通过这个对数据库进行读写。比如 EpisodicMemory 几乎所有的数据持久化函数都是通过外包给 doc_store 实现的;
vector_store:qdrant 连接器,可对 qdrant 进行读写。
embedder:嵌入模型,将文字转成向量,用于辅助 vector_store 的 qdrant 读写。
不过官方在教程文档中提到的 storage_backend 在官方库代码中已经没用了,我们默认采用 qdrant 而非自定义。
add 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 def add (self, memory_item: MemoryItem ) -> str : """添加情景记忆""" session_id = memory_item.metadata.get("session_id" , "default_session" ) context = memory_item.metadata.get("context" , {}) outcome = memory_item.metadata.get("outcome" ) participants = memory_item.metadata.get("participants" , []) tags = memory_item.metadata.get("tags" , []) episode = Episode( episode_id=memory_item.id , user_id=memory_item.user_id, session_id=session_id, timestamp=memory_item.timestamp, content=memory_item.content, context=context, outcome=outcome, importance=memory_item.importance ) self.episodes.append(episode) if session_id not in self.sessions: self.sessions[session_id] = [] self.sessions[session_id].append(episode.episode_id) ts_int = int (memory_item.timestamp.timestamp()) self.doc_store.add_memory( memory_id=memory_item.id , user_id=memory_item.user_id, content=memory_item.content, memory_type="episodic" , timestamp=ts_int, importance=memory_item.importance, properties={ "session_id" : session_id, "context" : context, "outcome" : outcome, "participants" : participants, "tags" : tags } ) try : embedding = self.embedder.encode(memory_item.content) if hasattr (embedding, "tolist" ): embedding = embedding.tolist() self.vector_store.add_vectors( vectors=[embedding], metadata=[{ "memory_id" : memory_item.id , "user_id" : memory_item.user_id, "memory_type" : "episodic" , "importance" : memory_item.importance, "session_id" : session_id, "content" : memory_item.content }], ids=[memory_item.id ] ) except Exception: pass return memory_item.id
分为两大步骤:1.sqlite 存储完整数据(官方称之为“权威存储”);2.qdrant 存储向量索引数据。
retrieve 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 def retrieve (self, query: str , limit: int = 5 , **kwargs ) -> List [MemoryItem]: """检索情景记忆(结构化过滤 + 语义向量检索)""" user_id = kwargs.get("user_id" ) session_id = kwargs.get("session_id" ) time_range: Optional [Tuple [datetime, datetime]] = kwargs.get("time_range" ) importance_threshold: Optional [float ] = kwargs.get("importance_threshold" ) candidate_ids: Optional [set ] = None if time_range is not None or importance_threshold is not None : start_ts = int (time_range[0 ].timestamp()) if time_range else None end_ts = int (time_range[1 ].timestamp()) if time_range else None docs = self.doc_store.search_memories( user_id=user_id, memory_type="episodic" , start_time=start_ts, end_time=end_ts, importance_threshold=importance_threshold, limit=1000 ) candidate_ids = {d["memory_id" ] for d in docs} try : query_vec = self.embedder.encode(query) if hasattr (query_vec, "tolist" ): query_vec = query_vec.tolist() where = {"memory_type" : "episodic" } if user_id: where["user_id" ] = user_id hits = self.vector_store.search_similar( query_vector=query_vec, limit=max (limit * 5 , 20 ), where=where ) except Exception: hits = [] now_ts = int (datetime.now().timestamp()) results: List [Tuple [float , MemoryItem]] = [] seen = set () for hit in hits: meta = hit.get("metadata" , {}) mem_id = meta.get("memory_id" ) if not mem_id or mem_id in seen: continue episode = next ((e for e in self.episodes if e.episode_id == mem_id), None ) if episode and episode.context.get("forgotten" , False ): continue if candidate_ids is not None and mem_id not in candidate_ids: continue if session_id and meta.get("session_id" ) != session_id: continue doc = self.doc_store.get_memory(mem_id) if not doc: continue vec_score = float (hit.get("score" , 0.0 )) age_days = max (0.0 , (now_ts - int (doc["timestamp" ])) / 86400.0 ) recency_score = 1.0 / (1.0 + age_days) imp = float (doc.get("importance" , 0.5 )) base_relevance = vec_score * 0.8 + recency_score * 0.2 importance_weight = 0.8 + (imp * 0.4 ) combined = base_relevance * importance_weight item = MemoryItem( id =doc["memory_id" ], content=doc["content" ], memory_type=doc["memory_type" ], user_id=doc["user_id" ], timestamp=datetime.fromtimestamp(doc["timestamp" ]), importance=doc.get("importance" , 0.5 ), metadata={ **doc.get("properties" , {}), "relevance_score" : combined, "vector_score" : vec_score, "recency_score" : recency_score } ) results.append((combined, item)) seen.add(mem_id) if not results: fallback = super ()._generate_id query_lower = query.lower() for ep in self._filter_episodes(user_id, session_id, time_range): if query_lower in ep.content.lower(): recency_score = 1.0 / (1.0 + max (0.0 , (now_ts - int (ep.timestamp.timestamp())) / 86400.0 )) keyword_score = 0.5 base_relevance = keyword_score * 0.8 + recency_score * 0.2 importance_weight = 0.8 + (ep.importance * 0.4 ) combined = base_relevance * importance_weight item = MemoryItem( id =ep.episode_id, content=ep.content, memory_type="episodic" , user_id=ep.user_id, timestamp=ep.timestamp, importance=ep.importance, metadata={ "session_id" : ep.session_id, "context" : ep.context, "outcome" : ep.outcome, "relevance_score" : combined } ) results.append((combined, item)) results.sort(key=lambda x: x[0 ], reverse=True ) return [it for _, it in results[:limit]]
查询依旧大工程这块。
我们从 SQLite 中提取出满足条件的初步候选,将其 id 存入一个 set,然后经过查询 qdrant 排除相关度低的候选,set 中剩余的 id 再经由 SQLite 查询完整的 MemoryItem,最排序输出。
不过,如果向量数据库的结果为空,可能是由于数据库连接失效,又或者真的相关度不高,那就退化为关键词匹配。
测试 我们把上面测试 MemoryTool 的程序的类型改成 semantic:
1 2 3 result2 = memory_tool.run({"action" : "add" , "content" : "李四是前端工程师,擅长React和Vue.js开发" , "memory_type" : "episodic" , "importance" : 0.7 }) print (f"记忆2: {result2} " )
结果:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 > python -m test.__testMemoryTool ✅ 记忆助手 初始化完成,工具调用: 禁用 [OK] SQLite 数据库表和索引创建完成 [OK] SQLite 文档存储初始化完成: ./memory_data\memory.db INFO:storage.qdrant_store:✅ 成功连接到Qdrant服务: http://localhost:6333 INFO:httpx:HTTP Request: GET http://localhost:6333/collections "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: GET http://localhost:6333/collections "HTTP/1.1 200 OK" INFO:storage.qdrant_store:✅ 使用现有Qdrant集合: helloAgent INFO:httpx:HTTP Request: PATCH http://localhost:6333/collections/helloAgent "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: GET http://localhost:6333 "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:memory.types.semanticMemory:✅ 嵌入模型就绪,维度: 1024 INFO:storage.qdrant_store:✅ 成功连接到本地Qdrant服务: localhost:6333 INFO:httpx:HTTP Request: GET http://localhost:6333/collections "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: GET http://localhost:6333/collections "HTTP/1.1 200 OK" INFO:storage.qdrant_store:✅ 使用现有Qdrant集合: helloAgent INFO:httpx:HTTP Request: PATCH http://localhost:6333/collections/helloAgent "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: GET http://localhost:6333 "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:storage.neo4j_store:✅ 成功连接到Neo4j服务: bolt://localhost:7687 INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX entity_id_index IF NOT EXISTS FOR (e:Entity) ON (e.id)' has no effect. The index or constraint specified by 'RANGE INDEX entity_id_index FOR (e:Entity) ON (e.id)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX entity_id_index IF NOT EXISTS FOR (e:Entity) ON (e.id)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX entity_name_index IF NOT EXISTS FOR (e:Entity) ON (e.name)' has no effect. The index or constraint specified by 'RANGE INDEX entity_name_index FOR (e:Entity) ON (e.name)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX entity_name_index IF NOT EXISTS FOR (e:Entity) ON (e.name)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX entity_type_index IF NOT EXISTS FOR (e:Entity) ON (e.type)' has no effect. The index or constraint specified by 'RANGE INDEX entity_type_index FOR (e:Entity) ON (e.type)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX entity_type_index IF NOT EXISTS FOR (e:Entity) ON (e.type)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX memory_id_index IF NOT EXISTS FOR (e:Memory) ON (e.id)' has no effect. The index or constraint specified by 'RANGE INDEX memory_id_index FOR (e:Memory) ON (e.id)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX memory_id_index IF NOT EXISTS FOR (m:Memory) ON (m.id)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX memory_type_index IF NOT EXISTS FOR (e:Memory) ON (e.memory_type)' has no effect. The index or constraint specified by 'RANGE INDEX memory_type_index FOR (e:Memory) ON (e.memory_type)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX memory_type_index IF NOT EXISTS FOR (m:Memory) ON (m.memory_type)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX memory_timestamp_index IF NOT EXISTS FOR (e:Memory) ON (e.timestamp)' has no effect. The index or constraint specified by 'RANGE INDEX memory_timestamp_index FOR (e:Memory) ON (e.timestamp)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX memory_timestamp_index IF NOT EXISTS FOR (m:Memory) ON (m.timestamp)' INFO:storage.neo4j_store:✅ Neo4j索引创建完成 INFO:memory.types.semanticMemory:✅ 加载中文spaCy模型: zh_core_web_sm INFO:memory.types.semanticMemory:✅ 加载英文spaCy模型: en_core_web_sm INFO:memory.types.semanticMemory:🎯 主要使用中文spaCy模型 INFO:memory.types.semanticMemory:📚 可用语言模型: 中文, 英文 INFO:memory.types.semanticMemory:增强语义记忆初始化完成(使用Qdrant+Neo4j专业数据库) INFO:memory.memoryManager:MemoryManager初始化完成,启用记忆类型: ['working', 'episodic', 'semantic'] ✅ 工具 'memory' 已注册。 === 添加多个记忆 === 记忆1: ✅ 记忆已添加 ID: 872d6d9f-7b37-4b0b-a409-5010349b01af 类型: working 内容: 用户张三是一名Python开发者,专注于机器学习和数据分析 INFO:storage.qdrant_store:[Qdrant] add_vectors start: n_vectors=1 n_meta=1 collection=helloAgent INFO:storage.qdrant_store:[Qdrant] upsert begin: points=1 INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/points?wait=true "HTTP/1.1 200 OK" INFO:storage.qdrant_store:[Qdrant] upsert done INFO:storage.qdrant_store:✅ 成功添加 1 个向量到Qdrant 记忆2: ✅ 记忆已添加 ID: 8bb07638-85b3-4355-a41f-4458dd31b5b1 类型: episodic 内容: 李四是前端工程师,擅长React和Vue.js开发 记忆3: ✅ 记忆已添加 ID: 9d3bd0f3-c214-4419-9831-53bffa07a651 类型: working 内容: 王五是产品经理,负责用户体验设计和需求分析 === 搜索特定记忆 === 🔍 搜索 '前端工程师': INFO:httpx:HTTP Request: POST http://localhost:6333/collections/helloAgent/points/query "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: POST http://localhost:6333/collections/helloAgent/points/query "HTTP/1.1 200 OK" INFO:memory.types.semanticMemory:✅ 检索到 1 条相关记忆 🔍 搜索 '前端工程师' 的结果(共 2 条): [1] [episodic] 李四是前端工程师,擅长React和Vue.js开发 (importance=0.70, id=8bb07638...) [2] [semantic] 李四是前端工程师,擅长React和Vue.js开发 (importance=0.70, id=5d3dbf35...) === 记忆摘要 === INFO:httpx:HTTP Request: GET http://localhost:6333/collections/helloAgent "HTTP/1.1 200 OK" 📊 记忆系统摘要: 用户: user123 总记忆数: 3 启用的记忆类型: working, episodic, semantic 配置: 最大容量=100, 重要性阈值=0.1, 衰减因子=0.95 [working] 活跃=2, 容量使用率=20.0% [episodic] 活跃=1, 容量使用率=0.0% [semantic] 活跃=0, 容量使用率=0.0%
语义记忆 SemanticMemory 官方定义;
语义记忆是记忆系统中最复杂的部分,它负责存储抽象的概念、规则和知识。
官方的设计思路是:
设计重点在于知识的结构化表示和智能推理能力。语义记忆采用了Neo4j图数据库和Qdrant向量数据库的混合架构,这种设计让系统既能进行快速的语义检索,又能利用知识图谱进行复杂的关系推理。
和情景记忆不同的是,语义记忆除了用到嵌入模型在 qdrant 中向量化存储,还要在 neo4j 中存储对象间关系等 绝对真理 。
这种设计让系统既能进行快速的语义检索,又能利用知识图谱进行复杂的关系推理。
实体与关系类 Entity & Relation neo4j 中需要存储 entity 之间的 relation,分别创建类:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 class Entity : """实体类""" def __init__ ( self, entity_id: str , name: str , entity_type: str = "MISC" , description: str = "" , properties: Dict [str , Any ] = None ): self.entity_id = entity_id self.name = name self.entity_type = entity_type self.description = description self.properties = properties or {} self.created_at = datetime.now() self.updated_at = datetime.now() self.frequency = 1 def to_dict (self ) -> Dict [str , Any ]: return { "entity_id" : self.entity_id, "name" : self.name, "entity_type" : self.entity_type, "description" : self.description, "properties" : self.properties, "frequency" : self.frequency }
实体类规定了一个实体的名称、类型、所含参数等,使用唯一标识符 entity_id 进行校准。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 class Relation : """关系类""" def __init__ ( self, from_entity: str , to_entity: str , relation_type: str , strength: float = 1.0 , evidence: str = "" , properties: Dict [str , Any ] = None ): self.from_entity = from_entity self.to_entity = to_entity self.relation_type = relation_type self.strength = strength self.evidence = evidence self.properties = properties or {} self.created_at = datetime.now() self.frequency = 1 def to_dict (self ) -> Dict [str , Any ]: return { "from_entity" : self.from_entity, "to_entity" : self.to_entity, "relation_type" : self.relation_type, "strength" : self.strength, "evidence" : self.evidence, "properties" : self.properties, "frequency" : self.frequency }
关系类与实体类不同,不包含一个 id,而是使用 from_entity 和 to_entity 进行唯一标识。同时,一个关系也展示了两个实体间的关系、关系强度、证据、参数等。
初始化 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 class SemanticMemory (Memory ): """增强语义记忆实现 特点: - 使用HuggingFace中文预训练模型进行文本嵌入 - 向量检索进行快速相似度匹配 - 知识图谱存储实体和关系 - 混合检索策略:向量+图+语义推理 """ def __init__ (self, config: MemoryConfig, storage_backend=None ): super ().__init__(config, storage_backend) self.embedder = get_text_embedder() self.entities: Dict [str , Entity] = {} self.relations: List [Relation] = [] self.nlp = None self._init_nlp() self.semantic_memories: List [MemoryItem] = [] self.memory_embeddings: Dict [str , np.ndarray] = {} logger.info("增强语义记忆初始化完成(使用Qdrant+Neo4j专业数据库)" ) def _init_embedding_model (self ): """初始化统一嵌入模型(由 embedding_provider 管理)。""" try : self.embedding_model = get_text_embedder() try : test_vec = self.embedding_model.encode("health_check" ) dim = getattr (self.embedding_model, "dimension" , len (test_vec)) logger.info(f"✅ 嵌入模型就绪,维度: {dim} " ) except Exception: logger.info("✅ 嵌入模型就绪" ) except Exception as e: logger.error(f"❌ 嵌入模型初始化失败: {e} " ) raise def _init_databases (self ): """初始化专业数据库存储""" try : qdrant_url = os.getenv("QDRANT_URI" ) qdrant_api_key = os.getenv("QDRANT_API_KEY" ) self.vector_store = QdrantConnectionManager.get_instance( url=qdrant_url, api_key=qdrant_api_key, collection_name=os.getenv("QDRANT_COLLECTION" , "hello_agents_vectors" ), vector_size=get_dimension(getattr (self.embedder, 'dimension' , 384 )), distance=os.getenv("QDRANT_DISTANCE" , "cosine" ) ) self.graph_store = Neo4jGraphStore( uri=os.getenv("NEO4J_URI" , "bolt://localhost:7687" ), username=os.getenv("NEO4J_USERNAME" , "neo4j" ), password=os.getenv("NEO4J_PASSWORD" , "password" ), database=os.getenv("NEO4J_DATABASE" , "neo4j" ) ) except Exception as e: logger.error(f"❌ 数据库初始化失败: {e} " ) logger.info("💡 请检查数据库配置和网络连接" ) logger.info("💡 参考 DATABASE_SETUP_GUIDE.md 进行配置" ) raise def _init_nlp (self ): """初始化NLP处理器 - 智能多语言支持""" try : self.nlp_models = {} models_to_try = [ ("zh_core_web_sm" , "中文" ), ("en_core_web_sm" , "英文" ) ] loaded_models = [] for model_name, lang_name in models_to_try: try : nlp = spacy.load(model_name) self.nlp_models[model_name] = nlp loaded_models.append(lang_name) logger.info(f"✅ 加载{lang_name} spaCy模型: {model_name} " ) except OSError: logger.warning(f"⚠️ {lang_name} spaCy模型不可用: {model_name} " ) if "zh_core_web_sm" in self.nlp_models: self.nlp = self.nlp_models["zh_core_web_sm" ] logger.info("🎯 主要使用中文spaCy模型" ) elif "en_core_web_sm" in self.nlp_models: self.nlp = self.nlp_models["en_core_web_sm" ] logger.info("🎯 主要使用英文spaCy模型" ) else : self.nlp = None logger.warning("⚠️ 无可用spaCy模型,实体提取将受限" ) if loaded_models: logger.info(f"📚 可用语言模型: {', ' .join(loaded_models)} " ) except ImportError: logger.warning("⚠️ spaCy不可用,实体提取将受限" ) self.nlp = None self.nlp_models = {}
add 自己看。
retrieve 自己看。
测试 我们把上面测试 MemoryTool 的程序的类型改成 semantic:
1 2 3 result2 = memory_tool.run({"action" : "add" , "content" : "李四是前端工程师,擅长React和Vue.js开发" , "memory_type" : "semantic" , "importance" : 0.7 }) print (f"记忆2: {result2} " )
再次跑:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 > python -m test.__testMemoryTool ✅ 记忆助手 初始化完成,工具调用: 禁用 [OK] SQLite 数据库表和索引创建完成 [OK] SQLite 文档存储初始化完成: ./memory_data\memory.db INFO:storage.qdrant_store:✅ 成功连接到Qdrant服务: http://localhost:6333 INFO:httpx:HTTP Request: GET http://localhost:6333/collections "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: GET http://localhost:6333/collections "HTTP/1.1 200 OK" INFO:storage.qdrant_store:✅ 使用现有Qdrant集合: helloAgent INFO:httpx:HTTP Request: PATCH http://localhost:6333/collections/helloAgent "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: GET http://localhost:6333 "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:memory.types.semanticMemory:✅ 嵌入模型就绪,维度: 1024 INFO:storage.qdrant_store:✅ 成功连接到本地Qdrant服务: localhost:6333 INFO:httpx:HTTP Request: GET http://localhost:6333/collections "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: GET http://localhost:6333/collections "HTTP/1.1 200 OK" INFO:storage.qdrant_store:✅ 使用现有Qdrant集合: helloAgent INFO:httpx:HTTP Request: PATCH http://localhost:6333/collections/helloAgent "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: GET http://localhost:6333 "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/index?wait=true "HTTP/1.1 200 OK" INFO:storage.neo4j_store:✅ 成功连接到Neo4j服务: bolt://localhost:7687 INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX entity_id_index IF NOT EXISTS FOR (e:Entity) ON (e.id)' has no effect. The index or constraint specified by 'RANGE INDEX entity_id_index FOR (e:Entity) ON (e.id)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX entity_id_index IF NOT EXISTS FOR (e:Entity) ON (e.id)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX entity_name_index IF NOT EXISTS FOR (e:Entity) ON (e.name)' has no effect. The index or constraint specified by 'RANGE INDEX entity_name_index FOR (e:Entity) ON (e.name)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX entity_name_index IF NOT EXISTS FOR (e:Entity) ON (e.name)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX entity_type_index IF NOT EXISTS FOR (e:Entity) ON (e.type)' has no effect. The index or constraint specified by 'RANGE INDEX entity_type_index FOR (e:Entity) ON (e.type)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX entity_type_index IF NOT EXISTS FOR (e:Entity) ON (e.type)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX memory_id_index IF NOT EXISTS FOR (e:Memory) ON (e.id)' has no effect. The index or constraint specified by 'RANGE INDEX memory_id_index FOR (e:Memory) ON (e.id)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX memory_id_index IF NOT EXISTS FOR (m:Memory) ON (m.id)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX memory_type_index IF NOT EXISTS FOR (e:Memory) ON (e.memory_type)' has no effect. The index or constraint specified by 'RANGE INDEX memory_type_index FOR (e:Memory) ON (e.memory_type)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX memory_type_index IF NOT EXISTS FOR (m:Memory) ON (m.memory_type)' INFO:neo4j.notifications:Received notification from DBMS server: <GqlStatusObject gql_status='00NA0', status_description="note: successful completion - index or constraint already exists. The command 'CREATE RANGE INDEX memory_timestamp_index IF NOT EXISTS FOR (e:Memory) ON (e.timestamp)' has no effect. The index or constraint specified by 'RANGE INDEX memory_timestamp_index FOR (e:Memory) ON (e.timestamp)' already exists.", position=None, raw_classification='SCHEMA', classification=<NotificationClassification.SCHEMA: 'SCHEMA'>, raw_severity='INFORMATION', severity=<NotificationSeverity.INFORMATION: 'INFORMATION'>, diagnostic_record={'_classification': 'SCHEMA', '_severity': 'INFORMATION', 'OPERATION': '', 'OPERATION_CODE': '0', 'CURRENT_SCHEMA': '/'}> for query: 'CREATE INDEX memory_timestamp_index IF NOT EXISTS FOR (m:Memory) ON (m.timestamp)' INFO:storage.neo4j_store:✅ Neo4j索引创建完成 INFO:memory.types.semanticMemory:✅ 加载中文spaCy模型: zh_core_web_sm INFO:memory.types.semanticMemory:✅ 加载英文spaCy模型: en_core_web_sm INFO:memory.types.semanticMemory:🎯 主要使用中文spaCy模型 INFO:memory.types.semanticMemory:📚 可用语言模型: 中文, 英文 INFO:memory.types.semanticMemory:增强语义记忆初始化完成(使用Qdrant+Neo4j专业数据库) INFO:memory.memoryManager:MemoryManager初始化完成,启用记忆类型: ['working', 'episodic', 'semantic'] ✅ 工具 'memory' 已注册。 === 添加多个记忆 === 记忆1: ✅ 记忆已添加 ID: 3ef16625-c077-4002-9725-e069e089f6d2 类型: working 内容: 用户张三是一名Python开发者,专注于机器学习和数据分析 INFO:storage.qdrant_store:[Qdrant] add_vectors start: n_vectors=1 n_meta=1 collection=helloAgent INFO:storage.qdrant_store:[Qdrant] upsert begin: points=1 INFO:httpx:HTTP Request: PUT http://localhost:6333/collections/helloAgent/points?wait=true "HTTP/1.1 200 OK" INFO:storage.qdrant_store:[Qdrant] upsert done INFO:storage.qdrant_store:✅ 成功添加 1 个向量到Qdrant INFO:memory.types.semanticMemory:✅ 添加语义记忆: 2个实体, 1个关系 记忆2: ✅ 记忆已添加 ID: 5d3dbf35-7f8f-46b3-b3f2-970d80640b17 类型: semantic 内容: 李四是前端工程师,擅长React和Vue.js开发 记忆3: ✅ 记忆已添加 ID: b5857a65-7ee6-4e47-b177-6dff006710cc 类型: working 内容: 王五是产品经理,负责用户体验设计和需求分析 === 搜索特定记忆 === 🔍 搜索 '前端工程师': INFO:httpx:HTTP Request: POST http://localhost:6333/collections/helloAgent/points/query "HTTP/1.1 200 OK" INFO:httpx:HTTP Request: POST http://localhost:6333/collections/helloAgent/points/query "HTTP/1.1 200 OK" INFO:memory.types.semanticMemory:✅ 检索到 1 条相关记忆 🔍 搜索 '前端工程师' 的结果(共 1 条): [1] [semantic] 李四是前端工程师,擅长React和Vue.js开发 (importance=0.70, id=5d3dbf35...) === 记忆摘要 === INFO:httpx:HTTP Request: GET http://localhost:6333/collections/helloAgent "HTTP/1.1 200 OK" 📊 记忆系统摘要: 用户: user123 总记忆数: 3 启用的记忆类型: working, episodic, semantic 配置: 最大容量=100, 重要性阈值=0.1, 衰减因子=0.95 [working] 活跃=2, 容量使用率=20.0% [episodic] 活跃=0, 容量使用率=0.0% [semantic] 活跃=1, 容量使用率=0.0%
成功了。
rag 检索增强生成 Retrieval-Augmented Generation,其核心思想就是:将知识库注入 prompt,补足 LLM 的知识水平或是让其遵守一定事实。
体验官方 RAG 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 from agent.simpleAgent import SimpleAgentfrom core.llm import LLMfrom tool.toolRegistry import ToolRegistryfrom hello_agents.tools import RAGToolllm = LLM() agent = SimpleAgent(name="知识助手" , llm=llm) rag_tool = RAGTool( knowledge_base_path="./knowledge_base" , collection_name="test_collection" , rag_namespace="test" ) tool_registry = ToolRegistry() tool_registry.register_tool(rag_tool) agent.tool_registry = tool_registry result1 = rag_tool.run({"action" :"add_text" , "text" :"Python是一种高级编程语言,由Guido van Rossum于1991年首次发布。Python的设计哲学强调代码的可读性和简洁的语法。" , "document_id" :"python_intro" }) print (f"知识1: {result1} " )result2 = rag_tool.run({"action" :"add_text" , "text" :"机器学习是人工智能的一个分支,通过算法让计算机从数据中学习模式。主要包括监督学习、无监督学习和强化学习三种类型。" , "document_id" :"ml_basics" }) print (f"知识2: {result2} " )result3 = rag_tool.run({"action" :"add_text" , "text" :"RAG(检索增强生成)是一种结合信息检索和文本生成的AI技术。它通过检索相关知识来增强大语言模型的生成能力。" , "document_id" :"rag_concept" }) print (f"知识3: {result3} " )print ("\n=== 搜索知识 ===" )result = rag_tool.run({"action" :"search" , "query" :"Python编程语言的历史" , "limit" :3 , "min_score" :0.1 }) print (result)print ("\n=== 知识库统计 ===" )result = rag_tool.run({"action" :"stats" }) print (result)
结果:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 > python -m test.__testRagTool ✅ 知识助手 初始化完成,工具调用: 禁用 INFO:hello_agents.memory.storage.qdrant_store:✅ 成功连接到Qdrant服务: http://localhost:6333 INFO:hello_agents.memory.storage.qdrant_store:✅ 使用现有Qdrant集合: test_collection ✅ RAG工具初始化成功: namespace=test, collection=test_collection ✅ 工具 'rag' 已注册。 INFO:hello_agents.memory.storage.qdrant_store:✅ 成功连接到Qdrant服务: http://localhost:6333 INFO:hello_agents.memory.storage.qdrant_store:✅ 使用现有Qdrant集合: test_collection [RAG] Universal loader start: files=1 chunk_size=800 overlap=100 ns=default [RAG] Processing: ./knowledge_base\python_intro.md C:\Users\a1829\AppData\Local\Programs\Python\Python312\Lib\site-packages\pydub\utils.py:170: RuntimeWarning: Couldn't find ffmpeg or avconv - defaulting to ffmpeg, but may not work warn("Couldn't find ffmpeg or avconv - defaulting to ffmpeg, but may not work", RuntimeWarning) [RAG] Universal loader done: total_chunks=1 [RAG] Embedding start: total_texts=1 batch_size=64 [RAG] Embedding progress: 1/1 [RAG] Qdrant upsert start: n=1 INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] add_vectors start: n_vectors=1 n_meta=1 collection=test_collection INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert begin: points=1 INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert done INFO:hello_agents.memory.storage.qdrant_store:✅ 成功添加 1 个向量到Qdrant [RAG] Qdrant upsert done: 1 vectors indexed 知识1: ✅ 文本已添加到知识库: python_intro 📊 分块数量: 1 ⏱️ 处理时间: 1520ms 📝 命名空间: default [RAG] Universal loader start: files=1 chunk_size=800 overlap=100 ns=default [RAG] Processing: ./knowledge_base\ml_basics.md [RAG] Universal loader done: total_chunks=1 [RAG] Embedding start: total_texts=1 batch_size=64 [RAG] Embedding progress: 1/1 [RAG] Qdrant upsert start: n=1 INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] add_vectors start: n_vectors=1 n_meta=1 collection=test_collection INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert begin: points=1 INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert done INFO:hello_agents.memory.storage.qdrant_store:✅ 成功添加 1 个向量到Qdrant [RAG] Qdrant upsert done: 1 vectors indexed 知识2: ✅ 文本已添加到知识库: ml_basics 📊 分块数量: 1 ⏱️ 处理时间: 510ms 📝 命名空间: default [RAG] Universal loader start: files=1 chunk_size=800 overlap=100 ns=default [RAG] Processing: ./knowledge_base\rag_concept.md [RAG] Universal loader done: total_chunks=1 [RAG] Embedding start: total_texts=1 batch_size=64 [RAG] Embedding progress: 1/1 [RAG] Qdrant upsert start: n=1 INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] add_vectors start: n_vectors=1 n_meta=1 collection=test_collection INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert begin: points=1 INFO:hello_agents.memory.storage.qdrant_store:[Qdrant] upsert done INFO:hello_agents.memory.storage.qdrant_store:✅ 成功添加 1 个向量到Qdrant [RAG] Qdrant upsert done: 1 vectors indexed 知识3: ✅ 文本已添加到知识库: rag_concept 📊 分块数量: 1 ⏱️ 处理时间: 470ms 📝 命名空间: default === 搜索知识 === 搜索结果: 1. 文档: **./knowledge_base\python_intro.md** (相似度: 0.702) Python是一种高级编程语言,由Guido van Rossum于1991年首次发布。Python的设计哲学强调代码的可读性和简洁的语法。... 2. 文档: **./knowledge_base\rag_concept.md** (相似度: 0.241) RAG(检索增强生成)是一种结合信息检索和文本生成的AI技术。它通过检索相关知识来增强大语言模型的生成能力。... 3. 文档: **./knowledge_base\ml_basics.md** (相似度: 0.218) 机器学习是人工智能的一个分支,通过算法让计算机从数据中学习模式。主要包括监督学习、无监督学习和强化学习三种类型。... === 知识库统计 === 📊 **RAG 知识库统计** 📝 命名空间: default 📋 集合名称: test_collection 📂 存储根路径: ./knowledge_base 📦 存储类型: qdrant 📊 文档分块数: 3 🔢 向量维度: 1024 📎 距离度量: Cosine 🟢 **系统状态** ✅ RAG 管道: 正常 ✅ LLM 连接: 正常
本节目标 1 2 3 4 5 6 7 8 9 HelloAgent/ │ ├── rag/ │ └── pipeline.py # 管线 │ ├── tool/ │ ├── toolAction.py # Memory 基类 │ └── builtin │ └── ragTool.py # RAG 工具包
Pipeline 基本流程 RAGTool 是一种工具,那么这就和前面提到的 MemoryTool 存在着一样的问题:RAGTool 是对外暴露的接口,它是否应该直接引用底层存储组件(例如QdrantVectorStore、Neo4j等数据库接口)?不要,我们需要构建中间层。不过,官方的意图不只是为了抽象封装,更主要的原因也是 RAGTool 到底层数据库之间的操作,其本身也便于梳理成一个工作流,也就是 Pipeline 的由来。
在官方的设计中,pipeline 是一个包含了若干数据处理的函数库,包括多模态文件统一转文本、数据分块、向量化等,其在 RAGTool 的调用流程大致是:
任意格式文档 → MarkItDown转换 → Markdown文本 → 智能分块 → 向量化 → 存储检索
该流程只涉及到了一些函数,官方提出来重点进行了介绍。
多模态文档载入 RAG 直接对接知识库,最直接最权威的知识来源就是各种文件,包括 pdf、txt 等文本文件,或者 mp3、avi、img 等多模态文件,我们希望 RAG 能够直接接受各种模态的文件。
Pipelin 使用 MarkItDown 作为统一的文档转换引擎,支持几乎所有常见的文档格式,无论输入是 PDF、Word、Excel、图片还是音频,最终都会转换为标准的Markdown格式,然后进入统一的分块、向量化和存储流程。
1 def _convert_to_markdown (path: str ) -> str :
函数返回一个存储了 md 格式原文的字符串。
智能分块 使用 MarkItDown 将文件转为 md 格式后,需要进行分块,可以利用Markdown的标题结构(#、##、###等)进行精确的语义分割。流程为:
1 2 3 4 标准Markdown文本 → 标题层次解析 → 段落语义分割 → Token计算分块 → 重叠策略优化 → 向量化准备 ↓ ↓ ↓ ↓ ↓ ↓ 统一格式 #/##/### 语义边界 大小控制 信息连续性 嵌入向量 结构清晰 层次识别 完整性保证 检索优化 上下文保持 相似度匹配W
首先,对完整 md 文本使用下列函数:
1 def _split_paragraphs_with_headings (text: str ) -> List [Dict ]:
这个函数能够完成 标题层次解析 和 段落语义分割,将完整的 md 文本转化为若干小块,分别代表局部完整的文本内容。
随后,对字典使用下列函数:
1 def _chunk_paragraphs (paragraphs: List [Dict ], chunk_tokens: int , overlap_tokens: int ) -> List [Dict ]:
该函数进一步完成后三个步骤,将每个字典进一步处理为分块文本,以能够直接向量化。
统一嵌入与向量存储 经过前面两个阶段的处理后,我们终于将各种文件,通过各种处理,将其变得能够存入向量数据库,也即加入知识库了。
嵌入模型负责将文本转换为高维向量,使得计算机能够理解和比较文本的语义相似性。
下列函数负责将分块文本存入知识库:
1 2 3 4 5 6 7 def index_chunks ( store = None , chunks: List [Dict ] = None , cache_db: Optional [str ] = None , batch_size: int = 64 , rag_namespace: str = "default" ) -> None :
Pipeline 高级扩展 上面的三大流程囊括了文件从输入到处理到存储的全过程,属于基础功能,而在此之外,官方的 pipeline 还提供了众多高级功能,这是为了解决很有可能出现的问题,官方将其描述为:
RAG系统的检索能力是其核心竞争力,在实际应用中,用户的查询表述与文档中的实际内容可能存在用词差异,导致相关文档无法被检索到。
官方提出了三种高级检索策略进行着重讲解:多查询扩展(MQE)、假设文档嵌入(HyDE)和统一的扩展检索框架。
这三种高级策略,有助于进一步提升 RAG 搜索的召回率。
多查询扩展 (Multi-Query Expansion, MQE) 该功能的核心思想是:同一个问题有多种不同的表述方式。我们知道 RAG 系统实际上并非直接进行逻辑判断,而是对 query 本身进行神经网络式的黑箱匹配,这就表示,同一个逻辑事件,采用多种表示方式,同时作为 query 进行查询,能够匹配出更多、更精确的结果。这就是 多查询扩展 。
思想很简单,实现上也很简单,多种方法表达让 LLM 进行就可以:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 def _prompt_mqe (query: str , n: int ) -> List [str ]: """使用LLM生成多样化的查询扩展""" try : llm = LLM() prompt = [ {"role" : "system" , "content" : "你是检索查询扩展助手。生成语义等价或互补的多样化查询。使用中文,简短,避免标点。" }, {"role" : "user" , "content" : f"原始查询:{query} \n请给出{n} 个不同表述的查询,每行一个。" } ] text = llm.think(prompt) lines = [ln.strip("- \t" ) for ln in (text or "" ).splitlines()] outs = [ln for ln in lines if ln] return outs[:n] or [query] except Exception: return [query]
假设文档嵌入(Hypothetical Document Embeddings ,HyDE) 这个方法的重要前提是,query 和 answer 之间存在一种现象,在语义空间中的分布往往存在差异 —— 问题通常是疑问句,而文档内容是陈述句。
这种现象带来一种问题,即从我们人类的思维上看,疑问对应解答,这是十分合理且理所应当的,但是在 RAG 系统中,疑问句与陈述句之间的语义特征几乎完全不同。
疑问句的向量,很大一部分维度被句法结构占据,在 RAG 看来:
疑问句 = 疑问词(什么/如何/为什么/what/how/why)+ 倒装语序(英文)+ 极短的长度 + 低占比的术语密度
而陈述句则几乎全部维度被语义内容占据,在 RAG 看来:
陈述句 = 少量语法结构词 + 高占比的术语密度
如果我们要搜索真正的术语,那么疑问句中的噪音未免过多,因为疑问句和陈述句完全是两个物种,即 句法差异贡献了主要方差。
而实际上,在 RAG 看来,疑问句并非直接用来存储知识库,反而是另一个陈述句的指针,比如疑问句 什么是机器学习? 反而是在引导系统向另一个陈述句 机器学习是人工智能的一个分支,它使计算机系统能够从数据中... 联想。
所以,如果利用 LLM 生成一个假想的陈述句回答(重点是陈述而非解答),例如 LLM 通过问题 什么是机器学习? 生成了解答 机器学习是人工智能的一个分支,它使计算机系统能够从数据中...,我们再通过 rag 搜索就很有可能搜索到相关度高的解答。
1 2 3 4 5 6 7 8 9 10 11 def _prompt_hyde (query: str ) -> Optional [str ]: """生成假设性文档用于改善检索""" try : llm = LLM() prompt = [ {"role" : "system" , "content" : "根据用户问题,先写一段可能的答案性段落,用于向量检索的查询文档(不要分析过程)。" }, {"role" : "user" , "content" : f"问题:{query} \n请直接写一段中等长度、客观、包含关键术语的段落。" } ] return llm.think(prompt) except Exception: return None
不过一个问题是,让 LLM 生成假想答案,如果问题是 什么是机器学习? 这种在网络中存在客观解答的问题,LLM 确实可以很简单地生成假想答案(这个答案甚至有可能比 RAG 知识库中还要准确);但如果是 我家的门牌号是多少? 这种网络知识库不可能知道的问题, LLM 应该生成什么?
LLM 很有可能生成三种回答:
1.胡说八道,例如 “住宅门牌号是用于邮政投递和地址识别的数字标识系统,通常由当地政府统一规划分配,门牌号的编码规则因地区而异,常见的有顺序编号和距离编号两种方式…”
2.拒绝回答,例如 “我无法知道您家的具体门牌号,因为这属于个人隐私信息。”
3.歧义回复,例如 “用户家的门牌号是 128 号,位于社区的核心区域。”
如果的第一种,看似牛头不对马嘴,实则反而会其效果,因为解答中存在 “住宅门牌号”、”邮政投递”、”码规则” 等和门牌号高度相关的名词,如果 RAG 中确实存在 用户家的门牌号是 xxx 号 这个权威结果,反而会提升搜索到的概率。
如果是第二种安全性拒绝,那就无功无过,至少未产生有害影响。
如果是第三种歧义回复,很有可能造成有害影响。我家的门牌号是多少? 这种问题还不是最狠的,更狠的是确实存在复数含义的名词,例如 Java 是什么?,LLM 极大概率生成 “Java 是一种广泛使用的编程语言,由 Sun Microsystems 于 1995 年发布…”,然而在 RAG 中,”Java 是隶属于印度尼西亚的一个岛屿。”,那么就会造成极其深刻的有害影响。
扩展检索框架 结合 MQE 和 HyDE 的搜索强化策略,通过 enable_mqe 和 enable_hyde 参数让用户可以根据具体场景选择启用哪些策略:对于需要高召回率的场景可以同时启用两种策略,对于性能敏感的场景可以只使用基础检索。
先使用 MQE 拓展为多个同义疑问句,再对每个疑问句采用 HyDE 改写为陈述句,最后合并排序返回最相关的 top-k 文档。
1 2 3 4 5 6 7 8 9 10 11 12 def search_vectors_expanded ( store = None , query: str = "" , top_k: int = 8 , rag_namespace: Optional [str ] = None , only_rag_data: bool = True , score_threshold: Optional [float ] = None , enable_mqe: bool = False , mqe_expansions: int = 2 , enable_hyde: bool = False , candidate_pool_multiplier: int = 4 , ) -> List [Dict ]:
RAGTool 需要注解 tool_action:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 from typing import Callable def tool_action (name: str = None , description: str = None ): """装饰器:标记一个方法为可展开的工具 action 用法: @tool_action("memory_add", "添加新记忆") def _add_memory(self, content: str, importance: float = 0.5) -> str: '''添加记忆 Args: content: 记忆内容 importance: 重要性分数 ''' ... Args: name: 工具名称(如果不提供,从方法名自动生成) description: 工具描述(如果不提供,从 docstring 提取) """ def decorator (func: Callable ): func._is_tool_action = True func._tool_name = name func._tool_description = description return func return decorator
RAGTool 是对外工具,其设计与 memoryTool 一样,使用统一接口输入操作符调用,官方演示:
1 2 3 4 5 6 7 8 9 10 11 12 result1 = rag_tool.run({ "action" :"add_text" , "text" :"Python是一种高级编程语言,由Guido van Rossum于1991年首次发布。Python的设计哲学强调代码的可读性和简洁的语法。" , "document_id" :"python_intro" }) result = rag_tool.run({"action" :"search" , "query" :"Python编程语言的历史" , "limit" :3 , "min_score" :0.1 })
我们通过 ragTool 的 run 函数,给定一个字典进行 RAG 操作,自然,run 函数就是统一接口和学习重点。
run 函数调用体系 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 def run (self, parameters: Dict [str , Any ] ) -> str : """执行工具(非展开模式) Args: parameters: 工具参数字典,必须包含action参数 Returns: 执行结果字符串 """ if not self.validate_parameters(parameters): return "❌ 参数验证失败:缺少必需的参数" if not self.initialized: return f"❌ RAG工具未正确初始化,请检查配置: {getattr (self, 'init_error' , '未知错误' )} action = parameters.get(" action") # 根据action调用对应的方法,传入提取的参数 try: if action == " add_document": return self._add_document( file_path=parameters.get(" file_path"), document_id=parameters.get(" document_id"), namespace=parameters.get(" namespace", " default"), chunk_size=parameters.get(" chunk_size", 800), chunk_overlap=parameters.get(" chunk_overlap", 100) ) elif action == " add_text": return self._add_text( text=parameters.get(" text"), # ······ else: return f" ❌ 不支持的操作: {action}" except Exception as e: return f" ❌ 执行操作 '{action}' 时发生错误: {str (e)}"
可以看到,run 函数内解析参数 action 后,就是将其路由到各自的函数,我们以 _add_document 为例,看看里面是什么:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 @tool_action("rag_add_document" , "添加文档到知识库(支持PDF、Word、Excel、PPT、图片、音频等多种格式)" ) def _add_document ( self, file_path: str , document_id: str = None , namespace: str = "default" , chunk_size: int = 800 , chunk_overlap: int = 100 ) -> str : """添加文档到知识库 Args: file_path: 文档文件路径 document_id: 文档ID(可选) namespace: 知识库命名空间(用于隔离不同项目) chunk_size: 分块大小 chunk_overlap: 分块重叠大小 Returns: 执行结果 """ try : if not file_path or not os.path.exists(file_path): return f"❌ 文件不存在: {file_path} " pipeline = self._get_pipeline(namespace) t0 = time.time() chunks_added = pipeline["add_documents" ]( file_paths=[file_path], chunk_size=chunk_size, chunk_overlap=chunk_overlap ) t1 = time.time() process_ms = int ((t1 - t0) * 1000 ) if chunks_added == 0 : return f"⚠️ 未能从文件解析内容: {os.path.basename(file_path)} " return ( f"✅ 文档已添加到知识库: {os.path.basename(file_path)} \n" f"📊 分块数量: {chunks_added} \n" f"⏱️ 处理时间: {process_ms} ms\n" f"📝 命名空间: {pipeline.get('namespace' , self.rag_namespace)} " ) except Exception as e: return f"❌ 添加文档失败: {str (e)} "
这里终于用到我们之前详细讲述的 pipeline 了。前面提到,pipeline 本质是一个闭包字典,字典所有的值都是对底层数据库的 crud,pipeline["add_documents"] 实际上就是一个函数,因为 pipeline 内部的一个字典是:
1 "add_documents" : add_documents,
add_documents 是 pipeline 内部的一个函数,直接操作底层数据库了。
前面提到,pipeline 的字典在大多数时候是固定的,即:
1 2 3 4 5 6 7 8 { "store" : store, "namespace" : rag_namespace, "add_documents" : add_documents, "search" : search, "search_advanced" : search_advanced, "get_stats" : get_stats }
我们上面提到的 RAGTool 函数 _add_document 当然直接由 pipeline:add_documents 代理了。但是,我们同时也注意到 RAGTool 里面同时存在同样为添加数据的函数 _add_text,而 pipeline 内部根本没有直接的代理函数,那么这个时候怎么办?
让我们看一看 _add_text 函数:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 @tool_action("rag_add_text" , "添加文本到知识库" ) def _add_text ( self, text: str , document_id: str = None , namespace: str = "default" , chunk_size: int = 800 , chunk_overlap: int = 100 ) -> str : """添加文本到知识库 Args: text: 要添加的文本内容 document_id: 文档ID(可选) namespace: 知识库命名空间 chunk_size: 分块大小 chunk_overlap: 分块重叠大小 Returns: 执行结果 """ metadata = None try : if not text or not text.strip(): return "❌ 文本内容不能为空" document_id = document_id or f"text_{abs (hash (text)) % 100000 } " tmp_path = os.path.join(self.knowledge_base_path, f"{document_id} .md" ) try : with open (tmp_path, 'w' , encoding='utf-8' ) as f: f.write(text) pipeline = self._get_pipeline(namespace) t0 = time.time() chunks_added = pipeline["add_documents" ]( file_paths=[tmp_path], chunk_size=chunk_size, chunk_overlap=chunk_overlap )
注意到,用的居然还是 add_documents。
不仅如此,让我们看看 RAGTool 里面的两个搜索函数 _search 和 _ask:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 @tool_action("rag_search" , "搜索知识库中的相关内容" ) def _search ( self, query: str , limit: int = 5 , min_score: float = 0.1 , enable_advanced_search: bool = True , max_chars: int = 1200 , include_citations: bool = True , namespace: str = "default" ) -> str : try : if not query or not query.strip(): return "❌ 搜索查询不能为空" pipeline = self._get_pipeline(namespace) if enable_advanced_search: results = pipeline["search_advanced" ]( query=query, top_k=limit, enable_mqe=True , enable_hyde=True , score_threshold=min_score if min_score > 0 else None ) else : results = pipeline["search" ]( query=query, top_k=limit, score_threshold=min_score if min_score > 0 else None ) @tool_action("rag_ask" , "基于知识库进行智能问答" ) def _ask ( self, question: str , limit: int = 5 , enable_advanced_search: bool = True , include_citations: bool = True , max_chars: int = 1200 , namespace: str = "default" ) -> str : try : if not question or not question.strip(): return "❌ 请提供要询问的问题" user_question = question.strip() print (f"🔍 智能问答: {user_question} " ) pipeline = self._get_pipeline(namespace) search_start = time.time() if enable_advanced_search: results = pipeline["search_advanced" ]( query=user_question, top_k=limit, enable_mqe=True , enable_hyde=True ) else : results = pipeline["search" ]( query=user_question, top_k=limit )
我们会发现,两个函数调用的均是 pipeline 的 search_advanced 和 search ,只是在调用顺序、后续处理的环节上存在较大差异。
所以我们会发现,pipeline 只负责 CRUD,只暴露出少量单一职责接口;而 RAGTool 负责基于单一职责接口进行功能拓展。
总结 本章将重心放在了 Agent 记忆构建,与记忆检索上,先后讲述了 hello agent 的记忆体系、记忆如何持久化、如何从持久化底层检索数据。
针对记忆体系,我们遵照了自下而上的构建原则,首先完成了记忆实体类和配置类;随后在此之上构建了抽象类 Memory,定义了记忆应该有的行为,并在此之上构建了三大记忆之一的工作记忆;而避免外部接口 MemoryTool 和底层记忆类 Memory 的直接接触,我们设计了 MemoryManager,供 MemoryTool 调用从而代理记忆管理。
在记忆体系构建完毕后,我们转向了持久化的设计,我们基于已有的 Qdrant、Neo4j、sqlite 三大数据库和 embedding 嵌入模型的库函数,设计了 storage 包,并完成了情景记忆与语义记忆这两个需要将记忆持久化的记忆类型,初步建立了 RAG 底层支持。
在存储体系趋于完善后,我们设计了 rag 包,其设计直接对标 memory 的工具体系,自下而上地构建了工具和底层直接的代理组件 pipeline 和直接对外暴露的统一 RAG 操作工具 RAGTool。我们首先学习了 pipeline 实现的主要职责与流程,其负责将文件统一转为文本格式,并通过分块、向量化等操作将文件存入向量数据库,初步实现 RAG 功能;随后我们还学习了官方介绍的两种提升 RAG 搜索精度的方法:MQE 和 HyDE;最后,我们详细学习了 RAGTool 的 run 函数设计,通过对同一个函数输入不同的操作符与参数,实现了一个函数统一路由到其他功能,也即仅对外暴露一个接口便实现全部功能的目的。
综上,我们通过构建记忆与检索体系,将 agent 的能力进一步增强。然而,记忆与搜索也仅仅是对 agent 的特定能力进行了增强,我们需要研究如何让 agent 以最正确、最及时、最有性价比地进行记忆与检索,这一点将在下一章进行。
本章代码 memory memoryItem.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 from datetime import datetimefrom pydantic import BaseModelfrom typing import Any , Dict class MemoryItem (BaseModel ): """记忆项数据结构""" id : str content: str memory_type: str user_id: str timestamp: datetime importance: float = 0.5 metadata: Dict [str , Any ] = {} arbitrary_types_allowed: bool = True
memoryConfig.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 from pydantic import BaseModelfrom typing import List class MemoryConfig (BaseModel ): """记忆系统配置""" storage_path: str = "./memory_data" max_capacity: int = 100 importance_threshold: float = 0.1 decay_factor: float = 0.95 working_memory_capacity: int = 10 working_memory_tokens: int = 2000 working_memory_ttl_minutes: int = 120 perceptual_memory_modalities: List [str ] = ["text" , "image" , "audio" , "video" ]
memory.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 from abc import ABC, abstractmethodfrom typing import List , Dict , Any from .memoryConfg import MemoryConfig from .memoryItem import MemoryItemclass Memory (ABC ): """记忆基类 定义所有记忆类型的通用接口和行为 """ def __init__ (self, config: MemoryConfig, storage_backend=None ): self.config = config self.storage = storage_backend self.memory_type = self.__class__.__name__.lower().replace("memory" , "" ) @abstractmethod def add (self, memory_item: MemoryItem ) -> str : """添加记忆项 Args: memory_item: 记忆项对象 Returns: 记忆ID """ pass @abstractmethod def retrieve (self, query: str , limit: int = 5 , **kwargs ) -> List [MemoryItem]: """检索相关记忆 Args: query: 查询内容 limit: 返回数量限制 **kwargs: 其他检索参数 Returns: 相关记忆列表 """ pass @abstractmethod def update (self, memory_id: str , content: str = None , importance: float = None , metadata: Dict [str , Any ] = None ) -> bool : """更新记忆 Args: memory_id: 记忆ID content: 新内容 importance: 新重要性 metadata: 新元数据 Returns: 是否更新成功 """ pass @abstractmethod def remove (self, memory_id: str ) -> bool : """删除记忆 Args: memory_id: 记忆ID Returns: 是否删除成功 """ pass @abstractmethod def has_memory (self, memory_id: str ) -> bool : """检查记忆是否存在 Args: memory_id: 记忆ID Returns: 是否存在 """ pass @abstractmethod def clear (self ): """清空所有记忆""" pass @abstractmethod def get_stats (self ) -> Dict [str , Any ]: """获取记忆统计信息 Returns: 统计信息字典 """ pass
types workingMemory.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 from typing import List , Dict , Any from memory.memoryConfg import MemoryConfigfrom memory.memoryItem import MemoryItemfrom memory.memory import Memoryfrom datetime import datetime, timedeltaclass WorkingMemory (Memory ): """工作记忆实现 特点: - 容量有限(默认50条)+ TTL自动清理 - 纯内存存储,访问速度极快 - 混合检索:TF-IDF向量化 + 关键词匹配 """ def __init__ (self, config: MemoryConfig, storage_backend=None ): super ().__init__(config, storage_backend) self.max_capacity = self.config.working_memory_capacity self.max_tokens = self.config.working_memory_tokens self.max_age_minutes = getattr (self.config, 'working_memory_ttl_minutes' , 120 ) self.current_tokens = 0 self.session_start = datetime.now() self.memories: List [MemoryItem] = [] def add (self, memory_item: MemoryItem ) -> str : """添加工作记忆""" self._expire_old_memories() priority = self._calculate_priority(memory_item) self.memories.append(memory_item) self.current_tokens += len (memory_item.content.split()) self._enforce_capacity_limits() return memory_item.id def retrieve (self, query: str , limit: int = 5 , user_id: str = None , **kwargs ) -> List [MemoryItem]: """检索工作记忆 - 混合语义向量检索和关键词匹配""" self._expire_old_memories() if not self.memories: return [] active_memories = self.memories filtered_memories = active_memories if user_id: filtered_memories = [m for m in active_memories if m.user_id == user_id] if not filtered_memories: return [] vector_scores = {} try : from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity import numpy as np documents = [query] + [m.content for m in filtered_memories] vectorizer = TfidfVectorizer(stop_words=None , lowercase=True ) tfidf_matrix = vectorizer.fit_transform(documents) query_vector = tfidf_matrix[0 :1 ] doc_vectors = tfidf_matrix[1 :] similarities = cosine_similarity(query_vector, doc_vectors).flatten() for i, memory in enumerate (filtered_memories): vector_scores[memory.id ] = similarities[i] except Exception as e: vector_scores = {} query_lower = query.lower() scored_memories = [] for memory in filtered_memories: content_lower = memory.content.lower() vector_score = vector_scores.get(memory.id , 0.0 ) keyword_score = 0.0 if query_lower in content_lower: keyword_score = len (query_lower) / len (content_lower) else : query_words = set (query_lower.split()) content_words = set (content_lower.split()) intersection = query_words.intersection(content_words) if intersection: keyword_score = len (intersection) / len (query_words.union(content_words)) * 0.8 if vector_score > 0 : base_relevance = vector_score * 0.7 + keyword_score * 0.3 else : base_relevance = keyword_score time_decay = self._calculate_time_decay(memory.timestamp) base_relevance *= time_decay importance_weight = 0.8 + (memory.importance * 0.4 ) final_score = base_relevance * importance_weight if final_score > 0 : scored_memories.append((final_score, memory)) scored_memories.sort(key=lambda x: x[0 ], reverse=True ) return [memory for _, memory in scored_memories[:limit]] def update ( self, memory_id: str , content: str = None , importance: float = None , metadata: Dict [str , Any ] = None ) -> bool : """更新工作记忆""" for memory in self.memories: if memory.id == memory_id: old_tokens = len (memory.content.split()) if content is not None : memory.content = content new_tokens = len (content.split()) self.current_tokens = self.current_tokens - old_tokens + new_tokens if importance is not None : memory.importance = importance if metadata is not None : memory.metadata.update(metadata) return True return False def remove (self, memory_id: str ) -> bool : """删除工作记忆""" for i, memory in enumerate (self.memories): if memory.id == memory_id: removed_memory = self.memories.pop(i) self.current_tokens -= len (removed_memory.content.split()) self.current_tokens = max (0 , self.current_tokens) return True return False def has_memory (self, memory_id: str ) -> bool : """检查记忆是否存在""" return any (memory.id == memory_id for memory in self.memories) def clear (self ): """清空所有工作记忆""" self.memories.clear() self.current_tokens = 0 def get_stats (self ) -> Dict [str , Any ]: """获取工作记忆统计信息""" self._expire_old_memories() active_memories = self.memories return { "count" : len (active_memories), "forgotten_count" : 0 , "total_count" : len (self.memories), "current_tokens" : self.current_tokens, "max_capacity" : self.max_capacity, "max_tokens" : self.max_tokens, "max_age_minutes" : self.max_age_minutes, "session_duration_minutes" : (datetime.now() - self.session_start).total_seconds() / 60 , "avg_importance" : sum (m.importance for m in active_memories) / len (active_memories) if active_memories else 0.0 , "capacity_usage" : len (active_memories) / self.max_capacity if self.max_capacity > 0 else 0.0 , "token_usage" : self.current_tokens / self.max_tokens if self.max_tokens > 0 else 0.0 , "memory_type" : "working" } def get_recent (self, limit: int = 10 ) -> List [MemoryItem]: """获取最近的记忆""" sorted_memories = sorted ( self.memories, key=lambda x: x.timestamp, reverse=True ) return sorted_memories[:limit] def get_important (self, limit: int = 10 ) -> List [MemoryItem]: """获取重要记忆""" sorted_memories = sorted ( self.memories, key=lambda x: x.importance, reverse=True ) return sorted_memories[:limit] def get_all (self ) -> List [MemoryItem]: """获取所有记忆""" return self.memories.copy() def get_context_summary (self, max_length: int = 500 ) -> str : """获取上下文摘要""" if not self.memories: return "No working memories available." sorted_memories = sorted ( self.memories, key=lambda m: (m.importance, m.timestamp), reverse=True ) summary_parts = [] current_length = 0 for memory in sorted_memories: content = memory.content if current_length + len (content) <= max_length: summary_parts.append(content) current_length += len (content) else : remaining = max_length - current_length if remaining > 50 : summary_parts.append(content[:remaining] + "..." ) break return "Working Memory Context:\n" + "\n" .join(summary_parts) def forget (self, strategy: str = "importance_based" , threshold: float = 0.1 , max_age_days: int = 1 ) -> int : """工作记忆遗忘机制""" forgotten_count = 0 current_time = datetime.now() to_remove = [] cutoff_ttl = current_time - timedelta(minutes=self.max_age_minutes) for memory in self.memories: if memory.timestamp < cutoff_ttl: to_remove.append(memory.id ) if strategy == "importance_based" : for memory in self.memories: if memory.importance < threshold: to_remove.append(memory.id ) elif strategy == "time_based" : cutoff_time = current_time - timedelta(hours=max_age_days * 24 ) for memory in self.memories: if memory.timestamp < cutoff_time: to_remove.append(memory.id ) elif strategy == "capacity_based" : if len (self.memories) > self.max_capacity: sorted_memories = sorted ( self.memories, key=lambda m: self._calculate_priority(m) ) excess_count = len (self.memories) - self.max_capacity for memory in sorted_memories[:excess_count]: to_remove.append(memory.id ) for memory_id in to_remove: if self.remove(memory_id): forgotten_count += 1 return forgotten_count def _calculate_priority (self, memory: MemoryItem ) -> float : """计算记忆优先级""" priority = memory.importance time_decay = self._calculate_time_decay(memory.timestamp) priority *= time_decay return priority def _calculate_time_decay (self, timestamp: datetime ) -> float : """计算时间衰减因子""" time_diff = datetime.now() - timestamp hours_passed = time_diff.total_seconds() / 3600 decay_factor = self.config.decay_factor ** (hours_passed / 6 ) return max (0.1 , decay_factor) def _enforce_capacity_limits (self ): """强制执行容量限制""" while len (self.memories) > self.max_capacity: self._remove_lowest_priority_memory() while self.current_tokens > self.max_tokens: self._remove_lowest_priority_memory() def _expire_old_memories (self ): """按TTL清理过期记忆,并同步更新堆与token计数""" if not self.memories: return cutoff_time = datetime.now() - timedelta(minutes=self.max_age_minutes) kept: List [MemoryItem] = [] removed_token_sum = 0 for m in self.memories: if m.timestamp >= cutoff_time: kept.append(m) else : removed_token_sum += len (m.content.split()) if len (kept) == len (self.memories): return self.memories = kept self.current_tokens = max (0 , self.current_tokens - removed_token_sum) def _remove_lowest_priority_memory (self ): """删除优先级最低的记忆""" if not self.memories: return lowest_priority = float ('inf' ) lowest_memory = None for memory in self.memories: priority = self._calculate_priority(memory) if priority < lowest_priority: lowest_priority = priority lowest_memory = memory if lowest_memory: self.remove(lowest_memory.id )
episodicMemory.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 from typing import List , Dict , Any , Optional , Tuple from datetime import datetime, timedeltaimport osimport loggingfrom memory.memoryConfg import MemoryConfigfrom memory.memoryItem import MemoryItemfrom memory.memory import Memoryfrom storage.document_store import SQLiteDocumentStorefrom storage.qdrant_store import QdrantConnectionManagerfrom storage.embedding import get_text_embedder, get_dimensionlogger = logging.getLogger(__name__) class Episode : """情景记忆中的单个情景""" def __init__ ( self, episode_id: str , user_id: str , session_id: str , timestamp: datetime, content: str , context: Dict [str , Any ], outcome: Optional [str ] = None , importance: float = 0.5 ): self.episode_id = episode_id self.user_id = user_id self.session_id = session_id self.timestamp = timestamp self.content = content self.context = context self.outcome = outcome self.importance = importance class EpisodicMemory (Memory ): """情景记忆实现 特点: - 存储具体的交互事件 - 包含丰富的上下文信息 - 按时间序列组织 - 支持模式识别和回溯 """ def __init__ (self, config: MemoryConfig, storage_backend=None ): super ().__init__(config, storage_backend) self.episodes: List [Episode] = [] self.sessions: Dict [str , List [str ]] = {} self.patterns_cache = {} self.last_pattern_analysis = None db_dir = self.config.storage_path if hasattr (self.config, 'storage_path' ) else "./memory_data" os.makedirs(db_dir, exist_ok=True ) db_path = os.path.join(db_dir, "memory.db" ) self.doc_store = SQLiteDocumentStore(db_path=db_path) self.embedder = get_text_embedder() qdrant_url = os.getenv("QDRANT_URL" ) qdrant_api_key = os.getenv("QDRANT_API_KEY" ) self.vector_store = QdrantConnectionManager.get_instance( url=qdrant_url, api_key=qdrant_api_key, collection_name=os.getenv("QDRANT_COLLECTION" , "hello_agents_vectors" ), vector_size=get_dimension(getattr (self.embedder, 'dimension' , 384 )), distance=os.getenv("QDRANT_DISTANCE" , "cosine" ) ) def add (self, memory_item: MemoryItem ) -> str : """添加情景记忆""" session_id = memory_item.metadata.get("session_id" , "default_session" ) context = memory_item.metadata.get("context" , {}) outcome = memory_item.metadata.get("outcome" ) participants = memory_item.metadata.get("participants" , []) tags = memory_item.metadata.get("tags" , []) episode = Episode( episode_id=memory_item.id , user_id=memory_item.user_id, session_id=session_id, timestamp=memory_item.timestamp, content=memory_item.content, context=context, outcome=outcome, importance=memory_item.importance ) self.episodes.append(episode) if session_id not in self.sessions: self.sessions[session_id] = [] self.sessions[session_id].append(episode.episode_id) ts_int = int (memory_item.timestamp.timestamp()) self.doc_store.add_memory( memory_id=memory_item.id , user_id=memory_item.user_id, content=memory_item.content, memory_type="episodic" , timestamp=ts_int, importance=memory_item.importance, properties={ "session_id" : session_id, "context" : context, "outcome" : outcome, "participants" : participants, "tags" : tags } ) try : embedding = self.embedder.encode(memory_item.content) if hasattr (embedding, "tolist" ): embedding = embedding.tolist() self.vector_store.add_vectors( vectors=[embedding], metadata=[{ "memory_id" : memory_item.id , "user_id" : memory_item.user_id, "memory_type" : "episodic" , "importance" : memory_item.importance, "session_id" : session_id, "content" : memory_item.content }], ids=[memory_item.id ] ) except Exception: pass return memory_item.id def retrieve (self, query: str , limit: int = 5 , **kwargs ) -> List [MemoryItem]: """检索情景记忆(结构化过滤 + 语义向量检索)""" user_id = kwargs.get("user_id" ) session_id = kwargs.get("session_id" ) time_range: Optional [Tuple [datetime, datetime]] = kwargs.get("time_range" ) importance_threshold: Optional [float ] = kwargs.get("importance_threshold" ) candidate_ids: Optional [set ] = None if time_range is not None or importance_threshold is not None : start_ts = int (time_range[0 ].timestamp()) if time_range else None end_ts = int (time_range[1 ].timestamp()) if time_range else None docs = self.doc_store.search_memories( user_id=user_id, memory_type="episodic" , start_time=start_ts, end_time=end_ts, importance_threshold=importance_threshold, limit=1000 ) candidate_ids = {d["memory_id" ] for d in docs} try : query_vec = self.embedder.encode(query) if hasattr (query_vec, "tolist" ): query_vec = query_vec.tolist() where = {"memory_type" : "episodic" } if user_id: where["user_id" ] = user_id hits = self.vector_store.search_similar( query_vector=query_vec, limit=max (limit * 5 , 20 ), where=where ) except Exception: hits = [] now_ts = int (datetime.now().timestamp()) results: List [Tuple [float , MemoryItem]] = [] seen = set () for hit in hits: meta = hit.get("metadata" , {}) mem_id = meta.get("memory_id" ) if not mem_id or mem_id in seen: continue episode = next ((e for e in self.episodes if e.episode_id == mem_id), None ) if episode and episode.context.get("forgotten" , False ): continue if candidate_ids is not None and mem_id not in candidate_ids: continue if session_id and meta.get("session_id" ) != session_id: continue doc = self.doc_store.get_memory(mem_id) if not doc: continue vec_score = float (hit.get("score" , 0.0 )) age_days = max (0.0 , (now_ts - int (doc["timestamp" ])) / 86400.0 ) recency_score = 1.0 / (1.0 + age_days) imp = float (doc.get("importance" , 0.5 )) base_relevance = vec_score * 0.8 + recency_score * 0.2 importance_weight = 0.8 + (imp * 0.4 ) combined = base_relevance * importance_weight item = MemoryItem( id =doc["memory_id" ], content=doc["content" ], memory_type=doc["memory_type" ], user_id=doc["user_id" ], timestamp=datetime.fromtimestamp(doc["timestamp" ]), importance=doc.get("importance" , 0.5 ), metadata={ **doc.get("properties" , {}), "relevance_score" : combined, "vector_score" : vec_score, "recency_score" : recency_score } ) results.append((combined, item)) seen.add(mem_id) if not results: fallback = super ()._generate_id query_lower = query.lower() for ep in self._filter_episodes(user_id, session_id, time_range): if query_lower in ep.content.lower(): recency_score = 1.0 / (1.0 + max (0.0 , (now_ts - int (ep.timestamp.timestamp())) / 86400.0 )) keyword_score = 0.5 base_relevance = keyword_score * 0.8 + recency_score * 0.2 importance_weight = 0.8 + (ep.importance * 0.4 ) combined = base_relevance * importance_weight item = MemoryItem( id =ep.episode_id, content=ep.content, memory_type="episodic" , user_id=ep.user_id, timestamp=ep.timestamp, importance=ep.importance, metadata={ "session_id" : ep.session_id, "context" : ep.context, "outcome" : ep.outcome, "relevance_score" : combined } ) results.append((combined, item)) results.sort(key=lambda x: x[0 ], reverse=True ) return [it for _, it in results[:limit]] def update ( self, memory_id: str , content: str = None , importance: float = None , metadata: Dict [str , Any ] = None ) -> bool : """更新情景记忆(SQLite为权威,Qdrant按需重嵌入)""" updated = False for episode in self.episodes: if episode.episode_id == memory_id: if content is not None : episode.content = content if importance is not None : episode.importance = importance if metadata is not None : episode.context.update(metadata.get("context" , {})) if "outcome" in metadata: episode.outcome = metadata["outcome" ] updated = True break doc_updated = self.doc_store.update_memory( memory_id=memory_id, content=content, importance=importance, properties=metadata ) if content is not None : try : embedding = self.embedder.encode(content) if hasattr (embedding, "tolist" ): embedding = embedding.tolist() doc = self.doc_store.get_memory(memory_id) payload = { "memory_id" : memory_id, "user_id" : doc["user_id" ] if doc else "" , "memory_type" : "episodic" , "importance" : (doc.get("importance" ) if doc else importance) or 0.5 , "session_id" : (doc.get("properties" , {}) or {}).get("session_id" ), "content" : content } self.vector_store.add_vectors( vectors=[embedding], metadata=[payload], ids=[memory_id] ) except Exception: pass return updated or doc_updated def remove (self, memory_id: str ) -> bool : """删除情景记忆(SQLite + Qdrant)""" removed = False for i, episode in enumerate (self.episodes): if episode.episode_id == memory_id: removed_episode = self.episodes.pop(i) session_id = removed_episode.session_id if session_id in self.sessions: self.sessions[session_id].remove(memory_id) if not self.sessions[session_id]: del self.sessions[session_id] removed = True break doc_deleted = self.doc_store.delete_memory(memory_id) try : self.vector_store.delete_memories([memory_id]) except Exception: pass return removed or doc_deleted def has_memory (self, memory_id: str ) -> bool : """检查记忆是否存在""" return any (episode.episode_id == memory_id for episode in self.episodes) def clear (self ): """清空所有情景记忆(仅清理episodic,不影响其他类型)""" self.episodes.clear() self.sessions.clear() self.patterns_cache.clear() docs = self.doc_store.search_memories(memory_type="episodic" , limit=10000 ) ids = [d["memory_id" ] for d in docs] for mid in ids: self.doc_store.delete_memory(mid) try : if ids: self.vector_store.delete_memories(ids) except Exception: pass def forget (self, strategy: str = "importance_based" , threshold: float = 0.1 , max_age_days: int = 30 ) -> int : """情景记忆遗忘机制(硬删除)""" forgotten_count = 0 current_time = datetime.now() to_remove = [] for episode in self.episodes: should_forget = False if strategy == "importance_based" : if episode.importance < threshold: should_forget = True elif strategy == "time_based" : cutoff_time = current_time - timedelta(days=max_age_days) if episode.timestamp < cutoff_time: should_forget = True elif strategy == "capacity_based" : if len (self.episodes) > self.config.max_capacity: sorted_episodes = sorted (self.episodes, key=lambda e: e.importance) excess_count = len (self.episodes) - self.config.max_capacity if episode in sorted_episodes[:excess_count]: should_forget = True if should_forget: to_remove.append(episode.episode_id) for episode_id in to_remove: if self.remove(episode_id): forgotten_count += 1 logger.info(f"情景记忆硬删除: {episode_id[:8 ]} ... (策略: {strategy} )" ) return forgotten_count def get_all (self ) -> List [MemoryItem]: """获取所有情景记忆(转换为MemoryItem格式)""" memory_items = [] for episode in self.episodes: memory_item = MemoryItem( id =episode.episode_id, content=episode.content, memory_type="episodic" , user_id=episode.user_id, timestamp=episode.timestamp, importance=episode.importance, metadata=episode.metadata ) memory_items.append(memory_item) return memory_items def get_stats (self ) -> Dict [str , Any ]: """获取情景记忆统计信息(合并SQLite与Qdrant)""" active_episodes = self.episodes db_stats = self.doc_store.get_database_stats() try : vs_stats = self.vector_store.get_collection_stats() except Exception: vs_stats = {"store_type" : "qdrant" } return { "count" : len (active_episodes), "forgotten_count" : 0 , "total_count" : len (self.episodes), "sessions_count" : len (self.sessions), "avg_importance" : sum (e.importance for e in active_episodes) / len (active_episodes) if active_episodes else 0.0 , "time_span_days" : self._calculate_time_span(), "memory_type" : "episodic" , "vector_store" : vs_stats, "document_store" : {k: v for k, v in db_stats.items() if k.endswith("_count" ) or k in ["store_type" , "db_path" ]} } def get_session_episodes (self, session_id: str ) -> List [Episode]: """获取指定会话的所有情景""" if session_id not in self.sessions: return [] episode_ids = self.sessions[session_id] return [e for e in self.episodes if e.episode_id in episode_ids] def find_patterns (self, user_id: str = None , min_frequency: int = 2 ) -> List [Dict [str , Any ]]: """发现用户行为模式""" cache_key = f"{user_id} _{min_frequency} " if (cache_key in self.patterns_cache and self.last_pattern_analysis and (datetime.now() - self.last_pattern_analysis).hours < 1 ): return self.patterns_cache[cache_key] episodes = [e for e in self.episodes if user_id is None or e.user_id == user_id] keyword_patterns = {} context_patterns = {} for episode in episodes: words = episode.content.lower().split() for word in words: if len (word) > 3 : keyword_patterns[word] = keyword_patterns.get(word, 0 ) + 1 for key, value in episode.context.items(): pattern_key = f"{key} :{value} " context_patterns[pattern_key] = context_patterns.get(pattern_key, 0 ) + 1 patterns = [] for keyword, frequency in keyword_patterns.items(): if frequency >= min_frequency: patterns.append({ "type" : "keyword" , "pattern" : keyword, "frequency" : frequency, "confidence" : frequency / len (episodes) }) for context_pattern, frequency in context_patterns.items(): if frequency >= min_frequency: patterns.append({ "type" : "context" , "pattern" : context_pattern, "frequency" : frequency, "confidence" : frequency / len (episodes) }) patterns.sort(key=lambda x: x["frequency" ], reverse=True ) self.patterns_cache[cache_key] = patterns self.last_pattern_analysis = datetime.now() return patterns def get_timeline (self, user_id: str = None , limit: int = 50 ) -> List [Dict [str , Any ]]: """获取时间线视图""" episodes = [e for e in self.episodes if user_id is None or e.user_id == user_id] episodes.sort(key=lambda x: x.timestamp, reverse=True ) timeline = [] for episode in episodes[:limit]: timeline.append({ "episode_id" : episode.episode_id, "timestamp" : episode.timestamp.isoformat(), "content" : episode.content[:100 ] + "..." if len (episode.content) > 100 else episode.content, "session_id" : episode.session_id, "importance" : episode.importance, "outcome" : episode.outcome }) return timeline def _filter_episodes ( self, user_id: str = None , session_id: str = None , time_range: Tuple [datetime, datetime] = None ) -> List [Episode]: """过滤情景""" filtered = self.episodes if user_id: filtered = [e for e in filtered if e.user_id == user_id] if session_id: filtered = [e for e in filtered if e.session_id == session_id] if time_range: start_time, end_time = time_range filtered = [e for e in filtered if start_time <= e.timestamp <= end_time] return filtered def _calculate_time_span (self ) -> float : """计算记忆时间跨度(天)""" if not self.episodes: return 0.0 timestamps = [e.timestamp for e in self.episodes] min_time = min (timestamps) max_time = max (timestamps) return (max_time - min_time).days def _persist_episode (self, episode: Episode ): """持久化情景到存储后端""" if self.storage and hasattr (self.storage, 'add_memory' ): self.storage.add_memory( memory_id=episode.episode_id, user_id=episode.user_id, content=episode.content, memory_type="episodic" , timestamp=int (episode.timestamp.timestamp()), importance=episode.importance, properties={ "session_id" : episode.session_id, "context" : episode.context, "outcome" : episode.outcome } ) def _remove_from_storage (self, memory_id: str ): """从存储后端删除""" if self.storage and hasattr (self.storage, 'delete_memory' ): self.storage.delete_memory(memory_id)
semanticMemory.py
memoryManager.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 from memory.memoryConfg import MemoryConfigfrom memory.memoryItem import MemoryItemfrom memory.types.workingMemory import WorkingMemoryfrom memory.types.episodicMemory import EpisodicMemoryfrom memory.types.semanticMemory import SemanticMemoryfrom typing import List , Dict , Any , Optional , Union from datetime import datetimeimport uuidimport logginglogger = logging.getLogger(__name__) class MemoryManager : """记忆管理器 - 统一的记忆操作接口 负责: - 记忆生命周期管理 - 记忆优先级和重要性评估 - 记忆遗忘和清理机制 - 多类型记忆的协调管理 """ def __init__ ( self, config: Optional [MemoryConfig] = None , user_id: str = "default_user" , enable_working: bool = True , enable_episodic: bool = True , enable_semantic: bool = True ): self.config = config or MemoryConfig() self.user_id = user_id self.memory_types = {} if enable_working: self.memory_types['working' ] = WorkingMemory(self.config) if enable_episodic: self.memory_types['episodic' ] = EpisodicMemory(self.config) if enable_semantic: self.memory_types['semantic' ] = SemanticMemory(self.config) logger.info(f"MemoryManager初始化完成,启用记忆类型: {list (self.memory_types.keys())} " ) def _classify_memory_type (self, content: str , metadata: Optional [Dict [str , Any ]] ) -> str : """自动分类记忆类型""" if metadata and metadata.get("type" ): return metadata["type" ] if self._is_episodic_content(content): return "episodic" elif self._is_semantic_content(content): return "semantic" else : return "working" def _is_episodic_content (self, content: str ) -> bool : """判断是否为情景记忆内容""" episodic_keywords = ["昨天" , "今天" , "明天" , "上次" , "记得" , "发生" , "经历" ] return any (keyword in content for keyword in episodic_keywords) def _is_semantic_content (self, content: str ) -> bool : """判断是否为语义记忆内容""" semantic_keywords = ["定义" , "概念" , "规则" , "知识" , "原理" , "方法" ] return any (keyword in content for keyword in semantic_keywords) def _calculate_importance (self, content: str , metadata: Optional [Dict [str , Any ]] ) -> float : """计算记忆重要性""" importance = 0.5 if len (content) > 100 : importance += 0.1 important_keywords = ["重要" , "关键" , "必须" , "注意" , "警告" , "错误" ] if any (keyword in content for keyword in important_keywords): importance += 0.2 if metadata: if metadata.get("priority" ) == "high" : importance += 0.3 elif metadata.get("priority" ) == "low" : importance -= 0.2 return max (0.0 , min (1.0 , importance)) def __str__ (self ) -> str : stats = self.get_memory_stats() return f"MemoryManager(user={self.user_id} , total={stats['total_memories' ]} )" def add_memory ( self, content: str , memory_type: str = "working" , importance: Optional [float ] = None , metadata: Optional [Dict [str , Any ]] = None , auto_classify: bool = True ) -> str : """添加记忆 Args: content: 记忆内容 memory_type: 记忆类型 importance: 重要性分数 (0-1) metadata: 元数据 auto_classify: 是否自动分类到合适的记忆类型 Returns: 记忆ID """ if auto_classify: memory_type = self._classify_memory_type(content, metadata) if importance is None : importance = self._calculate_importance(content, metadata) memory_item = MemoryItem( id =str (uuid.uuid4()), content=content, memory_type=memory_type, user_id=self.user_id, timestamp=datetime.now(), importance=importance, metadata=metadata or {} ) if memory_type in self.memory_types: memory_id = self.memory_types[memory_type].add(memory_item) logger.debug(f"添加记忆到 {memory_type} : {memory_id} " ) return memory_id else : raise ValueError(f"不支持的记忆类型: {memory_type} " ) def retrieve_memories ( self, query: str , memory_types: Optional [List [str ]] = None , limit: int = 10 , min_importance: float = 0.0 , time_range: Optional [tuple ] = None ) -> List [MemoryItem]: """检索记忆 Args: query: 查询内容 memory_types: 要检索的记忆类型列表 limit: 返回数量限制 min_importance: 最小重要性阈值 time_range: 时间范围 (start_time, end_time) Returns: 检索到的记忆列表 """ if memory_types is None : memory_types = list (self.memory_types.keys()) all_results = [] per_type_limit = max (1 , limit // len (memory_types)) for memory_type in memory_types: if memory_type in self.memory_types: memory_instance = self.memory_types[memory_type] try : type_results = memory_instance.retrieve( query=query, limit=per_type_limit, min_importance=min_importance, user_id=self.user_id ) all_results.extend(type_results) except Exception as e: logger.warning(f"检索 {memory_type} 记忆时出错: {e} " ) continue all_results.sort(key=lambda x: x.importance, reverse=True ) return all_results[:limit] def update_memory ( self, memory_id: str , content: Optional [str ] = None , importance: Optional [float ] = None , metadata: Optional [Dict [str , Any ]] = None ) -> bool : """更新记忆 Args: memory_id: 记忆ID content: 新内容 importance: 新重要性 metadata: 新元数据 Returns: 是否更新成功 """ for memory_type, memory_instance in self.memory_types.items(): if memory_instance.has_memory(memory_id): return memory_instance.update(memory_id, content, importance, metadata) logger.warning(f"未找到记忆: {memory_id} " ) return False def remove_memory (self, memory_id: str ) -> bool : """删除记忆 Args: memory_id: 记忆ID Returns: 是否删除成功 """ for memory_type, memory_instance in self.memory_types.items(): if memory_instance.has_memory(memory_id): return memory_instance.remove(memory_id) logger.warning(f"未找到记忆: {memory_id} " ) return False def forget_memories ( self, strategy: str = "importance_based" , threshold: float = 0.1 , max_age_days: int = 30 ) -> int : """记忆遗忘机制 Args: strategy: 遗忘策略 ("importance_based", "time_based", "capacity_based") threshold: 遗忘阈值 max_age_days: 最大保存天数 Returns: 遗忘的记忆数量 """ total_forgotten = 0 for memory_type, memory_instance in self.memory_types.items(): if hasattr (memory_instance, 'forget' ): forgotten = memory_instance.forget(strategy, threshold, max_age_days) total_forgotten += forgotten logger.info(f"记忆遗忘完成: {total_forgotten} 条记忆" ) return total_forgotten def consolidate_memories ( self, from_type: str = "working" , to_type: str = "episodic" , importance_threshold: float = 0.7 ) -> int : """记忆整合 - 将重要的短期记忆转换为长期记忆 Args: from_type: 源记忆类型 to_type: 目标记忆类型 importance_threshold: 重要性阈值 Returns: 整合的记忆数量 """ if from_type not in self.memory_types or to_type not in self.memory_types: logger.warning(f"记忆类型不存在: {from_type} -> {to_type} " ) return 0 source_memory = self.memory_types[from_type] target_memory = self.memory_types[to_type] all_memories = source_memory.get_all() candidates = [ m for m in all_memories if m.importance >= importance_threshold ] consolidated_count = 0 for memory in candidates: if source_memory.remove(memory.id ): memory.memory_type = to_type memory.importance *= 1.1 target_memory.add(memory) consolidated_count += 1 logger.info(f"记忆整合完成: {consolidated_count} 条记忆从 {from_type} 转移到 {to_type} " ) return consolidated_count def get_memory_stats (self ) -> Dict [str , Any ]: """获取记忆统计信息""" stats = { "user_id" : self.user_id, "enabled_types" : list (self.memory_types.keys()), "total_memories" : 0 , "memories_by_type" : {}, "config" : { "max_capacity" : self.config.max_capacity, "importance_threshold" : self.config.importance_threshold, "decay_factor" : self.config.decay_factor } } for memory_type, memory_instance in self.memory_types.items(): type_stats = memory_instance.get_stats() stats["memories_by_type" ][memory_type] = type_stats stats["total_memories" ] += type_stats.get("count" , 0 ) return stats def clear_all_memories (self ): """清空所有记忆""" for memory_type, memory_instance in self.memory_types.items(): memory_instance.clear() logger.info("所有记忆已清空" )
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 from typing import Callable def tool_action (name: str = None , description: str = None ): """装饰器:标记一个方法为可展开的工具 action 用法: @tool_action("memory_add", "添加新记忆") def _add_memory(self, content: str, importance: float = 0.5) -> str: '''添加记忆 Args: content: 记忆内容 importance: 重要性分数 ''' ... Args: name: 工具名称(如果不提供,从方法名自动生成) description: 工具描述(如果不提供,从 docstring 提取) """ def decorator (func: Callable ): func._is_tool_action = True func._tool_name = name func._tool_description = description return func return decorator
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 ### builtin #### memoryTool.py ```py """ MemoryTool - 记忆系统的统一工具接口 设计原则:统一入口,分发处理(Unified Entry, Dispatch to Handlers) 该工具将底层的 MemoryManager 封装为 LLM 可调用的 Tool, 使 Agent 具备持久化记忆和检索能力。 """ from typing import List, Dict, Any, Optional import json from ..tool import Tool from ..toolParameter import ToolParameter from memory.memoryManager import MemoryManager from memory.memoryConfg import MemoryConfig class MemoryTool(Tool): """记忆工具 —— 让 LLM 能够存取和管理记忆 作为记忆系统的统一接口,遵循"统一入口,分发处理"的架构模式。 MemoryTool 专注于参数解析和结果格式化,核心逻辑委托给 MemoryManager。 """ def __init__( self, user_id: str = "default_user", config: Optional[MemoryConfig] = None ): super().__init__( name="memory", description=( "记忆管理工具,用于存储、检索、更新和删除记忆。" "支持以下操作:\n" "- add: 添加一条新记忆。需要 content(记忆内容),可选 memory_type " "(working/episodic/semantic,默认自动推断)、importance(0-1,默认自动计算)\n" "- search: 按关键词检索记忆。需要 query(查询内容),可选 limit(返回数量,默认5)\n" "- update: 更新已有记忆。需要 memory_id,可选 content/importance\n" "- remove: 删除指定记忆。需要 memory_id\n" "- summary: 查看记忆系统的统计摘要(总数、各类型容量等)\n" "- forget: 触发记忆遗忘机制。可选 strategy(importance_based/time_based/capacity_based)、" "threshold(遗忘阈值,默认0.1)\n" "- consolidate: 将重要的短期记忆整合为长期记忆。" "可选 from_type(默认working)、to_type(默认episodic)、importance_threshold(默认0.7)\n\n" "记忆类型说明:\n" "- working: 工作记忆(当前会话上下文,容量有限,过期自动清理)\n" "- episodic: 情景记忆(重要交互记录,长期保留)\n" "- semantic: 语义记忆(抽象规则和用户偏好,永久保留)\n\n" "使用建议:\n" "- 用户说'记住XXX'时,根据内容类型选择合适的 memory_type\n" "- 用户偏好/规范类信息 → semantic,重要性 0.8+\n" "- 临时上下文/当前话题 → working,重要性 0.5\n" "- 重要事件/决定 → episodic,重要性 0.7+\n" "- 在回答用户问题前,先用 search 检索相关记忆" ) ) self.user_id = user_id self._manager = MemoryManager( config=config, user_id=user_id, enable_working=True, enable_episodic=True, enable_semantic=True ) # ------------------------------------------------------------------ # Tool 接口实现 # ------------------------------------------------------------------ def run(self, parameters: Dict[str, Any]) -> str: """统一入口:根据 action 分发给具体处理方法 Args: parameters: 包含 action 及其他所需参数的字典 Returns: 格式化后的结果字符串 """ action = parameters.get("action", "search") dispatch = { "add": self._handle_add, "search": self._handle_search, "update": self._handle_update, "remove": self._handle_remove, "summary": self._handle_summary, "forget": self._handle_forget, "consolidate": self._handle_consolidate, } handler = dispatch.get(action) if handler is None: return ( f"❌ 不支持的操作: '{action}'。" f"可用操作: {', '.join(dispatch.keys())}" ) try: return handler(parameters) except ValueError as e: return f"❌ 参数错误: {e}" except Exception as e: return f"❌ 操作失败: {e}" def get_parameters(self) -> List[ToolParameter]: """获取工具参数定义(用于 OpenAI function calling schema)""" return [ ToolParameter( name="action", type="string", description="操作类型: add, search, update, remove, summary, forget, consolidate", required=True ), ToolParameter( name="content", type="string", description="记忆内容文本(add/update 时使用)", required=False ), ToolParameter( name="query", type="string", description="检索查询关键词(search 时使用)", required=False ), ToolParameter( name="memory_type", type="string", description="记忆类型: working, episodic, semantic(add 时可选,默认自动推断)", required=False, default="working" ), ToolParameter( name="memory_id", type="string", description="记忆ID(update/remove 时使用)", required=False ), ToolParameter( name="importance", type="number", description="重要性分数 0-1(add/update 时可选,默认自动计算)", required=False ), ToolParameter( name="limit", type="integer", description="返回数量限制(search 时使用,默认5)", required=False, default=5 ), ToolParameter( name="strategy", type="string", description="遗忘策略: importance_based, time_based, capacity_based(forget 时使用,默认 importance_based)", required=False, default="importance_based" ), ToolParameter( name="threshold", type="number", description="遗忘阈值(forget 时使用,默认0.1)", required=False, default=0.1 ), ] # ------------------------------------------------------------------ # Action 处理器 # ------------------------------------------------------------------ def _handle_add(self, params: Dict[str, Any]) -> str: """处理 add 操作""" content = params.get("content") if not content: raise ValueError("add 操作需要 'content' 参数") memory_type = params.get("memory_type", "working") importance = params.get("importance") # 将 memory_type 注入 metadata,使 MemoryManager 优先采纳 LLM 的判断 metadata = params.get("metadata", {}) metadata["type"] = memory_type memory_id = self._manager.add_memory( content=content, memory_type=memory_type, importance=importance, metadata=metadata, auto_classify=True ) return ( f"✅ 记忆已添加\n" f" ID: {memory_id}\n" f" 类型: {memory_type}\n" f" 内容: {content[:200]}{'...' if len(content) > 200 else ''}" ) def _handle_search(self, params: Dict[str, Any]) -> str: """处理 search 操作""" query = params.get("query") if not query: raise ValueError("search 操作需要 'query' 参数") limit = params.get("limit", 5) memory_types = params.get("memory_types") # 可选:限定检索的记忆类型列表 results = self._manager.retrieve_memories( query=query, memory_types=memory_types, limit=limit ) if not results: return f"🔍 未找到与 '{query}' 相关的记忆。" lines = [f"🔍 搜索 '{query}' 的结果(共 {len(results)} 条):"] for i, item in enumerate(results, 1): lines.append( f" [{i}] [{item.memory_type}] {item.content[:150]}" f"{'...' if len(item.content) > 150 else ''}" f" (importance={item.importance:.2f}, id={item.id[:8]}...)" ) return "\n".join(lines) def _handle_update(self, params: Dict[str, Any]) -> str: """处理 update 操作""" memory_id = params.get("memory_id") if not memory_id: raise ValueError("update 操作需要 'memory_id' 参数") content = params.get("content") importance = params.get("importance") metadata = params.get("metadata") success = self._manager.update_memory( memory_id=memory_id, content=content, importance=importance, metadata=metadata ) if success: return f"✅ 记忆 {memory_id} 已更新" else: return f"❌ 未找到记忆: {memory_id}" def _handle_remove(self, params: Dict[str, Any]) -> str: """处理 remove 操作""" memory_id = params.get("memory_id") if not memory_id: raise ValueError("remove 操作需要 'memory_id' 参数") success = self._manager.remove_memory(memory_id) if success: return f"✅ 记忆 {memory_id} 已删除" else: return f"❌ 未找到记忆: {memory_id}" def _handle_summary(self, params: Dict[str, Any]) -> str: """处理 summary 操作""" stats = self._manager.get_memory_stats() lines = ["📊 记忆系统摘要:"] lines.append(f" 用户: {stats['user_id']}") lines.append(f" 总记忆数: {stats['total_memories']}") lines.append(f" 启用的记忆类型: {', '.join(stats['enabled_types'])}") lines.append(f" 配置: 最大容量={stats['config']['max_capacity']}, " f"重要性阈值={stats['config']['importance_threshold']}, " f"衰减因子={stats['config']['decay_factor']}") for mem_type, type_stats in stats.get("memories_by_type", {}).items(): lines.append(f" [{mem_type}] 活跃={type_stats.get('count', 0)}, " f"容量使用率={type_stats.get('capacity_usage', 0):.1%}") return "\n".join(lines) def _handle_forget(self, params: Dict[str, Any]) -> str: """处理 forget 操作""" strategy = params.get("strategy", "importance_based") threshold = params.get("threshold", 0.1) max_age_days = params.get("max_age_days", 30) forgotten = self._manager.forget_memories( strategy=strategy, threshold=threshold, max_age_days=max_age_days ) return f"🧹 记忆遗忘完成: {forgotten} 条记忆被清理(策略: {strategy})" def _handle_consolidate(self, params: Dict[str, Any]) -> str: """处理 consolidate 操作""" from_type = params.get("from_type", "working") to_type = params.get("to_type", "episodic") importance_threshold = params.get("importance_threshold", 0.7) # 检查目标记忆类型是否已启用 if to_type not in self._manager.memory_types: return ( f"⚠️ 目标记忆类型 '{to_type}' 尚未实现。" f"当前只启用了工作记忆(working),情景记忆和语义记忆有待开发。" ) count = self._manager.consolidate_memories( from_type=from_type, to_type=to_type, importance_threshold=importance_threshold ) return ( f"📦 记忆整合完成: {count} 条记忆从 {from_type} 转移到 {to_type} " f"(重要性阈值: {importance_threshold})" ) # ------------------------------------------------------------------ # 便捷方法(供程序化调用,非 LLM 路径) # ------------------------------------------------------------------ @property def manager(self) -> MemoryManager: """暴露底层 MemoryManager,便于程序化访问""" return self._manager
builtin/ 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 """ MemoryTool - 记忆系统的统一工具接口 设计原则:统一入口,分发处理(Unified Entry, Dispatch to Handlers) 该工具将底层的 MemoryManager 封装为 LLM 可调用的 Tool, 使 Agent 具备持久化记忆和检索能力。 """ from typing import List , Dict , Any , Optional import jsonfrom ..tool import Toolfrom ..toolParameter import ToolParameterfrom memory.memoryManager import MemoryManagerfrom memory.memoryConfg import MemoryConfigclass MemoryTool (Tool ): """记忆工具 —— 让 LLM 能够存取和管理记忆 作为记忆系统的统一接口,遵循"统一入口,分发处理"的架构模式。 MemoryTool 专注于参数解析和结果格式化,核心逻辑委托给 MemoryManager。 """ def __init__ ( self, user_id: str = "default_user" , config: Optional [MemoryConfig] = None ): super ().__init__( name="memory" , description=( "记忆管理工具,用于存储、检索、更新和删除记忆。" "支持以下操作:\n" "- add: 添加一条新记忆。需要 content(记忆内容),可选 memory_type " "(working/episodic/semantic,默认自动推断)、importance(0-1,默认自动计算)\n" "- search: 按关键词检索记忆。需要 query(查询内容),可选 limit(返回数量,默认5)\n" "- update: 更新已有记忆。需要 memory_id,可选 content/importance\n" "- remove: 删除指定记忆。需要 memory_id\n" "- summary: 查看记忆系统的统计摘要(总数、各类型容量等)\n" "- forget: 触发记忆遗忘机制。可选 strategy(importance_based/time_based/capacity_based)、" "threshold(遗忘阈值,默认0.1)\n" "- consolidate: 将重要的短期记忆整合为长期记忆。" "可选 from_type(默认working)、to_type(默认episodic)、importance_threshold(默认0.7)\n\n" "记忆类型说明:\n" "- working: 工作记忆(当前会话上下文,容量有限,过期自动清理)\n" "- episodic: 情景记忆(重要交互记录,长期保留)\n" "- semantic: 语义记忆(抽象规则和用户偏好,永久保留)\n\n" "使用建议:\n" "- 用户说'记住XXX'时,根据内容类型选择合适的 memory_type\n" "- 用户偏好/规范类信息 → semantic,重要性 0.8+\n" "- 临时上下文/当前话题 → working,重要性 0.5\n" "- 重要事件/决定 → episodic,重要性 0.7+\n" "- 在回答用户问题前,先用 search 检索相关记忆" ) ) self.user_id = user_id self._manager = MemoryManager( config=config, user_id=user_id, enable_working=True , enable_episodic=True , enable_semantic=True ) def run (self, parameters: Dict [str , Any ] ) -> str : """统一入口:根据 action 分发给具体处理方法 Args: parameters: 包含 action 及其他所需参数的字典 Returns: 格式化后的结果字符串 """ action = parameters.get("action" , "search" ) dispatch = { "add" : self._handle_add, "search" : self._handle_search, "update" : self._handle_update, "remove" : self._handle_remove, "summary" : self._handle_summary, "forget" : self._handle_forget, "consolidate" : self._handle_consolidate, } handler = dispatch.get(action) if handler is None : return ( f"❌ 不支持的操作: '{action} '。" f"可用操作: {', ' .join(dispatch.keys())} " ) try : return handler(parameters) except ValueError as e: return f"❌ 参数错误: {e} " except Exception as e: return f"❌ 操作失败: {e} " def get_parameters (self ) -> List [ToolParameter]: """获取工具参数定义(用于 OpenAI function calling schema)""" return [ ToolParameter( name="action" , type ="string" , description="操作类型: add, search, update, remove, summary, forget, consolidate" , required=True ), ToolParameter( name="content" , type ="string" , description="记忆内容文本(add/update 时使用)" , required=False ), ToolParameter( name="query" , type ="string" , description="检索查询关键词(search 时使用)" , required=False ), ToolParameter( name="memory_type" , type ="string" , description="记忆类型: working, episodic, semantic(add 时可选,默认自动推断)" , required=False , default="working" ), ToolParameter( name="memory_id" , type ="string" , description="记忆ID(update/remove 时使用)" , required=False ), ToolParameter( name="importance" , type ="number" , description="重要性分数 0-1(add/update 时可选,默认自动计算)" , required=False ), ToolParameter( name="limit" , type ="integer" , description="返回数量限制(search 时使用,默认5)" , required=False , default=5 ), ToolParameter( name="strategy" , type ="string" , description="遗忘策略: importance_based, time_based, capacity_based(forget 时使用,默认 importance_based)" , required=False , default="importance_based" ), ToolParameter( name="threshold" , type ="number" , description="遗忘阈值(forget 时使用,默认0.1)" , required=False , default=0.1 ), ] def _handle_add (self, params: Dict [str , Any ] ) -> str : """处理 add 操作""" content = params.get("content" ) if not content: raise ValueError("add 操作需要 'content' 参数" ) memory_type = params.get("memory_type" , "working" ) importance = params.get("importance" ) metadata = params.get("metadata" , {}) metadata["type" ] = memory_type memory_id = self._manager.add_memory( content=content, memory_type=memory_type, importance=importance, metadata=metadata, auto_classify=True ) return ( f"✅ 记忆已添加\n" f" ID: {memory_id} \n" f" 类型: {memory_type} \n" f" 内容: {content[:200 ]} {'...' if len (content) > 200 else '' } " ) def _handle_search (self, params: Dict [str , Any ] ) -> str : """处理 search 操作""" query = params.get("query" ) if not query: raise ValueError("search 操作需要 'query' 参数" ) limit = params.get("limit" , 5 ) memory_types = params.get("memory_types" ) results = self._manager.retrieve_memories( query=query, memory_types=memory_types, limit=limit ) if not results: return f"🔍 未找到与 '{query} ' 相关的记忆。" lines = [f"🔍 搜索 '{query} ' 的结果(共 {len (results)} 条):" ] for i, item in enumerate (results, 1 ): lines.append( f" [{i} ] [{item.memory_type} ] {item.content[:150 ]} " f"{'...' if len (item.content) > 150 else '' } " f" (importance={item.importance:.2 f} , id={item.id [:8 ]} ...)" ) return "\n" .join(lines) def _handle_update (self, params: Dict [str , Any ] ) -> str : """处理 update 操作""" memory_id = params.get("memory_id" ) if not memory_id: raise ValueError("update 操作需要 'memory_id' 参数" ) content = params.get("content" ) importance = params.get("importance" ) metadata = params.get("metadata" ) success = self._manager.update_memory( memory_id=memory_id, content=content, importance=importance, metadata=metadata ) if success: return f"✅ 记忆 {memory_id} 已更新" else : return f"❌ 未找到记忆: {memory_id} " def _handle_remove (self, params: Dict [str , Any ] ) -> str : """处理 remove 操作""" memory_id = params.get("memory_id" ) if not memory_id: raise ValueError("remove 操作需要 'memory_id' 参数" ) success = self._manager.remove_memory(memory_id) if success: return f"✅ 记忆 {memory_id} 已删除" else : return f"❌ 未找到记忆: {memory_id} " def _handle_summary (self, params: Dict [str , Any ] ) -> str : """处理 summary 操作""" stats = self._manager.get_memory_stats() lines = ["📊 记忆系统摘要:" ] lines.append(f" 用户: {stats['user_id' ]} " ) lines.append(f" 总记忆数: {stats['total_memories' ]} " ) lines.append(f" 启用的记忆类型: {', ' .join(stats['enabled_types' ])} " ) lines.append(f" 配置: 最大容量={stats['config' ]['max_capacity' ]} , " f"重要性阈值={stats['config' ]['importance_threshold' ]} , " f"衰减因子={stats['config' ]['decay_factor' ]} " ) for mem_type, type_stats in stats.get("memories_by_type" , {}).items(): lines.append(f" [{mem_type} ] 活跃={type_stats.get('count' , 0 )} , " f"容量使用率={type_stats.get('capacity_usage' , 0 ):.1 %} " ) return "\n" .join(lines) def _handle_forget (self, params: Dict [str , Any ] ) -> str : """处理 forget 操作""" strategy = params.get("strategy" , "importance_based" ) threshold = params.get("threshold" , 0.1 ) max_age_days = params.get("max_age_days" , 30 ) forgotten = self._manager.forget_memories( strategy=strategy, threshold=threshold, max_age_days=max_age_days ) return f"🧹 记忆遗忘完成: {forgotten} 条记忆被清理(策略: {strategy} )" def _handle_consolidate (self, params: Dict [str , Any ] ) -> str : """处理 consolidate 操作""" from_type = params.get("from_type" , "working" ) to_type = params.get("to_type" , "episodic" ) importance_threshold = params.get("importance_threshold" , 0.7 ) if to_type not in self._manager.memory_types: return ( f"⚠️ 目标记忆类型 '{to_type} ' 尚未实现。" f"当前只启用了工作记忆(working),情景记忆和语义记忆有待开发。" ) count = self._manager.consolidate_memories( from_type=from_type, to_type=to_type, importance_threshold=importance_threshold ) return ( f"📦 记忆整合完成: {count} 条记忆从 {from_type} 转移到 {to_type} " f"(重要性阈值: {importance_threshold} )" ) @property def manager (self ) -> MemoryManager: """暴露底层 MemoryManager,便于程序化访问""" return self._manager
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 from typing import Dict , Any , List , Optional import osimport timefrom tool.tool import Toolfrom tool.toolAction import tool_actionfrom tool.toolParameter import ToolParameterfrom rag.pipeline import create_rag_pipelinefrom core.llm import LLM"""RAG工具 - 检索增强生成 为HelloAgents框架提供简洁易用的RAG能力: - 🔄 数据流程:用户数据 → 文档解析 → 向量化存储 → 智能检索 → LLM增强问答 - 📚 多格式支持:PDF、Word、Excel、PPT、图片、音频、网页等 - 🧠 智能问答:自动检索相关内容,注入提示词,生成准确答案 - 🏷️ 命名空间:支持多项目隔离,便于管理不同知识库 使用示例: ```python # 1. 初始化RAG工具 rag = RAGTool() # 2. 添加文档 rag.run({"action": "add_document", "file_path": "document.pdf"}) # 3. 智能问答 answer = rag.run({"action": "ask", "question": "什么是机器学习?"})
“””
class RAGTool(Tool): “””RAG工具
提供完整的 RAG 能力:
- 添加多格式文档(PDF、Office、图片、音频等)
- 智能检索与召回
- LLM 增强问答
- 知识库管理
"""
def __init__(
self,
knowledge_base_path: str = "./knowledge_base",
qdrant_url: str = None,
qdrant_api_key: str = None,
collection_name: str = "rag_knowledge_base",
rag_namespace: str = "default",
expandable: bool = False
):
super().__init__(
name="rag",
description="RAG工具 - 支持多格式文档检索增强生成,提供智能问答能力",
expandable=expandable
)
self.knowledge_base_path = knowledge_base_path
self.qdrant_url = qdrant_url or os.getenv("QDRANT_URL")
self.qdrant_api_key = qdrant_api_key or os.getenv("QDRANT_API_KEY")
self.collection_name = collection_name
self.rag_namespace = rag_namespace
self._pipelines: Dict[str, Dict[str, Any]] = {}
# 确保知识库目录存在
os.makedirs(knowledge_base_path, exist_ok=True)
# 初始化组件
self._init_components()
def _init_components(self):
"""初始化RAG组件"""
try:
# 初始化默认命名空间的 RAG 管道
default_pipeline = create_rag_pipeline(
qdrant_url=self.qdrant_url,
qdrant_api_key=self.qdrant_api_key,
collection_name=self.collection_name,
rag_namespace=self.rag_namespace
)
self._pipelines[self.rag_namespace] = default_pipeline
# 初始化 LLM 用于回答生成
self.llm = LLM()
self.initialized = True
print(f"✅ RAG工具初始化成功: namespace={self.rag_namespace}, collection={self.collection_name}")
except Exception as e:
self.initialized = False
self.init_error = str(e)
print(f"❌ RAG工具初始化失败: {e}")
def _get_pipeline(self, namespace: Optional[str] = None) -> Dict[str, Any]:
"""获取指定命名空间的 RAG 管道,若不存在则自动创建"""
target_ns = namespace or self.rag_namespace
if target_ns in self._pipelines:
return self._pipelines[target_ns]
pipeline = create_rag_pipeline(
qdrant_url=self.qdrant_url,
qdrant_api_key=self.qdrant_api_key,
collection_name=self.collection_name,
rag_namespace=target_ns
)
self._pipelines[target_ns] = pipeline
return pipeline
def run(self, parameters: Dict[str, Any]) -> str:
"""执行工具(非展开模式)
Args:
parameters: 工具参数字典,必须包含action参数
Returns:
执行结果字符串
"""
if not self.validate_parameters(parameters):
return "❌ 参数验证失败:缺少必需的参数"
if not self.initialized:
return f"❌ RAG工具未正确初始化,请检查配置: {getattr(self, 'init_error', '未知错误')}"
action = parameters.get("action")
# 根据action调用对应的方法,传入提取的参数
try:
if action == "add_document":
return self._add_document(
file_path=parameters.get("file_path"),
document_id=parameters.get("document_id"),
namespace=parameters.get("namespace", "default"),
chunk_size=parameters.get("chunk_size", 800),
chunk_overlap=parameters.get("chunk_overlap", 100)
)
elif action == "add_text":
return self._add_text(
text=parameters.get("text"),
document_id=parameters.get("document_id"),
namespace=parameters.get("namespace", "default"),
chunk_size=parameters.get("chunk_size", 800),
chunk_overlap=parameters.get("chunk_overlap", 100)
)
elif action == "ask":
question = parameters.get("question") or parameters.get("query")
return self._ask(
question=question,
limit=parameters.get("limit", 5),
enable_advanced_search=parameters.get("enable_advanced_search", True),
include_citations=parameters.get("include_citations", True),
max_chars=parameters.get("max_chars", 1200),
namespace=parameters.get("namespace", "default")
)
elif action == "search":
return self._search(
query=parameters.get("query") or parameters.get("question"),
limit=parameters.get("limit", 5),
min_score=parameters.get("min_score", 0.1),
enable_advanced_search=parameters.get("enable_advanced_search", True),
max_chars=parameters.get("max_chars", 1200),
include_citations=parameters.get("include_citations", True),
namespace=parameters.get("namespace", "default")
)
elif action == "stats":
return self._get_stats(namespace=parameters.get("namespace", "default"))
elif action == "clear":
return self._clear_knowledge_base(
confirm=parameters.get("confirm", False),
namespace=parameters.get("namespace", "default")
)
else:
return f"❌ 不支持的操作: {action}"
except Exception as e:
return f"❌ 执行操作 '{action}' 时发生错误: {str(e)}"
def get_parameters(self) -> List[ToolParameter]:
"""获取工具参数定义 - Tool基类要求的接口"""
return [
# 核心操作参数
ToolParameter(
name="action",
type="string",
description="操作类型:add_document(添加文档), add_text(添加文本), ask(智能问答), search(搜索), stats(统计), clear(清空)",
required=True
),
# 内容参数
ToolParameter(
name="file_path",
type="string",
description="文档文件路径(支持PDF、Word、Excel、PPT、图片、音频等多种格式)",
required=False
),
ToolParameter(
name="text",
type="string",
description="要添加的文本内容",
required=False
),
ToolParameter(
name="question",
type="string",
description="用户问题(用于智能问答)",
required=False
),
ToolParameter(
name="query",
type="string",
description="搜索查询词(用于基础搜索)",
required=False
),
# 可选配置参数
ToolParameter(
name="namespace",
type="string",
description="知识库命名空间(用于隔离不同项目,默认:default)",
required=False,
default="default"
),
ToolParameter(
name="limit",
type="integer",
description="返回结果数量(默认:5)",
required=False,
default=5
),
ToolParameter(
name="include_citations",
type="boolean",
description="是否包含引用来源(默认:true)",
required=False,
default=True
)
]
@tool_action("rag_add_document", "添加文档到知识库(支持PDF、Word、Excel、PPT、图片、音频等多种格式)")
def _add_document(
self,
file_path: str,
document_id: str = None,
namespace: str = "default",
chunk_size: int = 800,
chunk_overlap: int = 100
) -> str:
"""添加文档到知识库
Args:
file_path: 文档文件路径
document_id: 文档ID(可选)
namespace: 知识库命名空间(用于隔离不同项目)
chunk_size: 分块大小
chunk_overlap: 分块重叠大小
Returns:
执行结果
"""
try:
if not file_path or not os.path.exists(file_path):
return f"❌ 文件不存在: {file_path}"
pipeline = self._get_pipeline(namespace)
t0 = time.time()
chunks_added = pipeline["add_documents"](
file_paths=[file_path],
chunk_size=chunk_size,
chunk_overlap=chunk_overlap
)
t1 = time.time()
process_ms = int((t1 - t0) * 1000)
if chunks_added == 0:
return f"⚠️ 未能从文件解析内容: {os.path.basename(file_path)}"
return (
f"✅ 文档已添加到知识库: {os.path.basename(file_path)}\n"
f"📊 分块数量: {chunks_added}\n"
f"⏱️ 处理时间: {process_ms}ms\n"
f"📝 命名空间: {pipeline.get('namespace', self.rag_namespace)}"
)
except Exception as e:
return f"❌ 添加文档失败: {str(e)}"
@tool_action("rag_add_text", "添加文本到知识库")
def _add_text(
self,
text: str,
document_id: str = None,
namespace: str = "default",
chunk_size: int = 800,
chunk_overlap: int = 100
) -> str:
"""添加文本到知识库
Args:
text: 要添加的文本内容
document_id: 文档ID(可选)
namespace: 知识库命名空间
chunk_size: 分块大小
chunk_overlap: 分块重叠大小
Returns:
执行结果
"""
metadata = None
try:
if not text or not text.strip():
return "❌ 文本内容不能为空"
# 创建临时文件
document_id = document_id or f"text_{abs(hash(text)) % 100000}"
tmp_path = os.path.join(self.knowledge_base_path, f"{document_id}.md")
try:
with open(tmp_path, 'w', encoding='utf-8') as f:
f.write(text)
pipeline = self._get_pipeline(namespace)
t0 = time.time()
chunks_added = pipeline["add_documents"](
file_paths=[tmp_path],
chunk_size=chunk_size,
chunk_overlap=chunk_overlap
)
t1 = time.time()
process_ms = int((t1 - t0) * 1000)
if chunks_added == 0:
return f"⚠️ 未能从文本生成有效分块"
return (
f"✅ 文本已添加到知识库: {document_id}\n"
f"📊 分块数量: {chunks_added}\n"
f"⏱️ 处理时间: {process_ms}ms\n"
f"📝 命名空间: {pipeline.get('namespace', self.rag_namespace)}"
)
finally:
# 清理临时文件
try:
if os.path.exists(tmp_path):
os.remove(tmp_path)
except Exception:
pass
except Exception as e:
return f"❌ 添加文本失败: {str(e)}"
@tool_action("rag_search", "搜索知识库中的相关内容")
def _search(
self,
query: str,
limit: int = 5,
min_score: float = 0.1,
enable_advanced_search: bool = True,
max_chars: int = 1200,
include_citations: bool = True,
namespace: str = "default"
) -> str:
"""搜索知识库
Args:
query: 搜索查询词
limit: 返回结果数量
min_score: 最低相关度分数
enable_advanced_search: 是否启用高级搜索(MQE、HyDE)
max_chars: 每个结果最大字符数
include_citations: 是否包含引用来源
namespace: 知识库命名空间
Returns:
搜索结果
"""
try:
if not query or not query.strip():
return "❌ 搜索查询不能为空"
# 使用统一 RAG 管道搜索
pipeline = self._get_pipeline(namespace)
if enable_advanced_search:
results = pipeline["search_advanced"](
query=query,
top_k=limit,
enable_mqe=True,
enable_hyde=True,
score_threshold=min_score if min_score > 0 else None
)
else:
results = pipeline["search"](
query=query,
top_k=limit,
score_threshold=min_score if min_score > 0 else None
)
if not results:
return f"🔍 未找到与 '{query}' 相关的内容"
# 格式化搜索结果
search_result = ["搜索结果:"]
for i, result in enumerate(results, 1):
meta = result.get("metadata", {})
score = result.get("score", 0.0)
content = meta.get("content", "")[:200] + "..."
source = meta.get("source_path", "unknown")
# 安全处理Unicode
def clean_text(text):
try:
return str(text).encode('utf-8', errors='ignore').decode('utf-8')
except Exception:
return str(text)
clean_content = clean_text(content)
clean_source = clean_text(source)
search_result.append(f"\n{i}. 文档: **{clean_source}** (相似度: {score:.3f})")
search_result.append(f" {clean_content}")
if include_citations and meta.get("heading_path"):
clean_heading = clean_text(str(meta['heading_path']))
search_result.append(f" 章节: {clean_heading}")
return "\n".join(search_result)
except Exception as e:
return f"❌ 搜索失败: {str(e)}"
@tool_action("rag_ask", "基于知识库进行智能问答")
def _ask(
self,
question: str,
limit: int = 5,
enable_advanced_search: bool = True,
include_citations: bool = True,
max_chars: int = 1200,
namespace: str = "default"
) -> str:
"""智能问答:检索 → 上下文注入 → LLM生成答案
Args:
question: 用户问题
limit: 检索结果数量
enable_advanced_search: 是否启用高级搜索
include_citations: 是否包含引用来源
max_chars: 每个结果最大字符数
namespace: 知识库命名空间
Returns:
智能问答结果
核心流程:
1. 解析用户问题
2. 智能检索相关内容
3. 构建上下文和提示词
4. LLM生成准确答案
5. 添加引用来源
"""
try:
# 验证问题
if not question or not question.strip():
return "❌ 请提供要询问的问题"
user_question = question.strip()
print(f"🔍 智能问答: {user_question}")
# 1. 检索相关内容
pipeline = self._get_pipeline(namespace)
search_start = time.time()
if enable_advanced_search:
results = pipeline["search_advanced"](
query=user_question,
top_k=limit,
enable_mqe=True,
enable_hyde=True
)
else:
results = pipeline["search"](
query=user_question,
top_k=limit
)
search_time = int((time.time() - search_start) * 1000)
if not results:
return (
f"🤔 抱歉,我在知识库中没有找到与「{user_question}」相关的信息。\n\n"
f"💡 建议:\n"
f"• 尝试使用更简洁的关键词\n"
f"• 检查是否已添加相关文档\n"
f"• 使用 stats 操作查看知识库状态"
)
# 2. 智能整理上下文
context_parts = []
citations = []
total_score = 0
for i, result in enumerate(results):
meta = result.get("metadata", {})
content = meta.get("content", "").strip()
source = meta.get("source_path", "unknown")
score = result.get("score", 0.0)
total_score += score
if content:
# 清理内容格式
cleaned_content = self._clean_content_for_context(content)
context_parts.append(f"片段 {i+1}:{cleaned_content}")
if include_citations:
citations.append({
"index": i+1,
"source": os.path.basename(source),
"score": score
})
# 3. 构建上下文(智能截断)
context = "\n\n".join(context_parts)
if len(context) > max_chars:
# 智能截断,保持完整性
context = self._smart_truncate_context(context, max_chars)
# 4. 构建增强提示词
system_prompt = self._build_system_prompt()
user_prompt = self._build_user_prompt(user_question, context)
enhanced_prompt = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
# 5. 调用 LLM 生成答案
llm_start = time.time()
answer = self.llm.think(enhanced_prompt)
llm_time = int((time.time() - llm_start) * 1000)
if not answer or not answer.strip():
return "❌ LLM未能生成有效答案,请稍后重试"
# 6. 构建最终回答
final_answer = self._format_final_answer(
question=user_question,
answer=answer.strip(),
citations=citations if include_citations else None,
search_time=search_time,
llm_time=llm_time,
avg_score=total_score / len(results) if results else 0
)
return final_answer
except Exception as e:
return f"❌ 智能问答失败: {str(e)}\n💡 请检查知识库状态或稍后重试"
def _clean_content_for_context(self, content: str) -> str:
"""清理内容用于上下文"""
# 移除过多的换行和空格
content = " ".join(content.split())
# 截断过长内容
if len(content) > 300:
content = content[:300] + "..."
return content
def _smart_truncate_context(self, context: str, max_chars: int) -> str:
"""智能截断上下文,保持段落完整性"""
if len(context) <= max_chars:
return context
# 寻找最近的段落分隔符
truncated = context[:max_chars]
last_break = truncated.rfind("\n\n")
if last_break > max_chars * 0.7: # 如果断点位置合理
return truncated[:last_break] + "\n\n[...更多内容被截断]"
else:
return truncated[:max_chars-20] + "...[内容被截断]"
def _build_system_prompt(self) -> str:
"""构建系统提示词"""
return (
"你是一个专业的知识助手,具备以下能力:\n"
"1. 📖 精准理解:仔细理解用户问题的核心意图\n"
"2. 🎯 可信回答:严格基于提供的上下文信息回答,不编造内容\n"
"3. 🔍 信息整合:从多个片段中提取关键信息,形成完整答案\n"
"4. 💡 清晰表达:用简洁明了的语言回答,适当使用结构化格式\n"
"5. 🚫 诚实表达:如果上下文不足以回答问题,请坦诚说明\n\n"
"回答格式要求:\n"
"• 直接回答核心问题\n"
"• 必要时使用要点或步骤\n"
"• 引用关键原文时使用引号\n"
"• 避免重复和冗余"
)
def _build_user_prompt(self, question: str, context: str) -> str:
"""构建用户提示词"""
return (
f"请基于以下上下文信息回答问题:\n\n"
f"【问题】{question}\n\n"
f"【相关上下文】\n{context}\n\n"
f"【要求】请提供准确、有帮助的回答。如果上下文信息不足,请说明需要什么额外信息。"
)
def _format_final_answer(self, question: str, answer: str, citations: Optional[List[Dict]] = None, search_time: int = 0, llm_time: int = 0, avg_score: float = 0) -> str:
"""格式化最终答案"""
result = [f"🤖 **智能问答结果**\n"]
result.append(answer)
if citations:
result.append("\n\n📚 **参考来源**")
for citation in citations:
score_emoji = "🟢" if citation["score"] > 0.8 else "🟡" if citation["score"] > 0.6 else "🔵"
result.append(f"{score_emoji} [{citation['index']}] {citation['source']} (相似度: {citation['score']:.3f})")
# 添加性能信息(调试模式)
result.append(f"\n⚡ 检索: {search_time}ms | 生成: {llm_time}ms | 平均相似度: {avg_score:.3f}")
return "\n".join(result)
@tool_action("rag_clear", "清空知识库(危险操作,请谨慎使用)")
def _clear_knowledge_base(self, confirm: bool = False, namespace: str = "default") -> str:
"""清空知识库
Args:
confirm: 确认执行(必须设置为True)
namespace: 知识库命名空间
Returns:
执行结果
"""
try:
if not confirm:
return (
"⚠️ 危险操作:清空知识库将删除所有数据!\n"
"请使用 confirm=true 参数确认执行。"
)
pipeline = self._get_pipeline(namespace)
store = pipeline.get("store")
namespace_id = pipeline.get("namespace", self.rag_namespace)
success = store.clear_collection() if store else False
if success:
# 重新初始化该命名空间
self._pipelines[namespace_id] = create_rag_pipeline(
qdrant_url=self.qdrant_url,
qdrant_api_key=self.qdrant_api_key,
collection_name=self.collection_name,
rag_namespace=namespace_id
)
return f"✅ 知识库已成功清空(命名空间:{namespace_id})"
else:
return "❌ 清空知识库失败"
except Exception as e:
return f"❌ 清空知识库失败: {str(e)}"
@tool_action("rag_stats", "获取知识库统计信息")
def _get_stats(self, namespace: str = "default") -> str:
"""获取知识库统计
Args:
namespace: 知识库命名空间
Returns:
统计信息
"""
try:
pipeline = self._get_pipeline(namespace)
stats = pipeline["get_stats"]()
stats_info = [
"📊 **RAG 知识库统计**",
f"📝 命名空间: {pipeline.get('namespace', self.rag_namespace)}",
f"📋 集合名称: {self.collection_name}",
f"📂 存储根路径: {self.knowledge_base_path}"
]
# 添加存储统计
if stats:
store_type = stats.get("store_type", "unknown")
total_vectors = (
stats.get("points_count") or
stats.get("vectors_count") or
stats.get("count") or 0
)
stats_info.extend([
f"📦 存储类型: {store_type}",
f"📊 文档分块数: {int(total_vectors)}",
])
if "config" in stats:
config = stats["config"]
if isinstance(config, dict):
vector_size = config.get("vector_size", "unknown")
distance = config.get("distance", "unknown")
stats_info.extend([
f"🔢 向量维度: {vector_size}",
f"📎 距离度量: {distance}"
])
# 添加系统状态
stats_info.extend([
"",
"🟢 **系统状态**",
f"✅ RAG 管道: {'正常' if self.initialized else '异常'}",
f"✅ LLM 连接: {'正常' if hasattr(self, 'llm') else '异常'}"
])
return "\n".join(stats_info)
except Exception as e:
return f"❌ 获取统计信息失败: {str(e)}"
def get_relevant_context(self, query: str, limit: int = 3, max_chars: int = 1200, namespace: Optional[str] = None) -> str:
"""为查询获取相关上下文
这个方法可以被Agent调用来获取相关的知识库上下文
"""
try:
if not query:
return ""
# 使用统一 RAG 管道搜索
pipeline = self._get_pipeline(namespace)
results = pipeline["search"](
query=query,
top_k=limit
)
if not results:
return ""
# 合并上下文
context_parts = []
for result in results:
content = result.get("metadata", {}).get("content", "")
if content:
context_parts.append(content)
merged_context = "\n\n".join(context_parts)
# 限制长度
if len(merged_context) > max_chars:
merged_context = merged_context[:max_chars] + "..."
return merged_context
except Exception as e:
return f"获取上下文失败: {str(e)}"
def batch_add_texts(self, texts: List[str], document_ids: Optional[List[str]] = None, chunk_size: int = 800, chunk_overlap: int = 100, namespace: Optional[str] = None) -> str:
"""批量添加文本"""
try:
if not texts:
return "❌ 文本列表不能为空"
if document_ids and len(document_ids) != len(texts):
return "❌ 文本数量和文档ID数量不匹配"
pipeline = self._get_pipeline(namespace)
t0 = time.time()
total_chunks = 0
successful_files = []
for i, text in enumerate(texts):
if not text or not text.strip():
continue
doc_id = document_ids[i] if document_ids else f"batch_text_{i}"
tmp_path = os.path.join(self.knowledge_base_path, f"{doc_id}.md")
try:
with open(tmp_path, 'w', encoding='utf-8') as f:
f.write(text)
chunks_added = pipeline["add_documents"](
file_paths=[tmp_path],
chunk_size=chunk_size,
chunk_overlap=chunk_overlap
)
total_chunks += chunks_added
successful_files.append(doc_id)
finally:
# 清理临时文件
try:
if os.path.exists(tmp_path):
os.remove(tmp_path)
except Exception:
pass
t1 = time.time()
process_ms = int((t1 - t0) * 1000)
return (
f"✅ 批量添加完成\n"
f"📊 成功文件: {len(successful_files)}/{len(texts)}\n"
f"📊 总分块数: {total_chunks}\n"
f"⏱️ 处理时间: {process_ms}ms"
)
except Exception as e:
return f"❌ 批量添加失败: {str(e)}"
def clear_all_namespaces(self) -> str:
"""清空当前工具管理的所有命名空间数据"""
try:
for ns, pipeline in self._pipelines.items():
store = pipeline.get("store")
if store:
store.clear_collection()
self._pipelines.clear()
# 重新初始化默认命名空间
self._init_components()
return "✅ 所有命名空间数据已清空并重新初始化"
except Exception as e:
return f"❌ 清空所有命名空间失败: {str(e)}"
# ========================================
# 便捷接口方法(简化用户调用)
# ========================================
def add_document(self, file_path: str, namespace: str = "default") -> str:
"""便捷方法:添加单个文档"""
return self.run({
"action": "add_document",
"file_path": file_path,
"namespace": namespace
})
def add_text(self, text: str, namespace: str = "default", document_id: str = None) -> str:
"""便捷方法:添加文本内容"""
return self.run({
"action": "add_text",
"text": text,
"namespace": namespace,
"document_id": document_id
})
def ask(self, question: str, namespace: str = "default", **kwargs) -> str:
"""便捷方法:智能问答"""
params = {
"action": "ask",
"question": question,
"namespace": namespace
}
params.update(kwargs)
return self.run(params)
def search(self, query: str, namespace: str = "default", **kwargs) -> str:
"""便捷方法:搜索知识库"""
params = {
"action": "search",
"query": query,
"namespace": namespace
}
params.update(kwargs)
return self.run(params)
def add_documents_batch(self, file_paths: List[str], namespace: str = "default") -> str:
"""批量添加多个文档"""
if not file_paths:
return "❌ 文件路径列表不能为空"
results = []
successful = 0
failed = 0
total_chunks = 0
start_time = time.time()
for i, file_path in enumerate(file_paths, 1):
print(f"📄 处理文档 {i}/{len(file_paths)}: {os.path.basename(file_path)}")
try:
result = self.add_document(file_path, namespace)
if "✅" in result:
successful += 1
# 提取分块数量
if "分块数量:" in result:
chunks = int(result.split("分块数量: ")[1].split("\n")[0])
total_chunks += chunks
else:
failed += 1
results.append(f"❌ {os.path.basename(file_path)}: 处理失败")
except Exception as e:
failed += 1
results.append(f"❌ {os.path.basename(file_path)}: {str(e)}")
process_time = int((time.time() - start_time) * 1000)
summary = [
"📊 **批量处理完成**",
f"✅ 成功: {successful}/{len(file_paths)} 个文档",
f"📊 总分块数: {total_chunks}",
f"⏱️ 总耗时: {process_time}ms",
f"📝 命名空间: {namespace}"
]
if failed > 0:
summary.append(f"❌ 失败: {failed} 个文档")
summary.append("\n**失败详情:**")
summary.extend(results)
return "\n".join(summary)
def add_texts_batch(self, texts: List[str], namespace: str = "default", document_ids: Optional[List[str]] = None) -> str:
"""批量添加多个文本"""
if not texts:
return "❌ 文本列表不能为空"
if document_ids and len(document_ids) != len(texts):
return "❌ 文本数量和文档ID数量不匹配"
results = []
successful = 0
failed = 0
total_chunks = 0
start_time = time.time()
for i, text in enumerate(texts):
doc_id = document_ids[i] if document_ids else f"batch_text_{i+1}"
print(f"📝 处理文本 {i+1}/{len(texts)}: {doc_id}")
try:
result = self.add_text(text, namespace, doc_id)
if "✅" in result:
successful += 1
# 提取分块数量
if "分块数量:" in result:
chunks = int(result.split("分块数量: ")[1].split("\n")[0])
total_chunks += chunks
else:
failed += 1
results.append(f"❌ {doc_id}: 处理失败")
except Exception as e:
failed += 1
results.append(f"❌ {doc_id}: {str(e)}")
process_time = int((time.time() - start_time) * 1000)
summary = [
"📊 **批量文本处理完成**",
f"✅ 成功: {successful}/{len(texts)} 个文本",
f"📊 总分块数: {total_chunks}",
f"⏱️ 总耗时: {process_time}ms",
f"📝 命名空间: {namespace}"
]
if failed > 0:
summary.append(f"❌ 失败: {failed} 个文本")
summary.append("\n**失败详情:**")
summary.extend(results)
return "\n".join(summary)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 ## storage ### document_store.py ```py """文档存储实现 支持多种文档数据库后端: - SQLite: 轻量级关系型数据库 - PostgreSQL: 企业级关系型数据库(可扩展) """ from abc import ABC, abstractmethod from typing import List, Dict, Any, Optional import sqlite3 import json import os import threading class DocumentStore(ABC): """文档存储基类""" @abstractmethod def add_memory( self, memory_id: str, user_id: str, content: str, memory_type: str, timestamp: int, importance: float, properties: Dict[str, Any] = None ) -> str: """添加记忆""" pass @abstractmethod def get_memory(self, memory_id: str) -> Optional[Dict[str, Any]]: """获取单个记忆""" pass @abstractmethod def search_memories( self, user_id: Optional[str] = None, memory_type: Optional[str] = None, start_time: Optional[int] = None, end_time: Optional[int] = None, importance_threshold: Optional[float] = None, limit: int = 10 ) -> List[Dict[str, Any]]: """搜索记忆""" pass @abstractmethod def update_memory( self, memory_id: str, content: str = None, importance: float = None, properties: Dict[str, Any] = None ) -> bool: """更新记忆""" pass @abstractmethod def delete_memory(self, memory_id: str) -> bool: """删除记忆""" pass @abstractmethod def get_database_stats(self) -> Dict[str, Any]: """获取数据库统计信息""" pass @abstractmethod def add_document(self, content: str, metadata: Dict[str, Any] = None) -> str: """添加文档""" pass @abstractmethod def get_document(self, document_id: str) -> Optional[Dict[str, Any]]: """获取文档""" pass class SQLiteDocumentStore(DocumentStore): """SQLite文档存储实现""" _instances = {} # 存储已创建的实例 _initialized_dbs = set() # 存储已初始化的数据库路径 def __new__(cls, db_path: str = "./memory.db"): """单例模式,同一路径只创建一个实例""" abs_path = os.path.abspath(db_path) if abs_path not in cls._instances: instance = super(SQLiteDocumentStore, cls).__new__(cls) cls._instances[abs_path] = instance return cls._instances[abs_path] def __init__(self, db_path: str = "./memory.db"): # 避免重复初始化 if hasattr(self, '_initialized'): return self.db_path = db_path self.local = threading.local() # 确保目录存在 os.makedirs(os.path.dirname(os.path.abspath(db_path)), exist_ok=True) # 初始化数据库(只初始化一次) abs_path = os.path.abspath(db_path) if abs_path not in self._initialized_dbs: self._init_database() self._initialized_dbs.add(abs_path) print(f"[OK] SQLite 文档存储初始化完成: {db_path}") self._initialized = True def _get_connection(self): """获取线程本地连接""" if not hasattr(self.local, 'connection'): self.local.connection = sqlite3.connect(self.db_path) self.local.connection.row_factory = sqlite3.Row # 使结果可以按列名访问 return self.local.connection def _init_database(self): """初始化数据库表""" conn = self._get_connection() cursor = conn.cursor() # 创建用户表 cursor.execute(""" CREATE TABLE IF NOT EXISTS users ( id TEXT PRIMARY KEY, name TEXT, properties TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) # 创建记忆表 cursor.execute(""" CREATE TABLE IF NOT EXISTS memories ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL, content TEXT NOT NULL, memory_type TEXT NOT NULL, timestamp INTEGER NOT NULL, importance REAL NOT NULL, properties TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (user_id) REFERENCES users (id) ) """) # 创建概念表 cursor.execute(""" CREATE TABLE IF NOT EXISTS concepts ( id TEXT PRIMARY KEY, name TEXT NOT NULL, description TEXT, properties TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) # 创建记忆-概念关联表 cursor.execute(""" CREATE TABLE IF NOT EXISTS memory_concepts ( memory_id TEXT NOT NULL, concept_id TEXT NOT NULL, relevance_score REAL DEFAULT 1.0, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY (memory_id, concept_id), FOREIGN KEY (memory_id) REFERENCES memories (id) ON DELETE CASCADE, FOREIGN KEY (concept_id) REFERENCES concepts (id) ON DELETE CASCADE ) """) # 创建概念关系表 cursor.execute(""" CREATE TABLE IF NOT EXISTS concept_relationships ( from_concept_id TEXT NOT NULL, to_concept_id TEXT NOT NULL, relationship_type TEXT NOT NULL, strength REAL DEFAULT 1.0, properties TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY (from_concept_id, to_concept_id, relationship_type), FOREIGN KEY (from_concept_id) REFERENCES concepts (id) ON DELETE CASCADE, FOREIGN KEY (to_concept_id) REFERENCES concepts (id) ON DELETE CASCADE ) """) # 创建索引 indexes = [ "CREATE INDEX IF NOT EXISTS idx_memories_user_id ON memories (user_id)", "CREATE INDEX IF NOT EXISTS idx_memories_type ON memories (memory_type)", "CREATE INDEX IF NOT EXISTS idx_memories_timestamp ON memories (timestamp)", "CREATE INDEX IF NOT EXISTS idx_memories_importance ON memories (importance)", "CREATE INDEX IF NOT EXISTS idx_memory_concepts_memory ON memory_concepts (memory_id)", "CREATE INDEX IF NOT EXISTS idx_memory_concepts_concept ON memory_concepts (concept_id)" ] for index_sql in indexes: cursor.execute(index_sql) conn.commit() print("[OK] SQLite 数据库表和索引创建完成") def add_memory( self, memory_id: str, user_id: str, content: str, memory_type: str, timestamp: int, importance: float, properties: Dict[str, Any] = None ) -> str: """添加记忆""" conn = self._get_connection() cursor = conn.cursor() # 确保用户存在 cursor.execute("INSERT OR IGNORE INTO users (id, name) VALUES (?, ?)", (user_id, user_id)) # 插入记忆 cursor.execute(""" INSERT OR REPLACE INTO memories (id, user_id, content, memory_type, timestamp, importance, properties, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP) """, ( memory_id, user_id, content, memory_type, timestamp, importance, json.dumps(properties) if properties else None )) conn.commit() return memory_id def get_memory(self, memory_id: str) -> Optional[Dict[str, Any]]: """获取单个记忆""" conn = self._get_connection() cursor = conn.cursor() cursor.execute(""" SELECT id, user_id, content, memory_type, timestamp, importance, properties, created_at FROM memories WHERE id = ? """, (memory_id,)) row = cursor.fetchone() if not row: return None return { "memory_id": row["id"], "user_id": row["user_id"], "content": row["content"], "memory_type": row["memory_type"], "timestamp": row["timestamp"], "importance": row["importance"], "properties": json.loads(row["properties"]) if row["properties"] else {}, "created_at": row["created_at"] } def search_memories( self, user_id: Optional[str] = None, memory_type: Optional[str] = None, start_time: Optional[int] = None, end_time: Optional[int] = None, importance_threshold: Optional[float] = None, limit: int = 10 ) -> List[Dict[str, Any]]: """搜索记忆""" conn = self._get_connection() cursor = conn.cursor() # 构建查询条件 where_conditions = [] params = [] if user_id: where_conditions.append("user_id = ?") params.append(user_id) if memory_type: where_conditions.append("memory_type = ?") params.append(memory_type) if start_time: where_conditions.append("timestamp >= ?") params.append(start_time) if end_time: where_conditions.append("timestamp <= ?") params.append(end_time) if importance_threshold: where_conditions.append("importance >= ?") params.append(importance_threshold) where_clause = "" if where_conditions: where_clause = "WHERE " + " AND ".join(where_conditions) cursor.execute(f""" SELECT id, user_id, content, memory_type, timestamp, importance, properties, created_at FROM memories {where_clause} ORDER BY importance DESC, timestamp DESC LIMIT ? """, params + [limit]) memories = [] for row in cursor.fetchall(): memories.append({ "memory_id": row["id"], "user_id": row["user_id"], "content": row["content"], "memory_type": row["memory_type"], "timestamp": row["timestamp"], "importance": row["importance"], "properties": json.loads(row["properties"]) if row["properties"] else {}, "created_at": row["created_at"] }) return memories def update_memory( self, memory_id: str, content: str = None, importance: float = None, properties: Dict[str, Any] = None ) -> bool: """更新记忆""" conn = self._get_connection() cursor = conn.cursor() # 构建更新字段 update_fields = [] params = [] if content is not None: update_fields.append("content = ?") params.append(content) if importance is not None: update_fields.append("importance = ?") params.append(importance) if properties is not None: update_fields.append("properties = ?") params.append(json.dumps(properties)) if not update_fields: return False update_fields.append("updated_at = CURRENT_TIMESTAMP") params.append(memory_id) cursor.execute(f""" UPDATE memories SET {', '.join(update_fields)} WHERE id = ? """, params) conn.commit() return cursor.rowcount > 0 def delete_memory(self, memory_id: str) -> bool: """删除记忆""" conn = self._get_connection() cursor = conn.cursor() cursor.execute("DELETE FROM memories WHERE id = ?", (memory_id,)) deleted_count = cursor.rowcount conn.commit() return deleted_count > 0 def get_database_stats(self) -> Dict[str, Any]: """获取数据库统计信息""" conn = self._get_connection() cursor = conn.cursor() stats = {} # 统计各表的记录数 tables = ["users", "memories", "concepts", "memory_concepts", "concept_relationships"] for table in tables: cursor.execute(f"SELECT COUNT(*) as count FROM {table}") stats[f"{table}_count"] = cursor.fetchone()["count"] # 统计记忆类型分布 cursor.execute(""" SELECT memory_type, COUNT(*) as count FROM memories GROUP BY memory_type """) memory_types = {} for row in cursor.fetchall(): memory_types[row["memory_type"]] = row["count"] stats["memory_types"] = memory_types # 统计用户分布 cursor.execute(""" SELECT user_id, COUNT(*) as count FROM memories GROUP BY user_id ORDER BY count DESC LIMIT 10 """) top_users = {} for row in cursor.fetchall(): top_users[row["user_id"]] = row["count"] stats["top_users"] = top_users stats["store_type"] = "sqlite" stats["db_path"] = self.db_path return stats def add_document(self, content: str, metadata: Dict[str, Any] = None) -> str: """添加文档""" import uuid import time doc_id = str(uuid.uuid4()) user_id = metadata.get("user_id", "system") if metadata else "system" return self.add_memory( memory_id=doc_id, user_id=user_id, content=content, memory_type="document", timestamp=int(time.time()), importance=0.5, properties=metadata or {} ) def get_document(self, document_id: str) -> Optional[Dict[str, Any]]: """获取文档""" return self.get_memory(document_id) def close(self): """关闭数据库连接""" if hasattr(self.local, 'connection'): self.local.connection.close() delattr(self.local, 'connection') print("[OK] SQLite 连接已关闭")
embedding.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 """统一嵌入模块(实现 + 提供器) 说明(中文): - 提供统一的文本嵌入接口与多实现:本地Transformer、DashScope(通义千问)、TF-IDF兜底。 - 暴露 get_text_embedder()/get_dimension()/refresh_embedder() 供各记忆类型统一使用。 - 通过环境变量优先级:dashscope > local > tfidf。 环境变量: - EMBED_MODEL_TYPE: "dashscope" | "local" | "tfidf"(默认 dashscope) - EMBED_MODEL_NAME: 模型名称(dashscope默认 text-embedding-v3;local默认 sentence-transformers/all-MiniLM-L6-v2) - EMBED_API_KEY: Embedding API Key(统一命名) - EMBED_BASE_URL: Embedding Base URL(统一命名,可选) """ from typing import List , Union , Optional import threadingimport osimport numpy as npclass EmbeddingModel : """嵌入模型基类(最小接口)""" def encode (self, texts: Union [str , List [str ]] ): raise NotImplementedError @property def dimension (self ) -> int : raise NotImplementedError class LocalTransformerEmbedding (EmbeddingModel ): """本地Transformer嵌入(优先 sentence-transformers,缺失回退 transformers+torch)""" def __init__ (self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2" ): self.model_name = model_name self._backend = None self._st_model = None self._hf_tokenizer = None self._hf_model = None self._dimension = None self._load_backend() def _load_backend (self ): try : from sentence_transformers import SentenceTransformer self._st_model = SentenceTransformer(self.model_name) test_vec = self._st_model.encode("test_text" ) self._dimension = len (test_vec) self._backend = "st" return except Exception: self._st_model = None try : from transformers import AutoTokenizer, AutoModel import torch self._hf_tokenizer = AutoTokenizer.from_pretrained(self.model_name) self._hf_model = AutoModel.from_pretrained(self.model_name) with torch.no_grad(): inputs = self._hf_tokenizer("test_text" , return_tensors="pt" , padding=True , truncation=True ) outputs = self._hf_model(**inputs) test_embedding = outputs.last_hidden_state.mean(dim=1 ) self._dimension = int (test_embedding.shape[1 ]) self._backend = "hf" return except Exception: self._hf_tokenizer = None self._hf_model = None raise ImportError("未找到可用的本地嵌入后端,请安装 sentence-transformers 或 transformers+torch" ) def encode (self, texts: Union [str , List [str ]] ): if isinstance (texts, str ): inputs = [texts] single = True else : inputs = list (texts) single = False if self._backend == "st" : vecs = self._st_model.encode(inputs) if hasattr (vecs, "tolist" ): vecs = [v for v in vecs] else : import torch tokenized = self._hf_tokenizer(inputs, return_tensors="pt" , padding=True , truncation=True , max_length=512 ) with torch.no_grad(): outputs = self._hf_model(**tokenized) embeddings = outputs.last_hidden_state.mean(dim=1 ).cpu().numpy() vecs = [v for v in embeddings] if single: return vecs[0 ] return vecs @property def dimension (self ) -> int : return int (self._dimension or 0 ) class TFIDFEmbedding (EmbeddingModel ): """TF-IDF 简易兜底(在无深度模型时保证可用)""" def __init__ (self, max_features: int = 1000 ): self.max_features = max_features self._vectorizer = None self._is_fitted = False self._dimension = max_features self._init_vectorizer() def _init_vectorizer (self ): try : from sklearn.feature_extraction.text import TfidfVectorizer self._vectorizer = TfidfVectorizer(max_features=self.max_features, stop_words='english' ) except ImportError: raise ImportError("请安装 scikit-learn: pip install scikit-learn" ) def fit (self, texts: List [str ] ): self._vectorizer.fit(texts) self._is_fitted = True self._dimension = len (self._vectorizer.get_feature_names_out()) def encode (self, texts: Union [str , List [str ]] ): if not self._is_fitted: raise ValueError("TF-IDF模型未训练,请先调用fit()方法" ) if isinstance (texts, str ): texts = [texts] single = True else : single = False tfidf_matrix = self._vectorizer.transform(texts) embeddings = tfidf_matrix.toarray() if single: return embeddings[0 ] return [e for e in embeddings] @property def dimension (self ) -> int : return self._dimension class DashScopeEmbedding (EmbeddingModel ): """阿里云 DashScope(通义千问)Embedding / OpenAI兼容REST 模式 行为: - 如提供 base_url,则优先使用 OpenAI 兼容的 REST 接口(POST {base_url}/embeddings)。 - 否则使用官方 dashscope SDK 的 TextEmbedding.call。 """ def __init__ (self, model_name: str = "text-embedding-v3" , api_key: Optional [str ] = None , base_url: Optional [str ] = None ): self.model_name = model_name self.api_key = api_key self.base_url = base_url self._dimension = None if not self.base_url: self._init_client() test = self.encode("health_check" ) self._dimension = len (test) def _init_client (self ): try : if self.api_key: os.environ["DASHSCOPE_API_KEY" ] = self.api_key import dashscope except ImportError: raise ImportError("请安装 dashscope: pip install dashscope" ) def encode (self, texts: Union [str , List [str ]] ): if isinstance (texts, str ): inputs = [texts] single = True else : inputs = list (texts) single = False if self.base_url: import requests url = self.base_url.rstrip("/" ) + "/embeddings" headers = { "Authorization" : f"Bearer {self.api_key} " if self.api_key else "" , "Content-Type" : "application/json" , } payload = {"model" : self.model_name, "input" : inputs} resp = requests.post(url, headers=headers, json=payload, timeout=30 ) if resp.status_code >= 400 : raise RuntimeError(f"Embedding REST 调用失败: {resp.status_code} {resp.text} " ) data = resp.json() items = data.get("data" ) or [] vecs = [np.array(item.get("embedding" )) for item in items] if single: return vecs[0 ] return vecs from dashscope import TextEmbedding rsp = TextEmbedding.call(model=self.model_name, input =inputs) embeddings_obj = None if isinstance (rsp, dict ): embeddings_obj = (rsp.get("output" ) or {}).get("embeddings" ) else : embeddings_obj = getattr (getattr (rsp, "output" , None ), "embeddings" , None ) if not embeddings_obj: raise RuntimeError("DashScope 返回为空或格式不匹配" ) vecs = [np.array(item.get("embedding" ) or item.get("vector" )) for item in embeddings_obj] if single: return vecs[0 ] return vecs @property def dimension (self ) -> int : return int (self._dimension or 0 ) def create_embedding_model (model_type: str = "local" , **kwargs ) -> EmbeddingModel: """创建嵌入模型实例 model_type: "dashscope" | "local" | "tfidf" kwargs: model_name, api_key """ if model_type in ("local" , "sentence_transformer" , "huggingface" ): return LocalTransformerEmbedding(**kwargs) elif model_type == "dashscope" : return DashScopeEmbedding(**kwargs) elif model_type == "tfidf" : return TFIDFEmbedding(**kwargs) else : raise ValueError(f"不支持的模型类型: {model_type} " ) def create_embedding_model_with_fallback (preferred_type: str = "dashscope" , **kwargs ) -> EmbeddingModel: """带回退的创建:dashscope -> local -> tfidf""" if preferred_type in ("sentence_transformer" , "huggingface" ): preferred_type = "local" fallback = ["dashscope" , "local" , "tfidf" ] if preferred_type in fallback: fallback.remove(preferred_type) fallback.insert(0 , preferred_type) for t in fallback: try : return create_embedding_model(t, **kwargs) except Exception: continue raise RuntimeError("所有嵌入模型都不可用,请安装依赖或检查配置" ) _lock = threading.RLock() _embedder: Optional [EmbeddingModel] = None def _build_embedder () -> EmbeddingModel: preferred = os.getenv("EMBED_MODEL_TYPE" , "dashscope" ).strip() default_model = "text-embedding-v3" if preferred == "dashscope" else "sentence-transformers/all-MiniLM-L6-v2" model_name = os.getenv("EMBED_MODEL_NAME" , default_model).strip() kwargs = {} if model_name: kwargs["model_name" ] = model_name api_key = os.getenv("EMBED_API_KEY" ) if api_key: kwargs["api_key" ] = api_key base_url = os.getenv("EMBED_BASE_URL" ) if base_url: kwargs["base_url" ] = base_url return create_embedding_model_with_fallback(preferred_type=preferred, **kwargs) def get_text_embedder () -> EmbeddingModel: """获取全局共享的文本嵌入实例(线程安全单例)""" global _embedder if _embedder is not None : return _embedder with _lock: if _embedder is None : _embedder = _build_embedder() return _embedder def get_dimension (default: int = 384 ) -> int : """获取统一向量维度(失败回退默认值)""" try : return int (getattr (get_text_embedder(), "dimension" , default)) except Exception: return int (default) def refresh_embedder () -> EmbeddingModel: """强制重建嵌入实例(可用于动态切换环境变量)""" global _embedder with _lock: _embedder = _build_embedder() return _embedder
neo4j_store.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 """ Neo4j图数据库存储实现 """ import loggingfrom typing import Dict , List , Optional , Any , Tuple from datetime import datetimetry : from neo4j import GraphDatabase from neo4j.exceptions import ServiceUnavailable, AuthError NEO4J_AVAILABLE = True except ImportError: NEO4J_AVAILABLE = False GraphDatabase = None logger = logging.getLogger(__name__) class Neo4jGraphStore : """Neo4j图数据库存储实现""" def __init__ ( self, uri: str = "bolt://localhost:7687" , username: str = "neo4j" , password: str = "hello-agents-password" , database: str = "neo4j" , max_connection_lifetime: int = 3600 , max_connection_pool_size: int = 50 , connection_acquisition_timeout: int = 60 , **kwargs ): """ 初始化Neo4j图存储 (支持云API) Args: uri: Neo4j连接URI (本地: bolt://localhost:7687, 云: neo4j+s://xxx.databases.neo4j.io) username: 用户名 password: 密码 database: 数据库名称 max_connection_lifetime: 最大连接生命周期(秒) max_connection_pool_size: 最大连接池大小 connection_acquisition_timeout: 连接获取超时(秒) """ if not NEO4J_AVAILABLE: raise ImportError( "neo4j未安装。请运行: pip install neo4j>=5.0.0" ) self.uri = uri self.username = username self.password = password self.database = database self.driver = None self._initialize_driver( max_connection_lifetime=max_connection_lifetime, max_connection_pool_size=max_connection_pool_size, connection_acquisition_timeout=connection_acquisition_timeout ) self._create_indexes() def _initialize_driver (self, **config ): """初始化Neo4j驱动""" try : self.driver = GraphDatabase.driver( self.uri, auth=(self.username, self.password), **config ) self.driver.verify_connectivity() if "neo4j.io" in self.uri or "aura" in self.uri.lower(): logger.info(f"✅ 成功连接到Neo4j云服务: {self.uri} " ) else : logger.info(f"✅ 成功连接到Neo4j服务: {self.uri} " ) except AuthError as e: logger.error(f"❌ Neo4j认证失败: {e} " ) logger.info("💡 请检查用户名和密码是否正确" ) raise except ServiceUnavailable as e: logger.error(f"❌ Neo4j服务不可用: {e} " ) if "localhost" in self.uri: logger.info("💡 本地连接失败,可以考虑使用Neo4j Aura云服务" ) logger.info("💡 或启动本地服务: docker run -p 7474:7474 -p 7687:7687 neo4j:5.14" ) else : logger.info("💡 请检查URL和网络连接" ) raise except Exception as e: logger.error(f"❌ Neo4j连接失败: {e} " ) raise def _create_indexes (self ): """创建必要的索引以提高查询性能""" indexes = [ "CREATE INDEX entity_id_index IF NOT EXISTS FOR (e:Entity) ON (e.id)" , "CREATE INDEX entity_name_index IF NOT EXISTS FOR (e:Entity) ON (e.name)" , "CREATE INDEX entity_type_index IF NOT EXISTS FOR (e:Entity) ON (e.type)" , "CREATE INDEX memory_id_index IF NOT EXISTS FOR (m:Memory) ON (m.id)" , "CREATE INDEX memory_type_index IF NOT EXISTS FOR (m:Memory) ON (m.memory_type)" , "CREATE INDEX memory_timestamp_index IF NOT EXISTS FOR (m:Memory) ON (m.timestamp)" , ] with self.driver.session(database=self.database) as session: for index_query in indexes: try : session.run(index_query) except Exception as e: logger.debug(f"索引创建跳过 (可能已存在): {e} " ) logger.info("✅ Neo4j索引创建完成" ) def add_entity (self, entity_id: str , name: str , entity_type: str , properties: Dict [str , Any ] = None ) -> bool : """ 添加实体节点 Args: entity_id: 实体ID name: 实体名称 entity_type: 实体类型 properties: 附加属性 Returns: bool: 是否成功 """ try : props = properties or {} props.update({ "id" : entity_id, "name" : name, "type" : entity_type, "created_at" : datetime.now().isoformat(), "updated_at" : datetime.now().isoformat() }) query = """ MERGE (e:Entity {id: $entity_id}) SET e += $properties RETURN e """ with self.driver.session(database=self.database) as session: result = session.run(query, entity_id=entity_id, properties=props) record = result.single() if record: logger.debug(f"✅ 添加实体: {name} ({entity_type} )" ) return True return False except Exception as e: logger.error(f"❌ 添加实体失败: {e} " ) return False def add_relationship ( self, from_entity_id: str , to_entity_id: str , relationship_type: str , properties: Dict [str , Any ] = None ) -> bool : """ 添加实体间关系 Args: from_entity_id: 源实体ID to_entity_id: 目标实体ID relationship_type: 关系类型 properties: 关系属性 Returns: bool: 是否成功 """ try : props = properties or {} props.update({ "type" : relationship_type, "created_at" : datetime.now().isoformat(), "updated_at" : datetime.now().isoformat() }) query = f""" MATCH (from:Entity {{id: $from_id}}) MATCH (to:Entity {{id: $to_id}}) MERGE (from)-[r:{relationship_type} ]->(to) SET r += $properties RETURN r """ with self.driver.session(database=self.database) as session: result = session.run( query, from_id=from_entity_id, to_id=to_entity_id, properties=props ) record = result.single() if record: logger.debug(f"✅ 添加关系: {from_entity_id} -{relationship_type} -> {to_entity_id} " ) return True return False except Exception as e: logger.error(f"❌ 添加关系失败: {e} " ) return False def find_related_entities ( self, entity_id: str , relationship_types: List [str ] = None , max_depth: int = 2 , limit: int = 50 ) -> List [Dict [str , Any ]]: """ 查找相关实体 Args: entity_id: 起始实体ID relationship_types: 关系类型过滤 max_depth: 最大搜索深度 limit: 结果限制 Returns: List[Dict]: 相关实体列表 """ try : rel_filter = "" if relationship_types: rel_types = "|" .join(relationship_types) rel_filter = f":{rel_types} " query = f""" MATCH path = (start:Entity {{id: $entity_id}})-[r{rel_filter} *1..{max_depth} ]-(related:Entity) WHERE start.id <> related.id RETURN DISTINCT related, length(path) as distance, [rel in relationships(path) | type(rel)] as relationship_path ORDER BY distance, related.name LIMIT $limit """ with self.driver.session(database=self.database) as session: result = session.run(query, entity_id=entity_id, limit=limit) entities = [] for record in result: entity_data = dict (record["related" ]) entity_data["distance" ] = record["distance" ] entity_data["relationship_path" ] = record["relationship_path" ] entities.append(entity_data) logger.debug(f"🔍 找到 {len (entities)} 个相关实体" ) return entities except Exception as e: logger.error(f"❌ 查找相关实体失败: {e} " ) return [] def search_entities_by_name (self, name_pattern: str , entity_types: List [str ] = None , limit: int = 20 ) -> List [Dict [str , Any ]]: """ 按名称搜索实体 Args: name_pattern: 名称模式 (支持部分匹配) entity_types: 实体类型过滤 limit: 结果限制 Returns: List[Dict]: 匹配的实体列表 """ try : type_filter = "" params = {"pattern" : f".*{name_pattern} .*" , "limit" : limit} if entity_types: type_filter = "AND e.type IN $types" params["types" ] = entity_types query = f""" MATCH (e:Entity) WHERE e.name =~ $pattern {type_filter} RETURN e ORDER BY e.name LIMIT $limit """ with self.driver.session(database=self.database) as session: result = session.run(query, **params) entities = [] for record in result: entity_data = dict (record["e" ]) entities.append(entity_data) logger.debug(f"🔍 按名称搜索到 {len (entities)} 个实体" ) return entities except Exception as e: logger.error(f"❌ 按名称搜索实体失败: {e} " ) return [] def get_entity_relationships (self, entity_id: str ) -> List [Dict [str , Any ]]: """ 获取实体的所有关系 Args: entity_id: 实体ID Returns: List[Dict]: 关系列表 """ try : query = """ MATCH (e:Entity {id: $entity_id})-[r]-(other:Entity) RETURN r, other, CASE WHEN startNode(r).id = $entity_id THEN 'outgoing' ELSE 'incoming' END as direction """ with self.driver.session(database=self.database) as session: result = session.run(query, entity_id=entity_id) relationships = [] for record in result: rel_data = dict (record["r" ]) other_data = dict (record["other" ]) relationship = { "relationship" : rel_data, "other_entity" : other_data, "direction" : record["direction" ] } relationships.append(relationship) return relationships except Exception as e: logger.error(f"❌ 获取实体关系失败: {e} " ) return [] def delete_entity (self, entity_id: str ) -> bool : """ 删除实体及其所有关系 Args: entity_id: 实体ID Returns: bool: 是否成功 """ try : query = """ MATCH (e:Entity {id: $entity_id}) DETACH DELETE e """ with self.driver.session(database=self.database) as session: result = session.run(query, entity_id=entity_id) summary = result.consume() deleted_count = summary.counters.nodes_deleted logger.info(f"✅ 删除实体: {entity_id} (删除 {deleted_count} 个节点)" ) return deleted_count > 0 except Exception as e: logger.error(f"❌ 删除实体失败: {e} " ) return False def clear_all (self ) -> bool : """ 清空所有数据 Returns: bool: 是否成功 """ try : query = "MATCH (n) DETACH DELETE n" with self.driver.session(database=self.database) as session: result = session.run(query) summary = result.consume() deleted_nodes = summary.counters.nodes_deleted deleted_relationships = summary.counters.relationships_deleted logger.info(f"✅ 清空Neo4j数据库: 删除 {deleted_nodes} 个节点, {deleted_relationships} 个关系" ) return True except Exception as e: logger.error(f"❌ 清空数据库失败: {e} " ) return False def get_stats (self ) -> Dict [str , Any ]: """ 获取图数据库统计信息 Returns: Dict: 统计信息 """ try : queries = { "total_nodes" : "MATCH (n) RETURN count(n) as count" , "total_relationships" : "MATCH ()-[r]->() RETURN count(r) as count" , "entity_nodes" : "MATCH (n:Entity) RETURN count(n) as count" , "memory_nodes" : "MATCH (n:Memory) RETURN count(n) as count" , } stats = {} with self.driver.session(database=self.database) as session: for key, query in queries.items(): result = session.run(query) record = result.single() stats[key] = record["count" ] if record else 0 return stats except Exception as e: logger.error(f"❌ 获取统计信息失败: {e} " ) return {} def health_check (self ) -> bool : """ 健康检查 Returns: bool: 服务是否健康 """ try : with self.driver.session(database=self.database) as session: result = session.run("RETURN 1 as health" ) record = result.single() return record["health" ] == 1 except Exception as e: logger.error(f"❌ Neo4j健康检查失败: {e} " ) return False def __del__ (self ): """析构函数,清理资源""" if hasattr (self, 'driver' ) and self.driver: try : self.driver.close() except : pass
qdrant_store.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 """ Qdrant向量数据库存储实现 使用专业的Qdrant向量数据库替代ChromaDB """ import loggingimport osimport uuidimport threadingfrom typing import Dict , List , Optional , Any , Union import numpy as npfrom datetime import datetimetry : from qdrant_client import QdrantClient from qdrant_client.http import models from qdrant_client.http.models import ( Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue, SearchRequest ) QDRANT_AVAILABLE = True except ImportError: QDRANT_AVAILABLE = False QdrantClient = None models = None logger = logging.getLogger(__name__) class QdrantConnectionManager : """Qdrant连接管理器 - 防止重复连接和初始化""" _instances = {} _lock = threading.Lock() @classmethod def get_instance ( cls, url: Optional [str ] = None , api_key: Optional [str ] = None , collection_name: str = "hello_agents_vectors" , vector_size: int = 384 , distance: str = "cosine" , timeout: int = 30 , **kwargs ) -> 'QdrantVectorStore' : """获取或创建Qdrant实例(单例模式)""" key = (url or "local" , collection_name) if key not in cls._instances: with cls._lock: if key not in cls._instances: logger.debug(f"🔄 创建新的Qdrant连接: {collection_name} " ) cls._instances[key] = QdrantVectorStore( url=url, api_key=api_key, collection_name=collection_name, vector_size=vector_size, distance=distance, timeout=timeout, **kwargs ) else : logger.debug(f"♻️ 复用现有Qdrant连接: {collection_name} " ) else : logger.debug(f"♻️ 复用现有Qdrant连接: {collection_name} " ) return cls._instances[key] class QdrantVectorStore : """Qdrant向量数据库存储实现""" def __init__ ( self, url: Optional [str ] = None , api_key: Optional [str ] = None , collection_name: str = "hello_agents_vectors" , vector_size: int = 384 , distance: str = "cosine" , timeout: int = 30 , **kwargs ): """ 初始化Qdrant向量存储 (支持云API) Args: url: Qdrant云服务URL (如果为None则使用本地) api_key: Qdrant云服务API密钥 collection_name: 集合名称 vector_size: 向量维度 distance: 距离度量方式 (cosine, dot, euclidean) timeout: 连接超时时间 """ if not QDRANT_AVAILABLE: raise ImportError( "qdrant-client未安装。请运行: pip install qdrant-client>=1.6.0" ) self.url = url self.api_key = api_key self.collection_name = collection_name self.vector_size = vector_size self.timeout = timeout try : self.hnsw_m = int (os.getenv("QDRANT_HNSW_M" , "32" )) except Exception: self.hnsw_m = 32 try : self.hnsw_ef_construct = int (os.getenv("QDRANT_HNSW_EF_CONSTRUCT" , "256" )) except Exception: self.hnsw_ef_construct = 256 try : self.search_ef = int (os.getenv("QDRANT_SEARCH_EF" , "128" )) except Exception: self.search_ef = 128 self.search_exact = os.getenv("QDRANT_SEARCH_EXACT" , "0" ) == "1" distance_map = { "cosine" : Distance.COSINE, "dot" : Distance.DOT, "euclidean" : Distance.EUCLID, } self.distance = distance_map.get(distance.lower(), Distance.COSINE) self.client = None self._initialize_client() def _initialize_client (self ): """初始化Qdrant客户端和集合""" try : if self.url and self.api_key: self.client = QdrantClient( url=self.url, api_key=self.api_key, timeout=self.timeout ) logger.info(f"✅ 成功连接到Qdrant云服务: {self.url} " ) elif self.url: self.client = QdrantClient( url=self.url, timeout=self.timeout ) logger.info(f"✅ 成功连接到Qdrant服务: {self.url} " ) else : self.client = QdrantClient( host="localhost" , port=6333 , timeout=self.timeout ) logger.info("✅ 成功连接到本地Qdrant服务: localhost:6333" ) collections = self.client.get_collections() self._ensure_collection() except Exception as e: logger.error(f"❌ Qdrant连接失败: {e} " ) if not self.url: logger.info("💡 本地连接失败,可以考虑使用Qdrant云服务" ) logger.info("💡 或启动本地服务: docker run -p 6333:6333 qdrant/qdrant" ) else : logger.info("💡 请检查URL和API密钥是否正确" ) raise def _ensure_collection (self ): """确保集合存在,不存在则创建""" try : collections = self.client.get_collections().collections collection_names = [c.name for c in collections] if self.collection_name not in collection_names: hnsw_cfg = None try : hnsw_cfg = models.HnswConfigDiff(m=self.hnsw_m, ef_construct=self.hnsw_ef_construct) except Exception: hnsw_cfg = None self.client.create_collection( collection_name=self.collection_name, vectors_config=VectorParams( size=self.vector_size, distance=self.distance ), hnsw_config=hnsw_cfg ) logger.info(f"✅ 创建Qdrant集合: {self.collection_name} " ) else : logger.info(f"✅ 使用现有Qdrant集合: {self.collection_name} " ) try : self.client.update_collection( collection_name=self.collection_name, hnsw_config=models.HnswConfigDiff(m=self.hnsw_m, ef_construct=self.hnsw_ef_construct) ) except Exception as ie: logger.debug(f"跳过更新HNSW配置: {ie} " ) self._ensure_payload_indexes() except Exception as e: logger.error(f"❌ 集合初始化失败: {e} " ) raise def _ensure_payload_indexes (self ): """为常用过滤字段创建payload索引""" try : index_fields = [ ("memory_type" , models.PayloadSchemaType.KEYWORD), ("user_id" , models.PayloadSchemaType.KEYWORD), ("memory_id" , models.PayloadSchemaType.KEYWORD), ("timestamp" , models.PayloadSchemaType.INTEGER), ("modality" , models.PayloadSchemaType.KEYWORD), ("source" , models.PayloadSchemaType.KEYWORD), ("external" , models.PayloadSchemaType.BOOL), ("namespace" , models.PayloadSchemaType.KEYWORD), ("is_rag_data" , models.PayloadSchemaType.BOOL), ("rag_namespace" , models.PayloadSchemaType.KEYWORD), ("data_source" , models.PayloadSchemaType.KEYWORD), ] for field_name, schema_type in index_fields: try : self.client.create_payload_index( collection_name=self.collection_name, field_name=field_name, field_schema=schema_type, ) except Exception as ie: logger.debug(f"索引 {field_name} 已存在或创建失败: {ie} " ) except Exception as e: logger.debug(f"创建payload索引时出错: {e} " ) def add_vectors ( self, vectors: List [List [float ]], metadata: List [Dict [str , Any ]], ids: Optional [List [str ]] = None ) -> bool : """ 添加向量到Qdrant Args: vectors: 向量列表 metadata: 元数据列表 ids: 可选的ID列表 Returns: bool: 是否成功 """ try : if not vectors: logger.warning("⚠️ 向量列表为空" ) return False if ids is None : ids = [f"vec_{i} _{int (datetime.now().timestamp() * 1000000 )} " for i in range (len (vectors))] logger.info(f"[Qdrant] add_vectors start: n_vectors={len (vectors)} n_meta={len (metadata)} collection={self.collection_name} " ) points = [] for i, (vector, meta, point_id) in enumerate (zip (vectors, metadata, ids)): try : vlen = len (vector) except Exception: logger.error(f"[Qdrant] 非法向量类型: index={i} type={type (vector)} value={vector} " ) continue if vlen != self.vector_size: logger.warning(f"⚠️ 向量维度不匹配: 期望{self.vector_size} , 实际{len (vector)} " ) continue meta_with_timestamp = meta.copy() meta_with_timestamp["timestamp" ] = int (datetime.now().timestamp()) meta_with_timestamp["added_at" ] = int (datetime.now().timestamp()) if "external" in meta_with_timestamp and not isinstance (meta_with_timestamp.get("external" ), bool ): val = meta_with_timestamp.get("external" ) meta_with_timestamp["external" ] = True if str (val).lower() in ("1" , "true" , "yes" ) else False safe_id: Any if isinstance (point_id, int ): safe_id = point_id elif isinstance (point_id, str ): try : uuid.UUID(point_id) safe_id = point_id except Exception: safe_id = str (uuid.uuid4()) else : safe_id = str (uuid.uuid4()) point = PointStruct( id =safe_id, vector=vector, payload=meta_with_timestamp ) points.append(point) if not points: logger.warning("⚠️ 没有有效的向量点" ) return False logger.info(f"[Qdrant] upsert begin: points={len (points)} " ) operation_info = self.client.upsert( collection_name=self.collection_name, points=points, wait=True ) logger.info("[Qdrant] upsert done" ) logger.info(f"✅ 成功添加 {len (points)} 个向量到Qdrant" ) return True except Exception as e: logger.error(f"❌ 添加向量失败: {e} " ) return False def search_similar ( self, query_vector: List [float ], limit: int = 10 , score_threshold: Optional [float ] = None , where: Optional [Dict [str , Any ]] = None ) -> List [Dict [str , Any ]]: """ 搜索相似向量 Args: query_vector: 查询向量 limit: 返回结果数量限制 score_threshold: 相似度阈值 where: 过滤条件 Returns: List[Dict]: 搜索结果 """ try : if len (query_vector) != self.vector_size: logger.error(f"❌ 查询向量维度错误: 期望{self.vector_size} , 实际{len (query_vector)} " ) return [] query_filter = None if where: conditions = [] for key, value in where.items(): if isinstance (value, (str , int , float , bool )): conditions.append( FieldCondition( key=key, match =MatchValue(value=value) ) ) if conditions: query_filter = Filter(must=conditions) search_params = None try : search_params = models.SearchParams(hnsw_ef=self.search_ef, exact=self.search_exact) except Exception: search_params = None try : response = self.client.query_points( collection_name=self.collection_name, query=query_vector, query_filter=query_filter, limit=limit, score_threshold=score_threshold, with_payload=True , with_vectors=False , search_params=search_params ) search_result = response.points except AttributeError: search_result = self.client.search( collection_name=self.collection_name, query_vector=query_vector, query_filter=query_filter, limit=limit, score_threshold=score_threshold, with_payload=True , with_vectors=False , search_params=search_params ) results = [] for hit in search_result: result = { "id" : hit.id , "score" : hit.score, "metadata" : hit.payload or {} } results.append(result) logger.debug(f"🔍 Qdrant搜索返回 {len (results)} 个结果" ) return results except Exception as e: logger.error(f"❌ 向量搜索失败: {e} " ) return [] def delete_vectors (self, ids: List [str ] ) -> bool : """ 删除向量 Args: ids: 要删除的向量ID列表 Returns: bool: 是否成功 """ try : if not ids: return True operation_info = self.client.delete( collection_name=self.collection_name, points_selector=models.PointIdsList( points=ids ), wait=True ) logger.info(f"✅ 成功删除 {len (ids)} 个向量" ) return True except Exception as e: logger.error(f"❌ 删除向量失败: {e} " ) return False def clear_collection (self ) -> bool : """ 清空集合 Returns: bool: 是否成功 """ try : self.client.delete_collection(collection_name=self.collection_name) self._ensure_collection() logger.info(f"✅ 成功清空Qdrant集合: {self.collection_name} " ) return True except Exception as e: logger.error(f"❌ 清空集合失败: {e} " ) return False def delete_memories (self, memory_ids: List [str ] ): """ 删除指定记忆(通过payload中的 memory_id 过滤删除) 注意:由于写入时可能将非UUID的点ID转换为UUID,这里不再依赖点ID, 而是通过payload中的memory_id来匹配删除,确保一致性。 """ try : if not memory_ids: return conditions = [ FieldCondition(key="memory_id" , match =MatchValue(value=mid)) for mid in memory_ids ] query_filter = Filter(should=conditions) self.client.delete( collection_name=self.collection_name, points_selector=models.FilterSelector(filter =query_filter), wait=True , ) logger.info(f"✅ 成功按memory_id删除 {len (memory_ids)} 个Qdrant向量" ) except Exception as e: logger.error(f"❌ 删除记忆失败: {e} " ) raise def get_collection_info (self ) -> Dict [str , Any ]: """ 获取集合信息 Returns: Dict: 集合信息 """ try : collection_info = self.client.get_collection(self.collection_name) count=0 if hasattr (collection_info, "points_count" ): count=collection_info.points_count elif hasattr (collection_info, "vectors_count" ): count=collection_info.vectors_count info = { "name" : self.collection_name, "points_count" : collection_info.points_count, "segments_count" : collection_info.segments_count, "indexed_vectors_count" : collection_info.indexed_vectors_count, "config" : { "vector_size" : self.vector_size, "distance" : self.distance.value, } } return info except Exception as e: logger.error(f"❌ 获取集合信息失败: {e} " ) return {} def get_collection_stats (self ) -> Dict [str , Any ]: """ 获取集合统计信息(兼容抽象接口) """ info = self.get_collection_info() if not info: return {"store_type" : "qdrant" , "name" : self.collection_name} info["store_type" ] = "qdrant" return info def health_check (self ) -> bool : """ 健康检查 Returns: bool: 服务是否健康 """ try : collections = self.client.get_collections() return True except Exception as e: logger.error(f"❌ Qdrant健康检查失败: {e} " ) return False def __del__ (self ): """析构函数,清理资源""" if hasattr (self, 'client' ) and self.client: try : self.client.close() except : pass
rag pipeline.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 from typing import List , Dict , Optional , Any import osimport hashlibfrom langdetect import detectfrom sentence_transformers import CrossEncoderfrom markitdown import MarkItDownfrom storage.embedding import get_dimension, get_text_embedderfrom storage.qdrant_store import QdrantVectorStorefrom core.llm import LLMdef _get_markitdown_instance (): """ Get a configured MarkItDown instance for document conversion. """ try : return MarkItDown() except ImportError: print ("[WARNING] MarkItDown not available. Install with: pip install markitdown" ) return None def _is_markitdown_supported_format (path: str ) -> bool : """ Check if the file format is supported by MarkItDown. Supports: PDF, Office docs (docx, xlsx, pptx), images (jpg, png, gif, bmp, tiff), audio (mp3, wav, m4a), HTML, text formats (txt, md, csv, json, xml), ZIP files, etc. """ ext = (os.path.splitext(path)[1 ] or '' ).lower() supported_formats = { '.pdf' , '.doc' , '.docx' , '.xls' , '.xlsx' , '.ppt' , '.pptx' , '.txt' , '.md' , '.csv' , '.json' , '.xml' , '.html' , '.htm' , '.jpg' , '.jpeg' , '.png' , '.gif' , '.bmp' , '.tiff' , '.tif' , '.webp' , '.mp3' , '.wav' , '.m4a' , '.aac' , '.flac' , '.ogg' , '.zip' , '.tar' , '.gz' , '.rar' , '.py' , '.js' , '.ts' , '.java' , '.cpp' , '.c' , '.h' , '.css' , '.scss' , '.log' , '.conf' , '.ini' , '.cfg' , '.yaml' , '.yml' , '.toml' } return ext in supported_formats def _convert_to_markdown (path: str ) -> str : """ Universal document reader using MarkItDown with enhanced PDF processing. Converts any supported file format to markdown text. """ if not os.path.exists(path): return "" ext = (os.path.splitext(path)[1 ] or '' ).lower() if ext == '.pdf' : return _enhanced_pdf_processing(path) md_instance = _get_markitdown_instance() if md_instance is None : return _fallback_text_reader(path) try : result = md_instance.convert(path) text = getattr (result, "text_content" , None ) if isinstance (text, str ) and text.strip(): return text return "" except Exception as e: print (f"[WARNING] MarkItDown failed for {path} : {e} " ) return _fallback_text_reader(path) def _enhanced_pdf_processing (path: str ) -> str : """ Enhanced PDF processing with post-processing cleanup. """ print (f"[RAG] Using enhanced PDF processing for: {path} " ) md_instance = _get_markitdown_instance() if md_instance is None : return _fallback_text_reader(path) try : result = md_instance.convert(path) raw_text = getattr (result, "text_content" , None ) if not raw_text or not raw_text.strip(): return "" cleaned_text = _post_process_pdf_text(raw_text) print (f"[RAG] PDF post-processing completed: {len (raw_text)} -> {len (cleaned_text)} chars" ) return cleaned_text except Exception as e: print (f"[WARNING] Enhanced PDF processing failed for {path} : {e} " ) return _fallback_text_reader(path) def _post_process_pdf_text (text: str ) -> str : """ Post-process PDF text to improve quality. """ import re lines = text.splitlines() cleaned_lines = [] for line in lines: line = line.strip() if not line: continue if len (line) <= 2 and not line.isdigit(): continue if re.match (r'^\d+$' , line): continue if line.lower() in ['github' , 'project' , 'forks' , 'stars' , 'language' ]: continue cleaned_lines.append(line) merged_lines = [] i = 0 while i < len (cleaned_lines): current_line = cleaned_lines[i] if len (current_line) < 60 and i + 1 < len (cleaned_lines): next_line = cleaned_lines[i + 1 ] if (not current_line.endswith(':' ) and not current_line.endswith(':' ) and not current_line.startswith('#' ) and not next_line.startswith('#' ) and len (next_line) < 120 ): merged_line = current_line + " " + next_line merged_lines.append(merged_line) i += 2 continue merged_lines.append(current_line) i += 1 paragraphs = [] current_paragraph = [] for line in merged_lines: if (line.startswith('#' ) or line.endswith(':' ) or line.endswith(':' ) or len (line) > 150 or not current_paragraph): if current_paragraph: paragraphs.append(' ' .join(current_paragraph)) current_paragraph = [] paragraphs.append(line) else : current_paragraph.append(line) if current_paragraph: paragraphs.append(' ' .join(current_paragraph)) return '\n\n' .join(paragraphs) def _fallback_text_reader (path: str ) -> str : """ Simple fallback reader for basic text files when MarkItDown is unavailable. """ try : with open (path, 'r' , encoding='utf-8' , errors='ignore' ) as f: return f.read() except Exception: try : with open (path, 'r' , encoding='latin-1' , errors='ignore' ) as f: return f.read() except Exception: return "" def _detect_lang (sample: str ) -> str : try : return detect(sample[:1000 ]) if sample else "unknown" except Exception: return "unknown" def _is_cjk (ch: str ) -> bool : code = ord (ch) return ( 0x4E00 <= code <= 0x9FFF or 0x3400 <= code <= 0x4DBF or 0x20000 <= code <= 0x2A6DF or 0x2A700 <= code <= 0x2B73F or 0x2B740 <= code <= 0x2B81F or 0x2B820 <= code <= 0x2CEAF or 0xF900 <= code <= 0xFAFF ) def _approx_token_len (text: str ) -> int : cjk = sum (1 for ch in text if _is_cjk(ch)) non_cjk_tokens = len ([t for t in text.split() if t]) return cjk + non_cjk_tokens def _split_paragraphs_with_headings (text: str ) -> List [Dict ]: lines = text.splitlines() heading_stack: List [str ] = [] paragraphs: List [Dict ] = [] buf: List [str ] = [] char_pos = 0 def flush_buf (end_pos: int ): if not buf: return content = "\n" .join(buf).strip() if not content: return paragraphs.append({ "content" : content, "heading_path" : " > " .join(heading_stack) if heading_stack else None , "start" : max (0 , end_pos - len (content)), "end" : end_pos, }) for ln in lines: raw = ln if raw.strip().startswith("#" ): flush_buf(char_pos) level = len (raw) - len (raw.lstrip('#' )) title = raw.lstrip('#' ).strip() if level <= 0 : level = 1 if level <= len (heading_stack): heading_stack = heading_stack[:level-1 ] heading_stack.append(title) char_pos += len (raw) + 1 continue if raw.strip() == "" : flush_buf(char_pos) buf = [] else : buf.append(raw) char_pos += len (raw) + 1 flush_buf(char_pos) if not paragraphs: paragraphs = [{"content" : text, "heading_path" : None , "start" : 0 , "end" : len (text)}] return paragraphs def _chunk_paragraphs (paragraphs: List [Dict ], chunk_tokens: int , overlap_tokens: int ) -> List [Dict ]: chunks: List [Dict ] = [] cur: List [Dict ] = [] cur_tokens = 0 i = 0 while i < len (paragraphs): p = paragraphs[i] p_tokens = _approx_token_len(p["content" ]) or 1 if cur_tokens + p_tokens <= chunk_tokens or not cur: cur.append(p) cur_tokens += p_tokens i += 1 else : content = "\n\n" .join(x["content" ] for x in cur) start = cur[0 ]["start" ] end = cur[-1 ]["end" ] heading_path = next ((x["heading_path" ] for x in reversed (cur) if x.get("heading_path" )), None ) chunks.append({ "content" : content, "start" : start, "end" : end, "heading_path" : heading_path, }) if overlap_tokens > 0 and cur: kept: List [Dict ] = [] kept_tokens = 0 for x in reversed (cur): t = _approx_token_len(x["content" ]) or 1 if kept_tokens + t > overlap_tokens: break kept.append(x) kept_tokens += t cur = list (reversed (kept)) cur_tokens = kept_tokens else : cur = [] cur_tokens = 0 if cur: content = "\n\n" .join(x["content" ] for x in cur) start = cur[0 ]["start" ] end = cur[-1 ]["end" ] heading_path = next ((x["heading_path" ] for x in reversed (cur) if x.get("heading_path" )), None ) chunks.append({ "content" : content, "start" : start, "end" : end, "heading_path" : heading_path, }) return chunks def load_and_chunk_texts (paths: List [str ], chunk_size: int = 800 , chunk_overlap: int = 100 , namespace: Optional [str ] = None , source_label: str = "rag" ) -> List [Dict ]: """ Universal document loader and chunker using MarkItDown. Converts all supported formats to markdown, then chunks intelligently. """ print (f"[RAG] Universal loader start: files={len (paths)} chunk_size={chunk_size} overlap={chunk_overlap} ns={namespace or 'default' } " ) chunks: List [Dict ] = [] seen_hashes = set () for path in paths: if not os.path.exists(path): print (f"[WARNING] File not found: {path} " ) continue print (f"[RAG] Processing: {path} " ) ext = (os.path.splitext(path)[1 ] or '' ).lower() markdown_text = _convert_to_markdown(path) if not markdown_text.strip(): print (f"[WARNING] No content extracted from: {path} " ) continue lang = _detect_lang(markdown_text) doc_id = hashlib.md5(f"{path} |{len (markdown_text)} " .encode('utf-8' )).hexdigest() para = _split_paragraphs_with_headings(markdown_text) token_chunks = _chunk_paragraphs(para, chunk_tokens=max (1 , chunk_size), overlap_tokens=max (0 , chunk_overlap)) for ch in token_chunks: content = ch["content" ] start = ch.get("start" , 0 ) end = ch.get("end" , start + len (content)) norm = content.strip() if not norm: continue content_hash = hashlib.md5(norm.encode('utf-8' )).hexdigest() if content_hash in seen_hashes: continue seen_hashes.add(content_hash) chunk_id = hashlib.md5(f"{doc_id} |{start} |{end} |{content_hash} " .encode('utf-8' )).hexdigest() chunks.append({ "id" : chunk_id, "content" : content, "metadata" : { "source_path" : path, "file_ext" : ext, "doc_id" : doc_id, "lang" : lang, "start" : start, "end" : end, "content_hash" : content_hash, "namespace" : namespace or "default" , "source" : source_label, "external" : True , "heading_path" : ch.get("heading_path" ), "format" : "markdown" , }, }) print (f"[RAG] Universal loader done: total_chunks={len (chunks)} " ) return chunks def build_graph_from_chunks (neo4j, chunks: List [Dict ] ) -> None : created_docs = set () for ch in chunks: mem_id = ch["id" ] meta = ch.get("metadata" , {}) source_path = meta.get("source_path" ) doc_id = meta.get("doc_id" ) if doc_id and doc_id not in created_docs: created_docs.add(doc_id) try : neo4j.add_entity( entity_id=doc_id, name=os.path.basename(source_path or doc_id), entity_type="Document" , properties={"source_path" : source_path, "lang" : meta.get("lang" )} ) except Exception: pass try : neo4j.add_entity(entity_id=mem_id, name=mem_id, entity_type="Memory" , properties={ "source_path" : source_path, "doc_id" : doc_id, "start" : meta.get("start" ), "end" : meta.get("end" ), }) except Exception: pass if doc_id: try : neo4j.add_relationship(from_id=doc_id, to_id=mem_id, rel_type="HAS_CHUNK" , properties={}) except Exception: pass def _preprocess_markdown_for_embedding (text: str ) -> str : """ Preprocess markdown text for better embedding quality. Removes excessive markup while preserving semantic content. """ import re text = re.sub(r'^#{1,6}\s+' , '' , text, flags=re.MULTILINE) text = re.sub(r'\[([^\]]+)\]\([^)]+\)' , r'\1' , text) text = re.sub(r'\*\*([^*]+)\*\*' , r'\1' , text) text = re.sub(r'\*([^*]+)\*' , r'\1' , text) text = re.sub(r'`([^`]+)`' , r'\1' , text) text = re.sub(r'```[^\n]*\n([\s\S]*?)```' , r'\1' , text) text = re.sub(r'\n\s*\n' , '\n\n' , text) text = re.sub(r'[ \t]+' , ' ' , text) return text.strip() def _create_default_vector_store (dimension: int = None ) -> QdrantVectorStore: """ Create default Qdrant vector store with RAG-optimized settings. 使用连接管理器避免重复连接。 """ if dimension is None : dimension = get_dimension(384 ) qdrant_url = os.getenv("QDRANT_URL" ) qdrant_api_key = os.getenv("QDRANT_API_KEY" ) from ..storage.qdrant_store import QdrantConnectionManager return QdrantConnectionManager.get_instance( url=qdrant_url, api_key=qdrant_api_key, collection_name="hello_agents_rag_vectors" , vector_size=dimension, distance="cosine" ) def index_chunks ( store = None , chunks: List [Dict ] = None , cache_db: Optional [str ] = None , batch_size: int = 64 , rag_namespace: str = "default" ) -> None : """ Index markdown chunks with unified embedding and Qdrant storage. Uses百炼 API with fallback to sentence-transformers. """ if not chunks: print ("[RAG] No chunks to index" ) return embedder = get_text_embedder() dimension = get_dimension(384 ) if store is None : store = _create_default_vector_store(dimension) print (f"[RAG] Created default Qdrant store with dimension {dimension} " ) processed_texts = [] for c in chunks: raw_content = c["content" ] processed_content = _preprocess_markdown_for_embedding(raw_content) processed_texts.append(processed_content) print (f"[RAG] Embedding start: total_texts={len (processed_texts)} batch_size={batch_size} " ) vecs: List [List [float ]] = [] for i in range (0 , len (processed_texts), batch_size): part = processed_texts[i:i+batch_size] try : part_vecs = embedder.encode(part) if not isinstance (part_vecs, list ): if hasattr (part_vecs, "tolist" ): part_vecs = [part_vecs.tolist()] else : part_vecs = [list (part_vecs)] else : if part_vecs and not isinstance (part_vecs[0 ], (list , tuple )) and hasattr (part_vecs[0 ], "__len__" ): normalized_vecs = [] for v in part_vecs: if hasattr (v, "tolist" ): normalized_vecs.append(v.tolist()) else : normalized_vecs.append(list (v)) part_vecs = normalized_vecs elif part_vecs and not isinstance (part_vecs[0 ], (list , tuple )): if hasattr (part_vecs, "tolist" ): part_vecs = [part_vecs.tolist()] else : part_vecs = [list (part_vecs)] for v in part_vecs: try : if hasattr (v, "tolist" ): v = v.tolist() v_norm = [float (x) for x in v] if len (v_norm) != dimension: print (f"[WARNING] 向量维度异常: 期望{dimension} , 实际{len (v_norm)} " ) if len (v_norm) < dimension: v_norm.extend([0.0 ] * (dimension - len (v_norm))) else : v_norm = v_norm[:dimension] vecs.append(v_norm) except Exception as e: print (f"[WARNING] 向量转换失败: {e} , 使用零向量" ) vecs.append([0.0 ] * dimension) except Exception as e: print (f"[WARNING] Batch {i} encoding failed: {e} " ) print (f"[RAG] Retrying batch {i} with smaller chunks..." ) success = False for j in range (0 , len (part), 8 ): small_part = part[j:j+8 ] try : import time time.sleep(2 ) small_vecs = embedder.encode(small_part) if isinstance (small_vecs, list ) and small_vecs and not isinstance (small_vecs[0 ], list ): small_vecs = [small_vecs] for v in small_vecs: if hasattr (v, "tolist" ): v = v.tolist() try : v_norm = [float (x) for x in v] if len (v_norm) != dimension: print (f"[WARNING] 向量维度异常: 期望{dimension} , 实际{len (v_norm)} " ) if len (v_norm) < dimension: v_norm.extend([0.0 ] * (dimension - len (v_norm))) else : v_norm = v_norm[:dimension] vecs.append(v_norm) success = True except Exception as e2: print (f"[WARNING] 小批次向量转换失败: {e2} " ) vecs.append([0.0 ] * dimension) except Exception as e2: print (f"[WARNING] 小批次 {j//8 } 仍然失败: {e2} " ) for _ in range (len (small_part)): vecs.append([0.0 ] * dimension) if not success: print (f"[ERROR] 批次 {i} 完全失败,使用零向量" ) print (f"[RAG] Embedding progress: {min (i+batch_size, len (processed_texts))} /{len (processed_texts)} " ) metas: List [Dict ] = [] ids: List [str ] = [] for ch in chunks: meta = { "memory_id" : ch["id" ], "user_id" : "rag_user" , "memory_type" : "rag_chunk" , "content" : ch["content" ], "data_source" : "rag_pipeline" , "rag_namespace" : rag_namespace, "is_rag_data" : True , } meta.update(ch.get("metadata" , {})) metas.append(meta) ids.append(ch["id" ]) print (f"[RAG] Qdrant upsert start: n={len (vecs)} " ) success = store.add_vectors(vectors=vecs, metadata=metas, ids=ids) if success: print (f"[RAG] Qdrant upsert done: {len (vecs)} vectors indexed" ) else : print (f"[RAG] Qdrant upsert failed" ) raise RuntimeError("Failed to index vectors to Qdrant" ) def embed_query (query: str ) -> List [float ]: """ Embed query using unified embedding (百炼 with fallback). """ embedder = get_text_embedder() dimension = get_dimension(384 ) try : vec = embedder.encode(query) if hasattr (vec, "tolist" ): vec = vec.tolist() if isinstance (vec, list ) and vec and isinstance (vec[0 ], (list , tuple )): vec = vec[0 ] result = [float (x) for x in vec] if len (result) != dimension: print (f"[WARNING] Query向量维度异常: 期望{dimension} , 实际{len (result)} " ) if len (result) < dimension: result.extend([0.0 ] * (dimension - len (result))) else : result = result[:dimension] return result except Exception as e: print (f"[WARNING] Query embedding failed: {e} " ) return [0.0 ] * dimension def search_vectors ( store = None , query: str = "" , top_k: int = 8 , rag_namespace: Optional [str ] = None , only_rag_data: bool = True , score_threshold: Optional [float ] = None ) -> List [Dict ]: """ Search RAG vectors using unified embedding and Qdrant. """ if not query: return [] if store is None : store = _create_default_vector_store() qv = embed_query(query) where = {"memory_type" : "rag_chunk" } if only_rag_data: where["is_rag_data" ] = True where["data_source" ] = "rag_pipeline" if rag_namespace: where["rag_namespace" ] = rag_namespace try : return store.search_similar( query_vector=qv, limit=top_k, score_threshold=score_threshold, where=where ) except Exception as e: print (f"[WARNING] RAG search failed: {e} " ) return [] def _prompt_mqe (query: str , n: int ) -> List [str ]: try : llm = LLM() prompt = [ {"role" : "system" , "content" : "你是检索查询扩展助手。生成语义等价或互补的多样化查询。使用中文,简短,避免标点。" }, {"role" : "user" , "content" : f"原始查询:{query} \n请给出{n} 个不同表述的查询,每行一个。" } ] text = llm.think(prompt) lines = [ln.strip("- \t" ) for ln in (text or "" ).splitlines()] outs = [ln for ln in lines if ln] return outs[:n] or [query] except Exception: return [query] def _prompt_hyde (query: str ) -> Optional [str ]: try : llm = LLM() prompt = [ {"role" : "system" , "content" : "根据用户问题,先写一段可能的答案性段落,用于向量检索的查询文档(不要分析过程)。" }, {"role" : "user" , "content" : f"问题:{query} \n请直接写一段中等长度、客观、包含关键术语的段落。" } ] return llm.think(prompt) except Exception: return None def search_vectors_expanded ( store = None , query: str = "" , top_k: int = 8 , rag_namespace: Optional [str ] = None , only_rag_data: bool = True , score_threshold: Optional [float ] = None , enable_mqe: bool = False , mqe_expansions: int = 2 , enable_hyde: bool = False , candidate_pool_multiplier: int = 4 , ) -> List [Dict ]: """ Search with query expansion using unified embedding and Qdrant. """ if not query: return [] if store is None : store = _create_default_vector_store() expansions: List [str ] = [query] if enable_mqe and mqe_expansions > 0 : expansions.extend(_prompt_mqe(query, mqe_expansions)) if enable_hyde: hyde_text = _prompt_hyde(query) if hyde_text: expansions.append(hyde_text) uniq: List [str ] = [] for e in expansions: if e and e not in uniq: uniq.append(e) expansions = uniq[: max (1 , len (uniq))] pool = max (top_k * candidate_pool_multiplier, 20 ) per = max (1 , pool // max (1 , len (expansions))) where = {"memory_type" : "rag_chunk" } if only_rag_data: where["is_rag_data" ] = True where["data_source" ] = "rag_pipeline" if rag_namespace: where["rag_namespace" ] = rag_namespace agg: Dict [str , Dict ] = {} for q in expansions: qv = embed_query(q) hits = store.search_similar(query_vector=qv, limit=per, score_threshold=score_threshold, where=where) for h in hits: mid = h.get("metadata" , {}).get("memory_id" , h.get("id" )) s = float (h.get("score" , 0.0 )) if mid not in agg or s > float (agg[mid].get("score" , 0.0 )): agg[mid] = h merged = list (agg.values()) merged.sort(key=lambda x: float (x.get("score" , 0.0 )), reverse=True ) return merged[:top_k] def _try_load_cross_encoder (model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2" ): try : return CrossEncoder(model_name) except Exception: return None def rerank_with_cross_encoder (query: str , items: List [Dict ], model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2" , top_k: int = 10 ) -> List [Dict ]: ce = _try_load_cross_encoder(model_name) if ce is None or not items: return items[:top_k] pairs = [[query, it.get("content" , "" )] for it in items] try : scores = ce.predict(pairs) for it, s in zip (items, scores): it["rerank_score" ] = float (s) items.sort(key=lambda x: x.get("rerank_score" , x.get("score" , 0.0 )), reverse=True ) return items[:top_k] except Exception: return items[:top_k] def compute_graph_signals_from_pool (vector_hits: List [Dict ], same_doc_weight: float = 1.0 , proximity_weight: float = 1.0 , proximity_window_chars: int = 1600 ) -> Dict [str , float ]: """ Compute graph signals with direct parameters instead of environment variables. """ by_doc: Dict [str , List [Dict ]] = {} for h in vector_hits: meta = h.get("metadata" , {}) did = meta.get("doc_id" ) if not did: did = meta.get("memory_id" ) or h.get("id" ) by_doc.setdefault(did, []).append(h) doc_counts = {d: len (arr) for d, arr in by_doc.items()} max_count = max (doc_counts.values()) if doc_counts else 1 graph_signal: Dict [str , float ] = {} for did, arr in by_doc.items(): arr.sort(key=lambda x: x.get("metadata" , {}).get("start" , 0 )) density = doc_counts.get(did, 1 ) / max_count for i, h in enumerate (arr): mid = h.get("metadata" , {}).get("memory_id" , h.get("id" )) pos_i = h.get("metadata" , {}).get("start" , 0 ) prox_acc = 0.0 j = i - 1 while j >= 0 : pos_j = arr[j].get("metadata" , {}).get("start" , 0 ) dist = abs (pos_i - pos_j) if dist > proximity_window_chars: break prox_acc += max (0.0 , 1.0 - (dist / max (1.0 , float (proximity_window_chars)))) j -= 1 j = i + 1 while j < len (arr): pos_j = arr[j].get("metadata" , {}).get("start" , 0 ) dist = abs (pos_i - pos_j) if dist > proximity_window_chars: break prox_acc += max (0.0 , 1.0 - (dist / max (1.0 , float (proximity_window_chars)))) j += 1 score = same_doc_weight * density + proximity_weight * prox_acc graph_signal[mid] = graph_signal.get(mid, 0.0 ) + score if graph_signal: max_v = max (graph_signal.values()) if max_v > 0 : for k in list (graph_signal.keys()): graph_signal[k] = graph_signal[k] / max_v return graph_signal def rank (vector_hits: List [Dict ], graph_signals: Optional [Dict [str , float ]] = None , w_vector: float = 0.7 , w_graph: float = 0.3 ) -> List [Dict ]: """ Rank results with direct weight parameters instead of environment variables. """ items: List [Dict ] = [] graph_signals = graph_signals or {} for h in vector_hits: mid = h.get("metadata" , {}).get("memory_id" , h.get("id" )) g = float (graph_signals.get(mid, 0.0 )) v = float (h.get("score" , 0.0 )) score = w_vector * v + w_graph * g items.append({ "memory_id" : mid, "score" : score, "vector_score" : v, "graph_score" : g, "content" : h.get("metadata" , {}).get("content" , "" ), "metadata" : h.get("metadata" , {}), }) items.sort(key=lambda x: x["score" ], reverse=True ) return items def merge_snippets (ranked_items: List [Dict ], max_chars: int = 1200 ) -> str : out: List [str ] = [] total = 0 for it in ranked_items: text = it.get("content" , "" ).strip() if not text: continue if total + len (text) > max_chars: remain = max_chars - total if remain <= 0 : break out.append(text[:remain]) total += remain break out.append(text) total += len (text) return "\n\n" .join(out) def expand_neighbors_from_pool (selected: List [Dict ], pool: List [Dict ], neighbors: int = 1 , max_additions: int = 5 ) -> List [Dict ]: if not selected or not pool or neighbors <= 0 : return selected by_doc: Dict [str , List [Dict ]] = {} for it in pool: meta = it.get("metadata" , {}) did = meta.get("doc_id" ) if not did: continue by_doc.setdefault(did, []).append(it) for did, arr in by_doc.items(): arr.sort(key=lambda x: (x.get("metadata" , {}).get("start" , 0 ))) selected_ids = set (it.get("memory_id" ) for it in selected) additions: List [Dict ] = [] for it in selected: meta = it.get("metadata" , {}) did = meta.get("doc_id" ) if not did or did not in by_doc: continue arr = by_doc[did] try : idx = next (i for i, x in enumerate (arr) if x.get("memory_id" ) == it.get("memory_id" )) except StopIteration: continue for offset in range (1 , neighbors + 1 ): for j in (idx - offset, idx + offset): if 0 <= j < len (arr): cand = arr[j] mid = cand.get("memory_id" ) if mid not in selected_ids: additions.append(cand) selected_ids.add(mid) if len (additions) >= max_additions: break if len (additions) >= max_additions: break if len (additions) >= max_additions: break extended = list (selected) + additions extended.sort(key=lambda x: (x.get("rerank_score" , x.get("score" , 0.0 ))), reverse=True ) return extended def merge_snippets_grouped (ranked_items: List [Dict ], max_chars: int = 1200 , include_citations: bool = True ) -> str : by_doc: Dict [str , List [Dict ]] = {} doc_score: Dict [str , float ] = {} for it in ranked_items: meta = it.get("metadata" , {}) did = meta.get("doc_id" ) or meta.get("source_path" ) or "unknown" by_doc.setdefault(did, []).append(it) doc_score[did] = doc_score.get(did, 0.0 ) + float (it.get("score" , 0.0 )) ordered_docs = sorted (by_doc.keys(), key=lambda d: doc_score.get(d, 0.0 ), reverse=True ) for d in ordered_docs: by_doc[d].sort(key=lambda x: (x.get("metadata" , {}).get("start" , 0 ))) out: List [str ] = [] citations: List [Dict ] = [] total = 0 cite_index = 1 for did in ordered_docs: parts = by_doc[did] for it in parts: text = (it.get("content" , "" ) or "" ).strip() if not text: continue suffix = "" if include_citations: suffix = f" [{cite_index} ]" need = len (text) + (len (suffix) if suffix else 0 ) if total + need > max_chars: remain = max_chars - total if remain <= 0 : break clipped = text[: max (0 , remain - len (suffix))] if clipped: out.append(clipped + suffix) total += len (clipped) + len (suffix) if include_citations: m = it.get("metadata" , {}) citations.append({ "index" : cite_index, "source_path" : m.get("source_path" ), "doc_id" : m.get("doc_id" ), "start" : m.get("start" ), "end" : m.get("end" ), "heading_path" : m.get("heading_path" ), }) cite_index += 1 break out.append(text + suffix) total += need if include_citations: m = it.get("metadata" , {}) citations.append({ "index" : cite_index, "source_path" : m.get("source_path" ), "doc_id" : m.get("doc_id" ), "start" : m.get("start" ), "end" : m.get("end" ), "heading_path" : m.get("heading_path" ), }) cite_index += 1 if total >= max_chars: break merged = "\n\n" .join(out) if include_citations and citations: lines: List [str ] = [merged, "" , "References:" ] for c in citations: loc = "" if c.get("start" ) is not None and c.get("end" ) is not None : loc = f" ({c['start' ]} -{c['end' ]} )" hp = f" – {c['heading_path' ]} " if c.get("heading_path" ) else "" sp = c.get("source_path" ) or c.get("doc_id" ) or "source" lines.append(f"[{c['index' ]} ] {sp} {loc} {hp} " ) return "\n" .join(lines) return merged def compress_ranked_items (ranked_items: List [Dict ], enable_compression: bool = True , max_per_doc: int = 2 , join_gap: int = 200 ) -> List [Dict ]: """ Compress ranked items with direct parameters instead of environment variables. """ if not enable_compression: return ranked_items by_doc_count: Dict [str , int ] = {} last_by_doc: Dict [str , Dict ] = {} new_items: List [Dict ] = [] for it in ranked_items: meta = it.get("metadata" , {}) did = meta.get("doc_id" ) or meta.get("source_path" ) or "unknown" start = int (meta.get("start" ) or 0 ) end = int (meta.get("end" ) or (start + len (it.get("content" , "" ) or "" ))) if did not in last_by_doc: last_by_doc[did] = it by_doc_count[did] = 1 new_items.append(it) continue last = last_by_doc[did] lmeta = last.get("metadata" , {}) lstart = int (lmeta.get("start" ) or 0 ) lend = int (lmeta.get("end" ) or (lstart + len (last.get("content" , "" ) or "" ))) if start - lend <= join_gap and start >= lstart: merged_text = (last.get("content" , "" ) or "" ).strip() add_text = (it.get("content" , "" ) or "" ).strip() if add_text: if merged_text: merged_text = merged_text + "\n\n" + add_text else : merged_text = add_text last["content" ] = merged_text lmeta["end" ] = max (lend, end) try : last["score" ] = max (float (last.get("score" , 0.0 )), float (it.get("score" , 0.0 ))) except Exception: pass last_by_doc[did] = last else : cnt = by_doc_count.get(did, 0 ) if cnt >= max_per_doc: continue new_items.append(it) last_by_doc[did] = it by_doc_count[did] = cnt + 1 return new_items def tldr_summarize (text: str , bullets: int = 3 ) -> Optional [str ]: try : if not text or len (text.strip()) == 0 : return None llm = LLM() prompt = [ {"role" : "system" , "content" : "请将以下内容概括为简洁的要点列表(最多3-5条),用中文,避免重复,突出关键信息。" }, {"role" : "user" , "content" : f"请用 {max (1 , min (5 , int (bullets)))} 条要点总结:\n\n{text} " }, ] out = llm.think(prompt) return out except Exception: return None def create_rag_pipeline ( qdrant_url: Optional [str ] = None , qdrant_api_key: Optional [str ] = None , collection_name: str = "hello_agents_rag_vectors" , rag_namespace: str = "default" ) -> Dict [str , Any ]: """ Create a complete RAG pipeline with Qdrant and unified embedding. Returns: Dict containing store, namespace, and helper functions """ dimension = get_dimension(384 ) store = QdrantVectorStore( url=qdrant_url, api_key=qdrant_api_key, collection_name=collection_name, vector_size=dimension, distance="cosine" ) def add_documents (file_paths: List [str ], chunk_size: int = 800 , chunk_overlap: int = 100 ): """Add documents to RAG pipeline""" chunks = load_and_chunk_texts( paths=file_paths, chunk_size=chunk_size, chunk_overlap=chunk_overlap, namespace=rag_namespace, source_label="rag" ) index_chunks( store=store, chunks=chunks, rag_namespace=rag_namespace ) return len (chunks) def search (query: str , top_k: int = 8 , score_threshold: Optional [float ] = None ): """Search RAG knowledge base""" return search_vectors( store=store, query=query, top_k=top_k, rag_namespace=rag_namespace, score_threshold=score_threshold ) def search_advanced ( query: str , top_k: int = 8 , enable_mqe: bool = False , enable_hyde: bool = False , score_threshold: Optional [float ] = None ): """Advanced search with query expansion""" return search_vectors_expanded( store=store, query=query, top_k=top_k, rag_namespace=rag_namespace, enable_mqe=enable_mqe, enable_hyde=enable_hyde, score_threshold=score_threshold ) def get_stats (): """Get pipeline statistics""" return store.get_collection_stats() return { "store" : store, "namespace" : rag_namespace, "add_documents" : add_documents, "search" : search, "search_advanced" : search_advanced, "get_stats" : get_stats }