sdpy-kit 0.1.0__py3-none-any.whl

kit/zvec_kit.py

@@ -0,0 +1,671 @@
+"""
+Copyright (c) 2021-2026 Clark Chang. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Project: sdpy
+Author: Clark Chang
+"""
+
+# date: 2026-08-05
+"""zvec 向量数据库工具类
+
+基于 zvec（本地嵌入式向量库）封装，提供开箱即用的向量数据服务方法。
+采用独立工具范式（类似 RedisKit），不依赖 kit.db 的 SQLAlchemy 体系。
+
+zvec 核心概念：
+- Collection：类比数据库表，数据落盘到指定目录，schema-first（需先定义结构）
+- Doc：类比一行记录，包含标量字段(fields) + 向量字段(vectors) + 唯一id
+
+两类使用模式：
+
+1. 单条操作（自动管理 Collection 句柄）::
+
+    # 业务方先定义 schema 并创建 collection
+    schema = ZvecKit.make_text_embedding_schema("docs", dim=768)
+    ZvecKit.create_collection("docs", schema)
+
+    # CRUD
+    ZvecKit.insert(
+        "docs",
+        {"id": "d1", "vectors": {"embedding": [0.1] * 768}, "fields": {"year": 2024}},
+    )
+    doc = ZvecKit.fetch("docs", "d1")
+    ZvecKit.update("docs", {"id": "d1", "fields": {"year": 2025}})
+    ZvecKit.delete("docs", "d1")
+
+2. 向量检索（阶段2补充）::
+
+    results = ZvecKit.vector_query("docs", field="embedding", vector=[0.3]*768, topk=10)
+
+注意：
+- zvec 的 zvec.init() 在进程内应只调用一次，本类已做幂等守卫。
+- 删除/更新后检索召回率受 index_completeness 影响，大批量写入后建议调用 optimize()。
+"""
+
+import json
+import os
+import threading
+from typing import Any
+
+import zvec
+from zvec import (
+    Collection,
+    CollectionSchema,
+    DataType,
+    Doc,
+    FieldSchema,
+    HnswIndexParam,
+    InvertIndexParam,
+    MetricType,
+    Query,
+    VectorSchema,
+)
+
+from kit.config.kit_config import get_settings
+from kit.config.log_config import get_default_logger
+
+logger = get_default_logger(__name__)
+
+
+class ZvecError(Exception):
+    """zvec 向量库操作异常基类
+
+    独立于 kit.db.exceptions.DBError（zvec 不走 SQLAlchemy 体系）。
+    """
+
+
+# ==================== 全局单例守卫 ====================
+
+_zvec_initialized: bool = False
+_init_lock = threading.Lock()
+
+# Collection 实例注册表：path -> Collection（线程安全单例，避免重复 open 同一目录）
+_collections: dict[str, Collection] = {}
+_collection_lock = threading.Lock()
+
+
+def _ensure_zvec_initialized() -> None:
+    """幂等初始化 zvec 引擎（线程安全单例）
+
+    zvec.init() 文档要求只调用一次，本函数用模块级标志位 + 锁保证幂等。
+    即使 zvec 内部允许重复调用，此守卫也能避免无谓的重复工作。
+    """
+    global _zvec_initialized
+    if _zvec_initialized:
+        return
+    with _init_lock:
+        if _zvec_initialized:
+            return
+        settings = get_settings()
+        # 确保 zvec 引擎日志目录存在（zvec 不会自动创建父目录）
+        log_dir = os.path.abspath(settings.ZVEC_LOG_DIR)
+        os.makedirs(log_dir, exist_ok=True)
+
+        # 映射日志级别字符串到 LogLevel 枚举
+        level_map = {
+            "debug": zvec.LogLevel.DEBUG,
+            "info": zvec.LogLevel.INFO,
+            "warn": zvec.LogLevel.WARN,
+            "error": zvec.LogLevel.ERROR,
+            "fatal": zvec.LogLevel.FATAL,
+        }
+        log_level = level_map.get(settings.ZVEC_LOG_LEVEL.lower(), zvec.LogLevel.WARN)
+
+        init_kwargs: dict[str, Any] = {
+            "log_type": zvec.LogType.FILE,
+            "log_dir": log_dir,
+            "log_level": log_level,
+        }
+        # 仅当配置非 None 时传入，让 zvec 使用内置默认值
+        if settings.ZVEC_QUERY_THREADS is not None:
+            init_kwargs["query_threads"] = settings.ZVEC_QUERY_THREADS
+        if settings.ZVEC_OPTIMIZE_THREADS is not None:
+            init_kwargs["optimize_threads"] = settings.ZVEC_OPTIMIZE_THREADS
+        if settings.ZVEC_MEMORY_LIMIT_MB is not None:
+            init_kwargs["memory_limit_mb"] = settings.ZVEC_MEMORY_LIMIT_MB
+
+        zvec.init(**init_kwargs)
+        _zvec_initialized = True
+        logger.info("✅ zvec 引擎已初始化")
+
+
+def _resolve_path(name_or_path: str) -> str:
+    """规范化 collection 路径
+
+    支持两种入参：
+    - 相对/绝对的 collection 名称（如 "docs"）→ 拼接到 ZVEC_DATA_DIR 下
+    - 已是绝对路径 → 原样使用
+    """
+    if os.path.isabs(name_or_path):
+        return os.path.abspath(name_or_path)
+    settings = get_settings()
+    return os.path.abspath(os.path.join(settings.ZVEC_DATA_DIR, name_or_path))
+
+
+def _to_doc(data: dict[str, Any]) -> Doc:
+    """将业务字典转换为 zvec.Doc
+
+    约定字典键：id（必需）、vectors（可选）、fields（可选）。
+    """
+    if "id" not in data:
+        raise ZvecError("文档字典必须包含 'id' 字段")
+    return Doc(
+        id=str(data["id"]),
+        vectors=data.get("vectors"),
+        fields=data.get("fields"),
+    )
+
+
+def _doc_to_dict(doc: Doc, include_vector: bool = True) -> dict[str, Any]:
+    """将 zvec.Doc 转换为业务字典（与 SQL Kit 返回字典列表风格一致）
+
+    Args:
+        doc: zvec 文档对象
+        include_vector: 是否包含向量数据
+    """
+    result: dict[str, Any] = {"id": doc.id}
+    if doc.score is not None:
+        result["score"] = doc.score
+    # 展开标量字段到顶层，便于业务直接取值
+    for name in doc.field_names():
+        result[name] = doc.field(name)
+    if include_vector:
+        for vname in doc.vector_names():
+            result[vname] = doc.vector(vname)
+    return result
+
+
+class ZvecKit:
+    """zvec 向量数据库工具类
+
+    采用独立工具范式（类似 RedisKit），自管理 Collection 实例池。
+    所有方法均为类方法，业务代码无需实例化。
+
+    与 SQL Kit（PgKit/MySqlKit/SQLiteKit）的一致性：
+    - 保留：connect / close / is_connected / count 等同名同语义方法
+    - 替换：无 SQL 相关方法（query/execute/transaction），
+            改为向量等价方法（vector_query 在阶段2实现）
+    - 变体：CRUD 入参无 ORM model，改为 collection 名称 + 业务字典
+    """
+
+    # ==================== 连接管理 ====================
+
+    @classmethod
+    def connect(cls) -> bool:
+        """初始化 zvec 引擎并验证数据目录可写
+
+        幂等：多次调用安全。zvec.init() 在进程内只执行一次。
+
+        Returns:
+            True 表示初始化成功
+        """
+        try:
+            _ensure_zvec_initialized()
+            # 验证数据目录可创建可写
+            settings = get_settings()
+            data_dir = os.path.abspath(settings.ZVEC_DATA_DIR)
+            os.makedirs(data_dir, exist_ok=True)
+            logger.info("✅ zvec 向量库连接成功")
+            return True
+        except Exception as e:
+            logger.error(f"❌ zvec 向量库连接失败: {e}")
+            return False
+
+    @classmethod
+    def close(cls) -> None:
+        """关闭所有已打开的 Collection 句柄
+
+        注意：zvec 的 zvec.init() 是进程级单例，无法重置；
+        本方法仅清理 Collection 注册表，让后续 open 重建句柄。
+        """
+        with _collection_lock:
+            for path, collection in _collections.items():
+                try:
+                    collection.flush()
+                except Exception as e:
+                    logger.warning(f"关闭 collection {path} 时 flush 失败: {e}")
+            _collections.clear()
+        logger.info("🔌 zvec 所有 collection 句柄已关闭")
+
+    @classmethod
+    def is_connected(cls) -> bool:
+        """检查 zvec 引擎是否已初始化"""
+        return _zvec_initialized
+
+    # ==================== Collection 管理 ====================
+
+    @classmethod
+    def create_collection(
+        cls,
+        name_or_path: str,
+        schema: CollectionSchema,
+    ) -> Collection:
+        """创建新的 Collection 并打开
+
+        Args:
+            name_or_path: collection 名称（拼接到 ZVEC_DATA_DIR）或绝对路径
+            schema: collection 结构定义
+
+        Returns:
+            打开的 Collection 实例
+        """
+        cls.connect()
+        path = _resolve_path(name_or_path)
+        if os.path.exists(path):
+            raise ZvecError(f"collection 目录已存在，请先 destroy 或换名: {path}")
+        try:
+            collection = zvec.create_and_open(path=path, schema=schema)
+            with _collection_lock:
+                _collections[path] = collection
+            logger.info(f"✅ 创建 collection: {path}")
+            return collection
+        except Exception as e:
+            logger.error(f"❌ 创建 collection 失败 {path}: {e}")
+            raise ZvecError(f"创建 collection 失败: {e}") from e
+
+    @classmethod
+    def open_collection(cls, name_or_path: str) -> Collection:
+        """打开已有 Collection（线程安全单例，重复打开返回同一实例）
+
+        Args:
+            name_or_path: collection 名称或绝对路径
+
+        Returns:
+            打开的 Collection 实例
+        """
+        cls.connect()
+        path = _resolve_path(name_or_path)
+        # 双重检查：已打开则直接返回，避免重复 open 导致资源占用
+        if path in _collections:
+            return _collections[path]
+        with _collection_lock:
+            if path in _collections:
+                return _collections[path]
+            if not os.path.exists(path):
+                raise ZvecError(f"collection 目录不存在: {path}")
+            try:
+                collection = zvec.open(path)
+                _collections[path] = collection
+                logger.info(f"✅ 打开 collection: {path}")
+                return collection
+            except Exception as e:
+                logger.error(f"❌ 打开 collection 失败 {path}: {e}")
+                raise ZvecError(f"打开 collection 失败: {e}") from e
+
+    @classmethod
+    def collection_exists(cls, name_or_path: str) -> bool:
+        """检查 Collection 是否存在（基于目录是否存在）"""
+        return os.path.exists(_resolve_path(name_or_path))
+
+    @classmethod
+    def list_collections(cls) -> list[str]:
+        """列出数据目录下的所有 collection 名称"""
+        settings = get_settings()
+        data_dir = os.path.abspath(settings.ZVEC_DATA_DIR)
+        if not os.path.exists(data_dir):
+            return []
+        return [
+            name
+            for name in os.listdir(data_dir)
+            if os.path.isdir(os.path.join(data_dir, name))
+        ]
+
+    @classmethod
+    def destroy_collection(cls, name_or_path: str) -> bool:
+        """永久删除 Collection（不可恢复）
+
+        Args:
+            name_or_path: collection 名称或绝对路径
+
+        Returns:
+            True 表示删除成功，False 表示不存在
+        """
+        path = _resolve_path(name_or_path)
+        # 先从注册表移除
+        with _collection_lock:
+            collection = _collections.pop(path, None)
+
+        try:
+            if collection is not None:
+                collection.destroy()
+            elif os.path.exists(path):
+                # 未在注册表中的，尝试 open 后 destroy
+                col = zvec.open(path)
+                col.destroy()
+            else:
+                return False
+            logger.info(f"✅ 已销毁 collection: {path}")
+            return True
+        except Exception as e:
+            logger.error(f"❌ 销毁 collection 失败 {path}: {e}")
+            raise ZvecError(f"销毁 collection 失败: {e}") from e
+
+    @classmethod
+    def get_schema(cls, name_or_path: str) -> CollectionSchema:
+        """获取 Collection 的 schema 定义"""
+        return cls.open_collection(name_or_path).schema
+
+    @classmethod
+    def get_stats(cls, name_or_path: str) -> dict[str, Any]:
+        """获取 Collection 的统计信息（解析 JSON 字符串为字典）
+
+        Returns:
+            统计字典，常见键：doc_count、index_completeness
+
+        Note:
+            zvec 返回的 stats 可能是原生 str、也可能是 C++ 绑定的字符串包装类
+            （isinstance(stats, str) 为 False 但 str(stats) 得到 JSON），
+            因此统一先 str() 再解析，保证健壮性。
+        """
+        collection = cls._get_collection(name_or_path)
+        raw = collection.stats
+        # 统一转字符串后解析 JSON（兼容 str 与字符串包装类两种情况）
+        text = raw if isinstance(raw, str) else str(raw)
+        try:
+            parsed = json.loads(text)
+        except (json.JSONDecodeError, TypeError):
+            return {"raw": text}
+        if not isinstance(parsed, dict):
+            return {"raw": text}
+        return parsed
+
+    @classmethod
+    def _get_collection(cls, name_or_path: str) -> Collection:
+        """内部方法：获取已注册的 Collection（不存在则 open）"""
+        path = _resolve_path(name_or_path)
+        if path in _collections:
+            return _collections[path]
+        return cls.open_collection(name_or_path)
+
+    # ==================== 文档 CRUD ====================
+
+    @classmethod
+    def insert(cls, name_or_path: str, data: dict[str, Any]) -> str:
+        """插入单条文档
+
+        Args:
+            name_or_path: collection 名称或路径
+            data: 文档字典，须含 'id'，可含 'vectors'/'fields'
+
+        Returns:
+            插入文档的 id
+        """
+        collection = cls._get_collection(name_or_path)
+        try:
+            doc = _to_doc(data)
+            collection.insert(doc)
+            logger.info(f"✅ 插入文档: {doc.id}")
+            return doc.id
+        except Exception as e:
+            logger.error(f"❌ 插入文档失败: {e}")
+            raise ZvecError(f"插入文档失败: {e}") from e
+
+    @classmethod
+    def insert_multi(
+        cls, name_or_path: str, data_list: list[dict[str, Any]]
+    ) -> list[str]:
+        """批量插入文档
+
+        Args:
+            name_or_path: collection 名称或路径
+            data_list: 文档字典列表
+
+        Returns:
+            插入成功的文档 id 列表
+        """
+        collection = cls._get_collection(name_or_path)
+        try:
+            docs = [_to_doc(d) for d in data_list]
+            collection.insert(docs)
+            ids = [d.id for d in docs]
+            logger.info(f"✅ 批量插入 {len(ids)} 条文档")
+            return ids
+        except Exception as e:
+            logger.error(f"❌ 批量插入文档失败: {e}")
+            raise ZvecError(f"批量插入文档失败: {e}") from e
+
+    @classmethod
+    def upsert(cls, name_or_path: str, data: dict[str, Any]) -> str:
+        """插入或更新单条文档（按 id 判断存在性）"""
+        collection = cls._get_collection(name_or_path)
+        try:
+            doc = _to_doc(data)
+            collection.upsert(doc)
+            logger.info(f"✅ upsert 文档: {doc.id}")
+            return doc.id
+        except Exception as e:
+            logger.error(f"❌ upsert 文档失败: {e}")
+            raise ZvecError(f"upsert 文档失败: {e}") from e
+
+    @classmethod
+    def update(cls, name_or_path: str, data: dict[str, Any]) -> str:
+        """更新单条文档（仅更新指定字段，其余不变）
+
+        Args:
+            data: 须含 'id'，'vectors'/'fields' 为可选的待更新字段
+        """
+        collection = cls._get_collection(name_or_path)
+        try:
+            doc = _to_doc(data)
+            collection.update(doc)
+            logger.info(f"✅ 更新文档: {doc.id}")
+            return doc.id
+        except Exception as e:
+            logger.error(f"❌ 更新文档失败: {e}")
+            raise ZvecError(f"更新文档失败: {e}") from e
+
+    @classmethod
+    def fetch(
+        cls,
+        name_or_path: str,
+        ids: str | list[str],
+        *,
+        include_vector: bool = True,
+    ) -> dict[str, dict[str, Any]]:
+        """按 id 获取文档
+
+        Args:
+            name_or_path: collection 名称或路径
+            ids: 单个 id 或 id 列表
+            include_vector: 是否包含向量数据
+
+        Returns:
+            id -> 文档字典 的映射（不存在的 id 会被省略）
+        """
+        collection = cls._get_collection(name_or_path)
+        try:
+            result = collection.fetch(ids, include_vector=include_vector)
+            return {
+                doc_id: _doc_to_dict(doc, include_vector)
+                for doc_id, doc in result.items()
+            }
+        except Exception as e:
+            logger.error(f"❌ 获取文档失败: {e}")
+            raise ZvecError(f"获取文档失败: {e}") from e
+
+    @classmethod
+    def delete(cls, name_or_path: str, ids: str | list[str]) -> bool:
+        """按 id 删除文档
+
+        Args:
+            ids: 单个 id 或 id 列表
+
+        Returns:
+            True 表示删除操作执行成功
+        """
+        collection = cls._get_collection(name_or_path)
+        try:
+            collection.delete(ids)
+            logger.info(f"✅ 删除文档: {ids}")
+            return True
+        except Exception as e:
+            logger.error(f"❌ 删除文档失败: {e}")
+            raise ZvecError(f"删除文档失败: {e}") from e
+
+    @classmethod
+    def delete_by_filter(cls, name_or_path: str, filter_expr: str) -> bool:
+        """按过滤表达式删除文档
+
+        Args:
+            filter_expr: 布尔表达式。zvec 过滤语法：相等用单等号 ``=````
+                （非 ``==``），数值比较用 ``year < 1950``，支持 AND/OR/IN 等。
+
+        Returns:
+            True 表示删除操作执行成功
+        """
+        collection = cls._get_collection(name_or_path)
+        try:
+            collection.delete_by_filter(filter=filter_expr)
+            logger.info(f"✅ 按条件删除文档: {filter_expr}")
+            return True
+        except Exception as e:
+            logger.error(f"❌ 按条件删除文档失败: {e}")
+            raise ZvecError(f"按条件删除文档失败: {e}") from e
+
+    @classmethod
+    def count(cls, name_or_path: str) -> int:
+        """获取 collection 中的文档数量
+
+        基于统计信息中的 doc_count 字段。
+        """
+        stats = cls.get_stats(name_or_path)
+        return int(stats.get("doc_count", 0))
+
+    @classmethod
+    def flush(cls, name_or_path: str) -> None:
+        """将待写入数据强制刷盘（保证持久化）"""
+        collection = cls._get_collection(name_or_path)
+        try:
+            collection.flush()
+            logger.info(f"✅ flush 完成: {name_or_path}")
+        except Exception as e:
+            logger.error(f"❌ flush 失败: {e}")
+            raise ZvecError(f"flush 失败: {e}") from e
+
+    # ==================== 向量检索（阶段2完整实现） ====================
+
+    @classmethod
+    def vector_query(
+        cls,
+        name_or_path: str,
+        *,
+        field: str,
+        vector: list[float],
+        topk: int = 10,
+        filter_expr: str | None = None,
+        include_vector: bool = False,
+    ) -> list[dict[str, Any]]:
+        """向量相似度检索
+
+        Args:
+            name_or_path: collection 名称或路径
+            field: 向量字段名
+            vector: 查询向量
+            topk: 返回近邻数量
+            filter_expr: 预过滤表达式。注意 zvec 过滤语法：相等用单等号
+                ``category = 'life'``（非 ``==``），数值比较用 ``year > 2020``，
+                支持 AND / OR / NOT / IN / BETWEEN / LIKE。
+            include_vector: 结果是否包含向量数据
+
+        Returns:
+            文档字典列表（含 id/score/fields），按相似度降序
+        """
+        collection = cls._get_collection(name_or_path)
+        try:
+            results = collection.query(
+                Query(field_name=field, vector=vector),
+                topk=topk,
+                filter=filter_expr,
+                include_vector=include_vector,
+            )
+            return [_doc_to_dict(doc, include_vector) for doc in results]
+        except Exception as e:
+            logger.error(f"❌ 向量检索失败: {e}")
+            raise ZvecError(f"向量检索失败: {e}") from e
+
+    @classmethod
+    def optimize(cls, name_or_path: str) -> None:
+        """优化 collection（合并段、重建向量索引）
+
+        大批量写入后建议调用，以加速后续检索并提升召回率。
+        属于重操作，不应在每条 insert 后调用。
+        """
+        collection = cls._get_collection(name_or_path)
+        try:
+            collection.optimize()
+            logger.info(f"✅ optimize 完成: {name_or_path}")
+        except Exception as e:
+            logger.error(f"❌ optimize 失败: {e}")
+            raise ZvecError(f"optimize 失败: {e}") from e
+
+    # ==================== Schema 工厂（降低使用门槛） ====================
+
+    @staticmethod
+    def make_text_embedding_schema(
+        name: str,
+        dim: int,
+        metric: str | None = None,
+        extra_fields: list[FieldSchema] | None = None,
+    ) -> CollectionSchema:
+        """构建常用的「文本 embedding」schema 工厂
+
+        默认包含：
+        - 向量字段 'embedding'（FP32, HNSW 索引）
+        - 标量字段 'publish_year'（INT32, 倒排索引，可选）
+
+        Args:
+            name: collection 名称
+            dim: 向量维度（如 768/1536）
+            metric: 距离度量（cosine/l2/ip），默认取配置 ZVEC_DEFAULT_METRIC
+            extra_fields: 额外的标量字段列表
+        """
+        settings = get_settings()
+        metric_str = (metric or settings.ZVEC_DEFAULT_METRIC).lower()
+        metric_map = {
+            "cosine": MetricType.COSINE,
+            "l2": MetricType.L2,
+            "ip": MetricType.IP,
+        }
+        metric_type = metric_map.get(metric_str, MetricType.COSINE)
+
+        vectors = [
+            VectorSchema(
+                name="embedding",
+                data_type=DataType.VECTOR_FP32,
+                dimension=dim,
+                index_param=HnswIndexParam(metric_type=metric_type),
+            )
+        ]
+        fields: list[FieldSchema] = [
+            FieldSchema(
+                name="publish_year",
+                data_type=DataType.INT32,
+                index_param=InvertIndexParam(enable_range_optimization=True),
+            )
+        ]
+        if extra_fields:
+            fields.extend(extra_fields)
+
+        return CollectionSchema(name=name, fields=fields, vectors=vectors)
+
+    # ==================== 测试辅助 ====================
+
+    @classmethod
+    def _reset_for_test(cls) -> None:
+        """重置内部状态（仅供单测隔离使用，业务代码勿调）
+
+        注意：zvec.init() 是进程级单例无法真正重置，此处仅清理 collection 注册表。
+        """
+        with _collection_lock:
+            _collections.clear()