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

jarvis/jarvis_rag/README_CONFIG.md

@@ -0,0 +1,314 @@
+# RAG 配置指南
+
+本文档介绍如何配置 RAG 系统的嵌入模型和重排模型。
+
+## 配置方式
+
+RAG 配置通过 `rag` 环境变量或 `GLOBAL_CONFIG_DATA` 中的 `rag` 键进行设置。
+
+## 配置结构
+
+### 基本配置
+
+```python
+rag = {
+    "embedding_model": "BAAI/bge-m3",           # 模型名称
+    "embedding_type": "LocalEmbeddingModel",     # 模型实现类型
+    "embedding_max_length": 512,                # 嵌入模型最大输入长度（token数）
+    "rerank_model": "BAAI/bge-reranker-v2-m3",  # 重排模型名称
+    "reranker_type": "LocalReranker",            # 重排模型实现类型
+    "reranker_max_length": 512,                 # 重排模型最大输入长度（token数）
+    "use_bm25": True,                            # 是否使用BM25
+    "use_rerank": True,                          # 是否使用重排
+}
+```
+
+### 高级配置（带额外参数）
+
+```python
+rag = {
+    "embedding_model": "BAAI/bge-m3",
+    "embedding_type": "LocalEmbeddingModel",
+    "embedding_config": {
+        # 额外的嵌入模型配置参数
+        # 这些参数会传递给模型的构造函数
+    },
+    "rerank_model": "bge-reranker-v2-m3",
+    "reranker_type": "EdgeFnReranker",
+    "reranker_config": {
+        # 额外的重排模型配置参数
+        "api_key": "your-api-key",  # 如果使用在线模型
+        "base_url": "https://api.example.com/v1/rerank",
+    },
+    "use_bm25": True,
+    "use_rerank": True,
+}
+```
+
+## 配置组
+
+类似于模型组，RAG 也支持配置组：
+
+```python
+rag_group = "text"  # 或 "code"
+
+rag_groups = [
+    {
+        "text": {
+            "embedding_model": "BAAI/bge-m3",
+            "embedding_type": "LocalEmbeddingModel",
+            "rerank_model": "BAAI/bge-reranker-v2-m3",
+            "reranker_type": "LocalReranker",
+            "use_bm25": True,
+            "use_rerank": True,
+        }
+    },
+    {
+        "code": {
+            "embedding_model": "Qodo/Qodo-Embed-1-1.5B",
+            "embedding_type": "LocalEmbeddingModel",
+            "use_bm25": False,
+            "use_rerank": False,
+        }
+    },
+]
+```
+
+## 可用的模型类型
+
+### 嵌入模型类型
+
+- `LocalEmbeddingModel` - 本地 HuggingFace 模型（默认）
+- `OpenAIEmbeddingModel` - OpenAI API
+- `CohereEmbeddingModel` - Cohere API
+- `EdgeFnEmbeddingModel` - EdgeFn API
+- 自定义模型（通过动态加载）
+
+### 重排模型类型
+
+- `LocalReranker` - 本地 CrossEncoder 模型（默认）
+- `CohereReranker` - Cohere API
+- `JinaReranker` - Jina API
+- `EdgeFnReranker` - EdgeFn API
+- 自定义模型（通过动态加载）
+
+## 配置示例
+
+### 示例 1: 使用本地模型（默认）
+
+```python
+# 使用默认配置，无需设置
+# 或显式设置：
+rag = {
+    "embedding_model": "BAAI/bge-m3",
+    "embedding_type": "LocalEmbeddingModel",
+    "embedding_max_length": 512,  # 根据模型调整
+    "rerank_model": "BAAI/bge-reranker-v2-m3",
+    "reranker_type": "LocalReranker",
+    "reranker_max_length": 512,  # 根据模型调整
+}
+```
+
+### 示例 2: 使用 OpenAI 嵌入模型
+
+```python
+rag = {
+    "embedding_model": "text-embedding-3-small",
+    "embedding_type": "OpenAIEmbeddingModel",
+    "embedding_config": {
+        "api_key": "your-openai-api-key",  # 或设置 OPENAI_API_KEY 环境变量
+    },
+    "rerank_model": "BAAI/bge-reranker-v2-m3",
+    "reranker_type": "LocalReranker",
+}
+```
+
+### 示例 3: 使用 EdgeFn 在线模型
+
+```python
+rag = {
+    "embedding_model": "BAAI/bge-m3",
+    "embedding_type": "EdgeFnEmbeddingModel",
+    "embedding_config": {
+        "api_key": "your-edgefn-api-key",  # 或设置 EDGEFN_API_KEY 环境变量
+    },
+    "rerank_model": "bge-reranker-v2-m3",
+    "reranker_type": "EdgeFnReranker",
+    "reranker_config": {
+        "api_key": "your-edgefn-api-key",
+    },
+}
+```
+
+### 示例 4: 混合使用（本地嵌入 + 在线重排）
+
+```python
+rag = {
+    "embedding_model": "BAAI/bge-m3",
+    "embedding_type": "LocalEmbeddingModel",
+    "rerank_model": "rerank-english-v3.0",
+    "reranker_type": "CohereReranker",
+    "reranker_config": {
+        "api_key": "your-cohere-api-key",  # 或设置 COHERE_API_KEY 环境变量
+    },
+}
+```
+
+### 示例 5: 使用自定义模型
+
+```python
+# 首先在 ~/.jarvis/embeddings/ 或 ~/.jarvis/rerankers/ 中创建自定义模型
+# 然后配置使用：
+
+rag = {
+    "embedding_model": "my-custom-model",
+    "embedding_type": "MyCustomEmbeddingModel",  # 自定义类名
+    "embedding_config": {
+        "api_key": "your-api-key",
+        "custom_param": "value",
+    },
+}
+```
+
+## 配置优先级
+
+配置的优先级顺序（从高到低）：
+
+1. **顶级 rag 设置** - 直接设置的配置项
+2. **rag_group 组配置** - 通过组名选择的配置
+3. **代码默认值** - 硬编码的默认值
+
+## 配置函数
+
+系统提供了以下配置函数：
+
+### 嵌入模型配置
+
+- `get_rag_embedding_model() -> str`: 获取嵌入模型名称
+- `get_rag_embedding_type() -> str`: 获取嵌入模型类型
+- `get_rag_embedding_max_length() -> int`: 获取嵌入模型最大输入长度（token数）
+- `get_rag_embedding_config() -> Dict[str, Any]`: 获取嵌入模型额外配置
+
+### 重排模型配置
+
+- `get_rag_rerank_model() -> str`: 获取重排模型名称
+- `get_rag_reranker_type() -> str`: 获取重排模型类型
+- `get_rag_reranker_max_length() -> int`: 获取重排模型最大输入长度（token数）
+- `get_rag_reranker_config() -> Dict[str, Any]`: 获取重排模型额外配置
+
+### 其他配置
+
+- `get_rag_use_bm25() -> bool`: 是否使用BM25
+- `get_rag_use_rerank() -> bool`: 是否使用重排
+- `get_rag_embedding_cache_path() -> str`: 嵌入缓存路径
+- `get_rag_vector_db_path() -> str`: 向量数据库路径
+
+## 在代码中使用配置
+
+### 方式 1: 使用 Registry 从配置创建
+
+```python
+from jarvis.jarvis_rag.embeddings import EmbeddingRegistry
+from jarvis.jarvis_rag.rerankers import RerankerRegistry
+
+# 从配置创建模型实例
+embedding = EmbeddingRegistry.create_from_config()
+reranker = RerankerRegistry.create_from_config()
+```
+
+### 方式 2: 手动创建（覆盖配置）
+
+```python
+from jarvis.jarvis_rag.embeddings import EmbeddingRegistry
+
+registry = EmbeddingRegistry.get_global_registry()
+embedding = registry.create_embedding(
+    "LocalEmbeddingModel",
+    model_name="custom-model",
+    cache_dir="/custom/cache"
+)
+```
+
+### 方式 3: 在 RAG Pipeline 中使用
+
+```python
+from jarvis.jarvis_rag import JarvisRAGPipeline
+
+# RAG Pipeline 会自动从配置创建模型
+pipeline = JarvisRAGPipeline()
+# 内部会调用 EmbeddingRegistry.create_from_config() 和 RerankerRegistry.create_from_config()
+```
+
+## 环境变量
+
+对于在线模型，建议使用环境变量存储 API 密钥：
+
+```bash
+export OPENAI_API_KEY="your-key"
+export COHERE_API_KEY="your-key"
+export EDGEFN_API_KEY="your-key"
+export JINA_API_KEY="your-key"
+```
+
+然后在配置中只需要指定类型，不需要在 `embedding_config` 或 `reranker_config` 中提供 API 密钥。
+
+## 向后兼容性
+
+如果不指定 `embedding_type` 或 `reranker_type`，系统会：
+
+1. 尝试从配置创建（如果配置了类型）
+2. 回退到使用 `LocalEmbeddingModel` 和 `LocalReranker`（传统方式）
+
+这确保了现有代码无需修改即可继续工作。
+
+## 调试
+
+如果模型创建失败，检查：
+
+1. 配置中的类型名称是否正确（区分大小写）
+2. 模型类型是否已在注册表中注册
+3. 配置参数是否正确（特别是 API 密钥）
+4. 查看控制台输出的错误信息
+
+## 完整配置示例
+
+```python
+# 在环境变量或配置文件中设置
+rag = {
+    # 嵌入模型配置
+    "embedding_model": "BAAI/bge-m3",
+    "embedding_type": "LocalEmbeddingModel",
+    "embedding_max_length": 512,  # 模型最大输入token数，用于文档分割
+    "embedding_config": {
+        # 本地模型通常不需要额外配置
+        # 在线模型可能需要 api_key 等
+    },
+    
+    # 重排模型配置
+    "rerank_model": "bge-reranker-v2-m3",
+    "reranker_type": "LocalReranker",
+    "reranker_max_length": 512,  # 模型最大输入token数
+    "reranker_config": {
+        # 配置参数
+    },
+    
+    # 功能开关
+    "use_bm25": True,
+    "use_rerank": True,
+}
+
+# 或使用配置组
+rag_group = "text"
+rag_groups = [
+    {
+        "text": {
+            "embedding_model": "BAAI/bge-m3",
+            "embedding_type": "LocalEmbeddingModel",
+            "rerank_model": "BAAI/bge-reranker-v2-m3",
+            "reranker_type": "LocalReranker",
+            "use_bm25": True,
+            "use_rerank": True,
+        }
+    },
+]
+```

