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

jarvis/jarvis_rag/embeddings/local.py

@@ -0,0 +1,260 @@
+"""
+本地嵌入模型实现。
+"""
+
+from jarvis.jarvis_utils.output import PrettyOutput
+
+import os
+from typing import List
+from typing import Optional
+from typing import cast
+
+import torch
+from huggingface_hub import snapshot_download
+from langchain_huggingface import HuggingFaceEmbeddings
+
+from ..cache import EmbeddingCache
+from ..embedding_interface import EmbeddingInterface
+
+
+class LocalEmbeddingModel(EmbeddingInterface):
+    """
+    管理本地嵌入模型的加载和使用，并带有缓存功能。
+
+    该类负责从Hugging Face加载指定的模型，并使用基于磁盘的缓存
+    来避免为相同文本重新计算嵌入。
+
+    这是 EmbeddingInterface 的本地实现，使用 HuggingFace 模型。
+    """
+
+    def __init__(
+        self, model_name: str, cache_dir: str, max_length: Optional[int] = None
+    ):
+        """
+        初始化LocalEmbeddingModel。
+
+        参数:
+            model_name: 要加载的Hugging Face模型的名称。
+            cache_dir: 用于存储嵌入缓存的目录。
+            max_length: 模型的最大输入长度（token数），用于文档分割。
+        """
+        self.model_name = model_name
+        self.max_length = max_length
+
+        PrettyOutput.auto_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:
+                    PrettyOutput.auto_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:
+            PrettyOutput.auto_print(f"❌ 加载嵌入模型 '{self.model_name}' 时出错: {e}")
+            PrettyOutput.auto_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:
+            PrettyOutput.auto_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:
+            PrettyOutput.auto_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)

jarvis/jarvis_rag/embeddings/openai.py

