jarvis-ai-assistant 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

jarvis/jarvis_rag/cli.py

@@ -1,7 +1,7 @@
 import os
 import sys
 from pathlib import Path
-from typing import Optional, List, Literal, cast
+from typing import Optional, List, Literal, cast, Tuple
 import mimetypes
 
 import pathspec  # type: ignore
@@ -15,6 +15,11 @@ from langchain_core.document_loaders.base import BaseLoader
 from rich.markdown import Markdown
 
 from jarvis.jarvis_utils.utils import init_env
+from jarvis.jarvis_utils.config import (
+    get_rag_embedding_model,
+    get_rag_use_bm25,
+    get_rag_use_rerank,
+)
 
 
 def is_likely_text_file(file_path: Path) -> bool:
@@ -65,9 +70,7 @@ class _CustomPlatformLLM(LLMInterface):
 
     def __init__(self, platform: BasePlatform):
         self.platform = platform
-        print(
-            f"✅ 使用自定义LLM: 平台='{platform.platform_name()}', 模型='{platform.name()}'"
-        )
+        print(f"✅ 使用自定义LLM: 平台='{platform.platform_name()}', 模型='{platform.name()}'")
 
     def generate(self, prompt: str, **kwargs) -> str:
         return self.platform.chat_until_success(prompt)
