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

jarvis/jarvis_rag/rerankers/edgefn.py

@@ -0,0 +1,140 @@
+"""
+EdgeFn 重排模型实现。
+"""
+
+from typing import List
+from typing import Optional
+
+from .base import OnlineReranker
+
+
+class EdgeFnReranker(OnlineReranker):
+    """
+    EdgeFn 重排模型的实现。
+
+    使用 EdgeFn 的 rerank API。
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model_name: str = "bge-reranker-v2-m3",
+        base_url: str = "https://api.edgefn.net/v1/rerank",
+        **kwargs,
+    ):
+        """
+        初始化 EdgeFn 重排模型。
+
+        参数:
+            api_key: EdgeFn API密钥。如果为None，将从EDGEFN_API_KEY环境变量读取。
+            model_name: 要使用的模型名称（如 'bge-reranker-v2-m3'）。
+            base_url: API的基础URL。
+            **kwargs: 传递给父类的其他参数。
+        """
+        super().__init__(
+            api_key=api_key,
+            api_key_env="EDGEFN_API_KEY",
+            base_url=base_url,
+            model_name=model_name,
+            **kwargs,
+        )
+
+    def _call_api(self, query: str, documents: List[str]) -> List[tuple[int, float]]:
+        """
+        调用 EdgeFn API 获取重排分数。
+
+        参数:
+            query: 查询文本。
+            documents: 文档文本列表。
+
+        返回:
+            包含 (索引, 分数) 元组的列表，按分数降序排序。
+        """
+        try:
+            import requests
+
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            }
+
+            payload = {
+                "model": self.model_name,
+                "query": query,
+                "documents": documents,
+            }
+
+            response = requests.post(
+                self.base_url, json=payload, headers=headers, timeout=60
+            )
+            response.raise_for_status()
+
+            data = response.json()
+
+            # EdgeFn API 返回格式可能是多种形式，需要处理不同的格式
+            # 可能的格式:
+            # 1. {"results": [{"index": 0, "score": 0.95}, ...]}
+            # 2. {"results": [{"index": 0, "relevance_score": 0.95}, ...]}
+            # 3. {"data": [{"index": 0, "score": 0.95}, ...]}
+            # 4. 直接返回列表: [{"index": 0, "score": 0.95}, ...]
+
+            results = []
+            if "results" in data:
+                # 格式1和2: 使用 results 字段
+                for item in data["results"]:
+                    index = item.get("index", item.get("rank", 0))
+                    # 尝试不同的分数字段名
+                    score = item.get(
+                        "score",
+                        item.get("relevance_score", item.get("relevance", 0.0)),
+                    )
+                    results.append((index, float(score)))
+            elif "data" in data:
+                # 格式3: 使用 data 字段
+                for item in data["data"]:
+                    index = item.get("index", item.get("rank", 0))
+                    score = item.get(
+                        "score",
+                        item.get("relevance_score", item.get("relevance", 0.0)),
+                    )
+                    results.append((index, float(score)))
+            elif isinstance(data, list):
+                # 格式4: 直接是列表
+                for item in data:
+                    if isinstance(item, dict):
+                        index = item.get("index", item.get("rank", 0))
+                        score = item.get(
+                            "score",
+                            item.get("relevance_score", item.get("relevance", 0.0)),
+                        )
+                        if score is None:
+                            score = 0.0
+                        results.append((index, float(score)))
+                    elif isinstance(item, (int, float)):
+                        # 如果直接是分数列表，使用索引作为位置
+                        results.append((len(results), float(item)))
+            else:
+                raise ValueError(
+                    f"EdgeFn API 返回了意外的格式。响应键: {list(data.keys()) if isinstance(data, dict) else '非字典类型'}"
+                )
+
+            # 按分数降序排序
+            results.sort(key=lambda x: x[1], reverse=True)
+
+            return results
+
+        except ImportError:
+            raise ImportError(
+                "使用 EdgeFnReranker 需要安装 requests 包: pip install requests"
+            )
+        except requests.exceptions.RequestException as e:
+            error_msg = str(e)
+            if hasattr(e, "response") and e.response is not None:
+                try:
+                    error_detail = e.response.json()
+                    error_msg = f"{error_msg} - 详情: {error_detail}"
+                except Exception:
+                    error_msg = f"{error_msg} - 响应状态码: {e.response.status_code}"
+            raise RuntimeError(f"调用 EdgeFn API 时出错: {error_msg}")
+        except Exception as e:
+            raise RuntimeError(f"处理 EdgeFn API 响应时出错: {e}")

jarvis/jarvis_rag/rerankers/jina.py

@@ -0,0 +1,79 @@
+"""
+Jina 重排模型实现。
+"""
+
+from typing import List
+from typing import Optional
+
+from .base import OnlineReranker
+
+
+class JinaReranker(OnlineReranker):
+    """
+    Jina 重排模型的实现。
+
+    使用 Jina 的 rerank API。
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model_name: str = "jina-reranker-v1-base-en",
+        base_url: str = "https://api.jina.ai/v1/rerank",
+        **kwargs,
+    ):
+        """
+        初始化 Jina 重排模型。
+
+        参数:
+            api_key: Jina API密钥。如果为None，将从JINA_API_KEY环境变量读取。
+            model_name: 要使用的模型名称。
+            base_url: API的基础URL。
+            **kwargs: 传递给父类的其他参数。
+        """
+        super().__init__(
+            api_key=api_key,
+            api_key_env="JINA_API_KEY",
+            base_url=base_url,
+            model_name=model_name,
+            **kwargs,
+        )
+
+    def _call_api(self, query: str, documents: List[str]) -> List[tuple[int, float]]:
+        """
+        调用 Jina API 获取重排分数。
+        """
+        try:
+            import requests
+
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.api_key}",
+            }
+
+            payload = {
+                "model": self.model_name,
+                "query": query,
+                "documents": documents,
+                "top_n": len(documents),
+            }
+
+            response = requests.post(self.base_url, json=payload, headers=headers)
+            response.raise_for_status()
+
+            data = response.json()
+
+            # Jina API 返回格式: {"results": [{"index": 0, "relevance_score": 0.95}, ...]}
+            results = [
+                (item["index"], item["relevance_score"])
+                for item in data.get("results", [])
+            ]
+            results.sort(key=lambda x: x[1], reverse=True)
+
+            return results
+        except ImportError:
+            raise ImportError(
+                "使用 JinaReranker 需要安装 requests 包: pip install requests"
+            )
+        except Exception as e:
+            raise RuntimeError(f"调用 Jina API 时出错: {e}")

