jarvis-ai-assistant 0.3.24__tar.gz → 0.3.26__tar.gz

{jarvis_ai_assistant-0.3.24/src/jarvis_ai_assistant.egg-info → jarvis_ai_assistant-0.3.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jarvis-ai-assistant
-Version: 0.3.24
+Version: 0.3.26
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jarvis-ai-assistant"
-version = "0.3.24"
+version = "0.3.26"
 description = "Jarvis: An AI assistant that uses tools to interact with the system"
 readme = "README.md"
 authors = [{ name = "skyfire", email = "skyfireitdiy@hotmail.com" }]

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages  # type: ignore
 
 setup(
     name="jarvis-ai-assistant",
-    version="0.3.24",
+    version="0.3.26",
     author="skyfire",
     author_email="skyfireitdiy@hotmail.com",
     description="An AI assistant that uses various tools to interact with the system",

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/src/jarvis/__init__.py

@@ -1,4 +1,4 @@
 # -*- coding: utf-8 -*-
 """Jarvis AI Assistant"""
 
-__version__ = "0.3.24"
+__version__ = "0.3.26"

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/src/jarvis/jarvis_agent/__init__.py

@@ -662,6 +662,10 @@ class Agent:
         if self.use_analysis:
             self.task_analyzer.analysis_task(satisfaction_feedback)
 
+        # 当开启强制保存记忆时，在分析步骤之后触发一次记忆保存
+        if self.force_save_memory:
+            self.memory_manager.prompt_memory_save()
+
         self._check_and_organize_memory()
 
         if self.need_summary:

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/src/jarvis/jarvis_agent/jarvis.py

@@ -384,21 +384,21 @@ def handle_builtin_config_selector(
                 table.add_column("描述", style="white")
 
                 for idx, opt in enumerate(options, 1):
-                    category = opt.get("category", "")
-                    name = opt.get("name", "")
-                    file_path = opt.get("file", "")
+                    category = str(opt.get("category", ""))
+                    name = str(opt.get("name", ""))
+                    file_path = str(opt.get("file", ""))
                     # 描述列显示配置描述；若为 roles 同时显示角色数量与列表
                     if category == "roles":
                         count = opt.get("roles_count")
-                        details = opt.get("details", "")
-                        parts = []
+                        details_val = opt.get("details", "")
+                        parts: list[str] = []
                         if isinstance(count, int) and count > 0:
                             parts.append(f"{count} 个角色")
-                        if details:
-                            parts.append(details)
+                        if isinstance(details_val, str) and details_val:
+                            parts.append(details_val)
                         desc_display = "\n".join(parts) if parts else ""
                     else:
-                        desc_display = opt.get("desc", "")
+                        desc_display = str(opt.get("desc", ""))
                     table.add_row(str(idx), category, name, file_path, desc_display)
 
                 Console().print(table)
@@ -412,29 +412,29 @@ def handle_builtin_config_selector(
                         index = int(choice.strip())
                         if 1 <= index <= len(options):
                             sel = options[index - 1]
-                            args = []
+                            args: list[str] = []
 
                             if sel["category"] == "agent":
                                 # jarvis-agent 支持 -f/--config（全局配置）与 -c/--agent-definition
-                                args = [sel["cmd"], "-c", sel["file"]]
+                                args = [str(sel["cmd"]), "-c", str(sel["file"])]
                                 if model_group:
-                                    args += ["-g", model_group]
+                                    args += ["-g", str(model_group)]
                                 if config_file:
-                                    args += ["-f", config_file]
+                                    args += ["-f", str(config_file)]
                                 if task:
-                                    args += ["--task", task]
+                                    args += ["--task", str(task)]
 
                             elif sel["category"] == "multi_agent":
                                 # jarvis-multi-agent 需要 -c/--config，用户输入通过 -i/--input 传递
-                                args = [sel["cmd"], "-c", sel["file"]]
+                                args = [str(sel["cmd"]), "-c", str(sel["file"])]
                                 if task:
-                                    args += ["-i", task]
+                                    args += ["-i", str(task)]
 
                             elif sel["category"] == "roles":
                                 # jarvis-platform-manager role 子命令，支持 -c/-t/-g
-                                args = [sel["cmd"], "role", "-c", sel["file"]]
+                                args = [str(sel["cmd"]), "role", "-c", str(sel["file"])]
                                 if model_group:
-                                    args += ["-g", model_group]
+                                    args += ["-g", str(model_group)]
 
                             if args:
                                 PrettyOutput.print(

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/src/jarvis/jarvis_data/config_schema.json

@@ -165,6 +165,13 @@
             "JARVIS_MAX_INPUT_TOKEN_COUNT": {
               "type": "number",
               "default": 32000
+            },
+            "ENV": {
+              "type": "object",
+              "description": "该模型组特定的环境变量，会覆盖全局 ENV 配置",
+              "additionalProperties": {
+                "type": "string"
+              }
             }
           },
           "required": [

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/src/jarvis/jarvis_git_utils/git_commiter.py

@@ -113,9 +113,9 @@ class GitCommitTool:
             # 获取文件列表
             files_cmd = ["git", "diff", "--cached", "--name-only"]
             process = subprocess.Popen(
-                files_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+                files_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
             )
-            files_output = process.communicate()[0].decode()
+            files_output = process.communicate()[0]
             files = [f.strip() for f in files_output.split("\n") if f.strip()]
             file_count = len(files)
 
@@ -124,8 +124,9 @@ class GitCommitTool:
                 ["git", "diff", "--cached", "--exit-code"],
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
+                text=True,
             )
-            diff = process.communicate()[0].decode(errors="ignore")
+            diff = process.communicate()[0]
 
             try:
                 temp_diff_file_path = None

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/src/jarvis/jarvis_platform/base.py

@@ -9,6 +9,7 @@ from typing_extensions import Self
 from rich import box  # type: ignore
 from rich.live import Live  # type: ignore
 from rich.panel import Panel  # type: ignore
+from rich.status import Status  # type: ignore
 from rich.text import Text  # type: ignore
 
 from jarvis.jarvis_utils.config import (
@@ -120,23 +121,44 @@ class BasePlatform(ABC):
         else:
             response = ""
 
-            text_content = Text(overflow="fold")  # 添加文本溢出处理
-            panel = Panel(
-                text_content,
-                title=f"[bold cyan]{self.name()}[/bold cyan]",
-                subtitle="[dim]思考中... (按 Ctrl+C 中断)[/dim]",
-                border_style="bright_blue",
-                box=box.ROUNDED,
-                expand=True  # 允许面板自动调整大小
-            )
-
             if not self.suppress_output:
                 if get_pretty_output():
-                    # 优化Live输出，减少闪烁
+                    chat_iterator = self.chat(message)
+                    first_chunk = None
+
+                    with Status(
+                        f"🤔 {self.name()} 正在思考中...", spinner="dots", console=console
+                    ):
+                        try:
+                            while True:
+                                first_chunk = next(chat_iterator)
+                                if first_chunk:
+                                    break
+                        except StopIteration:
+                            return ""
+
+                    text_content = Text(overflow="fold")
+                    panel = Panel(
+                        text_content,
+                        title=f"[bold cyan]{self.name()}[/bold cyan]",
+                        subtitle="[yellow]正在回答... (按 Ctrl+C 中断)[/yellow]",
+                        border_style="bright_blue",
+                        box=box.ROUNDED,
+                        expand=True,  # 允许面板自动调整大小
+                    )
+
                     buffer = []
                     buffer_count = 0
                     with Live(panel, refresh_per_second=4, transient=False) as live:
-                        for s in self.chat(message):
+                        # Process first chunk
+                        response += first_chunk
+                        buffer.append(first_chunk)
+                        buffer_count += 1
+
+                        # Process rest of the chunks
+                        for s in chat_iterator:
+                            if not s:
+                                continue
                             response += s  # Accumulate the full response string
                             buffer.append(s)
                             buffer_count += 1
@@ -144,13 +166,17 @@ class BasePlatform(ABC):
                             # 积累一定量或达到最后再更新，减少闪烁
                             if buffer_count >= 5 or s == "":
                                 # Append buffered content to the Text object
-                                text_content.append("".join(buffer), style="bright_white")
+                                text_content.append(
+                                    "".join(buffer), style="bright_white"
+                                )
                                 buffer.clear()
                                 buffer_count = 0
 
                                 # --- Scrolling Logic ---
                                 # Calculate available height in the panel
-                                max_text_height = console.height - 5  # Leave space for borders/titles
+                                max_text_height = (
+                                    console.height - 5
+                                )  # Leave space for borders/titles
                                 if max_text_height <= 0:
                                     max_text_height = 1
 
@@ -163,11 +189,15 @@ class BasePlatform(ABC):
                                 # If content overflows, truncate to show only the last few lines
                                 if len(lines) > max_text_height:
                                     new_text = "\n".join(
-                                        text_content.plain.splitlines()[-max_text_height:]
+                                        text_content.plain.splitlines()[
+                                            -max_text_height:
+                                        ]
                                     )
                                     text_content.plain = new_text
 
-                                panel.subtitle = "[yellow]正在回答... (按 Ctrl+C 中断)[/yellow]"
+                                panel.subtitle = (
+                                    "[yellow]正在回答... (按 Ctrl+C 中断)[/yellow]"
+                                )
                                 live.update(panel)
 
                             if is_immediate_abort() and get_interrupt():
@@ -175,7 +205,9 @@ class BasePlatform(ABC):
 
                         # Ensure any remaining content in the buffer is displayed
                         if buffer:
-                            text_content.append("".join(buffer), style="bright_white")
+                            text_content.append(
+                                "".join(buffer), style="bright_white"
+                            )
 
                         # At the end, display the entire response
                         text_content.plain = response

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/src/jarvis/jarvis_rag/cli.py

@@ -14,7 +14,7 @@ from langchain_community.document_loaders import (
 from langchain_core.document_loaders.base import BaseLoader
 from rich.markdown import Markdown
 
-from jarvis.jarvis_utils.utils import init_env
+from jarvis.jarvis_utils.utils import init_env, is_rag_installed, get_missing_rag_modules
 from jarvis.jarvis_utils.config import (
     get_rag_embedding_model,
     get_rag_use_bm25,
@@ -357,6 +357,12 @@ def retrieve(
         None, "--db-path", help="向量数据库的路径。覆盖全局配置。"
     ),
     n_results: int = typer.Option(5, "--top-n", help="要检索的文档数量。"),
+    rewrite: bool = typer.Option(
+        True,
+        "--rewrite/--no-rewrite",
+        help="是否对查询进行LLM重写以提升召回，默认开启。",
+        show_default=True,
+    ),
 ):
     """仅从RAG知识库检索文档并打印结果。"""
     try:
@@ -371,6 +377,7 @@ def retrieve(
             collection_name=collection_name,
             use_bm25=use_bm25,
             use_rerank=use_rerank,
+            use_query_rewrite=rewrite,
         )
 
         PrettyOutput.print(f"正在为问题检索文档: '{question}'", OutputType.INFO)
@@ -462,19 +469,15 @@ def query(
         raise typer.Exit(code=1)
 
 
-_RAG_INSTALLED = False
-try:
-    import langchain  # noqa
 
-    _RAG_INSTALLED = True
-except ImportError:
-    pass
 
 
 def _check_rag_dependencies():
-    if not _RAG_INSTALLED:
+    if not is_rag_installed():
+        missing = get_missing_rag_modules()
+        missing_str = f"缺少依赖: {', '.join(missing)}。" if missing else ""
         PrettyOutput.print(
-            "RAG依赖项未安装。请运行 'pip install \"jarvis-ai-assistant[rag]\"' 来使用此命令。",
+            f"RAG依赖项未安装或不完整。{missing_str}请运行 'pip install \"jarvis-ai-assistant[rag]\"' 后重试。",
             OutputType.ERROR,
         )
         raise typer.Exit(code=1)

jarvis_ai_assistant-0.3.26/src/jarvis/jarvis_rag/embedding_manager.py

@@ -0,0 +1,239 @@
+import torch
+import os
+from typing import List, cast
+from langchain_huggingface import HuggingFaceEmbeddings
+from huggingface_hub import snapshot_download
+
+
+from .cache import EmbeddingCache
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
+
+
+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
+
+        PrettyOutput.print(
+            f"初始化嵌入管理器, 模型: '{self.model_name}'...", OutputType.INFO
+        )
+
+        # 缓存的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:
+                from sentence_transformers import SentenceTransformer
+                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.print(
+                        "检测到本地模型路径但加载失败。为避免触发网络访问，已中止远程回退。\n"
+                        "请确认本地目录包含完整的 Transformers/Tokenizer 文件（如 config.json、model.safetensors、tokenizer.json/merges.txt 等），\n"
+                        "或在配置中将 embedding_model 设置为该本地目录，或将模型放置到默认的 Hugging Face 缓存目录（例如 ~/.cache/huggingface/hub）。",
+                        OutputType.ERROR,
+                    )
+                    raise
+
+                # 未发现任何本地候选，则保持原有行为：回退至远程下载
+                return HuggingFaceEmbeddings(
+                    model_name=self.model_name,
+                    model_kwargs=model_kwargs,
+                    encode_kwargs=encode_kwargs,
+                    show_progress=True,
+                )
+        except Exception as e:
+            PrettyOutput.print(
+                f"加载嵌入模型 '{self.model_name}' 时出错: {e}", OutputType.ERROR
+            )
+            PrettyOutput.print(
+                "请确保您已安装 'sentence_transformers' 和 'torch'。",
+                OutputType.WARNING,
+            )
+            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.print(
+                f"缓存未命中。正在为 {len(texts_to_embed)}/{len(texts)} 个文档计算嵌入。",
+                OutputType.INFO,
+            )
+            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.print(
+                f"缓存命中。所有 {len(texts)} 个文档的嵌入均从缓存中检索。",
+                OutputType.SUCCESS,
+            )
+
+        return cast(List[List[float]], cached_embeddings)
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        为单个查询计算嵌入。
+        查询通常不被缓存，但如果需要可以添加。
+        """
+        return self.model.embed_query(text)

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/src/jarvis/jarvis_rag/rag_pipeline.py

@@ -34,6 +34,7 @@ class JarvisRAGPipeline:
         collection_name: str = "jarvis_rag_collection",
         use_bm25: bool = True,
         use_rerank: bool = True,
+        use_query_rewrite: bool = True,
     ):
         """
         初始化RAG管道。
@@ -69,6 +70,8 @@ class JarvisRAGPipeline:
         self.collection_name = collection_name
         self.use_bm25 = use_bm25
         self.use_rerank = use_rerank
+        # 查询重写开关（默认开启，可由CLI控制）
+        self.use_query_rewrite = use_query_rewrite
 
         # 延迟加载的组件
         self._embedding_manager: Optional[EmbeddingManager] = None
@@ -229,8 +232,15 @@ class JarvisRAGPipeline:
         """
         # 0. 检测索引变更并可选更新（在重写query之前）
         self._pre_search_update_index_if_needed()
-        # 1. 将原始查询重写为多个查询
-        rewritten_queries = self._get_query_rewriter().rewrite(query_text)
+        # 1. 将原始查询重写为多个查询（可配置）
+        if self.use_query_rewrite:
+            rewritten_queries = self._get_query_rewriter().rewrite(query_text)
+        else:
+            PrettyOutput.print(
+                "已关闭查询重写，将直接使用原始查询进行检索。",
+                OutputType.INFO,
+            )
+            rewritten_queries = [query_text]
 
         # 2. 为每个重写的查询检索初始候选文档
         PrettyOutput.print(
@@ -303,8 +313,15 @@ class JarvisRAGPipeline:
         """
         # 0. 检测索引变更并可选更新（在重写query之前）
         self._pre_search_update_index_if_needed()
-        # 1. 重写查询
-        rewritten_queries = self._get_query_rewriter().rewrite(query_text)
+        # 1. 重写查询（可配置）
+        if self.use_query_rewrite:
+            rewritten_queries = self._get_query_rewriter().rewrite(query_text)
+        else:
+            PrettyOutput.print(
+                "已关闭查询重写，将直接使用原始查询进行检索。",
+                OutputType.INFO,
+            )
+            rewritten_queries = [query_text]
 
         # 2. 检索候选文档
         PrettyOutput.print(

{jarvis_ai_assistant-0.3.24 → jarvis_ai_assistant-0.3.26}/src/jarvis/jarvis_rag/reranker.py

@@ -1,9 +1,11 @@
 from typing import List
+import os
 
 from langchain.docstore.document import Document
 from sentence_transformers.cross_encoder import (  # type: ignore
     CrossEncoder,
 )
+from huggingface_hub import snapshot_download
 from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 
 
@@ -21,8 +23,28 @@ class Reranker:
             model_name (str): 要使用的Cross-Encoder模型的名称。
         """
         PrettyOutput.print(f"正在初始化重排模型: {model_name}...", OutputType.INFO)
-        self.model = CrossEncoder(model_name)
-        PrettyOutput.print("重排模型初始化成功。", OutputType.SUCCESS)
+        try:
+            local_dir = None
+
+            if os.path.isdir(model_name):
+                self.model = CrossEncoder(model_name)
+                PrettyOutput.print("重排模型初始化成功。", OutputType.SUCCESS)
+                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.print("重排模型初始化成功。", OutputType.SUCCESS)
+        except Exception as e:
+            PrettyOutput.print(f"初始化重排模型失败: {e}", OutputType.ERROR)
+            raise
 
     def rerank(
         self, query: str, documents: List[Document], top_n: int = 5