jarvis-ai-assistant 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

jarvis/jarvis_rag/rag_pipeline.py

@@ -30,6 +30,8 @@ class JarvisRAGPipeline:
         embedding_model: Optional[str] = None,
         db_path: Optional[str] = None,
         collection_name: str = "jarvis_rag_collection",
+        use_bm25: bool = True,
+        use_rerank: bool = True,
     ):
         """
         初始化RAG管道。
@@ -40,6 +42,8 @@ class JarvisRAGPipeline:
             embedding_model: 嵌入模型的名称。如果为None，则使用配置值。
             db_path: 持久化向量数据库的路径。如果为None，则使用配置值。
             collection_name: 向量数据库中集合的名称。
+            use_bm25: 是否在检索中使用BM25。
+            use_rerank: 是否在检索后使用重排器。
         """
         # 确定嵌入模型以隔离数据路径
         model_name = embedding_model or get_rag_embedding_model()
@@ -56,22 +60,87 @@ class JarvisRAGPipeline:
             get_rag_embedding_cache_path(), sanitized_model_name
         )
 
-        self.embedding_manager = EmbeddingManager(
-            model_name=model_name,
-            cache_dir=_final_cache_path,
+        # 存储初始化参数以供延迟加载
+        self.llm = llm if llm is not None else ToolAgent_LLM()
+        self.embedding_model_name = embedding_model or get_rag_embedding_model()
+        self.db_path = db_path
+        self.collection_name = collection_name
+        self.use_bm25 = use_bm25
+        self.use_rerank = use_rerank
+
+        # 延迟加载的组件
+        self._embedding_manager: Optional[EmbeddingManager] = None
+        self._retriever: Optional[ChromaRetriever] = None
+        self._reranker: Optional[Reranker] = None
+        self._query_rewriter: Optional[QueryRewriter] = None
+
+        print("✅ JarvisRAGPipeline 初始化成功 (模型按需加载).")
+
+    def _get_embedding_manager(self) -> EmbeddingManager:
+        if self._embedding_manager is None:
+            sanitized_model_name = self.embedding_model_name.replace("/", "_").replace(
+                "\\", "_"
+            )
+            _final_cache_path = os.path.join(
+                get_rag_embedding_cache_path(), sanitized_model_name
+            )
+            self._embedding_manager = EmbeddingManager(
+                model_name=self.embedding_model_name,
+                cache_dir=_final_cache_path,
+            )
+        return self._embedding_manager
+
+    def _get_retriever(self) -> ChromaRetriever:
+        if self._retriever is None:
+            sanitized_model_name = self.embedding_model_name.replace("/", "_").replace(
+                "\\", "_"
+            )
+            _final_db_path = (
+                str(self.db_path)
+                if self.db_path
+                else os.path.join(get_rag_vector_db_path(), sanitized_model_name)
+            )
+            self._retriever = ChromaRetriever(
+                embedding_manager=self._get_embedding_manager(),
+                db_path=_final_db_path,
+                collection_name=self.collection_name,
+            )
+        return self._retriever
+
+    def _get_collection(self):
+        """
+        在不加载嵌入模型的情况下，直接获取并返回Chroma集合对象。
+        这对于仅需要访问集合元数据（如列出文档）而无需嵌入功能的操作非常有用。
+        """
+        # 为了避免初始化embedding_manager，我们直接构建db_path
+        if self._retriever:
+            return self._retriever.collection
+
+        sanitized_model_name = self.embedding_model_name.replace("/", "_").replace(
+            "\\", "_"
         )
-        self.retriever = ChromaRetriever(
-            embedding_manager=self.embedding_manager,
-            db_path=_final_db_path,
-            collection_name=collection_name,
+        _final_db_path = (
+            str(self.db_path)
+            if self.db_path
+            else os.path.join(get_rag_vector_db_path(), sanitized_model_name)
         )
-        # 除非提供了特定的LLM，否则默认为ToolAgent_LLM
-        self.llm = llm if llm is not None else ToolAgent_LLM()
-        self.reranker = Reranker(model_name=get_rag_rerank_model())
-        # 使用标准LLM执行查询重写任务，而不是代理
-        self.query_rewriter = QueryRewriter(JarvisPlatform_LLM())
 
-        print("✅ JarvisRAGPipeline 初始化成功。")
+        # 直接创建ChromaRetriever所使用的chroma_client，但绕过embedding_manager
+        import chromadb
+
+        chroma_client = chromadb.PersistentClient(path=_final_db_path)
+        return chroma_client.get_collection(name=self.collection_name)
+
+    def _get_reranker(self) -> Reranker:
+        if self._reranker is None:
+            self._reranker = Reranker(model_name=get_rag_rerank_model())
+        return self._reranker
+
+    def _get_query_rewriter(self) -> QueryRewriter:
+        if self._query_rewriter is None:
+            # 使用标准LLM执行查询重写任务，而不是代理
+            self._query_rewriter = QueryRewriter(JarvisPlatform_LLM())
+        return self._query_rewriter
 
     def add_documents(self, documents: List[Document]):
         """
@@ -80,24 +149,21 @@ class JarvisRAGPipeline:
         参数:
             documents: 要添加的LangChain文档对象列表。
         """
-        self.retriever.add_documents(documents)
+        self._get_retriever().add_documents(documents)
 
-    def _create_prompt(
-        self, query: str, context_docs: List[Document], source_files: List[str]
-    ) -> str:
+    def _create_prompt(self, query: str, context_docs: List[Document]) -> str:
         """为LLM或代理创建最终的提示。"""
-        context = "\n\n".join([doc.page_content for doc in context_docs])
-        sources_text = "\n".join([f"- {source}" for source in source_files])
+        context_details = []
+        for doc in context_docs:
+            source = doc.metadata.get("source", "未知来源")
+            content = doc.page_content
+            context_details.append(f"来源: {source}\n\n---\n{content}\n---")
+        context = "\n\n".join(context_details)
 
         prompt_template = f"""
         你是一个专家助手。请根据用户的问题，结合下面提供的参考信息来回答。
 
-        **重要**: 提供的上下文和文件列表**仅供参考**，可能不完整或已过时。在回答前，你应该**优先使用工具（如 read_code）来获取最新、最准确的信息**。
-
-        参考文件列表:
-        ---
-        {sources_text}
-        ---
+        **重要**: 提供的上下文**仅供参考**，可能不完整或已过时。在回答前，你应该**优先使用工具（如 read_code）来获取最新、最准确的信息**。
 
         参考上下文:
         ---
@@ -122,13 +188,15 @@ class JarvisRAGPipeline:
             由LLM生成的答案。
         """
         # 1. 将原始查询重写为多个查询
-        rewritten_queries = self.query_rewriter.rewrite(query_text)
+        rewritten_queries = self._get_query_rewriter().rewrite(query_text)
 
         # 2. 为每个重写的查询检索初始候选文档
         all_candidate_docs = []
         for q in rewritten_queries:
             print(f"🔍 正在为查询变体 '{q}' 进行混合检索...")
-            candidates = self.retriever.retrieve(q, n_results=n_results * 2)
+            candidates = self._get_retriever().retrieve(
+                q, n_results=n_results * 2, use_bm25=self.use_bm25
+            )
             all_candidate_docs.extend(candidates)
 
         # 对候选文档进行去重
@@ -139,12 +207,13 @@ class JarvisRAGPipeline:
             return "我在提供的文档中找不到任何相关信息来回答您的问题。"
 
         # 3. 根据*原始*查询对统一的候选池进行重排
-        print(
-            f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）..."
-        )
-        retrieved_docs = self.reranker.rerank(
-            query_text, unique_candidate_docs, top_n=n_results
-        )
+        if self.use_rerank:
+            print(f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）...")
+            retrieved_docs = self._get_reranker().rerank(
+                query_text, unique_candidate_docs, top_n=n_results
+            )
+        else:
+            retrieved_docs = unique_candidate_docs[:n_results]
 
         if not retrieved_docs:
             return "我在提供的文档中找不到任何相关信息来回答您的问题。"
@@ -166,9 +235,49 @@ class JarvisRAGPipeline:
 
         # 4. 创建最终提示并生成答案
         # 我们使用原始的query_text作为给LLM的最终提示
-        prompt = self._create_prompt(query_text, retrieved_docs, sources)
+        prompt = self._create_prompt(query_text, retrieved_docs)
 
         print("🤖 正在从LLM生成答案...")
         answer = self.llm.generate(prompt)
 
         return answer
+
+    def retrieve_only(self, query_text: str, n_results: int = 5) -> List[Document]:
+        """
+        仅执行检索和重排，不生成答案。
+
+        参数:
+            query_text: 用户的原始问题。
+            n_results: 要检索的最终相关块的数量。
+
+        返回:
+            检索到的文档列表。
+        """
+        # 1. 重写查询
+        rewritten_queries = self._get_query_rewriter().rewrite(query_text)
+
+        # 2. 检索候选文档
+        all_candidate_docs = []
+        for q in rewritten_queries:
+            print(f"🔍 正在为查询变体 '{q}' 进行混合检索...")
+            candidates = self._get_retriever().retrieve(
+                q, n_results=n_results * 2, use_bm25=self.use_bm25
+            )
+            all_candidate_docs.extend(candidates)
+
+        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. 重排
+        if self.use_rerank:
+            print(f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排...")
+            retrieved_docs = self._get_reranker().rerank(
+                query_text, unique_candidate_docs, top_n=n_results
+            )
+        else:
+            retrieved_docs = unique_candidate_docs[:n_results]
+
+        return retrieved_docs

jarvis/jarvis_rag/retriever.py

@@ -39,9 +39,7 @@ class ChromaRetriever:
         self.collection = self.client.get_or_create_collection(
             name=self.collection_name
         )
-        print(
-            f"✅ ChromaDB 客户端已在 '{db_path}' 初始化，集合为 '{collection_name}'。"
-        )
+        print(f"✅ ChromaDB 客户端已在 '{db_path}' 初始化，集合为 '{collection_name}'。")
 
         # BM25索引设置
         self.bm25_index_path = os.path.join(self.db_path, f"{collection_name}_bm25.pkl")
@@ -107,7 +105,9 @@ class ChromaRetriever:
         self.bm25_index = BM25Okapi(self.bm25_corpus)
         self._save_bm25_index()
 
-    def retrieve(self, query: str, n_results: int = 5) -> List[Document]:
+    def retrieve(
+        self, query: str, n_results: int = 5, use_bm25: bool = True
+    ) -> List[Document]:
         """
         使用向量搜索和BM25执行混合检索，然后使用倒数排序融合（RRF）
         对结果进行融合。
@@ -121,7 +121,7 @@ class ChromaRetriever:
 
         # 2. 关键字搜索 (BM25)
         bm25_docs = []
-        if self.bm25_index:
+        if self.bm25_index and use_bm25:
             tokenized_query = query.split()
             doc_scores = self.bm25_index.get_scores(tokenized_query)
 

jarvis/jarvis_tools/generate_new_tool.py

@@ -12,7 +12,9 @@ class generate_new_tool:
     生成并注册新的Jarvis工具。该工具会在用户数据目录下创建新的工具文件，
     并自动注册到当前的工具注册表中。适用场景：1. 需要创建新的自定义工具；
     2. 扩展Jarvis功能；3. 自动化重复性操作；4. 封装特定领域的功能。
-    重要提示：在编写工具代码时，应尽量将工具执行的过程和结果打印出来，方便追踪工具的执行状态。
+    重要提示：
+    1. `tool_name` 参数必须与 `tool_code` 中定义的 `name` 属性完全一致。
+    2. 在编写工具代码时，应尽量将工具执行的过程和结果打印出来，方便追踪工具的执行状态。
     """
 
     parameters = {
@@ -75,6 +77,25 @@ class generate_new_tool:
                     "stderr": f"工具名称 '{tool_name}' 不是有效的Python标识符",
                 }
 
+            # 验证工具代码中的名称是否与tool_name一致
+            import re
+
+            match = re.search(r"^\s*name\s*=\s*[\"'](.+?)[\"']", tool_code, re.MULTILINE)
+            if not match:
+                return {
+                    "success": False,
+                    "stdout": "",
+                    "stderr": "无法在工具代码中找到 'name' 属性。请确保工具类中包含 'name = \"your_tool_name\"'。",
+                }
+
+            code_name = match.group(1)
+            if tool_name != code_name:
+                return {
+                    "success": False,
+                    "stdout": "",
+                    "stderr": f"工具名称不一致：参数 'tool_name' ('{tool_name}') 与代码中的 'name' 属性 ('{code_name}') 必须相同。",
+                }
+
             # 准备工具目录
             tools_dir = Path(get_data_dir()) / "tools"
             tools_dir.mkdir(parents=True, exist_ok=True)

jarvis/jarvis_utils/config.py

@@ -115,22 +115,26 @@ def get_shell_name() -> str:
     return os.path.basename(shell_path).lower()
 
 
-def _get_resolved_model_config(model_group_override: Optional[str] = None) -> Dict[str, Any]:
+def _get_resolved_model_config(
+    model_group_override: Optional[str] = None,
+) -> Dict[str, Any]:
     """
     解析并合并模型配置，处理模型组。
 
     优先级顺序:
     1. 单独的环境变量 (JARVIS_PLATFORM, JARVIS_MODEL, etc.)
-    2. JARVIS_MODEL_GROUP 中定义的组配置
+    2. JARVIS_LLM_GROUP 中定义的组配置
     3. 代码中的默认值
 
     返回:
         Dict[str, Any]: 解析后的模型配置字典
     """
     group_config = {}
-    model_group_name = model_group_override or GLOBAL_CONFIG_DATA.get("JARVIS_MODEL_GROUP")
+    model_group_name = model_group_override or GLOBAL_CONFIG_DATA.get(
+        "JARVIS_LLM_GROUP"
+    )
     # The format is a list of single-key dicts: [{'group_name': {...}}, ...]
-    model_groups = GLOBAL_CONFIG_DATA.get("JARVIS_MODEL_GROUPS", [])
+    model_groups = GLOBAL_CONFIG_DATA.get("JARVIS_LLM_GROUPS", [])
 
     if model_group_name and isinstance(model_groups, list):
         for group_item in model_groups:
@@ -202,7 +206,9 @@ def get_thinking_model_name(model_group_override: Optional[str] = None) -> str:
     """
     config = _get_resolved_model_config(model_group_override)
     # Fallback to normal model if thinking model is not specified
-    return config.get("JARVIS_THINKING_MODEL", get_normal_model_name(model_group_override))
+    return config.get(
+        "JARVIS_THINKING_MODEL", get_normal_model_name(model_group_override)
+    )
 
 
 def is_execute_tool_confirm() -> bool:
@@ -334,14 +340,65 @@ def get_mcp_config() -> List[Dict[str, Any]]:
 # ==============================================================================
 
 
-def get_rag_config() -> Dict[str, Any]:
+DEFAULT_RAG_GROUPS = [
+    {
+        "text": {
+            "embedding_model": "BAAI/bge-m3",
+            "rerank_model": "BAAI/bge-reranker-v2-m3",
+            "use_bm25": True,
+            "use_rerank": True,
+        }
+    },
+    {
+        "code": {
+            "embedding_model": "Qodo/Qodo-Embed-1-7B",
+            "use_bm25": False,
+            "use_rerank": False,
+        }
+    },
+]
+
+
+def _get_resolved_rag_config(
+    rag_group_override: Optional[str] = None,
+) -> Dict[str, Any]:
     """
-    获取RAG框架的配置。
+    解析并合并RAG配置，处理RAG组。
+
+    优先级顺序:
+    1. JARVIS_RAG 中的顶级设置 (embedding_model, etc.)
+    2. JARVIS_RAG_GROUP 中定义的组配置
+    3. 代码中的默认值
 
     返回:
-        Dict[str, Any]: RAG配置字典
+        Dict[str, Any]: 解析后的RAG配置字典
     """
-    return GLOBAL_CONFIG_DATA.get("JARVIS_RAG", {})
+    group_config = {}
+    rag_group_name = rag_group_override or GLOBAL_CONFIG_DATA.get("JARVIS_RAG_GROUP")
+    rag_groups = GLOBAL_CONFIG_DATA.get("JARVIS_RAG_GROUPS", DEFAULT_RAG_GROUPS)
+
+    if rag_group_name and isinstance(rag_groups, list):
+        for group_item in rag_groups:
+            if isinstance(group_item, dict) and rag_group_name in group_item:
+                group_config = group_item[rag_group_name]
+                break
+
+    # Start with group config
+    resolved_config = group_config.copy()
+
+    # Override with specific settings from the top-level JARVIS_RAG dict
+    top_level_rag_config = GLOBAL_CONFIG_DATA.get("JARVIS_RAG", {})
+    if isinstance(top_level_rag_config, dict):
+        for key in [
+            "embedding_model",
+            "rerank_model",
+            "use_bm25",
+            "use_rerank",
+        ]:
+            if key in top_level_rag_config:
+                resolved_config[key] = top_level_rag_config[key]
+
+    return resolved_config
 
 
 def get_rag_embedding_model() -> str:
@@ -351,7 +408,8 @@ def get_rag_embedding_model() -> str:
     返回:
         str: 嵌入模型的名称
     """
-    return get_rag_config().get("embedding_model", "BAAI/bge-base-zh-v1.5")
+    config = _get_resolved_rag_config()
+    return config.get("embedding_model", "BAAI/bge-m3")
 
 
 def get_rag_rerank_model() -> str:
@@ -361,7 +419,8 @@ def get_rag_rerank_model() -> str:
     返回:
         str: rerank模型的名称
     """
-    return get_rag_config().get("rerank_model", "BAAI/bge-reranker-base")
+    config = _get_resolved_rag_config()
+    return config.get("rerank_model", "BAAI/bge-reranker-v2-m3")
 
 
 def get_rag_embedding_cache_path() -> str:
@@ -382,3 +441,25 @@ def get_rag_vector_db_path() -> str:
         str: 数据库路径
     """
     return ".jarvis/rag/vectordb"
+
+
+def get_rag_use_bm25() -> bool:
+    """
+    获取RAG是否使用BM25。
+
+    返回:
+        bool: 如果使用BM25则返回True，默认为True
+    """
+    config = _get_resolved_rag_config()
+    return config.get("use_bm25", True) is True
+
+
+def get_rag_use_rerank() -> bool:
+    """
+    获取RAG是否使用rerank。
+
+    返回:
+        bool: 如果使用rerank则返回True，默认为True
+    """
+    config = _get_resolved_rag_config()
+    return config.get("use_rerank", True) is True

jarvis/jarvis_utils/globals.py

@@ -9,9 +9,11 @@
 """
 import os
 
-# 全局变量：保存最后一条消息
-last_message: str = ""
-from typing import Any, Dict, Set
+# 全局变量：保存消息历史
+from typing import Any, Dict, Set, List
+
+message_history: List[str] = []
+MAX_HISTORY_SIZE = 50
 
 import colorama
 from rich.console import Console
@@ -169,13 +171,18 @@ def get_interrupt() -> int:
 
 def set_last_message(message: str) -> None:
     """
-    设置最后一条消息。
+    将消息添加到历史记录中。
 
     参数:
         message: 要保存的消息
     """
-    global last_message
-    last_message = message
+    global message_history
+    if message:
+        # 避免重复添加
+        if not message_history or message_history[-1] != message:
+            message_history.append(message)
+            if len(message_history) > MAX_HISTORY_SIZE:
+                message_history.pop(0)
 
 
 def get_last_message() -> str:
@@ -183,6 +190,20 @@ def get_last_message() -> str:
     获取最后一条消息。
 
     返回:
-        str: 最后一条消息
+        str: 最后一条消息，如果历史记录为空则返回空字符串
+    """
+    global message_history
+    if message_history:
+        return message_history[-1]
+    return ""
+
+
+def get_message_history() -> List[str]:
+    """
+    获取完整的消息历史记录。
+
+    返回:
+        List[str]: 消息历史列表
     """
-    return last_message
+    global message_history
+    return message_history