mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-13 16:10:05 +08:00
♻️ refactor(core): 重构 AI 编排框架与记忆及 RAG 子系统 (#2149)
* ♻️ 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>
This commit is contained in:
co-authored by
webjoin111
pre-commit-ci[bot]
parent
922d092650
commit
52f7dbdedf
@@ -4,6 +4,7 @@ import threading
|
||||
from typing import Any, Literal, Protocol, runtime_checkable
|
||||
|
||||
from zhenxun.services.ai.core.messages import EmbedBatch
|
||||
from zhenxun.services.ai.core.options import LLMEmbeddingConfig
|
||||
from zhenxun.services.ai.llm.api import embed as api_embed
|
||||
from zhenxun.services.ai.message_builder import MessageBuilder
|
||||
from zhenxun.services.ai.utils.logger import log_rag as logger
|
||||
@@ -16,7 +17,7 @@ EmbedTaskType = Literal[
|
||||
@runtime_checkable
|
||||
class Embedder(Protocol):
|
||||
"""
|
||||
向量化引擎协议 (Callable Protocol)。
|
||||
向量化引擎协议。
|
||||
任何实现了异步 __call__ 的对象或闭包函数均可作为 Embedder。
|
||||
"""
|
||||
|
||||
@@ -32,7 +33,9 @@ class Embedder(Protocol):
|
||||
class DefaultEmbedder(Embedder):
|
||||
"""系统默认的向量化引擎,调用大模型底座 API"""
|
||||
|
||||
def __init__(self, model_name: str | None = None, config: Any = None):
|
||||
def __init__(
|
||||
self, model_name: str | None = None, config: LLMEmbeddingConfig | None = None
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.config = config
|
||||
|
||||
@@ -57,10 +60,28 @@ class BaseLocalEmbedder(Embedder, ABC):
|
||||
def __init__(self, model_name: str):
|
||||
self.model_name = model_name
|
||||
self._model_lock = threading.Lock()
|
||||
self._model: Any | None = None
|
||||
|
||||
def _ensure_model_loaded(self) -> None:
|
||||
"""线程安全的懒加载机制"""
|
||||
if self._model is None:
|
||||
with self._model_lock:
|
||||
if self._model is None:
|
||||
logger.info(
|
||||
f"正在后台加载本地向量模型: {self.model_name} ... "
|
||||
"(首次加载可能需要较长时间下载)"
|
||||
)
|
||||
self._model = self._load_model_impl()
|
||||
logger.info(f"本地向量模型 {self.model_name} 加载完毕!")
|
||||
|
||||
@abstractmethod
|
||||
def _load_model_impl(self) -> Any:
|
||||
"""子类实现:执行具体的依赖导入与模型实例化,并返回模型对象。"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _encode_texts(self, texts: list[str]) -> list[list[float]]:
|
||||
"""子类只需实现此同步的批量文本向量化方法即可。"""
|
||||
"""子类实现:执行同步的批量文本向量化方法。"""
|
||||
pass
|
||||
|
||||
async def __call__(
|
||||
@@ -93,7 +114,6 @@ class FastEmbedder(BaseLocalEmbedder):
|
||||
|
||||
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
|
||||
|
||||
@@ -102,29 +122,19 @@ class FastEmbedder(BaseLocalEmbedder):
|
||||
"⚠️ 使用 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 _load_model_impl(self) -> Any:
|
||||
try:
|
||||
from fastembed import TextEmbedding
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"⚠️ 使用 FastEmbed 需要额外依赖,请在终端执行: pip install fastembed"
|
||||
)
|
||||
return TextEmbedding(model_name=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)]
|
||||
assert self._model is not None
|
||||
return [vec.tolist() for vec in self._model.embed(texts)]
|
||||
|
||||
|
||||
class SentenceTransformerEmbedder(BaseLocalEmbedder):
|
||||
@@ -135,7 +145,6 @@ class SentenceTransformerEmbedder(BaseLocalEmbedder):
|
||||
|
||||
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
|
||||
|
||||
@@ -145,32 +154,18 @@ class SentenceTransformerEmbedder(BaseLocalEmbedder):
|
||||
"请在终端执行: 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 _load_model_impl(self) -> Any:
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"⚠️ 使用 SentenceTransformers 需要额外依赖,"
|
||||
"请在终端执行: pip install sentence-transformers"
|
||||
)
|
||||
return 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)
|
||||
assert self._model is not None
|
||||
embeddings = self._model.encode(texts)
|
||||
return embeddings.tolist()
|
||||
|
||||
@@ -11,6 +11,10 @@ from zhenxun.services.ai.context.rag.models import (
|
||||
SearchResult,
|
||||
)
|
||||
from zhenxun.services.ai.context.rag.retrieval import FilterEvaluator
|
||||
from zhenxun.services.ai.context.rag.utils import (
|
||||
InMemoryScorer,
|
||||
normalize_vector,
|
||||
)
|
||||
from zhenxun.services.ai.utils.logger import log_rag as logger
|
||||
from zhenxun.services.ai.utils.scope import ScopeSelector
|
||||
from zhenxun.services.db_context import Model
|
||||
@@ -24,9 +28,7 @@ class StorageBackend(Protocol):
|
||||
"""保存或更新数据块"""
|
||||
...
|
||||
|
||||
async def search(
|
||||
self, query: QueryRequest, scopes: list[str] | None = None
|
||||
) -> list[SearchResult]:
|
||||
async def search(self, request: QueryRequest) -> list[SearchResult]:
|
||||
"""按向量和前缀检索数据块"""
|
||||
...
|
||||
|
||||
@@ -49,22 +51,6 @@ class StorageBackend(Protocol):
|
||||
...
|
||||
|
||||
|
||||
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 存储实现"""
|
||||
|
||||
@@ -83,62 +69,36 @@ class DictStorageBackend(StorageBackend):
|
||||
else:
|
||||
self._vectors.pop(r.id, None)
|
||||
|
||||
async def search(
|
||||
self, query: QueryRequest, scopes: list[str] | None = None
|
||||
) -> list[SearchResult]:
|
||||
async def search(self, request: QueryRequest) -> list[SearchResult]:
|
||||
candidate_ids = []
|
||||
for record in self._records.values():
|
||||
if scopes is not None:
|
||||
if record.metadata.get("scope", "/") not in scopes:
|
||||
if request.scopes is not None:
|
||||
if record.metadata.get("scope", "/") not in request.scopes:
|
||||
continue
|
||||
if not FilterEvaluator.evaluate(record.metadata, query.metadata_filters):
|
||||
if not FilterEvaluator.evaluate(record.metadata, request.metadata_filters):
|
||||
continue
|
||||
if not query.embedding and query.text and query.text not in record.content:
|
||||
if (
|
||||
not request.embedding
|
||||
and request.text
|
||||
and request.text not in record.content
|
||||
):
|
||||
continue
|
||||
candidate_ids.append(record.id)
|
||||
|
||||
if not candidate_ids:
|
||||
return []
|
||||
|
||||
results = []
|
||||
if query.search_type == "sparse":
|
||||
import jieba
|
||||
records = [self._records[r_id] for r_id in candidate_ids]
|
||||
|
||||
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:
|
||||
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))
|
||||
if request.search_type == "sparse":
|
||||
results = InMemoryScorer.calculate_sparse_scores(request.text, records)
|
||||
elif request.search_type == "dense" and request.embedding:
|
||||
results = InMemoryScorer.calculate_dense_scores(request.embedding, records)
|
||||
else:
|
||||
for r_id in candidate_ids:
|
||||
results.append(SearchResult(record=self._records[r_id], score=0.1))
|
||||
results = [SearchResult(record=r, score=0.1) for r in records]
|
||||
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
return results[: query.limit]
|
||||
return results[: request.limit]
|
||||
|
||||
async def update(self, record: BaseRecord) -> None:
|
||||
if record.id in self._records:
|
||||
@@ -212,94 +172,48 @@ class TortoiseStorageBackend(StorageBackend):
|
||||
},
|
||||
)
|
||||
|
||||
async def search(
|
||||
self, query: QueryRequest, scopes: list[str] | None = None
|
||||
) -> list[SearchResult]:
|
||||
async def search(self, request: QueryRequest) -> list[SearchResult]:
|
||||
query_orm = self.model_class.all()
|
||||
if scopes is not None:
|
||||
query_orm = query_orm.filter(scope__in=scopes)
|
||||
if request.scopes is not None:
|
||||
query_orm = query_orm.filter(scope__in=request.scopes)
|
||||
|
||||
if query.search_type == "sparse" and query.text:
|
||||
if request.search_type == "sparse" and request.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]
|
||||
t for t in jieba.lcut_for_search(request.text) if len(t.strip()) > 1
|
||||
] or [request.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)
|
||||
elif request.search_type == "dense" and not request.embedding and request.text:
|
||||
query_orm = query_orm.filter(content__icontains=request.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):
|
||||
if not FilterEvaluator.evaluate(row_meta, request.metadata_filters):
|
||||
continue
|
||||
valid_rows.append(row)
|
||||
|
||||
if not valid_rows:
|
||||
return []
|
||||
|
||||
results = []
|
||||
if query.search_type == "sparse":
|
||||
import jieba
|
||||
records = [self._to_base_record(row) for row in valid_rows]
|
||||
|
||||
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:
|
||||
logger.warning(
|
||||
"⚠️ 数据库中缓存的向量维度与当前模型查询维度不匹配,"
|
||||
f"已安全跳过向量检索(降级为稀疏匹配)。原因: {e}"
|
||||
)
|
||||
|
||||
for row in missing_rows:
|
||||
results.append(
|
||||
SearchResult(record=self._to_base_record(row), score=0.1)
|
||||
)
|
||||
if request.search_type == "sparse":
|
||||
results = InMemoryScorer.calculate_sparse_scores(request.text, records)
|
||||
elif request.search_type == "dense" and request.embedding:
|
||||
results = InMemoryScorer.calculate_dense_scores(request.embedding, records)
|
||||
else:
|
||||
for row in valid_rows:
|
||||
results.append(
|
||||
SearchResult(record=self._to_base_record(row), score=0.1)
|
||||
)
|
||||
results = [SearchResult(record=r, score=0.1) for r in records]
|
||||
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
return results[: query.limit]
|
||||
return results[: request.limit]
|
||||
|
||||
async def update(self, record: BaseRecord) -> None:
|
||||
await self.model_class.filter(id=record.id).update(
|
||||
@@ -386,50 +300,50 @@ class QdrantStorageBackend(StorageBackend):
|
||||
)
|
||||
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:
|
||||
async def search(self, request: QueryRequest) -> list[SearchResult]:
|
||||
if request.search_type == "dense" and not request.embedding:
|
||||
return []
|
||||
if query.embedding:
|
||||
await self._ensure_collection(len(query.embedding))
|
||||
if request.embedding:
|
||||
await self._ensure_collection(len(request.embedding))
|
||||
|
||||
from qdrant_client.models import FieldCondition, Filter, MatchText, MatchValue
|
||||
|
||||
must_conditions = []
|
||||
|
||||
if scopes is not None:
|
||||
if request.scopes is not None:
|
||||
try:
|
||||
from qdrant_client.models import MatchAny
|
||||
|
||||
must_conditions.append(
|
||||
FieldCondition(key="metadata.scope", match=MatchAny(any=scopes))
|
||||
FieldCondition(
|
||||
key="metadata.scope", match=MatchAny(any=request.scopes)
|
||||
)
|
||||
)
|
||||
except ImportError:
|
||||
scope_conditions = [
|
||||
FieldCondition(key="metadata.scope", match=MatchValue(value=s))
|
||||
for s in scopes
|
||||
for s in request.scopes
|
||||
]
|
||||
must_conditions.append(Filter(should=scope_conditions))
|
||||
|
||||
if query.metadata_filters:
|
||||
for k, v in query.metadata_filters.items():
|
||||
if request.metadata_filters:
|
||||
for k, v in request.metadata_filters.items():
|
||||
must_conditions.append(
|
||||
FieldCondition(key=f"metadata.{k}", match=MatchValue(value=v))
|
||||
)
|
||||
|
||||
if query.search_type == "sparse":
|
||||
if request.search_type == "sparse":
|
||||
must_conditions.append(
|
||||
FieldCondition(key="content", match=MatchText(text=query.text))
|
||||
FieldCondition(key="content", match=MatchText(text=request.text))
|
||||
)
|
||||
|
||||
query_filter = Filter(must=must_conditions) if must_conditions else None
|
||||
|
||||
if query.search_type == "sparse":
|
||||
if request.search_type == "sparse":
|
||||
results = await self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=query_filter,
|
||||
limit=query.limit,
|
||||
limit=request.limit,
|
||||
with_payload=True,
|
||||
)
|
||||
return [
|
||||
@@ -446,8 +360,8 @@ class QdrantStorageBackend(StorageBackend):
|
||||
|
||||
results = await self.client.search( # type: ignore
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query.embedding,
|
||||
limit=query.limit,
|
||||
query_vector=request.embedding,
|
||||
limit=request.limit,
|
||||
query_filter=query_filter,
|
||||
)
|
||||
|
||||
@@ -559,27 +473,25 @@ class LanceDBStorageBackend(StorageBackend):
|
||||
else:
|
||||
self.db.open_table(self.table_name).add(data)
|
||||
|
||||
async def search(
|
||||
self, query: QueryRequest, scopes: list[str] | None = None
|
||||
) -> list[SearchResult]:
|
||||
async def search(self, request: QueryRequest) -> list[SearchResult]:
|
||||
if self.table_name not in self.db.table_names():
|
||||
return []
|
||||
if query.search_type == "dense" and not query.embedding:
|
||||
if request.search_type == "dense" and not request.embedding:
|
||||
return []
|
||||
|
||||
tbl = self.db.open_table(self.table_name)
|
||||
if query.search_type == "sparse":
|
||||
if request.search_type == "sparse":
|
||||
try:
|
||||
results = (
|
||||
tbl.search(query.text, query_type="fts")
|
||||
.limit(query.limit)
|
||||
tbl.search(request.text, query_type="fts")
|
||||
.limit(request.limit)
|
||||
.to_list()
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"LanceDB FTS 检索失败(可能是由于尚未创建FTS索引): {e}")
|
||||
return []
|
||||
else:
|
||||
results = tbl.search(query.embedding).limit(query.limit).to_list()
|
||||
results = tbl.search(request.embedding).limit(request.limit).to_list()
|
||||
|
||||
import ast
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@ from typing import Any
|
||||
|
||||
from zhenxun.services.ai.utils.logger import log_rag as logger
|
||||
|
||||
from .backends import StorageBackend
|
||||
from .backends import Embedder, StorageBackend
|
||||
from .configs import RAGConfig
|
||||
from .engine import ScopedRAGClient
|
||||
from .ingestion import (
|
||||
@@ -29,18 +29,25 @@ from .retrieval import (
|
||||
|
||||
class RAGBuilder:
|
||||
"""
|
||||
RAG 管线组装构建器 (Fluent Builder Pattern)。
|
||||
使用内部状态驱动模式 (The Memory Pattern),内部维护私有的 RAGConfig 实例。
|
||||
RAG 管线组装构建器。
|
||||
使用内部状态驱动模式,内部维护私有的 RAGConfig 实例。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, storage: StorageBackend | None = None, config: RAGConfig | None = None
|
||||
):
|
||||
"""
|
||||
初始化 RAG 管线组装构建器。
|
||||
|
||||
参数:
|
||||
storage: (可选) 底层存储后端实例。
|
||||
config: (可选) RAG 配置对象。若不提供,将新建默认的 RAGConfig 实例。
|
||||
"""
|
||||
self._config = config or RAGConfig()
|
||||
if storage is not None:
|
||||
self._config.storage = storage
|
||||
|
||||
def with_embedder(self, embedder: Any) -> "RAGBuilder":
|
||||
def with_embedder(self, embedder: Embedder) -> "RAGBuilder":
|
||||
"""
|
||||
设置向量化引擎。
|
||||
|
||||
|
||||
@@ -11,11 +11,13 @@ from .ingestion import (
|
||||
)
|
||||
from .models import (
|
||||
BaseRecord,
|
||||
QueryRequest,
|
||||
SearchResult,
|
||||
)
|
||||
from .retrieval import (
|
||||
BaseRetriever,
|
||||
)
|
||||
from .utils import normalize_query_text
|
||||
|
||||
|
||||
class ScopedRAGClient:
|
||||
@@ -55,7 +57,15 @@ class ScopedRAGClient:
|
||||
self.scope_prefix = self.scopes[0] if self.scopes else "/"
|
||||
|
||||
async def ingest(self, records: list[BaseRecord]) -> int:
|
||||
"""通过 RAG Ingestion Pipeline 处理并入库数据"""
|
||||
"""
|
||||
通过 RAG Ingestion Pipeline 处理并入库数据。
|
||||
|
||||
参数:
|
||||
records: 待导入的数据记录列表。
|
||||
|
||||
返回:
|
||||
int: 成功导入并存入的记录数量。
|
||||
"""
|
||||
for r in records:
|
||||
if "scope" not in r.metadata:
|
||||
r.metadata["scope"] = self.scope_prefix
|
||||
@@ -73,7 +83,16 @@ class ScopedRAGClient:
|
||||
"""
|
||||
多作用域联合切片视图检索 (Union Search)。
|
||||
并发向多个独立的作用域发起检索,并对结果进行合并、去重和重排。
|
||||
"""
|
||||
|
||||
参数:
|
||||
query: 检索的查询对象,可以是文本字符串、向量或高级查询对象。
|
||||
limit: 最大返回结果条数。
|
||||
scopes: (可选) 自定义检索的作用域。若不提供,则使用客户端初始化时设定的 scopes。
|
||||
**kwargs: 传递给底层检索器的其他关键字参数。
|
||||
|
||||
返回:
|
||||
list[SearchResult]: 检索到的相似度排序后的结果列表。
|
||||
""" # noqa: E501
|
||||
target_scopes = self.scopes
|
||||
if scopes is not None:
|
||||
target_scopes = (
|
||||
@@ -85,13 +104,40 @@ class ScopedRAGClient:
|
||||
if not target_scopes:
|
||||
return []
|
||||
|
||||
kwargs["scopes"] = target_scopes
|
||||
return await self.retriever.retrieve(query, limit=limit, **kwargs)
|
||||
text_query = normalize_query_text(query)
|
||||
metadata_filters = kwargs.pop("metadata_filters", None)
|
||||
|
||||
request = QueryRequest(
|
||||
text=text_query,
|
||||
limit=limit,
|
||||
scopes=target_scopes,
|
||||
metadata_filters=metadata_filters,
|
||||
extra=kwargs,
|
||||
)
|
||||
|
||||
return await self.retriever.retrieve(request)
|
||||
|
||||
async def update(self, record: BaseRecord) -> None:
|
||||
"""
|
||||
更新指定的已有数据记录。
|
||||
会自动强行附加当前客户端的单作用域前缀 (scope_prefix)。
|
||||
|
||||
参数:
|
||||
record: 待更新的完整数据记录对象。
|
||||
"""
|
||||
record.metadata["scope"] = self.scope_prefix
|
||||
await self.storage.update(record)
|
||||
|
||||
async def delete(self, record_ids: list[str] | None = None, **kwargs: Any) -> int:
|
||||
"""
|
||||
删除指定的已有数据记录,操作限定在当前客户端的单作用域前缀下。
|
||||
|
||||
参数:
|
||||
record_ids: (可选) 待删除的记录 ID 列表。
|
||||
**kwargs: 传递给底层存储后端的其他过滤或删除参数。
|
||||
|
||||
返回:
|
||||
int: 成功删除的记录条数。
|
||||
"""
|
||||
kwargs["scope_prefix"] = self.scope_prefix
|
||||
return await self.storage.delete(record_ids=record_ids, **kwargs)
|
||||
|
||||
@@ -42,6 +42,10 @@ class QueryRequest(BaseModel):
|
||||
"""元数据精确匹配字典"""
|
||||
limit: int = Field(default=10)
|
||||
"""返回的最大条数"""
|
||||
scopes: list[str] | None = Field(default=None)
|
||||
"""检索的数据隔离作用域列表"""
|
||||
extra: dict[str, Any] = Field(default_factory=dict)
|
||||
"""透传参数逃生舱"""
|
||||
|
||||
|
||||
StorageConfigType = dict[str, Any]
|
||||
|
||||
@@ -4,22 +4,15 @@ 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
|
||||
|
||||
|
||||
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):
|
||||
"""
|
||||
@@ -29,9 +22,7 @@ class BaseRetriever(Protocol):
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def retrieve(
|
||||
self, query: Any, limit: int = 10, **kwargs: Any
|
||||
) -> list[SearchResult]: ...
|
||||
async def retrieve(self, request: QueryRequest) -> list[SearchResult]: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
@@ -72,7 +63,7 @@ class VectorDBRetriever(BaseRetriever):
|
||||
def __init__(
|
||||
self,
|
||||
storage: "StorageBackend",
|
||||
embedder: Any,
|
||||
embedder: Embedder,
|
||||
scope_prefix: str | None = None,
|
||||
score_threshold: float = 0.4,
|
||||
):
|
||||
@@ -90,29 +81,25 @@ class VectorDBRetriever(BaseRetriever):
|
||||
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)
|
||||
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
|
||||
|
||||
vecs = await self.embedder(query, task="query")
|
||||
query_vec = vecs[0] if vecs else None
|
||||
|
||||
if not text_query.strip() and not query_vec:
|
||||
if not request.text.strip() and not request.embedding:
|
||||
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]
|
||||
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):
|
||||
@@ -136,24 +123,20 @@ class DatabaseSparseRetriever(BaseRetriever):
|
||||
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():
|
||||
async def retrieve(self, request: QueryRequest) -> list[SearchResult]:
|
||||
if not request.text.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]
|
||||
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):
|
||||
@@ -183,19 +166,17 @@ class RerankRetriever(BaseRetriever):
|
||||
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
|
||||
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 []
|
||||
|
||||
text_query = normalize_query_text(query)
|
||||
|
||||
docs: list[str | dict[str, str]] = [
|
||||
res.record.content for res in initial_results
|
||||
]
|
||||
@@ -204,14 +185,14 @@ class RerankRetriever(BaseRetriever):
|
||||
|
||||
try:
|
||||
reranked = await rerank(
|
||||
query=text_query,
|
||||
query=request.text,
|
||||
documents=docs,
|
||||
top_n=min(limit, self.top_n),
|
||||
top_n=min(request.limit, self.top_n),
|
||||
model=self.model_name,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Rerank 重排请求失败,将降级返回初筛结果: {e}")
|
||||
return initial_results[:limit]
|
||||
return initial_results[: request.limit]
|
||||
|
||||
final_results = []
|
||||
for rr in reranked:
|
||||
@@ -243,15 +224,11 @@ class PipelineRetriever(BaseRetriever):
|
||||
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)
|
||||
async def retrieve(self, request: QueryRequest) -> list[SearchResult]:
|
||||
requests_to_search = [request]
|
||||
|
||||
queries_to_search = [query]
|
||||
|
||||
if text_query.strip():
|
||||
processed_texts = [text_query]
|
||||
if request.text.strip():
|
||||
processed_texts = [request.text]
|
||||
for pp in self.pre_processors:
|
||||
new_texts = []
|
||||
for t in processed_texts:
|
||||
@@ -259,15 +236,23 @@ class PipelineRetriever(BaseRetriever):
|
||||
processed_texts = new_texts
|
||||
|
||||
if len(processed_texts) > 1 or (
|
||||
len(processed_texts) == 1 and processed_texts[0] != text_query
|
||||
len(processed_texts) == 1 and processed_texts[0] != request.text
|
||||
):
|
||||
queries_to_search.extend(processed_texts)
|
||||
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 q in queries_to_search:
|
||||
res = await self.base_retriever.retrieve(q, limit=limit * 2, **kwargs)
|
||||
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)
|
||||
@@ -276,9 +261,9 @@ class PipelineRetriever(BaseRetriever):
|
||||
results = sorted(all_results, key=lambda x: x.score, reverse=True)
|
||||
|
||||
for pp in self.post_processors:
|
||||
results = await pp.process(results, query)
|
||||
results = await pp.process(results, request.text)
|
||||
|
||||
return results[:limit]
|
||||
return results[: request.limit]
|
||||
|
||||
|
||||
class LifecyclePostProcessor(PostProcessor):
|
||||
@@ -367,14 +352,22 @@ class HybridRetriever(BaseRetriever):
|
||||
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)
|
||||
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(query, limit=oversample_limit, **kwargs),
|
||||
self.sparse_retriever.retrieve(query, limit=oversample_limit, **kwargs),
|
||||
self.dense_retriever.retrieve(dense_req),
|
||||
self.sparse_retriever.retrieve(sparse_req),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
@@ -435,6 +428,6 @@ class HybridRetriever(BaseRetriever):
|
||||
logger.debug(
|
||||
f"⚖️ [HybridSearch] 融合完成: "
|
||||
f"Dense({len(dense_res)}) + Sparse({len(sparse_res)}) "
|
||||
f"-> Merged({len(final_results)}), 截取 Top {limit}"
|
||||
f"-> Merged({len(final_results)}), 截取 Top {request.limit}"
|
||||
)
|
||||
return final_results[:limit]
|
||||
return final_results[: request.limit]
|
||||
|
||||
@@ -1,5 +1,10 @@
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from zhenxun.services.ai.context.rag.models import BaseRecord, SearchResult
|
||||
from zhenxun.services.ai.utils.logger import log_rag as logger
|
||||
|
||||
|
||||
def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
|
||||
"""使用 numpy 计算两组向量的余弦相似度"""
|
||||
@@ -10,3 +15,79 @@ def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
|
||||
if norm1 == 0 or norm2 == 0:
|
||||
return 0.0
|
||||
return float(np.dot(v1, v2) / (norm1 * norm2))
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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 ""
|
||||
|
||||
|
||||
class InMemoryScorer:
|
||||
"""提供内存级别的纯 Python 向量打分与 BM25 稀疏打分工具类"""
|
||||
|
||||
@staticmethod
|
||||
def calculate_sparse_scores(
|
||||
query: str, candidate_records: list[BaseRecord]
|
||||
) -> list[SearchResult]:
|
||||
import jieba
|
||||
|
||||
tokens = set(jieba.lcut_for_search(query.lower()))
|
||||
if not tokens:
|
||||
return [SearchResult(record=r, score=0.1) for r in candidate_records]
|
||||
|
||||
results = []
|
||||
for r in candidate_records:
|
||||
content_lower = r.content.lower()
|
||||
matched_count = sum(1 for t in tokens if t in content_lower)
|
||||
score = matched_count / len(tokens) if tokens else 0.1
|
||||
results.append(SearchResult(record=r, score=score))
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def calculate_dense_scores(
|
||||
query_embedding: list[float], candidate_records: list[BaseRecord]
|
||||
) -> list[SearchResult]:
|
||||
if not candidate_records or not query_embedding:
|
||||
return []
|
||||
|
||||
q_vec = normalize_vector(query_embedding)
|
||||
vec_records = [r for r in candidate_records if r.embedding]
|
||||
missing_records = [r for r in candidate_records if not r.embedding]
|
||||
|
||||
results = []
|
||||
if vec_records:
|
||||
try:
|
||||
raw_mat = np.array([r.embedding for r in vec_records], dtype=np.float32)
|
||||
norm_mat = normalize_matrix(raw_mat)
|
||||
scores = norm_mat @ q_vec
|
||||
for r, score in zip(vec_records, scores):
|
||||
results.append(SearchResult(record=r, score=float(score)))
|
||||
except ValueError as e:
|
||||
logger.warning(
|
||||
"⚠️ 维度不匹配,已安全跳过向量检索(降级)。原因: " + str(e)
|
||||
)
|
||||
|
||||
for r in missing_records:
|
||||
results.append(SearchResult(record=r, score=0.1))
|
||||
|
||||
return results
|
||||
|
||||
Reference in New Issue
Block a user