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

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()

jarvis/jarvis_rag/embedding_manager.py

@@ -0,0 +1,109 @@
+from typing import List, Literal, 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.
+    """
+
+    def __init__(
+        self,
+        mode: Literal["performance", "accuracy"],
+        cache_dir: str,
+    ):
+        """
+        Initializes the EmbeddingManager.
+
+        Args:
+            mode: The desired mode, either 'performance' or 'accuracy'.
+            cache_dir: The directory to store the embedding cache.
+        """
+        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"]
+
+        print(f"🚀 初始化嵌入管理器，模式: '{self.mode}', 模型: '{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))
+        self.model = self._load_model()
+
+    def _load_model(self) -> HuggingFaceEmbeddings:
+        """Loads the Hugging Face embedding model based on the configuration."""
+        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),
+            )
+        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]]:
+        """
+        Computes embeddings for a list of documents, using the cache.
+
+        Args:
+            texts: A list of documents (strings) to embed.
+
+        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 = []
+        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)
+
+        # 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:
+            print(f"✅ 缓存命中。所有 {len(texts)} 个文档的嵌入均从缓存中检索。")
+
+        return cast(List[List[float]], cached_embeddings)
+
+    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

@@ -0,0 +1,130 @@
+from abc import ABC, abstractmethod
+import os
+import os
+from abc import ABC, abstractmethod
+
+from jarvis.jarvis_agent import Agent as JarvisAgent
+from jarvis.jarvis_platform.base import BasePlatform
+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.
+    """
+
+    @abstractmethod
+    def generate(self, prompt: str, **kwargs) -> str:
+        """
+        Generates a response from the LLM based on a given prompt.
+
+        Args:
+            prompt: The input prompt to send to the LLM.
+            **kwargs: Additional keyword arguments for the LLM API call
+                      (e.g., temperature, max_tokens).
+
+        Returns:
+            The text response generated by the LLM.
+        """
+        pass
+
+
+class ToolAgent_LLM(LLMInterface):
+    """
+    An implementation of the LLMInterface that uses a tool-wielding JarvisAgent
+    to generate the final response.
+    """
+
+    def __init__(self):
+        """
+        Initializes the Tool-Agent LLM wrapper.
+        """
+        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>
+请为本次问答任务生成一个总结报告，包含以下内容：
+
+1. **原始问题**: 重述用户最开始提出的问题。
+2. **关键信息来源**: 总结你是基于哪些关键信息或文件得出的结论。
+3. **最终答案**: 给出最终的、精炼的回答。
+</report>
+"""
+
+    def generate(self, prompt: str, **kwargs) -> str:
+        """
+        Runs the JarvisAgent with a restricted toolset to generate an answer.
+
+        Args:
+            prompt: The full prompt, including context, to be sent to the agent.
+            **kwargs: Ignored, kept for interface compatibility.
+
+        Returns:
+            The final answer generated by the agent.
+        """
+        try:
+            # Initialize the agent with specific settings for RAG context
+            agent = JarvisAgent(
+                system_prompt=self.system_prompt,
+                use_tools=self.allowed_tools,
+                auto_complete=True,
+                use_methodology=False,
+                use_analysis=False,
+                need_summary=True,
+                summary_prompt=self.summary_prompt,
+            )
+
+            # The agent's run method expects the 'user_input' parameter
+            final_answer = agent.run(user_input=prompt)
+            return str(final_answer)
+
+        except Exception as e:
+            print(f"❌ Agent 在执行过程中发生错误: {e}")
+            return "错误: Agent 未能成功生成回答。"
+
+
+class JarvisPlatform_LLM(LLMInterface):
+    """
+    An implementation of the LLMInterface for the project's internal platform.
+
+    This class uses the PlatformRegistry to get the configured "normal" model.
+    """
+
+    def __init__(self):
+        """
+        Initializes the Jarvis Platform LLM client.
+        """
+        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}")
+            raise
+
+    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.
+
+        Returns:
+            The response generated by the platform model.
+        """
+        try:
+            # Use the robust chat_until_success method
+            return self.platform.chat_until_success(prompt)
+        except Exception as e:
+            print(f"❌ 调用 Jarvis 平台模型时发生错误: {e}")
+            return "错误: 无法从本地LLM获取响应。"

jarvis/jarvis_rag/query_rewriter.py

@@ -0,0 +1,63 @@
+from typing import List
+from .llm_interface import LLMInterface
+
+
+class QueryRewriter:
+    """
+    Uses an LLM to rewrite a user's query into multiple, diverse search
+    queries to enhance retrieval recall.
+    """
+
+    def __init__(self, llm: LLMInterface):
+        """
+        Initializes the QueryRewriter.
+
+        Args:
+            llm: An instance of a class implementing LLMInterface.
+        """
+        self.llm = llm
+        self.rewrite_prompt_template = self._create_prompt_template()
+
+    def _create_prompt_template(self) -> str:
+        """Creates the prompt template for the multi-query rewriting task."""
+        return """
+你是一个精通检索的AI助手。你的任务是将以下这个单一的用户问题，从不同角度改写成 3 个不同的、但语义上相关的搜索查询。这有助于在知识库中进行更全面的搜索。
+
+请遵循以下原则：
+1.  **多样性**：生成的查询应尝试使用不同的关键词和表述方式。
+2.  **保留核心意图**：所有查询都必须围绕原始问题的核心意图。
+3.  **简洁性**：每个查询都应该是独立的、可以直接用于搜索的短语或问题。
+4.  **格式要求**：请直接输出 3 个查询，每个查询占一行，用换行符分隔。不要添加任何编号、前缀或解释。
+
+原始问题:
+---
+{query}
+---
+
+3个改写后的查询 (每行一个):
+"""
+
+    def rewrite(self, query: str) -> List[str]:
+        """
+        Rewrites the user query into multiple queries using the LLM.
+
+        Args:
+            query: The original user query.
+
+        Returns:
+            A list of rewritten, search-optimized queries.
+        """
+        prompt = self.rewrite_prompt_template.format(query=query)
+        print(f"✍️  正在将原始查询重写为多个搜索查询...")
+
+        response_text = self.llm.generate(prompt)
+        rewritten_queries = [
+            line.strip() for line in response_text.strip().split("\n") if line.strip()
+        ]
+
+        # Also include the original query for robustness
+        if query not in rewritten_queries:
+            rewritten_queries.insert(0, query)
+
+        print(f"✅ 生成了 {len(rewritten_queries)} 个查询变体。")
+        return rewritten_queries