jarvis-ai-assistant 0.1.219__py3-none-any.whl → 0.1.221__py3-none-any.whl

jarvis/jarvis_agent/tool_executor.py

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+from typing import Any, Tuple, TYPE_CHECKING
+
+from jarvis.jarvis_utils.input import user_confirm
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
+
+if TYPE_CHECKING:
+    from jarvis.jarvis_agent import Agent
+
+
+def execute_tool_call(response: str, agent: "Agent") -> Tuple[bool, Any]:
+    """
+    Parses the model's response, identifies the appropriate tool, and executes it.
+
+    Args:
+        response: The response string from the model, potentially containing a tool call.
+        agent: The agent instance, providing context like output handlers and settings.
+
+    Returns:
+        A tuple containing:
+        - A boolean indicating if the tool's result should be returned to the user.
+        - The result of the tool execution or an error message.
+    """
+    tool_list = []
+    for handler in agent.output_handler:
+        if handler.can_handle(response):
+            tool_list.append(handler)
+
+    if len(tool_list) > 1:
+        error_message = (
+            f"操作失败：检测到多个操作。一次只能执行一个操作。"
+            f"尝试执行的操作：{', '.join([handler.name() for handler in tool_list])}"
+        )
+        PrettyOutput.print(error_message, OutputType.WARNING)
+        return False, error_message
+
+    if not tool_list:
+        return False, ""
+
+    tool_to_execute = tool_list[0]
+    if not agent.execute_tool_confirm or user_confirm(
+        f"需要执行{tool_to_execute.name()}确认执行？", True
+    ):
+        print(f"🔧 正在执行{tool_to_execute.name()}...")
+        result = tool_to_execute.handle(response, agent)
+        print(f"✅ {tool_to_execute.name()}执行完成")
+        return result
+
+    return False, ""

jarvis/jarvis_code_agent/code_agent.py

@@ -392,19 +392,19 @@ class CodeAgent:
             return
         # 用户确认最终结果
         if commited:
-            agent.prompt += final_ret
+            agent.session.prompt += final_ret
             return
         PrettyOutput.print(final_ret, OutputType.USER, lang="markdown")
         if not is_confirm_before_apply_patch() or user_confirm(
             "是否使用此回复？", default=True
         ):
-            agent.prompt += final_ret
+            agent.session.prompt += final_ret
             return
-        agent.prompt += final_ret
+        agent.session.prompt += final_ret
         custom_reply = get_multiline_input("请输入自定义回复")
         if custom_reply.strip():  # 如果自定义回复为空，返回空字符串
             agent.set_addon_prompt(custom_reply)
-        agent.prompt += final_ret
+        agent.session.prompt += final_ret
 
 
 def main() -> None:

jarvis/jarvis_data/config_schema.json

@@ -181,6 +181,26 @@
       "description": "是否打印提示",
       "default": false
     },
