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

jarvis/jarvis_rag/README_ONLINE_MODELS.md

@@ -0,0 +1,230 @@
+# 在线模型使用指南
+
+本文档介绍如何使用在线嵌入模型和重排模型。
+
+## 概述
+
+Jarvis RAG 框架现在支持通过抽象接口使用在线模型（API）和本地模型。这允许您：
+
+- 使用本地 HuggingFace 模型（默认）
+- 使用在线 API（如 OpenAI、Cohere、Jina 等）
+- 轻松切换或混合使用不同的模型
+
+## 接口抽象
+
+### EmbeddingInterface
+
+所有嵌入模型都实现 `EmbeddingInterface` 接口，提供以下方法：
+
+- `embed_documents(texts: List[str]) -> List[List[float]]`: 为文档列表计算嵌入
+- `embed_query(text: str) -> List[float]`: 为查询计算嵌入
+
+### RerankerInterface
+
+所有重排模型都实现 `RerankerInterface` 接口，提供以下方法：
+
+- `rerank(query: str, documents: List[Document], top_n: int = 5) -> List[Document]`: 对文档进行重排
+
+## 使用示例
+
+### 使用本地模型（默认）
+
+```python
+from jarvis.jarvis_rag import JarvisRAGPipeline
+
+# 使用默认的本地模型
+pipeline = JarvisRAGPipeline()
+```
+
+### 使用在线嵌入模型
+
+#### OpenAI 嵌入模型
+
+```python
+from jarvis.jarvis_rag import JarvisRAGPipeline
+from jarvis.jarvis_rag.online_embedding import OpenAIEmbeddingModel
+
+# 创建在线嵌入模型
+embedding_model = OpenAIEmbeddingModel(
+    api_key="your-api-key",  # 或设置 OPENAI_API_KEY 环境变量
+    model_name="text-embedding-3-small"
+)
+
+# 注意：目前 JarvisRAGPipeline 的构造函数还不直接支持传入嵌入模型实例
+# 您需要修改代码或等待后续更新
+```
+
+#### Cohere 嵌入模型
+
+```python
+from jarvis.jarvis_rag.online_embedding import CohereEmbeddingModel
+
+embedding_model = CohereEmbeddingModel(
+    api_key="your-api-key",  # 或设置 COHERE_API_KEY 环境变量
+    model_name="embed-english-v3.0"
+)
+```
+
+#### EdgeFn 嵌入模型
+
+```python
+from jarvis.jarvis_rag.online_embedding import EdgeFnEmbeddingModel
+
+embedding_model = EdgeFnEmbeddingModel(
+    api_key="your-api-key",  # 或设置 EDGEFN_API_KEY 环境变量
+    model_name="BAAI/bge-m3"  # 或其他支持的模型
+)
+```
+
+### 完整示例：使用 EdgeFn 嵌入和重排模型
+
+```python
+from jarvis.jarvis_rag.online_embedding import EdgeFnEmbeddingModel
+from jarvis.jarvis_rag.online_reranker import EdgeFnReranker
+from langchain.docstore.document import Document
+
+# 初始化 EdgeFn 嵌入模型
+embedding_model = EdgeFnEmbeddingModel(
+    api_key="your-api-key",  # 或设置 EDGEFN_API_KEY 环境变量
+    model_name="BAAI/bge-m3"
+)
+
+# 初始化 EdgeFn 重排模型
+reranker = EdgeFnReranker(
+    api_key="your-api-key",  # 或设置 EDGEFN_API_KEY 环境变量
+    model_name="bge-reranker-v2-m3"
+)
+
+# 使用嵌入模型
+documents = ["文档1内容", "文档2内容", "文档3内容"]
+embeddings = embedding_model.embed_documents(documents)
+query_embedding = embedding_model.embed_query("查询文本")
+
+# 使用重排模型
+doc_list = [
+    Document(page_content="apple", metadata={"source": "doc1"}),
+    Document(page_content="banana", metadata={"source": "doc2"}),
+    Document(page_content="fruit", metadata={"source": "doc3"}),
+    Document(page_content="vegetable", metadata={"source": "doc4"}),
+]
+reranked_docs = reranker.rerank(query="Apple", documents=doc_list, top_n=3)
+```
+
+### 使用在线重排模型
+
+#### Cohere 重排模型
+
+```python
+from jarvis.jarvis_rag.online_reranker import CohereReranker
+
+reranker = CohereReranker(
+    api_key="your-api-key",  # 或设置 COHERE_API_KEY 环境变量
+    model_name="rerank-english-v3.0"
+)
+```
+
+#### Jina 重排模型
+
+```python
+from jarvis.jarvis_rag.online_reranker import JinaReranker
+
+reranker = JinaReranker(
+    api_key="your-api-key",  # 或设置 JINA_API_KEY 环境变量
+    model_name="jina-reranker-v1-base-en"
+)
+```
+
+#### EdgeFn 重排模型
+
+```python
+from jarvis.jarvis_rag.online_reranker import EdgeFnReranker
+
+reranker = EdgeFnReranker(
+    api_key="your-api-key",  # 或设置 EDGEFN_API_KEY 环境变量
+    model_name="bge-reranker-v2-m3"  # 或其他支持的模型
+)
+```
+
+## 实现自定义在线模型
+
+### 实现自定义嵌入模型
+
+```python
+from jarvis.jarvis_rag.embedding_interface import EmbeddingInterface
+from typing import List
+
+class MyCustomEmbeddingModel(EmbeddingInterface):
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+        # 初始化您的API客户端
+    
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        # 实现您的API调用逻辑
+        embeddings = []
+        for text in texts:
+            # 调用您的API
+            embedding = self._call_your_api(text)
+            embeddings.append(embedding)
+        return embeddings
+    
+    def embed_query(self, text: str) -> List[float]:
+        # 实现查询嵌入逻辑
+        return self._call_your_api(text)
+```
+
+### 实现自定义重排模型
+
+```python
+from jarvis.jarvis_rag.reranker_interface import RerankerInterface
+from langchain.docstore.document import Document
+from typing import List
+
+class MyCustomReranker(RerankerInterface):
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+        # 初始化您的API客户端
+    
+    def rerank(
+        self, query: str, documents: List[Document], top_n: int = 5
+    ) -> List[Document]:
+        # 实现您的重排逻辑
+        doc_texts = [doc.page_content for doc in documents]
+        scores = self._call_your_api(query, doc_texts)
+        
+        # 根据分数排序
+        doc_with_scores = list(zip(documents, scores))
+        doc_with_scores.sort(key=lambda x: x[1], reverse=True)
+        
+        return [doc for doc, _ in doc_with_scores[:top_n]]
+```
+
+## 向后兼容性
+
+为了保持向后兼容性，原有的类名仍然可用：
+
+- `EmbeddingManager` 是 `LocalEmbeddingModel` 的别名
+- `Reranker` 是 `LocalReranker` 的别名
+
+现有代码无需修改即可继续工作。
+
+## 注意事项
+
+1. **API 密钥安全**: 建议使用环境变量存储 API 密钥，而不是硬编码在代码中。
+
+2. **成本考虑**: 在线模型通常按使用量收费，请根据您的需求选择合适的模型。
+
+3. **性能**: 在线模型需要网络请求，可能比本地模型慢，但通常提供更好的效果。
+
+4. **依赖安装**: 使用特定的在线模型需要安装相应的 Python 包：
+   - OpenAI: `pip install openai`
+   - Cohere: `pip install cohere`
+   - EdgeFn (嵌入和重排): `pip install requests` (已包含在大多数环境中)
+   - Jina: `pip install requests` (已包含在大多数环境中)
+
+## 未来扩展
+
+框架设计允许轻松添加新的在线模型提供商。如果您需要支持其他 API，可以：
+
+1. 继承 `OnlineEmbeddingModel` 或 `OnlineReranker` 基类
+2. 实现 `_call_api` 方法
+3. 在 `__init__.py` 中导出新类

