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

jarvis/__init__.py

@@ -1,4 +1,4 @@
 # -*- coding: utf-8 -*-
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.219"
+__version__ = "0.1.220"

jarvis/jarvis_agent/__init__.py

@@ -36,94 +36,49 @@ from jarvis.jarvis_utils.tag import ct, ot
 origin_agent_system_prompt = f"""
 <role>
 # 🤖 角色
-你是一个专业的任务执行助手，擅长根据用户需求生成详细的任务执行计划并执行。
+你是一个专业的任务执行助手，根据用户需求制定并执行详细的计划。
 </role>
 
-<requirements>
-# 🔥 绝对行动要求
-1. 每个响应必须包含且仅包含一个工具调用
-2. 唯一例外：任务结束
-3. 空响应会触发致命错误
-</requirements>
-
-<violations>
-# 🚫 违规示例
-- 没有工具调用的分析 → 永久挂起
-- 未选择的多选项 → 永久挂起
-- 请求用户确认 → 永久挂起
-</violations>
+<rules>
+# ❗ 核心规则
+1.  **单步操作**: 每个响应必须包含且仅包含一个工具调用。
+2.  **任务终结**: 当任务完成时，明确指出任务已完成。这是唯一可以不调用工具的例外。
+3.  **无响应错误**: 空响应或仅有分析无工具调用的响应是致命错误，会导致系统挂起。
+4.  **决策即工具**: 所有的决策和分析都必须通过工具调用来体现。
+5.  **等待结果**: 在继续下一步之前，必须等待当前工具的执行结果。
+6.  **持续推进**: 除非任务完成，否则必须生成可操作的下一步。
+7.  **记录沉淀**: 如果解决方案有普适价值，应记录为方法论。
+8.  **用户语言**: 始终使用用户的语言进行交流。
+</rules>
 
 <workflow>
-# 🔄 问题解决流程
-1. 问题分析
-   - 重述问题以确认理解
-   - 分析根本原因（针对问题分析任务）
-   - 定义清晰、可实现的目标
-   → 必须调用分析工具
-
-2. 解决方案设计
-   - 生成多个可执行的解决方案
-   - 评估并选择最优方案
-   - 使用PlantUML创建详细行动计划
-   → 必须调用设计工具
-
-3. 执行
-   - 一次执行一个步骤
-   - 每个步骤只使用一个工具
-   - 等待工具结果后再继续
-   - 监控结果并根据需要调整
-   → 必须调用执行工具
-
-4. 任务完成
-   - 验证目标完成情况
-   - 如有价值则记录方法论
+# 🔄 工作流程
+1.  **分析**: 理解和分析问题，定义清晰的目标。
+2.  **设计**: 设计解决方案并制定详细的行动计划。
+3.  **执行**: 按照计划，一次一个步骤地执行。
+4.  **完成**: 验证任务是否达成目标，并进行总结。
 </workflow>
 
-<principles>
-# ⚖️ 操作原则
-- 每个步骤一个操作
-- 下一步前必须等待结果
-- 除非任务完成否则必须生成可操作步骤
-- 根据反馈调整计划
-- 记录可复用的解决方案
-- 使用完成命令结束任务
-- 操作之间不能有中间思考状态
-- 所有决策必须表现为工具调用
-</principles>
-
-<rules>
-# ❗ 重要规则
-1. 每个步骤只能使用一个操作
-2. 必须等待操作执行结果
-3. 必须验证任务完成情况
-4. 必须生成可操作步骤
-5. 如果无需操作必须使用完成命令
-6. 永远不要使对话处于等待状态
-7. 始终使用用户语言交流
-8. 必须记录有价值的方法论
-9. 违反操作协议将导致系统崩溃
-10. 空响应会触发永久挂起
-</rules>
-
 <system_info>
-# 系统信息：
-{platform.platform()}
-{platform.version()}
-
-# 当前时间
-{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
+# 系统信息
+- OS: {platform.platform()} {platform.version()}
+- Time: {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
 </system_info>
 """
 
 
 class OutputHandlerProtocol(Protocol):
-    def name(self) -> str: ...
+    def name(self) -> str:
+        ...
 
-    def can_handle(self, response: str) -> bool: ...
+    def can_handle(self, response: str) -> bool:
+        ...
 
