mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-09-29 00:32:06 +08:00
52f7dbdedf
* ♻️ refactor(core): 重构 AI 编排框架与记忆及 RAG 子系统 - 【重构】重构 `BaseRunnable` 并引入统一的 `RunIntent` 意图载体，规范 Agent、Team 和 Workflow 的执行流 - 【解耦】将中期记忆槽和长期向量记忆从 `MemoryConfig` 中解耦，转为独立的能力组件与工具箱进行管理 - 【记忆】移除 `MemoryReader` 和 `MemoryWriter`，统一封装为 `SessionMemoryContext` 会话记忆门面 - 【RAG】重构检索器与存储后端接口，统一采用 `QueryRequest` 进行多维度联合检索，并引入 `InMemoryScorer` 提升打分性能 - 【事件】优化 `EventBus` 异步事件分发机制，引入队列机制确保事件按序处理，避免并发竞态问题 - 【依赖注入】移除 `memory` 注入项，优化 `DependencyInjector` 的签名解析缓存以提升性能 * 🚨 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>
434 lines
15 KiB
Python
434 lines
15 KiB
Python
from abc import abstractmethod
|
|
import asyncio
|
|
import time
|
|
from typing import TYPE_CHECKING, Any, Protocol, cast, runtime_checkable
|
|
|
|
from zhenxun.services.ai.utils.logger import log_rag as logger
|
|
from zhenxun.utils.pydantic_compat import model_copy
|
|
|
|
from .backends.embedders import Embedder
|
|
from .models import QueryRequest, SearchResult
|
|
|
|
if TYPE_CHECKING:
|
|
from .backends.storages import StorageBackend
|
|
|
|
|
|
@runtime_checkable
|
|
class BaseRetriever(Protocol):
|
|
"""
|
|
检索器核心协议。
|
|
任何实现了 retrieve 方法的对象均可作为检索器
|
|
（不仅限于向量检索，也可包含 BM25、SQL 搜索等）。
|
|
"""
|
|
|
|
@abstractmethod
|
|
async def retrieve(self, request: QueryRequest) -> 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: Embedder,
|
|
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, request: QueryRequest) -> list[SearchResult]:
|
|
if not request.embedding and request.text:
|
|
vecs = await self.embedder(request.text, task="query")
|
|
request.embedding = vecs[0] if vecs else None
|
|
|
|
if not request.text.strip() and not request.embedding:
|
|
return []
|
|
|
|
request.search_type = "dense"
|
|
if not request.scopes and self.scope_prefix:
|
|
request.scopes = [self.scope_prefix]
|
|
|
|
original_limit = request.limit
|
|
request.limit = original_limit * 2
|
|
|
|
results = await self.storage.search(request)
|
|
|
|
request.limit = original_limit
|
|
return [r for r in results if r.score >= self.score_threshold][: request.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, request: QueryRequest) -> list[SearchResult]:
|
|
if not request.text.strip():
|
|
return []
|
|
|
|
request.search_type = "sparse"
|
|
if not request.scopes and self.scope_prefix:
|
|
request.scopes = [self.scope_prefix]
|
|
|
|
original_limit = request.limit
|
|
request.limit = original_limit * 2
|
|
|
|
results = await self.storage.search(request)
|
|
request.limit = original_limit
|
|
return [r for r in results if r.score > self.score_threshold][: request.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, request: QueryRequest) -> list[SearchResult]:
|
|
req_clone = model_copy(request, deep=True)
|
|
req_clone.limit = max(
|
|
request.limit * self.oversample_factor, self.min_oversample
|
|
)
|
|
|
|
initial_results = await self.base_retriever.retrieve(req_clone)
|
|
|
|
if not initial_results:
|
|
return []
|
|
|
|
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=request.text,
|
|
documents=docs,
|
|
top_n=min(request.limit, self.top_n),
|
|
model=self.model_name,
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Rerank 重排请求失败，将降级返回初筛结果: {e}")
|
|
return initial_results[: request.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, request: QueryRequest) -> list[SearchResult]:
|
|
requests_to_search = [request]
|
|
|
|
if request.text.strip():
|
|
processed_texts = [request.text]
|
|
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] != request.text
|
|
):
|
|
requests_to_search = []
|
|
for pt in processed_texts:
|
|
new_req = model_copy(request, deep=True)
|
|
new_req.text = pt
|
|
requests_to_search.append(new_req)
|
|
|
|
all_results = []
|
|
seen_ids = set()
|
|
|
|
for req in requests_to_search:
|
|
original_limit = req.limit
|
|
req.limit = original_limit * 2
|
|
res = await self.base_retriever.retrieve(req)
|
|
req.limit = original_limit
|
|
|
|
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, request.text)
|
|
|
|
return results[: request.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, request: QueryRequest) -> list[SearchResult]:
|
|
oversample_limit = max(
|
|
request.limit * self.oversample_factor, self.min_oversample
|
|
)
|
|
|
|
dense_req = model_copy(request, deep=True)
|
|
dense_req.limit = oversample_limit
|
|
dense_req.search_type = "dense"
|
|
|
|
sparse_req = model_copy(request, deep=True)
|
|
sparse_req.limit = oversample_limit
|
|
sparse_req.search_type = "sparse"
|
|
|
|
results = await asyncio.gather(
|
|
self.dense_retriever.retrieve(dense_req),
|
|
self.sparse_retriever.retrieve(sparse_req),
|
|
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 {request.limit}"
|
|
)
|
|
return final_results[: request.limit]
|