jarvis-ai-assistant 0.1.222__py3-none-any.whl → 0.7.0__py3-none-any.whl

jarvis/jarvis_rag/retriever.py

@@ -1,6 +1,8 @@
 import os
 import pickle
-from typing import Any, Dict, List, cast
+import json
+import hashlib
+from typing import Any, Dict, List, Optional, cast
 
 import chromadb
 from langchain.docstore.document import Document
@@ -8,6 +10,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from rank_bm25 import BM25Okapi  # type: ignore
 
 from .embedding_manager import EmbeddingManager
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 
 
 class ChromaRetriever:
@@ -39,35 +42,271 @@ class ChromaRetriever:
         self.collection = self.client.get_or_create_collection(
             name=self.collection_name
         )
-        print(
-            f"✅ ChromaDB 客户端已在 '{db_path}' 初始化，集合为 '{collection_name}'。"
+        PrettyOutput.print(
+            f"ChromaDB 客户端已在 '{db_path}' 初始化，集合为 '{collection_name}'。",
+            OutputType.SUCCESS,
         )
 
         # BM25索引设置
         self.bm25_index_path = os.path.join(self.db_path, f"{collection_name}_bm25.pkl")
         self._load_or_initialize_bm25()
+        # 清单文件用于检测源文件的变更/删除
+        self.manifest_path = os.path.join(
+            self.db_path, f"{collection_name}_manifest.json"
+        )
 
     def _load_or_initialize_bm25(self):
         """从磁盘加载BM25索引或初始化一个新索引。"""
         if os.path.exists(self.bm25_index_path):
-            print("🔍 正在加载现有的 BM25 索引...")
+            PrettyOutput.print("正在加载现有的 BM25 索引...", OutputType.INFO)
             with open(self.bm25_index_path, "rb") as f:
                 data = pickle.load(f)
                 self.bm25_corpus = data["corpus"]
                 self.bm25_index = BM25Okapi(self.bm25_corpus)
-            print("✅ BM25 索引加载成功。")
+            PrettyOutput.print("BM25 索引加载成功。", OutputType.SUCCESS)
         else:
-            print("⚠️ 未找到 BM25 索引，将初始化一个新的。")
+            PrettyOutput.print(
+                "未找到 BM25 索引，将初始化一个新的。", OutputType.WARNING
+            )
             self.bm25_corpus = []
             self.bm25_index = None
 
     def _save_bm25_index(self):
         """将BM25索引保存到磁盘。"""
         if self.bm25_index:
-            print("💾 正在保存 BM25 索引...")
+            PrettyOutput.print("正在保存 BM25 索引...", OutputType.INFO)
             with open(self.bm25_index_path, "wb") as f:
                 pickle.dump({"corpus": self.bm25_corpus, "index": self.bm25_index}, f)
-            print("✅ BM25 索引保存成功。")
+            PrettyOutput.print("BM25 索引保存成功。", OutputType.SUCCESS)
+
+    def _load_manifest(self) -> Dict[str, Dict[str, Any]]:
+        """加载已索引文件清单，用于变更检测。"""
+        if os.path.exists(self.manifest_path):
+            try:
+                with open(self.manifest_path, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+                if isinstance(data, dict):
+                    return data  # type: ignore[return-value]
+            except Exception:
+                pass
+        return {}
+
+    def _save_manifest(self, manifest: Dict[str, Dict[str, Any]]) -> None:
+        """保存已索引文件清单。"""
+        try:
+            with open(self.manifest_path, "w", encoding="utf-8") as f:
+                json.dump(manifest, f, ensure_ascii=False, indent=2)
+        except Exception as e:
+            PrettyOutput.print(f"保存索引清单失败: {e}", OutputType.WARNING)
+
+    def _compute_md5(
+        self, file_path: str, chunk_size: int = 1024 * 1024
+    ) -> Optional[str]:
+        """流式计算文件的MD5，避免占用过多内存。失败时返回None。"""
+        try:
+            md5 = hashlib.md5()
+            with open(file_path, "rb") as f:
+                while True:
+                    data = f.read(chunk_size)
+                    if not data:
+                        break
+                    md5.update(data)
+            return md5.hexdigest()
+        except Exception:
+            return None
+
+    def _update_manifest_with_sources(self, sources: List[str]) -> None:
+        """根据本次新增文档的来源，更新索引清单（记录mtime与size）。"""
+        manifest = self._load_manifest()
+        updated = 0
+        for src in set(sources):
+            try:
+                if isinstance(src, str) and os.path.exists(src):
+                    st = os.stat(src)
+                    entry: Dict[str, Any] = {
+                        "mtime": float(st.st_mtime),
+                        "size": int(st.st_size),
+                    }
+                    md5sum = self._compute_md5(src)
+                    if md5sum:
+                        entry["md5"] = md5sum
+                    manifest[src] = entry  # type: ignore[dict-item]
+                    updated += 1
+            except Exception:
+                continue
+        if updated > 0:
+            self._save_manifest(manifest)
+            PrettyOutput.print(
+                f"已更新索引清单，记录 {updated} 个源文件状态。", OutputType.INFO
+            )
+
+    def _detect_changed_or_deleted(self) -> Dict[str, List[str]]:
+        """检测已记录的源文件是否发生变化或被删除。"""
+        manifest = self._load_manifest()
+        changed: List[str] = []
+        deleted: List[str] = []
+        for src, info in manifest.items():
+            try:
+                if not os.path.exists(src):
+                    deleted.append(src)
+                    continue
+                st = os.stat(src)
+                size_changed = int(info.get("size", -1)) != int(st.st_size)
+                if size_changed:
+                    changed.append(src)
+                    continue
+                md5_old = info.get("md5")
+                if md5_old:
+                    # 仅在mtime变化时计算md5以降低开销
+                    mtime_changed = (
+                        abs(float(info.get("mtime", 0.0)) - float(st.st_mtime)) >= 1e-6
+                    )
+                    if mtime_changed:
+                        md5_new = self._compute_md5(src)
+                        if not md5_new or md5_new != md5_old:
+                            changed.append(src)
+                else:
+                    # 没有记录md5，回退使用mtime判断
+                    mtime_changed = (
+                        abs(float(info.get("mtime", 0.0)) - float(st.st_mtime)) >= 1e-6
+                    )
+                    if mtime_changed:
+                        changed.append(src)
+            except Exception:
+                # 无法读取文件状态，视为发生变化
+                changed.append(src)
+        return {"changed": changed, "deleted": deleted}
+
+    def _warn_if_sources_changed(self) -> None:
+        """如发现已索引文件变化或删除，给出提醒。"""
+        result = self._detect_changed_or_deleted()
+        changed = result["changed"]
+        deleted = result["deleted"]
+        if not changed and not deleted:
+            return
+        # 为避免在循环中逐条打印，先拼接后统一打印
+        lines: List[str] = []
+        if changed:
+            lines.append(
+                f"检测到 {len(changed)} 个已索引文件发生变化，建议重新索引以保证检索准确性。"
+            )
+            lines.extend([f"  变更: {p}" for p in changed[:5]])
+            if len(changed) > 5:
+                lines.append(f"  ... 以及另外 {len(changed) - 5} 个文件")
+        if deleted:
+            lines.append(
+                f"检测到 {len(deleted)} 个已索引文件已被删除，建议清理并重新索引。"
+            )
+            lines.extend([f"  删除: {p}" for p in deleted[:5]])
+            if len(deleted) > 5:
+                lines.append(f"  ... 以及另外 {len(deleted) - 5} 个文件")
+        lines.append(
+            "提示：请使用 'jarvis-rag add <路径>' 重新索引相关文件，以更新向量库与BM25索引。"
+        )
+        PrettyOutput.print("\n".join(lines), OutputType.WARNING)
+
+    def detect_index_changes(self) -> Dict[str, List[str]]:
+        """
+        公共方法：检测索引变更（变更与删除）。
+        返回:
+            {'changed': List[str], 'deleted': List[str]}
+        """
+        return self._detect_changed_or_deleted()
+
+    def _remove_sources_from_manifest(self, sources: List[str]) -> None:
+        """从manifest中移除指定源文件记录并保存。"""
+        if not sources:
+            return
+        manifest = self._load_manifest()
+        removed = 0
+        for src in set(sources):
+            if src in manifest:
+                manifest.pop(src, None)
+                removed += 1
+        if removed > 0:
+            self._save_manifest(manifest)
+            PrettyOutput.print(
+                f"已从索引清单中移除 {removed} 个已删除的源文件记录。", OutputType.INFO
+            )
+
+    def update_index_for_changes(self, changed: List[str], deleted: List[str]) -> None:
+        """
+        公共方法：根据变更与删除列表更新索引。
+        - 对 deleted: 从向量库按 metadata.source 删除
+        - 对 changed: 先删除旧条目，再从源文件重建并添加
+        - 最后：从集合重建BM25索引，更新manifest
+        """
+        changed = list(
+            dict.fromkeys([p for p in (changed or []) if isinstance(p, str)])
+        )
+        deleted = list(
+            dict.fromkeys([p for p in (deleted or []) if isinstance(p, str)])
+        )
+
+        if not changed and not deleted:
+            return
+
+        # 先处理删除
+        delete_errors: List[str] = []
+        for src in deleted:
+            try:
+                self.collection.delete(where={"source": src})  # type: ignore[arg-type]
+            except Exception as e:
+                delete_errors.append(f"删除源 '{src}' 时出错: {e}")
+        if delete_errors:
+            PrettyOutput.print("\n".join(delete_errors), OutputType.WARNING)
+
+        # 再处理变更（重建）
+        docs_to_add: List[Document] = []
+        rebuild_errors: List[str] = []
+        for src in changed:
+            try:
+                # 删除旧条目
+                try:
+                    self.collection.delete(where={"source": src})  # type: ignore[arg-type]
+                except Exception:
+                    pass
+                # 读取源文件内容（作为单文档载入，由 add_documents 进行拆分与嵌入）
+                with open(src, "r", encoding="utf-8", errors="ignore") as f:
+                    content = f.read()
+                docs_to_add.append(
+                    Document(page_content=content, metadata={"source": src})
+                )
+            except Exception as e:
+                rebuild_errors.append(f"重建源 '{src}' 内容时出错: {e}")
+        if rebuild_errors:
+            PrettyOutput.print("\n".join(rebuild_errors), OutputType.WARNING)
+
+        if docs_to_add:
+            try:
+                # 复用现有拆分与嵌入逻辑
+                self.add_documents(docs_to_add)
+            except Exception as e:
+                PrettyOutput.print(f"添加变更文档到索引时出错: {e}", OutputType.ERROR)
+
+        # 重建BM25索引，确保删除后的语料被清理
+        try:
+            all_docs_in_collection = self.collection.get()
+            all_documents = all_docs_in_collection.get("documents") or []
+            self.bm25_corpus = [str(text).split() for text in all_documents if text]
+            self.bm25_index = BM25Okapi(self.bm25_corpus) if self.bm25_corpus else None
+            self._save_bm25_index()
+        except Exception as e:
+            PrettyOutput.print(f"重建BM25索引失败: {e}", OutputType.WARNING)
+
+        # 更新manifest：变更文件更新状态；删除文件从清单中移除
+        try:
+            if changed:
+                self._update_manifest_with_sources(changed)
+            if deleted:
+                self._remove_sources_from_manifest(deleted)
+        except Exception as e:
+            PrettyOutput.print(f"更新索引清单时出错: {e}", OutputType.WARNING)
+
+        PrettyOutput.print(
+            f"索引已更新：变更 {len(changed)} 个，删除 {len(deleted)} 个。",
+            OutputType.SUCCESS,
+        )
 
     def add_documents(
         self, documents: List[Document], chunk_size=1000, chunk_overlap=100
@@ -80,7 +319,10 @@ class ChromaRetriever:
         )
         chunks = text_splitter.split_documents(documents)
 
-        print(f"📄 已将 {len(documents)} 个文档拆分为 {len(chunks)} 个块。")
+        PrettyOutput.print(
+            f"已将 {len(documents)} 个文档拆分为 {len(chunks)} 个块。",
+            OutputType.INFO,
+        )
 
         if not chunks:
             return
@@ -99,19 +341,33 @@ class ChromaRetriever:
             documents=chunk_texts,
             metadatas=cast(Any, metadatas),
         )
-        print(f"✅ 成功将 {len(chunks)} 个块添加到 ChromaDB 集合中。")
+        PrettyOutput.print(
+            f"成功将 {len(chunks)} 个块添加到 ChromaDB 集合中。",
+            OutputType.SUCCESS,
+        )
 
         # 更新并保存BM25索引
         tokenized_chunks = [doc.split() for doc in chunk_texts]
         self.bm25_corpus.extend(tokenized_chunks)
         self.bm25_index = BM25Okapi(self.bm25_corpus)
         self._save_bm25_index()
-
-    def retrieve(self, query: str, n_results: int = 5) -> List[Document]:
+        # 更新索引清单（用于检测源文件变更/删除）
+        source_list = [
+            md.get("source")
+            for md in metadatas
+            if md and isinstance(md.get("source"), str)
+        ]
+        self._update_manifest_with_sources(cast(List[str], source_list))
+
+    def retrieve(
+        self, query: str, n_results: int = 5, use_bm25: bool = True
+    ) -> List[Document]:
         """
         使用向量搜索和BM25执行混合检索，然后使用倒数排序融合（RRF）
         对结果进行融合。
         """
+        # 在检索前检查源文件变更/删除并提醒
+        self._warn_if_sources_changed()
         # 1. 向量搜索 (ChromaDB)
         query_embedding = self.embedding_manager.embed_query(query)
         vector_results = self.collection.query(
@@ -121,7 +377,7 @@ class ChromaRetriever:
 
         # 2. 关键字搜索 (BM25)
         bm25_docs = []
-        if self.bm25_index:
+        if self.bm25_index and use_bm25:
             tokenized_query = query.split()
             doc_scores = self.bm25_index.get_scores(tokenized_query)
 
@@ -144,7 +400,7 @@ class ChromaRetriever:
                 ]
 
             # 按分数排序并取最高结果
-            bm25_results_with_docs.sort(key=lambda x: x[2], reverse=True)
+            bm25_results_with_docs.sort(key=lambda x: x[2], reverse=True)  # type: ignore
 
             for doc_text, metadata, _ in bm25_results_with_docs[: n_results * 2]:
                 bm25_docs.append(Document(page_content=doc_text, metadata=metadata))