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

jarvis/jarvis_rag/cli.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import Optional, List, Literal, cast
 import mimetypes
 
+import pathspec
 import typer
 from langchain.docstore.document import Document
 from langchain_community.document_loaders import (
@@ -18,29 +19,29 @@ from jarvis.jarvis_utils.utils import init_env
 
 def is_likely_text_file(file_path: Path) -> bool:
     """
-    Checks if a file is likely to be a text file by reading its beginning.
-    Avoids loading large binary files into memory.
+    通过读取文件开头部分，检查文件是否可能为文本文件。
+    此方法可以避免将大型二进制文件加载到内存中。
     """
     try:
-        # Heuristic 1: Check MIME type if available
+        # 启发式方法1：检查MIME类型（如果可用）
         mime_type, _ = mimetypes.guess_type(file_path)
         if mime_type and mime_type.startswith("text/"):
             return True
         if mime_type and any(x in mime_type for x in ["json", "xml", "javascript"]):
             return True
 
-        # Heuristic 2: Check for null bytes in the first few KB
+        # 启发式方法2：检查文件的前几KB中是否包含空字节
         with open(file_path, "rb") as f:
-            chunk = f.read(4096)  # Read first 4KB
+            chunk = f.read(4096)  # 读取前4KB
             if b"\x00" in chunk:
-                return False  # Null bytes are a strong indicator of a binary file
+                return False  # 空字节是二进制文件的强指示符
         return True
     except Exception:
         return False
 
 
-# Ensure the project root is in the Python path to allow absolute imports
-# This makes the script runnable as a module.
+# 确保项目根目录在Python路径中，以允许绝对导入
+# 这使得脚本可以作为模块运行。
 _project_root = os.path.abspath(
     os.path.join(os.path.dirname(__file__), "..", "..", "..")
 )
@@ -54,13 +55,13 @@ from jarvis.jarvis_rag.rag_pipeline import JarvisRAGPipeline
 
 app = typer.Typer(
     name="jarvis-rag",
-    help="A command-line tool to interact with the Jarvis RAG framework.",
+    help="一个与Jarvis RAG框架交互的命令行工具。",
     add_completion=False,
 )
 
 
 class _CustomPlatformLLM(LLMInterface):
-    """A simple wrapper to make a BasePlatform instance compatible with LLMInterface."""
+    """一个简单的包装器，使BasePlatform实例与LLMInterface兼容。"""
 
     def __init__(self, platform: BasePlatform):
         self.platform = platform
@@ -73,7 +74,7 @@ class _CustomPlatformLLM(LLMInterface):
 
 
 def _create_custom_llm(platform_name: str, model_name: str) -> Optional[LLMInterface]:
-    """Creates an LLM interface from a specific platform and model."""
+    """从指定的平台和模型创建LLM接口。"""
     if not platform_name or not model_name:
         return None
     try:
@@ -90,36 +91,70 @@ def _create_custom_llm(platform_name: str, model_name: str) -> Optional[LLMInter
         return None
 
 
+def _load_ragignore_spec() -> tuple[Optional[pathspec.PathSpec], Optional[Path]]:
+    """
+    从项目根目录加载忽略模式。
+    首先查找 `.jarvis/rag/.ragignore`，如果未找到，则回退到 `.gitignore`。
+    """
+    project_root_path = Path(_project_root)
+    ragignore_file = project_root_path / ".jarvis" / "rag" / ".ragignore"
+    gitignore_file = project_root_path / ".gitignore"
+
+    ignore_file_to_use = None
+    if ragignore_file.is_file():
+        ignore_file_to_use = ragignore_file
+    elif gitignore_file.is_file():
+        ignore_file_to_use = gitignore_file
+
+    if ignore_file_to_use:
+        try:
+            with open(ignore_file_to_use, "r", encoding="utf-8") as f:
+                patterns = f.read().splitlines()
+            spec = pathspec.PathSpec.from_lines("gitwildmatch", patterns)
+            print(f"✅ 加载忽略规则: {ignore_file_to_use}")
+            return spec, project_root_path
+        except Exception as e:
+            print(f"⚠️ 加载 {ignore_file_to_use.name} 文件失败: {e}")
+
+    return None, None
+
+
 @app.command(
     "add",
-    help="Add documents from files, directories, or glob patterns (e.g., 'src/**/*.py').",
+    help="从文件、目录或glob模式（例如 'src/**/*.py'）添加文档。",
 )
 def add_documents(
     paths: List[Path] = typer.Argument(
         ...,
-        help="File/directory paths or glob patterns. Shell expansion is supported.",
+        help="文件/目录路径或glob模式。支持Shell扩展。",
     ),
     collection_name: str = typer.Option(
         "jarvis_rag_collection",
         "--collection",
         "-c",
-        help="Name of the collection in the vector database.",
+        help="向量数据库中集合的名称。",
     ),
-    embedding_mode: Optional[str] = typer.Option(
+    embedding_model: Optional[str] = typer.Option(
         None,
-        "--embedding-mode",
+        "--embedding-model",
         "-e",
-        help="Embedding mode ('performance' or 'accuracy'). Overrides global config.",
+        help="嵌入模型的名称。覆盖全局配置。",
     ),
     db_path: Optional[Path] = typer.Option(
-        None, "--db-path", help="Path to the vector database. Overrides global config."
+        None, "--db-path", help="向量数据库的路径。覆盖全局配置。"
+    ),
+    batch_size: int = typer.Option(
+        500,
+        "--batch-size",
+        "-b",
+        help="单个批次中要处理的文档数。",
     ),
 ):
-    """Adds documents to the RAG knowledge base from various sources."""
+    """从不同来源向RAG知识库添加文档。"""
     files_to_process = set()
 
     for path_str in paths:
-        # Typer with List[Path] might not expand globs, so we do it manually
+        # Typer的List[Path]可能不会扩展glob，所以我们手动处理
         from glob import glob
 
         expanded_paths = glob(str(path_str), recursive=True)
@@ -141,59 +176,96 @@ def add_documents(
                     print(f"⚠️ 跳过可能的二进制文件: {path}")
 
     if not files_to_process:
-        print(f"⚠️ 在指定路径中未找到任何文本文件。")
+        print("⚠️ 在指定路径中未找到任何文本文件。")
+        return
+
+    # 使用 .ragignore 过滤文件
+    ragignore_spec, ragignore_root = _load_ragignore_spec()
+    if ragignore_spec and ragignore_root:
+        initial_count = len(files_to_process)
+        retained_files = set()
+        for file_path in files_to_process:
+            try:
+                # 将文件路径解析为绝对路径以确保正确比较
+                resolved_path = file_path.resolve()
+                relative_path = str(resolved_path.relative_to(ragignore_root))
+                if not ragignore_spec.match_file(relative_path):
+                    retained_files.add(file_path)
+            except ValueError:
+                # 文件不在项目根目录下，保留它
+                retained_files.add(file_path)
+
+        ignored_count = initial_count - len(retained_files)
+        if ignored_count > 0:
+            print(f"ℹ️ 根据 .ragignore 规则过滤掉 {ignored_count} 个文件。")
+        files_to_process = retained_files
+
+    if not files_to_process:
+        print("⚠️ 所有找到的文本文件都被忽略规则过滤掉了。")
         return
 
     print(f"✅ 发现 {len(files_to_process)} 个独立文件待处理。")
 
     try:
         pipeline = JarvisRAGPipeline(
-            embedding_mode=cast(
-                Optional[Literal["performance", "accuracy"]], embedding_mode
-            ),
+            embedding_model=embedding_model,
             db_path=str(db_path) if db_path else None,
             collection_name=collection_name,
         )
 
-        docs: List[Document] = []
+        docs_batch: List[Document] = []
+        total_docs_added = 0
         loader: BaseLoader
-        for file_path in sorted(list(files_to_process)):
+
+        sorted_files = sorted(list(files_to_process))
+        total_files = len(sorted_files)
+
+        for i, file_path in enumerate(sorted_files):
             try:
                 if file_path.suffix.lower() == ".md":
                     loader = UnstructuredMarkdownLoader(str(file_path))
-                else:  # Default to TextLoader for .txt and all code files
+                else:  # 对.txt和所有代码文件默认使用TextLoader
                     loader = TextLoader(str(file_path), encoding="utf-8")
 
-                docs.extend(loader.load())
-                print(f"✅ 已加载: {file_path}")
+                docs_batch.extend(loader.load())
+                print(f"✅ 已加载: {file_path} (文件 {i + 1}/{total_files})")
             except Exception as e:
                 print(f"⚠️ 加载失败 {file_path}: {e}")
 
-        if not docs:
+            # 当批处理已满或是最后一个文件时处理批处理
+            if docs_batch and (len(docs_batch) >= batch_size or (i + 1) == total_files):
+                print(f"⚙️ 正在处理批次，包含 {len(docs_batch)} 个文档...")
+                pipeline.add_documents(docs_batch)
+                total_docs_added += len(docs_batch)
+                print(f"✅ 成功添加 {len(docs_batch)} 个文档。")
+                docs_batch = []  # 清空批处理
+
+        if total_docs_added == 0:
             print("❌ 未能成功加载任何文档。")
             raise typer.Exit(code=1)
 
-        pipeline.add_documents(docs)
-        print(f"✅ 成功将 {len(docs)} 个文档的内容添加至集合 '{collection_name}'。")
+        print(
+            f"✅ 成功将 {total_docs_added} 个文档的内容添加至集合 '{collection_name}'。"
+        )
 
     except Exception as e:
         print(f"❌ 发生严重错误: {e}")
         raise typer.Exit(code=1)
 
 
-@app.command("list-docs", help="List all unique documents in the knowledge base.")
+@app.command("list-docs", help="列出知识库中所有唯一的文档。")
 def list_documents(
     collection_name: str = typer.Option(
         "jarvis_rag_collection",
         "--collection",
         "-c",
-        help="Name of the collection in the vector database.",
+        help="向量数据库中集合的名称。",
     ),
     db_path: Optional[Path] = typer.Option(
-        None, "--db-path", help="Path to the vector database. Overrides global config."
+        None, "--db-path", help="向量数据库的路径。覆盖全局配置。"
     ),
 ):
-    """Lists all unique documents in the specified collection."""
+    """列出指定集合中的所有唯一文档。"""
     try:
         pipeline = JarvisRAGPipeline(
             db_path=str(db_path) if db_path else None,
@@ -201,13 +273,13 @@ def list_documents(
         )
 
         collection = pipeline.retriever.collection
-        results = collection.get()  # Get all items in the collection
+        results = collection.get()  # 获取集合中的所有项目
 
         if not results or not results["metadatas"]:
             print("ℹ️ 知识库中没有找到任何文档。")
             return
 
-        # Extract unique source file paths from metadata
+        # 从元数据中提取唯一的源文件路径
         sources = set()
         for metadata in results["metadatas"]:
             if metadata:
@@ -228,38 +300,38 @@ def list_documents(
         raise typer.Exit(code=1)
 
 
-@app.command("query", help="Ask a question to the knowledge base.")
+@app.command("query", help="向知识库提问。")
 def query(
-    question: str = typer.Argument(..., help="The question to ask."),
+    question: str = typer.Argument(..., help="要提出的问题。"),
     collection_name: str = typer.Option(
         "jarvis_rag_collection",
         "--collection",
         "-c",
-        help="Name of the collection in the vector database.",
+        help="向量数据库中集合的名称。",
     ),
-    embedding_mode: Optional[str] = typer.Option(
+    embedding_model: Optional[str] = typer.Option(
         None,
-        "--embedding-mode",
+        "--embedding-model",
         "-e",
-        help="Embedding mode ('performance' or 'accuracy'). Overrides global config.",
+        help="嵌入模型的名称。覆盖全局配置。",
     ),
     db_path: Optional[Path] = typer.Option(
-        None, "--db-path", help="Path to the vector database. Overrides global config."
+        None, "--db-path", help="向量数据库的路径。覆盖全局配置。"
     ),
     platform: Optional[str] = typer.Option(
         None,
         "--platform",
         "-p",
-        help="Specify a platform name for the LLM. Overrides the default thinking model.",
+        help="为LLM指定平台名称。覆盖默认的思考模型。",
     ),
     model: Optional[str] = typer.Option(
         None,
         "--model",
         "-m",
-        help="Specify a model name for the LLM. Requires --platform.",
+        help="为LLM指定模型名称。需要 --platform。",
     ),
 ):
-    """Queries the RAG knowledge base and prints the answer."""
+    """查询RAG知识库并打印答案。"""
     if model and not platform:
         print("❌ 错误: --model 需要指定 --platform。")
         raise typer.Exit(code=1)
@@ -271,9 +343,7 @@ def query(
 
         pipeline = JarvisRAGPipeline(
             llm=custom_llm,
-            embedding_mode=cast(
-                Optional[Literal["performance", "accuracy"]], embedding_mode
-            ),
+            embedding_model=embedding_model,
             db_path=str(db_path) if db_path else None,
             collection_name=collection_name,
         )
@@ -282,7 +352,7 @@ def query(
         answer = pipeline.query(question)
 
         print("💬 答案:")
-        # We can still use rich.markdown.Markdown as PrettyOutput uses rich underneath
+        # 我们仍然可以使用 rich.markdown.Markdown，因为 PrettyOutput 底层使用了 rich
         from jarvis.jarvis_utils.globals import console
 
         console.print(Markdown(answer))
@@ -292,6 +362,25 @@ def query(
         raise typer.Exit(code=1)
 
 
+_RAG_INSTALLED = False
+try:
+    import langchain  # noqa
+
+    _RAG_INSTALLED = True
+except ImportError:
+    pass
+
+
+def _check_rag_dependencies():
+    if not _RAG_INSTALLED:
+        print(
+            "❌ RAG依赖项未安装。"
+            "请运行 'pip install \"jarvis-ai-assistant[rag]\"' 来使用此命令。"
+        )
+        raise typer.Exit(code=1)
+
+
 def main():
+    _check_rag_dependencies()
     init_env(welcome_str="Jarvis RAG")
     app()

jarvis/jarvis_rag/embedding_manager.py

@@ -1,59 +1,45 @@
-from typing import List, Literal, cast
+import torch
+from typing import List, cast
 from langchain_huggingface import HuggingFaceEmbeddings
 
-from jarvis.jarvis_utils.config import (
-    get_rag_embedding_models,
-    get_rag_embedding_cache_path,
-)
 from .cache import EmbeddingCache
 
 
 class EmbeddingManager:
     """
-    Manages the loading and usage of local embedding models with caching.
+    管理本地嵌入模型的加载和使用，并带有缓存功能。
 
-    This class handles the selection of embedding models based on a specified
-    mode ('performance' or 'accuracy'), loads the model from Hugging Face,
-    and uses a disk-based cache to avoid re-computing embeddings for the
-    same text.
+    该类负责从Hugging Face加载指定的模型，并使用基于磁盘的缓存
+    来避免为相同文本重新计算嵌入。
     """
 
-    def __init__(
-        self,
-        mode: Literal["performance", "accuracy"],
-        cache_dir: str,
-    ):
+    def __init__(self, model_name: str, cache_dir: str):
         """
-        Initializes the EmbeddingManager.
+        初始化EmbeddingManager。
 
-        Args:
-            mode: The desired mode, either 'performance' or 'accuracy'.
-            cache_dir: The directory to store the embedding cache.
+        参数:
+            model_name: 要加载的Hugging Face模型的名称。
+            cache_dir: 用于存储嵌入缓存的目录。
         """
-        self.mode = mode
-        self.embedding_models = get_rag_embedding_models()
-        if mode not in self.embedding_models:
-            raise ValueError(
-                f"Invalid mode '{mode}'. Must be one of {list(self.embedding_models.keys())}"
-            )
-
-        self.model_config = self.embedding_models[self.mode]
-        self.model_name = self.model_config["model_name"]
+        self.model_name = model_name
 
-        print(f"🚀 初始化嵌入管理器，模式: '{self.mode}', 模型: '{self.model_name}'...")
+        print(f"🚀 初始化嵌入管理器, 模型: '{self.model_name}'...")
 
-        # The salt for the cache is the model name to prevent collisions
-        self.cache = EmbeddingCache(cache_dir=cache_dir, salt=str(self.model_name))
+        # 缓存的salt是模型名称，以防止冲突
+        self.cache = EmbeddingCache(cache_dir=cache_dir, salt=self.model_name)
         self.model = self._load_model()
 
     def _load_model(self) -> HuggingFaceEmbeddings:
-        """Loads the Hugging Face embedding model based on the configuration."""
+        """根据配置加载Hugging Face嵌入模型。"""
+        model_kwargs = {"device": "cuda" if torch.cuda.is_available() else "cpu"}
+        encode_kwargs = {"normalize_embeddings": True}
+
         try:
             return HuggingFaceEmbeddings(
                 model_name=self.model_name,
-                model_kwargs=self.model_config.get("model_kwargs"),
-                encode_kwargs=self.model_config.get("encode_kwargs"),
-                show_progress=self.model_config.get("show_progress", False),
+                model_kwargs=model_kwargs,
+                encode_kwargs=encode_kwargs,
+                show_progress=True,
             )
         except Exception as e:
             print(f"❌ 加载嵌入模型 '{self.model_name}' 时出错: {e}")
@@ -62,18 +48,18 @@ class EmbeddingManager:
 
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
         """
-        Computes embeddings for a list of documents, using the cache.
+        使用缓存为文档列表计算嵌入。
 
-        Args:
-            texts: A list of documents (strings) to embed.
+        参数:
+            texts: 要嵌入的文档（字符串）列表。
 
-        Returns:
-            A list of embeddings, one for each document.
+        返回:
+            一个嵌入列表，每个文档对应一个嵌入。
         """
         if not texts:
             return []
 
-        # Check cache for existing embeddings
+        # 检查缓存中是否已存在嵌入
         cached_embeddings = self.cache.get_batch(texts)
 
         texts_to_embed = []
@@ -83,17 +69,17 @@ class EmbeddingManager:
                 texts_to_embed.append(text)
                 indices_to_embed.append(i)
 
-        # Compute embeddings for texts that were not in the cache
+        # 为不在缓存中的文本计算嵌入
         if texts_to_embed:
             print(
                 f"🔎 缓存未命中。正在为 {len(texts_to_embed)}/{len(texts)} 个文档计算嵌入。"
             )
             new_embeddings = self.model.embed_documents(texts_to_embed)
 
-            # Store new embeddings in the cache
+            # 将新的嵌入存储在缓存中
             self.cache.set_batch(texts_to_embed, new_embeddings)
 
-            # Place new embeddings back into the results list
+            # 将新的嵌入放回结果列表中
             for i, embedding in zip(indices_to_embed, new_embeddings):
                 cached_embeddings[i] = embedding
         else:
@@ -103,7 +89,7 @@ class EmbeddingManager:
 
     def embed_query(self, text: str) -> List[float]:
         """
-        Computes the embedding for a single query.
-        Queries are typically not cached, but we can add it if needed.
+        为单个查询计算嵌入。
+        查询通常不被缓存，但如果需要可以添加。
         """
         return self.model.embed_query(text)

jarvis/jarvis_rag/llm_interface.py

@@ -10,42 +10,40 @@ from jarvis.jarvis_platform.registry import PlatformRegistry
 
 class LLMInterface(ABC):
     """
-    Abstract Base Class for Large Language Model interfaces.
+    大型语言模型接口的抽象基类。
 
-    This class defines the standard interface for interacting with a remote LLM.
-    Any LLM provider (OpenAI, Anthropic, etc.) should be implemented as a
-    subclass of this interface.
+    该类定义了与远程LLM交互的标准接口。
+    任何LLM提供商（如OpenAI、Anthropic等）都应作为该接口的子类来实现。
     """
 
     @abstractmethod
     def generate(self, prompt: str, **kwargs) -> str:
         """
-        Generates a response from the LLM based on a given prompt.
+        根据给定的提示从LLM生成响应。
 
-        Args:
-            prompt: The input prompt to send to the LLM.
-            **kwargs: Additional keyword arguments for the LLM API call
-                      (e.g., temperature, max_tokens).
+        参数:
+            prompt: 发送给LLM的输入提示。
+            **kwargs: LLM API调用的其他关键字参数
+                      （例如，temperature, max_tokens）。
 
-        Returns:
-            The text response generated by the LLM.
+        返回:
+            由LLM生成的文本响应。
         """
         pass
 
 
 class ToolAgent_LLM(LLMInterface):
     """
-    An implementation of the LLMInterface that uses a tool-wielding JarvisAgent
-    to generate the final response.
+    LLMInterface的一个实现，它使用一个能操作工具的JarvisAgent来生成最终响应。
     """
 
     def __init__(self):
         """
-        Initializes the Tool-Agent LLM wrapper.
+        初始化工具-代理 LLM 包装器。
         """
         print("🤖 已初始化工具 Agent 作为最终应答者。")
         self.allowed_tools = ["read_code", "execute_script"]
-        # A generic system prompt for the agent
+        # 为代理提供一个通用的系统提示
         self.system_prompt = "You are a helpful assistant. Please answer the user's question based on the provided context. You can use tools to find more information if needed."
         self.summary_prompt = """
 <report>
@@ -59,17 +57,17 @@ class ToolAgent_LLM(LLMInterface):
 
     def generate(self, prompt: str, **kwargs) -> str:
         """
-        Runs the JarvisAgent with a restricted toolset to generate an answer.
+        使用受限的工具集运行JarvisAgent以生成答案。
 
-        Args:
-            prompt: The full prompt, including context, to be sent to the agent.
-            **kwargs: Ignored, kept for interface compatibility.
+        参数:
+            prompt: 要发送给代理的完整提示，包括上下文。
+            **kwargs: 已忽略，为保持接口兼容性而保留。
 
-        Returns:
-            The final answer generated by the agent.
+        返回:
+            由代理生成的最终答案。
         """
         try:
-            # Initialize the agent with specific settings for RAG context
+            # 使用RAG上下文的特定设置初始化代理
             agent = JarvisAgent(
                 system_prompt=self.system_prompt,
                 use_tools=self.allowed_tools,
@@ -80,7 +78,7 @@ class ToolAgent_LLM(LLMInterface):
                 summary_prompt=self.summary_prompt,
             )
 
-            # The agent's run method expects the 'user_input' parameter
+            # 代理的run方法需要'user_input'参数
             final_answer = agent.run(user_input=prompt)
             return str(final_answer)
 
@@ -91,21 +89,21 @@ class ToolAgent_LLM(LLMInterface):
 
 class JarvisPlatform_LLM(LLMInterface):
     """
-    An implementation of the LLMInterface for the project's internal platform.
+    项目内部平台的LLMInterface实现。
 
-    This class uses the PlatformRegistry to get the configured "normal" model.
+    该类使用PlatformRegistry来获取配置的“普通”模型。
     """
 
     def __init__(self):
         """
-        Initializes the Jarvis Platform LLM client.
+        初始化Jarvis平台LLM客户端。
         """
         try:
             self.registry = PlatformRegistry.get_global_platform_registry()
             self.platform: BasePlatform = self.registry.get_normal_platform()
             self.platform.set_suppress_output(
                 False
-            )  # Ensure no console output from the model
+            )  # 确保模型没有控制台输出
             print(f"🚀 已初始化 Jarvis 平台 LLM，模型: {self.platform.name()}")
         except Exception as e:
             print(f"❌ 初始化 Jarvis 平台 LLM 失败: {e}")
@@ -113,17 +111,17 @@ class JarvisPlatform_LLM(LLMInterface):
 
     def generate(self, prompt: str, **kwargs) -> str:
         """
-        Sends a prompt to the local platform model and returns the response.
+        向本地平台模型发送提示并返回响应。
 
-        Args:
-            prompt: The user's prompt.
-            **kwargs: Ignored, kept for interface compatibility.
+        参数:
+            prompt: 用户的提示。
+            **kwargs: 已忽略，为保持接口兼容性而保留。
 
-        Returns:
-            The response generated by the platform model.
+        返回:
+            由平台模型生成的响应。
         """
         try:
-            # Use the robust chat_until_success method
+            # 使用健壮的chat_until_success方法
             return self.platform.chat_until_success(prompt)
         except Exception as e:
             print(f"❌ 调用 Jarvis 平台模型时发生错误: {e}")