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

jarvis/jarvis_rag/embedding_manager.py

@@ -1,233 +1,10 @@
-import torch
-import os
-from typing import List, cast
-from langchain_huggingface import HuggingFaceEmbeddings
-from huggingface_hub import snapshot_download
+"""
+向后兼容模块：保持 EmbeddingManager 的导入路径。
 
+新的实现已移动到 embeddings/ 目录。
+"""
 
-from .cache import EmbeddingCache
+from .embeddings.local import LocalEmbeddingModel
 
-
-class EmbeddingManager:
-    """
-    管理本地嵌入模型的加载和使用，并带有缓存功能。
-
-    该类负责从Hugging Face加载指定的模型，并使用基于磁盘的缓存
-    来避免为相同文本重新计算嵌入。
-    """
-
-    def __init__(self, model_name: str, cache_dir: str):
-        """
-        初始化EmbeddingManager。
-
-        参数:
-            model_name: 要加载的Hugging Face模型的名称。
-            cache_dir: 用于存储嵌入缓存的目录。
-        """
-        self.model_name = model_name
-
-        print(
-            f"ℹ️ 初始化嵌入管理器, 模型: '{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:
-        """根据配置加载Hugging Face嵌入模型。"""
-        model_kwargs = {"device": "cuda" if torch.cuda.is_available() else "cpu"}
-        encode_kwargs = {"normalize_embeddings": True}
-
-        try:
-            # First try to load model from local cache without any network access
-            try:
-                local_dir = None
-                # Prefer explicit local dir via env or direct path
-
-                if os.path.isdir(self.model_name):
-                    return HuggingFaceEmbeddings(
-                        model_name=self.model_name,
-                        model_kwargs=model_kwargs,
-                        encode_kwargs=encode_kwargs,
-                        show_progress=False,
-                    )
-
-                # Try common local cache directories for sentence-transformers and HF hub
-                try:
-                    home = os.path.expanduser("~")
-                    st_home = os.path.join(home, ".cache", "sentence_transformers")
-                    torch_st_home = os.path.join(home, ".cache", "torch", "sentence_transformers")
-                    # Build common name variants found in local caches
-                    org, name = (
-                        self.model_name.split("/", 1)
-                        if "/" in self.model_name
-                        else ("", self.model_name)
-                    )
-                    san1 = self.model_name.replace("/", "_")
-                    san2 = self.model_name.replace("/", "__")
-                    san3 = self.model_name.replace("/", "--")
-                    # include plain 'name' for caches that drop org prefix
-                    name_variants = list(dict.fromkeys([self.model_name, san1, san2, san3, name]))
-                    candidates = []
-                    for base in [st_home, torch_st_home]:
-                        for nv in name_variants:
-                            p = os.path.join(base, nv)
-                            if os.path.isdir(p):
-                                candidates.append(p)
-                        # Fuzzy scan cache directory for entries that include variants
-                        try:
-                            for entry in os.listdir(base):
-                                ep = os.path.join(base, entry)
-                                if not os.path.isdir(ep):
-                                    continue
-                                if (
-                                    (org and entry.startswith(f"{org}__") and name in entry)
-                                    or (san1 in entry)
-                                    or (name in entry)
-                                ):
-                                    candidates.append(ep)
-                        except Exception:
-                            pass
-
-                    # Hugging Face Hub cache snapshots
-                    hf_cache = os.path.join(home, ".cache", "huggingface", "hub")
-                    if "/" in self.model_name:
-                        org, name = self.model_name.split("/", 1)
-                        models_dir = os.path.join(hf_cache, f"models--{org}--{name}", "snapshots")
-                        if os.path.isdir(models_dir):
-                            try:
-                                snaps = sorted(
-                                    [os.path.join(models_dir, d) for d in os.listdir(models_dir)],
-                                    key=lambda p: os.path.getmtime(p),
-                                    reverse=True,
-                                )
-                            except Exception:
-                                snaps = [os.path.join(models_dir, d) for d in os.listdir(models_dir)]
-                            for sp in snaps:
-                                if os.path.isdir(sp):
-                                    candidates.append(sp)
-                                    break
-
-                    for cand in candidates:
-                        try:
-                            return HuggingFaceEmbeddings(
-                                model_name=cand,
-                                model_kwargs=model_kwargs,
-                                encode_kwargs=encode_kwargs,
-                                show_progress=False,
-                            )
-                        except Exception:
-                            continue
-                except Exception:
-                    pass
-
-                try:
-                    # Try resolve local cached directory; do not hit network
-                    local_dir = snapshot_download(repo_id=self.model_name, local_files_only=True)
-                except Exception:
-                    local_dir = None
-
-                if local_dir:
-                    return HuggingFaceEmbeddings(
-                        model_name=local_dir,
-                        model_kwargs=model_kwargs,
-                        encode_kwargs=encode_kwargs,
-                        show_progress=False,
-                    )
-
-
-
-                # Fall back to remote download if local cache not found and not offline
-                return HuggingFaceEmbeddings(
-                    model_name=self.model_name,
-                    model_kwargs=model_kwargs,
-                    encode_kwargs=encode_kwargs,
-                    show_progress=True,
-                )
-            except Exception as _e:
-                # 如果已检测到本地候选路径（直接目录 / 本地缓存快照），则视为本地加载失败，
-                # 为避免在用户期望“本地优先不联网”的情况下触发联网，直接抛错并给出修复建议。
-                had_local_candidate = False
-                try:
-                    had_local_candidate = (
-                        os.path.isdir(self.model_name)
-                        # 如果上面 snapshot_download 命中了本地缓存，会将 local_dir 设为非 None
-                        or (locals().get("local_dir") is not None)
-                    )
-                except Exception:
-                    pass
-
-                if had_local_candidate:
-                    print(
-                        "❌ 检测到本地模型路径但加载失败。为避免触发网络访问，已中止远程回退。\n"
-                        "请确认本地目录包含完整的 Transformers/Tokenizer 文件（如 config.json、model.safetensors、tokenizer.json/merges.txt 等），\n"
-                        "或在配置中将 embedding_model 设置为该本地目录，或将模型放置到默认的 Hugging Face 缓存目录（例如 ~/.cache/huggingface/hub）。"
-                    )
-                    raise
-
-                # 未发现任何本地候选，则保持原有行为：回退至远程下载
-                return HuggingFaceEmbeddings(
-                    model_name=self.model_name,
-                    model_kwargs=model_kwargs,
-                    encode_kwargs=encode_kwargs,
-                    show_progress=True,
-                )
-        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]]:
-        """
-        使用缓存为文档列表计算嵌入。
-
-        参数:
-            texts: 要嵌入的文档（字符串）列表。
-
-        返回:
-            一个嵌入列表，每个文档对应一个嵌入。
-        """
-        if not texts:
-            return []
-
-        # 检查缓存中是否已存在嵌入
-        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)
-
-        # 为不在缓存中的文本计算嵌入
-        if texts_to_embed:
-            print(
-                f"ℹ️ 缓存未命中。正在为 {len(texts_to_embed)}/{len(texts)} 个文档计算嵌入。"
-            )
-            new_embeddings = self.model.embed_documents(texts_to_embed)
-
-            # 将新的嵌入存储在缓存中
-            self.cache.set_batch(texts_to_embed, new_embeddings)
-
-            # 将新的嵌入放回结果列表中
-            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]:
-        """
-        为单个查询计算嵌入。
-        查询通常不被缓存，但如果需要可以添加。
-        """
-        return self.model.embed_query(text)
+# 向后兼容：保持 EmbeddingManager 作为 LocalEmbeddingModel 的别名
+EmbeddingManager = LocalEmbeddingModel

