nex_docus/backend/app/services/zvec_service.py

"""
ZVec 本地向量化服务 - 使用本地 zvec 库 + 模型配置中的 embedding 模型
"""
import os
import hashlib
import logging
from pathlib import Path
from typing import Dict, List, Optional, Any

import zvec
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, delete

from app.models.llm_model_config import LLMModelConfig
from app.models.document_vector import DocumentVector
from app.services.llm_provider_service import LLMProviderService

logger = logging.getLogger(__name__)

def _resolve_zvec_data_dir() -> str:
    """解析 ZVec 向量库根目录。

    优先用 ZVEC_DATA_DIR 环境变量；否则放到 storage 存储区下的
    vector_index 目录，与 search_index 并列，便于统一备份/管理。
    """
    env_dir = os.getenv("ZVEC_DATA_DIR")
    if env_dir:
        return env_dir

    from app.core.config import settings

    storage_root = Path(settings.STORAGE_ROOT)
    if not storage_root.is_absolute():
        backend_dir = Path(__file__).parent.parent.parent
        storage_root = (backend_dir / storage_root).resolve()
    return str(storage_root / "vector_index")


ZVEC_DATA_DIR = _resolve_zvec_data_dir()
EMBEDDING_DIMENSION = int(os.getenv("ZVEC_EMBEDDING_DIM", "1536"))