-    def prompt(self) -> str: ...
+    def prompt(self) -> str:
+        ...
 
-    def handle(self, response: str, agent: Any) -> Tuple[bool, Any]: ...
+    def handle(self, response: str, agent: Any) -> Tuple[bool, Any]:
+        ...
 
 
 class Agent:
@@ -189,9 +144,7 @@ class Agent:
             if isinstance(platform, str):
                 self.model = PlatformRegistry().create_platform(platform)
                 if self.model is None:
-                    PrettyOutput.print(
-                        f"平台 {platform} 不存在，将使用普通模型", OutputType.WARNING
-                    )
+                    PrettyOutput.print(f"平台 {platform} 不存在，将使用普通模型", OutputType.WARNING)
                     self.model = PlatformRegistry().get_normal_platform()
             else:
                 self.model = platform
@@ -816,18 +769,14 @@ arguments:
 
                     if get_interrupt():
                         set_interrupt(False)
-                        user_input = self.multiline_inputer(
-                            f"模型交互期间被中断，请输入用户干预信息："
-                        )
+                        user_input = self.multiline_inputer(f"模型交互期间被中断，请输入用户干预信息：")
                         if user_input:
                             # 如果有工具调用且用户确认继续，则将干预信息和工具执行结果拼接为prompt
                             if any(
                                 handler.can_handle(current_response)
                                 for handler in self.output_handler
                             ):
-                                if user_confirm(
-                                    "检测到有工具调用，是否继续处理工具调用？", True
-                                ):
+                                if user_confirm("检测到有工具调用，是否继续处理工具调用？", True):
                                     self.prompt = f"{user_input}\n\n{current_response}"
                                     continue
                             self.prompt += f"{user_input}"
@@ -873,9 +822,7 @@ arguments:
             if self.use_methodology:
                 if not upload_methodology(self.model, other_files=self.files):
                     if self.files:
-                        PrettyOutput.print(
-                            "文件上传失败，将忽略文件列表", OutputType.WARNING
-                        )
+                        PrettyOutput.print("文件上传失败，将忽略文件列表", OutputType.WARNING)
                         # 上传失败则回退到本地加载
                     msg = self.prompt
                     for handler in self.input_handler:
@@ -883,14 +830,14 @@ arguments:
                     self.prompt = f"{self.prompt}\n\n以下是历史类似问题的执行经验，可参考：\n{load_methodology(msg, self.get_tool_registry())}"
                 else:
                     if self.files:
-                        self.prompt = f"{self.prompt}\n\n上传的文件包含历史对话信息和方法论文件，可以从中获取一些经验信息。"
+                        self.prompt = (
+                            f"{self.prompt}\n\n上传的文件包含历史对话信息和方法论文件，可以从中获取一些经验信息。"
+                        )
                     else:
                         self.prompt = f"{self.prompt}\n\n上传的文件包含历史对话信息，可以从中获取一些经验信息。"
             elif self.files:
                 if not self.model.upload_files(self.files):
-                    PrettyOutput.print(
-                        "文件上传失败，将忽略文件列表", OutputType.WARNING
-                    )
+                    PrettyOutput.print("文件上传失败，将忽略文件列表", OutputType.WARNING)
                 else:
                     self.prompt = f"{self.prompt}\n\n上传的文件包含历史对话信息，可以从中获取一些经验信息。"
         else:

jarvis/jarvis_data/config_schema.json

@@ -181,6 +181,36 @@
       "description": "是否打印提示",
       "default": false
     },
+    "JARVIS_RAG": {
+      "type": "object",
+      "description": "RAG框架的配置",
+      "properties": {
+        "embedding_mode": {
+          "type": "string",
+          "enum": [
+            "performance",
+            "accuracy"
+          ],
+          "default": "performance",
+          "description": "嵌入模型的模式, 'performance'表示性能优先, 'accuracy'表示准确度优先"
+        },
+        "embedding_cache_path": {
+          "type": "string",
+          "default": ".jarvis/rag/embeddings",
+          "description": "嵌入向量缓存的路径, 相对于当前工作目录"
+        },
+        "vector_db_path": {
+          "type": "string",
+          "default": ".jarvis/rag/vectordb",
+          "description": "向量数据库的持久化存储路径, 相对于当前工作目录"
+        }
+      },
+      "default": {
+        "embedding_mode": "performance",
+        "embedding_cache_path": ".jarvis/rag/embeddings",
+        "vector_db_path": ".jarvis/rag/vectordb"
+      }
+    },
     "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 Framework
