mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-12 23:40:03 +08:00
✨ feat!(llm): 重构并升级大语言模型服务为全新 AI 智能体框架 (#2146)
* ✨ feat!(llm): 重构并升级大语言模型服务为全新 AI 智能体框架 - 【重构】将原 services/llm 重构并迁移至全新的 services/ai 架构，提供向下兼容垫片 - 【新增】引入 Agent、Team、Workflow 三大智能体与工作流编排范式 - 【新增】引入基于 RAG 的长期向量记忆与中期槽位记忆系统 - 【新增】引入基于 Docker 的安全代码执行沙箱环境 - 【新增】支持 MCP 协议，允许动态管理和调用 MCP 服务 - 【新增】引入输入输出安全合规护栏与自愈反思机制 - 【优化】重构并优化多厂商 API 适配器 (Gemini, OpenAI, DeepSeek, GLM 等) - 【优化】优化日志脱敏与 Token 预估机制 - 【移除】移除旧版 llm default 和 llm reset-key 命令，新增 llm mcp 管理命令 * 🔧 chore(deps): 更新项目依赖与配置 - 添加 mcp、jieba 和 aiodocker 依赖到配置文件及 requirements.txt - 在 pyright 配置中设置 reportMissingImports 为 none - 调整 .gitignore 中 resources 目录的忽略规则 * ♻️ refactor(tools): 重构工具终止机制并清理知识库日志输出 - 统一使用 `context.state["__end_run__"]` 替代 `EndRunResult` 控制任务结束 - 移除文件系统和向量知识库检索工具中 `ToolResult` 的 `.with_log` 调用 - 调整指令处理器（Directive）的返回值为 `tool_res.output` - 修复部分类型检查警告并优化联合类型判断语法 * ♻️ refactor(tools): 重构工具副作用指令与控制流熔断机制 - 引入 `DirectivePayload` 及 `ToolResult` 的子类以结构化表达工具副作用 - 移除通过 `context.state` 传递魔术变量的隐式控制流设计 - 重构 `DirectiveManager` 处理器接口，直接在处理器中修改 `AgentState` 并构建 `AgentRunResult` - 在 `StandardAgentExecutor` 中统一通过 `directive_manager` 调度工具返回的副作用指令 - 补全 `MessageBuilder` 中部分核心方法的文档注释 * 🐛 fix(sandbox): 修复 Docker 沙箱容器状态检测与会话清理逻辑 -【修复】修正 `is_alive` 中直接读取私有属性的问题，改用 `show()` 返回值 -【修复】解决 `execute_code` 中缓存的执行器与当前会话不一致的问题 -【优化】在清理工作区前增加容器存活检测，避免向已死容器发送请求 -【优化】创建容器时增加运行状态校验，若已停止则自动从缓存中移除并重建 -【优化】优化容器销毁和清理逻辑，静默处理容器不存在 (404) 的异常 * 📝 docs(core): 补充核心模块初始化方法的文档注释 * 🚨 auto fix by pre-commit hooks --------- Co-authored-by: webjoin111 <455457521@qq.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Rumio
co-authored by
webjoin111
pre-commit-ci[bot]
webjoin111
pre-commit-ci[bot]
parent
bdc1374848
commit
80fc5b86a7
@@ -0,0 +1,20 @@
|
||||
"""
|
||||
Zhenxun AI - 上下文、记忆与知识管理子系统门面 (Context, Memory & Knowledge Facade)
|
||||
"""
|
||||
|
||||
from .knowledge import FileSystemKnowledge, VectorKnowledge
|
||||
from .memory import (
|
||||
AgentSessionFacade,
|
||||
MemoryBuilder,
|
||||
memory_manager,
|
||||
)
|
||||
from .rag import RAGBuilder
|
||||
|
||||
__all__ = [
|
||||
"AgentSessionFacade",
|
||||
"FileSystemKnowledge",
|
||||
"MemoryBuilder",
|
||||
"RAGBuilder",
|
||||
"VectorKnowledge",
|
||||
"memory_manager",
|
||||
]
|
||||
@@ -0,0 +1,7 @@
|
||||
from .filesystem import FileSystemKnowledge
|
||||
from .vector import VectorKnowledge
|
||||
|
||||
__all__ = [
|
||||
"FileSystemKnowledge",
|
||||
"VectorKnowledge",
|
||||
]
|
||||
@@ -0,0 +1,10 @@
|
||||
from zhenxun.services.ai.tools.core.toolkit import BaseToolkit
|
||||
|
||||
|
||||
class BaseKnowledge(BaseToolkit):
|
||||
"""
|
||||
知识库基类。
|
||||
继承自 BaseToolkit，允许知识库向 Agent 提供自定义的检索工具（如搜索、读取）。
|
||||
"""
|
||||
|
||||
pass
|
||||
@@ -0,0 +1,131 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from zhenxun.services.ai.tools.core.decorators import tool
|
||||
from zhenxun.services.ai.tools.models import ToolResult
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
from .base import BaseKnowledge
|
||||
|
||||
|
||||
class FileSystemKnowledge(BaseKnowledge):
|
||||
"""
|
||||
纯文本文件系统知识库。
|
||||
零依赖，无需向量数据库。通过大模型原生工具 (grep, list, read) 让其自主翻阅本地文件。
|
||||
"""
|
||||
|
||||
default_instructions = (
|
||||
"## 本地文件知识库\n"
|
||||
"你拥有访问本地专业文档的权限。在回答问题前，请遵循以下流程：\n"
|
||||
"1. **搜索**：优先使用 `search_knowledge_files` 通过关键词查找相关内容。\n"
|
||||
"2. **概览**：如果需要了解文件结构，使用 `list_knowledge_files`。\n"
|
||||
"3. **阅读**：找到目标后，使用 `read_knowledge_file` 获取完整上下文。\n"
|
||||
"**核心原则**：严禁凭空捏造事实，必须基于文件内容提取信息。"
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_dir: str | Path,
|
||||
allowed_extensions: tuple[str, ...] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
初始化文本文件系统知识库。
|
||||
|
||||
参数:
|
||||
base_dir: 知识库对应的本地根目录路径。
|
||||
allowed_extensions: 允许读取和检索的文件后缀元组，
|
||||
默认支持 txt, md, json, csv, yaml, log。
|
||||
**kwargs: 透传给父类 BaseKnowledge 的额外参数。
|
||||
"""
|
||||
super().__init__(**kwargs)
|
||||
self.base_dir = Path(base_dir).resolve()
|
||||
self.allowed_extensions = allowed_extensions or (
|
||||
".txt",
|
||||
".md",
|
||||
".json",
|
||||
".csv",
|
||||
".yaml",
|
||||
".log",
|
||||
)
|
||||
if not self.base_dir.exists():
|
||||
logger.warning(f"[FileSystemKnowledge] 警告：目录不存在 {self.base_dir}")
|
||||
self.base_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _is_safe_path(self, target_path: Path) -> bool:
|
||||
"""安全检查：防止跨目录访问"""
|
||||
try:
|
||||
return target_path.resolve().is_relative_to(self.base_dir)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@tool(
|
||||
name="search_knowledge_files",
|
||||
description="在知识库中搜索包含指定关键词的文件内容和上下文。",
|
||||
)
|
||||
async def search_knowledge_files(self, keyword: str) -> ToolResult:
|
||||
"""扫描所有文本文件，返回包含关键词的行及上下文。"""
|
||||
results = []
|
||||
try:
|
||||
pattern = re.compile(re.escape(keyword), re.IGNORECASE)
|
||||
except Exception:
|
||||
return ToolResult(output=f"无效的搜索关键词: {keyword}").as_error()
|
||||
|
||||
for root, _, files in os.walk(self.base_dir):
|
||||
for file in files:
|
||||
if not file.endswith(self.allowed_extensions):
|
||||
continue
|
||||
|
||||
file_path = Path(root) / file
|
||||
try:
|
||||
content = file_path.read_text(encoding="utf-8", errors="ignore")
|
||||
lines = content.splitlines()
|
||||
|
||||
matches = []
|
||||
for i, line in enumerate(lines):
|
||||
if pattern.search(line):
|
||||
start = max(0, i - 1)
|
||||
end = min(len(lines), i + 2)
|
||||
context = "\n".join(lines[start:end])
|
||||
matches.append(context)
|
||||
|
||||
if matches:
|
||||
rel_path = file_path.relative_to(self.base_dir).as_posix()
|
||||
match_text = "\n---\n".join(matches[:5])
|
||||
results.append(f"📁 文件: {rel_path}\n{match_text}")
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if not results:
|
||||
return ToolResult(output=f"未找到包含 '{keyword}' 的内容。尝试更换关键词。")
|
||||
|
||||
final_output = "\n\n======\n\n".join(results[:10])
|
||||
return ToolResult(output=final_output)
|
||||
|
||||
@tool(name="list_knowledge_files", description="列出知识库中所有可用的文档路径。")
|
||||
async def list_knowledge_files(self) -> ToolResult:
|
||||
files = []
|
||||
for root, _, filenames in os.walk(self.base_dir):
|
||||
for filename in filenames:
|
||||
file_path = Path(root) / filename
|
||||
rel_path = file_path.relative_to(self.base_dir).as_posix()
|
||||
files.append(rel_path)
|
||||
|
||||
if not files:
|
||||
return ToolResult(output="知识库当前为空。")
|
||||
return ToolResult(output="可用文件列表:\n" + "\n".join(files))
|
||||
|
||||
@tool(
|
||||
name="read_knowledge_file", description="读取知识库中指定文件的完整文本内容。"
|
||||
)
|
||||
async def read_knowledge_file(self, file_path: str) -> ToolResult:
|
||||
target = (self.base_dir / file_path).resolve()
|
||||
if not self._is_safe_path(target):
|
||||
return ToolResult(output="❌ 安全拦截：越权访问").as_error()
|
||||
if not target.is_file():
|
||||
return ToolResult(output=f"❌ 文件不存在: {file_path}").as_error()
|
||||
|
||||
content = target.read_text(encoding="utf-8", errors="ignore")
|
||||
return ToolResult(output=content)
|
||||
@@ -0,0 +1,74 @@
|
||||
import asyncio
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
from zhenxun.services.ai.context.rag.models import BaseRecord
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
|
||||
class BaseReader:
|
||||
"""读取器基类"""
|
||||
|
||||
async def read(self, file_path: Path) -> BaseRecord | None:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class TextReader(BaseReader):
|
||||
"""处理 .txt, .md, .json 等纯文本"""
|
||||
|
||||
async def read(self, file_path: Path) -> BaseRecord | None:
|
||||
try:
|
||||
import aiofiles
|
||||
|
||||
async with aiofiles.open(file_path, encoding="utf-8", errors="ignore") as f:
|
||||
content = await f.read()
|
||||
|
||||
enriched_content = f"文档名称：{file_path.stem}\n文档内容：\n{content}"
|
||||
return BaseRecord(
|
||||
content=enriched_content,
|
||||
metadata={
|
||||
"name": file_path.name,
|
||||
"extension": file_path.suffix.lower(),
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"[TextReader] 读取文件失败 {file_path}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
class CSVReader(BaseReader):
|
||||
"""
|
||||
处理 .csv 报表
|
||||
将其标准化为逗号分隔的文本块，便于后续的 RowChunking 处理。
|
||||
"""
|
||||
|
||||
async def read(self, file_path: Path) -> BaseRecord | None:
|
||||
try:
|
||||
|
||||
def _read_csv():
|
||||
with open(file_path, encoding="utf-8", errors="ignore") as f:
|
||||
reader = csv.reader(f)
|
||||
return [
|
||||
",".join(
|
||||
[
|
||||
str(cell).replace("\n", " ").replace("\r", "")
|
||||
for cell in row
|
||||
]
|
||||
)
|
||||
for row in reader
|
||||
]
|
||||
|
||||
lines = await asyncio.to_thread(_read_csv)
|
||||
content = "\n".join(lines)
|
||||
|
||||
enriched_content = f"数据表名称：{file_path.stem}\n数据内容：\n{content}"
|
||||
return BaseRecord(
|
||||
content=enriched_content,
|
||||
metadata={
|
||||
"name": file_path.name,
|
||||
"extension": file_path.suffix.lower(),
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"[CSVReader] 读取 CSV 失败 {file_path}: {e}")
|
||||
return None
|
||||
@@ -0,0 +1,319 @@
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import anyio
|
||||
from nonebot.adapters import Bot, Event
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from zhenxun.services.ai.context.knowledge.base import BaseKnowledge
|
||||
from zhenxun.services.ai.context.knowledge.readers import (
|
||||
BaseReader,
|
||||
CSVReader,
|
||||
TextReader,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.engine import ScopedRAGClient
|
||||
from zhenxun.services.ai.context.rag.models import BaseRecord
|
||||
from zhenxun.services.ai.core.messages import LLMMessage
|
||||
from zhenxun.services.ai.llm.api import generate_structured
|
||||
from zhenxun.services.ai.run import RunContext
|
||||
from zhenxun.services.ai.tools.core.decorators import tool
|
||||
from zhenxun.services.ai.tools.models import ToolkitConfig, ToolResult
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
|
||||
class QueryAnalysis(BaseModel):
|
||||
"""大模型结构化提取查询意图"""
|
||||
|
||||
keywords: list[str] = Field(
|
||||
description="提取出1~3个极其简短的搜索短语或名词，严格去除所有客套话、修饰词和标点。如果用户意图跨度较大，可以拆分为多个短语。"
|
||||
)
|
||||
|
||||
|
||||
class VectorKnowledge(BaseKnowledge):
|
||||
"""
|
||||
原生语义向量知识库。
|
||||
将长文档切分、向量化并存入关系型/向量数据库，
|
||||
向大模型提供语义检索 (Semantic Search) 工具。
|
||||
"""
|
||||
|
||||
default_instructions = (
|
||||
"## 语义知识库\n"
|
||||
"你拥有访问外部语义向量知识库的权限。请遵循以下规则：\n"
|
||||
"1. **优先检索**：在回答专业或背景问题时，务必使用 `search_knowledge` 工具。\n"
|
||||
"2. **语义搜索**：你可以直接输入完整的问题或描述作为检索词，"
|
||||
"系统会自动进行语义匹配。\n"
|
||||
"3. **精确过滤**：如果你需要查阅特定范围，可以在 filters 参数中"
|
||||
"传入 JSON 字典进行精确匹配（如 {'source': 'local_file'}）。\n"
|
||||
"4. **基于事实**：必须仅根据检索到的内容回答，严禁编造信息。"
|
||||
)
|
||||
|
||||
default_auto_inject_template = (
|
||||
"### 📚 [本地知识库自动检索结果]\n"
|
||||
"基于用户的最新提问，系统后台已自动为你检索了以下参考资料。"
|
||||
"请你务必优先结合以下资料回答用户的问题，严禁编造：\n\n"
|
||||
"{knowledge_text}"
|
||||
)
|
||||
|
||||
_global_storage: Any = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rag_client: ScopedRAGClient | None = None,
|
||||
injection_mode: Literal["tool", "auto", "smart"] = "tool",
|
||||
query_rewrite_model: str | None = None,
|
||||
auto_inject_template: str | None = None,
|
||||
query_rewrite_prompt: str | None = None,
|
||||
query_rewrite_instruction: str | None = None,
|
||||
search_limit: int = 8,
|
||||
inject_limit: int = 12,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
初始化向量语义知识库工具箱。
|
||||
|
||||
参数:
|
||||
rag_client: RAG 基础设施客户端实例，默认 None。
|
||||
injection_mode: 知识库的介入模式。
|
||||
- "tool": 纯工具模式 (默认)。大模型需自主思考并显式调用 `search_knowledge` 工具获取信息。
|
||||
- "auto": 自动注入模式。向大模型隐藏检索工具，直接使用用户的原始输入去数据库粗筛并静默注入。
|
||||
- "smart": 智能查询模式。向大模型隐藏检索工具，先利用 LLM 对用户的提问进行改写，再查库注入，准确率最高。
|
||||
query_rewrite_model: 在 "smart" 模式下，指定用于重写查询词的大模型名称（为空则跟随当前主模型）。
|
||||
auto_inject_template: 自动/智能注入模式下向大模型提示词注入的模板字符串，默认 None。
|
||||
query_rewrite_prompt: 智能模式下对查询词进行改写时的提示词，默认 None。
|
||||
query_rewrite_instruction: 智能模式下进行查询词改写的大模型 System 提示词说明，默认 None。
|
||||
search_limit: 单次库检索的返回记录数限制，默认 8。
|
||||
inject_limit: 最终合并去重后注入给大模型的上下文片段数上限，默认 12。
|
||||
**kwargs: 传递给父类的其他关键字参数。
|
||||
""" # noqa: E501
|
||||
self.injection_mode = injection_mode
|
||||
self.query_rewrite_model = query_rewrite_model
|
||||
self.auto_inject_template = (
|
||||
auto_inject_template or self.default_auto_inject_template
|
||||
)
|
||||
self.query_rewrite_prompt = (
|
||||
query_rewrite_prompt
|
||||
or "用户原始提问：{query}\n\n请提取核心搜索词用于专业知识库向量检索。"
|
||||
)
|
||||
self.query_rewrite_instruction = (
|
||||
query_rewrite_instruction or "你是一个资深的数据检索架构师。"
|
||||
)
|
||||
self.search_limit = search_limit
|
||||
self.inject_limit = inject_limit
|
||||
|
||||
if injection_mode in ("auto", "smart"):
|
||||
config = kwargs.get("config")
|
||||
if not config:
|
||||
config = ToolkitConfig()
|
||||
kwargs["config"] = config
|
||||
if config.exclude is None:
|
||||
config.exclude = []
|
||||
config.exclude.append("search_knowledge")
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
if rag_client is None:
|
||||
from zhenxun.services.ai.context.rag.backends import DictStorageBackend
|
||||
from zhenxun.services.ai.context.rag.builder import RAGBuilder
|
||||
|
||||
rag_client = RAGBuilder(DictStorageBackend()).build()
|
||||
|
||||
self.rag_client = rag_client
|
||||
self.readers: dict[str, BaseReader] = {}
|
||||
|
||||
txt_reader = TextReader()
|
||||
for ext in [".txt", ".md", ".json", ".log", ".yaml", ".yml", ".ini"]:
|
||||
self.readers[ext] = txt_reader
|
||||
self.readers[".csv"] = CSVReader()
|
||||
|
||||
@classmethod
|
||||
def from_event(
|
||||
cls,
|
||||
event: Event | None = None,
|
||||
bot: Bot | None = None,
|
||||
isolation: Literal["group", "user"] = "group",
|
||||
**kwargs,
|
||||
) -> "VectorKnowledge":
|
||||
"""
|
||||
根据 NoneBot 的 Event 自动推导并创建一个物理隔离的向量知识库实例。
|
||||
"""
|
||||
from zhenxun.services.ai.context.memory.types import Isolation
|
||||
from zhenxun.services.ai.context.rag.backends import DictStorageBackend
|
||||
from zhenxun.services.ai.context.rag.builder import RAGBuilder
|
||||
from zhenxun.services.ai.run.context import NoneBotDeps
|
||||
from zhenxun.services.ai.utils import ContextUtils
|
||||
|
||||
if not bot or not event:
|
||||
deps = NoneBotDeps.get_current()
|
||||
bot = bot or (deps.bot if deps else None)
|
||||
event = event or (deps.event if deps else None)
|
||||
|
||||
if not bot or not event:
|
||||
raise ValueError(
|
||||
"无法隐式获取当前对话的 Bot 或 Event 上下文,"
|
||||
"如果您在定时任务或后台线程中使用，请显式传入 bot 和 event 参数。"
|
||||
)
|
||||
|
||||
scope_builder = (
|
||||
Isolation.GROUP_SHARED()
|
||||
if isolation == "group"
|
||||
else Isolation.USER_GLOBAL()
|
||||
)
|
||||
session_meta = ContextUtils.generate_session_meta(
|
||||
bot=bot, event=event, scope_builder=scope_builder, namespace="auto_kb"
|
||||
)
|
||||
|
||||
if cls._global_storage is None:
|
||||
cls._global_storage = DictStorageBackend()
|
||||
|
||||
client = (
|
||||
RAGBuilder(cls._global_storage)
|
||||
.with_scope(session_meta.accessible_scopes)
|
||||
.build()
|
||||
)
|
||||
return cls(rag_client=client, **kwargs)
|
||||
|
||||
def register_reader(
|
||||
self, ext: str | list[str], reader: BaseReader
|
||||
) -> "VectorKnowledge":
|
||||
"""挂载自定义后缀文件解析器 (如 PDF, Docx)，支持链式调用"""
|
||||
exts = [ext] if isinstance(ext, str) else ext
|
||||
for e in exts:
|
||||
e = e.lower()
|
||||
if not e.startswith("."):
|
||||
e = f".{e}"
|
||||
self.readers[e] = reader
|
||||
return self
|
||||
|
||||
def get_instructions(self) -> str | None:
|
||||
"""
|
||||
如果是自动/智能注入模式，对大模型完全隐藏检索提示词，防止其误调用。
|
||||
"""
|
||||
if self.injection_mode != "tool":
|
||||
return None
|
||||
return super().get_instructions()
|
||||
|
||||
async def before_llm_request(
|
||||
self, context: RunContext, messages: list[Any]
|
||||
) -> None:
|
||||
"""
|
||||
生命周期钩子：在向底层 LLM 发起请求前触发。
|
||||
负责执行 "auto" 或 "smart" 模式下的前置主动检索与上下文注入。
|
||||
"""
|
||||
if self.injection_mode == "tool":
|
||||
return
|
||||
|
||||
user_input = context.run.user_input
|
||||
if not user_input:
|
||||
return
|
||||
|
||||
queries_to_search = [user_input]
|
||||
|
||||
if self.injection_mode == "smart":
|
||||
try:
|
||||
model_to_use = self.query_rewrite_model or context.run.current_model
|
||||
prompt = self.query_rewrite_prompt.format(query=user_input)
|
||||
res = await generate_structured(
|
||||
message=prompt,
|
||||
response_model=QueryAnalysis,
|
||||
model=model_to_use,
|
||||
instruction=self.query_rewrite_instruction,
|
||||
)
|
||||
if res.keywords:
|
||||
logger.info(
|
||||
f"✨ [Smart Knowledge] 搜索词改写成功: '{user_input}' -> "
|
||||
f"{res.keywords}"
|
||||
)
|
||||
queries_to_search = res.keywords
|
||||
except Exception as e:
|
||||
logger.warning(f"[Smart Knowledge] Query 改写失败，降级使用原词: {e}")
|
||||
|
||||
all_results = []
|
||||
seen_ids = set()
|
||||
for q in queries_to_search:
|
||||
results = await self.rag_client.search(query=q, limit=self.search_limit)
|
||||
for res in results:
|
||||
if res.record.id not in seen_ids:
|
||||
seen_ids.add(res.record.id)
|
||||
all_results.append(res)
|
||||
|
||||
if not all_results:
|
||||
return
|
||||
|
||||
all_results.sort(key=lambda x: x.score, reverse=True)
|
||||
all_results = all_results[: self.inject_limit]
|
||||
|
||||
formatted_results = []
|
||||
for result in all_results:
|
||||
doc_name = result.record.metadata.get("name", "未命名文档")
|
||||
formatted_results.append(
|
||||
f"📄 来源: {doc_name}\n内容片段:\n{result.record.content}"
|
||||
)
|
||||
|
||||
knowledge_text = "\n\n======\n\n".join(formatted_results)
|
||||
system_prompt = self.auto_inject_template.format(knowledge_text=knowledge_text)
|
||||
|
||||
messages.insert(0, LLMMessage.system(system_prompt))
|
||||
|
||||
async def add_document(self, document: BaseRecord) -> int:
|
||||
"""
|
||||
通过注入的 Ingestion Pipeline 处理并入库文档
|
||||
返回成功入库的 Chunk 数量。
|
||||
"""
|
||||
return await self.rag_client.ingest([document])
|
||||
|
||||
async def add_file(self, file_path: str | Path) -> int:
|
||||
"""
|
||||
读取并注入单个文件。
|
||||
"""
|
||||
aio_path = anyio.Path(file_path)
|
||||
std_path = Path(file_path)
|
||||
if not await aio_path.is_file():
|
||||
logger.error(f"[VectorKnowledge] 文件不存在: {std_path}")
|
||||
return 0
|
||||
|
||||
ext = std_path.suffix.lower()
|
||||
reader = self.readers.get(ext)
|
||||
if not reader:
|
||||
logger.warning(f"当前知识库未配置支持解析文件后缀: {ext}")
|
||||
return 0
|
||||
|
||||
doc = await reader.read(std_path)
|
||||
if not doc:
|
||||
return 0
|
||||
|
||||
return await self.rag_client.ingest([doc])
|
||||
|
||||
async def add_directory(self, dir_path: str | Path) -> int:
|
||||
"""扫描目录并注入所有支持的文件"""
|
||||
total_chunks = 0
|
||||
aio_path = anyio.Path(dir_path)
|
||||
async for p in aio_path.rglob("*"):
|
||||
if await p.is_file():
|
||||
total_chunks += await self.add_file(Path(p))
|
||||
return total_chunks
|
||||
|
||||
@tool(
|
||||
name="search_knowledge",
|
||||
description=(
|
||||
"在语义知识库中搜索最相关的内容片段。可以通过 filters 字典进行额外过滤。"
|
||||
),
|
||||
)
|
||||
async def search_knowledge(
|
||||
self, query: str, filters: dict[str, Any] | None = None, limit: int = 5
|
||||
) -> ToolResult:
|
||||
results = await self.rag_client.search(
|
||||
query=query, limit=limit, metadata_filters=filters
|
||||
)
|
||||
|
||||
if not results:
|
||||
return ToolResult(output=f"知识库中未找到与 '{query}' 紧密相关的内容。")
|
||||
|
||||
formatted_results = []
|
||||
for result in results:
|
||||
doc_name = result.record.metadata.get("name", "未命名文档")
|
||||
formatted_results.append(
|
||||
f"📄 来源: {doc_name}\n" f"片段内容:\n{result.record.content}"
|
||||
)
|
||||
|
||||
final_text = "\n\n======\n\n".join(formatted_results)
|
||||
return ToolResult(output=final_text)
|
||||
@@ -0,0 +1,23 @@
|
||||
from .builder import MemoryBuilder
|
||||
from .compression import MemoryPolicy
|
||||
from .facades import AgentSessionFacade
|
||||
from .manager import memory_manager
|
||||
from .models import (
|
||||
BaseMemoryIngestionMiddleware,
|
||||
MemoryConfig,
|
||||
)
|
||||
from .types import (
|
||||
Isolation,
|
||||
SessionMetadata,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AgentSessionFacade",
|
||||
"BaseMemoryIngestionMiddleware",
|
||||
"Isolation",
|
||||
"MemoryBuilder",
|
||||
"MemoryConfig",
|
||||
"MemoryPolicy",
|
||||
"SessionMetadata",
|
||||
"memory_manager",
|
||||
]
|
||||
@@ -0,0 +1,270 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing_extensions import Self
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from zhenxun.services.ai.context.memory.models import MemorySlot
|
||||
from zhenxun.services.ai.context.memory.storage.interfaces import (
|
||||
BaseChatContext,
|
||||
BaseMemoryIngestionMiddleware,
|
||||
BaseSlotContext,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.backends import Embedder, StorageBackend
|
||||
from zhenxun.services.ai.context.rag.engine import ScopedRAGClient
|
||||
|
||||
from zhenxun.services.ai.context.memory.compression import MemoryPolicy
|
||||
from zhenxun.services.ai.context.memory.models import (
|
||||
ContextCompressionConfig,
|
||||
IngestionConfig,
|
||||
LongTermConfig,
|
||||
MemoryConfig,
|
||||
ShortTermConfig,
|
||||
SlotMemoryConfig,
|
||||
)
|
||||
from zhenxun.services.ai.context.memory.types import (
|
||||
AutoRecallPolicy,
|
||||
)
|
||||
from zhenxun.services.ai.utils.scope import ScopeBuilder
|
||||
|
||||
|
||||
class MemoryBuilder:
|
||||
"""
|
||||
记忆配置的链式构建器 (Fluent Builder)。
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
初始化 MemoryBuilder 实例。
|
||||
|
||||
创建一个默认关闭短期和长期记忆，并包含默认上下文压缩配置的构建器。
|
||||
"""
|
||||
self._config = MemoryConfig(
|
||||
short_term=ShortTermConfig(enable=False),
|
||||
slots=SlotMemoryConfig(enable=False),
|
||||
long_term=LongTermConfig(enable=False),
|
||||
compression=ContextCompressionConfig(),
|
||||
ingestion=IngestionConfig(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def auto(cls) -> "MemoryBuilder":
|
||||
"""
|
||||
创建一个开箱即用的默认记忆配置构建器。
|
||||
|
||||
默认开启隔离的短期记忆，并使用 LLM 对话摘要进行上下文压缩。
|
||||
"""
|
||||
return cls().with_short_term(enable=True).with_llm_summary()
|
||||
|
||||
@classmethod
|
||||
def resolve(
|
||||
cls, memory: bool | MemoryConfig | "MemoryBuilder" | None
|
||||
) -> MemoryConfig:
|
||||
if isinstance(memory, MemoryConfig):
|
||||
return memory
|
||||
if isinstance(memory, cls):
|
||||
return memory.build()
|
||||
if isinstance(memory, bool):
|
||||
return MemoryConfig(
|
||||
short_term=ShortTermConfig(enable=memory),
|
||||
long_term=LongTermConfig(enable=memory),
|
||||
)
|
||||
|
||||
return MemoryConfig(short_term=ShortTermConfig(enable=False))
|
||||
|
||||
def with_base_isolation(self, isolation: ScopeBuilder) -> Self:
|
||||
"""设置顶层基准隔离级别，短期/中期/长期记忆将默认继承此级别"""
|
||||
self._config.base_isolation = isolation
|
||||
self._config.short_term.isolation = isolation
|
||||
return self
|
||||
|
||||
def with_short_term(
|
||||
self,
|
||||
enable: bool = True,
|
||||
isolation: ScopeBuilder | None = None,
|
||||
backend: "str | BaseChatContext | None" = None,
|
||||
) -> Self:
|
||||
"""
|
||||
配置短期对话历史记忆。
|
||||
|
||||
参数:
|
||||
enable: 是否开启短期记忆。
|
||||
isolation: 记忆隔离级别 (ScopeBuilder)，决定会话历史记录的区分范围。
|
||||
backend: 短期记忆存储后端实例，如果为 None 则使用全局默认后端。
|
||||
"""
|
||||
self._config.short_term.enable = enable
|
||||
if isolation is not None:
|
||||
self._config.base_isolation = isolation
|
||||
self._config.short_term.isolation = isolation
|
||||
if backend is not None:
|
||||
self._config.short_term.backend = backend
|
||||
return self
|
||||
|
||||
def with_slots(
|
||||
self,
|
||||
enable: bool = True,
|
||||
scopes: dict[str, ScopeBuilder] | None = None,
|
||||
default_slots: list["MemorySlot"] | None = None,
|
||||
backend: "str | BaseSlotContext | None" = None,
|
||||
instructions: str | None = None,
|
||||
) -> Self:
|
||||
"""
|
||||
配置核心槽位记忆 (Memory Slots)。
|
||||
|
||||
参数:
|
||||
enable: 是否启用槽位记忆。
|
||||
scopes: 语义化作用域映射字典。如果只有一个键值对，则大模型不可见该参数。
|
||||
default_slots: 首次初始化时自动写入的默认槽位列表。
|
||||
backend: 槽位记忆存储后端，如果为 None 则使用全局默认后端。
|
||||
instructions: 覆写内置槽位管理工具箱的默认系统提示词规则。
|
||||
"""
|
||||
self._config.slots.enable = enable
|
||||
if scopes is not None:
|
||||
self._config.slots.scopes = scopes
|
||||
if default_slots is not None:
|
||||
self._config.slots.default_slots = default_slots
|
||||
if backend is not None:
|
||||
self._config.slots.backend = backend
|
||||
if instructions is not None:
|
||||
self._config.slots.instructions = instructions
|
||||
return self
|
||||
|
||||
def with_long_term(
|
||||
self,
|
||||
enable: bool = True,
|
||||
scopes: dict[str, ScopeBuilder] | None = None,
|
||||
engine: "ScopedRAGClient | None" = None,
|
||||
backend: "str | StorageBackend | None" = None,
|
||||
embedder: "Embedder | str | None" = None,
|
||||
agentic: bool = True,
|
||||
auto_recall: AutoRecallPolicy = False,
|
||||
instructions: str | None = None,
|
||||
) -> Self:
|
||||
"""
|
||||
配置长期向量记忆与 RAG 设定。
|
||||
|
||||
参数:
|
||||
enable: 是否启用长期记忆。
|
||||
scopes: 语义化作用域映射字典。如果只有一个键值对，则大模型不可见该参数。
|
||||
engine: 高级 RAG 检索引擎实例 (推荐)。若提供，将接管记忆的底层检索、混合与重排。
|
||||
backend: 长期记忆存储后端。
|
||||
embedder: 用于向量化的文本嵌入模型实例。
|
||||
agentic: 是否开启主动智能体记忆管理 (增删改查工具自动注入)。
|
||||
auto_recall: 长期记忆的自动召回策略，支持 bool 或 Callable 函数。
|
||||
instructions: 覆写内置长期记忆工具箱的默认系统提示词规则。
|
||||
""" # noqa: E501
|
||||
self._config.long_term.enable = enable
|
||||
self._config.long_term.engine = engine
|
||||
if scopes is not None:
|
||||
self._config.long_term.scopes = scopes
|
||||
self._config.long_term.backend = backend
|
||||
self._config.long_term.embedder = embedder
|
||||
self._config.long_term.agentic = agentic
|
||||
self._config.long_term.auto_recall = auto_recall
|
||||
if instructions is not None:
|
||||
self._config.long_term.instructions = instructions
|
||||
return self
|
||||
|
||||
def with_multimodal_window(self, window_size: int = 5) -> Self:
|
||||
"""
|
||||
配置多模态历史视窗大小。
|
||||
|
||||
超出此窗口的图片/视频等富媒体消息会自动转换为纯文本占位符，以节省 Token 预算。
|
||||
|
||||
参数:
|
||||
window_size: 允许保留多模态信息的最新的对话轮数。
|
||||
"""
|
||||
self._config.compression.vision_window = window_size
|
||||
return self
|
||||
|
||||
def with_llm_summary(
|
||||
self,
|
||||
trigger_tokens: int = 4000,
|
||||
max_turns: int = 0,
|
||||
keep_recent_turns: int = 0,
|
||||
summarization_model: str | None = None,
|
||||
summarization_prompt: str = "请概括以下对话内容，保留关键的约束条件、用户偏好、已完成的任务状态和未解决的问题。", # noqa: E501
|
||||
) -> Self:
|
||||
"""
|
||||
配置使用大模型自然语言总结作为上下文压缩策略。
|
||||
|
||||
参数:
|
||||
trigger_tokens: 触发压缩的 Token 门槛。
|
||||
max_turns: 压缩策略作用的最大历史对话轮数上限。
|
||||
keep_recent_turns: 在大模型总结之外，强制保留的最近原始对话轮数。
|
||||
summarization_model: 负责生成总结的大模型名称。
|
||||
summarization_prompt: 生成总结时所使用的系统提示词。
|
||||
"""
|
||||
self._config.compression.policy = MemoryPolicy.llm_summarize(
|
||||
trigger_tokens=trigger_tokens,
|
||||
max_turns=max_turns,
|
||||
keep_recent_turns=keep_recent_turns,
|
||||
summarization_model=summarization_model,
|
||||
summarization_prompt=summarization_prompt,
|
||||
)
|
||||
return self
|
||||
|
||||
def with_structured_summary(
|
||||
self,
|
||||
trigger_tokens: int = 4000,
|
||||
max_turns: int = 0,
|
||||
keep_recent_turns: int = 0,
|
||||
summarization_model: str | None = None,
|
||||
response_model: type[BaseModel] | None = None,
|
||||
prompt_template: str | None = None,
|
||||
format_callback: Callable[[Any], str] | None = None,
|
||||
) -> Self:
|
||||
"""
|
||||
配置使用自定义结构化 JSON 提取作为上下文压缩策略。
|
||||
|
||||
参数:
|
||||
trigger_tokens: 触发压缩的 Token 门槛。
|
||||
max_turns: 压缩策略作用的最大历史对话轮数上限。
|
||||
keep_recent_turns: 强制保留的最近原始对话轮数。
|
||||
summarization_model: 负责生成结构化总结的大模型名称。
|
||||
response_model: (可选) 自定义的 Pydantic 数据模型，用于指导提取的结构。
|
||||
prompt_template: (可选) 提取提示词模板，
|
||||
支持 {prev_summary} 和 {dialogue} 变量。
|
||||
format_callback: (可选) 将提取出的 Pydantic 实例格式化为字符串的回调函数。
|
||||
"""
|
||||
self._config.compression.policy = MemoryPolicy.structured_summarize(
|
||||
trigger_tokens=trigger_tokens,
|
||||
max_turns=max_turns,
|
||||
keep_recent_turns=keep_recent_turns,
|
||||
summarization_model=summarization_model,
|
||||
response_model=response_model,
|
||||
prompt_template=prompt_template,
|
||||
format_callback=format_callback,
|
||||
)
|
||||
return self
|
||||
|
||||
def unlimited(self) -> Self:
|
||||
"""
|
||||
配置为不进行任何截断和压缩的策略。
|
||||
适用于短程会话或者具备超长上下文窗口的底层语言模型。
|
||||
"""
|
||||
self._config.compression.policy = MemoryPolicy.unlimited()
|
||||
return self
|
||||
|
||||
def with_ingestion_middlewares(
|
||||
self, *middlewares: "BaseMemoryIngestionMiddleware"
|
||||
) -> Self:
|
||||
"""
|
||||
配置记忆入库管线中间件。
|
||||
用于在消息正式落盘前进行实体消解、隐私脱敏、自动打标签等操作。
|
||||
"""
|
||||
self._config.ingestion.middlewares.extend(middlewares)
|
||||
return self
|
||||
|
||||
def build(self) -> MemoryConfig:
|
||||
"""
|
||||
生成最终构建好的 MemoryConfig 配置对象。
|
||||
"""
|
||||
if not self._config.slots.scopes:
|
||||
self._config.slots.scopes = {"私有": self._config.base_isolation}
|
||||
if not self._config.long_term.scopes:
|
||||
self._config.long_term.scopes = {"私有": self._config.base_isolation}
|
||||
return self._config
|
||||
@@ -0,0 +1,52 @@
|
||||
from typing import Any
|
||||
|
||||
from zhenxun.services.ai.capabilities.base import AbstractCapability
|
||||
from zhenxun.services.ai.context.memory.models import MemoryConfig
|
||||
from zhenxun.services.ai.run.context import RunContext
|
||||
from zhenxun.services.ai.tools.providers.builtin.memory import MemoryManagementToolkit
|
||||
|
||||
|
||||
class AgenticMemoryCapability(AbstractCapability):
|
||||
"""
|
||||
智能体主动记忆管理能力 (Agentic Memory Management)。
|
||||
当 `MemoryConfig.long_term.enable == True` 且 `agentic == True` 时隐式挂载，
|
||||
在运行时动态组装并向大模型提供 `MemoryManagementToolkit` 工具箱。
|
||||
"""
|
||||
|
||||
def __init__(self, memory_config: MemoryConfig, namespace: str):
|
||||
self.memory_config = memory_config
|
||||
self.namespace = namespace
|
||||
|
||||
async def get_tools(self, context: RunContext) -> list[Any]:
|
||||
kwargs = self.memory_config.long_term.toolkit_kwargs.copy()
|
||||
kwargs["memory_config"] = self.memory_config
|
||||
kwargs["namespace"] = self.namespace
|
||||
if self.memory_config.long_term.instructions is not None:
|
||||
kwargs["instructions"] = self.memory_config.long_term.instructions
|
||||
|
||||
toolkit = MemoryManagementToolkit(**kwargs)
|
||||
return [toolkit]
|
||||
|
||||
|
||||
class SlotMemoryCapability(AbstractCapability):
|
||||
"""
|
||||
槽位记忆能力组件。
|
||||
当 `MemoryConfig.slots.enable == True` 时隐式挂载，
|
||||
在运行时动态组装并向大模型提供 `MemorySlotToolkit` 工具箱。
|
||||
"""
|
||||
|
||||
def __init__(self, memory_config: MemoryConfig, namespace: str):
|
||||
self.memory_config = memory_config
|
||||
self.namespace = namespace
|
||||
|
||||
async def get_tools(self, context: RunContext) -> list[Any]:
|
||||
from zhenxun.services.ai.tools.providers.builtin.slots import MemorySlotToolkit
|
||||
|
||||
kwargs = self.memory_config.slots.toolkit_kwargs.copy()
|
||||
kwargs["memory_config"] = self.memory_config
|
||||
kwargs["namespace"] = self.namespace
|
||||
if self.memory_config.slots.instructions is not None:
|
||||
kwargs["instructions"] = self.memory_config.slots.instructions
|
||||
|
||||
toolkit = MemorySlotToolkit(**kwargs)
|
||||
return [toolkit]
|
||||
@@ -0,0 +1,688 @@
|
||||
from abc import abstractmethod
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from zhenxun.services.ai.context.memory.storage.interfaces import (
|
||||
BaseMemoryReducer,
|
||||
)
|
||||
from zhenxun.services.ai.core.engine.token_counter import token_counter
|
||||
from zhenxun.services.ai.core.messages import (
|
||||
AudioPart,
|
||||
FilePart,
|
||||
ImagePart,
|
||||
LLMMessage,
|
||||
SystemMessage,
|
||||
TextPart,
|
||||
VideoPart,
|
||||
)
|
||||
from zhenxun.services.ai.llm.manager import get_default_model
|
||||
from zhenxun.services.log import logger
|
||||
from zhenxun.utils.pydantic_compat import model_copy
|
||||
|
||||
|
||||
class MultimodalPlaceholderReducer(BaseMemoryReducer):
|
||||
"""视觉媒体降级：将超过一定轮数的老图片/视频替换为 <图片> 占位符文本"""
|
||||
|
||||
def __init__(self, window_size: int = 5):
|
||||
"""
|
||||
初始化多模态占位符修剪器。
|
||||
|
||||
参数：
|
||||
window_size: 多模态视窗大小。在保留最近的指定数量的多模态消息对后，
|
||||
超出的旧消息中多模态内容（如图片、视频等）将被替换为占位文本。
|
||||
"""
|
||||
self.window_size = window_size
|
||||
|
||||
@staticmethod
|
||||
def apply_multimodal_placeholder(message: LLMMessage) -> LLMMessage:
|
||||
sanitized_message = model_copy(message, deep=False)
|
||||
new_content_parts = []
|
||||
|
||||
for part in sanitized_message.content:
|
||||
if isinstance(part, ImagePart):
|
||||
new_content_parts.append(TextPart(text="<图片>"))
|
||||
elif isinstance(part, AudioPart):
|
||||
new_content_parts.append(TextPart(text="<音频>"))
|
||||
elif isinstance(part, VideoPart):
|
||||
new_content_parts.append(TextPart(text="<视频>"))
|
||||
elif isinstance(part, FilePart):
|
||||
new_content_parts.append(TextPart(text="<文件>"))
|
||||
elif isinstance(part, TextPart) and "[多模态内容:" in part.text:
|
||||
new_content_parts.append(TextPart(text="<图片>"))
|
||||
else:
|
||||
new_content_parts.append(part)
|
||||
|
||||
merged_parts = []
|
||||
for part in new_content_parts:
|
||||
if (
|
||||
isinstance(part, TextPart)
|
||||
and merged_parts
|
||||
and isinstance(merged_parts[-1], TextPart)
|
||||
):
|
||||
new_text = (merged_parts[-1].text or "") + " " + (part.text or "")
|
||||
merged_parts[-1] = TextPart(text=new_text.strip())
|
||||
else:
|
||||
merged_parts.append(part)
|
||||
|
||||
sanitized_message.content = merged_parts
|
||||
sanitized_message.token_cost = None
|
||||
return sanitized_message
|
||||
|
||||
async def reduce(self, messages, current_tokens, model_name, base_overhead=0):
|
||||
if self.window_size <= 0:
|
||||
return messages, False, current_tokens
|
||||
|
||||
processed_messages = []
|
||||
user_multimodal_count = 0
|
||||
changed = False
|
||||
|
||||
for msg in reversed(messages):
|
||||
has_multimodal = False
|
||||
if isinstance(msg.content, list):
|
||||
has_multimodal = any(
|
||||
isinstance(p, ImagePart | AudioPart | VideoPart | FilePart)
|
||||
or (isinstance(p, TextPart) and "[多模态内容:" in p.text)
|
||||
for p in msg.content
|
||||
)
|
||||
|
||||
if has_multimodal:
|
||||
if msg.role == "user":
|
||||
user_multimodal_count += 1
|
||||
if user_multimodal_count > self.window_size:
|
||||
processed_messages.append(self.apply_multimodal_placeholder(msg))
|
||||
changed = True
|
||||
else:
|
||||
processed_messages.append(msg)
|
||||
else:
|
||||
processed_messages.append(msg)
|
||||
|
||||
if not changed:
|
||||
return messages, False, current_tokens
|
||||
|
||||
processed_messages.reverse()
|
||||
new_tokens = token_counter.count_context(
|
||||
processed_messages, model_name, base_overhead
|
||||
)
|
||||
return processed_messages, True, new_tokens
|
||||
|
||||
|
||||
class MessageDropper(BaseMemoryReducer):
|
||||
"""消息丢弃器：在 Token 超过阈值时丢弃最早的非置顶消息对。"""
|
||||
|
||||
def __init__(self, trigger_tokens: int = 4000):
|
||||
"""
|
||||
初始化消息丢弃器。
|
||||
|
||||
参数：
|
||||
trigger_tokens: 触发丢弃策略的 Token 阈值上限。
|
||||
当当前对话 Token 总数超过此值时，将触发硬截断。
|
||||
"""
|
||||
self.trigger_tokens = trigger_tokens
|
||||
|
||||
async def reduce(self, messages, current_tokens, model_name, base_overhead=0):
|
||||
if current_tokens <= self.trigger_tokens:
|
||||
return messages, False, current_tokens
|
||||
|
||||
logger.info(
|
||||
"✂️ [MemoryCompression] 触发硬截断丢弃策略 | 原因: "
|
||||
f"当前 Token 预估 ({current_tokens}) 仍超过硬性上限 ({self.trigger_tokens})，" # noqa: E501
|
||||
"开始丢弃最旧的历史对话..."
|
||||
)
|
||||
new_messages = list(messages)
|
||||
changed = False
|
||||
|
||||
while current_tokens > self.trigger_tokens:
|
||||
user_indices = [
|
||||
i
|
||||
for i, m in enumerate(new_messages)
|
||||
if m.role == "user"
|
||||
and not (m.metadata and m.metadata.get("pinned", False))
|
||||
]
|
||||
|
||||
if len(user_indices) < 2:
|
||||
break
|
||||
|
||||
start_idx = user_indices[0]
|
||||
end_idx = user_indices[1]
|
||||
|
||||
del new_messages[start_idx:end_idx]
|
||||
changed = True
|
||||
|
||||
current_tokens = token_counter.count_context(
|
||||
new_messages, model_name, base_overhead
|
||||
)
|
||||
return new_messages, changed, current_tokens
|
||||
|
||||
|
||||
class ToolPrunerReducer(BaseMemoryReducer):
|
||||
"""工具结果修剪器：纯粹计算工具输出的 Token 和轮数，超标时剔除老旧工具返回结果"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
keep_recent_turns: int = 3,
|
||||
trigger_tokens: int = 4000,
|
||||
max_turns: int = 0,
|
||||
):
|
||||
"""
|
||||
初始化工具结果修剪器。
|
||||
|
||||
参数：
|
||||
keep_recent_turns: 保留最近的工具调用轮数（不进行内容截断的轮数）。
|
||||
trigger_tokens: 触发工具修剪策略的工具总 Token 阈值上限。
|
||||
max_turns: 触发工具修剪的最大工具调用轮数上限。若为 0，则不限制轮数。
|
||||
"""
|
||||
self.keep_recent_turns = keep_recent_turns
|
||||
self.trigger_tokens = trigger_tokens
|
||||
self.max_turns = max_turns
|
||||
|
||||
async def reduce(self, messages, current_tokens, model_name, base_overhead=0):
|
||||
tool_msgs = [m for m in messages if m.role == "tool"]
|
||||
tool_turns = len(tool_msgs)
|
||||
|
||||
if tool_turns == 0:
|
||||
return messages, False, current_tokens
|
||||
|
||||
tool_tokens = sum(token_counter.count_message(m, model_name) for m in tool_msgs)
|
||||
|
||||
is_token_exceeded = tool_tokens > self.trigger_tokens
|
||||
is_turn_exceeded = self.max_turns > 0 and tool_turns > self.max_turns
|
||||
|
||||
if not (is_token_exceeded or is_turn_exceeded):
|
||||
return messages, False, current_tokens
|
||||
|
||||
reasons = []
|
||||
if is_token_exceeded:
|
||||
reasons.append(f"工具Token超标 ({tool_tokens} > {self.trigger_tokens})")
|
||||
if is_turn_exceeded:
|
||||
reasons.append(f"工具调用轮数超限 ({tool_turns} > {self.max_turns})")
|
||||
|
||||
logger.info(
|
||||
f"✂️ [MemoryCompression] 触发工具结果修剪策略 | 原因: {' 且 '.join(reasons)}"
|
||||
)
|
||||
|
||||
from zhenxun.services.ai.core.messages import ToolReturnPart
|
||||
|
||||
new_messages = []
|
||||
tools_kept = 0
|
||||
changed = False
|
||||
|
||||
for msg in reversed(messages):
|
||||
if msg.role != "tool":
|
||||
new_messages.append(msg)
|
||||
continue
|
||||
|
||||
if tools_kept < self.keep_recent_turns:
|
||||
tools_kept += 1
|
||||
new_messages.append(msg)
|
||||
continue
|
||||
|
||||
new_content = []
|
||||
part_changed = False
|
||||
for p in msg.content:
|
||||
if isinstance(p, ToolReturnPart):
|
||||
old_len = len(str(p.output))
|
||||
new_p = model_copy(
|
||||
p,
|
||||
update={
|
||||
"output": f"[数据过载自动截断 - 原长度: {old_len} 字符]"
|
||||
},
|
||||
)
|
||||
new_content.append(new_p)
|
||||
part_changed = True
|
||||
changed = True
|
||||
else:
|
||||
new_content.append(p)
|
||||
|
||||
if part_changed:
|
||||
new_msg = model_copy(
|
||||
msg, update={"content": new_content, "token_cost": None}
|
||||
)
|
||||
new_messages.append(new_msg)
|
||||
else:
|
||||
new_messages.append(msg)
|
||||
|
||||
if not changed:
|
||||
return messages, False, current_tokens
|
||||
|
||||
new_messages.reverse()
|
||||
new_total = token_counter.count_context(new_messages, model_name, base_overhead)
|
||||
return new_messages, True, new_total
|
||||
|
||||
|
||||
class AbstractSummarizerReducer(BaseMemoryReducer):
|
||||
"""抽象总结压缩器：提取阈值判断与上下文分流的公共逻辑"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
strategy_name: str,
|
||||
keep_recent_turns: int = 0,
|
||||
trigger_tokens: int = 4000,
|
||||
max_turns: int | None = None,
|
||||
summarization_model: str | None = None,
|
||||
):
|
||||
"""
|
||||
初始化抽象总结压缩基类。
|
||||
|
||||
参数：
|
||||
strategy_name: 压缩策略的名称，用于日志输出和追踪。
|
||||
keep_recent_turns: 压缩时需要保留的最新的对话轮数（不参与总结的轮数）。
|
||||
trigger_tokens: 触发总结策略的 Token 阈值上限。
|
||||
max_turns: 触发总结策略的最大对话轮数上限。
|
||||
summarization_model: 用于执行总结压缩大模型请求的模型名称，若为 None 则使用默认模型。
|
||||
""" # noqa: E501
|
||||
self.strategy_name = strategy_name
|
||||
self.keep_recent_turns = keep_recent_turns
|
||||
self.trigger_tokens = trigger_tokens
|
||||
self.max_turns = max_turns
|
||||
self.summarization_model = summarization_model
|
||||
|
||||
@abstractmethod
|
||||
async def _execute_summarization(
|
||||
self, to_summarize: list[LLMMessage], prev_summary: str
|
||||
) -> LLMMessage | None:
|
||||
"""由子类实现具体的 LLM 调用逻辑，返回新的总结消息"""
|
||||
pass
|
||||
|
||||
async def reduce(self, messages, current_tokens, model_name, base_overhead=0):
|
||||
user_turns = sum(
|
||||
1
|
||||
for m in messages
|
||||
if m.role == "user"
|
||||
and not (m.metadata and m.metadata.get("is_summary", False))
|
||||
)
|
||||
is_token_exceeded = current_tokens > self.trigger_tokens
|
||||
is_turn_exceeded = (
|
||||
self.max_turns is not None
|
||||
and self.max_turns > 0
|
||||
and user_turns > self.max_turns
|
||||
)
|
||||
|
||||
if not (is_token_exceeded or is_turn_exceeded):
|
||||
return messages, False, current_tokens
|
||||
|
||||
reasons = []
|
||||
if is_token_exceeded:
|
||||
reasons.append(f"Token 预估超限 ({current_tokens} > {self.trigger_tokens})")
|
||||
if is_turn_exceeded:
|
||||
reasons.append(f"有效对话轮次超限 ({user_turns} > {self.max_turns})")
|
||||
logger.info(
|
||||
f"🔄 [MemoryCompression] 触发{self.strategy_name}策略 | 原因: "
|
||||
f"{' 且 '.join(reasons)}"
|
||||
)
|
||||
|
||||
pinned_msgs, working_msgs, prev_summary = [], [], ""
|
||||
for msg in messages:
|
||||
is_pinned = isinstance(msg, SystemMessage) or (
|
||||
msg.metadata and msg.metadata.get("pinned", False)
|
||||
)
|
||||
if msg.metadata and msg.metadata.get("is_summary", False):
|
||||
prev_summary = msg.extract_text
|
||||
elif is_pinned:
|
||||
pinned_msgs.append(msg)
|
||||
else:
|
||||
working_msgs.append(msg)
|
||||
|
||||
user_indices = [i for i, m in enumerate(working_msgs) if m.role == "user"]
|
||||
|
||||
if len(user_indices) <= self.keep_recent_turns:
|
||||
return messages, False, current_tokens
|
||||
|
||||
split_idx = (
|
||||
user_indices[-self.keep_recent_turns]
|
||||
if self.keep_recent_turns > 0
|
||||
else len(working_msgs)
|
||||
)
|
||||
to_summarize = working_msgs[:split_idx]
|
||||
to_keep = working_msgs[split_idx:]
|
||||
|
||||
new_summary_msg = await self._execute_summarization(to_summarize, prev_summary)
|
||||
if not new_summary_msg:
|
||||
return messages, False, current_tokens
|
||||
|
||||
new_messages = [*pinned_msgs, new_summary_msg, *to_keep]
|
||||
return (
|
||||
new_messages,
|
||||
True,
|
||||
token_counter.count_context(new_messages, model_name, base_overhead),
|
||||
)
|
||||
|
||||
|
||||
class LLMSummarizerReducer(AbstractSummarizerReducer):
|
||||
"""大模型总结压缩器：将较早的历史对话记录通过 LLM 压缩合并为一段文本摘要。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
keep_recent_turns: int = 0,
|
||||
trigger_tokens: int = 4000,
|
||||
max_turns: int | None = None,
|
||||
summarization_model: str | None = None,
|
||||
summarization_prompt: str = (
|
||||
"请概括以下对话内容，保留关键的约束条件、用户偏好、"
|
||||
"已完成的任务状态和未解决的问题。"
|
||||
),
|
||||
):
|
||||
"""
|
||||
初始化大模型总结压缩器。
|
||||
|
||||
参数：
|
||||
keep_recent_turns: 压缩时需要保留的最新的对话轮数。
|
||||
trigger_tokens: 触发总结策略的 Token 阈值上限。
|
||||
max_turns: 触发总结策略的最大对话轮数上限。
|
||||
summarization_model: 用于执行总结压缩的大模型名称。
|
||||
summarization_prompt: 发送给大模型的总结引导 Prompt 提示词。
|
||||
"""
|
||||
super().__init__(
|
||||
strategy_name="历史对话合并总结",
|
||||
keep_recent_turns=keep_recent_turns,
|
||||
trigger_tokens=trigger_tokens,
|
||||
max_turns=max_turns,
|
||||
summarization_model=summarization_model,
|
||||
)
|
||||
self.summarization_prompt = summarization_prompt
|
||||
|
||||
async def _execute_summarization(
|
||||
self, to_summarize: list[LLMMessage], prev_summary: str
|
||||
) -> LLMMessage | None:
|
||||
prompt_text = f"### 📋 [对话摘要任务]\n{self.summarization_prompt}\n\n"
|
||||
if prev_summary:
|
||||
prompt_text += "#### önceki_summary (参考先前的快照):\n"
|
||||
prompt_text += f"> {prev_summary}\n\n"
|
||||
prompt_text += "#### 待处理的历史消息流：\n"
|
||||
for m in to_summarize:
|
||||
c_str = m.extract_text[:1500]
|
||||
speaker = m.source_name if m.source_name else m.role.capitalize()
|
||||
prompt_text += f"[{speaker}]: {c_str}\n"
|
||||
prompt_text += "</需要合并的旧对话记录>\n"
|
||||
|
||||
from zhenxun.services.ai.llm.api import chat
|
||||
|
||||
try:
|
||||
model_to_use = self.summarization_model or get_default_model("chat")
|
||||
response = await chat(
|
||||
prompt_text,
|
||||
model=model_to_use,
|
||||
instruction="你是后台记忆整理引擎。请客观、简明输出当前对话全局摘要。",
|
||||
)
|
||||
new_summary_msg = LLMMessage.assistant_text_response(
|
||||
f"【历史对话摘要记忆】\n{response.text}"
|
||||
)
|
||||
new_summary_msg.metadata = {"is_summary": True, "pinned": True}
|
||||
return new_summary_msg
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[{self.__class__.__name__}] 压缩总结调用失败，已跳过本次压缩: {e}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
_T_Summary = TypeVar("_T_Summary", bound=BaseModel)
|
||||
|
||||
|
||||
class StructuredSummaryReducer(AbstractSummarizerReducer, Generic[_T_Summary]):
|
||||
"""结构化总结压缩器：基于 JSON Schema 格式化抽取长上下文状态信息并合并"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
response_model: type[_T_Summary],
|
||||
prompt_template: str,
|
||||
format_callback: Callable[[_T_Summary], str],
|
||||
keep_recent_turns: int = 0,
|
||||
trigger_tokens: int = 4000,
|
||||
max_turns: int | None = None,
|
||||
summarization_model: str | None = None,
|
||||
instruction: str = (
|
||||
"请提取并合并先前的状态和最新的对话内容，保持精简，不要编造事实"
|
||||
),
|
||||
):
|
||||
"""
|
||||
初始化结构化总结压缩器。
|
||||
|
||||
参数：
|
||||
response_model: 接收结构化输出的 Pydantic 模型类，需继承自 BaseModel。
|
||||
prompt_template: 用于抽取合并状态的 Prompt 模板，包含 {prev_summary} 和 {dialogue} 占位符。
|
||||
format_callback: 格式化回调函数，用于将结构化 Pydantic 响应对象转换为便于大模型阅读的字符串。
|
||||
keep_recent_turns: 压缩时需要保留的最新的对话轮数。
|
||||
trigger_tokens: 触发总结策略的 Token 阈值上限。
|
||||
max_turns: 触发总结策略的最大对话轮数上限。
|
||||
summarization_model: 用于执行总结压缩的大模型名称。
|
||||
instruction: 指导大模型生成结构化数据时的系统指令说明。
|
||||
""" # noqa: E501
|
||||
super().__init__(
|
||||
strategy_name="结构化状态抽取压缩",
|
||||
keep_recent_turns=keep_recent_turns,
|
||||
trigger_tokens=trigger_tokens,
|
||||
max_turns=max_turns,
|
||||
summarization_model=summarization_model,
|
||||
)
|
||||
self.response_model = response_model
|
||||
self.prompt_template = prompt_template
|
||||
self.format_callback = format_callback
|
||||
self.instruction = instruction
|
||||
|
||||
async def _execute_summarization(
|
||||
self, to_summarize: list[LLMMessage], prev_summary: str
|
||||
) -> LLMMessage | None:
|
||||
dialogue_text = ""
|
||||
for m in to_summarize:
|
||||
c_str = m.extract_text[:1500]
|
||||
speaker = m.source_name if m.source_name else m.role.capitalize()
|
||||
dialogue_text += f"[{speaker}]: {c_str}\n"
|
||||
|
||||
prompt_text = self.prompt_template.format(
|
||||
prev_summary=prev_summary, dialogue=dialogue_text
|
||||
)
|
||||
|
||||
from zhenxun.services.ai.llm.api import generate_structured
|
||||
|
||||
try:
|
||||
model_to_use = self.summarization_model or get_default_model("chat")
|
||||
summary_obj = await generate_structured(
|
||||
prompt_text,
|
||||
response_model=self.response_model,
|
||||
model=model_to_use,
|
||||
instruction=self.instruction,
|
||||
)
|
||||
|
||||
summary_text = self.format_callback(summary_obj)
|
||||
|
||||
new_summary_msg = LLMMessage.assistant_text_response(
|
||||
f"【历史状态摘要记忆】\n{summary_text}"
|
||||
)
|
||||
new_summary_msg.metadata = {"is_summary": True, "pinned": True}
|
||||
return new_summary_msg
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[{self.__class__.__name__}] 结构化总结失败，已跳过本次压缩: {e}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class CondenserPipeline:
|
||||
"""上下文压缩流水线：按顺序依次执行各阶段的记忆压缩减项。"""
|
||||
|
||||
def __init__(self, reducers: list[BaseMemoryReducer]):
|
||||
"""
|
||||
初始化上下文压缩流水线。
|
||||
|
||||
参数：
|
||||
reducers: 压缩减项器列表，将按顺序对记忆进行多阶段修剪和压缩。
|
||||
"""
|
||||
self.reducers = reducers
|
||||
|
||||
@classmethod
|
||||
def create_from_configs(
|
||||
cls, memory_config: Any, model_name: str
|
||||
) -> "CondenserPipeline":
|
||||
"""基于全局和局部配置组装压缩管线工厂方法"""
|
||||
from zhenxun.services.ai.config import get_llm_config
|
||||
from zhenxun.services.ai.llm.system.capabilities import get_model_capabilities
|
||||
|
||||
config = get_llm_config().context_settings
|
||||
pipeline_reducers = []
|
||||
caps = get_model_capabilities(model_name)
|
||||
|
||||
vw = config.vision_window_size
|
||||
if memory_config and memory_config.compression.vision_window is not None:
|
||||
vw = memory_config.compression.vision_window
|
||||
if vw > 0:
|
||||
pipeline_reducers.append(MultimodalPlaceholderReducer(window_size=vw))
|
||||
|
||||
tp = config.tool_pruning
|
||||
if tp.enable:
|
||||
tp_limit = (
|
||||
int(caps.max_input_tokens * tp.trigger_threshold)
|
||||
if tp.trigger_threshold <= 1.0
|
||||
else int(tp.trigger_threshold)
|
||||
)
|
||||
pipeline_reducers.append(
|
||||
ToolPrunerReducer(
|
||||
keep_recent_turns=tp.keep_recent_turns,
|
||||
trigger_tokens=tp_limit,
|
||||
max_turns=tp.max_history_turns,
|
||||
)
|
||||
)
|
||||
|
||||
policy = memory_config.compression.policy if memory_config else None
|
||||
if policy is not None:
|
||||
pipeline_reducers.extend(policy)
|
||||
else:
|
||||
threshold = config.llm_summary.trigger_threshold
|
||||
if memory_config and memory_config.compression.threshold is not None:
|
||||
threshold = memory_config.compression.threshold
|
||||
|
||||
limit = (
|
||||
int(caps.max_input_tokens * threshold)
|
||||
if threshold <= 1.0
|
||||
else int(threshold)
|
||||
)
|
||||
|
||||
max_turns = config.llm_summary.max_history_turns
|
||||
if (
|
||||
memory_config
|
||||
and memory_config.compression.max_history_turns is not None
|
||||
):
|
||||
max_turns = memory_config.compression.max_history_turns
|
||||
|
||||
if config.llm_summary.enable:
|
||||
pipeline_reducers.extend(
|
||||
MemoryPolicy.llm_summarize(
|
||||
trigger_tokens=limit,
|
||||
max_turns=max_turns,
|
||||
keep_recent_turns=config.llm_summary.keep_recent_turns,
|
||||
summarization_model=config.llm_summary.summarization_model,
|
||||
summarization_prompt=config.llm_summary.summarization_prompt,
|
||||
)
|
||||
)
|
||||
else:
|
||||
pipeline_reducers.extend(MemoryPolicy.unlimited())
|
||||
|
||||
return cls(pipeline_reducers)
|
||||
|
||||
async def run(
|
||||
self, messages, model_name, base_overhead=0
|
||||
) -> tuple[list[LLMMessage], bool]:
|
||||
current_tokens = token_counter.count_context(
|
||||
messages, model_name, base_overhead
|
||||
)
|
||||
|
||||
current_messages = messages
|
||||
any_changed = False
|
||||
for reducer in self.reducers:
|
||||
current_messages, changed, current_tokens = await reducer.reduce(
|
||||
current_messages,
|
||||
current_tokens,
|
||||
model_name,
|
||||
base_overhead,
|
||||
)
|
||||
if changed:
|
||||
any_changed = True
|
||||
return current_messages, any_changed
|
||||
|
||||
|
||||
class MemoryPolicy:
|
||||
"""
|
||||
记忆策略工厂 (Strategy Factory Facade)。
|
||||
为开发者提供开箱即用的上下文压缩管线组装方案。
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def unlimited() -> list[BaseMemoryReducer]:
|
||||
"""无限制模式。不进行任何形式的截断和总结，适用于短对话或纯 Agent 内部流转。"""
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def llm_summarize(
|
||||
trigger_tokens: int = 4000,
|
||||
max_turns: int | None = None,
|
||||
keep_recent_turns: int = 0,
|
||||
summarization_model: str | None = None,
|
||||
summarization_prompt: str = (
|
||||
"请概括以下对话内容，保留关键的约束条件、用户偏好、"
|
||||
"已完成的任务状态和未解决的问题。"
|
||||
),
|
||||
) -> list[BaseMemoryReducer]:
|
||||
"""LLM 总结压缩模式。Token 达标后，自动将历史对话合并为一段 Summary。"""
|
||||
return [
|
||||
LLMSummarizerReducer(
|
||||
keep_recent_turns=keep_recent_turns,
|
||||
trigger_tokens=trigger_tokens,
|
||||
max_turns=max_turns,
|
||||
summarization_model=summarization_model,
|
||||
summarization_prompt=summarization_prompt,
|
||||
),
|
||||
MessageDropper(trigger_tokens=trigger_tokens),
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def structured_summarize(
|
||||
trigger_tokens: int = 4000,
|
||||
max_turns: int | None = None,
|
||||
keep_recent_turns: int = 0,
|
||||
summarization_model: str | None = None,
|
||||
response_model: type[BaseModel] | None = None,
|
||||
prompt_template: str | None = None,
|
||||
format_callback: Callable[[Any], str] | None = None,
|
||||
) -> list[BaseMemoryReducer]:
|
||||
"""结构化总结压缩模式。使用 JSON Schema 强制大模型提取核心状态。"""
|
||||
|
||||
class DefaultStateSummary(BaseModel):
|
||||
user_context: str = Field(
|
||||
description="用户的核心意图、诉求、人设或长期记忆规则。"
|
||||
)
|
||||
completed_tasks: str = Field(description="已完成的操作或已经确认的情节。")
|
||||
pending_tasks: str = Field(description="正在进行中的任务或尚未解答的问题。")
|
||||
current_state: str = Field(
|
||||
description="当前状态，如重要变量、玩家血量、关键物品坐标等。"
|
||||
)
|
||||
|
||||
def default_format(obj: DefaultStateSummary) -> str:
|
||||
return (
|
||||
f"👤 用户上下文: {obj.user_context}\n"
|
||||
f"✅ 已完成/确认: {obj.completed_tasks}\n"
|
||||
f"⏳ 待处理/疑问: {obj.pending_tasks}\n"
|
||||
f"📌 当前状态: {obj.current_state}"
|
||||
)
|
||||
|
||||
default_prompt = (
|
||||
"你是一个专门用于长上下文状态压缩的引擎。请阅读以下先前的总结和旧对话，"
|
||||
"提取核心状态信息，并合并它们。\n\n"
|
||||
"<之前的状态摘要>\n{prev_summary}\n</之前的状态摘要>\n\n"
|
||||
"<需要合并的旧对话记录>\n"
|
||||
"{dialogue}"
|
||||
"</需要合并的旧对话记录>\n"
|
||||
)
|
||||
|
||||
return [
|
||||
StructuredSummaryReducer(
|
||||
response_model=response_model or DefaultStateSummary,
|
||||
prompt_template=prompt_template or default_prompt,
|
||||
format_callback=format_callback or default_format,
|
||||
keep_recent_turns=keep_recent_turns,
|
||||
trigger_tokens=trigger_tokens,
|
||||
max_turns=max_turns,
|
||||
summarization_model=summarization_model,
|
||||
),
|
||||
MessageDropper(trigger_tokens=trigger_tokens),
|
||||
]
|
||||
@@ -0,0 +1,260 @@
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, cast
|
||||
|
||||
from zhenxun.services.ai.context.memory.compression import (
|
||||
CondenserPipeline,
|
||||
)
|
||||
from zhenxun.services.ai.context.memory.manager import memory_manager
|
||||
from zhenxun.services.ai.context.memory.models import MemoryConfig
|
||||
from zhenxun.services.ai.context.memory.types import SessionMetadata
|
||||
from zhenxun.services.ai.core.engine.context_renderer import ContextConverter
|
||||
from zhenxun.services.ai.core.messages import AgentMessage
|
||||
from zhenxun.services.log import logger
|
||||
from zhenxun.utils.pydantic_compat import model_copy
|
||||
|
||||
|
||||
class MemoryReader:
|
||||
"""
|
||||
记忆读取器 (Memory Reader)。
|
||||
负责从数据库中提取短期上下文历史，召回长期的背景知识，并执行自动压缩。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, session_meta: SessionMetadata, memory_config: MemoryConfig | None
|
||||
):
|
||||
"""
|
||||
初始化记忆读取器。
|
||||
|
||||
参数:
|
||||
session_meta: 会话元数据，包含 Namespace 与作用域映射等上下文信息。
|
||||
memory_config: 记忆系统的配置对象，控制长期、短期及槽位记忆的启用与逻辑。
|
||||
"""
|
||||
self.session_meta = session_meta
|
||||
self.memory_config = memory_config
|
||||
|
||||
async def get_long_term_context(self, user_input: str) -> str:
|
||||
"""
|
||||
基于用户输入召回长期记忆（RAG），返回格式化后的背景提示词。
|
||||
"""
|
||||
if (
|
||||
not self.memory_config
|
||||
or not self.memory_config.long_term.enable
|
||||
or not user_input
|
||||
):
|
||||
return ""
|
||||
|
||||
policy = self.memory_config.long_term.auto_recall
|
||||
should_recall = False
|
||||
|
||||
if isinstance(policy, bool):
|
||||
should_recall = policy
|
||||
elif callable(policy):
|
||||
import inspect
|
||||
|
||||
try:
|
||||
res = policy(user_input, self.session_meta)
|
||||
if inspect.isawaitable(res):
|
||||
should_recall = await res
|
||||
else:
|
||||
should_recall = bool(res)
|
||||
except Exception as e:
|
||||
logger.error(f"[MemoryReader] 自定义 auto_recall 函数执行失败: {e}")
|
||||
should_recall = False
|
||||
|
||||
if not should_recall:
|
||||
return ""
|
||||
|
||||
ltm_scope = memory_manager.get_long_term_memory(
|
||||
self.memory_config,
|
||||
namespace=self.session_meta.selector.namespace or "global",
|
||||
)
|
||||
if not ltm_scope:
|
||||
return ""
|
||||
|
||||
matches = await ltm_scope.recall(session=self.session_meta, query=user_input)
|
||||
if matches:
|
||||
logger.debug(f"🧠 [MemoryReader] 长期记忆召回详情 (Query: '{user_input}'):")
|
||||
for i, m in enumerate(matches):
|
||||
logger.debug(
|
||||
f" [{i + 1}] 得分: {m.score:.4f} | 内容: {m.record.content}"
|
||||
)
|
||||
|
||||
threshold = self.memory_config.long_term.recall_threshold
|
||||
valid_matches = [m for m in matches if m.score >= threshold]
|
||||
if not valid_matches:
|
||||
logger.debug("🧠 [MemoryReader] 召回的记忆均未达到相关性阈值，已丢弃。")
|
||||
return ""
|
||||
|
||||
fact_str = "\n".join(f"- {m.record.content}" for m in valid_matches)
|
||||
logger.debug(
|
||||
f"🧠 [MemoryReader]"
|
||||
f"成功截取并注入 {len(valid_matches)} 条高价值长期记忆。"
|
||||
)
|
||||
return f"[系统补充：有关用户的长期记忆设定]\n{fact_str}"
|
||||
return ""
|
||||
|
||||
async def get_slots_context(self) -> str:
|
||||
"""
|
||||
读取并组装核心槽位记忆 (Memory Slots)，返回 XML 格式字符串供大模型使用。
|
||||
"""
|
||||
if not self.memory_config or not self.memory_config.slots.enable:
|
||||
return ""
|
||||
slot_ctx = memory_manager.get_slot_context(
|
||||
self.memory_config,
|
||||
namespace=self.session_meta.selector.namespace or "global",
|
||||
)
|
||||
if not slot_ctx:
|
||||
return ""
|
||||
|
||||
if self.memory_config.slots.default_slots:
|
||||
for default_slot in self.memory_config.slots.default_slots:
|
||||
existing = await slot_ctx.get_slot(
|
||||
self.session_meta, default_slot.label
|
||||
)
|
||||
if not existing:
|
||||
await slot_ctx.set_slot(self.session_meta, default_slot)
|
||||
|
||||
slots = await slot_ctx.list_pinned_slots(self.session_meta)
|
||||
if not slots:
|
||||
return ""
|
||||
|
||||
show_scope = False
|
||||
if (
|
||||
self.memory_config
|
||||
and self.memory_config.slots.scopes
|
||||
and len(self.memory_config.slots.scopes) > 1
|
||||
):
|
||||
show_scope = True
|
||||
|
||||
xml_parts = ["<memory_slots>"]
|
||||
for slot in slots:
|
||||
if show_scope:
|
||||
semantic_name = self.session_meta.scope_name_mapping.get(
|
||||
slot.scope, "未知"
|
||||
)
|
||||
xml_parts.append(
|
||||
f' <slot name="{slot.label}" scope="{semantic_name}">\n'
|
||||
f" {slot.content}\n"
|
||||
" </slot>"
|
||||
)
|
||||
else:
|
||||
xml_parts.append(
|
||||
f' <slot name="{slot.label}">\n {slot.content}\n </slot>'
|
||||
)
|
||||
xml_parts.append("</memory_slots>")
|
||||
return "\n".join(xml_parts)
|
||||
|
||||
async def get_short_term_context(
|
||||
self,
|
||||
model_name: str,
|
||||
override_history: Sequence[AgentMessage] | None = None,
|
||||
) -> list[AgentMessage]:
|
||||
"""
|
||||
拉取短期对话历史，并执行 Token 压缩。
|
||||
"""
|
||||
current_history: list[AgentMessage] = []
|
||||
if override_history is not None:
|
||||
current_history = list(override_history)
|
||||
|
||||
chat_context = memory_manager.get_chat_context(
|
||||
self.memory_config,
|
||||
namespace=self.session_meta.selector.namespace or "global",
|
||||
)
|
||||
|
||||
if self.memory_config and self.memory_config.short_term.enable and chat_context:
|
||||
if override_history is not None:
|
||||
flattened_override = ContextConverter.flatten_to_llm_messages(
|
||||
override_history
|
||||
)
|
||||
await chat_context.set_messages(self.session_meta, flattened_override)
|
||||
else:
|
||||
current_history = cast(
|
||||
list[AgentMessage],
|
||||
await chat_context.get_messages(self.session_meta),
|
||||
)
|
||||
|
||||
pipeline = CondenserPipeline.create_from_configs(
|
||||
self.memory_config, model_name
|
||||
)
|
||||
if pipeline.reducers:
|
||||
flattened_to_reduce = ContextConverter.flatten_to_llm_messages(
|
||||
current_history
|
||||
)
|
||||
new_history, changed = await pipeline.run(
|
||||
flattened_to_reduce, model_name=model_name, base_overhead=0
|
||||
)
|
||||
if changed:
|
||||
await chat_context.set_messages(self.session_meta, new_history)
|
||||
logger.info(
|
||||
"💾 [MemoryReader] 压缩截断完毕，已同步覆写数据库。"
|
||||
f"压缩后条数: {len(new_history)}"
|
||||
)
|
||||
current_history = cast(list[AgentMessage], new_history)
|
||||
|
||||
return current_history
|
||||
|
||||
|
||||
class MemoryWriter:
|
||||
"""
|
||||
记忆写入器 (Memory Writer)。
|
||||
负责将对话增量安全地写入数据库。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_meta: SessionMetadata,
|
||||
memory_config: MemoryConfig | None,
|
||||
context: Any = None,
|
||||
):
|
||||
"""
|
||||
初始化记忆写入器。
|
||||
|
||||
参数:
|
||||
session_meta: 会话元数据，包含 Namespace 与作用域映射等上下文信息。
|
||||
memory_config: 记忆系统的配置对象，控制记忆存入的逻辑。
|
||||
context: 运行时上下文环境，作为可选参数传入，供中间件使用，默认 None。
|
||||
"""
|
||||
self.session_meta = session_meta
|
||||
self.memory_config = memory_config
|
||||
self.context = context
|
||||
|
||||
async def save_new_messages(
|
||||
self,
|
||||
new_messages: Sequence[AgentMessage],
|
||||
):
|
||||
"""将新产生的对话增量保存到数据库"""
|
||||
if not new_messages:
|
||||
return
|
||||
|
||||
messages_to_save = new_messages
|
||||
|
||||
if self.memory_config and self.memory_config.ingestion.middlewares:
|
||||
messages_to_save = [model_copy(m, deep=True) for m in new_messages]
|
||||
|
||||
for middleware in self.memory_config.ingestion.middlewares:
|
||||
try:
|
||||
messages_to_save = await middleware.process(
|
||||
messages_to_save, self.context
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[MemoryIngestion] 中间件 {middleware.__class__.__name__} "
|
||||
f"执行失败: {e}",
|
||||
e=e,
|
||||
)
|
||||
|
||||
if not messages_to_save:
|
||||
return
|
||||
|
||||
chat_ctx = memory_manager.get_chat_context(
|
||||
self.memory_config,
|
||||
namespace=self.session_meta.selector.namespace or "global",
|
||||
)
|
||||
|
||||
flattened_msgs = ContextConverter.flatten_to_llm_messages(
|
||||
messages_to_save, self.context
|
||||
)
|
||||
|
||||
if chat_ctx and self.memory_config and self.memory_config.short_term.enable:
|
||||
if flattened_msgs:
|
||||
await chat_ctx.add_messages(self.session_meta, flattened_msgs)
|
||||
@@ -0,0 +1,128 @@
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
from zhenxun.services.ai.context.memory.types import (
|
||||
MemorySlot,
|
||||
SessionMetadata,
|
||||
)
|
||||
from zhenxun.services.ai.core.messages import AgentMessage, LLMMessage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from zhenxun.services.ai.context.memory.manager import GlobalMemoryManager
|
||||
from zhenxun.services.ai.context.memory.storage.interfaces import (
|
||||
BaseChatContext,
|
||||
BaseSlotContext,
|
||||
)
|
||||
|
||||
|
||||
class ChatHistoryFacade:
|
||||
"""短期对话历史门面"""
|
||||
|
||||
def __init__(self, manager: "GlobalMemoryManager", session_meta: SessionMetadata):
|
||||
self.manager = manager
|
||||
self.session_meta = session_meta
|
||||
|
||||
@property
|
||||
def _backend(self) -> "BaseChatContext | None":
|
||||
return self.manager.get_chat_context(
|
||||
None, self.session_meta.namespace or "global"
|
||||
)
|
||||
|
||||
async def get(self, limit: int | None = None) -> list[LLMMessage]:
|
||||
"""获取当前会话的历史消息"""
|
||||
if not self._backend:
|
||||
return []
|
||||
msgs = await self._backend.get_messages(self.session_meta)
|
||||
return msgs[-limit:] if limit else msgs
|
||||
|
||||
async def add(self, messages: Sequence[AgentMessage] | AgentMessage) -> None:
|
||||
"""向当前会话追加一条或多条历史消息"""
|
||||
if not self._backend:
|
||||
return
|
||||
from zhenxun.services.ai.core.engine.context_renderer import ContextConverter
|
||||
|
||||
msgs = messages if isinstance(messages, Sequence) else [messages]
|
||||
flattened = ContextConverter.flatten_to_llm_messages(msgs)
|
||||
if flattened:
|
||||
await self._backend.add_messages(self.session_meta, flattened)
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""清空当前会话的短期对话历史"""
|
||||
if not self._backend:
|
||||
return
|
||||
await self._backend.clear(self.session_meta)
|
||||
|
||||
|
||||
class SlotFacade:
|
||||
"""中期记忆槽门面"""
|
||||
|
||||
def __init__(self, manager: "GlobalMemoryManager", session_meta: SessionMetadata):
|
||||
self.manager = manager
|
||||
self.session_meta = session_meta
|
||||
|
||||
@property
|
||||
def _backend(self) -> "BaseSlotContext | None":
|
||||
"""获取底层槽位存储后端"""
|
||||
return self.manager.get_slot_context(
|
||||
None, self.session_meta.namespace or "global"
|
||||
)
|
||||
|
||||
async def get(self, label: str) -> str | None:
|
||||
"""获取指定标识的槽位记忆内容"""
|
||||
if not self._backend:
|
||||
return None
|
||||
slot = await self._backend.get_slot(self.session_meta, label)
|
||||
return slot.content if slot else None
|
||||
|
||||
async def set(
|
||||
self,
|
||||
label: str,
|
||||
content: str,
|
||||
scope: Literal["session", "global"] = "session",
|
||||
size_limit: int = 2000,
|
||||
pinned: bool = True,
|
||||
) -> None:
|
||||
"""设置或更新指定的槽位记忆"""
|
||||
if not self._backend:
|
||||
return
|
||||
slot = MemorySlot(
|
||||
label=label,
|
||||
content=content,
|
||||
scope=scope,
|
||||
size_limit=size_limit,
|
||||
pinned=pinned,
|
||||
)
|
||||
await self._backend.set_slot(self.session_meta, slot)
|
||||
|
||||
async def delete(
|
||||
self, label: str, scope: Literal["session", "global"] = "session"
|
||||
) -> None:
|
||||
"""删除指定的槽位记忆"""
|
||||
if not self._backend:
|
||||
return
|
||||
await self._backend.delete_slot(self.session_meta, label, scope)
|
||||
|
||||
async def list_all(self) -> dict[str, str]:
|
||||
"""获取当前会话所有被置顶的槽位记忆"""
|
||||
if not self._backend:
|
||||
return {}
|
||||
slots = await self._backend.list_pinned_slots(self.session_meta)
|
||||
return {s.label: s.content for s in slots}
|
||||
|
||||
|
||||
class AgentSessionFacade:
|
||||
"""
|
||||
提供给第三方开发者的会话记忆访问聚合门面 (Facade)。
|
||||
"""
|
||||
|
||||
def __init__(self, manager: "GlobalMemoryManager", session_meta: SessionMetadata):
|
||||
self.manager = manager
|
||||
self.session_meta = session_meta
|
||||
self.history = ChatHistoryFacade(manager, session_meta)
|
||||
self.slots = SlotFacade(manager, session_meta)
|
||||
|
||||
async def clear_all(self) -> None:
|
||||
"""一键清空当前会话下的短期对话历史与记忆槽"""
|
||||
cleaner = self.manager.cleaner().session(self.session_meta.session_id)
|
||||
await cleaner.clear_short_term()
|
||||
await cleaner.clear_slots()
|
||||
@@ -0,0 +1,211 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
|
||||
from zhenxun.services.ai.context.memory.models import MemoryConfig
|
||||
from zhenxun.services.ai.context.memory.storage.backends import (
|
||||
InMemoryChatContext,
|
||||
MemoryScope,
|
||||
)
|
||||
from zhenxun.services.ai.context.memory.storage.interfaces import (
|
||||
BaseChatContext,
|
||||
BaseSlotContext,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.backends import Embedder, StorageBackend
|
||||
from zhenxun.services.ai.utils.scope import BaseScopeBuilder
|
||||
from zhenxun.utils.utils import infer_plugin_namespace
|
||||
|
||||
|
||||
class MemoryCleaner(BaseScopeBuilder["MemoryCleaner"]):
|
||||
"""
|
||||
声明式记忆清理构建器 (Query Builder)。
|
||||
为第三方开发者提供极端友好的链式 API，彻底屏蔽底层前缀逻辑。
|
||||
"""
|
||||
|
||||
def __init__(self, manager: "GlobalMemoryManager"):
|
||||
super().__init__()
|
||||
self.manager = manager
|
||||
self._config: Any = None
|
||||
|
||||
def config(self, cfg: Any):
|
||||
"""指定私有记忆配置（自动识别未全局注册 of 第三方私有数据库实例）"""
|
||||
self._config = cfg.build() if hasattr(cfg, "build") else cfg
|
||||
return self
|
||||
|
||||
async def clear_short_term(self):
|
||||
"""一键清理目标范围下的短期对话历史记忆"""
|
||||
if self._config and self._config.short_term.backend:
|
||||
await self._config.short_term.backend.clear_by_query(self._selector)
|
||||
else:
|
||||
for backend in self.manager._chat_backends.values():
|
||||
await backend.clear_by_query(self._selector)
|
||||
|
||||
async def clear_slots(self):
|
||||
"""一键清理目标范围下的中期记忆槽 (Memory Slots)"""
|
||||
if self._config and self._config.slots.backend:
|
||||
await self._config.slots.backend.clear_by_query(self._selector)
|
||||
else:
|
||||
for backend in self.manager._slot_backends.values():
|
||||
await backend.clear_by_query(self._selector)
|
||||
|
||||
async def clear_long_term(self):
|
||||
"""一键清理目标范围下的长期向量记忆 (RAG Vector Database)"""
|
||||
if self._config and self._config.long_term.backend:
|
||||
from zhenxun.services.ai.context.rag.backends import StorageBackend
|
||||
|
||||
storage = cast(StorageBackend, self._config.long_term.backend)
|
||||
await storage.clear_by_query(self._selector)
|
||||
else:
|
||||
for factory in self.manager._storage_factories.values():
|
||||
storage = factory()
|
||||
if hasattr(storage, "clear_by_query"):
|
||||
await storage.clear_by_query(self._selector)
|
||||
else:
|
||||
await storage.delete(scope_prefix=self._selector.scope_prefix)
|
||||
|
||||
async def clear_all(self):
|
||||
"""一键清理指定范围下的所有生命周期记忆（对话、槽位、RAG）"""
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
await self.clear_short_term()
|
||||
await self.clear_slots()
|
||||
await self.clear_long_term()
|
||||
logger.info(
|
||||
f"🧹 [MemoryCleaner] 成功清理作用域 '{self._selector.scope_prefix}'"
|
||||
"下的所有记忆痕迹！"
|
||||
)
|
||||
|
||||
|
||||
class GlobalMemoryManager:
|
||||
"""
|
||||
全局记忆大管家 (IoC 容器)。
|
||||
使用现代化依赖注入机制管理短/长期记忆引擎的默认实例。
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._chat_backends: dict[str, BaseChatContext] = {
|
||||
"global": InMemoryChatContext()
|
||||
}
|
||||
self._slot_backends: dict[str, BaseSlotContext] = {}
|
||||
|
||||
from zhenxun.services.ai.context.rag.backends import DictStorageBackend
|
||||
|
||||
self._storage_factories: dict[str, Callable[[], StorageBackend]] = {
|
||||
"global": lambda: DictStorageBackend()
|
||||
}
|
||||
|
||||
def register_chat_backend(
|
||||
self, backend: BaseChatContext, scope: str | None = None
|
||||
) -> None:
|
||||
"""注册特定命名空间的短期记忆存储后端。"""
|
||||
ns = scope if scope is not None else infer_plugin_namespace()
|
||||
self._chat_backends[ns] = backend
|
||||
|
||||
def register_slot_backend(
|
||||
self, backend: BaseSlotContext, scope: str | None = None
|
||||
) -> None:
|
||||
"""注册特定命名空间的中期记忆槽存储后端。"""
|
||||
ns = scope if scope is not None else infer_plugin_namespace()
|
||||
self._slot_backends[ns] = backend
|
||||
|
||||
def register_storage_factory(
|
||||
self, factory: Callable[[], StorageBackend], scope: str | None = None
|
||||
) -> None:
|
||||
"""注册特定命名空间的长期记忆向量存储工厂。"""
|
||||
ns = scope if scope is not None else infer_plugin_namespace()
|
||||
self._storage_factories[ns] = factory
|
||||
|
||||
def cleaner(self) -> MemoryCleaner:
|
||||
"""获取声明式记忆清理构建器，供第三方开发者极速清理指定记忆"""
|
||||
return MemoryCleaner(self)
|
||||
|
||||
def get_embedder(self, embedder_val: "Embedder | str | None") -> Embedder | None:
|
||||
"""获取向量化引擎实例。如果传入的是字符串，则视为 API 模型名称。"""
|
||||
if not embedder_val:
|
||||
return None
|
||||
|
||||
if isinstance(embedder_val, str):
|
||||
from zhenxun.services.ai.context.rag.backends.embedders import (
|
||||
DefaultEmbedder,
|
||||
)
|
||||
|
||||
return DefaultEmbedder(model_name=embedder_val)
|
||||
|
||||
return embedder_val
|
||||
|
||||
def get_chat_context(
|
||||
self, config: MemoryConfig | None, namespace: str = "global"
|
||||
) -> BaseChatContext | None:
|
||||
"""根据配置分配对应的短期对话历史实例"""
|
||||
if not config or not config.short_term.enable:
|
||||
return None
|
||||
|
||||
backend_cfg = config.short_term.backend
|
||||
if backend_cfg is not None:
|
||||
return cast(BaseChatContext, backend_cfg)
|
||||
|
||||
return self._chat_backends.get(namespace) or self._chat_backends["global"]
|
||||
|
||||
def get_slot_context(
|
||||
self, config: MemoryConfig | None, namespace: str = "global"
|
||||
) -> BaseSlotContext | None:
|
||||
"""根据配置分配对应的槽位记忆实例"""
|
||||
if not config or not config.slots.enable:
|
||||
return None
|
||||
|
||||
backend_cfg = config.slots.backend
|
||||
if backend_cfg is not None:
|
||||
return cast(BaseSlotContext, backend_cfg)
|
||||
|
||||
return self._slot_backends.get(namespace) or self._slot_backends["global"]
|
||||
|
||||
def get_long_term_memory(
|
||||
self, config: MemoryConfig | None, namespace: str = "global"
|
||||
) -> MemoryScope | None:
|
||||
"""根据声明式配置动态组装长期向量记忆实例"""
|
||||
if not config or not config.long_term.enable:
|
||||
return None
|
||||
|
||||
if config.long_term.engine is not None:
|
||||
return MemoryScope(
|
||||
rag_client=config.long_term.engine,
|
||||
)
|
||||
|
||||
storage_instance = None
|
||||
backend_cfg = config.long_term.backend
|
||||
if backend_cfg is not None:
|
||||
storage_instance = cast(StorageBackend, backend_cfg)
|
||||
else:
|
||||
factory = (
|
||||
self._storage_factories.get(namespace)
|
||||
or self._storage_factories["global"]
|
||||
)
|
||||
storage_instance = factory()
|
||||
|
||||
embedder = self.get_embedder(config.long_term.embedder)
|
||||
|
||||
from zhenxun.services.ai.context.rag.builder import RAGBuilder
|
||||
|
||||
builder = RAGBuilder(storage_instance).with_scope("/")
|
||||
if embedder:
|
||||
builder.with_embedder(embedder)
|
||||
|
||||
from zhenxun.services.ai.context.memory.models import MemoryScoringConfig
|
||||
|
||||
scoring_cfg = MemoryScoringConfig()
|
||||
|
||||
builder.enable_lifecycle_scoring(
|
||||
half_life_days=scoring_cfg.recency_half_life_days,
|
||||
decay_weight=scoring_cfg.recency_weight,
|
||||
semantic_weight=scoring_cfg.semantic_weight,
|
||||
importance_weight=scoring_cfg.importance_weight,
|
||||
reinforcement_weight=scoring_cfg.reinforcement_weight,
|
||||
)
|
||||
|
||||
client = builder.build()
|
||||
|
||||
return MemoryScope(
|
||||
rag_client=client,
|
||||
)
|
||||
|
||||
|
||||
memory_manager = GlobalMemoryManager()
|
||||
@@ -0,0 +1,173 @@
|
||||
"""
|
||||
记忆域类型定义
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from zhenxun.services.ai.context.memory.storage.interfaces import (
|
||||
BaseChatContext,
|
||||
BaseMemoryIngestionMiddleware,
|
||||
BaseMemoryReducer,
|
||||
BaseSlotContext,
|
||||
)
|
||||
from zhenxun.services.ai.context.memory.types import (
|
||||
AutoRecallPolicy,
|
||||
Isolation,
|
||||
MemorySlot,
|
||||
SessionMetadata,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.backends import Embedder, StorageBackend
|
||||
from zhenxun.services.ai.context.rag.engine import ScopedRAGClient
|
||||
from zhenxun.services.ai.utils.scope import ScopeBuilder
|
||||
|
||||
|
||||
class SlotMemoryConfig(BaseModel):
|
||||
"""槽位记忆 (Memory Slots) 配置"""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
enable: bool = Field(default=False)
|
||||
"""是否启用中期记忆槽"""
|
||||
scopes: dict[str, ScopeBuilder] | None = Field(default=None)
|
||||
"""语义化作用域映射字典，供大模型作为 Literal 选择。如果只有一个，则自动隐藏参数"""
|
||||
default_slots: list[MemorySlot] = Field(default_factory=list)
|
||||
"""首次初始化时自动写入的默认槽位列表"""
|
||||
backend: str | BaseSlotContext | None = Field(default=None)
|
||||
"""
|
||||
指定底层槽位记忆数据库注册名称，或直接传入 BaseSlotContext 实例。
|
||||
为空则使用全局默认
|
||||
"""
|
||||
instructions: str | None = Field(default=None)
|
||||
"""覆写内置槽位管理工具箱的系统提示词"""
|
||||
toolkit_kwargs: dict[str, Any] = Field(default_factory=dict)
|
||||
"""透传给底层 MemorySlotToolkit 的高级参数 (如 prefix, exclude, shared_options)"""
|
||||
|
||||
|
||||
class MemoryScoringConfig(BaseModel):
|
||||
"""长期记忆的复合打分与检索配置"""
|
||||
|
||||
recency_weight: float = Field(default=0.3)
|
||||
"""时间衰减权重"""
|
||||
semantic_weight: float = Field(default=0.5)
|
||||
"""语义相似度权重"""
|
||||
importance_weight: float = Field(default=0.2)
|
||||
"""重要性权重"""
|
||||
recency_half_life_days: int = Field(default=30)
|
||||
"""时间衰减的半衰期(天)"""
|
||||
|
||||
reinforcement_weight: float = Field(default=0.2)
|
||||
"""访问强化的加权权重 (被检索越多得分越高)"""
|
||||
|
||||
|
||||
class ShortTermConfig(BaseModel):
|
||||
"""短期对话记忆配置"""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
enable: bool = Field(default=True)
|
||||
"""是否启用短期对话记忆上下文"""
|
||||
backend: str | BaseChatContext | None = Field(default=None)
|
||||
"""
|
||||
指定底层短期记忆数据库注册名称，或直接传入 BaseChatContext 实例。
|
||||
为空则使用全局默认
|
||||
"""
|
||||
isolation: ScopeBuilder = Field(default_factory=Isolation.AGENT_USER)
|
||||
"""单一的记忆隔离级别 (ScopeBuilder)，决定短期记忆存储边界"""
|
||||
|
||||
|
||||
class LongTermConfig(BaseModel):
|
||||
"""长期向量记忆配置"""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
enable: bool = Field(default=False)
|
||||
"""是否启用长期记忆（开启后自动赋予 Agent 存取记忆的工具，并附加 RAG 召回能力）"""
|
||||
engine: ScopedRAGClient | None = Field(default=None)
|
||||
"""
|
||||
[推荐] 指定底层的高级 RAG 检索引擎实例。若传入此项，将覆盖默认的 backend
|
||||
和 embedder 配置。
|
||||
"""
|
||||
backend: str | StorageBackend | None = Field(default=None)
|
||||
"""
|
||||
指定底层长期向量数据库 (Storage) 注册名称，或直接传入 StorageBackend 实例。
|
||||
为空则使用全局默认
|
||||
"""
|
||||
scopes: dict[str, ScopeBuilder] | None = Field(default=None)
|
||||
"""语义化作用域映射字典，决定长期记忆存储边界。如果只有一个，则自动隐藏参数"""
|
||||
embedder: str | Embedder | None = Field(default=None)
|
||||
"""
|
||||
指定底层向量化引擎 (Embedder) 实例，若为字符串则视为 API 模型名称。
|
||||
为空则使用全局默认
|
||||
"""
|
||||
|
||||
agentic: bool = Field(default=True)
|
||||
"""是否赋予大模型主动管理记忆的能力 (Agentic Memory)"""
|
||||
auto_recall: AutoRecallPolicy = Field(default=False)
|
||||
"""长期记忆的自动召回策略，默认 False (从不自动召回)，由大模型自主
|
||||
决定调用搜索工具"""
|
||||
recall_threshold: float = Field(default=0.5)
|
||||
"""长期记忆召回的最低余弦相似度要求"""
|
||||
instructions: str | None = Field(default=None)
|
||||
"""覆写内置长期记忆管理工具箱的系统提示词"""
|
||||
toolkit_kwargs: dict[str, Any] = Field(default_factory=dict)
|
||||
"""透传给底层 MemoryManagementToolkit 的高级参数"""
|
||||
|
||||
|
||||
class ContextCompressionConfig(BaseModel):
|
||||
"""上下文压缩与管理配置"""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
threshold: float | None = Field(default=None)
|
||||
"""(局部重写) 触发记忆压缩的 Token 阈值"""
|
||||
max_history_turns: int | None = Field(default=None)
|
||||
"""(局部重写) 触发记忆压缩的对话轮数上限。设为 0 表示不限制轮数。"""
|
||||
vision_window: int | None = Field(default=None)
|
||||
"""多模态滑动窗口大小。0表示关闭该功能，>0表示仅保留最近N轮包含多模态数据的消息，None表示跟随全局配置。"""
|
||||
policy: list[BaseMemoryReducer] | None = Field(default=None)
|
||||
"""核心记忆压缩策略管线 (List[BaseMemoryReducer])。为 None 时将应用全局默认策略。"""
|
||||
|
||||
|
||||
class IngestionConfig(BaseModel):
|
||||
"""记忆入库管线配置"""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
middlewares: list[BaseMemoryIngestionMiddleware] = Field(default_factory=list)
|
||||
"""入库中间件列表（按顺序依次执行清洗过滤）"""
|
||||
|
||||
|
||||
class MemoryConfig(BaseModel):
|
||||
"""统一的记忆配置项声明"""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
base_isolation: ScopeBuilder = Field(default_factory=Isolation.AGENT_USER)
|
||||
"""顶层基准隔离级别，短期/中期/长期记忆将默认继承此级别"""
|
||||
short_term: ShortTermConfig = Field(default_factory=ShortTermConfig)
|
||||
"""短期对话记忆配置"""
|
||||
slots: SlotMemoryConfig = Field(default_factory=SlotMemoryConfig)
|
||||
"""槽位记忆配置"""
|
||||
long_term: LongTermConfig = Field(default_factory=LongTermConfig)
|
||||
"""长期向量记忆配置"""
|
||||
compression: ContextCompressionConfig = Field(
|
||||
default_factory=ContextCompressionConfig
|
||||
)
|
||||
"""上下文压缩与管理配置"""
|
||||
ingestion: IngestionConfig = Field(default_factory=IngestionConfig)
|
||||
"""记忆入库前的清洗与过滤管线配置"""
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AutoRecallPolicy",
|
||||
"BaseMemoryIngestionMiddleware",
|
||||
"ContextCompressionConfig",
|
||||
"IngestionConfig",
|
||||
"Isolation",
|
||||
"LongTermConfig",
|
||||
"MemoryConfig",
|
||||
"MemoryScoringConfig",
|
||||
"SessionMetadata",
|
||||
"ShortTermConfig",
|
||||
]
|
||||
@@ -0,0 +1,21 @@
|
||||
from .backends import (
|
||||
AbstractMemoryRecord,
|
||||
AbstractSlotRecord,
|
||||
InMemoryChatContext,
|
||||
MemoryScope,
|
||||
TortoiseChatContext,
|
||||
TortoiseSlotContext,
|
||||
get_orm_chat_context,
|
||||
get_orm_slot_context,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AbstractMemoryRecord",
|
||||
"AbstractSlotRecord",
|
||||
"InMemoryChatContext",
|
||||
"MemoryScope",
|
||||
"TortoiseChatContext",
|
||||
"TortoiseSlotContext",
|
||||
"get_orm_chat_context",
|
||||
"get_orm_slot_context",
|
||||
]
|
||||
@@ -0,0 +1,465 @@
|
||||
import asyncio
|
||||
import base64
|
||||
from collections.abc import Callable
|
||||
import datetime
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from zhenxun.services.ai.context.rag.engine import ScopedRAGClient
|
||||
|
||||
from nonebot.utils import is_coroutine_callable
|
||||
from tortoise import fields
|
||||
from tortoise.timezone import now
|
||||
|
||||
from zhenxun.services.ai.context.memory.storage.interfaces import (
|
||||
BaseChatContext,
|
||||
BaseSlotContext,
|
||||
)
|
||||
from zhenxun.services.ai.context.memory.types import (
|
||||
MemorySlot,
|
||||
SessionMetadata,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.models import BaseRecord, SearchResult
|
||||
from zhenxun.services.ai.core.messages import (
|
||||
AssistantMessage,
|
||||
LLMContentPart,
|
||||
LLMMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
UserMessage,
|
||||
)
|
||||
from zhenxun.services.ai.utils.scope import ScopeSelector
|
||||
from zhenxun.services.db_context import Model
|
||||
from zhenxun.utils.pydantic_compat import TypeAdapter, model_dump
|
||||
|
||||
|
||||
class DBMessageSerializer:
|
||||
"""将 LLMMessage 与数据库 JSON 格式进行序列化/反序列化的帮助类"""
|
||||
|
||||
@staticmethod
|
||||
def deserialize_content(content_raw: Any) -> list[LLMContentPart]:
|
||||
from zhenxun.services.ai.core.messages import TextPart
|
||||
|
||||
content_parts: list[LLMContentPart] = []
|
||||
if isinstance(content_raw, list):
|
||||
adapter = TypeAdapter(LLMContentPart)
|
||||
for p in content_raw:
|
||||
if isinstance(p, dict):
|
||||
for k in list(p.keys()):
|
||||
if k.startswith("_is_b64_"):
|
||||
orig_k = k[8:]
|
||||
if orig_k in p and isinstance(p[orig_k], str):
|
||||
p[orig_k] = base64.b64decode(p[orig_k])
|
||||
p.pop(k, None)
|
||||
content_parts.append(adapter.validate_python(p))
|
||||
elif isinstance(content_raw, str):
|
||||
content_parts.append(TextPart(text=content_raw))
|
||||
return content_parts
|
||||
|
||||
@staticmethod
|
||||
def serialize_content(content_payload: Any) -> list[dict[str, Any]]:
|
||||
from pathlib import Path
|
||||
|
||||
if isinstance(content_payload, str):
|
||||
return [{"type": "text", "text": content_payload}]
|
||||
elif isinstance(content_payload, list):
|
||||
processed_content = []
|
||||
for p in content_payload:
|
||||
p_dump = (
|
||||
model_dump(p, exclude_none=True)
|
||||
if hasattr(p, "model_dump")
|
||||
else (p.copy() if isinstance(p, dict) else p)
|
||||
)
|
||||
if isinstance(p_dump, dict):
|
||||
for k, v in list(p_dump.items()):
|
||||
if isinstance(v, bytes):
|
||||
p_dump[k] = base64.b64encode(v).decode("utf-8")
|
||||
p_dump[f"_is_b64_{k}"] = True
|
||||
elif isinstance(v, Path):
|
||||
p_dump[k] = str(v)
|
||||
processed_content.append(p_dump)
|
||||
return (
|
||||
processed_content
|
||||
if processed_content
|
||||
else [
|
||||
{"type": "text", "text": "[仅包含思维链或工具调度，无实质文本输出]"}
|
||||
]
|
||||
)
|
||||
return []
|
||||
|
||||
|
||||
class MemoryScope:
|
||||
"""长期记忆的作用域视图与 RAG 管线。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rag_client: "ScopedRAGClient",
|
||||
):
|
||||
self.rag_client = rag_client
|
||||
self._background_tasks: set[Any] = set()
|
||||
|
||||
async def remember(
|
||||
self,
|
||||
session: SessionMetadata,
|
||||
content: str,
|
||||
importance: float = 0.5,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""通过 RAG Ingestion Pipeline 完成记忆落盘"""
|
||||
meta = metadata.copy() if metadata else {}
|
||||
meta.update(
|
||||
{
|
||||
"scope": session.scope_prefix,
|
||||
"importance": importance,
|
||||
"created_at": time.time(),
|
||||
}
|
||||
)
|
||||
record = BaseRecord(content=content, metadata=meta)
|
||||
|
||||
await self.rag_client.ingest([record])
|
||||
|
||||
async def recall(
|
||||
self,
|
||||
session: SessionMetadata,
|
||||
query: str,
|
||||
limit: int = 10,
|
||||
metadata_filter: dict[str, Any] | None = None,
|
||||
) -> list[SearchResult]:
|
||||
"""委托至 Retriever 检索与重排，并触发读时惰性强化"""
|
||||
matches = await self.rag_client.search(
|
||||
query=query,
|
||||
limit=limit,
|
||||
scopes=session.accessible_scopes,
|
||||
metadata_filters=metadata_filter,
|
||||
)
|
||||
if matches:
|
||||
task = asyncio.create_task(
|
||||
self._reinforce_memories([m.record for m in matches])
|
||||
)
|
||||
self._background_tasks.add(task)
|
||||
task.add_done_callback(self._background_tasks.discard)
|
||||
return matches
|
||||
|
||||
async def update(
|
||||
self,
|
||||
session: SessionMetadata,
|
||||
record_id: str,
|
||||
new_content: str,
|
||||
importance: float = 0.5,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""原子更新：通过先删后插，确保底层向量(Embedding)能根据新文本被正确刷新"""
|
||||
deleted_count = await self.forget(session, record_ids=[record_id])
|
||||
if deleted_count > 0:
|
||||
await self.remember(
|
||||
session=session,
|
||||
content=new_content,
|
||||
importance=importance,
|
||||
metadata=metadata,
|
||||
)
|
||||
return True
|
||||
return False
|
||||
|
||||
async def forget(
|
||||
self, session: SessionMetadata, record_ids: list[str] | None = None
|
||||
) -> int:
|
||||
return await self.rag_client.delete(
|
||||
record_ids=record_ids,
|
||||
)
|
||||
|
||||
async def _reinforce_memories(self, records: list[BaseRecord]):
|
||||
import time
|
||||
|
||||
now = time.time()
|
||||
for r in records:
|
||||
r.metadata["access_count"] = r.metadata.get("access_count", 0) + 1
|
||||
r.metadata["last_accessed_at"] = now
|
||||
await self.rag_client.storage.update(r)
|
||||
|
||||
|
||||
class InMemoryChatContext(BaseChatContext):
|
||||
def __init__(self):
|
||||
self._messages: dict[str, list[LLMMessage]] = {}
|
||||
|
||||
async def get_messages(self, session: SessionMetadata) -> list[LLMMessage]:
|
||||
return list(self._messages.get(session.session_id, []))
|
||||
|
||||
async def search(
|
||||
self, query: str, session: SessionMetadata, limit: int = 10
|
||||
) -> list[LLMMessage]:
|
||||
results = []
|
||||
for msg in self._messages.get(session.session_id, []):
|
||||
if query in msg.extract_text:
|
||||
results.append(msg)
|
||||
if len(results) >= limit:
|
||||
break
|
||||
return results
|
||||
|
||||
async def add_messages(
|
||||
self, session: SessionMetadata, messages: list[LLMMessage]
|
||||
) -> None:
|
||||
if session.session_id not in self._messages:
|
||||
self._messages[session.session_id] = []
|
||||
self._messages[session.session_id].extend(messages)
|
||||
|
||||
async def set_messages(
|
||||
self, session: SessionMetadata, messages: list[LLMMessage]
|
||||
) -> None:
|
||||
self._messages[session.session_id] = list(messages)
|
||||
|
||||
async def clear(self, session: SessionMetadata) -> None:
|
||||
self._messages.pop(session.session_id, None)
|
||||
|
||||
async def clear_by_query(self, query: ScopeSelector) -> None:
|
||||
"""内存级：前缀匹配清理所有符合要求的短期会话"""
|
||||
scope_prefix = query.scope_prefix
|
||||
keys_to_delete = [
|
||||
sid for sid in self._messages.keys() if sid.startswith(scope_prefix)
|
||||
]
|
||||
for sid in keys_to_delete:
|
||||
self._messages.pop(sid, None)
|
||||
|
||||
|
||||
class AbstractMemoryRecord(Model):
|
||||
"""Tortoise ORM 短期记忆持久化基类 (Mixin)。"""
|
||||
|
||||
id = fields.UUIDField(pk=True, description="主键")
|
||||
session_id = fields.CharField(max_length=255, index=True)
|
||||
role = fields.CharField(max_length=32)
|
||||
content = fields.JSONField()
|
||||
api_context = fields.JSONField(null=True)
|
||||
created_at = fields.DatetimeField(auto_now_add=True)
|
||||
metadata = fields.JSONField(null=True)
|
||||
|
||||
class Meta: # type: ignore
|
||||
abstract = True
|
||||
|
||||
|
||||
class TortoiseChatContext(BaseChatContext):
|
||||
def __init__(
|
||||
self,
|
||||
model_class: type[AbstractMemoryRecord],
|
||||
custom_save_hook: Callable[
|
||||
[AbstractMemoryRecord, LLMMessage, SessionMetadata], Any
|
||||
]
|
||||
| None = None,
|
||||
):
|
||||
self.model_class = model_class
|
||||
self.custom_save_hook = custom_save_hook
|
||||
|
||||
def _row_to_message(self, row: AbstractMemoryRecord) -> LLMMessage:
|
||||
content_parts = DBMessageSerializer.deserialize_content(row.content)
|
||||
metadata: dict[str, Any] | None = (
|
||||
row.metadata if isinstance(row.metadata, dict) else None
|
||||
)
|
||||
kwargs = {
|
||||
"content": content_parts,
|
||||
"metadata": metadata,
|
||||
"created_at": row.created_at.timestamp() if row.created_at else time.time(),
|
||||
}
|
||||
role = row.role
|
||||
if role == "system":
|
||||
return cast(LLMMessage, SystemMessage(**kwargs))
|
||||
elif role == "user":
|
||||
return cast(LLMMessage, UserMessage(**kwargs))
|
||||
elif role == "assistant":
|
||||
return cast(LLMMessage, AssistantMessage(**kwargs))
|
||||
elif role == "tool":
|
||||
return cast(LLMMessage, ToolMessage(**kwargs))
|
||||
return cast(LLMMessage, LLMMessage(role=role, **kwargs))
|
||||
|
||||
async def get_messages(self, session: SessionMetadata) -> list[LLMMessage]:
|
||||
rows = (
|
||||
await self.model_class.filter(session_id=session.session_id)
|
||||
.order_by("created_at")
|
||||
.all()
|
||||
)
|
||||
return [self._row_to_message(row) for row in rows]
|
||||
|
||||
async def search(
|
||||
self, query: str, session: SessionMetadata, limit: int = 10
|
||||
) -> list[LLMMessage]:
|
||||
rows = (
|
||||
await self.model_class.filter(
|
||||
session_id=session.session_id, content__icontains=query
|
||||
)
|
||||
.order_by("-created_at")
|
||||
.limit(limit)
|
||||
.all()
|
||||
)
|
||||
return [self._row_to_message(row) for row in reversed(rows)]
|
||||
|
||||
async def add_messages(
|
||||
self, session: SessionMetadata, messages: list[LLMMessage]
|
||||
) -> None:
|
||||
if not messages:
|
||||
return
|
||||
|
||||
base_time = now()
|
||||
|
||||
last_msg = (
|
||||
await self.model_class.filter(session_id=session.session_id)
|
||||
.order_by("-created_at")
|
||||
.first()
|
||||
)
|
||||
if last_msg and last_msg.created_at and last_msg.created_at >= base_time:
|
||||
base_time = last_msg.created_at + datetime.timedelta(milliseconds=10)
|
||||
|
||||
orm_objects = []
|
||||
for i, msg in enumerate(messages):
|
||||
content_payload = DBMessageSerializer.serialize_content(msg.content)
|
||||
|
||||
msg_time = base_time + datetime.timedelta(milliseconds=i * 10)
|
||||
orm_obj = self.model_class(
|
||||
session_id=session.session_id,
|
||||
role=msg.role,
|
||||
content=content_payload,
|
||||
api_context=None,
|
||||
metadata=msg.metadata,
|
||||
created_at=msg_time,
|
||||
)
|
||||
if self.custom_save_hook:
|
||||
if is_coroutine_callable(self.custom_save_hook):
|
||||
await self.custom_save_hook(orm_obj, msg, session)
|
||||
else:
|
||||
self.custom_save_hook(orm_obj, msg, session)
|
||||
orm_objects.append(orm_obj)
|
||||
if orm_objects:
|
||||
await self.model_class.bulk_create(orm_objects)
|
||||
|
||||
async def set_messages(
|
||||
self, session: SessionMetadata, messages: list[LLMMessage]
|
||||
) -> None:
|
||||
await self.clear(session)
|
||||
await self.add_messages(session, messages)
|
||||
|
||||
async def clear(self, session: SessionMetadata) -> None:
|
||||
await self.model_class.filter(session_id=session.session_id).delete()
|
||||
|
||||
async def clear_by_query(self, query: ScopeSelector) -> None:
|
||||
"""ORM 级：利用数据库 startswith 原生语法批量级联删除短期记忆"""
|
||||
scope_prefix = query.scope_prefix
|
||||
await self.model_class.filter(session_id__startswith=scope_prefix).delete()
|
||||
|
||||
|
||||
def get_orm_chat_context(
|
||||
model_class: type[AbstractMemoryRecord],
|
||||
custom_save_hook: Callable[[AbstractMemoryRecord, LLMMessage, SessionMetadata], Any]
|
||||
| None = None,
|
||||
) -> TortoiseChatContext:
|
||||
"""
|
||||
[工厂方法] 供第三方开发者调用，
|
||||
将 Tortoise ORM 表直接包装为对话历史记录系统。
|
||||
"""
|
||||
return TortoiseChatContext(
|
||||
model_class=model_class, custom_save_hook=custom_save_hook
|
||||
)
|
||||
|
||||
|
||||
class AbstractSlotRecord(Model):
|
||||
"""Tortoise ORM 记忆槽持久化基类 (Mixin)。"""
|
||||
|
||||
id = fields.CharField(
|
||||
pk=True, max_length=128, description="复合主键: session_id + label"
|
||||
)
|
||||
session_id = fields.CharField(max_length=255, index=True)
|
||||
label = fields.CharField(max_length=64, index=True)
|
||||
content = fields.TextField()
|
||||
size_limit = fields.IntField(default=2000)
|
||||
pinned = fields.BooleanField(default=True)
|
||||
scope = fields.CharField(max_length=255)
|
||||
description = fields.CharField(max_length=255, default="")
|
||||
created_at = fields.FloatField()
|
||||
updated_at = fields.FloatField()
|
||||
|
||||
class Meta: # type: ignore
|
||||
abstract = True
|
||||
|
||||
|
||||
class TortoiseSlotContext(BaseSlotContext):
|
||||
def __init__(self, model_class: type[AbstractSlotRecord]):
|
||||
self.model_class = model_class
|
||||
|
||||
def _row_to_slot(self, row: AbstractSlotRecord) -> MemorySlot:
|
||||
return MemorySlot(
|
||||
label=row.label,
|
||||
content=row.content,
|
||||
size_limit=row.size_limit,
|
||||
pinned=row.pinned,
|
||||
scope=row.scope,
|
||||
description=row.description,
|
||||
created_at=row.created_at,
|
||||
updated_at=row.updated_at,
|
||||
)
|
||||
|
||||
async def get_slot(self, session: SessionMetadata, label: str) -> MemorySlot | None:
|
||||
rows = await self.model_class.filter(
|
||||
session_id__in=session.accessible_scopes, label=label
|
||||
).all()
|
||||
|
||||
row_map = {r.session_id: r for r in rows}
|
||||
for scope in reversed(session.accessible_scopes):
|
||||
if scope in row_map:
|
||||
return self._row_to_slot(row_map[scope])
|
||||
return None
|
||||
|
||||
async def set_slot(self, session: SessionMetadata, slot: MemorySlot) -> None:
|
||||
composite_id = f"{slot.scope}_{slot.label}"
|
||||
|
||||
await self.model_class.update_or_create(
|
||||
id=composite_id,
|
||||
defaults={
|
||||
"session_id": slot.scope,
|
||||
"label": slot.label,
|
||||
"content": slot.content,
|
||||
"size_limit": slot.size_limit,
|
||||
"pinned": slot.pinned,
|
||||
"scope": slot.scope,
|
||||
"description": slot.description,
|
||||
"created_at": slot.created_at,
|
||||
"updated_at": slot.updated_at,
|
||||
},
|
||||
)
|
||||
|
||||
async def delete_slot(
|
||||
self, session: SessionMetadata, label: str, scope: str
|
||||
) -> None:
|
||||
composite_id = f"{scope}_{label}"
|
||||
await self.model_class.filter(id=composite_id).delete()
|
||||
|
||||
async def list_pinned_slots(self, session: SessionMetadata) -> list[MemorySlot]:
|
||||
rows = await self.model_class.filter(
|
||||
session_id__in=session.accessible_scopes, pinned=True
|
||||
).all()
|
||||
|
||||
merged = {}
|
||||
for scope in session.accessible_scopes:
|
||||
for row in rows:
|
||||
if row.session_id == scope:
|
||||
merged[row.label] = self._row_to_slot(row)
|
||||
|
||||
return [s for s in merged.values() if s.content.strip()]
|
||||
|
||||
async def list_all_slots(self, session: SessionMetadata) -> list[MemorySlot]:
|
||||
rows = await self.model_class.filter(
|
||||
session_id__in=session.accessible_scopes
|
||||
).all()
|
||||
|
||||
merged = {}
|
||||
for scope in session.accessible_scopes:
|
||||
for row in rows:
|
||||
if row.session_id == scope:
|
||||
merged[row.label] = self._row_to_slot(row)
|
||||
return list(merged.values())
|
||||
|
||||
async def clear_by_query(self, query: ScopeSelector) -> None:
|
||||
scope_prefix = query.scope_prefix
|
||||
await self.model_class.filter(session_id__startswith=scope_prefix).delete()
|
||||
|
||||
|
||||
def get_orm_slot_context(model_class: type[AbstractSlotRecord]) -> TortoiseSlotContext:
|
||||
"""
|
||||
[工厂方法] 供第三方开发者调用，将 Tortoise ORM 表直接包装为记忆槽存储系统。
|
||||
"""
|
||||
return TortoiseSlotContext(model_class=model_class)
|
||||
@@ -0,0 +1,110 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Sequence
|
||||
|
||||
from zhenxun.services.ai.context.memory.types import (
|
||||
MemorySlot,
|
||||
SessionMetadata,
|
||||
)
|
||||
from zhenxun.services.ai.core.messages import AgentMessage, LLMMessage
|
||||
from zhenxun.services.ai.run.context import RunContext
|
||||
from zhenxun.services.ai.utils.scope import ScopeSelector
|
||||
|
||||
|
||||
class BaseChatContext(ABC):
|
||||
"""短期对话历史记忆接口"""
|
||||
|
||||
@abstractmethod
|
||||
async def get_messages(self, session: SessionMetadata) -> list[LLMMessage]:
|
||||
"""获取当前会话的所有历史消息。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
self, query: str, session: SessionMetadata, limit: int = 10
|
||||
) -> list[LLMMessage]:
|
||||
"""根据查询词搜索当前会话的历史消息。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def add_messages(
|
||||
self, session: SessionMetadata, messages: list[LLMMessage]
|
||||
) -> None:
|
||||
"""向当前会话追加消息。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def set_messages(
|
||||
self, session: SessionMetadata, messages: list[LLMMessage]
|
||||
) -> None:
|
||||
"""重置并设置当前会话的消息列表。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def clear(self, session: SessionMetadata) -> None:
|
||||
"""清空当前会话的历史消息。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def clear_by_query(self, query: ScopeSelector) -> None:
|
||||
"""根据条件领域查询对象清理对话历史。"""
|
||||
...
|
||||
|
||||
|
||||
class BaseSlotContext(ABC):
|
||||
"""中期记忆槽持久化接口"""
|
||||
|
||||
@abstractmethod
|
||||
async def get_slot(self, session: SessionMetadata, label: str) -> MemorySlot | None:
|
||||
"""获取指定会话下特定标签的记忆槽。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def set_slot(self, session: SessionMetadata, slot: MemorySlot) -> None:
|
||||
"""设置或更新指定会话下的记忆槽。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_slot(
|
||||
self, session: SessionMetadata, label: str, scope: str
|
||||
) -> None:
|
||||
"""删除指定会话下特定标签和作用域的记忆槽。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def list_pinned_slots(self, session: SessionMetadata) -> list[MemorySlot]:
|
||||
"""列出当前会话所有固定的记忆槽。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def list_all_slots(self, session: SessionMetadata) -> list[MemorySlot]:
|
||||
"""列出当前会话的所有记忆槽（包括未置顶的）。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def clear_by_query(self, query: ScopeSelector) -> None:
|
||||
"""根据条件领域查询对象清理记忆槽。"""
|
||||
...
|
||||
|
||||
|
||||
class BaseMemoryReducer(ABC):
|
||||
"""记忆压缩器基类"""
|
||||
|
||||
@abstractmethod
|
||||
async def reduce(
|
||||
self,
|
||||
messages: list[LLMMessage],
|
||||
current_tokens: int,
|
||||
model_name: str,
|
||||
base_overhead: int = 0,
|
||||
) -> tuple[list[LLMMessage], bool, int]:
|
||||
"""对消息列表进行压缩处理。"""
|
||||
...
|
||||
|
||||
|
||||
class BaseMemoryIngestionMiddleware(ABC):
|
||||
"""记忆入库中间件基类，在写入数据库前拦截并修改/清洗消息"""
|
||||
|
||||
@abstractmethod
|
||||
async def process(
|
||||
self, messages: Sequence[AgentMessage], context: RunContext
|
||||
) -> list[AgentMessage]: ...
|
||||
@@ -0,0 +1,103 @@
|
||||
from collections.abc import Awaitable, Callable
|
||||
import time
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from zhenxun.services.ai.utils.scope import ScopeBuilder, ScopeSelector
|
||||
|
||||
AutoRecallPolicy = bool | Callable[[str, "SessionMetadata"], Awaitable[bool] | bool]
|
||||
"""长期记忆的自动召回策略"""
|
||||
|
||||
|
||||
class MemorySlot(BaseModel):
|
||||
"""可编辑的持久化记忆槽 (Mid-Term Memory)"""
|
||||
|
||||
label: str = Field(...)
|
||||
"""槽位唯一标签标识 (如 persona, preferences)"""
|
||||
content: str = Field(default="")
|
||||
"""槽位存储的具体文本内容"""
|
||||
size_limit: int = Field(default=2000)
|
||||
"""槽位内容的最大字符数限制"""
|
||||
pinned: bool = Field(default=True)
|
||||
"""是否固定注入到大模型的每次系统提示词中"""
|
||||
scope: str = Field(...)
|
||||
"""作用域：表示隔离的路径前缀 (scope_prefix)"""
|
||||
description: str = Field(default="")
|
||||
"""该记忆槽的用途说明，便于大模型理解"""
|
||||
created_at: float = Field(default_factory=time.time)
|
||||
"""创建时间戳"""
|
||||
updated_at: float = Field(default_factory=time.time)
|
||||
"""最近更新时间戳"""
|
||||
|
||||
|
||||
class Isolation:
|
||||
"""预设策略工厂，提供友好的隔离级别声明式 API"""
|
||||
|
||||
@staticmethod
|
||||
def _base() -> ScopeBuilder:
|
||||
"""获取底座通用隔离级别（包含Bot、平台、命名空间、智能体）。"""
|
||||
return ScopeBuilder().bot().platform().namespace().agent()
|
||||
|
||||
@classmethod
|
||||
def GROUP_SHARED(cls) -> ScopeBuilder:
|
||||
"""群组共享隔离：同群内共享金库与记忆。"""
|
||||
return cls._base().group()
|
||||
|
||||
@classmethod
|
||||
def USER_GLOBAL(cls) -> ScopeBuilder:
|
||||
"""用户全局隔离：跨群、跨插件共享用户记忆。"""
|
||||
return cls._base().user()
|
||||
|
||||
@classmethod
|
||||
def GROUP_USER(cls) -> ScopeBuilder:
|
||||
"""群组用户隔离：单群内单用户独立隔离。"""
|
||||
return cls._base().group().user()
|
||||
|
||||
@classmethod
|
||||
def AGENT_USER(cls) -> ScopeBuilder:
|
||||
"""智能体用户隔离：单智能体单用户物理隔离。"""
|
||||
return cls.GROUP_USER()
|
||||
|
||||
|
||||
class SessionMetadata(BaseModel):
|
||||
"""结构化会话元数据"""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
session_id: str = Field(...)
|
||||
"""核心会话标识符。"""
|
||||
selector: ScopeSelector = Field(default_factory=ScopeSelector)
|
||||
"""统一的作用域与实体资源选择器。"""
|
||||
isolation_level: ScopeBuilder | None = Field(default=None)
|
||||
"""生成此会话时的隔离级别。"""
|
||||
scope_prefix: str = Field(default="/")
|
||||
"""基于隔离级别生成的路径作用域，用于长期记忆 (RAG) 的向量检索前缀过滤。"""
|
||||
accessible_scopes: list[str] = Field(default_factory=lambda: ["/"])
|
||||
"""
|
||||
当前会话有权访问的作用域列表，用于 Slice 联合检索。
|
||||
"""
|
||||
scope_name_mapping: dict[str, str] = Field(default_factory=dict)
|
||||
"""物理路径到语义化名称的逆向映射字典，供大模型友好阅读"""
|
||||
|
||||
@property
|
||||
def platform(self) -> str | None:
|
||||
return self.selector.platform
|
||||
|
||||
@property
|
||||
def group_id(self) -> str | None:
|
||||
return self.selector.group_id
|
||||
|
||||
@property
|
||||
def user_id(self) -> str | None:
|
||||
return self.selector.user_id
|
||||
|
||||
@property
|
||||
def namespace(self) -> str | None:
|
||||
return self.selector.namespace
|
||||
|
||||
@property
|
||||
def agent_name(self) -> str | None:
|
||||
return self.selector.agent_name
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.session_id
|
||||
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
Zhenxun AI - RAG (检索增强生成) 基础设施层
|
||||
"""
|
||||
|
||||
from .backends.embedders import Embedder
|
||||
from .backends.storages import (
|
||||
AbstractVectorRecord,
|
||||
StorageBackend,
|
||||
TortoiseStorageBackend,
|
||||
)
|
||||
from .builder import RAGBuilder
|
||||
from .configs import RAGConfig
|
||||
from .engine import ScopedRAGClient
|
||||
from .models import BaseRecord, SearchResult
|
||||
|
||||
__all__ = [
|
||||
"AbstractVectorRecord",
|
||||
"BaseRecord",
|
||||
"Embedder",
|
||||
"RAGBuilder",
|
||||
"RAGConfig",
|
||||
"ScopedRAGClient",
|
||||
"SearchResult",
|
||||
"StorageBackend",
|
||||
"TortoiseStorageBackend",
|
||||
]
|
||||
@@ -0,0 +1,20 @@
|
||||
from .embedders import DefaultEmbedder, Embedder
|
||||
from .storages import (
|
||||
AbstractVectorRecord,
|
||||
DictStorageBackend,
|
||||
LanceDBStorageBackend,
|
||||
QdrantStorageBackend,
|
||||
StorageBackend,
|
||||
TortoiseStorageBackend,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AbstractVectorRecord",
|
||||
"DefaultEmbedder",
|
||||
"DictStorageBackend",
|
||||
"Embedder",
|
||||
"LanceDBStorageBackend",
|
||||
"QdrantStorageBackend",
|
||||
"StorageBackend",
|
||||
"TortoiseStorageBackend",
|
||||
]
|
||||
@@ -0,0 +1,176 @@
|
||||
from abc import ABC, abstractmethod
|
||||
import asyncio
|
||||
import threading
|
||||
from typing import Any, Literal, Protocol, runtime_checkable
|
||||
|
||||
from zhenxun.services.ai.core.messages import EmbedBatch
|
||||
from zhenxun.services.ai.llm.api import embed as api_embed
|
||||
from zhenxun.services.ai.message_builder import MessageBuilder
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
EmbedTaskType = Literal[
|
||||
"general", "query", "document", "similarity", "classification", "clustering"
|
||||
]
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Embedder(Protocol):
|
||||
"""
|
||||
向量化引擎协议 (Callable Protocol)。
|
||||
任何实现了异步 __call__ 的对象或闭包函数均可作为 Embedder。
|
||||
"""
|
||||
|
||||
async def __call__(
|
||||
self, input_batch: Any, task: EmbedTaskType = "general", **kwargs
|
||||
) -> list[list[float]]:
|
||||
"""
|
||||
将文本、多模态或预构建的 EmbedBatch 转换为向量列表。
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class DefaultEmbedder(Embedder):
|
||||
"""系统默认的向量化引擎，调用大模型底座 API"""
|
||||
|
||||
def __init__(self, model_name: str | None = None, config: Any = None):
|
||||
self.model_name = model_name
|
||||
self.config = config
|
||||
|
||||
async def __call__(
|
||||
self, input_batch: Any, task: EmbedTaskType = "general", **kwargs
|
||||
) -> list[list[float]]:
|
||||
if not input_batch:
|
||||
return []
|
||||
try:
|
||||
res = await api_embed(
|
||||
input_batch, model=self.model_name, task=task, config=self.config
|
||||
)
|
||||
return res.embeddings
|
||||
except Exception as e:
|
||||
logger.error(f"DefaultEmbedder 向量化失败: {e}", e=e)
|
||||
return []
|
||||
|
||||
|
||||
class BaseLocalEmbedder(Embedder, ABC):
|
||||
"""本地向量化引擎基类，统一处理多模态降级与同步推理由协程包裹逻辑。"""
|
||||
|
||||
def __init__(self, model_name: str):
|
||||
self.model_name = model_name
|
||||
self._model_lock = threading.Lock()
|
||||
|
||||
@abstractmethod
|
||||
def _encode_texts(self, texts: list[str]) -> list[list[float]]:
|
||||
"""子类只需实现此同步的批量文本向量化方法即可。"""
|
||||
pass
|
||||
|
||||
async def __call__(
|
||||
self, input_batch: Any, task: EmbedTaskType = "general", **kwargs
|
||||
) -> list[list[float]]:
|
||||
if not input_batch:
|
||||
return []
|
||||
|
||||
if isinstance(input_batch, EmbedBatch):
|
||||
batch = input_batch
|
||||
else:
|
||||
batch = await MessageBuilder.normalize_to_embed_batch(input_batch)
|
||||
|
||||
texts = batch.to_text_only(f"本地模型 {self.model_name}")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
def _sync_embed():
|
||||
return self._encode_texts(texts)
|
||||
|
||||
return await asyncio.to_thread(_sync_embed)
|
||||
|
||||
|
||||
class FastEmbedder(BaseLocalEmbedder):
|
||||
"""
|
||||
基于 FastEmbed 的轻量级本地向量化引擎。
|
||||
零 PyTorch 依赖，CPU 推理极快。
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str | None = None):
|
||||
super().__init__(model_name or "BAAI/bge-small-zh-v1.5")
|
||||
self.model = None
|
||||
|
||||
import importlib.util
|
||||
|
||||
if importlib.util.find_spec("fastembed") is None:
|
||||
raise ImportError(
|
||||
"⚠️ 使用 FastEmbed 需要额外依赖，请在终端执行: pip install fastembed"
|
||||
)
|
||||
|
||||
def _ensure_model_loaded(self):
|
||||
"""线程安全的懒加载机制"""
|
||||
if self.model is None:
|
||||
with self._model_lock:
|
||||
if self.model is None:
|
||||
try:
|
||||
from fastembed import TextEmbedding
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"⚠️ 使用 FastEmbed 需要额外依赖，"
|
||||
"请在终端执行: pip install fastembed"
|
||||
)
|
||||
logger.info(
|
||||
f"正在后台加载 FastEmbed 本地模型: {self.model_name} ... "
|
||||
"(首次加载可能需要极长时间下载)"
|
||||
)
|
||||
self.model = TextEmbedding(model_name=self.model_name)
|
||||
logger.info(f"FastEmbed 模型 {self.model_name} 加载完毕！")
|
||||
|
||||
def _encode_texts(self, texts: list[str]) -> list[list[float]]:
|
||||
self._ensure_model_loaded()
|
||||
assert self.model is not None
|
||||
return [vec.tolist() for vec in self.model.embed(texts)]
|
||||
|
||||
|
||||
class SentenceTransformerEmbedder(BaseLocalEmbedder):
|
||||
"""
|
||||
基于 Sentence-Transformers 的本地向量化引擎。
|
||||
支持 GPU 加速，适合重度用户。
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str | None = None):
|
||||
super().__init__(model_name or "BAAI/bge-small-zh-v1.5")
|
||||
self.model = None
|
||||
|
||||
import importlib.util
|
||||
|
||||
if importlib.util.find_spec("sentence_transformers") is None:
|
||||
raise ImportError(
|
||||
"⚠️ 使用 SentenceTransformers 需要额外依赖，"
|
||||
"请在终端执行: pip install sentence-transformers"
|
||||
)
|
||||
|
||||
def _ensure_model_loaded(self):
|
||||
"""线程安全的懒加载机制"""
|
||||
if self.model is None:
|
||||
with self._model_lock:
|
||||
if self.model is None:
|
||||
try:
|
||||
from sentence_transformers import (
|
||||
SentenceTransformer,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"⚠️ 使用 SentenceTransformers 需要额外依赖，"
|
||||
"请在终端执行: pip install sentence-transformers"
|
||||
)
|
||||
logger.info(
|
||||
"正在后台加载 SentenceTransformer "
|
||||
f"本地模型: {self.model_name} ... "
|
||||
"(首次加载可能需要极长时间下载)"
|
||||
)
|
||||
self.model = SentenceTransformer(self.model_name)
|
||||
logger.info(
|
||||
f"SentenceTransformer 模型 {self.model_name} 加载完毕！"
|
||||
)
|
||||
|
||||
def _encode_texts(self, texts: list[str]) -> list[list[float]]:
|
||||
self._ensure_model_loaded()
|
||||
assert self.model is not None
|
||||
embeddings = self.model.encode(texts)
|
||||
return embeddings.tolist()
|
||||
@@ -0,0 +1,628 @@
|
||||
import os
|
||||
from typing import ClassVar, Protocol, runtime_checkable
|
||||
import uuid
|
||||
|
||||
import numpy as np
|
||||
from tortoise import fields
|
||||
|
||||
from zhenxun.services.ai.context.rag.models import (
|
||||
BaseRecord,
|
||||
QueryRequest,
|
||||
SearchResult,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.retrieval import FilterEvaluator
|
||||
from zhenxun.services.ai.utils.scope import ScopeSelector
|
||||
from zhenxun.services.db_context import Model
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class StorageBackend(Protocol):
|
||||
"""纯粹的向量存储后端协议"""
|
||||
|
||||
async def save(self, records: list[BaseRecord]) -> None:
|
||||
"""保存或更新数据块"""
|
||||
...
|
||||
|
||||
async def search(
|
||||
self, query: QueryRequest, scopes: list[str] | None = None
|
||||
) -> list[SearchResult]:
|
||||
"""按向量和前缀检索数据块"""
|
||||
...
|
||||
|
||||
async def update(self, record: BaseRecord) -> None:
|
||||
"""更新已有数据块"""
|
||||
...
|
||||
|
||||
async def delete(
|
||||
self, record_ids: list[str] | None = None, scope_prefix: str | None = None
|
||||
) -> int:
|
||||
"""删除数据块"""
|
||||
...
|
||||
|
||||
async def clear_by_query(self, query: ScopeSelector) -> int:
|
||||
"""根据统一领域查询对象清理数据块（在各实现中回退到 delete）"""
|
||||
...
|
||||
|
||||
async def get_all(self, scope_prefix: str | None = None) -> list[BaseRecord]:
|
||||
"""获取作用域下所有记录（用于容量控制）"""
|
||||
...
|
||||
|
||||
|
||||
def normalize_vector(vec: list[float] | np.ndarray) -> np.ndarray:
|
||||
"""将一维向量转化为 float32 数组并进行 L2 归一化"""
|
||||
v = np.array(vec, dtype=np.float32)
|
||||
norm = np.linalg.norm(v)
|
||||
if norm == 0:
|
||||
return v
|
||||
return v / norm
|
||||
|
||||
|
||||
def normalize_matrix(mat: np.ndarray) -> np.ndarray:
|
||||
"""将二维矩阵的每一行进行 L2 归一化"""
|
||||
norms = np.linalg.norm(mat, axis=1, keepdims=True)
|
||||
norms[norms == 0] = 1.0
|
||||
return mat / norms
|
||||
|
||||
|
||||
class DictStorageBackend(StorageBackend):
|
||||
"""基于内存字典的轻量级纯净 RAG 存储实现"""
|
||||
|
||||
_shared_records: ClassVar[dict[str, BaseRecord]] = {}
|
||||
_shared_vectors: ClassVar[dict[str, np.ndarray]] = {}
|
||||
|
||||
def __init__(self):
|
||||
self._records = self._shared_records
|
||||
self._vectors = self._shared_vectors
|
||||
|
||||
async def save(self, records: list[BaseRecord]) -> None:
|
||||
for r in records:
|
||||
self._records[r.id] = r
|
||||
if r.embedding:
|
||||
self._vectors[r.id] = normalize_vector(r.embedding)
|
||||
else:
|
||||
self._vectors.pop(r.id, None)
|
||||
|
||||
async def search(
|
||||
self, query: QueryRequest, scopes: list[str] | None = None
|
||||
) -> list[SearchResult]:
|
||||
candidate_ids = []
|
||||
for record in self._records.values():
|
||||
if scopes is not None:
|
||||
if record.metadata.get("scope", "/") not in scopes:
|
||||
continue
|
||||
if not FilterEvaluator.evaluate(record.metadata, query.metadata_filters):
|
||||
continue
|
||||
if not query.embedding and query.text and query.text not in record.content:
|
||||
continue
|
||||
candidate_ids.append(record.id)
|
||||
|
||||
if not candidate_ids:
|
||||
return []
|
||||
|
||||
results = []
|
||||
if query.search_type == "sparse":
|
||||
import jieba
|
||||
|
||||
tokens = set(jieba.lcut_for_search(query.text.lower()))
|
||||
for r_id in candidate_ids:
|
||||
record = self._records[r_id]
|
||||
content = record.content.lower()
|
||||
matched_count = sum(1 for t in tokens if t in content)
|
||||
if matched_count > 0:
|
||||
score = matched_count / len(tokens)
|
||||
results.append(SearchResult(record=record, score=score))
|
||||
elif query.search_type == "dense" and query.embedding:
|
||||
q_vec = normalize_vector(query.embedding)
|
||||
valid_ids = [r_id for r_id in candidate_ids if r_id in self._vectors]
|
||||
|
||||
if valid_ids:
|
||||
try:
|
||||
mat = np.array([self._vectors[r_id] for r_id in valid_ids])
|
||||
scores = mat @ q_vec
|
||||
for r_id, score in zip(valid_ids, scores):
|
||||
results.append(
|
||||
SearchResult(record=self._records[r_id], score=float(score))
|
||||
)
|
||||
except ValueError as e:
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
logger.warning(
|
||||
"⚠️ DictStorage 中缓存的向量维度与当前查询维度不匹配，"
|
||||
f"跳过向量检索。原因: {e}"
|
||||
)
|
||||
|
||||
missing_ids = [r_id for r_id in candidate_ids if r_id not in self._vectors]
|
||||
for r_id in missing_ids:
|
||||
results.append(SearchResult(record=self._records[r_id], score=0.1))
|
||||
else:
|
||||
for r_id in candidate_ids:
|
||||
results.append(SearchResult(record=self._records[r_id], score=0.1))
|
||||
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
return results[: query.limit]
|
||||
|
||||
async def update(self, record: BaseRecord) -> None:
|
||||
if record.id in self._records:
|
||||
self._records[record.id] = record
|
||||
if record.embedding:
|
||||
self._vectors[record.id] = normalize_vector(record.embedding)
|
||||
else:
|
||||
self._vectors.pop(record.id, None)
|
||||
|
||||
async def delete(
|
||||
self, record_ids: list[str] | None = None, scope_prefix: str | None = None
|
||||
) -> int:
|
||||
to_delete = []
|
||||
for r_id, r in self._records.items():
|
||||
if scope_prefix is not None:
|
||||
if not r.metadata.get("scope", "/").startswith(scope_prefix):
|
||||
continue
|
||||
if record_ids and r_id not in record_ids:
|
||||
continue
|
||||
to_delete.append(r_id)
|
||||
for r_id in to_delete:
|
||||
del self._records[r_id]
|
||||
self._vectors.pop(r_id, None)
|
||||
return len(to_delete)
|
||||
|
||||
async def clear_by_query(self, query: ScopeSelector) -> int:
|
||||
return await self.delete(scope_prefix=query.scope_prefix)
|
||||
|
||||
async def get_all(self, scope_prefix: str | None = None) -> list[BaseRecord]:
|
||||
res = []
|
||||
for r in self._records.values():
|
||||
if scope_prefix is not None:
|
||||
if r.metadata.get("scope", "/") != scope_prefix:
|
||||
continue
|
||||
res.append(r)
|
||||
return res
|
||||
|
||||
|
||||
class AbstractVectorRecord(Model):
|
||||
id = fields.CharField(pk=True, max_length=64)
|
||||
scope = fields.CharField(max_length=255, index=True)
|
||||
content = fields.TextField()
|
||||
embedding = fields.JSONField(null=True)
|
||||
meta_data = fields.JSONField(null=True)
|
||||
|
||||
class Meta: # type: ignore
|
||||
abstract = True
|
||||
|
||||
|
||||
class TortoiseStorageBackend(StorageBackend):
|
||||
def __init__(self, model_class: type[AbstractVectorRecord]):
|
||||
self.model_class = model_class
|
||||
|
||||
def _to_base_record(self, row: AbstractVectorRecord) -> BaseRecord:
|
||||
return BaseRecord(
|
||||
id=row.id,
|
||||
content=row.content,
|
||||
embedding=row.embedding if isinstance(row.embedding, list) else None,
|
||||
metadata=row.meta_data if isinstance(row.meta_data, dict) else {},
|
||||
)
|
||||
|
||||
async def save(self, records: list[BaseRecord]) -> None:
|
||||
for r in records:
|
||||
await self.model_class.update_or_create(
|
||||
id=r.id,
|
||||
defaults={
|
||||
"content": r.content,
|
||||
"scope": r.metadata.get("scope", "/"),
|
||||
"embedding": r.embedding,
|
||||
"meta_data": r.metadata,
|
||||
},
|
||||
)
|
||||
|
||||
async def search(
|
||||
self, query: QueryRequest, scopes: list[str] | None = None
|
||||
) -> list[SearchResult]:
|
||||
query_orm = self.model_class.all()
|
||||
if scopes is not None:
|
||||
query_orm = query_orm.filter(scope__in=scopes)
|
||||
|
||||
if query.search_type == "sparse" and query.text:
|
||||
import jieba
|
||||
from tortoise.expressions import Q
|
||||
|
||||
tokens = [
|
||||
t for t in jieba.lcut_for_search(query.text) if len(t.strip()) > 1
|
||||
] or [query.text]
|
||||
q_expr = Q()
|
||||
for token in tokens:
|
||||
q_expr |= Q(content__icontains=token)
|
||||
query_orm = query_orm.filter(q_expr)
|
||||
elif query.search_type == "dense" and not query.embedding and query.text:
|
||||
query_orm = query_orm.filter(content__icontains=query.text)
|
||||
|
||||
rows = await query_orm
|
||||
|
||||
valid_rows = []
|
||||
for row in rows:
|
||||
row_meta = row.meta_data if isinstance(row.meta_data, dict) else {}
|
||||
if not FilterEvaluator.evaluate(row_meta, query.metadata_filters):
|
||||
continue
|
||||
valid_rows.append(row)
|
||||
|
||||
if not valid_rows:
|
||||
return []
|
||||
|
||||
results = []
|
||||
if query.search_type == "sparse":
|
||||
import jieba
|
||||
|
||||
tokens = set(jieba.lcut_for_search(query.text.lower()))
|
||||
for row in valid_rows:
|
||||
content = row.content.lower()
|
||||
matched_count = sum(1 for t in tokens if t in content)
|
||||
score = matched_count / len(tokens) if tokens else 0.1
|
||||
results.append(
|
||||
SearchResult(record=self._to_base_record(row), score=score)
|
||||
)
|
||||
elif query.search_type == "dense" and query.embedding:
|
||||
q_vec = normalize_vector(query.embedding)
|
||||
vec_rows = []
|
||||
missing_rows = []
|
||||
|
||||
for row in valid_rows:
|
||||
if isinstance(row.embedding, list):
|
||||
vec_rows.append(row)
|
||||
else:
|
||||
missing_rows.append(row)
|
||||
|
||||
if vec_rows:
|
||||
try:
|
||||
raw_mat = np.array(
|
||||
[r.embedding for r in vec_rows], dtype=np.float32
|
||||
)
|
||||
norm_mat = normalize_matrix(raw_mat)
|
||||
scores = norm_mat @ q_vec
|
||||
|
||||
for row, score in zip(vec_rows, scores):
|
||||
results.append(
|
||||
SearchResult(
|
||||
record=self._to_base_record(row), score=float(score)
|
||||
)
|
||||
)
|
||||
except ValueError as e:
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
logger.warning(
|
||||
"⚠️ 数据库中缓存的向量维度与当前模型查询维度不匹配，"
|
||||
f"已安全跳过向量检索(降级为稀疏匹配)。原因: {e}"
|
||||
)
|
||||
|
||||
for row in missing_rows:
|
||||
results.append(
|
||||
SearchResult(record=self._to_base_record(row), score=0.1)
|
||||
)
|
||||
else:
|
||||
for row in valid_rows:
|
||||
results.append(
|
||||
SearchResult(record=self._to_base_record(row), score=0.1)
|
||||
)
|
||||
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
return results[: query.limit]
|
||||
|
||||
async def update(self, record: BaseRecord) -> None:
|
||||
await self.model_class.filter(id=record.id).update(
|
||||
content=record.content,
|
||||
scope=record.metadata.get("scope", "/"),
|
||||
embedding=record.embedding,
|
||||
meta_data=record.metadata,
|
||||
)
|
||||
|
||||
async def delete(
|
||||
self, record_ids: list[str] | None = None, scope_prefix: str | None = None
|
||||
) -> int:
|
||||
query = self.model_class.all()
|
||||
if scope_prefix is not None:
|
||||
query = query.filter(scope__startswith=scope_prefix)
|
||||
if record_ids is not None:
|
||||
if not record_ids:
|
||||
return 0
|
||||
query = query.filter(id__in=record_ids)
|
||||
|
||||
return await query.delete()
|
||||
|
||||
async def clear_by_query(self, query: ScopeSelector) -> int:
|
||||
return await self.delete(scope_prefix=query.scope_prefix)
|
||||
|
||||
async def get_all(self, scope_prefix: str | None = None) -> list[BaseRecord]:
|
||||
query = self.model_class.all()
|
||||
if scope_prefix is not None:
|
||||
query = query.filter(scope=scope_prefix)
|
||||
rows = await query
|
||||
return [self._to_base_record(row) for row in rows]
|
||||
|
||||
|
||||
class QdrantStorageBackend(StorageBackend):
|
||||
"""Qdrant 向量数据库可选存储后端"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
location: str = ":memory:",
|
||||
collection_name: str = "zhenxun_rag",
|
||||
**kwargs,
|
||||
):
|
||||
try:
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"缺少 Qdrant 依赖！请执行 `pip install qdrant-client` 安装"
|
||||
)
|
||||
|
||||
self.client = AsyncQdrantClient(location=location, **kwargs)
|
||||
self.collection_name = collection_name
|
||||
self._initialized = False
|
||||
|
||||
async def _ensure_collection(self, dim: int):
|
||||
if self._initialized:
|
||||
return
|
||||
from qdrant_client.models import Distance, VectorParams
|
||||
|
||||
if not await self.client.collection_exists(self.collection_name):
|
||||
await self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
|
||||
)
|
||||
self._initialized = True
|
||||
|
||||
async def save(self, records: list[BaseRecord]) -> None:
|
||||
if not records:
|
||||
return
|
||||
dim = len(records[0].embedding) if records[0].embedding else 1536
|
||||
await self._ensure_collection(dim)
|
||||
|
||||
from qdrant_client.models import PointStruct
|
||||
|
||||
points = []
|
||||
for r in records:
|
||||
points.append(
|
||||
PointStruct(
|
||||
id=r.id
|
||||
if len(r.id) == 36
|
||||
else str(uuid.uuid5(uuid.NAMESPACE_DNS, r.id)),
|
||||
vector=r.embedding or [],
|
||||
payload={"content": r.content, "metadata": r.metadata},
|
||||
)
|
||||
)
|
||||
await self.client.upsert(collection_name=self.collection_name, points=points)
|
||||
|
||||
async def search(
|
||||
self, query: QueryRequest, scopes: list[str] | None = None
|
||||
) -> list[SearchResult]:
|
||||
if query.search_type == "dense" and not query.embedding:
|
||||
return []
|
||||
if query.embedding:
|
||||
await self._ensure_collection(len(query.embedding))
|
||||
|
||||
from qdrant_client.models import FieldCondition, Filter, MatchText, MatchValue
|
||||
|
||||
must_conditions = []
|
||||
|
||||
if scopes is not None:
|
||||
try:
|
||||
from qdrant_client.models import MatchAny
|
||||
|
||||
must_conditions.append(
|
||||
FieldCondition(key="metadata.scope", match=MatchAny(any=scopes))
|
||||
)
|
||||
except ImportError:
|
||||
scope_conditions = [
|
||||
FieldCondition(key="metadata.scope", match=MatchValue(value=s))
|
||||
for s in scopes
|
||||
]
|
||||
must_conditions.append(Filter(should=scope_conditions))
|
||||
|
||||
if query.metadata_filters:
|
||||
for k, v in query.metadata_filters.items():
|
||||
must_conditions.append(
|
||||
FieldCondition(key=f"metadata.{k}", match=MatchValue(value=v))
|
||||
)
|
||||
|
||||
if query.search_type == "sparse":
|
||||
must_conditions.append(
|
||||
FieldCondition(key="content", match=MatchText(text=query.text))
|
||||
)
|
||||
|
||||
query_filter = Filter(must=must_conditions) if must_conditions else None
|
||||
|
||||
if query.search_type == "sparse":
|
||||
results = await self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=query_filter,
|
||||
limit=query.limit,
|
||||
with_payload=True,
|
||||
)
|
||||
return [
|
||||
SearchResult(
|
||||
record=BaseRecord(
|
||||
id=str(r.id),
|
||||
content=(r.payload or {}).get("content", ""),
|
||||
metadata=(r.payload or {}).get("metadata", {}),
|
||||
),
|
||||
score=1.0,
|
||||
)
|
||||
for r in results[0]
|
||||
]
|
||||
|
||||
results = await self.client.search( # type: ignore
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query.embedding,
|
||||
limit=query.limit,
|
||||
query_filter=query_filter,
|
||||
)
|
||||
|
||||
return [
|
||||
SearchResult(
|
||||
record=BaseRecord(
|
||||
id=str(r.id),
|
||||
content=r.payload.get("content", ""),
|
||||
metadata=r.payload.get("metadata", {}),
|
||||
),
|
||||
score=r.score,
|
||||
)
|
||||
for r in results
|
||||
]
|
||||
|
||||
async def update(self, record: BaseRecord) -> None:
|
||||
await self.save([record])
|
||||
|
||||
async def delete(
|
||||
self, record_ids: list[str] | None = None, scope_prefix: str | None = None
|
||||
) -> int:
|
||||
if not await self.client.collection_exists(self.collection_name):
|
||||
return 0
|
||||
from qdrant_client.models import FieldCondition, Filter, MatchText
|
||||
|
||||
query_filter = None
|
||||
if scope_prefix is not None:
|
||||
query_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="metadata.scope", match=MatchText(text=scope_prefix)
|
||||
)
|
||||
]
|
||||
)
|
||||
if query_filter:
|
||||
await self.client.delete(
|
||||
collection_name=self.collection_name, points_selector=query_filter
|
||||
)
|
||||
return 1
|
||||
|
||||
async def clear_by_query(self, query: ScopeSelector) -> int:
|
||||
return await self.delete(scope_prefix=query.scope_prefix)
|
||||
|
||||
async def get_all(self, scope_prefix: str | None = None) -> list[BaseRecord]:
|
||||
if not await self.client.collection_exists(self.collection_name):
|
||||
return []
|
||||
from qdrant_client.models import FieldCondition, Filter, MatchText
|
||||
|
||||
q_filter = None
|
||||
if scope_prefix and scope_prefix != "/":
|
||||
q_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="metadata.scope", match=MatchText(text=scope_prefix)
|
||||
)
|
||||
]
|
||||
)
|
||||
res = await self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=q_filter,
|
||||
limit=10000,
|
||||
with_payload=True,
|
||||
)
|
||||
return [
|
||||
BaseRecord(
|
||||
id=str(r.id),
|
||||
content=(r.payload or {}).get("content", ""),
|
||||
metadata=(r.payload or {}).get("metadata", {}),
|
||||
)
|
||||
for r in res[0]
|
||||
]
|
||||
|
||||
|
||||
class LanceDBStorageBackend(StorageBackend):
|
||||
"""LanceDB 向量数据库可选存储后端"""
|
||||
|
||||
def __init__(
|
||||
self, uri: str = "./data/lancedb", table_name: str = "zhenxun_rag", **kwargs
|
||||
):
|
||||
try:
|
||||
import lancedb
|
||||
except ImportError:
|
||||
raise ImportError("缺少 LanceDB 依赖！请执行 `pip install lancedb` 安装")
|
||||
|
||||
os.makedirs(
|
||||
os.path.dirname(uri) if os.path.dirname(uri) else ".", exist_ok=True
|
||||
)
|
||||
self.db = lancedb.connect(uri)
|
||||
self.table_name = table_name
|
||||
|
||||
async def save(self, records: list[BaseRecord]) -> None:
|
||||
if not records:
|
||||
return
|
||||
data = []
|
||||
dim = len(records[0].embedding) if records[0].embedding else 0
|
||||
|
||||
for r in records:
|
||||
data.append(
|
||||
{
|
||||
"id": r.id,
|
||||
"vector": r.embedding or [0.0] * dim,
|
||||
"content": r.content,
|
||||
"metadata": str(r.metadata),
|
||||
}
|
||||
)
|
||||
|
||||
if self.table_name not in self.db.table_names():
|
||||
self.db.create_table(self.table_name, data=data)
|
||||
else:
|
||||
self.db.open_table(self.table_name).add(data)
|
||||
|
||||
async def search(
|
||||
self, query: QueryRequest, scopes: list[str] | None = None
|
||||
) -> list[SearchResult]:
|
||||
if self.table_name not in self.db.table_names():
|
||||
return []
|
||||
if query.search_type == "dense" and not query.embedding:
|
||||
return []
|
||||
|
||||
tbl = self.db.open_table(self.table_name)
|
||||
if query.search_type == "sparse":
|
||||
try:
|
||||
results = (
|
||||
tbl.search(query.text, query_type="fts")
|
||||
.limit(query.limit)
|
||||
.to_list()
|
||||
)
|
||||
except Exception as e:
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
logger.warning(f"LanceDB FTS 检索失败(可能是由于尚未创建FTS索引): {e}")
|
||||
return []
|
||||
else:
|
||||
results = tbl.search(query.embedding).limit(query.limit).to_list()
|
||||
|
||||
import ast
|
||||
|
||||
return [
|
||||
SearchResult(
|
||||
record=BaseRecord(
|
||||
id=r["id"],
|
||||
content=r["content"],
|
||||
metadata=ast.literal_eval(r["metadata"]) if "metadata" in r else {},
|
||||
),
|
||||
score=1.0 - r.get("_distance", 0.0),
|
||||
)
|
||||
for r in results
|
||||
]
|
||||
|
||||
async def update(self, record: BaseRecord) -> None:
|
||||
pass
|
||||
|
||||
async def delete(
|
||||
self, record_ids: list[str] | None = None, scope_prefix: str | None = None
|
||||
) -> int:
|
||||
return 0
|
||||
|
||||
async def clear_by_query(self, query: ScopeSelector) -> int:
|
||||
return await self.delete(scope_prefix=query.scope_prefix)
|
||||
|
||||
async def get_all(self, scope_prefix: str | None = None) -> list[BaseRecord]:
|
||||
if self.table_name not in self.db.table_names():
|
||||
return []
|
||||
tbl = self.db.open_table(self.table_name)
|
||||
df = tbl.to_pandas()
|
||||
import ast
|
||||
|
||||
res = []
|
||||
for _, row in df.iterrows():
|
||||
meta = ast.literal_eval(row["metadata"]) if "metadata" in row else {}
|
||||
if scope_prefix is not None:
|
||||
if meta.get("scope", "/") != scope_prefix:
|
||||
continue
|
||||
res.append(BaseRecord(id=row["id"], content=row["content"], metadata=meta))
|
||||
return res
|
||||
@@ -0,0 +1,272 @@
|
||||
from typing import Any
|
||||
|
||||
from zhenxun.services.ai.context.rag.backends import StorageBackend
|
||||
from zhenxun.services.ai.context.rag.configs import RAGConfig
|
||||
from zhenxun.services.ai.context.rag.engine import ScopedRAGClient
|
||||
from zhenxun.services.ai.context.rag.ingestion import (
|
||||
ChunkingStrategy,
|
||||
DedupNode,
|
||||
DocumentChunking,
|
||||
DynamicChunkingNode,
|
||||
EmbeddingNode,
|
||||
IndexPipeline,
|
||||
StorageCommitNode,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.retrieval import (
|
||||
BaseRetriever,
|
||||
DatabaseSparseRetriever,
|
||||
HybridRetriever,
|
||||
LifecyclePostProcessor,
|
||||
PipelineRetriever,
|
||||
PostProcessor,
|
||||
PreProcessor,
|
||||
RerankRetriever,
|
||||
VectorDBRetriever,
|
||||
)
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
|
||||
class RAGBuilder:
|
||||
"""
|
||||
RAG 管线组装构建器 (Fluent Builder Pattern)。
|
||||
使用内部状态驱动模式 (The Memory Pattern)，内部维护私有的 RAGConfig 实例。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, storage: StorageBackend | None = None, config: RAGConfig | None = None
|
||||
):
|
||||
self._config = config or RAGConfig()
|
||||
if storage is not None:
|
||||
self._config.storage = storage
|
||||
|
||||
def with_embedder(self, embedder: Any) -> "RAGBuilder":
|
||||
"""
|
||||
设置向量化引擎。
|
||||
|
||||
参数:
|
||||
embedder: 实现了向量化协议的引擎实例 (如 DefaultEmbedder, FastEmbedder)。
|
||||
"""
|
||||
self._config.embedder = embedder
|
||||
return self
|
||||
|
||||
def with_retriever(self, retriever: BaseRetriever) -> "RAGBuilder":
|
||||
"""
|
||||
替换底层的向量库查表算法，注入自定义召回器。
|
||||
|
||||
参数:
|
||||
retriever: 实现了 BaseRetriever 协议的自定义召回器实例。
|
||||
"""
|
||||
self._config.custom_retriever = retriever
|
||||
return self
|
||||
|
||||
def with_scope(self, scopes: str | list[str]) -> "RAGBuilder":
|
||||
"""
|
||||
设置数据隔离作用域。
|
||||
|
||||
参数:
|
||||
scopes: 支持单作用域前缀(字符串)或联合检索多作用域(列表)。
|
||||
如 "/" 或 ["/group_1", "/user_2"]。
|
||||
"""
|
||||
self._config.scopes = scopes
|
||||
return self
|
||||
|
||||
def with_chunking(self, strategy: ChunkingStrategy) -> "RAGBuilder":
|
||||
"""
|
||||
设置文档切块策略。
|
||||
|
||||
参数:
|
||||
strategy: 切块策略实例 (如 DocumentChunking, RecursiveCharacterChunking)。
|
||||
"""
|
||||
self._config.chunking.strategy = strategy
|
||||
return self
|
||||
|
||||
def enable_dedup(self, threshold: float = 0.98) -> "RAGBuilder":
|
||||
"""
|
||||
开启批处理入库去重。
|
||||
在文本分块入库前，通过对比向量相似度拦截高度重复的内容。
|
||||
|
||||
参数:
|
||||
threshold: 去重相似度阈值，大于该值的块将被判定为重复并丢弃。
|
||||
"""
|
||||
self._config.dedup.enable = True
|
||||
self._config.dedup.threshold = threshold
|
||||
return self
|
||||
|
||||
def disable_dedup(self) -> "RAGBuilder":
|
||||
"""
|
||||
关闭批处理入库去重。
|
||||
"""
|
||||
self._config.dedup.enable = False
|
||||
return self
|
||||
|
||||
def enable_rerank(self, model_name: str, top_n: int = 5) -> "RAGBuilder":
|
||||
"""
|
||||
开启大模型交叉注意力重排。
|
||||
对初筛召回的结果使用专用的 Rerank 模型进行二次排序，极大提升召回准确率。
|
||||
|
||||
参数:
|
||||
model_name: 用于重排序的模型名称 (如 'BAAI/bge-reranker-v2-m3')。
|
||||
top_n: 重排后最终保留并返回的文档数量。
|
||||
"""
|
||||
self._config.rerank.enable = True
|
||||
self._config.rerank.model_name = model_name
|
||||
self._config.rerank.top_n = top_n
|
||||
return self
|
||||
|
||||
def enable_hybrid_search(
|
||||
self, dense_weight: float = 0.7, sparse_weight: float = 0.3
|
||||
) -> "RAGBuilder":
|
||||
"""
|
||||
开启双轨混合检索 (Hybrid Search) 及 RRF 融合。
|
||||
并发调用向量检索 (Dense) 和关键词检索 (Sparse)，结合两者优势。
|
||||
|
||||
参数:
|
||||
dense_weight: 稠密向量检索的分数计算权重。
|
||||
sparse_weight: 稀疏关键词(BM25等)检索的分数计算权重。
|
||||
"""
|
||||
self._config.hybrid.enable = True
|
||||
self._config.hybrid.dense_weight = dense_weight
|
||||
self._config.hybrid.sparse_weight = sparse_weight
|
||||
return self
|
||||
|
||||
def enable_lifecycle_scoring(
|
||||
self,
|
||||
half_life_days: int = 30,
|
||||
decay_weight: float = 0.3,
|
||||
semantic_weight: float = 0.7,
|
||||
importance_weight: float = 0.0,
|
||||
reinforcement_weight: float = 0.2,
|
||||
) -> "RAGBuilder":
|
||||
"""
|
||||
开启生命周期打分后处理。
|
||||
综合考虑记忆的时间新鲜度、基础语义相似度、重要性以及访问频次，模拟人类记忆遗忘曲线。
|
||||
|
||||
参数:
|
||||
half_life_days: 时间衰减的半衰期(天)。经过这么多天后，时间得分衰减一半。
|
||||
decay_weight: 时间衰减得分所占权重。
|
||||
semantic_weight: 基础语义相似度所占权重。
|
||||
importance_weight: 客观重要性所占权重。
|
||||
reinforcement_weight: 访问强化(被回想次数越多越容易想起)所占权重。
|
||||
"""
|
||||
self._config.lifecycle.enable = True
|
||||
self._config.lifecycle.half_life_days = half_life_days
|
||||
self._config.lifecycle.decay_weight = decay_weight
|
||||
self._config.lifecycle.semantic_weight = semantic_weight
|
||||
self._config.lifecycle.importance_weight = importance_weight
|
||||
self._config.lifecycle.reinforcement_weight = reinforcement_weight
|
||||
return self
|
||||
|
||||
def add_pre_processor(self, processor: PreProcessor) -> "RAGBuilder":
|
||||
"""
|
||||
挂载自定义查询前处理器。
|
||||
|
||||
参数:
|
||||
processor: 实现了 PreProcessor 协议的处理器实例。
|
||||
"""
|
||||
self._config.pre_processors.append(processor)
|
||||
return self
|
||||
|
||||
def add_post_processor(self, processor: PostProcessor) -> "RAGBuilder":
|
||||
"""
|
||||
挂载自定义检索后处理器。
|
||||
|
||||
参数:
|
||||
processor: 实现了 PostProcessor 协议的处理器实例。
|
||||
"""
|
||||
self._config.post_processors.append(processor)
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
def resolve(cls, config: RAGConfig | dict | Any | None) -> RAGConfig:
|
||||
"""
|
||||
解析多种格式的配置，统一转换为 RAGConfig 对象。
|
||||
|
||||
参数:
|
||||
config: RAGConfig 实例、字典或其他类型配置。
|
||||
"""
|
||||
if isinstance(config, RAGConfig):
|
||||
return config
|
||||
if isinstance(config, dict):
|
||||
return RAGConfig(**config)
|
||||
return RAGConfig()
|
||||
|
||||
def build(self) -> ScopedRAGClient:
|
||||
"""
|
||||
完成所有积木组装，输出最终的 RAG 客户端实体。
|
||||
"""
|
||||
cfg = self._config
|
||||
storage = cfg.storage
|
||||
if not storage:
|
||||
raise ValueError("RAGBuilder 必须配置 storage 才能 build。")
|
||||
|
||||
embedder = cfg.embedder
|
||||
if not embedder:
|
||||
from zhenxun.services.ai.context.rag.backends import DefaultEmbedder
|
||||
from zhenxun.services.ai.llm.manager import get_default_model
|
||||
|
||||
embedder = DefaultEmbedder(model_name=get_default_model("embedding"))
|
||||
logger.debug("RAGBuilder: 未指定 Embedder，已使用系统默认 Embedder。")
|
||||
|
||||
chunking_strategy = cfg.chunking.strategy or DocumentChunking()
|
||||
|
||||
scopes = cfg.scopes
|
||||
nodes = [
|
||||
DynamicChunkingNode(chunking_strategy),
|
||||
EmbeddingNode(embedder),
|
||||
]
|
||||
if cfg.dedup.enable:
|
||||
nodes.append(DedupNode(cfg.dedup.threshold))
|
||||
|
||||
nodes.append(StorageCommitNode(storage))
|
||||
pipeline = IndexPipeline(nodes)
|
||||
|
||||
base_retriever: BaseRetriever = cfg.custom_retriever or VectorDBRetriever(
|
||||
storage,
|
||||
embedder,
|
||||
scopes[0] if isinstance(scopes, list) else scopes,
|
||||
)
|
||||
|
||||
if cfg.hybrid.enable:
|
||||
database_sparse_retriever = DatabaseSparseRetriever(
|
||||
storage, scopes[0] if isinstance(scopes, list) else scopes
|
||||
)
|
||||
base_retriever = HybridRetriever(
|
||||
dense_retriever=base_retriever,
|
||||
sparse_retriever=database_sparse_retriever,
|
||||
dense_weight=cfg.hybrid.dense_weight,
|
||||
sparse_weight=cfg.hybrid.sparse_weight,
|
||||
)
|
||||
|
||||
if cfg.rerank.enable and cfg.rerank.model_name:
|
||||
base_retriever = RerankRetriever(
|
||||
base_retriever, cfg.rerank.model_name, cfg.rerank.top_n
|
||||
)
|
||||
|
||||
pre_processors = list(cfg.pre_processors)
|
||||
|
||||
post_processors = list(cfg.post_processors)
|
||||
if cfg.lifecycle.enable:
|
||||
post_processors.append(
|
||||
LifecyclePostProcessor(
|
||||
half_life_days=cfg.lifecycle.half_life_days,
|
||||
decay_weight=cfg.lifecycle.decay_weight,
|
||||
semantic_weight=cfg.lifecycle.semantic_weight,
|
||||
importance_weight=cfg.lifecycle.importance_weight,
|
||||
reinforcement_weight=cfg.lifecycle.reinforcement_weight,
|
||||
)
|
||||
)
|
||||
|
||||
if pre_processors or post_processors:
|
||||
base_retriever = PipelineRetriever(
|
||||
base_retriever,
|
||||
pre_processors=pre_processors,
|
||||
post_processors=post_processors,
|
||||
)
|
||||
|
||||
return ScopedRAGClient(
|
||||
storage=storage,
|
||||
retriever=base_retriever,
|
||||
pipeline=pipeline,
|
||||
scopes=scopes,
|
||||
config=cfg,
|
||||
)
|
||||
@@ -0,0 +1,86 @@
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from zhenxun.services.ai.context.rag.backends import Embedder, StorageBackend
|
||||
from zhenxun.services.ai.context.rag.ingestion import ChunkingStrategy
|
||||
from zhenxun.services.ai.context.rag.retrieval import (
|
||||
BaseRetriever,
|
||||
PostProcessor,
|
||||
PreProcessor,
|
||||
)
|
||||
|
||||
StorageConfigType = dict[str, Any]
|
||||
|
||||
|
||||
class ChunkingConfig(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
strategy: ChunkingStrategy | None = None
|
||||
"""文档切块策略"""
|
||||
|
||||
|
||||
class DedupConfig(BaseModel):
|
||||
enable: bool = True
|
||||
"""是否开启入库去重"""
|
||||
threshold: float = 0.98
|
||||
"""去重相似度阈值"""
|
||||
|
||||
|
||||
class RerankConfig(BaseModel):
|
||||
enable: bool = False
|
||||
"""是否开启重排"""
|
||||
model_name: str | None = None
|
||||
"""重排使用的模型名称"""
|
||||
top_n: int = 5
|
||||
"""重排后保留的文档数量"""
|
||||
|
||||
|
||||
class HybridSearchConfig(BaseModel):
|
||||
enable: bool = False
|
||||
"""是否开启双轨混合检索及 RRF 融合"""
|
||||
dense_weight: float = 0.7
|
||||
"""向量检索权重"""
|
||||
sparse_weight: float = 0.3
|
||||
"""BM25 检索权重"""
|
||||
|
||||
|
||||
class LifecycleConfig(BaseModel):
|
||||
enable: bool = False
|
||||
"""是否开启生命周期打分(时间衰减+访问强化)后处理"""
|
||||
half_life_days: int = 30
|
||||
"""半衰期天数"""
|
||||
decay_weight: float = 0.3
|
||||
"""时间衰减权重"""
|
||||
semantic_weight: float = 0.7
|
||||
"""语义相似度权重"""
|
||||
importance_weight: float = 0.0
|
||||
"""重要性打分权重"""
|
||||
reinforcement_weight: float = 0.2
|
||||
"""访问强化得分权重"""
|
||||
|
||||
|
||||
class RAGConfig(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
storage: StorageBackend | None = None
|
||||
"""存储后端"""
|
||||
embedder: Embedder | None = None
|
||||
"""向量化引擎"""
|
||||
custom_retriever: BaseRetriever | None = None
|
||||
"""自定义召回器"""
|
||||
scopes: str | list[str] = "/"
|
||||
"""数据隔离作用域"""
|
||||
chunking: ChunkingConfig = Field(default_factory=ChunkingConfig)
|
||||
"""文档切块配置"""
|
||||
dedup: DedupConfig = Field(default_factory=DedupConfig)
|
||||
"""去重配置"""
|
||||
|
||||
rerank: RerankConfig = Field(default_factory=RerankConfig)
|
||||
"""重排配置"""
|
||||
hybrid: HybridSearchConfig = Field(default_factory=HybridSearchConfig)
|
||||
"""混合检索与 RRF 融合配置"""
|
||||
lifecycle: LifecycleConfig = Field(default_factory=LifecycleConfig)
|
||||
"""生命周期打分配置"""
|
||||
pre_processors: list[PreProcessor] = Field(default_factory=list)
|
||||
"""自定义查询前处理器列表"""
|
||||
post_processors: list[PostProcessor] = Field(default_factory=list)
|
||||
"""自定义检索后处理器列表"""
|
||||
@@ -0,0 +1,96 @@
|
||||
from typing import Any
|
||||
|
||||
from zhenxun.services.ai.context.rag.backends import (
|
||||
StorageBackend,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.configs import RAGConfig
|
||||
from zhenxun.services.ai.context.rag.ingestion import (
|
||||
IndexPipeline,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.models import (
|
||||
BaseRecord,
|
||||
SearchResult,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.retrieval import (
|
||||
BaseRetriever,
|
||||
)
|
||||
from zhenxun.services.ai.utils.scope import normalize_scope_path
|
||||
|
||||
|
||||
class ScopedRAGClient:
|
||||
"""
|
||||
RAG 基础设施门面。
|
||||
封装了存储后端、检索器和写入管线，将所有操作透明地限定在指定的作用域前缀下。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
storage: StorageBackend,
|
||||
retriever: BaseRetriever,
|
||||
pipeline: IndexPipeline,
|
||||
scopes: str | list[str] = "/",
|
||||
config: RAGConfig | None = None,
|
||||
):
|
||||
"""
|
||||
初始化 ScopedRAGClient 实例。
|
||||
|
||||
参数:
|
||||
storage: 底层向量/文档存储后端。
|
||||
retriever: 数据召回检索器。
|
||||
pipeline: 数据入库和索引分块处理管线。
|
||||
scopes: 数据隔离作用域，支持单作用域前缀或多作用域前缀列表。
|
||||
config: 全局的 RAG 配置项。
|
||||
"""
|
||||
self.storage = storage
|
||||
self.retriever = retriever
|
||||
self.pipeline = pipeline
|
||||
self.config = config or RAGConfig()
|
||||
self._background_tasks = set()
|
||||
if isinstance(scopes, str):
|
||||
self.scopes = [normalize_scope_path(scopes)]
|
||||
else:
|
||||
self.scopes = [normalize_scope_path(s) for s in scopes]
|
||||
|
||||
self.scope_prefix = self.scopes[0] if self.scopes else "/"
|
||||
|
||||
async def ingest(self, records: list[BaseRecord]) -> int:
|
||||
"""通过 RAG Ingestion Pipeline 处理并入库数据"""
|
||||
for r in records:
|
||||
if "scope" not in r.metadata:
|
||||
r.metadata["scope"] = self.scope_prefix
|
||||
|
||||
res = await self.pipeline.run(records)
|
||||
return len(res)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query: Any,
|
||||
limit: int = 10,
|
||||
scopes: str | list[str] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> list[SearchResult]:
|
||||
"""
|
||||
多作用域联合切片视图检索 (Union Search)。
|
||||
并发向多个独立的作用域发起检索，并对结果进行合并、去重和重排。
|
||||
"""
|
||||
target_scopes = self.scopes
|
||||
if scopes is not None:
|
||||
target_scopes = (
|
||||
[normalize_scope_path(scopes)]
|
||||
if isinstance(scopes, str)
|
||||
else [normalize_scope_path(s) for s in scopes]
|
||||
)
|
||||
|
||||
if not target_scopes:
|
||||
return []
|
||||
|
||||
kwargs["scopes"] = target_scopes
|
||||
return await self.retriever.retrieve(query, limit=limit, **kwargs)
|
||||
|
||||
async def update(self, record: BaseRecord) -> None:
|
||||
record.metadata["scope"] = self.scope_prefix
|
||||
await self.storage.update(record)
|
||||
|
||||
async def delete(self, record_ids: list[str] | None = None, **kwargs: Any) -> int:
|
||||
kwargs["scope_prefix"] = self.scope_prefix
|
||||
return await self.storage.delete(record_ids=record_ids, **kwargs)
|
||||
@@ -0,0 +1,521 @@
|
||||
from abc import ABC, abstractmethod
|
||||
import asyncio
|
||||
import re
|
||||
|
||||
from zhenxun.services.ai.context.rag.models import BaseRecord
|
||||
from zhenxun.services.ai.context.rag.utils import cosine_similarity
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
|
||||
class ChunkingStrategy(ABC):
|
||||
@abstractmethod
|
||||
def chunk(self, record: BaseRecord) -> list[BaseRecord]:
|
||||
raise NotImplementedError
|
||||
|
||||
def clean_text(self, text: str) -> str:
|
||||
cleaned_text = re.sub(r"\n+", "\n", text)
|
||||
cleaned_text = re.sub(r"[ \t]+", " ", cleaned_text)
|
||||
return cleaned_text.strip()
|
||||
|
||||
def _create_chunk_record(
|
||||
self, original_record: BaseRecord, chunk_number: int, content: str
|
||||
) -> BaseRecord:
|
||||
meta_data = original_record.metadata.copy()
|
||||
meta_data["chunk_index"] = chunk_number
|
||||
meta_data["chunk_size"] = len(content)
|
||||
meta_data["parent_id"] = original_record.id
|
||||
return BaseRecord(
|
||||
id=f"{original_record.id}_{chunk_number}",
|
||||
content=content,
|
||||
metadata=meta_data,
|
||||
)
|
||||
|
||||
|
||||
class DocumentChunking(ChunkingStrategy):
|
||||
"""段落语义分块策略 (按双换行切分)"""
|
||||
|
||||
def __init__(self, chunk_size: int = 1000):
|
||||
"""
|
||||
初始化段落语义分块策略。
|
||||
|
||||
参数:
|
||||
chunk_size: 单个分块的最大字符长度限制，默认 1000。
|
||||
"""
|
||||
self.chunk_size = chunk_size
|
||||
|
||||
def chunk(self, record: BaseRecord) -> list[BaseRecord]:
|
||||
if len(record.content) <= self.chunk_size:
|
||||
return [
|
||||
self._create_chunk_record(record, 0, self.clean_text(record.content))
|
||||
]
|
||||
|
||||
raw_paragraphs = record.content.split("\n\n")
|
||||
paragraphs = [self.clean_text(para) for para in raw_paragraphs if para.strip()]
|
||||
|
||||
chunks: list[BaseRecord] = []
|
||||
current_chunk_texts = []
|
||||
current_length = 0
|
||||
chunk_index = 0
|
||||
|
||||
for para in paragraphs:
|
||||
para_len = len(para)
|
||||
if current_length + para_len > self.chunk_size and current_chunk_texts:
|
||||
chunk_content = "\n\n".join(current_chunk_texts)
|
||||
chunks.append(
|
||||
self._create_chunk_record(record, chunk_index, chunk_content)
|
||||
)
|
||||
chunk_index += 1
|
||||
current_chunk_texts = []
|
||||
current_length = 0
|
||||
|
||||
current_chunk_texts.append(para)
|
||||
current_length += para_len + 2
|
||||
|
||||
if current_chunk_texts:
|
||||
chunk_content = "\n\n".join(current_chunk_texts)
|
||||
chunks.append(self._create_chunk_record(record, chunk_index, chunk_content))
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
class RecursiveCharacterChunking(ChunkingStrategy):
|
||||
"""递归字符分块策略"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_size: int = 1000,
|
||||
overlap: int = 100,
|
||||
separators: list[str] | None = None,
|
||||
):
|
||||
"""
|
||||
初始化递归字符分块策略。
|
||||
|
||||
参数:
|
||||
chunk_size: 单个分块的最大字符长度限制，默认 1000。
|
||||
overlap: 相邻分块之间的重叠字符长度，默认 100。
|
||||
separators: 用于切分文本的候选分隔符列表，按优先级从高到低尝试，
|
||||
默认包含段落、句子和常见标点。
|
||||
"""
|
||||
if overlap >= chunk_size:
|
||||
raise ValueError(f"重叠长度 ({overlap}) 必须小于分块大小 ({chunk_size})")
|
||||
self.chunk_size = chunk_size
|
||||
self.overlap = overlap
|
||||
self.separators = separators or [
|
||||
"\n\n",
|
||||
"\n",
|
||||
"。",
|
||||
"！",
|
||||
"？",
|
||||
"；",
|
||||
"，",
|
||||
" ",
|
||||
"",
|
||||
]
|
||||
|
||||
def _split_text(self, text: str, separators: list[str]) -> list[str]:
|
||||
"""核心递归切分逻辑"""
|
||||
final_chunks = []
|
||||
separator = separators[-1]
|
||||
new_separators = []
|
||||
|
||||
for i, _s in enumerate(separators):
|
||||
if _s == "":
|
||||
separator = _s
|
||||
break
|
||||
if _s in text:
|
||||
separator = _s
|
||||
new_separators = separators[i + 1 :]
|
||||
break
|
||||
|
||||
if separator:
|
||||
splits = [s for s in text.split(separator) if s]
|
||||
else:
|
||||
splits = list(text)
|
||||
|
||||
good_splits = []
|
||||
for s in splits:
|
||||
if len(s) < self.chunk_size:
|
||||
good_splits.append(s)
|
||||
else:
|
||||
if good_splits:
|
||||
merged_chunks = self._merge_splits(good_splits, separator)
|
||||
final_chunks.extend(merged_chunks)
|
||||
good_splits = []
|
||||
if new_separators:
|
||||
final_chunks.extend(self._split_text(s, new_separators))
|
||||
else:
|
||||
for i in range(0, len(s), self.chunk_size):
|
||||
final_chunks.append(s[i : i + self.chunk_size])
|
||||
|
||||
if good_splits:
|
||||
merged_chunks = self._merge_splits(good_splits, separator)
|
||||
final_chunks.extend(merged_chunks)
|
||||
|
||||
return final_chunks
|
||||
|
||||
def _merge_splits(self, splits: list[str], separator: str) -> list[str]:
|
||||
"""将零散的切片合并为符合 chunk_size 的块，并处理 Overlap"""
|
||||
chunks = []
|
||||
current_chunk = []
|
||||
current_length = 0
|
||||
|
||||
for split in splits:
|
||||
split_len = len(split)
|
||||
sep_len = len(separator) if current_chunk else 0
|
||||
|
||||
if current_length + sep_len + split_len > self.chunk_size and current_chunk:
|
||||
chunk_str = separator.join(current_chunk)
|
||||
chunks.append(chunk_str)
|
||||
|
||||
while current_length > self.overlap or (
|
||||
current_length + sep_len + split_len > self.chunk_size
|
||||
and len(current_chunk) > 0
|
||||
):
|
||||
popped = current_chunk.pop(0)
|
||||
current_length -= len(popped) + (
|
||||
len(separator) if current_chunk else 0
|
||||
)
|
||||
sep_len = len(separator) if current_chunk else 0
|
||||
|
||||
current_chunk.append(split)
|
||||
current_length += sep_len + split_len
|
||||
|
||||
if current_chunk:
|
||||
chunk_str = separator.join(current_chunk)
|
||||
chunks.append(chunk_str)
|
||||
|
||||
return chunks
|
||||
|
||||
def chunk(self, record: BaseRecord) -> list[BaseRecord]:
|
||||
content = record.content.strip()
|
||||
|
||||
if len(content) <= self.chunk_size:
|
||||
return [self._create_chunk_record(record, 0, content)]
|
||||
|
||||
text_chunks = self._split_text(content, self.separators)
|
||||
|
||||
chunks: list[BaseRecord] = []
|
||||
for i, text_chunk in enumerate(text_chunks):
|
||||
clean_chunk = text_chunk.strip()
|
||||
if clean_chunk:
|
||||
chunks.append(self._create_chunk_record(record, i, clean_chunk))
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
class RowChunking(ChunkingStrategy):
|
||||
"""
|
||||
行数据分块策略 (专为 CSV/表格设计)
|
||||
核心特性：自动识别表头，并将其附加到每一个被切分的 Chunk 首部，防止上下文丢失。
|
||||
"""
|
||||
|
||||
def __init__(self, rows_per_chunk: int = 50):
|
||||
"""
|
||||
初始化表格行数据分块策略。
|
||||
|
||||
参数:
|
||||
rows_per_chunk: 每个分块包含的数据行数（不含表头），默认 50。
|
||||
"""
|
||||
self.rows_per_chunk = rows_per_chunk
|
||||
|
||||
def chunk(self, record: BaseRecord) -> list[BaseRecord]:
|
||||
lines = record.content.splitlines()
|
||||
lines = [line for line in lines if line.strip()]
|
||||
|
||||
if not lines:
|
||||
return []
|
||||
|
||||
header = lines[0]
|
||||
data_lines = lines[1:]
|
||||
|
||||
if not data_lines:
|
||||
return [self._create_chunk_record(record, 0, header)]
|
||||
|
||||
chunks: list[BaseRecord] = []
|
||||
chunk_index = 0
|
||||
|
||||
for i in range(0, len(data_lines), self.rows_per_chunk):
|
||||
chunk_lines = [header, *data_lines[i : i + self.rows_per_chunk]]
|
||||
chunk_content = "\n".join(chunk_lines)
|
||||
chunks.append(self._create_chunk_record(record, chunk_index, chunk_content))
|
||||
chunk_index += 1
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
class DeduplicationProcessor:
|
||||
"""
|
||||
入库批处理去重器 (Intra-batch Deduplication)。
|
||||
在 Chunk 存入数据库前，通过对比向量相似度，拦截高度重复的内容（如群聊复读机内容）。
|
||||
"""
|
||||
|
||||
def __init__(self, threshold: float = 0.98):
|
||||
"""
|
||||
初始化入库批处理去重处理器。
|
||||
|
||||
参数:
|
||||
threshold: 余弦相似度重复阈值，超过该阈值的块将被判定为重复并过滤，
|
||||
默认 0.98。
|
||||
"""
|
||||
self.threshold = threshold
|
||||
|
||||
async def process(self, records: list[BaseRecord]) -> list[BaseRecord]:
|
||||
if not records or len(records) <= 1:
|
||||
return records
|
||||
|
||||
kept_records: list[BaseRecord] = []
|
||||
dropped_count = 0
|
||||
|
||||
for record in records:
|
||||
if not record.embedding:
|
||||
kept_records.append(record)
|
||||
continue
|
||||
|
||||
is_duplicate = False
|
||||
for kept in kept_records:
|
||||
if not kept.embedding:
|
||||
continue
|
||||
sim = cosine_similarity(record.embedding, kept.embedding)
|
||||
if sim >= self.threshold:
|
||||
is_duplicate = True
|
||||
dropped_count += 1
|
||||
break
|
||||
|
||||
if not is_duplicate:
|
||||
kept_records.append(record)
|
||||
|
||||
if dropped_count > 0:
|
||||
logger.debug(
|
||||
f"🧹 [入库管线] 触发批处理去重，已拦截 {dropped_count} "
|
||||
f"个高度重复的 Chunk (阈值: {self.threshold})"
|
||||
)
|
||||
|
||||
return kept_records
|
||||
|
||||
|
||||
class BaseBatchNode(ABC):
|
||||
"""批处理节点基类：一次性接收并处理全部记录"""
|
||||
|
||||
@abstractmethod
|
||||
async def process_batch(self, records: list[BaseRecord]) -> list[BaseRecord]: ...
|
||||
|
||||
|
||||
class BaseMapNode(ABC):
|
||||
"""单映射节点基类：接收单条记录，引擎负责并发调度，返回None代表丢弃该数据"""
|
||||
|
||||
@abstractmethod
|
||||
async def process_one(
|
||||
self, record: BaseRecord
|
||||
) -> BaseRecord | list[BaseRecord] | None: ...
|
||||
|
||||
|
||||
class DynamicChunkingNode(BaseBatchNode):
|
||||
"""智能路由切块节点。根据记录的扩展名动态选择切块策略。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
default_strategy: ChunkingStrategy,
|
||||
custom_strategies: dict[str, ChunkingStrategy] | None = None,
|
||||
):
|
||||
"""
|
||||
初始化智能路由切块节点。
|
||||
|
||||
参数:
|
||||
default_strategy: 默认的切块策略。
|
||||
custom_strategies: 针对特定文件后缀的自定义切块策略映射表，
|
||||
默认 CSV 文件使用 RowChunking。
|
||||
"""
|
||||
self.default_strategy = default_strategy
|
||||
self.strategies = custom_strategies or {".csv": RowChunking(rows_per_chunk=30)}
|
||||
|
||||
async def process_batch(self, records: list[BaseRecord]) -> list[BaseRecord]:
|
||||
chunks = []
|
||||
for record in records:
|
||||
ext = record.metadata.get("extension", "")
|
||||
strategy = self.strategies.get(ext, self.default_strategy)
|
||||
chunks.extend(strategy.chunk(record))
|
||||
return chunks
|
||||
|
||||
|
||||
class BaseEmbeddingBatchNode(BaseBatchNode):
|
||||
"""批量向量化抽象基类：提取文本、分批请求 API 并将结果写回的公共逻辑"""
|
||||
|
||||
def __init__(self, embedder, batch_size: int = 80):
|
||||
"""
|
||||
初始化批量向量化抽象基类。
|
||||
|
||||
参数:
|
||||
embedder: 向量嵌入模型/函数，用于将文本生成向量。
|
||||
batch_size: 向量化请求的单批次大小限制，默认 80。
|
||||
"""
|
||||
self.embedder = embedder
|
||||
self.batch_size = batch_size
|
||||
|
||||
@abstractmethod
|
||||
def _filter_target_records(self, records: list[BaseRecord]) -> list[BaseRecord]:
|
||||
"""由子类实现：筛选出本次需要进行向量化的目标记录"""
|
||||
pass
|
||||
|
||||
async def process_batch(self, records: list[BaseRecord]) -> list[BaseRecord]:
|
||||
if not self.embedder or not records:
|
||||
return records
|
||||
|
||||
target_records = self._filter_target_records(records)
|
||||
if not target_records:
|
||||
return records
|
||||
|
||||
texts = [r.content for r in target_records]
|
||||
try:
|
||||
vecs = []
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch_texts = texts[i : i + self.batch_size]
|
||||
batch_vecs = await self.embedder(batch_texts, task="document")
|
||||
vecs.extend(batch_vecs)
|
||||
|
||||
for i, r in enumerate(target_records):
|
||||
if vecs and i < len(vecs) and vecs[i]:
|
||||
r.embedding = vecs[i]
|
||||
except Exception as e:
|
||||
logger.error(f"[{self.__class__.__name__}] 批量向量化失败: {e}")
|
||||
|
||||
return records
|
||||
|
||||
|
||||
class EmbeddingNode(BaseEmbeddingBatchNode):
|
||||
"""并发向量化初次构建节点"""
|
||||
|
||||
def _filter_target_records(self, records: list[BaseRecord]) -> list[BaseRecord]:
|
||||
return [r for r in records if r.content.strip()]
|
||||
|
||||
|
||||
class DedupNode(BaseBatchNode):
|
||||
"""批次内查重节点"""
|
||||
|
||||
def __init__(self, threshold: float):
|
||||
"""
|
||||
初始化批次内查重节点。
|
||||
|
||||
参数:
|
||||
threshold: 余弦相似度重复阈值，超过该阈值的块将被判定为重复并过滤。
|
||||
"""
|
||||
self.processor = DeduplicationProcessor(threshold=threshold)
|
||||
|
||||
async def process_batch(self, records: list[BaseRecord]) -> list[BaseRecord]:
|
||||
return await self.processor.process(records)
|
||||
|
||||
|
||||
class StorageCommitNode(BaseBatchNode):
|
||||
"""持久化事务提交节点 (Reduce)。统一收集意图并执行并发数据库 I/O。"""
|
||||
|
||||
def __init__(self, storage):
|
||||
"""
|
||||
初始化持久化事务提交节点。
|
||||
|
||||
参数:
|
||||
storage: 存储后端，负责将记录存入或删除。
|
||||
"""
|
||||
self.storage = storage
|
||||
|
||||
async def process_batch(self, records: list[BaseRecord]) -> list[BaseRecord]:
|
||||
if not records:
|
||||
return records
|
||||
|
||||
to_delete = set()
|
||||
to_update = []
|
||||
to_insert = []
|
||||
|
||||
for record in records:
|
||||
if record.action == "delete":
|
||||
to_delete.add(record.id)
|
||||
elif record.action == "update":
|
||||
to_update.append(record)
|
||||
elif record.action == "insert":
|
||||
to_insert.append(record)
|
||||
|
||||
if to_delete:
|
||||
await self.storage.delete(record_ids=list(to_delete))
|
||||
|
||||
if to_update:
|
||||
await asyncio.gather(*[self.storage.update(r) for r in to_update])
|
||||
|
||||
if to_insert:
|
||||
await self.storage.save(to_insert)
|
||||
|
||||
logger.debug(
|
||||
"💾 RAG 事务提交完成：插入 "
|
||||
f"{len(to_insert)} 条, 更新 {len(to_update)} 条, "
|
||||
f"删除 {len(to_delete)} 条。"
|
||||
)
|
||||
return to_insert + to_update
|
||||
|
||||
|
||||
class IndexPipeline:
|
||||
"""统一入库流水线 (Map-Reduce 范式并发调度引擎)"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
nodes: list[BaseBatchNode | BaseMapNode] | None = None,
|
||||
max_workers: int = 5,
|
||||
):
|
||||
"""
|
||||
初始化统一入库流水线。
|
||||
|
||||
参数:
|
||||
nodes: 管道节点列表，按顺序执行数据处理，默认 None。
|
||||
max_workers: 最大并发工作协程数，用于 Map 节点的并发调度，默认 5。
|
||||
"""
|
||||
self.nodes = nodes or []
|
||||
self.max_workers = max_workers
|
||||
|
||||
def add_node(self, node: BaseBatchNode | BaseMapNode):
|
||||
self.nodes.append(node)
|
||||
|
||||
async def run(self, records: list[BaseRecord]) -> list[BaseRecord]:
|
||||
if not records:
|
||||
return []
|
||||
|
||||
current_records = records
|
||||
for node in self.nodes:
|
||||
if not current_records:
|
||||
break
|
||||
|
||||
if isinstance(node, BaseBatchNode):
|
||||
current_records = await node.process_batch(current_records)
|
||||
elif isinstance(node, BaseMapNode):
|
||||
sem = asyncio.Semaphore(self.max_workers)
|
||||
map_node: BaseMapNode = node
|
||||
|
||||
async def _process_with_sem(r: BaseRecord):
|
||||
async with sem:
|
||||
return await map_node.process_one(r)
|
||||
|
||||
tasks = [_process_with_sem(r) for r in current_records]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
next_records = []
|
||||
for r in results:
|
||||
if isinstance(r, list):
|
||||
next_records.extend(r)
|
||||
elif r is not None:
|
||||
next_records.append(r)
|
||||
current_records = next_records
|
||||
else:
|
||||
raise ValueError(f"未知的管道节点类型: {type(node)}")
|
||||
|
||||
return current_records
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseBatchNode",
|
||||
"BaseMapNode",
|
||||
"ChunkingStrategy",
|
||||
"DedupNode",
|
||||
"DeduplicationProcessor",
|
||||
"DocumentChunking",
|
||||
"DynamicChunkingNode",
|
||||
"EmbeddingNode",
|
||||
"IndexPipeline",
|
||||
"RecursiveCharacterChunking",
|
||||
"RowChunking",
|
||||
"StorageCommitNode",
|
||||
]
|
||||
@@ -0,0 +1,47 @@
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class BaseRecord(BaseModel):
|
||||
"""RAG 基础记录载体，没有任何业务属性"""
|
||||
|
||||
id: str = Field(default_factory=lambda: __import__("uuid").uuid4().hex)
|
||||
"""记录的唯一标识符"""
|
||||
content: str = Field(...)
|
||||
"""数据块的文本内容"""
|
||||
embedding: list[float] | None = Field(default=None)
|
||||
"""数据块对应的向量嵌入"""
|
||||
metadata: dict[str, Any] = Field(default_factory=dict)
|
||||
"""数据块的元数据字典"""
|
||||
action: Literal["insert", "update", "delete", "ignore"] = Field(default="insert")
|
||||
"""数据块在索引管线中的操作意图"""
|
||||
|
||||
|
||||
class SearchResult(BaseModel):
|
||||
"""搜索结果"""
|
||||
|
||||
record: BaseRecord
|
||||
"""检索到的基础记录"""
|
||||
score: float
|
||||
"""检索相似度得分"""
|
||||
|
||||
|
||||
class QueryRequest(BaseModel):
|
||||
"""通用检索请求"""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
text: str = Field(default="")
|
||||
"""原始查询文本"""
|
||||
embedding: list[float] | None = Field(default=None)
|
||||
"""用于向量检索的数组"""
|
||||
search_type: Literal["dense", "sparse", "hybrid"] = Field(default="dense")
|
||||
"""检索类型标识：稠密向量、稀疏关键词或混合"""
|
||||
metadata_filters: dict[str, Any] | None = Field(default=None)
|
||||
"""元数据精确匹配字典"""
|
||||
limit: int = Field(default=10)
|
||||
"""返回的最大条数"""
|
||||
|
||||
|
||||
StorageConfigType = dict[str, Any]
|
||||
@@ -0,0 +1,439 @@
|
||||
from abc import abstractmethod
|
||||
import asyncio
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any, Protocol, cast, runtime_checkable
|
||||
|
||||
from zhenxun.services.ai.context.rag.models import QueryRequest, SearchResult
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from zhenxun.services.ai.context.rag.backends.storages import StorageBackend
|
||||
|
||||
|
||||
def normalize_query_text(query: Any) -> str:
|
||||
"""辅助函数：提取各种输入形式（如字符串、平台Message对象）的纯文本用于检索"""
|
||||
if isinstance(query, str):
|
||||
return query
|
||||
if hasattr(query, "extract_plain_text"):
|
||||
return query.extract_plain_text()
|
||||
return str(query) if query is not None else ""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class BaseRetriever(Protocol):
|
||||
"""
|
||||
检索器核心协议。
|
||||
任何实现了 retrieve 方法的对象均可作为检索器
|
||||
（不仅限于向量检索，也可包含 BM25、SQL 搜索等）。
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def retrieve(
|
||||
self, query: Any, limit: int = 10, **kwargs: Any
|
||||
) -> list[SearchResult]: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class PostProcessor(Protocol):
|
||||
"""后处理器协议（如重排、时间衰减打分等）。"""
|
||||
|
||||
@abstractmethod
|
||||
async def process(
|
||||
self, results: list[SearchResult], query: str
|
||||
) -> list[SearchResult]: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class PreProcessor(Protocol):
|
||||
"""预处理器协议（如 LLM Query 改写、意图提取等）。"""
|
||||
|
||||
@abstractmethod
|
||||
async def process(self, query: str) -> list[str]:
|
||||
"""接收原始查询，返回一个或多个处理/改写后的查询词"""
|
||||
...
|
||||
|
||||
|
||||
class FilterEvaluator:
|
||||
"""纯 Python 内存求值器，用于为轻量级 Storage 提供字典精确匹配过滤"""
|
||||
|
||||
@classmethod
|
||||
def evaluate(
|
||||
cls, metadata: dict[str, Any], filter_dict: dict[str, Any] | None
|
||||
) -> bool:
|
||||
if filter_dict is None:
|
||||
return True
|
||||
return all(metadata.get(k) == v for k, v in filter_dict.items())
|
||||
|
||||
|
||||
class VectorDBRetriever(BaseRetriever):
|
||||
"""基于向量数据库的标准检索器"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
storage: "StorageBackend",
|
||||
embedder: Any,
|
||||
scope_prefix: str | None = None,
|
||||
score_threshold: float = 0.4,
|
||||
):
|
||||
"""
|
||||
初始化向量数据库检索器。
|
||||
|
||||
参数:
|
||||
storage: 存储后端，用于执行向量相似度搜索。
|
||||
embedder: 向量嵌入模型/函数，用于将文本转换为向量。
|
||||
scope_prefix: 作用域前缀，用于限制检索范围，默认 None。
|
||||
score_threshold: 分数阈值，过滤掉相似度低于该值的检索结果，默认 0.4。
|
||||
"""
|
||||
self.storage = storage
|
||||
self.embedder = embedder
|
||||
self.scope_prefix = scope_prefix
|
||||
self.score_threshold = score_threshold
|
||||
|
||||
async def retrieve(
|
||||
self, query: Any, limit: int = 10, **kwargs: Any
|
||||
) -> list[SearchResult]:
|
||||
text_query = normalize_query_text(query)
|
||||
|
||||
vecs = await self.embedder(query, task="query")
|
||||
query_vec = vecs[0] if vecs else None
|
||||
|
||||
if not text_query.strip() and not query_vec:
|
||||
return []
|
||||
|
||||
req = QueryRequest(
|
||||
text=text_query,
|
||||
embedding=query_vec,
|
||||
limit=limit * 2,
|
||||
search_type="dense",
|
||||
metadata_filters=kwargs.get("metadata_filters"),
|
||||
)
|
||||
effective_scopes = kwargs.get(
|
||||
"scopes", [self.scope_prefix] if self.scope_prefix else None
|
||||
)
|
||||
results = await self.storage.search(req, scopes=effective_scopes)
|
||||
return [r for r in results if r.score >= self.score_threshold][:limit]
|
||||
|
||||
|
||||
class DatabaseSparseRetriever(BaseRetriever):
|
||||
"""纯数据库下沉的稀疏检索器 (Keyword/FTS)"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
storage: "StorageBackend",
|
||||
scope_prefix: str | None = None,
|
||||
score_threshold: float = 0.0,
|
||||
):
|
||||
"""
|
||||
初始化数据库稀疏检索器。
|
||||
|
||||
参数:
|
||||
storage: 存储后端，用于执行全文检索/关键词检索。
|
||||
scope_prefix: 作用域前缀，用于限制检索范围，默认 None。
|
||||
score_threshold: 分数阈值，过滤掉相关度低于该值的检索结果，默认 0.0。
|
||||
"""
|
||||
self.storage = storage
|
||||
self.scope_prefix = scope_prefix
|
||||
self.score_threshold = score_threshold
|
||||
|
||||
async def retrieve(
|
||||
self, query: Any, limit: int = 10, **kwargs: Any
|
||||
) -> list[SearchResult]:
|
||||
text_query = normalize_query_text(query)
|
||||
if not text_query.strip():
|
||||
return []
|
||||
|
||||
req = QueryRequest(
|
||||
text=text_query,
|
||||
limit=limit * 2,
|
||||
search_type="sparse",
|
||||
metadata_filters=kwargs.get("metadata_filters"),
|
||||
)
|
||||
effective_scopes = kwargs.get(
|
||||
"scopes", [self.scope_prefix] if self.scope_prefix else None
|
||||
)
|
||||
results = await self.storage.search(req, scopes=effective_scopes)
|
||||
return [r for r in results if r.score > self.score_threshold][:limit]
|
||||
|
||||
|
||||
class RerankRetriever(BaseRetriever):
|
||||
"""带大模型交叉注意力重排的高阶检索器 (Decorator Pattern)"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_retriever: BaseRetriever,
|
||||
model_name: str | None = None,
|
||||
top_n: int = 5,
|
||||
oversample_factor: int = 2,
|
||||
min_oversample: int = 20,
|
||||
):
|
||||
"""
|
||||
初始化重排检索器。
|
||||
|
||||
参数:
|
||||
base_retriever: 基础检索器，用于初筛。
|
||||
model_name: 重排模型的名称，默认 None。
|
||||
top_n: 重排后保留的前 N 个文档数，默认 5。
|
||||
oversample_factor: 过采样系数，决定初筛检索的文档数量倍数，默认 2。
|
||||
min_oversample: 最小过采样文档数，默认 20。
|
||||
"""
|
||||
self.base_retriever = base_retriever
|
||||
self.model_name = model_name
|
||||
self.top_n = top_n
|
||||
self.oversample_factor = oversample_factor
|
||||
self.min_oversample = min_oversample
|
||||
|
||||
async def retrieve(
|
||||
self, query: Any, limit: int = 10, **kwargs: Any
|
||||
) -> list[SearchResult]:
|
||||
oversample_limit = max(limit * self.oversample_factor, self.min_oversample)
|
||||
initial_results = await self.base_retriever.retrieve(
|
||||
query, limit=oversample_limit, **kwargs
|
||||
)
|
||||
|
||||
if not initial_results:
|
||||
return []
|
||||
|
||||
text_query = normalize_query_text(query)
|
||||
|
||||
docs: list[str | dict[str, str]] = [
|
||||
res.record.content for res in initial_results
|
||||
]
|
||||
|
||||
from zhenxun.services.ai.llm.api import rerank
|
||||
|
||||
try:
|
||||
reranked = await rerank(
|
||||
query=text_query,
|
||||
documents=docs,
|
||||
top_n=min(limit, self.top_n),
|
||||
model=self.model_name,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Rerank 重排请求失败，将降级返回初筛结果: {e}")
|
||||
return initial_results[:limit]
|
||||
|
||||
final_results = []
|
||||
for rr in reranked:
|
||||
original_res = initial_results[rr.index]
|
||||
original_res.score = rr.relevance_score
|
||||
final_results.append(original_res)
|
||||
|
||||
return final_results
|
||||
|
||||
|
||||
class PipelineRetriever(BaseRetriever):
|
||||
"""支持挂载多个后处理器的流水线检索器"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_retriever: BaseRetriever,
|
||||
post_processors: list[PostProcessor] | None = None,
|
||||
pre_processors: list[PreProcessor] | None = None,
|
||||
):
|
||||
"""
|
||||
初始化流水线检索器。
|
||||
|
||||
参数:
|
||||
base_retriever: 基础检索器，执行最初的检索过程。
|
||||
post_processors: 后处理器列表，用于对检索到的结果进行重排、过滤等后处理，默认 None。
|
||||
pre_processors: 预处理器列表，用于对查询词进行改写、扩展等预处理，默认 None。
|
||||
""" # noqa: E501
|
||||
self.base_retriever = base_retriever
|
||||
self.post_processors = post_processors or []
|
||||
self.pre_processors = pre_processors or []
|
||||
|
||||
async def retrieve(
|
||||
self, query: Any, limit: int = 10, **kwargs: Any
|
||||
) -> list[SearchResult]:
|
||||
text_query = normalize_query_text(query)
|
||||
|
||||
queries_to_search = [query]
|
||||
|
||||
if text_query.strip():
|
||||
processed_texts = [text_query]
|
||||
for pp in self.pre_processors:
|
||||
new_texts = []
|
||||
for t in processed_texts:
|
||||
new_texts.extend(await pp.process(t))
|
||||
processed_texts = new_texts
|
||||
|
||||
if len(processed_texts) > 1 or (
|
||||
len(processed_texts) == 1 and processed_texts[0] != text_query
|
||||
):
|
||||
queries_to_search.extend(processed_texts)
|
||||
|
||||
all_results = []
|
||||
seen_ids = set()
|
||||
|
||||
for q in queries_to_search:
|
||||
res = await self.base_retriever.retrieve(q, limit=limit * 2, **kwargs)
|
||||
for r in res:
|
||||
if r.record.id not in seen_ids:
|
||||
seen_ids.add(r.record.id)
|
||||
all_results.append(r)
|
||||
|
||||
results = sorted(all_results, key=lambda x: x.score, reverse=True)
|
||||
|
||||
for pp in self.post_processors:
|
||||
results = await pp.process(results, query)
|
||||
|
||||
return results[:limit]
|
||||
|
||||
|
||||
class LifecyclePostProcessor(PostProcessor):
|
||||
"""生命周期后处理器（融合时间衰减与惰性访问强化）"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
half_life_days: int = 30,
|
||||
decay_weight: float = 0.3,
|
||||
semantic_weight: float = 0.7,
|
||||
importance_weight: float = 0.0,
|
||||
reinforcement_weight: float = 0.2,
|
||||
):
|
||||
"""
|
||||
初始化生命周期后处理器。
|
||||
|
||||
参数:
|
||||
half_life_days: 记忆衰减半衰期天数，控制信息随时间的降权速度，默认 30。
|
||||
decay_weight: 时间衰减得分的权重，默认 0.3。
|
||||
semantic_weight: 语义相关度得分的权重，默认 0.7。
|
||||
importance_weight: 信息重要性得分的权重，默认 0.0。
|
||||
reinforcement_weight: 惰性访问强化（如访问次数得分）的权重，默认 0.2。
|
||||
"""
|
||||
self.half_life_days = half_life_days
|
||||
self.decay_weight = decay_weight
|
||||
self.semantic_weight = semantic_weight
|
||||
self.importance_weight = importance_weight
|
||||
self.reinforcement_weight = reinforcement_weight
|
||||
|
||||
async def process(
|
||||
self, results: list[SearchResult], query: str
|
||||
) -> list[SearchResult]:
|
||||
now = time.time()
|
||||
import math
|
||||
|
||||
for res in results:
|
||||
created_at = res.record.metadata.get("created_at", now)
|
||||
importance = res.record.metadata.get("importance", 0.5)
|
||||
access_count = res.record.metadata.get("access_count", 0)
|
||||
last_accessed_at = res.record.metadata.get("last_accessed_at", created_at)
|
||||
age_days = max(0.0, (now - last_accessed_at) / 86400.0)
|
||||
decay = 0.5 ** (age_days / self.half_life_days)
|
||||
access_score = min(1.0, math.log1p(access_count) / 5.0)
|
||||
res.score = (
|
||||
(self.semantic_weight * res.score)
|
||||
+ (self.decay_weight * decay)
|
||||
+ (self.importance_weight * importance)
|
||||
+ (self.reinforcement_weight * access_score)
|
||||
)
|
||||
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
return results
|
||||
|
||||
|
||||
class HybridRetriever(BaseRetriever):
|
||||
"""
|
||||
双轨混合检索器 (Hybrid Search Engine)。
|
||||
并发调用 Dense (VectorDB) 和 Sparse (BM25)，并使用倒数秩融合 (RRF) 算法合并结果。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dense_retriever: BaseRetriever,
|
||||
sparse_retriever: BaseRetriever,
|
||||
dense_weight: float = 0.7,
|
||||
sparse_weight: float = 0.3,
|
||||
rrf_k: int = 60,
|
||||
oversample_factor: int = 2,
|
||||
min_oversample: int = 20,
|
||||
):
|
||||
"""
|
||||
初始化双轨混合检索器。
|
||||
|
||||
参数:
|
||||
dense_retriever: 稠密向量检索器，用于语义召回。
|
||||
sparse_retriever: 稀疏文本检索器，用于关键词召回（如 BM25）。
|
||||
dense_weight: 稠密向量检索的加权权重，默认 0.7。
|
||||
sparse_weight: 稀疏文本检索的加权权重，默认 0.3。
|
||||
rrf_k: 倒数秩融合(RRF)算法中的常数参数，默认 60。
|
||||
"""
|
||||
self.dense_retriever = dense_retriever
|
||||
self.sparse_retriever = sparse_retriever
|
||||
self.dense_weight = dense_weight
|
||||
self.sparse_weight = sparse_weight
|
||||
self.rrf_k = rrf_k
|
||||
self.oversample_factor = oversample_factor
|
||||
self.min_oversample = min_oversample
|
||||
|
||||
async def retrieve(
|
||||
self, query: Any, limit: int = 10, **kwargs: Any
|
||||
) -> list[SearchResult]:
|
||||
oversample_limit = max(limit * self.oversample_factor, self.min_oversample)
|
||||
|
||||
results = await asyncio.gather(
|
||||
self.dense_retriever.retrieve(query, limit=oversample_limit, **kwargs),
|
||||
self.sparse_retriever.retrieve(query, limit=oversample_limit, **kwargs),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
for res in results:
|
||||
if isinstance(res, ImportError):
|
||||
raise res
|
||||
|
||||
dense_res = (
|
||||
cast(list[SearchResult], results[0])
|
||||
if not isinstance(results[0], BaseException)
|
||||
else []
|
||||
)
|
||||
sparse_res = (
|
||||
cast(list[SearchResult], results[1])
|
||||
if not isinstance(results[1], BaseException)
|
||||
else []
|
||||
)
|
||||
|
||||
if isinstance(results[0], BaseException):
|
||||
logger.error(f"[HybridSearch] 向量检索异常: {results[0]}")
|
||||
if isinstance(results[1], BaseException):
|
||||
logger.error(f"[HybridSearch] BM25 检索异常: {results[1]}")
|
||||
|
||||
rrf_scores: dict[str, float] = {}
|
||||
merged_records = {}
|
||||
|
||||
for rank, res in enumerate(dense_res):
|
||||
record_id = res.record.id
|
||||
merged_records[record_id] = res.record
|
||||
rrf_score = 1.0 / (self.rrf_k + rank + 1)
|
||||
rrf_scores[record_id] = rrf_scores.get(record_id, 0.0) + (
|
||||
self.dense_weight * rrf_score
|
||||
)
|
||||
|
||||
for rank, res in enumerate(sparse_res):
|
||||
record_id = res.record.id
|
||||
merged_records[record_id] = res.record
|
||||
rrf_score = 1.0 / (self.rrf_k + rank + 1)
|
||||
rrf_scores[record_id] = rrf_scores.get(record_id, 0.0) + (
|
||||
self.sparse_weight * rrf_score
|
||||
)
|
||||
|
||||
max_possible_score = (self.dense_weight * (1.0 / (self.rrf_k + 1))) + (
|
||||
self.sparse_weight * (1.0 / (self.rrf_k + 1))
|
||||
)
|
||||
|
||||
final_results = []
|
||||
for record_id, score in sorted(
|
||||
rrf_scores.items(), key=lambda x: x[1], reverse=True
|
||||
):
|
||||
normalized_score = (
|
||||
score / max_possible_score if max_possible_score > 0 else 0.0
|
||||
)
|
||||
final_results.append(
|
||||
SearchResult(record=merged_records[record_id], score=normalized_score)
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"⚖️ [HybridSearch] 融合完成: "
|
||||
f"Dense({len(dense_res)}) + Sparse({len(sparse_res)}) "
|
||||
f"-> Merged({len(final_results)}), 截取 Top {limit}"
|
||||
)
|
||||
return final_results[:limit]
|
||||
@@ -0,0 +1,12 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
|
||||
"""使用 numpy 计算两组向量的余弦相似度"""
|
||||
if not vec1 or not vec2 or len(vec1) != len(vec2):
|
||||
return 0.0
|
||||
v1, v2 = np.array(vec1), np.array(vec2)
|
||||
norm1, norm2 = np.linalg.norm(v1), np.linalg.norm(v2)
|
||||
if norm1 == 0 or norm2 == 0:
|
||||
return 0.0
|
||||
return float(np.dot(v1, v2) / (norm1 * norm2))
|
||||
Reference in New Issue
Block a user