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

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}")

jarvis/jarvis_rag/query_rewriter.py

@@ -4,22 +4,21 @@ 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.
+    使用LLM将用户的查询重写为多个不同的搜索查询，以提高检索召回率。
     """
 
     def __init__(self, llm: LLMInterface):
         """
-        Initializes the QueryRewriter.
+        初始化QueryRewriter。
 
-        Args:
-            llm: An instance of a class implementing LLMInterface.
+        参数:
+            llm: 实现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 个不同的、但语义上相关的搜索查询。这有助于在知识库中进行更全面的搜索。
 
@@ -39,13 +38,13 @@ class QueryRewriter:
 
     def rewrite(self, query: str) -> List[str]:
         """
-        Rewrites the user query into multiple queries using the LLM.
+        使用LLM将用户查询重写为多个查询。
 
-        Args:
-            query: The original user query.
+        参数:
+            query: 原始用户查询。
 
-        Returns:
-            A list of rewritten, search-optimized queries.
+        返回:
+            一个经过重写、搜索优化的查询列表。
         """
         prompt = self.rewrite_prompt_template.format(query=query)
         print(f"✍️  正在将原始查询重写为多个搜索查询...")
@@ -55,7 +54,7 @@ class QueryRewriter:
             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)
 

jarvis/jarvis_rag/rag_pipeline.py

@@ -1,5 +1,5 @@
 import os
-from typing import List, Literal, Optional, cast
+from typing import List, Optional
 
 from langchain.docstore.document import Document
 
@@ -9,57 +9,55 @@ 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_embedding_model,
+    get_rag_rerank_model,
     get_rag_vector_db_path,
     get_rag_embedding_cache_path,
-    get_rag_embedding_models,
 )
 
 
 class JarvisRAGPipeline:
     """
-    The main orchestrator for the RAG pipeline.
+    RAG管道的主要协调器。
 
-    This class integrates the embedding manager, retriever, and LLM to provide
-    a complete pipeline for adding documents and querying them.
+    该类集成了嵌入管理器、检索器和LLM，为添加文档和查询
+    提供了一个完整的管道。
     """
 
     def __init__(
         self,
         llm: Optional[LLMInterface] = None,
-        embedding_mode: Optional[Literal["performance", "accuracy"]] = None,
+        embedding_model: Optional[str] = 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.
+        初始化RAG管道。
+
+        参数:
+            llm: 实现LLMInterface接口的类的实例。
+                 如果为None，则默认为ToolAgent_LLM。
+            embedding_model: 嵌入模型的名称。如果为None，则使用配置值。
+            db_path: 持久化向量数据库的路径。如果为None，则使用配置值。
+            collection_name: 向量数据库中集合的名称。
         """
-        # 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"]
+        # 确定嵌入模型以隔离数据路径
+        model_name = embedding_model or get_rag_embedding_model()
         sanitized_model_name = model_name.replace("/", "_").replace("\\", "_")
 
-        # If a specific db_path is given, use it. Otherwise, create a model-specific path.
+        # 如果给定了特定的db_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),
+            model_name=model_name,
             cache_dir=_final_cache_path,
         )
         self.retriever = ChromaRetriever(
@@ -67,27 +65,27 @@ class JarvisRAGPipeline:
             db_path=_final_db_path,
             collection_name=collection_name,
         )
-        # Default to the ToolAgent_LLM unless a specific LLM is provided
+        # 除非提供了特定的LLM，否则默认为ToolAgent_LLM
         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.reranker = Reranker(model_name=get_rag_rerank_model())
+        # 使用标准LLM执行查询重写任务，而不是代理
         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.
+        参数:
+            documents: 要添加的LangChain文档对象列表。
         """
         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."""
+        """为LLM或代理创建最终的提示。"""
         context = "\n\n".join([doc.page_content for doc in context_docs])
         sources_text = "\n".join([f"- {source}" for source in source_files])
 
@@ -114,34 +112,33 @@ class JarvisRAGPipeline:
 
     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.
+        参数:
+            query_text: 用户的原始问题。
+            n_results: 要检索的最终相关块的数量。
 
-        Returns:
-            The answer generated by the LLM.
+        返回:
+            由LLM生成的答案。
         """
-        # 1. Rewrite the original query into multiple queries
+        # 1. 将原始查询重写为多个查询
         rewritten_queries = self.query_rewriter.rewrite(query_text)
 
-        # 2. Retrieve initial candidates for each rewritten query
+        # 2. 为每个重写的查询检索初始候选文档
         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
+        # 3. 根据*原始*查询对统一的候选池进行重排
         print(
             f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）..."
         )
@@ -152,7 +149,7 @@ class JarvisRAGPipeline:
         if not retrieved_docs:
             return "我在提供的文档中找不到任何相关信息来回答您的问题。"
 
-        # Print the sources of the final retrieved documents
+        # 打印最终检索到的文档的来源
         sources = sorted(
             list(
                 {
@@ -167,8 +164,8 @@ class JarvisRAGPipeline:
             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
+        # 4. 创建最终提示并生成答案
+        # 我们使用原始的query_text作为给LLM的最终提示
         prompt = self._create_prompt(query_text, retrieved_docs, sources)
 
         print("🤖 正在从LLM生成答案...")

jarvis/jarvis_rag/reranker.py

@@ -8,16 +8,16 @@ from sentence_transformers.cross_encoder import (  # type: ignore
 
 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.
+    一个重排器类，使用Cross-Encoder模型根据文档与给定查询的相关性
+    对文档进行重新评分和排序。
     """
 
-    def __init__(self, model_name: str = "BAAI/bge-reranker-base"):
+    def __init__(self, model_name: str):
         """
-        Initializes the Reranker.
+        初始化重排器。
 
-        Args:
-            model_name (str): The name of the Cross-Encoder model to use.
+        参数:
+            model_name (str): 要使用的Cross-Encoder模型的名称。
         """
         print(f"🔍 正在初始化重排模型: {model_name}...")
         self.model = CrossEncoder(model_name)
@@ -27,30 +27,30 @@ class Reranker:
         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.
+        参数:
+            query (str): 用户的查询。
+            documents (List[Document]): 从初始搜索中检索到的文档列表。
+            top_n (int): 重排后要返回的顶部文档数。
 
-        Returns:
-            List[Document]: A sorted list of the most relevant documents.
+        返回:
+            List[Document]: 一个已排序的最相关文档列表。
         """
         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
+        # 从Cross-Encoder模型获取分数
         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)
+        doc_with_scores.sort(key=lambda x: x[1], reverse=True)  # type: ignore
 
-        # Return the top N documents
+        # 返回前N个文档
         reranked_docs = [doc for doc, score in doc_with_scores[:top_n]]
 
         return reranked_docs