""" 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()