@@ -0,0 +1,62 @@
+"""
+OpenAI 嵌入模型实现。
+"""
+
+from typing import List
+from typing import Optional
+
+from .base import OnlineEmbeddingModel
+
+
+class OpenAIEmbeddingModel(OnlineEmbeddingModel):
+    """
+    OpenAI 嵌入模型的实现。
+
+    使用 OpenAI 的 embeddings API。
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model_name: str = "text-embedding-3-small",
+        base_url: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        初始化 OpenAI 嵌入模型。
+
+        参数:
+            api_key: OpenAI API密钥。如果为None，将从OPENAI_API_KEY环境变量读取。
+            model_name: 要使用的模型名称（如 'text-embedding-3-small'）。
+            base_url: API的基础URL（用于自定义端点）。
+            **kwargs: 传递给父类的其他参数。
+        """
+        super().__init__(
+            api_key=api_key,
+            api_key_env="OPENAI_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]]:
+        """
+        调用 OpenAI API 获取嵌入。
+        """
+        try:
+            from openai import OpenAI
+
+            client = OpenAI(api_key=self.api_key, base_url=self.base_url)
+
+            response = client.embeddings.create(
+                model=self.model_name,
+                input=texts,
+            )
+
+            return [item.embedding for item in response.data]
+        except ImportError:
+            raise ImportError(
+                "使用 OpenAIEmbeddingModel 需要安装 openai 包: pip install openai"
+            )
+        except Exception as e:
+            raise RuntimeError(f"调用 OpenAI API 时出错: {e}")

jarvis/jarvis_rag/embeddings/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 ..embedding_interface import EmbeddingInterface
+
+
+class EmbeddingRegistry:
+    """嵌入模型注册表，支持动态加载自定义嵌入模型实现"""
+
+    global_registry: Optional["EmbeddingRegistry"] = None
+
+    @staticmethod
+    def get_embedding_dir() -> str:
+        """获取用户自定义嵌入模型目录"""
+        embedding_dir = os.path.join(get_data_dir(), "embeddings")
+        if not os.path.exists(embedding_dir):
+            try:
+                os.makedirs(embedding_dir)
+                # 创建 __init__.py 使其成为 Python 包
+                with open(
+                    os.path.join(embedding_dir, "__init__.py"), "w", errors="ignore"
+                ):
+                    pass
+            except Exception as e:
+                PrettyOutput.auto_print(f"❌ 创建嵌入模型目录失败: {str(e)}")
+                return ""
+        return embedding_dir
+
+    @staticmethod
+    def check_embedding_implementation(
+        embedding_class: Type[EmbeddingInterface],
+    ) -> bool:
+        """检查嵌入模型类是否实现了所有必需的方法
+
+        参数:
+            embedding_class: 要检查的嵌入模型类
+
+        返回:
+            bool: 是否实现了所有必需的方法
+        """
+        required_methods = [
+            ("embed_documents", ["texts"]),
+            ("embed_query", ["text"]),
+        ]
+
+        missing_methods = []
+
+        for method_name, params in required_methods:
+            if not hasattr(embedding_class, method_name):
+                missing_methods.append(method_name)
+                continue
+
+            method = getattr(embedding_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"]
+            if len(method_params) != len(params):
+                missing_methods.append(f"{method_name}(parameter mismatch)")
+
+        if missing_methods:
+            PrettyOutput.auto_print(
+                f"⚠️ 嵌入模型 {embedding_class.__name__} 缺少必要的方法: {', '.join(missing_methods)}"
+            )
+            return False
+
+        return True
+
+    @staticmethod
+    def load_embeddings_from_dir(
+        directory: str,
+    ) -> Dict[str, Type[EmbeddingInterface]]:
+        """从指定目录加载嵌入模型
+
+        参数:
+            directory: 嵌入模型目录路径
+
+        返回:
+            Dict[str, Type[EmbeddingInterface]]: 嵌入模型名称到类的映射
+        """
+        embeddings: Dict[str, Type[EmbeddingInterface]] = {}
+
+        # 确保目录存在
+        if not os.path.exists(directory):
+            PrettyOutput.auto_print(f"⚠️ 嵌入模型目录不存在: {directory}")
+            return embeddings
+
+        # 获取目录的包名
+        package_name = None
+        if directory == os.path.dirname(__file__):
+            package_name = "jarvis.jarvis_rag.embeddings"
+
+        # 添加目录到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):
+                        # 检查是否是EmbeddingInterface的子类，但不是EmbeddingInterface本身
+                        if (
+                            inspect.isclass(obj)
+                            and issubclass(obj, EmbeddingInterface)
+                            and obj != EmbeddingInterface
+                        ):
+                            # 检查嵌入模型实现
+                            if not EmbeddingRegistry.check_embedding_implementation(
+                                obj
+                            ):
+                                continue
+                            try:
+                                # 使用类名作为注册名（可以后续扩展为使用类方法获取名称）
+                                embedding_name = obj.__name__
+                                embeddings[embedding_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 embeddings
+
+    @staticmethod
+    def get_global_registry() -> "EmbeddingRegistry":
+        """获取全局嵌入模型注册表"""
+        if EmbeddingRegistry.global_registry is None:
+            EmbeddingRegistry.global_registry = EmbeddingRegistry()
+        return EmbeddingRegistry.global_registry
+
+    def __init__(self) -> None:
+        """初始化嵌入模型注册表"""
+        self.embeddings: Dict[str, Type[EmbeddingInterface]] = {}
+
+        # 从用户自定义目录加载额外嵌入模型
+        embedding_dir = EmbeddingRegistry.get_embedding_dir()
+        if embedding_dir and os.path.exists(embedding_dir):
+            for (
+                embedding_name,
+                embedding_class,
+            ) in EmbeddingRegistry.load_embeddings_from_dir(embedding_dir).items():
+                self.register_embedding(embedding_name, embedding_class)
+
+        # 从内置目录加载嵌入模型
+        embedding_dir = os.path.dirname(__file__)
+        if embedding_dir and os.path.exists(embedding_dir):
+            for (
+                embedding_name,
+                embedding_class,
+            ) in EmbeddingRegistry.load_embeddings_from_dir(embedding_dir).items():
+                self.register_embedding(embedding_name, embedding_class)
+
+    def register_embedding(
+        self, name: str, embedding_class: Type[EmbeddingInterface]
+    ) -> None:
+        """注册嵌入模型类
+
+        参数:
+            name: 嵌入模型名称
+            embedding_class: 嵌入模型类
+        """
+        self.embeddings[name] = embedding_class
+
+    def create_embedding(
+        self, name: str, *args, **kwargs
+    ) -> Optional[EmbeddingInterface]:
+        """创建嵌入模型实例
+
+        参数:
+            name: 嵌入模型名称
+            *args: 传递给构造函数的参数
+            **kwargs: 传递给构造函数的关键字参数
+
+        返回:
+            EmbeddingInterface: 嵌入模型实例
+        """
+        if name not in self.embeddings:
+            PrettyOutput.auto_print(f"⚠️ 未找到嵌入模型: {name}")
+            return None
+
+        try:
+            embedding = self.embeddings[name](*args, **kwargs)
+            return embedding
+        except Exception as e:
+            PrettyOutput.auto_print(f"❌ 创建嵌入模型失败: {str(e)}")
+            return None
+
+    def get_available_embeddings(self) -> List[str]:
+        """获取可用的嵌入模型列表"""
+        return list(self.embeddings.keys())
+
+    @staticmethod
+    def create_from_config() -> Optional[EmbeddingInterface]:
+        """从配置创建嵌入模型实例
+
+        从配置系统读取embedding_type、embedding_model和embedding_config，
+        然后创建相应的嵌入模型实例。
+
+        返回:
+            Optional[EmbeddingInterface]: 嵌入模型实例，如果创建失败则返回None
+        """
+        from jarvis.jarvis_utils.config import get_rag_embedding_cache_path
+        from jarvis.jarvis_utils.config import get_rag_embedding_config
+        from jarvis.jarvis_utils.config import get_rag_embedding_model
+        from jarvis.jarvis_utils.config import get_rag_embedding_type
+
+        embedding_type = get_rag_embedding_type()
+        model_name = get_rag_embedding_model()
+        embedding_config = get_rag_embedding_config()
+
+        registry = EmbeddingRegistry.get_global_registry()
+
+        from jarvis.jarvis_utils.config import get_rag_embedding_max_length
+
+        # 构建创建参数
+        create_kwargs = {"model_name": model_name}
+        create_kwargs.update(embedding_config)
+
+        # 将配置中的键名映射到 api_key 参数
+        # 支持的键名：openai_api_key, cohere_api_key, edgefn_api_key, jina_api_key
+        api_key_mapping = {
+            "openai_api_key": "api_key",
+            "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 = {
+            "openai_api_base": "base_url",
+            "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_embedding_max_length())
+
+        # 如果是LocalEmbeddingModel，需要添加cache_dir
+        if embedding_type == "LocalEmbeddingModel":
+            create_kwargs["cache_dir"] = get_rag_embedding_cache_path()
+
+        return registry.create_embedding(embedding_type, **create_kwargs)

jarvis/jarvis_rag/llm_interface.py

@@ -1,8 +1,10 @@
-from abc import ABC, abstractmethod
+from abc import ABC
+from abc import abstractmethod
 
 from jarvis.jarvis_agent import Agent as JarvisAgent
 from jarvis.jarvis_platform.base import BasePlatform
 from jarvis.jarvis_platform.registry import PlatformRegistry
+from jarvis.jarvis_utils.output import PrettyOutput
 
 
 class LLMInterface(ABC):
@@ -38,7 +40,7 @@ class ToolAgent_LLM(LLMInterface):
         """
         初始化工具-代理 LLM 包装器。
         """
-        print("ℹ️ 已初始化工具 Agent 作为最终应答者。")
+        PrettyOutput.auto_print("ℹ️ 已初始化工具 Agent 作为最终应答者。")
         self.allowed_tools = ["read_code", "execute_script"]
         # 为代理提供一个通用的系统提示
         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."
@@ -78,7 +80,7 @@ class ToolAgent_LLM(LLMInterface):
             return str(final_answer)
 
         except Exception as e:
-            print(f"❌ Agent 在执行过程中发生错误: {e}")
+            PrettyOutput.auto_print(f"❌ Agent 在执行过程中发生错误: {e}")
             return "错误: Agent 未能成功生成回答。"
 
 
@@ -98,11 +100,11 @@ class JarvisPlatform_LLM(LLMInterface):
             self.registry = PlatformRegistry.get_global_platform_registry()
             self.platform: BasePlatform = self.registry.get_cheap_platform()
             self.platform.set_suppress_output(False)  # 确保模型没有控制台输出
-            print(
+            PrettyOutput.auto_print(
                 f"ℹ️ 已初始化 Jarvis 平台 LLM（cheap），模型: {self.platform.name()}"
             )
         except Exception as e:
-            print(f"❌ 初始化 Jarvis 平台 LLM 失败: {e}")
+            PrettyOutput.auto_print(f"❌ 初始化 Jarvis 平台 LLM 失败: {e}")
             raise
 
     def generate(self, prompt: str, **kwargs) -> str:
@@ -120,5 +122,5 @@ class JarvisPlatform_LLM(LLMInterface):
             # 使用健壮的chat_until_success方法
             return self.platform.chat_until_success(prompt)
         except Exception as e:
-            print(f"❌ 调用 Jarvis 平台模型时发生错误: {e}")
+            PrettyOutput.auto_print(f"❌ 调用 Jarvis 平台模型时发生错误: {e}")
             return "错误: 无法从本地LLM获取响应。"