jarvis/jarvis_rag/embeddings/__init__.py

@@ -0,0 +1,41 @@
+"""
+嵌入模型实现模块。
+
+包含本地和在线嵌入模型的实现，支持动态加载自定义模型。
+"""
+
+from .base import OnlineEmbeddingModel  # noqa: F401
+from .local import LocalEmbeddingModel
+from .registry import EmbeddingRegistry
+
+# 向后兼容别名
+EmbeddingManager = LocalEmbeddingModel
+
+# 在线模型实现（可选导入）
+try:
+    from .cohere import CohereEmbeddingModel  # noqa: F401
+    from .edgefn import EdgeFnEmbeddingModel  # noqa: F401
+    from .openai import OpenAIEmbeddingModel  # noqa: F401
+
+    _base_exports = [
+        "OnlineEmbeddingModel",
+        "LocalEmbeddingModel",
+        "EmbeddingManager",  # 向后兼容
+        "EmbeddingRegistry",
+        "OpenAIEmbeddingModel",
+        "CohereEmbeddingModel",
+        "EdgeFnEmbeddingModel",
+    ]
+except ImportError:
+    _base_exports = [
+        "OnlineEmbeddingModel",
+        "LocalEmbeddingModel",
+        "EmbeddingManager",  # 向后兼容
+        "EmbeddingRegistry",
+    ]
+
+# 动态加载的模型（通过 registry）
+_registry = EmbeddingRegistry.get_global_registry()
+_dynamic_exports = _registry.get_available_embeddings()
+
+__all__ = _base_exports + _dynamic_exports

