jarvis-ai-assistant 0.7.8__py3-none-any.whl → 1.0.2__py3-none-any.whl

jarvis/jarvis_rag/query_rewriter.py

@@ -1,4 +1,7 @@
 from typing import List
+
+from jarvis.jarvis_utils.output import PrettyOutput
+
 from .llm_interface import LLMInterface
 
 
@@ -56,9 +59,7 @@ English version of the query
             一个经过重写、搜索优化的查询列表。
         """
         prompt = self.rewrite_prompt_template.format(query=query)
-        print(
-            "ℹ️ 正在将原始查询重写为多个搜索查询..."
-        )
+        PrettyOutput.auto_print("ℹ️ 正在将原始查询重写为多个搜索查询...")
 
         import re
 
@@ -77,18 +78,18 @@ English version of the query
                 rewritten_queries = [
                     line.strip() for line in content.split("\n") if line.strip()
                 ]
-                print(
+                PrettyOutput.auto_print(
                     f"✅ 成功从LLM响应中提取到内容 (尝试 {attempts}/{max_retries})。"
                 )
                 break  # 提取成功，退出循环
             else:
-                print(
+                PrettyOutput.auto_print(
                     f"⚠️ 未能从LLM响应中提取内容。正在重试... ({attempts}/{max_retries})"
                 )
 
         # 如果所有重试都失败，则跳过重写步骤
         if not rewritten_queries:
-            print(
+            PrettyOutput.auto_print(
                 "❌ 所有重试均失败。跳过查询重写，将仅使用原始查询。"
             )
 
@@ -96,7 +97,5 @@ English version of the query
         if query not in rewritten_queries:
             rewritten_queries.insert(0, query)
 
-        print(
-            f"✅ 生成了 {len(rewritten_queries)} 个查询变体。"
-        )
+        PrettyOutput.auto_print(f"✅ 生成了 {len(rewritten_queries)} 个查询变体。")
         return rewritten_queries

jarvis/jarvis_rag/rag_pipeline.py

@@ -1,20 +1,27 @@
 import os
-from typing import List, Optional
+from typing import List
+from typing import Optional
 
 from langchain.docstore.document import Document
 
-from .embedding_manager import EmbeddingManager
-from .llm_interface import JarvisPlatform_LLM, LLMInterface, ToolAgent_LLM
+from jarvis.jarvis_utils.config import get_rag_embedding_cache_path
+from jarvis.jarvis_utils.config import get_rag_embedding_model
+from jarvis.jarvis_utils.config import get_rag_rerank_model
+from jarvis.jarvis_utils.config import get_rag_vector_db_path
+from jarvis.jarvis_utils.output import PrettyOutput
+from jarvis.jarvis_utils.utils import get_yes_no
+
+from .embedding_interface import EmbeddingInterface
+from .embeddings import EmbeddingManager
+from .embeddings import EmbeddingRegistry
+from .llm_interface import JarvisPlatform_LLM
+from .llm_interface import LLMInterface
+from .llm_interface import ToolAgent_LLM
 from .query_rewriter import QueryRewriter
-from .reranker import Reranker
+from .reranker_interface import RerankerInterface
+from .rerankers import Reranker
+from .rerankers import RerankerRegistry
 from .retriever import ChromaRetriever
-from jarvis.jarvis_utils.config import (
-    get_rag_embedding_model,
-    get_rag_rerank_model,
-    get_rag_vector_db_path,
-    get_rag_embedding_cache_path,
-)
-from jarvis.jarvis_utils.utils import get_yes_no
 
 
 class JarvisRAGPipeline:
@@ -73,27 +80,31 @@ class JarvisRAGPipeline:
         self.use_query_rewrite = use_query_rewrite
 
         # 延迟加载的组件
-        self._embedding_manager: Optional[EmbeddingManager] = None
+        self._embedding_manager: Optional[EmbeddingInterface] = None
         self._retriever: Optional[ChromaRetriever] = None
-        self._reranker: Optional[Reranker] = None
+        self._reranker: Optional[RerankerInterface] = None
         self._query_rewriter: Optional[QueryRewriter] = None
 
-        print(
-            "✅ JarvisRAGPipeline 初始化成功 (模型按需加载)."
-        )
+        PrettyOutput.auto_print("✅ JarvisRAGPipeline 初始化成功 (模型按需加载).")
 
-    def _get_embedding_manager(self) -> EmbeddingManager:
+    def _get_embedding_manager(self) -> EmbeddingInterface:
         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,
-            )
+            # 尝试从配置创建模型
+            embedding_from_config = EmbeddingRegistry.create_from_config()
+            if embedding_from_config:
+                self._embedding_manager = embedding_from_config
+            else:
+                # 回退到传统方式（向后兼容）
+                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:
@@ -137,9 +148,15 @@ class JarvisRAGPipeline:
         chroma_client = chromadb.PersistentClient(path=_final_db_path)
         return chroma_client.get_collection(name=self.collection_name)
 
-    def _get_reranker(self) -> Reranker:
+    def _get_reranker(self) -> RerankerInterface:
         if self._reranker is None:
-            self._reranker = Reranker(model_name=get_rag_rerank_model())
+            # 尝试从配置创建模型
+            reranker_from_config = RerankerRegistry.create_from_config()
+            if reranker_from_config:
+                self._reranker = reranker_from_config
+            else:
+                # 回退到传统方式（向后兼容）
+                self._reranker = Reranker(model_name=get_rag_rerank_model())
         return self._reranker
 
     def _get_query_rewriter(self) -> QueryRewriter:
@@ -171,19 +188,19 @@ class JarvisRAGPipeline:
                 lines.extend([f"  变更: {p}" for p in changed[:3]])
             if deleted:
                 lines.extend([f"  删除: {p}" for p in deleted[:3]])
-            joined_lines = '\n'.join(lines)
-            print(f"⚠️ {joined_lines}")
+            joined_lines = "\n".join(lines)
+            PrettyOutput.auto_print(f"⚠️ {joined_lines}")
             # 询问用户
             if get_yes_no(
                 "检测到索引变更，是否现在更新索引后再开始检索？", default=True
             ):
                 retriever.update_index_for_changes(changed, deleted)
             else:
-                print(
+                PrettyOutput.auto_print(
                     "ℹ️ 已跳过索引更新，将直接使用当前索引进行检索。"
                 )
         except Exception as e:
-            print(f"⚠️ 检索前索引检查失败：{e}")
+            PrettyOutput.auto_print(f"⚠️ 检索前索引检查失败：{e}")
 
     def add_documents(self, documents: List[Document]):
         """