+    "JARVIS_RAG": {
+      "type": "object",
+      "description": "RAG框架的配置",
+      "properties": {
+        "embedding_model": {
+          "type": "string",
+          "default": "BAAI/bge-base-zh-v1.5",
+          "description": "用于RAG的嵌入模型的名称, 默认为 'BAAI/bge-base-zh-v1.5'"
+        },
+        "rerank_model": {
+          "type": "string",
+          "default": "BAAI/bge-reranker-base",
+          "description": "用于RAG的rerank模型的名称, 默认为 'BAAI/bge-reranker-base'"
+        }
+      },
+      "default": {
+        "embedding_model": "BAAI/bge-base-zh-v1.5",
+        "rerank_model": "BAAI/bge-reranker-base"
+      }
+    },
     "JARVIS_REPLACE_MAP": {
       "type": "object",
       "description": "自定义替换映射表配置",

jarvis/jarvis_platform/yuanbao.py

@@ -38,7 +38,9 @@ class YuanbaoPlatform(BasePlatform):
         self.agent_id = "naQivTmsDa"
 
         if not self.cookies:
-            PrettyOutput.print("YUANBAO_COOKIES 未设置", OutputType.WARNING)
+            raise ValueError(
+                "YUANBAO_COOKIES environment variable not set. Please provide your cookies to use the Yuanbao platform."
+            )
 
         self.system_message = ""  # 系统消息，用于初始化对话
         self.first_chat = True  # 标识是否为第一次对话

jarvis/jarvis_rag/__init__.py

@@ -0,0 +1,11 @@
+"""
+Jarvis RAG 框架
+
+一个灵活的RAG管道，具有可插拔的远程LLM和本地带缓存的嵌入模型。
+"""
+
+from .rag_pipeline import JarvisRAGPipeline
+from .llm_interface import LLMInterface
+from .embedding_manager import EmbeddingManager
+
+__all__ = ["JarvisRAGPipeline", "LLMInterface", "EmbeddingManager"]

jarvis/jarvis_rag/cache.py

@@ -0,0 +1,85 @@
+import hashlib
+from typing import List, Optional, Any
+
+from diskcache import Cache
+
+
+class EmbeddingCache:
+    """
+    一个用于存储和检索文本嵌入的基于磁盘的缓存。
+
+    该类使用diskcache创建一个持久化的本地缓存。它根据每个文本内容的
+    SHA256哈希值为其生成一个键，使得查找过程具有确定性和高效性。
+    """
+
+    def __init__(self, cache_dir: str, salt: str = ""):
+        """
+        初始化EmbeddingCache。
+
+        参数:
+            cache_dir (str): 缓存将要存储的目录。
+            salt (str): 添加到哈希中的盐值。这对于确保由不同模型生成的
+                        嵌入不会发生冲突至关重要。例如，可以使用模型名称作为盐值。
+        """
+        self.cache = Cache(cache_dir)
+        self.salt = salt
+
+    def _get_key(self, text: str) -> str:
+        """为一个给定的文本和盐值生成一个唯一的缓存键。"""
+        hash_object = hashlib.sha256((self.salt + text).encode("utf-8"))
+        return hash_object.hexdigest()
+
+    def get(self, text: str) -> Optional[Any]:
+        """
+        从缓存中检索一个嵌入。
+
+        参数:
+            text (str): 要查找的文本。
+
+        返回:
+            缓存的嵌入，如果不在缓存中则返回None。
+        """
+        key = self._get_key(text)
+        return self.cache.get(key)
+
+    def set(self, text: str, embedding: Any) -> None:
+        """
+        在缓存中存储一个嵌入。
+
+        参数:
+            text (str): 与嵌入相对应的文本。
+            embedding (Any): 要存储的嵌入向量。
+        """
+        key = self._get_key(text)
+        self.cache.set(key, embedding)
+
+    def get_batch(self, texts: List[str]) -> List[Optional[Any]]:
+        """
+        从缓存中检索一批嵌入。
+
+        参数:
+            texts (List[str]): 要查找的文本列表。
+
+        返回:
+            一个列表，其中包含缓存的嵌入，对于缓存未命中的情况则为None。
+        """
+        return [self.get(text) for text in texts]
+
+    def set_batch(self, texts: List[str], embeddings: List[Any]) -> None:
+        """
+        在缓存中存储一批嵌入。
+
+        参数:
+            texts (List[str]): 文本列表。
+            embeddings (List[Any]): 相应的嵌入列表。
+        """
+        if len(texts) != len(embeddings):
+            raise ValueError("Length of texts and embeddings must be the same.")
+
+        with self.cache.transact():
+            for text, embedding in zip(texts, embeddings):
+                self.set(text, embedding)
+
+    def close(self):
+        """关闭缓存连接。"""
+        self.cache.close()

jarvis/jarvis_rag/cli.py

@@ -0,0 +1,386 @@
+import os
+import sys
+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 (
+    TextLoader,
+    UnstructuredMarkdownLoader,
+)
+from langchain_core.document_loaders.base import BaseLoader
+from rich.markdown import Markdown
+
+from jarvis.jarvis_utils.utils import init_env
+
+
+def is_likely_text_file(file_path: Path) -> bool:
+    """
+    通过读取文件开头部分，检查文件是否可能为文本文件。
+    此方法可以避免将大型二进制文件加载到内存中。
+    """
+    try:
+        # 启发式方法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
+
+        # 启发式方法2：检查文件的前几KB中是否包含空字节
+        with open(file_path, "rb") as f:
+            chunk = f.read(4096)  # 读取前4KB
+            if b"\x00" in chunk:
+                return False  # 空字节是二进制文件的强指示符
+        return True
+    except Exception:
+        return False
+
+
+# 确保项目根目录在Python路径中，以允许绝对导入
+# 这使得脚本可以作为模块运行。
+_project_root = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), "..", "..", "..")
+)
+if _project_root not in sys.path:
+    sys.path.insert(0, _project_root)
+
+from jarvis.jarvis_platform.base import BasePlatform
+from jarvis.jarvis_platform.registry import PlatformRegistry
+from jarvis.jarvis_rag.llm_interface import LLMInterface
+from jarvis.jarvis_rag.rag_pipeline import JarvisRAGPipeline
+
+app = typer.Typer(
+    name="jarvis-rag",
+    help="一个与Jarvis RAG框架交互的命令行工具。",
+    add_completion=False,
+)
+
+
+class _CustomPlatformLLM(LLMInterface):
+    """一个简单的包装器，使BasePlatform实例与LLMInterface兼容。"""
+
+    def __init__(self, platform: BasePlatform):
+        self.platform = platform
+        print(
+            f"✅ 使用自定义LLM: 平台='{platform.platform_name()}', 模型='{platform.name()}'"
+        )
+
+    def generate(self, prompt: str, **kwargs) -> str:
+        return self.platform.chat_until_success(prompt)
+
+
+def _create_custom_llm(platform_name: str, model_name: str) -> Optional[LLMInterface]:
+    """从指定的平台和模型创建LLM接口。"""
+    if not platform_name or not model_name:
+        return None
+    try:
+        registry = PlatformRegistry.get_global_platform_registry()
+        platform_instance = registry.create_platform(platform_name)
+        if not platform_instance:
+            print(f"❌ 错误: 平台 '{platform_name}' 未找到。")
+            return None
+        platform_instance.set_model_name(model_name)
+        platform_instance.set_suppress_output(True)
+        return _CustomPlatformLLM(platform_instance)
+    except Exception as e:
+        print(f"❌ 创建自定义LLM时出错: {e}")
+        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="从文件、目录或glob模式（例如 'src/**/*.py'）添加文档。",
+)
+def add_documents(
+    paths: List[Path] = typer.Argument(
+        ...,
+        help="文件/目录路径或glob模式。支持Shell扩展。",
+    ),
+    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="向量数据库的路径。覆盖全局配置。"
+    ),
+    batch_size: int = typer.Option(
+        500,
+        "--batch-size",
+        "-b",
+        help="单个批次中要处理的文档数。",
+    ),
+):
+    """从不同来源向RAG知识库添加文档。"""
+    files_to_process = set()
+
+    for path_str in paths:
+        # Typer的List[Path]可能不会扩展glob，所以我们手动处理
+        from glob import glob
+
+        expanded_paths = glob(str(path_str), recursive=True)
+
+        for p_str in expanded_paths:
+            path = Path(p_str)
+            if not path.exists():
+                continue
+
+            if path.is_dir():
+                print(f"🔍 正在扫描目录: {path}")
+                for item in path.rglob("*"):
+                    if item.is_file() and is_likely_text_file(item):
+                        files_to_process.add(item)
+            elif path.is_file():
+                if is_likely_text_file(path):
+                    files_to_process.add(path)
+                else:
+                    print(f"⚠️ 跳过可能的二进制文件: {path}")
+
+    if not files_to_process:
+        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_model=embedding_model,
+            db_path=str(db_path) if db_path else None,
+            collection_name=collection_name,
+        )
+
+        docs_batch: List[Document] = []
+        total_docs_added = 0
+        loader: BaseLoader
+
+        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:  # 对.txt和所有代码文件默认使用TextLoader
+                    loader = TextLoader(str(file_path), encoding="utf-8")
+
+                docs_batch.extend(loader.load())
+                print(f"✅ 已加载: {file_path} (文件 {i + 1}/{total_files})")
+            except Exception as e:
+                print(f"⚠️ 加载失败 {file_path}: {e}")
+
+            # 当批处理已满或是最后一个文件时处理批处理
+            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)
+
+        print(
+            f"✅ 成功将 {total_docs_added} 个文档的内容添加至集合 '{collection_name}'。"
+        )
+
+    except Exception as e:
+        print(f"❌ 发生严重错误: {e}")
+        raise typer.Exit(code=1)
+
+
+@app.command("list-docs", help="列出知识库中所有唯一的文档。")
+def list_documents(
+    collection_name: str = typer.Option(
+        "jarvis_rag_collection",
+        "--collection",
+        "-c",
+        help="向量数据库中集合的名称。",
+    ),
+    db_path: Optional[Path] = typer.Option(
+        None, "--db-path", help="向量数据库的路径。覆盖全局配置。"
+    ),
+):
+    """列出指定集合中的所有唯一文档。"""
+    try:
+        pipeline = JarvisRAGPipeline(
+            db_path=str(db_path) if db_path else None,
+            collection_name=collection_name,
+        )
+
+        collection = pipeline.retriever.collection
+        results = collection.get()  # 获取集合中的所有项目
+
+        if not results or not results["metadatas"]:
+            print("ℹ️ 知识库中没有找到任何文档。")
+            return
+
+        # 从元数据中提取唯一的源文件路径
+        sources = set()
+        for metadata in results["metadatas"]:
+            if metadata:
+                source = metadata.get("source")
+                if isinstance(source, str):
+                    sources.add(source)
+
+        if not sources:
+            print("ℹ️ 知识库中没有找到任何带有源信息的文档。")
+            return
+
+        print(f"📚 知识库 '{collection_name}' 中共有 {len(sources)} 个独立文档:")
+        for i, source in enumerate(sorted(list(sources)), 1):
+            print(f"  {i}. {source}")
+
+    except Exception as e:
+        print(f"❌ 发生错误: {e}")
+        raise typer.Exit(code=1)
+
+
+@app.command("query", help="向知识库提问。")
+def query(
+    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="向量数据库的路径。覆盖全局配置。"
+    ),
+    platform: Optional[str] = typer.Option(
+        None,
+        "--platform",
+        "-p",
+        help="为LLM指定平台名称。覆盖默认的思考模型。",
+    ),
+    model: Optional[str] = typer.Option(
+        None,
+        "--model",
+        "-m",
+        help="为LLM指定模型名称。需要 --platform。",
+    ),
+):
+    """查询RAG知识库并打印答案。"""
+    if model and not platform:
+        print("❌ 错误: --model 需要指定 --platform。")
+        raise typer.Exit(code=1)
+
+    try:
+        custom_llm = _create_custom_llm(platform, model) if platform and model else None
+        if (platform or model) and not custom_llm:
+            raise typer.Exit(code=1)
+
+        pipeline = JarvisRAGPipeline(
+            llm=custom_llm,
+            embedding_model=embedding_model,
+            db_path=str(db_path) if db_path else None,
+            collection_name=collection_name,
+        )
+
+        print(f"🤔 正在查询: '{question}'")
+        answer = pipeline.query(question)
+
+        print("💬 答案:")
+        # 我们仍然可以使用 rich.markdown.Markdown，因为 PrettyOutput 底层使用了 rich
+        from jarvis.jarvis_utils.globals import console
+
+        console.print(Markdown(answer))
+
+    except Exception as e:
+        print(f"❌ 发生错误: {e}")
+        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