jarvis/jarvis_rag/__init__.py

@@ -1,11 +1,64 @@
 """
 Jarvis RAG 框架
 
-一个灵活的RAG管道，具有可插拔的远程LLM和本地带缓存的嵌入模型。
+一个灵活的RAG管道，具有可插拔的远程LLM和本地/在线嵌入模型和重排模型。
 """
 
-from .rag_pipeline import JarvisRAGPipeline
+from .embedding_interface import EmbeddingInterface
+
+# 从新的目录结构导入
+from .embeddings import EmbeddingManager  # 向后兼容别名
+from .embeddings import EmbeddingRegistry
+from .embeddings import LocalEmbeddingModel
 from .llm_interface import LLMInterface
-from .embedding_manager import EmbeddingManager
+from .rag_pipeline import JarvisRAGPipeline
+from .reranker_interface import RerankerInterface
+from .rerankers import LocalReranker
+from .rerankers import Reranker  # 向后兼容别名
+from .rerankers import RerankerRegistry
+
+# 在线模型实现（可选导入）
+try:
+    from .embeddings import CohereEmbeddingModel
+    from .embeddings import EdgeFnEmbeddingModel
+    from .embeddings import OnlineEmbeddingModel
+    from .embeddings import OpenAIEmbeddingModel
+    from .rerankers import CohereReranker
+    from .rerankers import EdgeFnReranker
+    from .rerankers import JinaReranker
+    from .rerankers import OnlineReranker
 