jarvis/jarvis_rag/README_DYNAMIC_LOADING.md

@@ -0,0 +1,311 @@
+# 动态加载自定义模型指南
+
+本文档介绍如何使用动态加载功能添加自定义嵌入模型和重排模型。
+
+## 概述
+
+Jarvis RAG 框架支持动态加载自定义模型实现，类似于 `jarvis_platform` 的注册表机制。您可以在用户数据目录中创建自定义模型实现，系统会自动发现并加载它们。
+
+## 目录结构
+
+自定义模型应该放在以下目录：
+
+- **嵌入模型**: `~/.jarvis/embeddings/` (或通过 `get_data_dir()` 获取的目录下的 `embeddings/`)
+- **重排模型**: `~/.jarvis/rerankers/` (或通过 `get_data_dir()` 获取的目录下的 `rerankers/`)
+
+这些目录会在首次使用时自动创建。
+
+## 实现自定义嵌入模型
+
+### 1. 创建模型文件
+
+在 `~/.jarvis/embeddings/` 目录下创建一个 Python 文件，例如 `my_custom_embedding.py`:
+
+```python
+from typing import List, Optional
+from jarvis.jarvis_rag.embedding_interface import EmbeddingInterface
+from jarvis.jarvis_rag.embeddings.base import OnlineEmbeddingModel
+
+class MyCustomEmbeddingModel(OnlineEmbeddingModel):
+    """
+    自定义嵌入模型实现示例。
+    """
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model_name: str = "my-model",
+        base_url: str = "https://api.example.com/v1/embeddings",
+        **kwargs,
+    ):
+        super().__init__(
+            api_key=api_key,
+            api_key_env="MY_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]]:
+        """
+        实现您的 API 调用逻辑。
+        """
+        import requests
+        
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        
+        payload = {
+            "model": self.model_name,
+            "input": texts if len(texts) > 1 else texts[0],
+        }
+        
+        response = requests.post(
+            self.base_url, headers=headers, json=payload, timeout=60
+        )
+        response.raise_for_status()
+        
+        data = response.json()
+        # 根据您的 API 响应格式解析嵌入向量
+        return data["embeddings"]
+```
+
+### 2. 必需的方法
+
+自定义嵌入模型必须实现以下方法：
+
+- `embed_documents(texts: List[str]) -> List[List[float]]`: 为文档列表计算嵌入
+- `embed_query(text: str) -> List[float]`: 为查询计算嵌入
+
+如果继承自 `OnlineEmbeddingModel`，只需要实现 `_call_api` 方法。
+
+### 3. 使用自定义模型
+
+```python
+from jarvis.jarvis_rag.embeddings import EmbeddingRegistry
+
+# 获取注册表
+registry = EmbeddingRegistry.get_global_registry()
+
+# 查看可用的模型
+print(registry.get_available_embeddings())
+
+# 创建自定义模型实例
+embedding = registry.create_embedding(
+    "MyCustomEmbeddingModel",
+    api_key="your-api-key",
+    model_name="my-model"
+)
+
+# 使用模型
+embeddings = embedding.embed_documents(["文档1", "文档2"])
+query_embedding = embedding.embed_query("查询文本")
+```
+
+## 实现自定义重排模型
+
+### 1. 创建模型文件
+
+在 `~/.jarvis/rerankers/` 目录下创建一个 Python 文件，例如 `my_custom_reranker.py`:
+
+```python
+from typing import List, Optional
+from langchain.docstore.document import Document
+from jarvis.jarvis_rag.reranker_interface import RerankerInterface
+from jarvis.jarvis_rag.rerankers.base import OnlineReranker
+
+class MyCustomReranker(OnlineReranker):
+    """
+    自定义重排模型实现示例。
+    """
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model_name: str = "my-reranker",
+        base_url: str = "https://api.example.com/v1/rerank",
+        **kwargs,
+    ):
+        super().__init__(
+            api_key=api_key,
+            api_key_env="MY_API_KEY",
+            base_url=base_url,
+            model_name=model_name,
+            **kwargs,
+        )
+    
+    def _call_api(
+        self, query: str, documents: List[str]
+    ) -> List[tuple[int, float]]:
+        """
+        实现您的 API 调用逻辑。
+        返回 (索引, 分数) 元组列表，按分数降序排序。
+        """
+        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, headers=headers, json=payload, timeout=60
+        )
+        response.raise_for_status()
+        
+        data = response.json()
+        
+        # 解析结果，返回 (索引, 分数) 列表
+        results = [
+            (item["index"], item["score"])
+            for item in data["results"]
+        ]
+        results.sort(key=lambda x: x[1], reverse=True)
+        
+        return results
+```
+
+### 2. 必需的方法
+
+自定义重排模型必须实现以下方法：
+
+- `rerank(query: str, documents: List[Document], top_n: int = 5) -> List[Document]`: 对文档进行重排
+
+如果继承自 `OnlineReranker`，只需要实现 `_call_api` 方法。
+
+### 3. 使用自定义模型
+
+```python
+from jarvis.jarvis_rag.rerankers import RerankerRegistry
+from langchain.docstore.document import Document
+
+# 获取注册表
+registry = RerankerRegistry.get_global_registry()
+
+# 查看可用的模型
+print(registry.get_available_rerankers())
+
+# 创建自定义模型实例
+reranker = registry.create_reranker(
+    "MyCustomReranker",
+    api_key="your-api-key",
+    model_name="my-reranker"
+)
+
+# 使用模型
+documents = [
+    Document(page_content="文档1", metadata={"source": "doc1"}),
+    Document(page_content="文档2", metadata={"source": "doc2"}),
+]
+reranked = reranker.rerank("查询", documents, top_n=3)
+```
+
+## 实现本地模型
+
+您也可以实现完全自定义的本地模型（不继承基类）：
+
+```python
+from typing import List
+from jarvis.jarvis_rag.embedding_interface import EmbeddingInterface
+
+class MyLocalEmbeddingModel(EmbeddingInterface):
+    """
+    完全自定义的本地嵌入模型。
+    """
+    
+    def __init__(self, model_path: str):
+        # 加载您的本地模型
+        self.model = load_your_model(model_path)
+    
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        # 实现文档嵌入逻辑
+        return [self.model.encode(text) for text in texts]
+    
+    def embed_query(self, text: str) -> List[float]:
+        # 实现查询嵌入逻辑
+        return self.model.encode(text)
+```
+
+## 验证实现
+
+系统会自动验证您的实现：
+
+1. **方法检查**: 确保实现了所有必需的方法
+2. **参数检查**: 确保方法参数匹配
+3. **类型检查**: 确保类继承自正确的接口
+
+如果验证失败，系统会输出错误信息，但不会中断其他模型的加载。
+
+## 注册表 API
+
+### EmbeddingRegistry
+
+- `get_global_registry() -> EmbeddingRegistry`: 获取全局注册表实例
+- `get_available_embeddings() -> List[str]`: 获取所有可用的嵌入模型名称
+- `create_embedding(name: str, *args, **kwargs) -> Optional[EmbeddingInterface]`: 创建嵌入模型实例
+- `register_embedding(name: str, embedding_class: Type[EmbeddingInterface])`: 手动注册模型类
+
+### RerankerRegistry
+
+- `get_global_registry() -> RerankerRegistry`: 获取全局注册表实例
+- `get_available_rerankers() -> List[str]`: 获取所有可用的重排模型名称
+- `create_reranker(name: str, *args, **kwargs) -> Optional[RerankerInterface]`: 创建重排模型实例
+- `register_reranker(name: str, reranker_class: Type[RerankerInterface])`: 手动注册模型类
+
+## 注意事项
+
+1. **文件命名**: 文件名应该使用下划线命名（如 `my_custom_model.py`），类名使用驼峰命名（如 `MyCustomModel`）
+
+2. **导入依赖**: 确保您的自定义模型所需的依赖已安装
+
+3. **错误处理**: 建议在 `_call_api` 方法中添加适当的错误处理
+
+4. **性能考虑**: 对于批量处理，考虑实现批处理逻辑以提高效率
+
+5. **缓存**: 如果需要缓存，可以在自定义类中添加缓存逻辑
+
+## 示例：完整工作流
+
+```python
+# 1. 创建自定义模型文件 ~/.jarvis/embeddings/my_api.py
+# (参考上面的示例代码)
+
+# 2. 在代码中使用
+from jarvis.jarvis_rag.embeddings import EmbeddingRegistry
+
+registry = EmbeddingRegistry.get_global_registry()
+
+# 检查模型是否已加载
+if "MyApiEmbeddingModel" in registry.get_available_embeddings():
+    # 创建实例
+    embedding = registry.create_embedding(
+        "MyApiEmbeddingModel",
+        api_key="your-key"
+    )
+    
+    # 使用模型
+    result = embedding.embed_query("测试查询")
+    print(result)
+else:
+    print("模型未找到，请检查文件是否正确放置")
+```
+
+## 调试
+
+如果模型未加载，检查：
+
+1. 文件是否在正确的目录（`~/.jarvis/embeddings/` 或 `~/.jarvis/rerankers/`）
+2. 文件是否以 `.py` 结尾且不以 `__` 开头
+3. 类是否继承自正确的接口
+4. 是否实现了所有必需的方法
+5. 查看控制台输出的错误信息