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

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

jarvis/jarvis_rag/rag_pipeline.py

@@ -0,0 +1,177 @@
+import os
+from typing import List, Literal, Optional, cast
+
+from langchain.docstore.document import Document
+
+from .embedding_manager import EmbeddingManager
+from .llm_interface import JarvisPlatform_LLM, LLMInterface, ToolAgent_LLM
+from .query_rewriter import QueryRewriter
+from .reranker import Reranker
+from .retriever import ChromaRetriever
+from jarvis.jarvis_utils.config import (
+    get_rag_embedding_mode,
+    get_rag_vector_db_path,
+    get_rag_embedding_cache_path,
+    get_rag_embedding_models,
+)
+
+
+class JarvisRAGPipeline:
+    """
+    The main orchestrator for the RAG pipeline.
+
+    This class integrates the embedding manager, retriever, and LLM to provide
+    a complete pipeline for adding documents and querying them.
+    """
+
+    def __init__(
+        self,
+        llm: Optional[LLMInterface] = None,
+        embedding_mode: Optional[Literal["performance", "accuracy"]] = None,
+        db_path: Optional[str] = None,
+        collection_name: str = "jarvis_rag_collection",
+    ):
+        """
+        Initializes the RAG pipeline.
+
+        Args:
+            llm: An instance of a class implementing LLMInterface.
+                 If None, defaults to the ToolAgent_LLM.
+            embedding_mode: The mode for the local embedding model. If None, uses config value.
+            db_path: Path to the persistent vector database. If None, uses config value.
+            collection_name: Name of the collection in the vector database.
+        """
+        # Determine the embedding model to isolate data paths
+        _embedding_mode = embedding_mode or get_rag_embedding_mode()
+        embedding_models = get_rag_embedding_models()
+        model_name = embedding_models[_embedding_mode]["model_name"]
+        sanitized_model_name = model_name.replace("/", "_").replace("\\", "_")
+
+        # If a specific db_path is given, use it. Otherwise, create a model-specific path.
+        _final_db_path = (
+            str(db_path)
+            if db_path
+            else os.path.join(get_rag_vector_db_path(), sanitized_model_name)
+        )
+        # Always create a model-specific cache path.
+        _final_cache_path = os.path.join(
+            get_rag_embedding_cache_path(), sanitized_model_name
+        )
+
+        self.embedding_manager = EmbeddingManager(
+            mode=cast(Literal["performance", "accuracy"], _embedding_mode),
+            cache_dir=_final_cache_path,
+        )
+        self.retriever = ChromaRetriever(
+            embedding_manager=self.embedding_manager,
+            db_path=_final_db_path,
+            collection_name=collection_name,
+        )
+        # Default to the ToolAgent_LLM unless a specific LLM is provided
+        self.llm = llm if llm is not None else ToolAgent_LLM()
+        self.reranker = Reranker()
+        # Use a standard LLM for the query rewriting task, not the agent
+        self.query_rewriter = QueryRewriter(JarvisPlatform_LLM())
+
+        print("✅ JarvisRAGPipeline 初始化成功。")
+
+    def add_documents(self, documents: List[Document]):
+        """
+        Adds documents to the vector knowledge base.
+
+        Args:
+            documents: A list of LangChain Document objects to add.
+        """
+        self.retriever.add_documents(documents)
+
+    def _create_prompt(
+        self, query: str, context_docs: List[Document], source_files: List[str]
+    ) -> str:
+        """Creates the final prompt for the LLM or Agent."""
+        context = "\n\n".join([doc.page_content for doc in context_docs])
+        sources_text = "\n".join([f"- {source}" for source in source_files])
+
+        prompt_template = f"""
+        你是一个专家助手。请根据用户的问题，结合下面提供的参考信息来回答。
+
+        **重要**: 提供的上下文和文件列表**仅供参考**，可能不完整或已过时。在回答前，你应该**优先使用工具（如 read_code）来获取最新、最准确的信息**。
+
+        参考文件列表:
+        ---
+        {sources_text}
+        ---
+
+        参考上下文:
+        ---
+        {context}
+        ---
+
+        问题: {query}
+
+        回答:
+        """
+        return prompt_template.strip()
+
+    def query(self, query_text: str, n_results: int = 5) -> str:
+        """
+        Performs a query against the knowledge base using a multi-query
+        retrieval and reranking pipeline.
+
+        Args:
+            query_text: The user's original question.
+            n_results: The number of final relevant chunks to retrieve.
+
+        Returns:
+            The answer generated by the LLM.
+        """
+        # 1. Rewrite the original query into multiple queries
+        rewritten_queries = self.query_rewriter.rewrite(query_text)
+
+        # 2. Retrieve initial candidates for each rewritten query
+        all_candidate_docs = []
+        for q in rewritten_queries:
+            print(f"🔍 正在为查询变体 '{q}' 进行混合检索...")
+            candidates = self.retriever.retrieve(q, n_results=n_results * 2)
+            all_candidate_docs.extend(candidates)
+
+        # De-duplicate the candidate documents
+        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. Rerank the unified candidate pool against the *original* query
+        print(
+            f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）..."
+        )
+        retrieved_docs = self.reranker.rerank(
+            query_text, unique_candidate_docs, top_n=n_results
+        )
+
+        if not retrieved_docs:
+            return "我在提供的文档中找不到任何相关信息来回答您的问题。"
+
+        # Print the sources of the final retrieved documents
+        sources = sorted(
+            list(
+                {
+                    doc.metadata["source"]
+                    for doc in retrieved_docs
+                    if "source" in doc.metadata
+                }
+            )
+        )
+        if sources:
+            print(f"📚 根据以下文档回答:")
+            for source in sources:
+                print(f"  - {source}")
+
+        # 4. Create the final prompt and generate the answer
+        # We use the original query_text for the final prompt to the LLM
+        prompt = self._create_prompt(query_text, retrieved_docs, sources)
+
+        print("🤖 正在从LLM生成答案...")
+        answer = self.llm.generate(prompt)
+
+        return answer

