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

jarvis/jarvis_rag/query_rewriter.py

@@ -1,5 +1,6 @@
 from typing import List
 from .llm_interface import LLMInterface
+from jarvis.jarvis_utils.output import PrettyOutput, OutputType
 
 
 class QueryRewriter:
@@ -20,20 +21,29 @@ class QueryRewriter:
     def _create_prompt_template(self) -> str:
         """为多查询重写任务创建提示模板。"""
         return """
-你是一个精通检索的AI助手。你的任务是将以下这个单一的用户问题，从不同角度改写成 3 个不同的、但语义上相关的搜索查询。这有助于在知识库中进行更全面的搜索。
+你是一个精通检索和语言的AI助手。你的任务是将以下这个单一的用户问题，改写为几个语义相关但表达方式不同的搜索查询，并提供英文翻译。这有助于在多语言知识库中进行更全面的搜索。
 
 请遵循以下原则：
-1.  **多样性**：生成的查询应尝试使用不同的关键词和表述方式。
-2.  **保留核心意图**：所有查询都必须围绕原始问题的核心意图。
-3.  **简洁性**：每个查询都应该是独立的、可以直接用于搜索的短语或问题。
-4.  **格式要求**：请直接输出 3 个查询，每个查询占一行，用换行符分隔。不要添加任何编号、前缀或解释。
+1.  **保留核心意图**: 所有查询都必须围绕原始问题的核心意图。
+2.  **查询类型**:
+    - **同义词/相关术语查询**: 使用原始语言，通过替换同义词或相关术语来生成1-2个新的查询。
+    - **英文翻译查询**: 将原始问题翻译成一个简洁的英文搜索查询。
+3.  **简洁性**: 每个查询都应该是独立的、可以直接用于搜索的短语或问题。
+4.  **严格格式要求**: 你必须将所有重写后的查询放置在 `<REWRITE>` 和 `</REWRITE>` 标签之间。每个查询占一行。不要在标签内外添加任何编号、前缀或解释。
+
+示例输出格式:
+<REWRITE>
+使用不同表述的中文查询
+另一个中文查询
+English version of the query
+</REWRITE>
 
 原始问题:
 ---
 {query}
 ---
 
-3个改写后的查询 (每行一个):
+请将改写后的查询包裹在 `<REWRITE>` 标签内:
 """
 
     def rewrite(self, query: str) -> List[str]:
@@ -47,16 +57,57 @@ class QueryRewriter:
             一个经过重写、搜索优化的查询列表。
         """
         prompt = self.rewrite_prompt_template.format(query=query)
-        print(f"✍️  正在将原始查询重写为多个搜索查询...")
+        PrettyOutput.print(
+            "正在将原始查询重写为多个搜索查询...",
+            output_type=OutputType.INFO,
+            timestamp=False,
+        )
+
+        import re
+
+        max_retries = 3
+        attempts = 0
+        rewritten_queries = []
+        response_text = ""
+
+        while attempts < max_retries:
+            attempts += 1
+            response_text = self.llm.generate(prompt)
+            match = re.search(r"<REWRITE>(.*?)</REWRITE>", response_text, re.DOTALL)
+
+            if match:
+                content = match.group(1).strip()
+                rewritten_queries = [
+                    line.strip() for line in content.split("\n") if line.strip()
+                ]
+                PrettyOutput.print(
+                    f"成功从LLM响应中提取到内容 (尝试 {attempts}/{max_retries})。",
+                    output_type=OutputType.SUCCESS,
+                    timestamp=False,
+                )
+                break  # 提取成功，退出循环
+            else:
+                PrettyOutput.print(
+                    f"未能从LLM响应中提取内容。正在重试... ({attempts}/{max_retries})",
+                    output_type=OutputType.WARNING,
+                    timestamp=False,
+                )
 
-        response_text = self.llm.generate(prompt)
-        rewritten_queries = [
-            line.strip() for line in response_text.strip().split("\n") if line.strip()
-        ]
+        # 如果所有重试都失败，则跳过重写步骤
+        if not rewritten_queries:
+            PrettyOutput.print(
+                "所有重试均失败。跳过查询重写，将仅使用原始查询。",
+                output_type=OutputType.ERROR,
+                timestamp=False,
+            )
 
         # 同时包含原始查询以保证鲁棒性
         if query not in rewritten_queries:
             rewritten_queries.insert(0, query)
 
-        print(f"✅ 生成了 {len(rewritten_queries)} 个查询变体。")
+        PrettyOutput.print(
+            f"生成了 {len(rewritten_queries)} 个查询变体。",
+            output_type=OutputType.SUCCESS,
+            timestamp=False,
+        )
         return rewritten_queries

jarvis/jarvis_rag/rag_pipeline.py

@@ -8,12 +8,14 @@ from .llm_interface import JarvisPlatform_LLM, LLMInterface, ToolAgent_LLM
 from .query_rewriter import QueryRewriter
 from .reranker import Reranker
 from .retriever import ChromaRetriever
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 from jarvis.jarvis_utils.config import (
     get_rag_embedding_model,
     get_rag_rerank_model,
     get_rag_vector_db_path,
     get_rag_embedding_cache_path,
 )
+from jarvis.jarvis_utils.utils import get_yes_no
 
 
 class JarvisRAGPipeline:
@@ -30,6 +32,9 @@ class JarvisRAGPipeline:
         embedding_model: Optional[str] = None,
         db_path: Optional[str] = None,
         collection_name: str = "jarvis_rag_collection",
+        use_bm25: bool = True,
+        use_rerank: bool = True,
+        use_query_rewrite: bool = True,
     ):
         """
         初始化RAG管道。