@@ -236,16 +253,12 @@ class JarvisRAGPipeline:
         if self.use_query_rewrite:
             rewritten_queries = self._get_query_rewriter().rewrite(query_text)
         else:
-            print(
-                "ℹ️ 已关闭查询重写，将直接使用原始查询进行检索。"
-            )
+            PrettyOutput.auto_print("ℹ️ 已关闭查询重写，将直接使用原始查询进行检索。")
             rewritten_queries = [query_text]
 
         # 2. 为每个重写的查询检索初始候选文档
-        query_lines = '\n'.join([f'  - {q}' for q in rewritten_queries])
-        print(
-            f"ℹ️ 将为以下查询变体进行混合检索:\n{query_lines}"
-        )
+        query_lines = "\n".join([f"  - {q}" for q in rewritten_queries])
+        PrettyOutput.auto_print(f"ℹ️ 将为以下查询变体进行混合检索:\n{query_lines}")
         all_candidate_docs = []
         for q in rewritten_queries:
             candidates = self._get_retriever().retrieve(
@@ -262,7 +275,7 @@ class JarvisRAGPipeline:
 
         # 3. 根据*原始*查询对统一的候选池进行重排
         if self.use_rerank:
-            print(
+            PrettyOutput.auto_print(
                 f"ℹ️ 正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）..."
             )
             retrieved_docs = self._get_reranker().rerank(
@@ -287,14 +300,14 @@ class JarvisRAGPipeline:
         if sources:
             # 合并来源列表后一次性打印，避免多次加框
             lines = ["根据以下文档回答:"] + [f"  - {source}" for source in sources]
-            joined_lines = '\n'.join(lines)
-            print(f"ℹ️ {joined_lines}")
+            joined_lines = "\n".join(lines)
+            PrettyOutput.auto_print(f"ℹ️ {joined_lines}")
 
         # 4. 创建最终提示并生成答案
         # 我们使用原始的query_text作为给LLM的最终提示
         prompt = self._create_prompt(query_text, retrieved_docs)
 
-        print("ℹ️ 正在从LLM生成答案...")
+        PrettyOutput.auto_print("ℹ️ 正在从LLM生成答案...")
         answer = self.llm.generate(prompt)
 
         return answer
@@ -316,16 +329,12 @@ class JarvisRAGPipeline:
         if self.use_query_rewrite:
             rewritten_queries = self._get_query_rewriter().rewrite(query_text)
         else:
-            print(
-                "ℹ️ 已关闭查询重写，将直接使用原始查询进行检索。"
-            )
+            PrettyOutput.auto_print("ℹ️ 已关闭查询重写，将直接使用原始查询进行检索。")
             rewritten_queries = [query_text]
 
         # 2. 检索候选文档
-        query_lines = '\n'.join([f'  - {q}' for q in rewritten_queries])
-        print(
-            f"ℹ️ 将为以下查询变体进行混合检索:\n{query_lines}"
-        )
+        query_lines = "\n".join([f"  - {q}" for q in rewritten_queries])
+        PrettyOutput.auto_print(f"ℹ️ 将为以下查询变体进行混合检索:\n{query_lines}")
         all_candidate_docs = []
         for q in rewritten_queries:
             candidates = self._get_retriever().retrieve(
@@ -341,7 +350,7 @@ class JarvisRAGPipeline:
 
         # 3. 重排
         if self.use_rerank:
-            print(
+            PrettyOutput.auto_print(
                 f"ℹ️ 正在对 {len(unique_candidate_docs)} 个候选文档进行重排..."
             )
             retrieved_docs = self._get_reranker().rerank(

jarvis/jarvis_rag/reranker.py

@@ -1,78 +1,10 @@
-from typing import List
-import os
+"""
+向后兼容模块：保持 Reranker 的导入路径。
 
-from langchain.docstore.document import Document
-from sentence_transformers.cross_encoder import (  # type: ignore
-    CrossEncoder,
-)
-from huggingface_hub import snapshot_download
+新的实现已移动到 rerankers/ 目录。
+"""
 
+from .rerankers.local import LocalReranker
 
-class Reranker:
-    """
-    一个重排器类，使用Cross-Encoder模型根据文档与给定查询的相关性
-    对文档进行重新评分和排序。
-    """
-
-    def __init__(self, model_name: str):
-        """
-        初始化重排器。
-
-        参数:
-            model_name (str): 要使用的Cross-Encoder模型的名称。
-        """
-        print(f"ℹ️ 正在初始化重排模型: {model_name}...")
-        try:
-            local_dir = None
-
-            if os.path.isdir(model_name):
-                self.model = CrossEncoder(model_name)
-                print("✅ 重排模型初始化成功。")
-                return
-            try:
-                # Prefer local cache; avoid any network access
-                local_dir = snapshot_download(repo_id=model_name, local_files_only=True)
-            except Exception:
-                local_dir = None
-
-            if local_dir:
-                self.model = CrossEncoder(local_dir)
-            else:
-                self.model = CrossEncoder(model_name)
-
-            print("✅ 重排模型初始化成功。")
-        except Exception as e:
-            print(f"❌ 初始化重排模型失败: {e}")
-            raise
-
-    def rerank(
-        self, query: str, documents: List[Document], top_n: int = 5
-    ) -> List[Document]:
-        """
-        根据文档与查询的相关性对文档列表进行重排。
-
-        参数:
-            query (str): 用户的查询。
-            documents (List[Document]): 从初始搜索中检索到的文档列表。
-            top_n (int): 重排后要返回的顶部文档数。
-
-        返回:
-            List[Document]: 一个已排序的最相关文档列表。
-        """
-        if not documents:
-            return []
-
-        # 创建 [查询, 文档内容] 对用于评分
-        pairs = [[query, doc.page_content] for doc in documents]
-
-        # 从Cross-Encoder模型获取分数
-        scores = self.model.predict(pairs)
-
-        # 将文档与它们的分数结合并排序
-        doc_with_scores = list(zip(documents, scores))
-        doc_with_scores.sort(key=lambda x: x[1], reverse=True)  # type: ignore
-
-        # 返回前N个文档
-        reranked_docs = [doc for doc, score in doc_with_scores[:top_n]]
-
-        return reranked_docs
+# 向后兼容：保持 Reranker 作为 LocalReranker 的别名
+Reranker = LocalReranker

jarvis/jarvis_rag/reranker_interface.py

@@ -0,0 +1,32 @@
+from abc import ABC
+from abc import abstractmethod
+from typing import List
+
+from langchain.docstore.document import Document
+
+
+class RerankerInterface(ABC):
+    """
+    重排模型接口的抽象基类。
+
+    该类定义了重排模型的标准接口，支持本地模型和在线模型（API）的实现。
+    任何重排模型提供商（如sentence-transformers CrossEncoder、Cohere Rerank API等）
+    都应作为该接口的子类来实现。
+    """
+
+    @abstractmethod
+    def rerank(
+        self, query: str, documents: List[Document], top_n: int = 5
+    ) -> List[Document]:
+        """
+        根据文档与查询的相关性对文档列表进行重排。
+
+        参数:
+            query: 用户的查询。
+            documents: 从初始搜索中检索到的文档列表。
+            top_n: 重排后要返回的顶部文档数。
+
+        返回:
+            一个已排序的最相关文档列表。
+        """
+        pass

jarvis/jarvis_rag/rerankers/__init__.py

@@ -0,0 +1,41 @@
+"""
+重排模型实现模块。
+
+包含本地和在线重排模型的实现，支持动态加载自定义模型。
+"""
+
+from .base import OnlineReranker  # noqa: F401
+from .local import LocalReranker
+from .registry import RerankerRegistry
+
+# 向后兼容别名
+Reranker = LocalReranker
+
+# 在线模型实现（可选导入）
+try:
+    from .cohere import CohereReranker  # noqa: F401
+    from .edgefn import EdgeFnReranker  # noqa: F401
+    from .jina import JinaReranker  # noqa: F401
+
+    _base_exports = [
+        "OnlineReranker",
+        "LocalReranker",
+        "Reranker",  # 向后兼容
+        "RerankerRegistry",
+        "CohereReranker",
+        "JinaReranker",
+        "EdgeFnReranker",
+    ]
+except ImportError:
+    _base_exports = [
+        "OnlineReranker",
+        "LocalReranker",
+        "Reranker",  # 向后兼容
+        "RerankerRegistry",
+    ]
+
+# 动态加载的模型（通过 registry）
+_registry = RerankerRegistry.get_global_registry()
+_dynamic_exports = _registry.get_available_rerankers()
+
+__all__ = _base_exports + _dynamic_exports

jarvis/jarvis_rag/rerankers/base.py

@@ -0,0 +1,109 @@
+"""
+在线重排模型的基类实现。
+"""
+
+import os
+from typing import List
+from typing import Optional
+
+from langchain.docstore.document import Document
+
+from ..reranker_interface import RerankerInterface
+
+
+class OnlineReranker(RerankerInterface):
+    """
+    在线重排模型的基类实现。
+
+    这是一个抽象基类，定义了在线重排模型的基本结构。
+    子类需要实现具体的API调用逻辑。
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        api_key_env: Optional[str] = None,
+        base_url: Optional[str] = None,
+        model_name: Optional[str] = None,
+        top_n: int = 5,
+        max_length: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        初始化在线重排模型。
+
+        参数:
+            api_key: API密钥。如果为None，将从环境变量中读取。
+            api_key_env: 用于读取API密钥的环境变量名。
+            base_url: API的基础URL。
+            model_name: 要使用的模型名称。
+            top_n: 默认返回的顶部文档数。
+            max_length: 模型的最大输入长度（token数），用于文档处理。
+            **kwargs: 其他配置参数（可能包含从 reranker_config 传入的配置）。
+        """
+        # 优先从 kwargs 中读取 api_key（可能来自 reranker_config）
+        # 如果没有，则使用传入的 api_key 参数
+        # 最后才从环境变量读取（向后兼容）
+        self.api_key = (
+            kwargs.get("api_key")
+            or api_key
+            or (os.getenv(api_key_env) if api_key_env else None)
+        )
+        if not self.api_key:
+            raise ValueError(
+                f"API密钥未提供。请通过api_key参数、reranker_config配置或环境变量{api_key_env}提供。"
+            )
+
+        # 如果 base_url 在 kwargs 中，优先使用
+        if "base_url" in kwargs and kwargs["base_url"]:
+            self.base_url = kwargs["base_url"]
+        else:
+            self.base_url = base_url
+
+        self.model_name = model_name
+        self.default_top_n = top_n
+        self.max_length = max_length
+
+    def _call_api(self, query: str, documents: List[str]) -> List[tuple[int, float]]:
+        """
+        调用在线API获取重排分数。
+
+        参数:
+            query: 查询文本。
+            documents: 文档文本列表。
+
+        返回:
+            包含 (索引, 分数) 元组的列表，按分数降序排序。
+
+        注意:
+            子类必须实现此方法。
+        """
+        raise NotImplementedError("子类必须实现 _call_api 方法")
+
+    def rerank(
+        self, query: str, documents: List[Document], top_n: int = 5
+    ) -> List[Document]:
+        """
+        根据文档与查询的相关性对文档列表进行重排。
+
+        参数:
+            query: 用户的查询。
+            documents: 从初始搜索中检索到的文档列表。
+            top_n: 重排后要返回的顶部文档数。
+
+        返回:
+            一个已排序的最相关文档列表。
+        """
+        if not documents:
+            return []
+
+        # 提取文档内容
+        doc_texts = [doc.page_content for doc in documents]
+
+        # 调用API获取重排分数
+        scored_indices = self._call_api(query, doc_texts)
+
+        # 根据分数排序并返回前top_n个文档
+        reranked_docs = [documents[idx] for idx, _ in scored_indices[:top_n]]
+
+        return reranked_docs

jarvis/jarvis_rag/rerankers/cohere.py

@@ -0,0 +1,67 @@
+"""
+Cohere 重排模型实现。
+"""
+
+from typing import List
+from typing import Optional
+
+from .base import OnlineReranker
+
+
+class CohereReranker(OnlineReranker):
+    """
+    Cohere 重排模型的实现。
+
+    使用 Cohere 的 rerank API。
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model_name: str = "rerank-english-v3.0",
+        **kwargs,
+    ):
+        """
+        初始化 Cohere 重排模型。
+
+        参数:
+            api_key: Cohere API密钥。如果为None，将从COHERE_API_KEY环境变量读取。
+            model_name: 要使用的模型名称（如 'rerank-english-v3.0'）。
+            **kwargs: 传递给父类的其他参数。
+        """
+        super().__init__(
+            api_key=api_key,
+            api_key_env="COHERE_API_KEY",
+            model_name=model_name,
+            **kwargs,
+        )
+
+    def _call_api(self, query: str, documents: List[str]) -> List[tuple[int, float]]:
+        """
+        调用 Cohere API 获取重排分数。
+        """
+        try:
+            import cohere
+
+            client = cohere.Client(api_key=self.api_key)
+
+            response = client.rerank(
+                model=self.model_name,
+                query=query,
+                documents=documents,
+                top_n=len(documents),  # 获取所有文档的分数
+            )
+
+            # 返回 (索引, 分数) 元组列表，按分数降序排序
+            results = [
+                (result.index, result.relevance_score) for result in response.results
+            ]
+            results.sort(key=lambda x: x[1], reverse=True)
+
+            return results
+        except ImportError:
+            raise ImportError(
+                "使用 CohereReranker 需要安装 cohere 包: pip install cohere"
+            )
+        except Exception as e:
+            raise RuntimeError(f"调用 Cohere API 时出错: {e}")