jarvis-ai-assistant 0.1.219__tar.gz → 0.1.220__tar.gz

{jarvis_ai_assistant-0.1.219/src/jarvis_ai_assistant.egg-info → jarvis_ai_assistant-0.1.220}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jarvis-ai-assistant
-Version: 0.1.219
+Version: 0.1.220
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire
@@ -33,11 +33,11 @@ Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: POSIX :: Linux
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Python: >=3.8
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: requests==2.32.3
@@ -56,9 +56,19 @@ Requires-Dist: tabulate==0.9.0
 Requires-Dist: pyte==0.8.2
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: pyyaml>=5.3.1
-Requires-Dist: ddgs==6.1.8
-Requires-Dist: beautifulsoup4==4.12.3
-Requires-Dist: lxml==5.2.2
+Requires-Dist: ddgs==9.0.2
+Requires-Dist: beautifulsoup4==4.13.4
+Requires-Dist: lxml==6.0.0
+Requires-Dist: langchain
+Requires-Dist: langchain-community
+Requires-Dist: langchain-huggingface
+Requires-Dist: chromadb
+Requires-Dist: diskcache
+Requires-Dist: sentence-transformers==2.7.0
+Requires-Dist: torch>=2.6
+Requires-Dist: typer
+Requires-Dist: unstructured[md]
+Requires-Dist: rank-bm25
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"
@@ -553,7 +563,76 @@ OPENAI_API_BASE: https://api.openai.com/v1
 | search_web | 使用互联网搜索 |
 | virtual_tty | 控制虚拟终端执行操作 |
 
-### 2. 命令替换功能
+### 2. 🧠 RAG增强知识库 (`jarvis-rag`)
+
+`jarvis-rag` 是一个强大的命令行工具，用于构建、管理和查询本地化的RAG（检索增强生成）知识库。它允许您将自己的文档（代码、笔记、文章等）作为外部知识源，让AI能够基于这些具体内容进行回答，而不是仅仅依赖其通用训练数据。
+
+#### 核心功能
+
+- **本地化知识库**：所有文档索引和数据都存储在本地，确保数据隐私和安全。
+- **智能文档加载**：自动识别并加载多种文本文件，无需关心文件后缀。
+- **灵活的查询**：支持使用项目默认的“思考”模型或动态指定任意已注册的LLM进行查询。
+- **配置驱动**：通过 `config.toml` 或 `.jarvis/config.yaml` 进行集中配置。
+
+#### 子命令说明
+
+##### 2.1 添加文档到知识库 (`add`)
+
+此命令用于将文件、目录或符合通配符模式的文档添加到知识库中。
+
+```bash
+# 基本用法
+jarvis-rag add <文件路径/目录路径/通配符模式>...
+
+# 示例
+# 1. 添加单个文件
+jarvis-rag add ./docs/my_document.md
+
+# 2. 添加整个目录（将递归扫描所有文本文件）
+jarvis-rag add ./src/
+
+# 3. 使用通配符添加所有Python文件
+jarvis-rag add 'src/**/*.py'
+
+# 4. 混合添加
+jarvis-rag add README.md ./docs/ 'src/jarvis/jarvis_rag/*.py'
+```
+
+**参数与选项:**
+
+| 参数/选项 | 描述 |
+|---|---|
+| `paths` | **[必需]** 一个或多个文件路径、目录路径或用引号包裹的通配符模式。 |
+| `--collection` | 指定知识库的集合名称（默认为 `jarvis_rag_collection`）。 |
+| `--embedding-mode` | 覆盖全局配置，强制使用特定嵌入模式 (`performance` 或 `accuracy`)。 |
+| `--db-path` | 覆盖全局配置，指定向量数据库的存储路径。 |
+
+##### 2.2 查询知识库 (`query`)
+
+此命令用于向已建立的知识库提出问题。
+
+```bash
+# 基本用法
+jarvis-rag query "你的问题"
+
+# 示例
+# 1. 使用默认配置进行查询
+jarvis-rag query "请总结一下我添加的文档的核心内容"
+
+# 2. 指定使用Kimi模型进行查询
+jarvis-rag query "代码中的 'PlatformRegistry' 类是做什么用的？" --platform kimi --model moonshot-v1-8k
+```
+
+**参数与选项:**
+
+| 参数/选项 | 描述 |
+|---|---|
+| `question` | **[必需]** 你要向知识库提出的问题。 |
+| `--collection` | 指定要查询的知识库集合名称。 |
+| `--platform` | 指定一个平台名称来回答问题，覆盖默认的“思考”模型。 |
+| `--model` | 指定一个模型名称来回答问题，需要与 `--platform` 同时使用。 |
+
+### 3. 命令替换功能
 支持使用特殊标记`'<tag>'`触发命令替换功能：
 
 | 标记 | 功能 |