jarvis/jarvis_rag/reranker.py

@@ -0,0 +1,56 @@
+from typing import List
+
+from langchain.docstore.document import Document
+from sentence_transformers.cross_encoder import (  # type: ignore
+    CrossEncoder,
+)
+
+
+class Reranker:
+    """
+    A reranker class that uses a Cross-Encoder model to re-score and sort
+    documents based on their relevance to a given query.
+    """
+
+    def __init__(self, model_name: str = "BAAI/bge-reranker-base"):
+        """
+        Initializes the Reranker.
+
+        Args:
+            model_name (str): The name of the Cross-Encoder model to use.
+        """
+        print(f"🔍 正在初始化重排模型: {model_name}...")
+        self.model = CrossEncoder(model_name)
+        print("✅ 重排模型初始化成功。")
+
+    def rerank(
+        self, query: str, documents: List[Document], top_n: int = 5
+    ) -> List[Document]:
+        """
+        Reranks a list of documents based on their relevance to the query.
+
+        Args:
+            query (str): The user's query.
+            documents (List[Document]): The list of documents retrieved from the initial search.
+            top_n (int): The number of top documents to return after reranking.
+
+        Returns:
+            List[Document]: A sorted list of the most relevant documents.
+        """
+        if not documents:
+            return []
+
+        # Create pairs of [query, document_content] for scoring
+        pairs = [[query, doc.page_content] for doc in documents]
+
+        # Get scores from the Cross-Encoder model
+        scores = self.model.predict(pairs)
+
+        # Combine documents with their scores and sort
+        doc_with_scores = list(zip(documents, scores))
+        doc_with_scores.sort(key=lambda x: x[1], reverse=True)
+
+        # Return the top N documents
+        reranked_docs = [doc for doc, score in doc_with_scores[:top_n]]
+
+        return reranked_docs