@@ -40,6 +45,8 @@ class JarvisRAGPipeline:
             embedding_model: 嵌入模型的名称。如果为None，则使用配置值。
             db_path: 持久化向量数据库的路径。如果为None，则使用配置值。
             collection_name: 向量数据库中集合的名称。
+            use_bm25: 是否在检索中使用BM25。
+            use_rerank: 是否在检索后使用重排器。
         """
         # 确定嵌入模型以隔离数据路径
         model_name = embedding_model or get_rag_embedding_model()
@@ -56,22 +63,127 @@ class JarvisRAGPipeline:
             get_rag_embedding_cache_path(), sanitized_model_name
         )
 
-        self.embedding_manager = EmbeddingManager(
-            model_name=model_name,
-            cache_dir=_final_cache_path,
+        # 存储初始化参数以供延迟加载
+        self.llm = llm if llm is not None else ToolAgent_LLM()
+        self.embedding_model_name = embedding_model or get_rag_embedding_model()
+        self.db_path = db_path
+        self.collection_name = collection_name
+        self.use_bm25 = use_bm25
+        self.use_rerank = use_rerank
+        # 查询重写开关（默认开启，可由CLI控制）
+        self.use_query_rewrite = use_query_rewrite
+
+        # 延迟加载的组件
+        self._embedding_manager: Optional[EmbeddingManager] = None
+        self._retriever: Optional[ChromaRetriever] = None
+        self._reranker: Optional[Reranker] = None
+        self._query_rewriter: Optional[QueryRewriter] = None
+
+        PrettyOutput.print(
+            "JarvisRAGPipeline 初始化成功 (模型按需加载).", OutputType.SUCCESS
+        )
+
+    def _get_embedding_manager(self) -> EmbeddingManager:
+        if self._embedding_manager is None:
+            sanitized_model_name = self.embedding_model_name.replace("/", "_").replace(
+                "\\", "_"
+            )
+            _final_cache_path = os.path.join(
+                get_rag_embedding_cache_path(), sanitized_model_name
+            )
+            self._embedding_manager = EmbeddingManager(
+                model_name=self.embedding_model_name,
+                cache_dir=_final_cache_path,
+            )
+        return self._embedding_manager
+
+    def _get_retriever(self) -> ChromaRetriever:
+        if self._retriever is None:
+            sanitized_model_name = self.embedding_model_name.replace("/", "_").replace(
+                "\\", "_"
+            )
+            _final_db_path = (
+                str(self.db_path)
+                if self.db_path
+                else os.path.join(get_rag_vector_db_path(), sanitized_model_name)
+            )
+            self._retriever = ChromaRetriever(
+                embedding_manager=self._get_embedding_manager(),
+                db_path=_final_db_path,
+                collection_name=self.collection_name,
+            )
+        return self._retriever
+
+    def _get_collection(self):
+        """
+        在不加载嵌入模型的情况下，直接获取并返回Chroma集合对象。
+        这对于仅需要访问集合元数据（如列出文档）而无需嵌入功能的操作非常有用。
+        """
+        # 为了避免初始化embedding_manager，我们直接构建db_path
+        if self._retriever:
+            return self._retriever.collection
+
+        sanitized_model_name = self.embedding_model_name.replace("/", "_").replace(
+            "\\", "_"
         )
-        self.retriever = ChromaRetriever(
-            embedding_manager=self.embedding_manager,
-            db_path=_final_db_path,
-            collection_name=collection_name,
+        _final_db_path = (
+            str(self.db_path)
+            if self.db_path
+            else os.path.join(get_rag_vector_db_path(), sanitized_model_name)
         )
-        # 除非提供了特定的LLM，否则默认为ToolAgent_LLM
-        self.llm = llm if llm is not None else ToolAgent_LLM()
-        self.reranker = Reranker(model_name=get_rag_rerank_model())
-        # 使用标准LLM执行查询重写任务，而不是代理
-        self.query_rewriter = QueryRewriter(JarvisPlatform_LLM())
 
-        print("✅ JarvisRAGPipeline 初始化成功。")
+        # 直接创建ChromaRetriever所使用的chroma_client，但绕过embedding_manager
+        import chromadb
+
+        chroma_client = chromadb.PersistentClient(path=_final_db_path)
+        return chroma_client.get_collection(name=self.collection_name)
+
+    def _get_reranker(self) -> Reranker:
+        if self._reranker is None:
+            self._reranker = Reranker(model_name=get_rag_rerank_model())
+        return self._reranker
+
+    def _get_query_rewriter(self) -> QueryRewriter:
+        if self._query_rewriter is None:
+            # 使用标准LLM执行查询重写任务，而不是代理
+            self._query_rewriter = QueryRewriter(JarvisPlatform_LLM())
+        return self._query_rewriter
+
+    def _pre_search_update_index_if_needed(self) -> None:
+        """
+        在重写query之前执行：
+        - 检测索引变更（变更/删除）
+        - 询问用户是否立即更新索引
+        - 如确认，则执行增量更新并重建BM25
+        """
+        try:
+            retriever = self._get_retriever()
+            result = retriever.detect_index_changes()
+            changed = result.get("changed", [])
+            deleted = result.get("deleted", [])
+            if not changed and not deleted:
+                return
+            # 打印摘要
+            # 先拼接列表信息再统一打印，避免循环中逐条打印
+            lines = [
+                f"检测到索引可能不一致：变更 {len(changed)} 个，删除 {len(deleted)} 个。"
+            ]
+            if changed:
+                lines.extend([f"  变更: {p}" for p in changed[:3]])
+            if deleted:
+                lines.extend([f"  删除: {p}" for p in deleted[:3]])
+            PrettyOutput.print("\n".join(lines), OutputType.WARNING)
+            # 询问用户
+            if get_yes_no(
+                "检测到索引变更，是否现在更新索引后再开始检索？", default=True
+            ):
+                retriever.update_index_for_changes(changed, deleted)
+            else:
+                PrettyOutput.print(
+                    "已跳过索引更新，将直接使用当前索引进行检索。", OutputType.INFO
+                )
+        except Exception as e:
+            PrettyOutput.print(f"检索前索引检查失败：{e}", OutputType.WARNING)
 
     def add_documents(self, documents: List[Document]):
         """