jarvis/jarvis_rag/embeddings/base.py

@@ -0,0 +1,114 @@
+"""
+在线嵌入模型的基类实现。
+"""
+
+import os
+from typing import List
+from typing import Optional
+
+from ..embedding_interface import EmbeddingInterface
+
+
+class OnlineEmbeddingModel(EmbeddingInterface):
+    """
+    在线嵌入模型的基类实现。
+
+    这是一个抽象基类，定义了在线嵌入模型的基本结构。
+    子类需要实现具体的API调用逻辑。
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        api_key_env: Optional[str] = None,
+        base_url: Optional[str] = None,
+        model_name: str = "text-embedding-3-small",
+        batch_size: int = 100,
+        max_length: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        初始化在线嵌入模型。
+
+        参数:
+            api_key: API密钥。如果为None，将从环境变量中读取。
+            api_key_env: 用于读取API密钥的环境变量名。
+            base_url: API的基础URL。
+            model_name: 要使用的模型名称。
+            batch_size: 批量处理时的批次大小。
+            max_length: 模型的最大输入长度（token数），用于文档分割。
+            **kwargs: 其他配置参数（可能包含从 embedding_config 传入的配置）。
+        """
+        # 优先从 kwargs 中读取 api_key（可能来自 embedding_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参数、embedding_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.batch_size = batch_size
+        self.max_length = max_length
+
+    def _call_api(self, texts: List[str], is_query: bool = False) -> List[List[float]]:
+        """
+        调用在线API获取嵌入。
+
+        参数:
+            texts: 要嵌入的文本列表。
+            is_query: 是否为查询（某些API对查询和文档有不同的端点）。
+
+        返回:
+            嵌入向量列表。
+
+        注意:
+            子类必须实现此方法。
+        """
+        raise NotImplementedError("子类必须实现 _call_api 方法")
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        为文档列表计算嵌入。
+
+        参数:
+            texts: 要嵌入的文档（字符串）列表。
+
+        返回:
+            一个嵌入列表，每个文档对应一个嵌入向量。
+        """
+        if not texts:
+            return []
+
+        # 批量处理以优化API调用
+        all_embeddings = []
+        for i in range(0, len(texts), self.batch_size):
+            batch = texts[i : i + self.batch_size]
+            batch_embeddings = self._call_api(batch, is_query=False)
+            all_embeddings.extend(batch_embeddings)
+
+        return all_embeddings
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        为单个查询计算嵌入。
+
+        参数:
+            text: 要嵌入的查询文本。
+
+        返回:
+            查询的嵌入向量。
+        """
+        embeddings = self._call_api([text], is_query=True)
+        return embeddings[0] if embeddings else []

jarvis/jarvis_rag/embeddings/cohere.py