-__all__ = ["JarvisRAGPipeline", "LLMInterface", "EmbeddingManager"]
+    __all__ = [
+        "JarvisRAGPipeline",
+        "LLMInterface",
+        "EmbeddingInterface",
+        "EmbeddingManager",  # 向后兼容别名
+        "LocalEmbeddingModel",
+        "EmbeddingRegistry",
+        "RerankerInterface",
+        "Reranker",  # 向后兼容别名
+        "LocalReranker",
+        "RerankerRegistry",
+        "OnlineEmbeddingModel",
+        "OpenAIEmbeddingModel",
+        "CohereEmbeddingModel",
+        "EdgeFnEmbeddingModel",
+        "OnlineReranker",
+        "CohereReranker",
+        "JinaReranker",
+        "EdgeFnReranker",
+    ]
+except ImportError:
+    # 如果在线模型依赖未安装，只导出基础接口
+    __all__ = [
+        "JarvisRAGPipeline",
+        "LLMInterface",
+        "EmbeddingInterface",
+        "EmbeddingManager",  # 向后兼容别名
+        "LocalEmbeddingModel",
+        "EmbeddingRegistry",
+        "RerankerInterface",
+        "Reranker",  # 向后兼容别名
+        "LocalReranker",
+        "RerankerRegistry",
+    ]

jarvis/jarvis_rag/cache.py

@@ -1,5 +1,7 @@
 import hashlib
-from typing import List, Optional, Any
+from typing import Any
+from typing import List
+from typing import Optional
 
 from diskcache import Cache
 

jarvis/jarvis_rag/cli.py

@@ -1,25 +1,26 @@
+import mimetypes
 import os
 import sys
 from pathlib import Path
-from typing import Optional, List, Tuple
-import mimetypes
+from typing import List
+from typing import Optional
+from typing import Tuple
 
-import pathspec  # type: ignore
+import pathspec
 import typer
 from langchain.docstore.document import Document
-from langchain_community.document_loaders import (
-    TextLoader,
-    UnstructuredMarkdownLoader,
-)
+from langchain_community.document_loaders import TextLoader
+from langchain_community.document_loaders import UnstructuredMarkdownLoader
 from langchain_core.document_loaders.base import BaseLoader
 from rich.markdown import Markdown
 