@@ -80,24 +192,21 @@ class JarvisRAGPipeline:
         参数:
             documents: 要添加的LangChain文档对象列表。
         """
-        self.retriever.add_documents(documents)
+        self._get_retriever().add_documents(documents)
 
-    def _create_prompt(
-        self, query: str, context_docs: List[Document], source_files: List[str]
-    ) -> str:
+    def _create_prompt(self, query: str, context_docs: List[Document]) -> str:
         """为LLM或代理创建最终的提示。"""
-        context = "\n\n".join([doc.page_content for doc in context_docs])
-        sources_text = "\n".join([f"- {source}" for source in source_files])
+        context_details = []
+        for doc in context_docs:
+            source = doc.metadata.get("source", "未知来源")
+            content = doc.page_content
+            context_details.append(f"来源: {source}\n\n---\n{content}\n---")
+        context = "\n\n".join(context_details)
 
         prompt_template = f"""
         你是一个专家助手。请根据用户的问题，结合下面提供的参考信息来回答。
 
-        **重要**: 提供的上下文和文件列表**仅供参考**，可能不完整或已过时。在回答前，你应该**优先使用工具（如 read_code）来获取最新、最准确的信息**。
-
-        参考文件列表:
-        ---
-        {sources_text}
-        ---
+        **重要**: 提供的上下文**仅供参考**，可能不完整或已过时。在回答前，你应该**优先使用工具（如 read_code）来获取最新、最准确的信息**。
 
         参考上下文:
         ---