+
+A flexible RAG pipeline with pluggable remote LLMs and local, cache-enabled embedding models.
+"""
+
+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,87 @@
+import hashlib
+from typing import List, Optional, Any
+
+from diskcache import Cache
+
+
+class EmbeddingCache:
+    """
+    A disk-based cache for storing and retrieving text embeddings.
+
+    This class uses diskcache to create a persistent, local cache. It generates
+    a key for each text content based on its SHA256 hash, making lookups
+    deterministic and efficient.
+    """
+
+    def __init__(self, cache_dir: str, salt: str = ""):
+        """
+        Initializes the EmbeddingCache.
+
+        Args:
+            cache_dir (str): The directory where the cache will be stored.
+            salt (str): A salt to be added to the hash. This is crucial for
+                        ensuring that embeddings generated by different models
+                        do not collide. For example, use the model name as a salt.
+        """
+        self.cache = Cache(cache_dir)
+        self.salt = salt
+
+    def _get_key(self, text: str) -> str:
+        """Generates a unique cache key for a given text and salt."""
+        hash_object = hashlib.sha256((self.salt + text).encode("utf-8"))
+        return hash_object.hexdigest()
+
+    def get(self, text: str) -> Optional[Any]:
+        """
+        Retrieves an embedding from the cache.
+
+        Args:
+            text (str): The text to look up.
+
+        Returns:
+            The cached embedding, or None if it's not in the cache.
+        """
+        key = self._get_key(text)
+        return self.cache.get(key)
+
+    def set(self, text: str, embedding: Any) -> None:
+        """
+        Stores an embedding in the cache.
+
+        Args:
+            text (str): The text corresponding to the embedding.
+            embedding (Any): The embedding vector to store.
+        """
+        key = self._get_key(text)
+        self.cache.set(key, embedding)
+
+    def get_batch(self, texts: List[str]) -> List[Optional[Any]]:
+        """
+        Retrieves a batch of embeddings from the cache.
+
+        Args:
+            texts (List[str]): A list of texts to look up.
+
+        Returns:
+            A list containing cached embeddings or None for cache misses.
+        """
+        return [self.get(text) for text in texts]
+
+    def set_batch(self, texts: List[str], embeddings: List[Any]) -> None:
+        """
+        Stores a batch of embeddings in the cache.
+
+        Args:
+            texts (List[str]): The list of texts.
+            embeddings (List[Any]): The list of corresponding embeddings.
+        """
+        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):
+        """Closes the cache connection."""
+        self.cache.close()

jarvis/jarvis_rag/cli.py

@@ -0,0 +1,297 @@
+import os
+import sys
+from pathlib import Path
+from typing import Optional, List, Literal, cast
+import mimetypes
+
+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:
+    """
+    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
+        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
+        with open(file_path, "rb") as f:
+            chunk = f.read(4096)  # Read first 4KB
+            if b"\x00" in chunk:
+                return False  # Null bytes are a strong indicator of a binary file
+        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.
+_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="A command-line tool to interact with the Jarvis RAG framework.",
+    add_completion=False,
+)
+
+
+class _CustomPlatformLLM(LLMInterface):
+    """A simple wrapper to make a BasePlatform instance compatible with 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]:
+    """Creates an LLM interface from a specific platform and model."""
+    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
+
+
+@app.command(
+    "add",
+    help="Add documents from files, directories, or glob patterns (e.g., 'src/**/*.py').",
+)
+def add_documents(
+    paths: List[Path] = typer.Argument(
+        ...,
+        help="File/directory paths or glob patterns. Shell expansion is supported.",
+    ),
+    collection_name: str = typer.Option(
+        "jarvis_rag_collection",
+        "--collection",
+        "-c",
+        help="Name of the collection in the vector database.",
+    ),
+    embedding_mode: Optional[str] = typer.Option(
+        None,
+        "--embedding-mode",
+        "-e",
+        help="Embedding mode ('performance' or 'accuracy'). Overrides global config.",
+    ),
+    db_path: Optional[Path] = typer.Option(
+        None, "--db-path", help="Path to the vector database. Overrides global config."
+    ),
+):
+    """Adds documents to the RAG knowledge base from various sources."""
+    files_to_process = set()
+
+    for path_str in paths:
+        # Typer with List[Path] might not expand globs, so we do it manually
+        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(f"⚠️ 在指定路径中未找到任何文本文件。")
+        return
+
+    print(f"✅ 发现 {len(files_to_process)} 个独立文件待处理。")
+
+    try:
+        pipeline = JarvisRAGPipeline(
+            embedding_mode=cast(
+                Optional[Literal["performance", "accuracy"]], embedding_mode
+            ),
+            db_path=str(db_path) if db_path else None,
+            collection_name=collection_name,
+        )
+
+        docs: List[Document] = []
+        loader: BaseLoader
+        for file_path in sorted(list(files_to_process)):
+            try:
+                if file_path.suffix.lower() == ".md":
+                    loader = UnstructuredMarkdownLoader(str(file_path))
+                else:  # Default to TextLoader for .txt and all code files
+                    loader = TextLoader(str(file_path), encoding="utf-8")
+
+                docs.extend(loader.load())
+                print(f"✅ 已加载: {file_path}")
+            except Exception as e:
+                print(f"⚠️ 加载失败 {file_path}: {e}")
+
+        if not docs:
+            print("❌ 未能成功加载任何文档。")
+            raise typer.Exit(code=1)
+
+        pipeline.add_documents(docs)
+        print(f"✅ 成功将 {len(docs)} 个文档的内容添加至集合 '{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.")
+def list_documents(
+    collection_name: str = typer.Option(
+        "jarvis_rag_collection",
+        "--collection",
+        "-c",
+        help="Name of the collection in the vector database.",
+    ),
+    db_path: Optional[Path] = typer.Option(
+        None, "--db-path", help="Path to the vector database. Overrides global config."
+    ),
+):
+    """Lists all unique documents in the specified collection."""
+    try:
+        pipeline = JarvisRAGPipeline(
+            db_path=str(db_path) if db_path else None,
+            collection_name=collection_name,
+        )
+
+        collection = pipeline.retriever.collection
+        results = collection.get()  # Get all items in the collection
+
+        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:
+                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="Ask a question to the knowledge base.")
+def query(
+    question: str = typer.Argument(..., help="The question to ask."),
+    collection_name: str = typer.Option(
+        "jarvis_rag_collection",
+        "--collection",
+        "-c",
+        help="Name of the collection in the vector database.",
+    ),
+    embedding_mode: Optional[str] = typer.Option(
+        None,
+        "--embedding-mode",
+        "-e",
+        help="Embedding mode ('performance' or 'accuracy'). Overrides global config.",
+    ),
+    db_path: Optional[Path] = typer.Option(
+        None, "--db-path", help="Path to the vector database. Overrides global config."
+    ),
+    platform: Optional[str] = typer.Option(
+        None,
+        "--platform",
+        "-p",
+        help="Specify a platform name for the LLM. Overrides the default thinking model.",
+    ),
+    model: Optional[str] = typer.Option(
+        None,
+        "--model",
+        "-m",
+        help="Specify a model name for the LLM. Requires --platform.",
+    ),
+):
+    """Queries the RAG knowledge base and prints the answer."""
+    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_mode=cast(
+                Optional[Literal["performance", "accuracy"]], embedding_mode
+            ),
+            db_path=str(db_path) if db_path else None,
+            collection_name=collection_name,
+        )
+
+        print(f"🤔 正在查询: '{question}'")
+        answer = pipeline.query(question)
+
+        print("💬 答案:")
+        # We can still use rich.markdown.Markdown as PrettyOutput uses rich underneath
+        from jarvis.jarvis_utils.globals import console
+
+        console.print(Markdown(answer))
+
+    except Exception as e:
+        print(f"❌ 发生错误: {e}")
+        raise typer.Exit(code=1)
+
+
+def main():
+    init_env(welcome_str="Jarvis RAG")
+    app()