mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-12 15:30:04 +08:00
✨ feat!(llm): 重构并升级大语言模型服务为全新 AI 智能体框架 (#2146)
* ✨ feat!(llm): 重构并升级大语言模型服务为全新 AI 智能体框架 - 【重构】将原 services/llm 重构并迁移至全新的 services/ai 架构，提供向下兼容垫片 - 【新增】引入 Agent、Team、Workflow 三大智能体与工作流编排范式 - 【新增】引入基于 RAG 的长期向量记忆与中期槽位记忆系统 - 【新增】引入基于 Docker 的安全代码执行沙箱环境 - 【新增】支持 MCP 协议，允许动态管理和调用 MCP 服务 - 【新增】引入输入输出安全合规护栏与自愈反思机制 - 【优化】重构并优化多厂商 API 适配器 (Gemini, OpenAI, DeepSeek, GLM 等) - 【优化】优化日志脱敏与 Token 预估机制 - 【移除】移除旧版 llm default 和 llm reset-key 命令，新增 llm mcp 管理命令 * 🔧 chore(deps): 更新项目依赖与配置 - 添加 mcp、jieba 和 aiodocker 依赖到配置文件及 requirements.txt - 在 pyright 配置中设置 reportMissingImports 为 none - 调整 .gitignore 中 resources 目录的忽略规则 * ♻️ refactor(tools): 重构工具终止机制并清理知识库日志输出 - 统一使用 `context.state["__end_run__"]` 替代 `EndRunResult` 控制任务结束 - 移除文件系统和向量知识库检索工具中 `ToolResult` 的 `.with_log` 调用 - 调整指令处理器（Directive）的返回值为 `tool_res.output` - 修复部分类型检查警告并优化联合类型判断语法 * ♻️ refactor(tools): 重构工具副作用指令与控制流熔断机制 - 引入 `DirectivePayload` 及 `ToolResult` 的子类以结构化表达工具副作用 - 移除通过 `context.state` 传递魔术变量的隐式控制流设计 - 重构 `DirectiveManager` 处理器接口，直接在处理器中修改 `AgentState` 并构建 `AgentRunResult` - 在 `StandardAgentExecutor` 中统一通过 `directive_manager` 调度工具返回的副作用指令 - 补全 `MessageBuilder` 中部分核心方法的文档注释 * 🐛 fix(sandbox): 修复 Docker 沙箱容器状态检测与会话清理逻辑 -【修复】修正 `is_alive` 中直接读取私有属性的问题，改用 `show()` 返回值 -【修复】解决 `execute_code` 中缓存的执行器与当前会话不一致的问题 -【优化】在清理工作区前增加容器存活检测，避免向已死容器发送请求 -【优化】创建容器时增加运行状态校验，若已停止则自动从缓存中移除并重建 -【优化】优化容器销毁和清理逻辑，静默处理容器不存在 (404) 的异常 * 📝 docs(core): 补充核心模块初始化方法的文档注释 * 🚨 auto fix by pre-commit hooks --------- Co-authored-by: webjoin111 <455457521@qq.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Rumio
co-authored by
webjoin111
pre-commit-ci[bot]
webjoin111
pre-commit-ci[bot]
parent
bdc1374848
commit
80fc5b86a7
@@ -0,0 +1,20 @@
|
||||
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
|
||||
Reference in New Issue
Block a user