jarvis/jarvis_rag/rerankers/local.py

@@ -0,0 +1,89 @@
+"""
+本地重排模型实现。
+"""
+
+from jarvis.jarvis_utils.output import PrettyOutput
+
+import os
+from typing import List
+from typing import Optional
+
+from huggingface_hub import snapshot_download
+from langchain.docstore.document import Document
+from sentence_transformers.cross_encoder import CrossEncoder
+
+from ..reranker_interface import RerankerInterface
+
+
+class LocalReranker(RerankerInterface):
+    """
+    一个重排器类，使用Cross-Encoder模型根据文档与给定查询的相关性
+    对文档进行重新评分和排序。
+
+    这是 RerankerInterface 的本地实现，使用 sentence-transformers CrossEncoder。
+    """
+
+    def __init__(self, model_name: str, max_length: Optional[int] = None):
+        """
+        初始化重排器。
+
+        参数:
+            model_name (str): 要使用的Cross-Encoder模型的名称。
+            max_length: 模型的最大输入长度（token数），用于文档处理。
+        """
+        self.max_length = max_length
+        PrettyOutput.auto_print(f"ℹ️ 正在初始化重排模型: {model_name}...")
+        try:
+            local_dir = None
+
+            if os.path.isdir(model_name):
+                self.model = CrossEncoder(model_name)
+                PrettyOutput.auto_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)
+
+            PrettyOutput.auto_print("✅ 重排模型初始化成功。")
+        except Exception as e:
+            PrettyOutput.auto_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)
+
+        # 返回前N个文档
+        reranked_docs = [doc for doc, score in doc_with_scores[:top_n]]
+
+        return reranked_docs

jarvis/jarvis_rag/rerankers/registry.py