@@ -121,14 +230,28 @@ class JarvisRAGPipeline:
         返回:
             由LLM生成的答案。
         """
-        # 1. 将原始查询重写为多个查询
-        rewritten_queries = self.query_rewriter.rewrite(query_text)
+        # 0. 检测索引变更并可选更新（在重写query之前）
+        self._pre_search_update_index_if_needed()
+        # 1. 将原始查询重写为多个查询（可配置）
+        if self.use_query_rewrite:
+            rewritten_queries = self._get_query_rewriter().rewrite(query_text)
+        else:
+            PrettyOutput.print(
+                "已关闭查询重写，将直接使用原始查询进行检索。",
+                OutputType.INFO,
+            )
+            rewritten_queries = [query_text]
 
         # 2. 为每个重写的查询检索初始候选文档
+        PrettyOutput.print(
+            "将为以下查询变体进行混合检索:\n" + "\n".join([f"  - {q}" for q in rewritten_queries]),
+            OutputType.INFO,
+        )
         all_candidate_docs = []
         for q in rewritten_queries:
-            print(f"🔍 正在为查询变体 '{q}' 进行混合检索...")
-            candidates = self.retriever.retrieve(q, n_results=n_results * 2)
+            candidates = self._get_retriever().retrieve(
+                q, n_results=n_results * 2, use_bm25=self.use_bm25
+            )
             all_candidate_docs.extend(candidates)
 
         # 对候选文档进行去重
@@ -139,12 +262,16 @@ class JarvisRAGPipeline:
             return "我在提供的文档中找不到任何相关信息来回答您的问题。"
 
         # 3. 根据*原始*查询对统一的候选池进行重排
-        print(
-            f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）..."
-        )
-        retrieved_docs = self.reranker.rerank(
-            query_text, unique_candidate_docs, top_n=n_results
-        )
+        if self.use_rerank:
+            PrettyOutput.print(
+                f"正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）...",
+                OutputType.INFO,
+            )
+            retrieved_docs = self._get_reranker().rerank(
+                query_text, unique_candidate_docs, top_n=n_results
+            )
+        else:
+            retrieved_docs = unique_candidate_docs[:n_results]
 
         if not retrieved_docs:
             return "我在提供的文档中找不到任何相关信息来回答您的问题。"
@@ -160,15 +287,70 @@ class JarvisRAGPipeline:
             )
         )
         if sources:
-            print(f"📚 根据以下文档回答:")
-            for source in sources:
-                print(f"  - {source}")
+            # 合并来源列表后一次性打印，避免多次加框
+            lines = ["根据以下文档回答:"] + [f"  - {source}" for source in sources]
+            PrettyOutput.print("\n".join(lines), OutputType.INFO)
 
         # 4. 创建最终提示并生成答案
         # 我们使用原始的query_text作为给LLM的最终提示
-        prompt = self._create_prompt(query_text, retrieved_docs, sources)
+        prompt = self._create_prompt(query_text, retrieved_docs)
 
-        print("🤖 正在从LLM生成答案...")
+        PrettyOutput.print("正在从LLM生成答案...", OutputType.INFO)
         answer = self.llm.generate(prompt)
 
         return answer
+
+    def retrieve_only(self, query_text: str, n_results: int = 5) -> List[Document]:
+        """
+        仅执行检索和重排，不生成答案。
+
+        参数:
+            query_text: 用户的原始问题。
+            n_results: 要检索的最终相关块的数量。
+
+        返回:
+            检索到的文档列表。
+        """
+        # 0. 检测索引变更并可选更新（在重写query之前）
+        self._pre_search_update_index_if_needed()
+        # 1. 重写查询（可配置）
+        if self.use_query_rewrite:
+            rewritten_queries = self._get_query_rewriter().rewrite(query_text)
+        else:
+            PrettyOutput.print(
+                "已关闭查询重写，将直接使用原始查询进行检索。",
+                OutputType.INFO,
+            )
+            rewritten_queries = [query_text]
+
+        # 2. 检索候选文档
+        PrettyOutput.print(
+            "将为以下查询变体进行混合检索:\n" + "\n".join([f"  - {q}" for q in rewritten_queries]),
+            OutputType.INFO,
+        )
+        all_candidate_docs = []
+        for q in rewritten_queries:
+            candidates = self._get_retriever().retrieve(
+                q, n_results=n_results * 2, use_bm25=self.use_bm25
+            )
+            all_candidate_docs.extend(candidates)
+
+        unique_docs_dict = {doc.page_content: doc for doc in all_candidate_docs}
+        unique_candidate_docs = list(unique_docs_dict.values())
+
+        if not unique_candidate_docs:
+            return []
+
+        # 3. 重排
+        if self.use_rerank:
+            PrettyOutput.print(
+                f"正在对 {len(unique_candidate_docs)} 个候选文档进行重排...",
+                OutputType.INFO,
+            )
+            retrieved_docs = self._get_reranker().rerank(
+                query_text, unique_candidate_docs, top_n=n_results
+            )
+        else:
+            retrieved_docs = unique_candidate_docs[:n_results]
+
+        return retrieved_docs

jarvis/jarvis_rag/reranker.py

@@ -1,9 +1,12 @@
 from typing import List
+import os
 
 from langchain.docstore.document import Document
 from sentence_transformers.cross_encoder import (  # type: ignore
     CrossEncoder,
 )
+from huggingface_hub import snapshot_download
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 
 
 class Reranker:
@@ -19,9 +22,29 @@ class Reranker:
         参数:
             model_name (str): 要使用的Cross-Encoder模型的名称。
         """
-        print(f"🔍 正在初始化重排模型: {model_name}...")
-        self.model = CrossEncoder(model_name)
-        print("✅ 重排模型初始化成功。")
+        PrettyOutput.print(f"正在初始化重排模型: {model_name}...", OutputType.INFO)
+        try:
+            local_dir = None
+
+            if os.path.isdir(model_name):
+                self.model = CrossEncoder(model_name)
+                PrettyOutput.print("重排模型初始化成功。", OutputType.SUCCESS)
+                return
+            try:
+                # Prefer local cache; avoid any network access
+                local_dir = snapshot_download(repo_id=model_name, local_files_only=True)
+            except Exception:
+                local_dir = None
+
+            if local_dir:
+                self.model = CrossEncoder(local_dir)
+            else:
+                self.model = CrossEncoder(model_name)
+
+            PrettyOutput.print("重排模型初始化成功。", OutputType.SUCCESS)
+        except Exception as e:
+            PrettyOutput.print(f"初始化重排模型失败: {e}", OutputType.ERROR)
+            raise
 
     def rerank(
         self, query: str, documents: List[Document], top_n: int = 5