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,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