@@ -0,0 +1,95 @@
+import torch
+from typing import List, cast
+from langchain_huggingface import HuggingFaceEmbeddings
+
+from .cache import EmbeddingCache
+
+
+class EmbeddingManager:
+    """
+    管理本地嵌入模型的加载和使用，并带有缓存功能。
+
+    该类负责从Hugging Face加载指定的模型，并使用基于磁盘的缓存
+    来避免为相同文本重新计算嵌入。
+    """
+
+    def __init__(self, model_name: str, cache_dir: str):
+        """
+        初始化EmbeddingManager。
+
+        参数:
+            model_name: 要加载的Hugging Face模型的名称。
+            cache_dir: 用于存储嵌入缓存的目录。
+        """
+        self.model_name = model_name
+
+        print(f"🚀 初始化嵌入管理器, 模型: '{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:
+        """根据配置加载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=model_kwargs,
+                encode_kwargs=encode_kwargs,
+                show_progress=True,
+            )
+        except Exception as e:
+            print(f"❌ 加载嵌入模型 '{self.model_name}' 时出错: {e}")
+            print("请确保您已安装 'sentence_transformers' 和 'torch'。")
+            raise
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        使用缓存为文档列表计算嵌入。
+
+        参数:
+            texts: 要嵌入的文档（字符串）列表。
+
+        返回:
+            一个嵌入列表，每个文档对应一个嵌入。
+        """
+        if not texts:
+            return []
+
+        # 检查缓存中是否已存在嵌入
+        cached_embeddings = self.cache.get_batch(texts)
+
+        texts_to_embed = []
+        indices_to_embed = []
+        for i, (text, cached) in enumerate(zip(texts, cached_embeddings)):
+            if cached is None:
+                texts_to_embed.append(text)
+                indices_to_embed.append(i)
+
+        # 为不在缓存中的文本计算嵌入
+        if texts_to_embed:
+            print(
+                f"🔎 缓存未命中。正在为 {len(texts_to_embed)}/{len(texts)} 个文档计算嵌入。"
+            )
+            new_embeddings = self.model.embed_documents(texts_to_embed)
+
+            # 将新的嵌入存储在缓存中
+            self.cache.set_batch(texts_to_embed, new_embeddings)
+
+            # 将新的嵌入放回结果列表中
+            for i, embedding in zip(indices_to_embed, new_embeddings):
+                cached_embeddings[i] = embedding
+        else:
+            print(f"✅ 缓存命中。所有 {len(texts)} 个文档的嵌入均从缓存中检索。")
+
+        return cast(List[List[float]], cached_embeddings)
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        为单个查询计算嵌入。
+        查询通常不被缓存，但如果需要可以添加。
+        """
+        return self.model.embed_query(text)