@@ -568,7 +647,7 @@ OPENAI_API_BASE: https://api.openai.com/v1
 | `'Check'` | 执行静态代码检查，包括错误和风格问题 |
 | `'SaveSession'` | 保存当前会话并退出 |
 
-### 3. 自定义替换配置
+### 4. 自定义替换配置
 在`~/.jarvis/config.yaml`中添加：
 ```yaml
 JARVIS_REPLACE_MAP:

{jarvis_ai_assistant-0.1.219 → jarvis_ai_assistant-0.1.220}/README.md

@@ -480,7 +480,76 @@ OPENAI_API_BASE: https://api.openai.com/v1
 | search_web | 使用互联网搜索 |
 | virtual_tty | 控制虚拟终端执行操作 |
 
-### 2. 命令替换功能
+### 2. 🧠 RAG增强知识库 (`jarvis-rag`)
+
+`jarvis-rag` 是一个强大的命令行工具，用于构建、管理和查询本地化的RAG（检索增强生成）知识库。它允许您将自己的文档（代码、笔记、文章等）作为外部知识源，让AI能够基于这些具体内容进行回答，而不是仅仅依赖其通用训练数据。
+
+#### 核心功能
+
+- **本地化知识库**：所有文档索引和数据都存储在本地，确保数据隐私和安全。
+- **智能文档加载**：自动识别并加载多种文本文件，无需关心文件后缀。
+- **灵活的查询**：支持使用项目默认的“思考”模型或动态指定任意已注册的LLM进行查询。
+- **配置驱动**：通过 `config.toml` 或 `.jarvis/config.yaml` 进行集中配置。
+
+#### 子命令说明
+
+##### 2.1 添加文档到知识库 (`add`)
+
+此命令用于将文件、目录或符合通配符模式的文档添加到知识库中。
+
+```bash
+# 基本用法
+jarvis-rag add <文件路径/目录路径/通配符模式>...
+
+# 示例
+# 1. 添加单个文件
+jarvis-rag add ./docs/my_document.md
+
+# 2. 添加整个目录（将递归扫描所有文本文件）
+jarvis-rag add ./src/
+
+# 3. 使用通配符添加所有Python文件
+jarvis-rag add 'src/**/*.py'
+
+# 4. 混合添加
+jarvis-rag add README.md ./docs/ 'src/jarvis/jarvis_rag/*.py'
+```
+
+**参数与选项:**
+
+| 参数/选项 | 描述 |
+|---|---|
+| `paths` | **[必需]** 一个或多个文件路径、目录路径或用引号包裹的通配符模式。 |
+| `--collection` | 指定知识库的集合名称（默认为 `jarvis_rag_collection`）。 |
+| `--embedding-mode` | 覆盖全局配置，强制使用特定嵌入模式 (`performance` 或 `accuracy`)。 |
+| `--db-path` | 覆盖全局配置，指定向量数据库的存储路径。 |
+
+##### 2.2 查询知识库 (`query`)
+
+此命令用于向已建立的知识库提出问题。
+
+```bash
+# 基本用法
+jarvis-rag query "你的问题"
+
+# 示例
+# 1. 使用默认配置进行查询
+jarvis-rag query "请总结一下我添加的文档的核心内容"
+
+# 2. 指定使用Kimi模型进行查询
+jarvis-rag query "代码中的 'PlatformRegistry' 类是做什么用的？" --platform kimi --model moonshot-v1-8k
+```
+
+**参数与选项:**
+
+| 参数/选项 | 描述 |
+|---|---|
+| `question` | **[必需]** 你要向知识库提出的问题。 |
+| `--collection` | 指定要查询的知识库集合名称。 |
+| `--platform` | 指定一个平台名称来回答问题，覆盖默认的“思考”模型。 |
+| `--model` | 指定一个模型名称来回答问题，需要与 `--platform` 同时使用。 |
+
+### 3. 命令替换功能
 支持使用特殊标记`'<tag>'`触发命令替换功能：
 
 | 标记 | 功能 |
@@ -495,7 +564,7 @@ OPENAI_API_BASE: https://api.openai.com/v1
 | `'Check'` | 执行静态代码检查，包括错误和风格问题 |
 | `'SaveSession'` | 保存当前会话并退出 |
 
-### 3. 自定义替换配置
+### 4. 自定义替换配置
 在`~/.jarvis/config.yaml`中添加：
 ```yaml
 JARVIS_REPLACE_MAP:

{jarvis_ai_assistant-0.1.219 → jarvis_ai_assistant-0.1.220}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jarvis-ai-assistant"
-version = "0.1.219"
+version = "0.1.220"
 description = "Jarvis: An AI assistant that uses tools to interact with the system"
 readme = "README.md"
 authors = [{ name = "skyfire", email = "skyfireitdiy@hotmail.com" }]
@@ -15,10 +15,10 @@ classifiers = [
     "License :: OSI Approved :: MIT License",
     "Operating System :: POSIX :: Linux",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
 ]
 keywords = ["jarvis", "ai", "assistant", "tools", "automation"]
 dependencies = [
@@ -38,11 +38,22 @@ dependencies = [
     "pyte==0.8.2",
     "httpx>=0.28.1",
     "pyyaml>=5.3.1",
-    "ddgs==6.1.8",
-    "beautifulsoup4==4.12.3",
-    "lxml==5.2.2",
+    "ddgs==9.0.2",
+    "beautifulsoup4==4.13.4",
+    "lxml==6.0.0",
+    # RAG dependencies
+    "langchain",
+    "langchain-community",
+    "langchain-huggingface",
+    "chromadb",
+    "diskcache",
+    "sentence-transformers==2.7.0",
+    "torch>=2.6",
+    "typer",
+    "unstructured[md]",
+    "rank-bm25",
 ]
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 
 [project.optional-dependencies]
 dev = ["pytest", "black", "isort", "mypy", "build", "twine"]
@@ -67,3 +78,4 @@ jarvis-agent = "jarvis.jarvis_agent.main:main"
 jarvis-tool = "jarvis.jarvis_tools.cli.main:main"
 jarvis-git-details = "jarvis.jarvis_git_details.main:main"
 jarvis-methodology = "jarvis.jarvis_methodology.main:main"
+jarvis-rag = "jarvis.jarvis_rag.cli:main"

{jarvis_ai_assistant-0.1.219 → jarvis_ai_assistant-0.1.220}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages  # type: ignore
 
 setup(
     name="jarvis-ai-assistant",
-    version="0.1.219",
+    version="0.1.220",
     author="skyfire",
     author_email="skyfireitdiy@hotmail.com",
     description="An AI assistant that uses various tools to interact with the system",
@@ -33,12 +33,22 @@ setup(
         "pyte==0.8.2",
         "httpx>=0.28.1",
         "pyyaml>=5.3.1",
-        "ddgs==6.1.8",
-        "beautifulsoup4==4.12.3",
-        "lxml==5.2.2",
+        "ddgs==9.0.2",
+        "beautifulsoup4==4.13.4",
+        "lxml==6.0.0",
+        # RAG dependencies
+        "langchain",
+        "langchain-community",
+        "langchain-huggingface",
+        "chromadb",
+        "diskcache",
+        "sentence-transformers==2.7.0",
+        "torch>=2.6",
+        "typer",
+        "unstructured[md]",
+        "rank-bm25",
     ],
-    extras_require={"dev": ["pytest", "black",
-                            "isort", "mypy", "build", "twine"]},
+    extras_require={"dev": ["pytest", "black", "isort", "mypy", "build", "twine"]},
     entry_points={
         "console_scripts": [
             "jarvis=jarvis.jarvis_agent.jarvis:main",
@@ -57,18 +67,19 @@ setup(
             "jarvis-tool=jarvis.jarvis_tools.cli.main:main",
             "jarvis-git-details=jarvis.jarvis_git_details.main:main",
             "jarvis-methodology=jarvis.jarvis_methodology.main:main",
+            "jarvis-rag=jarvis.jarvis_rag.cli:main",
         ],
     },
-    python_requires=">=3.8",
+    python_requires=">=3.9",
     classifiers=[
         "Development Status :: 3 - Alpha",
         "Intended Audience :: Developers",
         "License :: OSI Approved :: MIT License",
         "Operating System :: POSIX :: Linux",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
     ],
 )

{jarvis_ai_assistant-0.1.219 → jarvis_ai_assistant-0.1.220}/src/jarvis/__init__.py

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

{jarvis_ai_assistant-0.1.219 → jarvis_ai_assistant-0.1.220}/src/jarvis/jarvis_agent/__init__.py

@@ -36,94 +36,49 @@ from jarvis.jarvis_utils.tag import ct, ot
 origin_agent_system_prompt = f"""
 <role>
 # 🤖 角色
-你是一个专业的任务执行助手，擅长根据用户需求生成详细的任务执行计划并执行。
+你是一个专业的任务执行助手，根据用户需求制定并执行详细的计划。
 </role>
 
-<requirements>
-# 🔥 绝对行动要求
-1. 每个响应必须包含且仅包含一个工具调用
-2. 唯一例外：任务结束
-3. 空响应会触发致命错误
-</requirements>
-
-<violations>
-# 🚫 违规示例
-- 没有工具调用的分析 → 永久挂起
-- 未选择的多选项 → 永久挂起
-- 请求用户确认 → 永久挂起
-</violations>
+<rules>
+# ❗ 核心规则
+1.  **单步操作**: 每个响应必须包含且仅包含一个工具调用。
+2.  **任务终结**: 当任务完成时，明确指出任务已完成。这是唯一可以不调用工具的例外。
+3.  **无响应错误**: 空响应或仅有分析无工具调用的响应是致命错误，会导致系统挂起。
+4.  **决策即工具**: 所有的决策和分析都必须通过工具调用来体现。
+5.  **等待结果**: 在继续下一步之前，必须等待当前工具的执行结果。
+6.  **持续推进**: 除非任务完成，否则必须生成可操作的下一步。
+7.  **记录沉淀**: 如果解决方案有普适价值，应记录为方法论。
+8.  **用户语言**: 始终使用用户的语言进行交流。
+</rules>
 
 <workflow>
-# 🔄 问题解决流程
-1. 问题分析
-   - 重述问题以确认理解
-   - 分析根本原因（针对问题分析任务）
-   - 定义清晰、可实现的目标
-   → 必须调用分析工具
-
-2. 解决方案设计
-   - 生成多个可执行的解决方案
-   - 评估并选择最优方案
-   - 使用PlantUML创建详细行动计划
-   → 必须调用设计工具
-
-3. 执行
-   - 一次执行一个步骤
-   - 每个步骤只使用一个工具
-   - 等待工具结果后再继续
-   - 监控结果并根据需要调整
-   → 必须调用执行工具
-
-4. 任务完成
-   - 验证目标完成情况
-   - 如有价值则记录方法论
+# 🔄 工作流程
+1.  **分析**: 理解和分析问题，定义清晰的目标。
+2.  **设计**: 设计解决方案并制定详细的行动计划。
+3.  **执行**: 按照计划，一次一个步骤地执行。
+4.  **完成**: 验证任务是否达成目标，并进行总结。
 </workflow>
 
-<principles>
-# ⚖️ 操作原则
-- 每个步骤一个操作
-- 下一步前必须等待结果
-- 除非任务完成否则必须生成可操作步骤
-- 根据反馈调整计划
-- 记录可复用的解决方案
-- 使用完成命令结束任务
-- 操作之间不能有中间思考状态
-- 所有决策必须表现为工具调用
-</principles>
-
-<rules>
-# ❗ 重要规则
-1. 每个步骤只能使用一个操作
-2. 必须等待操作执行结果
-3. 必须验证任务完成情况
-4. 必须生成可操作步骤
-5. 如果无需操作必须使用完成命令
-6. 永远不要使对话处于等待状态
-7. 始终使用用户语言交流
-8. 必须记录有价值的方法论
-9. 违反操作协议将导致系统崩溃
-10. 空响应会触发永久挂起
-</rules>
-
 <system_info>
-# 系统信息：
-{platform.platform()}
-{platform.version()}
-
-# 当前时间
-{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
+# 系统信息
+- OS: {platform.platform()} {platform.version()}
+- Time: {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
 </system_info>
 """
 
 
 class OutputHandlerProtocol(Protocol):
-    def name(self) -> str: ...
+    def name(self) -> str:
+        ...
 
-    def can_handle(self, response: str) -> bool: ...
+    def can_handle(self, response: str) -> bool:
+        ...
 
-    def prompt(self) -> str: ...
+    def prompt(self) -> str:
+        ...
 
-    def handle(self, response: str, agent: Any) -> Tuple[bool, Any]: ...
+    def handle(self, response: str, agent: Any) -> Tuple[bool, Any]:
+        ...
 
 
 class Agent:
@@ -189,9 +144,7 @@ class Agent:
             if isinstance(platform, str):
                 self.model = PlatformRegistry().create_platform(platform)
                 if self.model is None:
-                    PrettyOutput.print(
-                        f"平台 {platform} 不存在，将使用普通模型", OutputType.WARNING
-                    )
+                    PrettyOutput.print(f"平台 {platform} 不存在，将使用普通模型", OutputType.WARNING)
                     self.model = PlatformRegistry().get_normal_platform()
             else:
                 self.model = platform
@@ -816,18 +769,14 @@ arguments:
 
                     if get_interrupt():
                         set_interrupt(False)
-                        user_input = self.multiline_inputer(
-                            f"模型交互期间被中断，请输入用户干预信息："
-                        )
+                        user_input = self.multiline_inputer(f"模型交互期间被中断，请输入用户干预信息：")
                         if user_input:
                             # 如果有工具调用且用户确认继续，则将干预信息和工具执行结果拼接为prompt
                             if any(
                                 handler.can_handle(current_response)
                                 for handler in self.output_handler
                             ):
-                                if user_confirm(
-                                    "检测到有工具调用，是否继续处理工具调用？", True
-                                ):
+                                if user_confirm("检测到有工具调用，是否继续处理工具调用？", True):
                                     self.prompt = f"{user_input}\n\n{current_response}"
                                     continue
                             self.prompt += f"{user_input}"
@@ -873,9 +822,7 @@ arguments:
             if self.use_methodology:
                 if not upload_methodology(self.model, other_files=self.files):
                     if self.files:
-                        PrettyOutput.print(
-                            "文件上传失败，将忽略文件列表", OutputType.WARNING
-                        )
+                        PrettyOutput.print("文件上传失败，将忽略文件列表", OutputType.WARNING)
                         # 上传失败则回退到本地加载
                     msg = self.prompt
                     for handler in self.input_handler:
@@ -883,14 +830,14 @@ arguments:
                     self.prompt = f"{self.prompt}\n\n以下是历史类似问题的执行经验，可参考：\n{load_methodology(msg, self.get_tool_registry())}"
                 else:
                     if self.files:
-                        self.prompt = f"{self.prompt}\n\n上传的文件包含历史对话信息和方法论文件，可以从中获取一些经验信息。"
+                        self.prompt = (
+                            f"{self.prompt}\n\n上传的文件包含历史对话信息和方法论文件，可以从中获取一些经验信息。"
+                        )
                     else:
                         self.prompt = f"{self.prompt}\n\n上传的文件包含历史对话信息，可以从中获取一些经验信息。"
             elif self.files:
                 if not self.model.upload_files(self.files):
-                    PrettyOutput.print(
-                        "文件上传失败，将忽略文件列表", OutputType.WARNING
-                    )
+                    PrettyOutput.print("文件上传失败，将忽略文件列表", OutputType.WARNING)
                 else:
                     self.prompt = f"{self.prompt}\n\n上传的文件包含历史对话信息，可以从中获取一些经验信息。"
         else:

{jarvis_ai_assistant-0.1.219 → jarvis_ai_assistant-0.1.220}/src/jarvis/jarvis_data/config_schema.json

@@ -181,6 +181,36 @@
       "description": "是否打印提示",
       "default": false
     },
+    "JARVIS_RAG": {
+      "type": "object",
+      "description": "RAG框架的配置",
+      "properties": {
+        "embedding_mode": {
+          "type": "string",
+          "enum": [
+            "performance",
+            "accuracy"
+          ],
+          "default": "performance",
+          "description": "嵌入模型的模式, 'performance'表示性能优先, 'accuracy'表示准确度优先"
+        },
+        "embedding_cache_path": {
+          "type": "string",
+          "default": ".jarvis/rag/embeddings",
+          "description": "嵌入向量缓存的路径, 相对于当前工作目录"
+        },
+        "vector_db_path": {
+          "type": "string",
+          "default": ".jarvis/rag/vectordb",
+          "description": "向量数据库的持久化存储路径, 相对于当前工作目录"
+        }
+      },
+      "default": {
+        "embedding_mode": "performance",
+        "embedding_cache_path": ".jarvis/rag/embeddings",
+        "vector_db_path": ".jarvis/rag/vectordb"
+      }
+    },
     "JARVIS_REPLACE_MAP": {
       "type": "object",
       "description": "自定义替换映射表配置",

{jarvis_ai_assistant-0.1.219 → jarvis_ai_assistant-0.1.220}/src/jarvis/jarvis_platform/yuanbao.py

@@ -38,7 +38,9 @@ class YuanbaoPlatform(BasePlatform):
         self.agent_id = "naQivTmsDa"
 
         if not self.cookies:
-            PrettyOutput.print("YUANBAO_COOKIES 未设置", OutputType.WARNING)
+            raise ValueError(
+                "YUANBAO_COOKIES environment variable not set. Please provide your cookies to use the Yuanbao platform."
+            )
 
         self.system_message = ""  # 系统消息，用于初始化对话
         self.first_chat = True  # 标识是否为第一次对话

jarvis_ai_assistant-0.1.220/src/jarvis/jarvis_rag/__init__.py

@@ -0,0 +1,11 @@
+"""
+Jarvis RAG Framework
+
+A flexible RAG pipeline with pluggable remote LLMs and local, cache-enabled embedding models.
+"""
+
+from .rag_pipeline import JarvisRAGPipeline
+from .llm_interface import LLMInterface
+from .embedding_manager import EmbeddingManager
+
+__all__ = ["JarvisRAGPipeline", "LLMInterface", "EmbeddingManager"]

jarvis_ai_assistant-0.1.220/src/jarvis/jarvis_rag/cache.py

@@ -0,0 +1,87 @@
+import hashlib
+from typing import List, Optional, Any
+
+from diskcache import Cache
+
+
+class EmbeddingCache:
+    """
+    A disk-based cache for storing and retrieving text embeddings.
+
+    This class uses diskcache to create a persistent, local cache. It generates
+    a key for each text content based on its SHA256 hash, making lookups
+    deterministic and efficient.
+    """
+
+    def __init__(self, cache_dir: str, salt: str = ""):
+        """
+        Initializes the EmbeddingCache.
+
+        Args:
+            cache_dir (str): The directory where the cache will be stored.
+            salt (str): A salt to be added to the hash. This is crucial for
+                        ensuring that embeddings generated by different models
+                        do not collide. For example, use the model name as a salt.
+        """
+        self.cache = Cache(cache_dir)
+        self.salt = salt
+
+    def _get_key(self, text: str) -> str:
+        """Generates a unique cache key for a given text and salt."""
+        hash_object = hashlib.sha256((self.salt + text).encode("utf-8"))
+        return hash_object.hexdigest()
+
+    def get(self, text: str) -> Optional[Any]:
+        """
+        Retrieves an embedding from the cache.
+
+        Args:
+            text (str): The text to look up.
+
+        Returns:
+            The cached embedding, or None if it's not in the cache.
+        """
+        key = self._get_key(text)
+        return self.cache.get(key)
+
+    def set(self, text: str, embedding: Any) -> None:
+        """
+        Stores an embedding in the cache.
+
+        Args:
+            text (str): The text corresponding to the embedding.
+            embedding (Any): The embedding vector to store.
+        """
+        key = self._get_key(text)
+        self.cache.set(key, embedding)
+
+    def get_batch(self, texts: List[str]) -> List[Optional[Any]]:
+        """
+        Retrieves a batch of embeddings from the cache.
+
+        Args:
+            texts (List[str]): A list of texts to look up.
+
+        Returns:
+            A list containing cached embeddings or None for cache misses.
+        """
+        return [self.get(text) for text in texts]
+
+    def set_batch(self, texts: List[str], embeddings: List[Any]) -> None:
+        """
+        Stores a batch of embeddings in the cache.
+
+        Args:
+            texts (List[str]): The list of texts.
+            embeddings (List[Any]): The list of corresponding embeddings.
+        """
+        if len(texts) != len(embeddings):
+            raise ValueError("Length of texts and embeddings must be the same.")
+
+        with self.cache.transact():
+            for text, embedding in zip(texts, embeddings):
+                self.set(text, embedding)
+
+    def close(self):
+        """Closes the cache connection."""
+        self.cache.close()