@@ -0,0 +1,66 @@
+"""
+Cohere 嵌入模型实现。
+"""
+
+from typing import List
+from typing import Optional
+from typing import cast
+
+from .base import OnlineEmbeddingModel
+
+
+class CohereEmbeddingModel(OnlineEmbeddingModel):
+    """
+    Cohere 嵌入模型的实现。
+
+    使用 Cohere 的 embed API。
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model_name: str = "embed-english-v3.0",
+        input_type: str = "search_document",
+        **kwargs,
+    ):
+        """
+        初始化 Cohere 嵌入模型。
+
+        参数:
+            api_key: Cohere API密钥。如果为None，将从COHERE_API_KEY环境变量读取。
+            model_name: 要使用的模型名称。
+            input_type: 输入类型（'search_document' 或 'search_query'）。
+            **kwargs: 传递给父类的其他参数。
+        """
+        super().__init__(
+            api_key=api_key,
+            api_key_env="COHERE_API_KEY",
+            model_name=model_name,
+            **kwargs,
+        )
+        self.input_type = input_type
+
+    def _call_api(self, texts: List[str], is_query: bool = False) -> List[List[float]]:
+        """
+        调用 Cohere API 获取嵌入。
+        """
+        try:
+            import cohere
+
+            client = cohere.Client(api_key=self.api_key)
+
+            input_type = "search_query" if is_query else "search_document"
+
+            response = client.embed(
+                texts=texts,
+                model=self.model_name,
+                input_type=input_type,
+            )
+
+            return cast(List[List[float]], response.embeddings)
+        except ImportError:
+            raise ImportError(
+                "使用 CohereEmbeddingModel 需要安装 cohere 包: pip install cohere"
+            )
+        except Exception as e:
+            raise RuntimeError(f"调用 Cohere API 时出错: {e}")

jarvis/jarvis_rag/embeddings/edgefn.py

@@ -0,0 +1,117 @@
+"""
+EdgeFn 嵌入模型实现。
+"""
+
+from typing import List
+from typing import Optional
+
+from .base import OnlineEmbeddingModel
+
+
+class EdgeFnEmbeddingModel(OnlineEmbeddingModel):
+    """
+    EdgeFn 嵌入模型的实现。
+
+    使用 EdgeFn 的 embeddings API。
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model_name: str = "BAAI/bge-m3",
+        base_url: str = "https://api.edgefn.net/v1/embeddings",
+        **kwargs,
+    ):
+        """
+        初始化 EdgeFn 嵌入模型。
+
+        参数:
+            api_key: EdgeFn API密钥。如果为None，将从EDGEFN_API_KEY环境变量读取。
+            model_name: 要使用的模型名称（如 'BAAI/bge-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, texts: List[str], is_query: bool = False) -> List[List[float]]:
+        """
+        调用 EdgeFn API 获取嵌入。
+
+        参数:
+            texts: 要嵌入的文本列表。
+            is_query: 是否为查询（EdgeFn API 不区分查询和文档）。
+
+        返回:
+            嵌入向量列表。
+        """
+        try:
+            import requests
+
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            }
+
+            # EdgeFn API 支持批量处理
+            # 如果只有一个文本，可以发送字符串；多个文本发送列表
+            payload = {
+                "model": self.model_name,
+                "input": texts[0] if len(texts) == 1 else texts,
+            }
+
+            response = requests.post(
+                self.base_url, headers=headers, json=payload, timeout=60
+            )
+            response.raise_for_status()
+
+            data = response.json()
+
+            # EdgeFn API 返回格式通常是: {"data": [{"embedding": [...]}, ...]}
+            # 处理不同的返回格式
+            if "data" in data:
+                # 标准格式：data 是列表，每个元素包含 embedding
+                if isinstance(data["data"], list) and len(data["data"]) > 0:
+                    if (
+                        isinstance(data["data"][0], dict)
+                        and "embedding" in data["data"][0]
+                    ):
+                        # 格式: {"data": [{"embedding": [...]}, ...]}
+                        embeddings = [item["embedding"] for item in data["data"]]
+                    else:
+                        # 格式: {"data": [[...], [...]]} - 直接是嵌入向量列表
+                        embeddings = data["data"]
+                    return embeddings
+                else:
+                    raise ValueError("EdgeFn API 返回的 data 格式不正确")
+            elif "embedding" in data:
+                # 单个嵌入的情况（非标准格式）
+                return [data["embedding"]]
+            else:
+                # 尝试直接返回（如果 API 直接返回嵌入列表）
+                if isinstance(data, list):
+                    return data
+                raise ValueError(
+                    f"EdgeFn API 返回了意外的格式。响应键: {list(data.keys()) if isinstance(data, dict) else '非字典类型'}"
+                )
+
+        except ImportError:
+            raise ImportError(
+                "使用 EdgeFnEmbeddingModel 需要安装 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}")