class ZVecService:
    """基于本地 zvec 库的向量存储和检索服务"""

    _collections: Dict[int, zvec.Collection] = {}

    @classmethod
    def _get_collection_path(cls, project_id: int) -> str:
        """返回项目 collection 路径。

        只确保父目录存在；叶子目录由 zvec.create_and_open 自行创建，
        预先创建叶子目录会让 create_and_open 报 "path exists"。
        """
        path = Path(ZVEC_DATA_DIR) / str(project_id)
        path.parent.mkdir(parents=True, exist_ok=True)
        return str(path)

    @classmethod
    def _open_collection(cls, project_id: int) -> Optional[zvec.Collection]:
        """打开已存在的 collection；不存在则返回 None"""
        if project_id in cls._collections:
            return cls._collections[project_id]

        collection_path = cls._get_collection_path(project_id)
        try:
            collection = zvec.open(path=collection_path)
        except Exception:
            return None

        cls._collections[project_id] = collection
        return collection

    @classmethod
    def _get_or_create_collection(cls, project_id: int, dimension: int) -> zvec.Collection:
        """按指定维度获取或创建 collection。

        ZVec 的 collection 维度创建后固定，若已存在的维度与当前 embedding
        模型不一致，则重建该项目的 collection（旧向量作废，需重新向量化）。
        """
        collection = cls._open_collection(project_id)
        if collection is not None:
            existing_dim = cls._collection_dimension(collection)
            if existing_dim is not None and existing_dim != dimension:
                logger.warning(
                    "Project %s collection dim %s != model dim %s, rebuilding",
                    project_id, existing_dim, dimension,
                )
                cls._drop_collection(project_id)
                collection = None
            else:
                return collection

        collection_path = cls._get_collection_path(project_id)
        # ZVec 的 create_and_open 要求目标路径不存在，清掉可能残留的空/损坏目录，保证幂等
        leaf = Path(ZVEC_DATA_DIR) / str(project_id)
        if leaf.exists():
            import shutil
            shutil.rmtree(leaf, ignore_errors=True)
        schema = zvec.CollectionSchema(
            name=f"project_{project_id}",
            vectors=zvec.VectorSchema(
                "embedding", zvec.DataType.VECTOR_FP32, dimension
            ),
        )
        collection = zvec.create_and_open(path=collection_path, schema=schema)
        cls._collections[project_id] = collection
        return collection

    @staticmethod
    def _collection_dimension(collection: zvec.Collection) -> Optional[int]:
        """尽力读取 collection 的向量维度，读取失败返回 None"""
        try:
            schema = collection.schema
            vectors = getattr(schema, "vectors", None)
            if vectors is None:
                return None
            vector_schema = vectors[0] if isinstance(vectors, (list, tuple)) else vectors
            return getattr(vector_schema, "dimension", None)
        except Exception:
            return None

    @classmethod
    def _drop_collection(cls, project_id: int) -> None:
        """删除项目 collection 的本地数据并清理缓存"""
        cls._collections.pop(project_id, None)
        import shutil

        path = Path(ZVEC_DATA_DIR) / str(project_id)
        if path.exists():
            shutil.rmtree(path, ignore_errors=True)

    @classmethod
    async def get_embedding_config(cls, db: AsyncSession) -> Optional[LLMModelConfig]:
        """获取系统配置的 embedding 模型（按 model_type 精确查询）"""
        stmt = (
            select(LLMModelConfig)
            .where(
                LLMModelConfig.is_active == True,
                LLMModelConfig.model_type == "embedding",
            )
            .order_by(
                LLMModelConfig.is_default.desc(),
                LLMModelConfig.updated_at.desc(),
                LLMModelConfig.config_id.desc(),
            )
        )
        result = await db.execute(stmt)
        return result.scalars().first()

    @classmethod
    async def generate_embedding(
        cls, db: AsyncSession, text: str
    ) -> Optional[List[float]]:
        """调用配置的 embedding 模型生成向量"""
        if not text or not text.strip():
            return None

        config = await cls.get_embedding_config(db)
        if not config:
            logger.warning("No embedding model configured, skipping vectorization")
            return None

        try:
            return await LLMProviderService.generate_embedding(
                provider=config.provider,
                endpoint_url=config.endpoint_url,
                api_key=config.api_key,
                llm_model_name=config.llm_model_name,
                text=text,
                timeout=config.llm_timeout or 60,
                dimension=config.embedding_dimension,
            )
        except Exception as exc:
            logger.error(f"Embedding generation failed: {exc}")
            return None

    @staticmethod
    def _doc_id(file_path: str, chunk_index: int = 0) -> str:
        """生成 ZVec 合法的 doc id。

        ZVec 的 doc id 不允许中文等非 ASCII 字符，故对文件路径做 MD5 哈希。
        分块后每个 chunk 需独立 doc id，故追加 chunk 序号。
        file_path 的反查通过 document_vector 表的 zvec_id 字段完成。
        """
        base = hashlib.md5(file_path.encode("utf-8")).hexdigest()
        return f"{base}#{chunk_index}"

    @staticmethod
    def _content_hash(content: str) -> str:
        return hashlib.sha256((content or "").encode("utf-8")).hexdigest()

    @staticmethod
    def _make_anchor(text: str, max_chars: int = 120) -> str:
        """从 chunk 文本中提取定位锚点：首个非空文本行、去除常见 markdown 标记。

        锚点供旧引用定位兼容使用，故取纯文本片段。
        """
        import re

        for raw_line in (text or "").splitlines():
            line = raw_line.strip()
            if not line:
                continue
            # 去除标题井号、列表符号、引用符号等行首标记
            line = re.sub(r"^\s*#{1,6}\s+", "", line)
            line = re.sub(r"^\s*[-*+]\s+", "", line)
            line = re.sub(r"^\s*>\s+", "", line)
            line = re.sub(r"^\s*\d+\.\s+", "", line)
            # 去除行内强调/代码标记
            line = re.sub(r"[*_`~]", "", line)
            line = line.strip()
            if line:
                return line[:max_chars]
        return (text or "").strip()[:max_chars]

    @classmethod
    def _chunk_text(
        cls,
        content: str,
        chunk_size: int,
        overlap: int,
    ) -> List[Dict[str, Any]]:
        """按字符滑动窗口分块。

        返回 [{index, text, anchor}]。相邻分块重叠 overlap 个字符，
        以避免语义在分块边界被切断。空内容返回空列表。
        """
        text = content or ""
        if not text.strip():
            return []

        chunk_size = max(1, int(chunk_size))
        overlap = max(0, min(int(overlap), chunk_size - 1))
        step = chunk_size - overlap

        chunks: List[Dict[str, Any]] = []
        start = 0
        length = len(text)
        index = 0
        while start < length:
            piece = text[start:start + chunk_size]
            if piece.strip():
                chunks.append({
                    "index": index,
                    "text": piece,
                    "anchor": cls._make_anchor(piece),
                })
                index += 1
            start += step

        return chunks

    @classmethod
    async def _upsert_vector_record(
        cls,
        db: AsyncSession,
        project_id: int,
        file_path: str,
        *,
        status: str,
        chunk_index: int = 0,
        chunk_text: Optional[str] = None,
        content_hash: Optional[str] = None,
        zvec_id: Optional[str] = None,
        error_message: Optional[str] = None,
    ) -> None:
        """写入/更新 document_vector 表中某个分块的向量化状态。

        按 (project_id, file_path, chunk_index) 定位记录。失败状态通常写
        chunk_index=0 的一条即可，成功状态则逐分块写入。
        """
        stmt = select(DocumentVector).where(
            DocumentVector.project_id == project_id,
            DocumentVector.file_path == file_path,
            DocumentVector.chunk_index == chunk_index,
        )
        result = await db.execute(stmt)
        record = result.scalar_one_or_none()

        if record is None:
            record = DocumentVector(
                project_id=project_id,
                file_path=file_path,
                chunk_index=chunk_index,
                chunk_text=chunk_text,
                status=status,
                content_hash=content_hash,
                zvec_id=zvec_id,
                error_message=(error_message or "")[:500] or None,
            )
            db.add(record)
        else:
            record.status = status
            if chunk_text is not None:
                record.chunk_text = chunk_text
            if content_hash is not None:
                record.content_hash = content_hash
            if zvec_id is not None:
                record.zvec_id = zvec_id
            record.error_message = (error_message or "")[:500] or None

    @classmethod
    async def _purge_file_chunks(
        cls,
        db: AsyncSession,
        project_id: int,
        file_path: str,
    ) -> None:
        """删除某文件在 ZVec 与 document_vector 表中的所有分块记录（不 commit）。"""
        # 先查出该文件所有分块的 zvec_id，用于从 ZVec 集合中删除
        result = await db.execute(
            select(DocumentVector.zvec_id).where(
                DocumentVector.project_id == project_id,
                DocumentVector.file_path == file_path,
            )
        )
        zvec_ids = [z for (z,) in result.all() if z]
        if zvec_ids:
            try:
                collection = cls._open_collection(project_id)
                if collection is not None:
                    collection.delete(zvec_ids)
            except Exception as exc:
                logger.error(f"ZVec delete chunks failed for {file_path}: {exc}")

        await db.execute(
            delete(DocumentVector).where(
                DocumentVector.project_id == project_id,
                DocumentVector.file_path == file_path,
            )
        )

    @classmethod
    async def vectorize_markdown(
        cls,
        db: AsyncSession,
        project_id: int,
        file_path: str,
        content: str,
    ) -> bool:
        """对 MD 文件分块向量化并存入 ZVec，同时逐分块记录状态到 document_vector。

        流程：删除旧分块 → 分块 → 逐块生成 embedding、写入 ZVec、写记录。
        """
        from app.core.config import settings

        content_hash = cls._content_hash(content)

        # 先清理该文件的所有旧分块（增量重建），保证不残留过期向量
        await cls._purge_file_chunks(db, project_id, file_path)

        chunks = cls._chunk_text(content, settings.CHUNK_SIZE, settings.CHUNK_OVERLAP)
        if not chunks:
            # 空文件：清理后直接提交，视为成功（无可向量化内容）
            await db.commit()
            return True

        success_count = 0
        last_error: Optional[str] = None
        for chunk in chunks:
            embedding = await cls.generate_embedding(db, chunk["text"])
            if not embedding:
                last_error = "未配置可用的 embedding 模型或向量生成失败"
                continue

            doc_id = cls._doc_id(file_path, chunk["index"])
            try:
                collection = cls._get_or_create_collection(project_id, len(embedding))
                collection.insert([
                    zvec.Doc(id=doc_id, vectors={"embedding": embedding})
                ])
            except Exception as exc:
                logger.error(
                    f"ZVec insert failed for {file_path} chunk {chunk['index']}: {exc}"
                )
                last_error = str(exc)
                continue

            await cls._upsert_vector_record(
                db, project_id, file_path,
                status="success",
                chunk_index=chunk["index"],
                chunk_text=chunk["text"],
                content_hash=content_hash,
                zvec_id=doc_id,
            )
            success_count += 1

        if success_count == 0:
            # 全部分块失败：写一条 failed 记录（chunk_index=0）便于进度展示
            await cls._upsert_vector_record(
                db, project_id, file_path,
                status="failed",
                chunk_index=0,
                content_hash=content_hash,
                error_message=last_error or "向量化失败",
            )
            await db.commit()
            return False

        await db.commit()
        return True

    @classmethod
    async def delete_vector(
        cls,
        db: AsyncSession,
        project_id: int,
        file_path: str,
    ) -> bool:
        """从 ZVec 删除文档的所有分块向量，并清除 document_vector 记录"""
        try:
            await cls._purge_file_chunks(db, project_id, file_path)
            await db.commit()
        except Exception as exc:
            logger.error(f"Delete document_vector record failed for {file_path}: {exc}")
            await db.rollback()
        return True

    # 兼容 project_file_service 中的调用别名
    @classmethod
    async def delete_vectors(
        cls,
        db: AsyncSession,
        project_id: int,
        file_path: str,
    ) -> bool:
        """delete_vector 的别名（兼容文件服务调用）"""
        return await cls.delete_vector(db, project_id, file_path)

    @classmethod
    async def sync_vector_move(
        cls,
        db: AsyncSession,
        project_id: int,
        old_path: str,
        new_path: str,
        content: Optional[str] = None,
    ) -> bool:
        """文件移动/重命名时同步向量。

        content 为空时会尝试读取新路径文件内容再向量化。
        """
        await cls.delete_vector(db, project_id, old_path)
        if content is None:
            return await cls.revectorize_path(db, project_id, new_path)
        return await cls.vectorize_markdown(db, project_id, new_path, content)

    @classmethod
    async def update_vector_path(
        cls,
        db: AsyncSession,
        project_id: int,
        old_path: str,
        new_path: str,
        content: str,
    ) -> bool:
        """文件重命名/移动时更新向量（删除旧的，插入新的）"""
        await cls.delete_vector(db, project_id, old_path)
        return await cls.vectorize_markdown(db, project_id, new_path, content)

    @classmethod
    async def revectorize_path(
        cls,
        db: AsyncSession,
        project_id: int,
        file_path: str,
    ) -> bool:
        """读取项目中指定文件内容并重新向量化"""
        content = await cls._read_project_file(db, project_id, file_path)
        if content is None:
            return False
        return await cls.vectorize_markdown(db, project_id, file_path, content)

    @staticmethod
    async def _read_project_file(
        db: AsyncSession,
        project_id: int,
        file_path: str,
    ) -> Optional[str]:
        """读取项目内某个相对路径文件的文本内容"""
        from app.models.project import Project
        from app.services.storage import storage_service

        result = await db.execute(select(Project).where(Project.id == project_id))
        project = result.scalar_one_or_none()
        if not project:
            return None
        try:
            full_path = storage_service.get_secure_path(project.storage_key, file_path)
            return await storage_service.read_file(full_path)
        except Exception as exc:
            logger.error(f"Read project file failed {file_path}: {exc}")
            return None

    @classmethod
    def _list_markdown_files(cls, storage_key: str) -> List[str]:
        """列出项目下所有 MD 文件的相对路径（排除 _assets 资源目录）"""
        from app.services.storage import storage_service

        root = storage_service.get_secure_path(storage_key)
        if not root.exists():
            return []
        files: List[str] = []
        for md_path in root.rglob("*.md"):
            try:
                rel = md_path.relative_to(root).as_posix()
            except ValueError:
                continue
            # 跳过 _assets 资源目录
            if rel.startswith("_assets/") or "/_assets/" in rel:
                continue
            # 跳过 . 开头的隐藏文件或位于隐藏目录下的文件（如 .git/.obsidian 等）
            if any(part.startswith(".") for part in rel.split("/")):
                continue
            files.append(rel)
        return files

    @staticmethod
    def _aggregate_file_records(
        rows: List[DocumentVector],
    ) -> Dict[str, Dict[str, Any]]:
        """将多分块记录按文件聚合为文件级状态。

        返回 {file_path: {status, content_hash, error_message}}。
        任一分块 failed 则该文件视为 failed；否则若有 success 则 success。
        """
        by_file: Dict[str, Dict[str, Any]] = {}
        for r in rows:
            agg = by_file.get(r.file_path)
            if agg is None:
                agg = {"status": r.status, "content_hash": r.content_hash, "error_message": r.error_message}
                by_file[r.file_path] = agg
            else:
                # failed 优先级最高，用于暴露问题
                if r.status == "failed":
                    agg["status"] = "failed"
                    agg["error_message"] = r.error_message or agg.get("error_message")
                elif agg["status"] != "failed" and r.status == "success":
                    agg["status"] = "success"
                if agg.get("content_hash") is None and r.content_hash:
                    agg["content_hash"] = r.content_hash
        return by_file

    @classmethod
    async def get_progress(
        cls,
        db: AsyncSession,
        project_id: int,
        storage_key: str,
    ) -> Dict[str, Any]:
        """统计项目向量化进度：已成功 / 失败 / 待处理 / 总 MD 文件数"""
        md_files = cls._list_markdown_files(storage_key)
        total = len(md_files)

        result = await db.execute(
            select(DocumentVector).where(DocumentVector.project_id == project_id)
        )
        records = cls._aggregate_file_records(list(result.scalars().all()))

        success = 0
        failed_items: List[Dict[str, str]] = []
        pending_items: List[str] = []
        for path in md_files:
            record = records.get(path)
            if record and record["status"] == "success":
                success += 1
            elif record and record["status"] == "failed":
                failed_items.append({
                    "file_path": path,
                    "error": record.get("error_message") or "向量化失败",
                })
            else:
                pending_items.append(path)

        percent = int(success / total * 100) if total else 0
        return {
            "total": total,
            "success": success,
            "failed": len(failed_items),
            "pending": len(pending_items),
            "percent": percent,
            "failed_items": failed_items[:50],
            "pending_items": pending_items[:50],
            "embedding_ready": await cls.get_embedding_config(db) is not None,
        }

    @classmethod
    async def vectorize_project(
        cls,
        db: AsyncSession,
        project_id: int,
        storage_key: str,
        *,
        force: bool = False,
    ) -> Dict[str, Any]:
        """批量向量化项目下所有 MD 文件。

        force=False 时跳过内容未变化且已成功的文件（增量）；
        force=True 时全部重新向量化（全量）。
        """
        md_files = cls._list_markdown_files(storage_key)
        md_file_set = set(md_files)

        result = await db.execute(
            select(DocumentVector).where(DocumentVector.project_id == project_id)
        )
        records = cls._aggregate_file_records(list(result.scalars().all()))

        if force:
            cls._drop_collection(project_id)
            await db.execute(
                delete(DocumentVector).where(DocumentVector.project_id == project_id)
            )
            await db.commit()
            records = {}
        else:
            for stale_path in set(records) - md_file_set:
                await cls.delete_vector(db, project_id, stale_path)
                records.pop(stale_path, None)

        processed = 0
        skipped = 0
        failed = 0
        for path in md_files:
            content = await cls._read_project_file(db, project_id, path)
            if content is None:
                failed += 1
                continue

            if not force:
                record = records.get(path)
                if (
                    record
                    and record["status"] == "success"
                    and record.get("content_hash") == cls._content_hash(content)
                ):
                    skipped += 1
                    continue

            ok = await cls.vectorize_markdown(db, project_id, path, content)
            if ok:
                processed += 1
            else:
                failed += 1

        return {
            "total": len(md_files),
            "processed": processed,
            "skipped": skipped,
            "failed": failed,
        }

    @classmethod
    async def search_similar(
        cls,
        db: AsyncSession,
        project_id: int,
        query: str,
        top_k: int = 5,
    ) -> List[Dict[str, Any]]:
        """基于查询文本检索相似文档"""
        query_embedding = await cls.generate_embedding(db, query)
        if not query_embedding:
            return []

        try:
            collection = cls._open_collection(project_id)
            if collection is None:
                return []  # 项目尚未向量化
            results = collection.query(
                zvec.VectorQuery("embedding", vector=query_embedding),
                topk=top_k,
            )

            # doc id 是「路径哈希#分块序号」，通过 document_vector 表的 zvec_id
            # 反查真实 file_path、chunk_index 及分块文本 chunk_text
            zvec_ids = [r.id for r in results]
            if not zvec_ids:
                return []

            db_result = await db.execute(
                select(DocumentVector).where(
                    DocumentVector.project_id == project_id,
                    DocumentVector.zvec_id.in_(zvec_ids),
                )
            )
            record_by_id = {r.zvec_id: r for r in db_result.scalars().all()}

            matched_docs = []
            for result in results:
                record = record_by_id.get(result.id)
                if not record:
                    continue
                matched_docs.append({
                    "file_path": record.file_path,
                    "chunk_index": record.chunk_index,
                    "chunk_text": record.chunk_text or "",
                    "score": result.score if hasattr(result, "score") else 0.0,
                })

            return matched_docs

        except Exception as exc:
            logger.error(f"ZVec search failed: {exc}")
            return []

    @classmethod
    def close_collection(cls, project_id: int) -> None:
        """关闭项目的 ZVec collection"""
        if project_id in cls._collections:
            del cls._collections[project_id]


zvec_service = ZVecService()