-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,
-    get_rag_use_rerank,
-)
+from jarvis.jarvis_utils.config import get_rag_embedding_model
+from jarvis.jarvis_utils.config import get_rag_use_bm25
+from jarvis.jarvis_utils.config import get_rag_use_rerank
+from jarvis.jarvis_utils.output import PrettyOutput
+from jarvis.jarvis_utils.utils import get_missing_rag_modules
+from jarvis.jarvis_utils.utils import init_env
+from jarvis.jarvis_utils.utils import is_rag_installed
 
 
 def is_likely_text_file(file_path: Path) -> bool:
@@ -70,7 +71,7 @@ class _CustomPlatformLLM(LLMInterface):
 
     def __init__(self, platform: BasePlatform):
         self.platform = platform
-        print(
+        PrettyOutput.auto_print(
             f"ℹ️ 使用自定义LLM: 平台='{platform.platform_name()}', 模型='{platform.name()}'"
         )
 
@@ -86,15 +87,13 @@ def _create_custom_llm(platform_name: str, model_name: str) -> Optional[LLMInter
         registry = PlatformRegistry.get_global_platform_registry()
         platform_instance = registry.create_platform(platform_name)
         if not platform_instance:
-            print(
-                f"❌ 错误: 平台 '{platform_name}' 未找到。"
-            )
+            PrettyOutput.auto_print(f"❌ 错误: 平台 '{platform_name}' 未找到。")
             return None
         platform_instance.set_model_name(model_name)
         platform_instance.set_suppress_output(True)
         return _CustomPlatformLLM(platform_instance)
     except Exception as e:
-        print(f"❌ 创建自定义LLM时出错: {e}")
+        PrettyOutput.auto_print(f"❌ 创建自定义LLM时出错: {e}")
         return None
 
 
@@ -118,14 +117,10 @@ def _load_ragignore_spec() -> Tuple[Optional[pathspec.PathSpec], Optional[Path]]
             with open(ignore_file_to_use, "r", encoding="utf-8") as f:
                 patterns = f.read().splitlines()
             spec = pathspec.PathSpec.from_lines("gitwildmatch", patterns)
-            print(
-                f"✅ 加载忽略规则: {ignore_file_to_use}"
-            )
+            PrettyOutput.auto_print(f"✅ 加载忽略规则: {ignore_file_to_use}")
             return spec, project_root_path
         except Exception as e:
-            print(
-                f"⚠️ 加载 {ignore_file_to_use.name} 文件失败: {e}"
-            )
+            PrettyOutput.auto_print(f"⚠️ 加载 {ignore_file_to_use.name} 文件失败: {e}")
 
     return None, None
 
@@ -176,7 +171,7 @@ def add_documents(
                 continue
 
             if path.is_dir():
-                print(f"ℹ️ 正在扫描目录: {path}")
+                PrettyOutput.auto_print(f"ℹ️ 正在扫描目录: {path}")
                 for item in path.rglob("*"):
                     if item.is_file() and is_likely_text_file(item):
                         files_to_process.add(item)
@@ -184,12 +179,10 @@ def add_documents(
                 if is_likely_text_file(path):
                     files_to_process.add(path)
                 else:
-                    print(
-                        f"⚠️ 跳过可能的二进制文件: {path}"
-                    )
+                    PrettyOutput.auto_print(f"⚠️ 跳过可能的二进制文件: {path}")
 
     if not files_to_process:
-        print("⚠️ 在指定路径中未找到任何文本文件。")
+        PrettyOutput.auto_print("⚠️ 在指定路径中未找到任何文本文件。")
         return
 
     # 使用 .ragignore 过滤文件
@@ -210,20 +203,16 @@ def add_documents(
 
         ignored_count = initial_count - len(retained_files)
         if ignored_count > 0:
-            print(
+            PrettyOutput.auto_print(
                 f"ℹ️ 根据 .ragignore 规则过滤掉 {ignored_count} 个文件。"
             )
         files_to_process = retained_files
 
     if not files_to_process:
-        print(
-            "⚠️ 所有找到的文本文件都被忽略规则过滤掉了。"
-        )
+        PrettyOutput.auto_print("⚠️ 所有找到的文本文件都被忽略规则过滤掉了。")
         return
 
-    print(
-        f"ℹ️ 发现 {len(files_to_process)} 个独立文件待处理。"
-    )
+    PrettyOutput.auto_print(f"ℹ️ 发现 {len(files_to_process)} 个独立文件待处理。")
 
     try:
         pipeline = JarvisRAGPipeline(
@@ -250,38 +239,36 @@ def add_documents(
                 docs_batch.extend(loader.load())
                 loaded_msgs.append(f"已加载: {file_path} (文件 {i + 1}/{total_files})")
             except Exception as e:
-                print(f"⚠️ 加载失败 {file_path}: {e}")
+                PrettyOutput.auto_print(f"⚠️ 加载失败 {file_path}: {e}")
 
             # 当批处理已满或是最后一个文件时处理批处理
             if docs_batch and (len(docs_batch) >= batch_size or (i + 1) == total_files):
                 if loaded_msgs:
-                    joined_msgs = '\n'.join(loaded_msgs)
-                    print(f"ℹ️ {joined_msgs}")
+                    joined_msgs = "\n".join(loaded_msgs)
+                    PrettyOutput.auto_print(f"ℹ️ {joined_msgs}")
                     loaded_msgs = []
-                print(
+                PrettyOutput.auto_print(
                     f"ℹ️ 正在处理批次，包含 {len(docs_batch)} 个文档..."
                 )
                 pipeline.add_documents(docs_batch)
                 total_docs_added += len(docs_batch)
-                print(
-                    f"✅ 成功添加 {len(docs_batch)} 个文档。"
-                )
+                PrettyOutput.auto_print(f"✅ 成功添加 {len(docs_batch)} 个文档。")
                 docs_batch = []  # 清空批处理
 
         # 最后统一打印可能残留的"已加载"信息
         if loaded_msgs:
-            print(f"ℹ️ {chr(10).join(loaded_msgs)}")
+            PrettyOutput.auto_print(f"ℹ️ {chr(10).join(loaded_msgs)}")
             loaded_msgs = []
         if total_docs_added == 0:
-            print("❌ 未能成功加载任何文档。")
+            PrettyOutput.auto_print("❌ 未能成功加载任何文档。")
             raise typer.Exit(code=1)
 
-        print(
+        PrettyOutput.auto_print(
             f"✅ 成功将 {total_docs_added} 个文档的内容添加至集合 '{collection_name}'。"
         )
 
     except Exception as e:
-        print(f"❌ 发生严重错误: {e}")
+        PrettyOutput.auto_print(f"❌ 发生严重错误: {e}")
         raise typer.Exit(code=1)
 
 
@@ -308,7 +295,7 @@ def list_documents(
         results = collection.get()  # 获取集合中的所有项目
 
         if not results or not results["metadatas"]:
-            print("ℹ️ 知识库中没有找到任何文档。")
+            PrettyOutput.auto_print("ℹ️ 知识库中没有找到任何文档。")
             return
 
         # 从元数据中提取唯一的源文件路径
@@ -320,20 +307,18 @@ def list_documents(
                     sources.add(source)
 
         if not sources:
-            print(
-                "ℹ️ 知识库中没有找到任何带有源信息的文档。"
-            )
+            PrettyOutput.auto_print("ℹ️ 知识库中没有找到任何带有源信息的文档。")
             return
 
         # 避免在循环中逐条打印，先拼接后统一打印
         lines = [f"知识库 '{collection_name}' 中共有 {len(sources)} 个独立文档:"]
         for i, source in enumerate(sorted(list(sources)), 1):
             lines.append(f"  {i}. {source}")
-        joined_lines = '\n'.join(lines)
-        print(f"ℹ️ {joined_lines}")
+        joined_lines = "\n".join(lines)
+        PrettyOutput.auto_print(f"ℹ️ {joined_lines}")
 
     except Exception as e:
-        print(f"❌ 发生错误: {e}")
+        PrettyOutput.auto_print(f"❌ 发生错误: {e}")
         raise typer.Exit(code=1)
 
 
@@ -379,16 +364,14 @@ def retrieve(
             use_query_rewrite=rewrite,
         )
 
-        print(f"ℹ️ 正在为问题检索文档: '{question}'")
+        PrettyOutput.auto_print(f"ℹ️ 正在为问题检索文档: '{question}'")
         retrieved_docs = pipeline.retrieve_only(question, n_results=n_results)
 
         if not retrieved_docs:
-            print("ℹ️ 未找到相关文档。")
+            PrettyOutput.auto_print("ℹ️ 未找到相关文档。")
             return
 
-        print(
-            f"✅ 成功检索到 {len(retrieved_docs)} 个文档:"
-        )
+        PrettyOutput.auto_print(f"✅ 成功检索到 {len(retrieved_docs)} 个文档:")
         from jarvis.jarvis_utils.globals import console
 
         for i, doc in enumerate(retrieved_docs, 1):
@@ -399,7 +382,7 @@ def retrieve(
             console.print(Markdown(f"```\n{content}\n```"))
 
     except Exception as e:
-        print(f"❌ 发生错误: {e}")
+        PrettyOutput.auto_print(f"❌ 发生错误: {e}")
         raise typer.Exit(code=1)
 
 
@@ -436,7 +419,7 @@ def query(
 ):
     """查询RAG知识库并打印答案。"""
     if model and not platform:
-        print("❌ 错误: --model 需要指定 --platform。")
+        PrettyOutput.auto_print("❌ 错误: --model 需要指定 --platform。")
         raise typer.Exit(code=1)
 
     try:
@@ -458,24 +441,21 @@ def query(
             use_rerank=use_rerank,
         )
 
-        print(f"ℹ️ 正在查询: '{question}'")
+        PrettyOutput.auto_print(f"ℹ️ 正在查询: '{question}'")
         answer = pipeline.query(question)
 
-        print(f"✅ {answer}")
+        PrettyOutput.auto_print(f"✅ {answer}")
 
     except Exception as e:
-        print(f"❌ 发生错误: {e}")
+        PrettyOutput.auto_print(f"❌ 发生错误: {e}")
         raise typer.Exit(code=1)
 
 
-
-
-
 def _check_rag_dependencies():
     if not is_rag_installed():
         missing = get_missing_rag_modules()
         missing_str = f"缺少依赖: {', '.join(missing)}。" if missing else ""
-        print(
+        PrettyOutput.auto_print(
             f"❌ RAG依赖项未安装或不完整。{missing_str}请运行 'pip install \"jarvis-ai-assistant[rag]\"' 后重试。"
         )
         raise typer.Exit(code=1)

jarvis/jarvis_rag/embedding_interface.py

@@ -0,0 +1,39 @@
+from abc import ABC
+from abc import abstractmethod
+from typing import List
+
+
+class EmbeddingInterface(ABC):
+    """
+    嵌入模型接口的抽象基类。
+
+    该类定义了嵌入模型的标准接口，支持本地模型和在线模型（API）的实现。
+    任何嵌入模型提供商（如HuggingFace本地模型、OpenAI API、Cohere API等）
+    都应作为该接口的子类来实现。
+    """
+
+    @abstractmethod
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        为文档列表计算嵌入。
+
+        参数:
+            texts: 要嵌入的文档（字符串）列表。
+
+        返回:
+            一个嵌入列表，每个文档对应一个嵌入向量。
+        """
+        pass
+
+    @abstractmethod
+    def embed_query(self, text: str) -> List[float]:
+        """
+        为单个查询计算嵌入。
+
+        参数:
+            text: 要嵌入的查询文本。
+
+        返回:
+            查询的嵌入向量。
+        """
+        pass