@@ -91,7 +94,7 @@ def _create_custom_llm(platform_name: str, model_name: str) -> Optional[LLMInter
         return None
 
 
-def _load_ragignore_spec() -> tuple[Optional[pathspec.PathSpec], Optional[Path]]:
+def _load_ragignore_spec() -> Tuple[Optional[pathspec.PathSpec], Optional[Path]]:
     """
     从项目根目录加载忽略模式。
     首先查找 `.jarvis/rag/.ragignore`，如果未找到，则回退到 `.gitignore`。
@@ -140,9 +143,7 @@ def add_documents(
         "-e",
         help="嵌入模型的名称。覆盖全局配置。",
     ),
-    db_path: Optional[Path] = typer.Option(
-        None, "--db-path", help="向量数据库的路径。覆盖全局配置。"
-    ),
+    db_path: Optional[Path] = typer.Option(None, "--db-path", help="向量数据库的路径。覆盖全局配置。"),
     batch_size: int = typer.Option(
         500,
         "--batch-size",
@@ -244,9 +245,7 @@ def add_documents(
             print("❌ 未能成功加载任何文档。")
             raise typer.Exit(code=1)
 
-        print(
-            f"✅ 成功将 {total_docs_added} 个文档的内容添加至集合 '{collection_name}'。"
-        )
+        print(f"✅ 成功将 {total_docs_added} 个文档的内容添加至集合 '{collection_name}'。")
 
     except Exception as e:
         print(f"❌ 发生严重错误: {e}")
@@ -261,9 +260,7 @@ def list_documents(
         "-c",
         help="向量数据库中集合的名称。",
     ),
-    db_path: Optional[Path] = typer.Option(
-        None, "--db-path", help="向量数据库的路径。覆盖全局配置。"
-    ),
+    db_path: Optional[Path] = typer.Option(None, "--db-path", help="向量数据库的路径。覆盖全局配置。"),
 ):
     """列出指定集合中的所有唯一文档。"""
     try:
@@ -272,7 +269,7 @@ def list_documents(
             collection_name=collection_name,
         )
 
-        collection = pipeline.retriever.collection
+        collection = pipeline._get_collection()
         results = collection.get()  # 获取集合中的所有项目
 
         if not results or not results["metadatas"]:
@@ -300,6 +297,63 @@ def list_documents(
         raise typer.Exit(code=1)
 
 
+@app.command("retrieve", help="仅从知识库检索相关文档，不生成答案。")
+def retrieve(
+    question: str = typer.Argument(..., help="要提出的问题。"),
+    collection_name: str = typer.Option(
+        "jarvis_rag_collection",
+        "--collection",
+        "-c",
+        help="向量数据库中集合的名称。",
+    ),
+    embedding_model: Optional[str] = typer.Option(
+        None,
+        "--embedding-model",
+        "-e",
+        help="嵌入模型的名称。覆盖全局配置。",
+    ),
+    db_path: Optional[Path] = typer.Option(None, "--db-path", help="向量数据库的路径。覆盖全局配置。"),
+    n_results: int = typer.Option(5, "--top-n", help="要检索的文档数量。"),
+):
+    """仅从RAG知识库检索文档并打印结果。"""
+    try:
+        # 如果未在命令行中指定，则从配置中加载RAG设置
+        final_embedding_model = embedding_model or get_rag_embedding_model()
+        use_bm25 = get_rag_use_bm25()
+        use_rerank = get_rag_use_rerank()
+
+        pipeline = JarvisRAGPipeline(
+            embedding_model=final_embedding_model,
+            db_path=str(db_path) if db_path else None,
+            collection_name=collection_name,
+            use_bm25=use_bm25,
+            use_rerank=use_rerank,
+        )
+
+        print(f"🤔 正在为问题检索文档: '{question}'")
+        retrieved_docs = pipeline.retrieve_only(question, n_results=n_results)
+
+        if not retrieved_docs:
+            print("ℹ️ 未找到相关文档。")
+            return
+
+        print(f"✅ 成功检索到 {len(retrieved_docs)} 个文档:")
+        from jarvis.jarvis_utils.globals import console
+
+        for i, doc in enumerate(retrieved_docs, 1):
+            source = doc.metadata.get("source", "未知来源")
+            content = doc.page_content
+            panel_title = f"文档 {i} | 来源: {source}"
+            console.print(
+                f"\n[bold magenta]{panel_title}[/bold magenta]"
+            )
+            console.print(Markdown(f"```\n{content}\n```"))
+
+    except Exception as e:
+        print(f"❌ 发生错误: {e}")
+        raise typer.Exit(code=1)
+
+
 @app.command("query", help="向知识库提问。")
 def query(
     question: str = typer.Argument(..., help="要提出的问题。"),
@@ -315,9 +369,7 @@ def query(
         "-e",
         help="嵌入模型的名称。覆盖全局配置。",
     ),
-    db_path: Optional[Path] = typer.Option(
-        None, "--db-path", help="向量数据库的路径。覆盖全局配置。"
-    ),
+    db_path: Optional[Path] = typer.Option(None, "--db-path", help="向量数据库的路径。覆盖全局配置。"),
     platform: Optional[str] = typer.Option(
         None,
         "--platform",
@@ -341,11 +393,18 @@ def query(
         if (platform or model) and not custom_llm:
             raise typer.Exit(code=1)
 
+        # 如果未在命令行中指定，则从配置中加载RAG设置
+        final_embedding_model = embedding_model or get_rag_embedding_model()
+        use_bm25 = get_rag_use_bm25()
+        use_rerank = get_rag_use_rerank()
+
         pipeline = JarvisRAGPipeline(
             llm=custom_llm,
-            embedding_model=embedding_model,
+            embedding_model=final_embedding_model,
             db_path=str(db_path) if db_path else None,
             collection_name=collection_name,
+            use_bm25=use_bm25,
+            use_rerank=use_rerank,
         )
 
         print(f"🤔 正在查询: '{question}'")
@@ -373,10 +432,7 @@ except ImportError:
 
 def _check_rag_dependencies():
     if not _RAG_INSTALLED:
-        print(
-            "❌ RAG依赖项未安装。"
-            "请运行 'pip install \"jarvis-ai-assistant[rag]\"' 来使用此命令。"
-        )
+        print("❌ RAG依赖项未安装。" "请运行 'pip install \"jarvis-ai-assistant[rag]\"' 来使用此命令。")
         raise typer.Exit(code=1)
 
 

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,55 @@ 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

@@ -30,6 +30,8 @@ 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,
     ):
         """
         初始化RAG管道。
@@ -40,6 +42,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 +60,87 @@ 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
+
+        # 延迟加载的组件
+        self._embedding_manager: Optional[EmbeddingManager] = None
+        self._retriever: Optional[ChromaRetriever] = None
+        self._reranker: Optional[Reranker] = None
+        self._query_rewriter: Optional[QueryRewriter] = None
+
+        print("✅ JarvisRAGPipeline 初始化成功 (模型按需加载).")
+
+    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 add_documents(self, documents: List[Document]):
         """
@@ -80,24 +149,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）来获取最新、最准确的信息**。
 
         参考上下文:
         ---
@@ -122,13 +188,15 @@ class JarvisRAGPipeline:
             由LLM生成的答案。
         """
         # 1. 将原始查询重写为多个查询
-        rewritten_queries = self.query_rewriter.rewrite(query_text)
+        rewritten_queries = self._get_query_rewriter().rewrite(query_text)
 
         # 2. 为每个重写的查询检索初始候选文档
         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 +207,13 @@ 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:
+            print(f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）...")
+            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 "我在提供的文档中找不到任何相关信息来回答您的问题。"
@@ -166,9 +235,49 @@ class JarvisRAGPipeline:
 
         # 4. 创建最终提示并生成答案
         # 我们使用原始的query_text作为给LLM的最终提示
-        prompt = self._create_prompt(query_text, retrieved_docs, sources)
+        prompt = self._create_prompt(query_text, retrieved_docs)
 
         print("🤖 正在从LLM生成答案...")
         answer = self.llm.generate(prompt)
 
         return answer
+
+    def retrieve_only(self, query_text: str, n_results: int = 5) -> List[Document]:
+        """
+        仅执行检索和重排，不生成答案。
+
+        参数:
+            query_text: 用户的原始问题。
+            n_results: 要检索的最终相关块的数量。
+
+        返回:
+            检索到的文档列表。
+        """
+        # 1. 重写查询
+        rewritten_queries = self._get_query_rewriter().rewrite(query_text)
+
+        # 2. 检索候选文档
+        all_candidate_docs = []
+        for q in rewritten_queries:
+            print(f"🔍 正在为查询变体 '{q}' 进行混合检索...")
+            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:
+            print(f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排...")
+            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/retriever.py

@@ -39,9 +39,7 @@ class ChromaRetriever:
         self.collection = self.client.get_or_create_collection(
             name=self.collection_name
         )
-        print(
-            f"✅ ChromaDB 客户端已在 '{db_path}' 初始化，集合为 '{collection_name}'。"
-        )
+        print(f"✅ ChromaDB 客户端已在 '{db_path}' 初始化，集合为 '{collection_name}'。")
 
         # BM25索引设置
         self.bm25_index_path = os.path.join(self.db_path, f"{collection_name}_bm25.pkl")
@@ -107,7 +105,9 @@ class ChromaRetriever:
         self.bm25_index = BM25Okapi(self.bm25_corpus)
         self._save_bm25_index()
 
-    def retrieve(self, query: str, n_results: int = 5) -> List[Document]:
+    def retrieve(
+        self, query: str, n_results: int = 5, use_bm25: bool = True
+    ) -> List[Document]:
         """
         使用向量搜索和BM25执行混合检索，然后使用倒数排序融合（RRF）
         对结果进行融合。
@@ -121,7 +121,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 +144,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))

jarvis/jarvis_smart_shell/main.py

@@ -170,8 +170,8 @@ def process_request(request: str) -> Optional[str]:
 4. 多个命令用&&连接
 
 # 示例
-输入: "查找Python文件"
-输出: find . -name "*.py"
+输入: "显示当前目录内容"
+输出: ls -la
 """
         model.set_system_prompt(system_message)
 

jarvis/jarvis_stats/__init__.py

@@ -0,0 +1,13 @@
+"""
+Jarvis统计模块
+
+提供指标统计、数据持久化、可视化展示等功能
+"""
+
+from jarvis.jarvis_stats.stats import StatsManager
+from jarvis.jarvis_stats.storage import StatsStorage
+from jarvis.jarvis_stats.visualizer import StatsVisualizer
+
+__all__ = ["StatsManager", "StatsStorage", "StatsVisualizer"]
+
+__version__ = "1.0.0"