@@ -0,0 +1,293 @@
+"""
+重排模型注册表，支持动态加载自定义重排模型实现。
+"""
+
+import importlib
+import inspect
+
+from jarvis.jarvis_utils.output import PrettyOutput
+
+# -*- coding: utf-8 -*-
+import os
+import sys
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Type
+
+from jarvis.jarvis_utils.config import get_data_dir
+
+from ..reranker_interface import RerankerInterface
+
+
+class RerankerRegistry:
+    """重排模型注册表，支持动态加载自定义重排模型实现"""
+
+    global_registry: Optional["RerankerRegistry"] = None
+
+    @staticmethod
+    def get_reranker_dir() -> str:
+        """获取用户自定义重排模型目录"""
+        reranker_dir = os.path.join(get_data_dir(), "rerankers")
+        if not os.path.exists(reranker_dir):
+            try:
+                os.makedirs(reranker_dir)
+                # 创建 __init__.py 使其成为 Python 包
+                with open(
+                    os.path.join(reranker_dir, "__init__.py"), "w", errors="ignore"
+                ):
+                    pass
+            except Exception as e:
+                PrettyOutput.auto_print(f"❌ 创建重排模型目录失败: {str(e)}")
+                return ""
+        return reranker_dir
+
+    @staticmethod
+    def check_reranker_implementation(
+        reranker_class: Type[RerankerInterface],
+    ) -> bool:
+        """检查重排模型类是否实现了所有必需的方法
+
+        参数:
+            reranker_class: 要检查的重排模型类
+
+        返回:
+            bool: 是否实现了所有必需的方法
+        """
+        required_methods = [
+            ("rerank", ["query", "documents", "top_n"]),
+        ]
+
+        missing_methods = []
+
+        for method_name, params in required_methods:
+            if not hasattr(reranker_class, method_name):
+                missing_methods.append(method_name)
+                continue
+
+            method = getattr(reranker_class, method_name)
+            if not callable(method):
+                missing_methods.append(method_name)
+                continue
+
+            # 检查方法参数（允许有默认值）
+            sig = inspect.signature(method)
+            method_params = [p for p in sig.parameters if p != "self"]
+            # 检查必需参数数量（不考虑有默认值的参数）
+            required_params = [
+                p
+                for p in method_params
+                if sig.parameters[p].default == inspect.Parameter.empty
+            ]
+            # rerank 方法需要至少 2 个必需参数（query 和 documents），top_n 可以有默认值
+            if method_name == "rerank":
+                if len(required_params) < 2:
+                    missing_methods.append(f"{method_name}(parameter mismatch)")
+            elif len(required_params) < len(params):
+                missing_methods.append(f"{method_name}(parameter mismatch)")
+
+        if missing_methods:
+            PrettyOutput.auto_print(
+                f"⚠️ 重排模型 {reranker_class.__name__} 缺少必要的方法: {', '.join(missing_methods)}"
+            )
+            return False
+
+        return True
+
+    @staticmethod
+    def load_rerankers_from_dir(
+        directory: str,
+    ) -> Dict[str, Type[RerankerInterface]]:
+        """从指定目录加载重排模型
+
+        参数:
+            directory: 重排模型目录路径
+
+        返回:
+            Dict[str, Type[RerankerInterface]]: 重排模型名称到类的映射
+        """
+        rerankers: Dict[str, Type[RerankerInterface]] = {}
+
+        # 确保目录存在
+        if not os.path.exists(directory):
+            PrettyOutput.auto_print(f"⚠️ 重排模型目录不存在: {directory}")
+            return rerankers
+
+        # 获取目录的包名
+        package_name = None
+        if directory == os.path.dirname(__file__):
+            package_name = "jarvis.jarvis_rag.rerankers"
+
+        # 添加目录到Python路径
+        if directory not in sys.path:
+            sys.path.append(directory)
+
+        error_lines = []
+        # 遍历目录下的所有.py文件
+        for filename in os.listdir(directory):
+            if filename.endswith(".py") and not filename.startswith("__"):
+                module_name = filename[:-3]  # 移除.py后缀
+                try:
+                    # 导入模块
+                    if package_name:
+                        module = importlib.import_module(
+                            f"{package_name}.{module_name}"
+                        )
+                    else:
+                        module = importlib.import_module(module_name)
+
+                    # 遍历模块中的所有类
+                    for _, obj in inspect.getmembers(module):
+                        # 检查是否是RerankerInterface的子类，但不是RerankerInterface本身
+                        if (
+                            inspect.isclass(obj)
+                            and issubclass(obj, RerankerInterface)
+                            and obj != RerankerInterface
+                        ):
+                            # 检查重排模型实现
+                            if not RerankerRegistry.check_reranker_implementation(obj):
+                                continue
+                            try:
+                                # 使用类名作为注册名（可以后续扩展为使用类方法获取名称）
+                                reranker_name = obj.__name__
+                                rerankers[reranker_name] = obj
+                            except Exception as e:
+                                error_lines.append(
+                                    f"注册重排模型失败 {obj.__name__}: {str(e)}"
+                                )
+                except Exception as e:
+                    error_lines.append(f"加载重排模型 {module_name} 失败: {str(e)}")
+
+        if error_lines:
+            joined_errors = "\n".join(error_lines)
+            PrettyOutput.auto_print(f"❌ {joined_errors}")
+        return rerankers
+
+    @staticmethod
+    def get_global_registry() -> "RerankerRegistry":
+        """获取全局重排模型注册表"""
+        if RerankerRegistry.global_registry is None:
+            RerankerRegistry.global_registry = RerankerRegistry()
+        return RerankerRegistry.global_registry
+
+    def __init__(self) -> None:
+        """初始化重排模型注册表"""
+        self.rerankers: Dict[str, Type[RerankerInterface]] = {}
+
+        # 从用户自定义目录加载额外重排模型
+        reranker_dir = RerankerRegistry.get_reranker_dir()
+        if reranker_dir and os.path.exists(reranker_dir):
+            for (
+                reranker_name,
+                reranker_class,
+            ) in RerankerRegistry.load_rerankers_from_dir(reranker_dir).items():
+                self.register_reranker(reranker_name, reranker_class)
+
+        # 从内置目录加载重排模型
+        reranker_dir = os.path.dirname(__file__)
+        if reranker_dir and os.path.exists(reranker_dir):
+            for (
+                reranker_name,
+                reranker_class,
+            ) in RerankerRegistry.load_rerankers_from_dir(reranker_dir).items():
+                self.register_reranker(reranker_name, reranker_class)
+
+    def register_reranker(
+        self, name: str, reranker_class: Type[RerankerInterface]
+    ) -> None:
+        """注册重排模型类
+
+        参数:
+            name: 重排模型名称
+            reranker_class: 重排模型类
+        """
+        self.rerankers[name] = reranker_class
+
+    def create_reranker(
+        self, name: str, *args, **kwargs
+    ) -> Optional[RerankerInterface]:
+        """创建重排模型实例
+
+        参数:
+            name: 重排模型名称
+            *args: 传递给构造函数的参数
+            **kwargs: 传递给构造函数的关键字参数
+
+        返回:
+            RerankerInterface: 重排模型实例
+        """
+        if name not in self.rerankers:
+            PrettyOutput.auto_print(f"⚠️ 未找到重排模型: {name}")
+            return None
+
+        try:
+            reranker = self.rerankers[name](*args, **kwargs)
+            return reranker
+        except Exception as e:
+            PrettyOutput.auto_print(f"❌ 创建重排模型失败: {str(e)}")
+            return None
+
+    def get_available_rerankers(self) -> List[str]:
+        """获取可用的重排模型列表"""
+        return list(self.rerankers.keys())
+
+    @staticmethod
+    def create_from_config() -> Optional[RerankerInterface]:
+        """从配置创建重排模型实例
+
+        从配置系统读取reranker_type、rerank_model和reranker_config，
+        然后创建相应的重排模型实例。
+
+        返回:
+            Optional[RerankerInterface]: 重排模型实例，如果创建失败则返回None
+        """
+        from jarvis.jarvis_utils.config import get_rag_rerank_model
+        from jarvis.jarvis_utils.config import get_rag_reranker_config
+        from jarvis.jarvis_utils.config import get_rag_reranker_type
+
+        reranker_type = get_rag_reranker_type()
+        model_name = get_rag_rerank_model()
+        reranker_config = get_rag_reranker_config()
+
+        registry = RerankerRegistry.get_global_registry()
+
+        from jarvis.jarvis_utils.config import get_rag_reranker_max_length
+
+        # 构建创建参数
+        create_kwargs = {"model_name": model_name}
+        create_kwargs.update(reranker_config)
+
+        # 将配置中的键名映射到 api_key 参数
+        # 支持的键名：cohere_api_key, edgefn_api_key, jina_api_key
+        api_key_mapping = {
+            "cohere_api_key": "api_key",
+            "edgefn_api_key": "api_key",
+            "jina_api_key": "api_key",
+        }
+        for config_key, param_key in api_key_mapping.items():
+            if config_key in create_kwargs:
+                # 如果还没有设置 api_key，则使用配置中的值
+                if param_key not in create_kwargs:
+                    create_kwargs[param_key] = create_kwargs.pop(config_key)
+                else:
+                    # 如果已经设置了 api_key，移除配置中的键
+                    create_kwargs.pop(config_key)
+
+        # 同样处理 base_url
+        base_url_mapping = {
+            "cohere_api_base": "base_url",
+            "edgefn_api_base": "base_url",
+            "jina_api_base": "base_url",
+        }
+        for config_key, param_key in base_url_mapping.items():
+            if config_key in create_kwargs:
+                if param_key not in create_kwargs:
+                    create_kwargs[param_key] = create_kwargs.pop(config_key)
+                else:
+                    create_kwargs.pop(config_key)
+
+        # 添加max_length（如果配置中没有指定，使用配置系统的默认值）
+        if "max_length" not in create_kwargs:
+            create_kwargs["max_length"] = str(get_rag_reranker_max_length())
+
+        return registry.create_reranker(reranker_type, **create_kwargs)