diff --git a/docs/deprecated.txt b/docs/deprecated.txt
index 6d7cb14f..06f9648f 100644
--- a/docs/deprecated.txt
+++ b/docs/deprecated.txt
@@ -1,7 +1,9 @@
from loguru import logger
-用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码
-用英文注释，完善module/class/function docstring，要一句话简洁，代码要简洁，符合pep和pylint规范
+用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码逻辑，符合pep和pylint规范，使用list而不是typing.List/Dict，不使用typing.Union
+
+看看代码有什么问题
+用英文注释，完善module/class/function docstring，要一句话简洁，代码要简洁，符合pep和pylint规范，使用list而不是typing.List，不使用typing.Union
C0114: Missing module docstring (missing-module-docstring)
C0115: Missing class docstring (missing-class-docstring)
C0116: Missing function or method docstring (missing-function-docstring)
diff --git a/docs/reme_v2_design.md b/docs/reme_v2_design.md
index 18d14657..6fa78140 100644
--- a/docs/reme_v2_design.md
+++ b/docs/reme_v2_design.md
@@ -1,378 +1,735 @@
-# ReMeV2 - agent memory kit design
+# ReMeV2 深度设计文档：渐进式 Agentic Memory 方案
-ReMeV2 是一个面向智能体（Agent）的分层记忆系统，用于从连续的对话与交互中，逐步沉淀出稳定、可检索、可更新的长期记忆。
-系统强调从**原始轨迹 → 事件索引 → 高维抽象记忆**的分层存储，并通过**检索器（Retriever）**与**总结器（Summarizer）**实现动态更新。
+## 一、 背景与现状分析
+
+### 1.1 当前面临的挑战
+
+* **外功修炼（接口易用性）**：现有的 `server-client` 模式对新手开发者不够友好，集成成本高，需要更直观、纯 Pythonic 的调用方式。
+* **内功修炼（架构深度）**：受 `skills` 和 `agentic memory` 启发，现有的存储检索较为机械。我们需要一种基于**渐进式检索（Progressive Retrieval）**与**渐进式总结（Progressive Summarization）**的智能体记忆方案。
+
+### 1.2 核心目标
+
+1. **极简开发体验**：开发者友好，全异步接口，支持本地直接运行与 CLI 体验。
+2. **认知架构升级**：引入 渐进式检索 & 渐进式总结 的 Agentic 模式，融合多种记忆，让记忆的存取具备“思考”过程。
+3. **生态融合**：原生支持 AgentScope、LangChain 等主流框架。
---
-## 设计目标
+## 二、 竞品调研与启示
-本设计旨在构建一个可嵌入任意智能体框架（如对话 Agent、任务 Agent 等）的记忆子系统，使其具备以下能力：
+### 2.1 主流竞品深度对比
+
+| 产品 | 设计哲学 | 核心优势 | 局限性 |
+|-------------|----------|---------------------------------------------|-------------------|
+| **mem0** | 智能便签本 | 原子事实提取，极高 Token 效率。 | 缺乏对复杂逻辑链条的支持。 |
+| **Letta** | 带硬盘的 CPU | 模拟计算机三级存储（Core/Recall/Archival），Agent 自主控存。 | 状态机管理相对复杂。 |
+| **MIRIX** | 认知架构图谱 | 实体-关系双引擎，支持记忆“进化”与“固化”。 | 侧重研究，落地集成门槛较高。 |
+| **LangMem** | 用户档案系统 | 异步 Compaction（压缩），Schema 驱动，强一致性。 | 偏向 SaaS 应用，灵活性略逊。 |
+
+### 2.2 mem0
+- https://github.com/mem0ai/mem0
+- https://docs.mem0.ai/core-concepts/memory-operations/add
+- https://docs.mem0.ai/core-concepts/memory-operations/search
+- https://docs.mem0.ai/core-concepts/memory-operations/update
+- https://docs.mem0.ai/core-concepts/memory-operations/delete
+
+#### 2.2.1 API Reference
+| 接口名称 | 核心输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) |
+| --- | --- | --- | --- |
+| **Add** | `messages` (文本/对话), `user_id`, `metadata` | `id`, `event` (ADD/UPDATE), `data` | **提取与合并**：LLM 提取事实，自动去重并更新已有记忆，而非简单堆叠。 |
+| **Search** | `query` (自然语言), `filters`, `limit` | `id`, `memory` (事实文本), `score`, `metadata` | **语义检索**：基于向量相似度查找最相关的“原子事实”，支持多维过滤。 |
+| **Update** | `memory_id` (必填), `data` (新内容) | 操作状态 (Success/Fail) | **手动干预**：允许开发者对特定的事实进行精确修正。 |
+| **Delete** | `memory_id` 或 `user_id` (清空) | 操作状态 (Success/Fail) | **遗忘机制**：物理删除或逻辑移除不再需要的信息。 |
+
+#### 2.2.2 Tech Strategy & Benefits
+| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) |
+| --- | --- | --- |
+| **存储架构** | **混合存储**：向量数据库 (Vector) + 图数据库 (Graph) + 关系型元数据。 | **多维关联**：不仅能搜到相似内容，还能理解实体间的逻辑关系（如“父子”、“因果”）。 |
+| **数据处理** | **原子化事实提取**：利用 LLM 将长篇对话压缩为简短的 Fact。 | **极高 Token 效率**：注入 Prompt 的内容更精炼，减少 90% 以上的冗余信息，大幅降本。 |
+| **管理层级** | **多级联动**：User (长期) Agent (专业) Session (短期)。 | **个性化定制**：实现跨会话的“长效记忆”，AI 能记住用户一个月前说过的偏好。 |
+| **冲突处理** | **自适应更新算法**：新信息进入时自动比对旧记忆。 | **数据一致性**：自动处理矛盾信息（如用户更换了住址），确保记忆库始终是“最新真理”。 |
+| **兼容性** | **解耦设计**：支持多种 Embedding 模型与向量数据库后端。 | **快速集成**：几行代码即可为现有 LLM 应用增加记忆层，适配各种生产环境。 |
-- **支持多天、多会话的长期记忆积累**，不丢失原始细节。
-- **在推理时按需检索**与当前问题高度相关的记忆片段。
-- **在后台持续做分层总结与压缩**，提炼出抽象的“世界观 / 自我认知 / 用户画像 / 任务知识 / 工具经验”。
-- **支持记忆的增删改**，能够对过时或冲突信息进行更新，保持长期记忆的时效性与一致性。
-- **支持渐进式检索**，从高层抽象记忆逐步回溯到中层事件索引与底层原始轨迹，按需展开细节、控制上下文长度与推理成本。
-- **支持渐进式总结**，从最新对话开始按块增量总结并更新索引与主题记忆，而非每次对全量历史做重复总结。
---
-## 架构概览
+### 2.3 Letta
+- https://github.com/letta-ai/letta
+- https://docs.letta.com/guides/agents/archival-memory/
+- https://docs.letta.com/guides/agents/archival-search/
-系统整体可以分为两个维度来看：
+#### 2.3.1 存储架构层级 (Memory Tiering)
-- **静态维度**：记忆如何被分层存储（底层消息块 → 中层索引 → 顶层主题记忆）。
-- **动态维度**：记忆如何被检索与更新（上行 Summarizer、下行 Retriever）。
+Letta 将记忆分为三个物理/逻辑层，模拟计算机的存储架构：
-核心组件包括：
+| 记忆层级 | 存储介质 | 访问方式 | 核心作用 |
+| --- | --- | --- | --- |
+| **Core Memory** | **上下文窗口 (Prompt)** | 直接读写 | **即时意识**：包含 `Persona`（AI 设定）和 `Human`（用户信息）。Agent 随时可见，响应最快。 |
+| **Recall Memory** | **关系型数据库 (SQL)** | 分页检索 | **短期/历史回顾**：存储完整的对话流（Messages）。用于回答“你刚才说了什么”。 |
+| **Archival Memory** | **向量数据库 (Vector)** | 语义搜索 | **长期知识库**：存储海量事实或文档。Agent 通过工具自主检索或存入。 |
-- **底层记忆层（Raw Dialogue Store）**：多天对话轨迹切分成 `messages block`。
-- **中层索引层（Index & Summary Store）**：为每个 `messages block` 生成 `block_id: time + summary`。
-- **顶层主题记忆层（High-level Memories）**：自我认知、任务记忆、工具记忆、用户画像等高维抽象。
-- **Meta Agent（记忆中枢）**：持有高层记忆视图与索引，负责路由检索与更新请求。
-- **Sub Agents（记忆子智能体）**：面向具体主题（如某个用户、某类任务、某组工具）的记忆管理单元。
-- **Retriever（检索器）**：从高层到低层，渐进式查找相关记忆。
-- **Summarizer（总结器）**：从低层到高层，渐进式抽象和更新长期记忆。
+#### 2.3.2 核心操作接口 (API & Tool Reference)
+
+在 Letta 中，记忆的操作通常封装为 **Tools**，由 Agent 根据推理需求主动调用。
+
+| 接口/工具名称 | 输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) |
+| --- | --- | --- | --- |
+| **`core_memory_update`** | `section`, `new_content` | 更新后的段落内容 | **原子替换**：直接修改 System Prompt 中的特定块（如：更新用户的职业或 AI 的性格偏好）。 |
+| **`archival_memory_insert`** | `content` (字符串) | 写入状态/ID | **知识沉淀**：将当前对话中的重要信息或外部文件片段“持久化”到向量数据库。 |
+| **`archival_memory_search`** | `query`, `page` | 匹配的文本块列表 | **主动 RAG**：Agent 意识到知识不足时，自主发起向量检索，并将结果拉入临时上下文。 |
+| **`conversation_search`** | `query`, `start_date` | 历史消息记录 | **全文检索**：在 Recall Memory 中根据关键词或时间戳查找历史对话详情。 |
+| **`send_message`** | `message`, `agent_id` | 响应流/状态更新 | **状态循环**：这是主入口，触发 Agent 的“思考-行动-观察”循环，自动处理内存同步。 |
+
+#### 2.3.3 技术策略与核心优势 (Tech Strategy & Benefits)
+
+| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) |
+| --- | --- | --- |
+| **状态持久化** | **Agent State Snapshot**：将 Agent 的所有内存、工具定义和历史记录打包存入数据库。 | **无限存续**：Agent 不再是无状态的 API 调用。重启服务器后，Agent 依然记得所有细节。 |
+| **自主演进** | **Self-Editing Loop**：Agent 拥有修改自己 Core Memory 的权限（通过函数调用）。 | **认知闭环**：AI 能在交流中发现矛盾并自我更正，例如发现用户搬家后自动更新 `Human` 模块。 |
+| **算力调度** | **OOC (Out-of-Context) 管理**：当对话过长，系统自动将旧消息从 Core 移入 Recall。 | **突破 Context 限制**：在 8k 窗口的模型上也能处理相当于 1M 窗口的逻辑量，且成本更低。 |
+| **多代理协同** | **Letta Server 中控**：统一管理多个 Agent 的状态机与资源访问权限。 | **企业级扩展**：支持创建 Agent 团队，每个 Agent 拥有独立的记忆空间但可共享 Archival 库。 |
+| **解耦灵活性** | **Provider Agnostic**：后端支持 Postgres/Chroma，前端支持 OpenAI/Anthropic/Local LLMs。 | **无缝迁移**：不绑定特定模型，开发者可以根据成本或能力随时更换底座。 |
+
+#### 2.3.4 与 mem0 的深度对比
+
+* **设计哲学**：
+* **mem0** 像是一个**“智能记事本”**，它在后台默默地帮你总结事实。
+* **Letta** 像是一个**“带硬盘的 CPU”**，它把记忆管理完全交给了 Agent 自己的逻辑推理。
+
+
+* **交互模式**：
+* **mem0** 通常是外部干预（Add/Search）。
+* **Letta** 强调 **Agentic Control**（Agent 意识到需要搜索时才去搜索），这种模式更接近人类的思维过程。
---
-## 静态记忆三层架构
+### 2.4 MIRIX
+- https://github.com/Mirix-AI/MIRIX
+- https://docs.mirix.io/
-### 垂直层级划分
+#### 2.4.1 API Reference
-系统将记忆分为自下而上的三个层级，抽象程度逐级升高：
+| 接口名称 | 核心输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) |
+| --- | --- | --- | --- |
+| **Add** | `content` (观察/对话), `agent_id`, `context_type` (如任务/闲聊) | `memory_id`, `graph_nodes`, `status` | **实体建模**：不只是提取事实，而是将信息拆解为实体（Entities）与关系（Relations），并挂载到智能体的知识图谱中。 |
+| **Query** | `query` (意图), `scope` (全局/局部), `top_k` | `retrieved_memories`, `relation_paths`, `score` | **混合检索**：结合向量（Vector）的语义相关性和图（Graph）的拓扑连接性，寻找具有逻辑深度背景的记忆。 |
+| **Evolve** | `target_memories` (可选), `agent_id` | `optimized_structure`, `merged_nodes` | **记忆固化/压缩**：模仿人类大脑的“睡眠”机制，自动合并碎片化记忆，将短期经验转化为长期的结构化知识。 |
+| **Observe** | `interaction_stream`, `feedback` | `insights`, `priority_update` | **实时学习**：根据用户反馈或环境变化，动态调整记忆的权重（Importance）和置信度。 |
-- **底层：记忆碎片（历史对话 / 轨迹）**
- - 存放原始的、未经压缩的对话与交互记录。
- - **规则**：不去重、不解冲突，只做追加与切片。
+#### 2.4.2 Tech Strategy & Benefits
-- **中层：记忆索引（事件总结）**
- - 对原始消息块做首轮总结，形成“时间 + 简要摘要”的事件索引。
- - 连接底层细节与顶层抽象。
- - **规则**：允许存在部分冗余与矛盾，主要目标是加速检索。
+| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) |
+| --- | --- | --- |
+| **存储架构** | **语义-关系双引擎**：向量索引（Vector Index）+ 属性图（Property Graph）。 | **深度上下文**：不仅知道“是什么”，还能通过图路径推理出“为什么”，有效解决 LLM 幻觉问题。 |
+| **记忆层级** | **三层架构**：感知记忆 (Perception) -> 语义记忆 (Semantic) -> 经验记忆 (Episodic)。 | **任务适应性**：不同任务自动匹配不同的记忆深度，短期任务关注细节，长期任务关注模式。 |
+| **演化机制** | **自主固化 (Self-Consolidation)**：通过 LLM 定期对冗余、矛盾信息进行清洗和逻辑抽象。 | **永久生命力**：解决随时间推移记忆库膨胀导致的检索噪声，确保记忆库“越用越聪明”。 |
+| **推理增强** | **基于记忆的 RAG+**：在检索到的事实基础上，额外提供关联的逻辑链条（Logic Chains）。 | **辅助决策**：为 Agent 提供决策支撑，使其在处理复杂流程时具备类似“长期经验值”的直觉。 |
+| **多代理协同** | **内存共享协议**：支持 Agent 之间的记忆交换与知识同步。 | **群体智能**：多个 Agent 可以共享同一套底层知识体系，同时保留各自的私有工作记忆。 |
-- **顶层：主题记忆（高维抽象知识）**
- - 聚合并抽象出关于“自我、任务、工具、用户”等长期知识。
- - **规则**：需要去重（Deduplication）与解冲突（Resolve Conflict），追求紧凑、一致、可解释。
+#### 2.4.3 与 mem0 的主要区别
-### 底层：原始对话与轨迹（Raw Dialogue）
-
-- 按时间顺序堆叠的对话与行为轨迹：
- - `Day1 Messages → Day2 Messages → ... → DayN Messages`
-- 通过时间窗口、业务逻辑、固定长度（可重叠）等策略切分为多个 `messages block`：
- - `Day1 Messages Block1`
- - `Day1 Messages Block2`
- - `DayN Messages BlockN`
-- 特性：
- - 不丢失任何历史细节。
- - 支持后续回溯、审计与再训练。
-
-### 中层：记忆索引与事件总结（Index & Summary）
-
-- 对每个 `messages block` 生成一个结构化索引：
- - 格式：`block_id: time + summary`
-- 示例：
- - `block_id1: 2025-11-11 + 用户提到他非常喜欢吃西瓜。`
- - `block_id2: 2025-11-12 + 用户表达自己很孤独。`
-- 作用：
- - 作为底层原始消息与顶层抽象记忆之间的“桥梁”与“目录”。
- - 为时间感知检索、语义检索提供最小粒度的事件单元。
-
-### 顶层：主题记忆（High-Dimensional Abstraction）
-
-顶层记忆由多个“记忆气泡”（Memory Bubble）构成，每个气泡代表一类主题知识，并通过关联的 `block_id` 与中层索引相连。主要包括：
-
-- **自我认知记忆（Self-Cognition Memory）**
- - 示例：
- - “我叫 Remy”
- - “我每天需要上班”
- - “我是一名面向长期陪伴的智能助手”
- - 位置：处于记忆拓扑中心，与其他记忆类别广泛连接。
-
-- **金融任务记忆（Financial Task Memory）**
- - 示例：
- - “做财报分析前需要调研企业的经营数据”
- - “估值时可以优先考虑 forward PE”
- - 作用：支撑在特定领域（如金融）的任务推理能力。
-
-- **工具记忆（Tool Memory）**
- - 示例：
- - “`bing_search` 很不稳定，经常失败”
- - “`tongyi_search` 的结果质量较好，常需要配合 `web_extract` 使用”
- - 作用：让智能体在工具选择、调用顺序和错误预期上具有“经验”。
-
-- **用户画像记忆（Per-User Memory，如用户 A / 用户 B）**
- - 对用户 A 的记忆示例：
- - “A 很喜欢吃西瓜”
- - “A 非常有礼貌”
- - “A 计划下周出行”
- - 与中层事件关联：
- - 顶层的“喜欢吃西瓜”记忆与 `block_id1` 建立逻辑连接。
- - 每个用户独立维护一套长期画像，避免混淆。
+* **Mem0** 侧重于**个性化偏好存储**（Personalization），核心是记住“用户喜欢什么”。
+* **MIRIX** 侧重于**智能体认知架构**（Agent Cognition），核心是让 Agent 具备类似人类的“知识归纳”和“逻辑推理”记忆能力。
---
-## 动态流程一：ReMeV2 Retriever（检索流程）
+### 2.5 LangMem
+- https://github.com/langchain-ai/langmem
+- https://langchain-ai.github.io/langmem/
-检索流程是一个**自上而下**的渐进式过程：从高维抽象出发，逐步定位到具体事件甚至原始对话块，为当前问答或决策提供证据。
+#### 2.5.1 API Reference
-```mermaid
-flowchart TD
- op0[op: 当前查询 / 上下文输入]
- op1[op: Meta Agent 读取自我认知与主题索引]
- op2[op: Meta Agent 基于 topic/intent 选择记忆子空间]
- op3[op: 调度对应 Sub Agent 实例]
- op4[op: 在顶层/中层记忆空间检索相关 Memories / 索引条目]
- op5[op: 根据需要回溯到底层 messages block（retrieve_block）]
- op6[op: 汇总得到的 messages block 列表，构造检索上下文]
- op7[op: 将上下文交给上层 Agent 做回答 / 决策]
+| 接口名称 | 核心输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) |
+| --- | --- | --- | --- |
+| **Add Messages** | `thread_id`, `messages` (List), `user_id` | 操作确认 / 任务 ID | **流式注入**：将原始对话追加到指定的 Thread。LangMem 会自动关联用户上下文，准备进行后续的异步处理。 |
+| **Query Memory** | `user_id`, `query` (语义描述), `namespace` | 结构化记忆对象 (JSON / Text) | **多维检索**：不仅支持向量相似度搜索，还能根据定义的 Schema 返回结构化的用户画像或知识状态。 |
+| **Trigger Logic** | `thread_id`, `memory_type` | 更新后的 Memory State | **异步固化**：后台启动 LLM 任务，将长篇对话“压缩”并“提取”到长期存储中。支持自定义提取逻辑（如更新用户信息）。 |
+| **Manage State** | `user_id`, `patch_data` (增量更新) | 成功/失败 状态 | **精确受控**：开发者可以直接修改持久化的状态（State），支持类似于 Git 的状态管理。 |
- op0 --> op1 --> op2 --> op3 --> op4 --> op5 --> op6 --> op7
+#### 2.5.2 Tech Strategy & Benefits
+
+| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) |
+| --- | --- | --- |
+| **存储架构** | **Stateful Persistence**：基于关系型数据库 (Postgres) + 向量索引。 | **强一致性**：利用数据库事务确保记忆更新的可靠性，支持复杂的结构化查询与过滤。 |
+| **数据处理** | **异步化 Compaction (压缩)**：在对话间隙通过后台 Worker 提取知识。 | **无感延迟**：核心对话流程不被记忆提取阻塞，通过定时或事件驱动完成“记忆固化”，优化用户体验。 |
+| **管理层级** | **Thread -> User -> Organization**：三层级联记忆。 | **上下文隔离**：完美适配 SaaS 应用场景，既能记住单次对话（Thread），也能沉淀用户习惯（User）。 |
+| **逻辑引擎** | **Schema-Driven (模式驱动)**：允许定义 JSON Schema 来规范记忆内容。 | **高度可预测**：输出不再是散乱的句子，而是结构化的字段，方便下游程序直接调用逻辑（如自动填充表单）。 |
+| **集成生态** | **LangGraph 原生集成**：作为 Checkpointer 或存储节点直接接入。 | **生态协同**：如果你已经在用 LangChain，LangMem 可以无缝接管状态流转，无需重写底层存储逻辑。 |
+
+#### 2.5.3 与 mem0 的核心差异
+
+* **mem0** 像是一个**“便签本”**：它擅长从每一句话里抠出零散的事实（如“我喜欢吃苹果”），然后把它们存成一条条语义片段。
+* **LangMem** 像是一个**“用户档案系统”**：它更擅长分析一整段对话，然后更新一个复杂的 JSON 档案（如更新用户的偏好模型、性格标签、历史任务状态）。
+
+---
+
+## 三、 ReMeV2 API 接口设计
+
+### 3.1 Long-Term Memory (长期记忆)
+
+#### 3.1.1 Basic Usage (基础用法)
+
+The most straightforward way to use ReMe for long-term memory management. Supports basic summary and retrieval operations.
+
+```python
+import os
+from reme_ai import ReMe
+
+os.environ["REME_LLM_API_KEY"] = "sk-..."
+os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+os.environ["REME_EMBEDDING_API_KEY"] = "sk-..."
+os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+
+memory = ReMe(
+ memory_space="remy", # workspace identifier
+ llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6},
+ embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024},
+ vector_store={"backend": "local_file"}, # supported: local_file, chromadb, qdrant, etc.
+)
+
+# Summarize conversation into memory
+result = await memory.summary(
+ messages=[
+ {"role": "user", "content": "I'm travelling to SF"},
+ {"role": "assistant", "content": "That's great to hear!"}
+ ],
+ user_id="Alice",
+ # memory_type="auto" # default: auto (auto, personal, procedural, tool)
+)
+
+# Retrieve relevant memories
+memories = await memory.retrieve(
+ query="what is your travel plan?",
+ limit=3,
+ user_id="Alice",
+ # memory_type="auto" # default: auto
+)
+memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"])
+print(memories_str)
```
-### Meta Agent：记忆入口与路由
+#### 3.1.2 CLI Chat Application (命令行聊天应用)
-- 持有：
- - **常驻的自我认知记忆**（总是加载在工作内存中）。
- - 各类主题记忆的**名称 + 描述索引**，例如：
- - `1. 金融任务记忆`
- - `2. 工具记忆`
- - `3. 用户A记忆`
- - `4. 用户B记忆`
-- 职责：
- - 根据当前查询的 **topic / intent**，选择合适的记忆子空间。
- - 按需创建或调度具体的 `Sub_Agent` 来完成深入检索。
+A complete example demonstrating how to build a memory-enhanced chatbot with CLI interface.
-### Sub Agent：按主题检索
+```python
+import os
+from reme_ai import ReMe
+from openai import OpenAI
-- **预定义子智能体**（如：工具记忆 Agent、某个用户画像 Agent 等）：
- - 支持接口：
- - `retrieve_by_query(query)`：基于语义相似度的检索。
- - `time_aware_retrieve(time)`：基于时间的窗口检索（适合用户时间线类记忆）。
-- **自定义子智能体**：
- - 可以实现更复杂的检索逻辑（例如结合时间衰减、多源融合等）。
+os.environ["REME_LLM_API_KEY"] = "sk-..."
+os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+os.environ["REME_EMBEDDING_API_KEY"] = "sk-..."
+os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
-### 渐进式检索与回溯
+memory = ReMe(
+ memory_space="remy",
+ llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6},
+ embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024},
+ vector_store={"backend": "local_file"},
+)
-- 子智能体首先在顶层 / 中层记忆空间中找到语义相关的 **Memories / 索引条目**：
- - 每条记忆通常携带 `block_id` 与简要 summary。
-- 若需要更细节的证据，系统可以进一步调用：
- - `retrieve_block(block_id)`：回溯到底层 `messages block`。
- - `time_match_memories(time_range)`：根据时间窗口获取原始对话块。
-- 最终输出：
- - 一组与当前问题高度相关的 `messages block List`，作为上下文，供 Agent 生成回答或做决策。
+os.environ["OPENAI_API_KEY"] = "sk-..."
+os.environ["OPENAI_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+openai_client = OpenAI()
----
+def chat_with_memories(
+ query: str,
+ history_messages: list[dict],
+ user_name: str = "",
+ start_summary_size: int = 2,
+ keep_size: int = 0
+) -> str:
+ # Retrieve relevant memories for the query
+ memories = memory.retrieve(query=query, user_id=user_name, limit=3)
-## 动态流程二：ReMeV2 Summarizer（总结流程）
+ # Build system prompt with memories
+ system_prompt = (
+ "You are a helpful AI named `Remy`. Use the user memories to answer the question. "
+ "If you don't know the answer, just say you don't know. Don't try to make up an answer.\n"
+ )
+ if memories:
+ memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"])
+ system_prompt += f"User Memories:\n{memories_str}\n"
-总结流程是一个**自下而上**的学习与更新过程：将新发生的对话转化为结构化事件，再进一步纳入长期记忆。
+ # Generate response
+ system_message = {"role": "system", "content": system_prompt}
+ history_messages.append({"role": "user", "content": query})
+ response = openai_client.chat.completions.create(
+ model="qwen-plus",
+ messages=[system_message] + history_messages
+ )
+ history_messages.append({"role": "assistant", "content": response.choices[0].message.content})
-```mermaid
-flowchart TD
- op0[op: 收集 DayN 的原始对话与行为轨迹]
- op1[op: 按时间/业务逻辑/固定长度切片为 messages block]
- op2[op: 对每个 block 做摘要，生成 block_id: time + summary]
- op3[op: 将 block summary 流交给 Meta Agent]
- op4[op: Meta Agent 判断主题并路由/创建对应 Sub Agent]
- op5[op: Sub Agent 加载当前主题的长期记忆与新 block]
- op6[op: 通过 recall_similar_memory 检索历史相似记忆]
- op7[op: 对比新旧信息，决定 delete/new/update 记忆操作]
- op8[op: 写回更新后的高层主题记忆，维护 block_id 链接]
- op9[op: 将最新记忆视图回流给 Meta Agent，用于后续检索]
+ # Summarize history when it gets too long
+ if len(history_messages) >= start_summary_size:
+ memory.summary(history_messages[:-keep_size], user_id=user_name)
+ print("Current memories: " + memory.list_memories(user_id=user_name))
+ history_messages = history_messages[-keep_size:]
- op0 --> op1 --> op2 --> op3 --> op4 --> op5 --> op6 --> op7 --> op8 --> op9
+ return history_messages[-1]["content"]
+
+def main():
+ user_name = input("Enter your name: ").strip()
+ print("Chat with Remy (type 'exit' to quit)")
+
+ messages = []
+ while True:
+ user_input = input(f"{user_name}: ").strip()
+ if user_input.lower() == 'exit':
+ print("Goodbye!")
+ break
+
+ print(f"Remy: {chat_with_memories(user_input, messages, user_name)}")
+
+ # Cleanup
+ memory.delete_all_memories(user_id=user_name)
+ print("All memories deleted")
+
+if __name__ == "__main__":
+ main()
```
-### 输入与切片
+#### 3.1.3 Advanced Usage (高级用法)
-- 输入源：
- - `DayN Messages / Trajectory`（第 N 天的完整对话与行为轨迹）。
-- 切片策略：
- - 按时间窗口切片（例如每 5 分钟一个 block）。
- - 按业务逻辑边界切片（如一个任务完成即成一个 block）。
- - 按固定长度（可带重叠）切片，兼顾局部上下文。
-- 结果：
- - `messages block1, block2, ... blockN`
+For advanced users who want to customize retriever and summarizer behavior with Agentic mode.
-### 初步总结：生成索引条目
+```python
+import os
+from reme_ai import ReMe
+from reme_ai.retriever import AgenticRetriever
+from reme_ai.summarizer import AgenticSummarizer
+from reme_ai.tools import ATool, BTool, CTool
-- 对每个 `messages block` 进行摘要，形成：
- - `block_id: time + summary`
-- 示例：
- - `id: 20251111_01 - 用户询问“什么是 AI”并表达对技术的好奇。`
-- 这一阶段的目标：
- - 为中层索引层提供结构化事件单元，尚不过度抽象。
+os.environ["REME_LLM_API_KEY"] = "sk-..."
+os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+os.environ["REME_EMBEDDING_API_KEY"] = "sk-..."
+os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
-### Meta Agent 路由与子智能体更新
+memory = ReMe(
+ memory_space="remy",
+ llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6},
+ embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024},
+ vector_store={"backend": "local_file"},
+ use_agentic_mode=True,
+)
-- Meta Agent 接收新的 `block summary` 流：
- - 根据内容判断其归属主题（例如用户相关、任务相关、工具体验相关等）。
- - 调用 `create_sub_agent(topic)` 或路由到已有的 `Sub_Agent`。
-- 子智能体内的处理步骤：
- - **Step 1：加载上下文记忆**
- - 加载与当前主题相关的长期记忆（如“对用户 A 的记忆”）与当前 `messages block`。
- - **Step 2：回忆（Recall）**
- - 通过 `recall_similar_memory(query)` 拉取历史相似记忆，用于对比新旧信息。
- - **Step 3：更新（Update）**
- - 基于新旧信息的对比，决定对顶层记忆的操作：
- - `delete_memory(mem_id)`：删除已过时或被否定的记忆。
- - `new_memory(mem_content)`：新增一条新的长期记忆。
- - `update_memory(mem_id, mem_content)`：对既有记忆做内容更新或细化。
+# Customize retriever and summarizer with custom tools and prompts
+memory.set_retriever(
+ AgenticRetriever(tools=[ATool(), BTool(), CTool()]),
+ system_prompt="Custom retrieval instructions..."
+)
+memory.set_summarizer(
+ AgenticSummarizer(tools=[ATool(), BTool(), CTool()])
+)
+
+# Use the customized memory system
+result = memory.summary(
+ messages=[
+ {"role": "user", "content": "I'm travelling to SF"},
+ {"role": "assistant", "content": "That's great to hear!"}
+ ],
+ user_id="Alice",
+ memory_type="auto", # auto, personal, procedural, tool
+)
+
+memories = memory.retrieve(
+ query="what is your travel plan?",
+ limit=3,
+ user_id="Alice",
+ memory_type="auto",
+)
+memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"])
+print(memories_str)
+```
+
+### 3.2 Short-Term Memory (短期记忆)
+
+#### 3.2.1 Basic Usage (基础用法)
+
+Context offload/reload API for managing short-term conversational memory within a session.
+
+```python
+import os
+from reme_ai import ReMe
+
+os.environ["REME_LLM_API_KEY"] = "sk-..."
+os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+os.environ["REME_EMBEDDING_API_KEY"] = "sk-..."
+os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+
+memory = ReMe(
+ memory_space="remy",
+ llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6},
+ embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024},
+ vector_store={"backend": "local_file"},
+)
+
+# Offload context when conversation gets too long
+result = memory.offload_context(
+ messages=[
+ {"role": "user", "content": "I'm travelling to SF"},
+ {"role": "assistant", "content": "That's great to hear!"}
+ ],
+)
+
+# Reload relevant context when needed
+memories = memory.reload_context(
+ query="what is your travel plan?",
+ limit=3,
+)
+memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"])
+print(memories_str)
+```
+
+### 3.3 Framework Integration (框架集成)
+
+#### 3.3.1 Integration with AgentScope
+
+Integration example for AgentScope ReActAgent with long-term memory support.
+
+```python
+# TODO: Provide AgentScope integration example
+```
+
+#### 3.3.2 Integration with LangChain
+
+Integration example for LangChain agents with ReMe memory layer.
+
+```python
+# TODO: Provide LangChain integration example
+```
+
+### 3.4 OpenAI Compatible Interface
+
+OpenAI-compatible API interface for seamless integration with existing OpenAI-based applications.
+
+```python
+# TODO: Research and implement OpenAI-compatible interface
+# - Support for threads and assistants API
+# - Compatible with OpenAI SDK
+# - Support for streaming responses
+```
-### 输出与闭环
-- Summarizer 输出一批新的或更新后的 **高层记忆项（New Memories）**，并附带 `block_id` 链接。
-- 这些记忆将：
- - 被 Meta Agent 纳入其主题索引中。
- - 为后续 Retriever 流程提供可检索的抽象知识。
-- 至此，形成从“原始对话 → 索引总结 → 高维记忆 → 检索使用 → 再次总结更新”的完整闭环。
---
-## 总结
+## 四、核心方案设计
-- **分层存储**：底层保留原始细节，中层作为事件索引，顶层沉淀抽象认知；中层通过 `block_id` 将两端连通。
-- **按需调用**：检索时从高层主题记忆出发，由 Meta Agent 路由到 Sub Agent，再逐级回溯到底层对话。
-- **动态增删改**：新对话先被切片与总结，再与旧记忆对比，通过新增、更新、删除操作保持长期记忆“准确 + 精简 + 有时序感”。
-- **渐进式检索与总结**：在检索方向自上而下逐级展开，在学习方向自下而上逐级抽象，实现“用多少展开多少、学多少沉淀多少”的流式记忆管理。
+### 4.1 设计概述
+ReMeV2 采用简洁的架构设计，核心理念为：**ReMeV2 = Tool(s) + Agent(s)**
-``` python code
-base_memory中增加
-# block_id: str = Field(default=..., description="block_id")
-block_ids: List[str] = Field(default=..., description="block_id")
+- **Tool层**：提供原子化的记忆操作能力，包括增删改查、检索、元数据管理等基础操作
+- **Agent层**：基于Tool层构建的智能代理，负责复杂的记忆管理逻辑，如分类总结、渐进式检索等
+- **Runtime层**：内部调度机制，协调Tool和Agent的交互流程
+### 4.2 Tool层设计
-from typing import List, Tuple
+Tool层提供装饰器形式的记忆操作工具，每个工具类通过 `@tool` 装饰器注册，明确定义初始化参数和调用参数。
-from flowllm.core.schema import Message
-from pydantic import BaseModel, Field
+#### 4.2.1 基类：BaseMemoryToolOp
-class MessageModel(Message):
+**初始化参数：**
+- `enable_multiple` (bool): Enable multi-item operation mode. Default: `True`
+- `enable_thinking_params` (bool): Include thinking parameter in tool schema for model reasoning. Default: `False`
+- `memory_metadata_dir` (str): Directory path for storing memory metadata. Default: `"./memory_metadata"`
+
+#### 4.2.2 Tool操作列表
+
+以下是所有Tool操作的完整定义，包括继承关系、初始化参数和调用参数：
+
+| Tool类 | 继承自 | 初始化参数（除基类外） | Tool Call参数（单项模式） | Tool Call参数（多项模式） |
+|----------------------------|------------------|------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------|
+| **AddMemoryOp** | BaseMemoryToolOp | `add_when_to_use` (bool, 默认: False)
`add_metadata` (bool, 默认: True) | `when_to_use` (str, 可选)
`memory_content` (str, 必需)
`metadata` (dict, 可选) | `memories` (array, 必需):
- `when_to_use` (str, 可选)
- `memory_content` (str, 必需)
- `metadata` (dict, 可选) |
+| **UpdateMemoryOp** | BaseMemoryToolOp | 无 | `memory_id` (str, 必需)
`memory_content` (str, 必需)
`metadata` (dict, 可选) | `memories` (array, 必需):
- `memory_id` (str, 必需)
- `memory_content` (str, 必需)
- `metadata` (dict, 可选) |
+| **DeleteMemoryOp** | BaseMemoryToolOp | 无 | `memory_id` (str, 必需) | `memory_ids` (array[str], 必需) |
+| **VectorRetrieveMemoryOp** | BaseMemoryToolOp | `enable_summary_memory` (bool, 默认: False)
`add_memory_type_target` (bool, 默认: False)
`top_k` (int, 默认: 20) | `query` (str, 必需)
`memory_type` (str, 可选, 枚举: [identity, personal, procedural])
`memory_target` (str, 可选) | `query_items` (array, 必需):
- `query` (str, 必需)
- `memory_type` (str, 可选)
- `memory_target` (str, 可选) |
+| **AddMetaMemoryOp** | BaseMemoryToolOp | 无 | `memory_type` (str, 必需, 枚举: [personal, procedural])
`memory_target` (str, 必需) | `meta_memories` (array, 必需):
- `memory_type` (str, 必需)
- `memory_target` (str, 必需) |
+| **ReadMetaMemoryOp** | BaseMemoryToolOp | `enable_tool_memory` (bool, 默认: False)
`enable_identity_memory` (bool, 默认: False) | 无（无输入schema） | N/A (enable_multiple=False) |
+| **AddHistoryMemoryOp** | BaseMemoryToolOp | 无 | `messages` (array[object], 必需) | N/A (enable_multiple=False) |
+| **ReadHistoryMemoryOp** | BaseMemoryToolOp | 无 | `memory_id` (str, 必需) | `memory_ids` (array[str], 必需) |
+| **AddSummaryMemoryOp** | AddMemoryOp | 无（继承自AddMemoryOp） | `summary_memory` (str, 必需)
`metadata` (dict, 可选) | N/A (enable_multiple=False) |
+| **ReadIdentityMemoryOp** | BaseMemoryToolOp | 无 | 无（无输入schema） | N/A (enable_multiple=False) |
+| **UpdateIdentityMemoryOp** | BaseMemoryToolOp | 无 | `identity_memory` (str, 必需) | N/A (enable_multiple=False) |
+| **ThinkToolOp** | BaseAsyncToolOp | `add_output_reflection` (bool, 默认: False) | `reflection` (str, 必需) | N/A |
+| **HandsOffOp** | BaseMemoryToolOp | 无 | `memory_type` (str, 必需, 枚举: [identity, personal, procedural, tool])
`memory_target` (str, 必需) | `memory_tasks` (array, 必需):
- `memory_type` (str, 必需)
- `memory_target` (str, 必需) |
+
+### 4.3 Agent层设计
+
+#### 4.3.1 基类：BaseMemoryAgentOp
+
+Agent层构建在Tool层之上，封装复杂的记忆管理逻辑。每个Agent通过组合多个Tool实现特定的记忆管理任务。
+
+#### 4.3.2 Agent操作列表
+
+以下是所有Agent操作的完整定义，包括初始化参数、调用参数和可用工具：
+
+| Agent类 | 继承自 | 初始化参数（基类外） | Tool Call参数 | 可用工具 |
+|--------------------------------|-------------------|------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------|
+| **PersonalSummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`memory_target` (str, required)
`query` (str, optional)
`messages` (array, optional)
`ref_memory_id` (str, required) | add_memory
update_memory
delete_memory
vector_retrieve_memory |
+| **ProceduralSummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`memory_target` (str, required)
`query` (str, optional)
`messages` (array, optional)
`ref_memory_id` (str, required) | add_memory
update_memory
delete_memory
vector_retrieve_memory |
+| **ToolSummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`memory_target` (str, required)
`query` (str, optional)
`messages` (array, optional)
`ref_memory_id` (str, required) | add_memory
update_memory
vector_retrieve_memory |
+| **IdentitySummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | read_identity_memory
update_identity_memory |
+| **ReMeSummaryAgentV1Op** | BaseMemoryAgentOp | `enable_tool_memory` (bool, 默认: True)
`enable_identity_memory` (bool, 默认: True) | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | add_meta_memory
add_summary_memory
hands_off
(内部调用: add_history_memory, read_identity_memory, read_meta_memory) |
+| **ReMeRetrieveAgentV1Op** | BaseMemoryAgentOp | `enable_tool_memory` (bool, 默认: True) | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | vector_retrieve_memory
read_history_memory
(内部调用: read_meta_memory) |
+| **ReMyAgentV1Op** | BaseMemoryAgentOp | `enable_tool_memory` (bool, 默认: True)
`enable_identity_memory` (bool, 默认: True) | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | vector_retrieve_memory
read_history_memory
(内部调用: read_identity_memory, read_meta_memory) |
+
+### 4.4 Runtime层设计（内部实现）
+
+Runtime层负责协调Tool和Agent的调用流程，实现记忆的渐进式处理。
+
+#### 4.4.1 渐进式总结流程（Summary）
+
+总结流程采用分层处理策略，首先保存历史对话，读取元信息，然后由主Agent协调多个专用Agent完成分类总结。
+
+**流程结构：**
+
+```python
+# Step 1: Save conversation history
+AddHistoryMemoryOp()
+
+# Step 2: Load meta information (memory types and targets)
+ReadMetaMemoryOp()
+
+# Step 3: Progressive summarization with delegation
+ReMeSummaryAgentV1Op(tools=[
+ # Add meta memory entries for new memory types/targets
+ AddMetaMemoryOp(list(memory_type, memory_target)),
+
+ # Add general summary memory as fallback
+ AddSummaryMemoryOp(summary_memory),
+
+ # Delegate to specialized summary agents
+ HandsOffOp(list(memory_type, memory_target), agents=[
+ PersonalSummaryAgentV1Op, # Summarize personal memories
+ ProceduralSummaryAgentV1Op, # Summarize procedural memories
+ ToolSummaryAgentV1Op, # Summarize tool-related memories
+ IdentitySummaryAgentV1Op # Update identity memory
+ ]),
+])
+
+# Specialized agents and their available tools
+PersonalSummaryAgentV1Op(tools=[AddMemoryOp, UpdateMemoryOp, DeleteMemoryOp, VectorRetrieveMemoryOp])
+ProceduralSummaryAgentV1Op(tools=[AddMemoryOp, UpdateMemoryOp, DeleteMemoryOp, VectorRetrieveMemoryOp])
+ToolSummaryAgentV1Op(tools=[AddMemoryOp, UpdateMemoryOp, VectorRetrieveMemoryOp])
+IdentitySummaryAgentV1Op(tools=[ReadIdentityMemoryOp, UpdateIdentityMemoryOp])
+```
+
+#### 4.4.2 渐进式检索流程（Retrieve）
+
+检索流程采用三层检索策略，类似于技能系统的加载机制，逐层加载和过滤记忆。
+
+**流程结构：**
+
+```python
+# Progressive retrieval with three layers
+ReMeRetrieveAgentV1Op(tools=[
+ # Layer 0: Load meta memory (all available memory types and targets)
+ ReadMetaMemoryOp(),
+ # Output format example:
+ # - personal(jinli): Information about Jinli's personal life and preferences
+ # - personal(jiaji): Information about Jiaji's background and interests
+ # - personal(jinli&jiaji): Shared memories between Jinli and Jiaji
+ # - procedural(appworld): Procedural knowledge for AppWorld tasks
+ # - procedural(bfcl-v3): Procedural knowledge for BFCL-v3 benchmark
+ # - tool(tool_guidelines): Guidelines for tool usage
+ # - identity(self): Agent's self-identity information
+
+ # Layer 1+2: Vector-based retrieval on structured memories
+ VectorRetrieveMemoryOp(list(memory_type, memory_target, query)),
+
+ # Layer 3: Load full conversation history for specific memory
+ ReadHistoryMemoryOp(ref_memory_id),
+])
+```
+
+**与技能系统的类比：**
+
+```python
+# Skill system hierarchy (for reference)
+load_meta_skills # Load skill metadata
+load_skills # Load skill implementations
+load_reference_skills # Load detailed skill documentation
+execute_shell # Execute actual commands
+```
+
+## 五、扩展设计与实验方向
+
+### 5.1 Summary Memory机制
+
+Summary Memory作为通用维度的记忆类型，提供兜底的原始对话索引能力。
+
+**工作流程示例：**
+
+```txt
+Step 1: Progressive summarization across sessions
+ session1: List[Message] -> session2: List[Message] -> session3: List[Message] -> ...
+summary ✓ (always) ✓ (always) ✓ (always)
+personal ✗ ✗ ✓ (when applicable)
+procedural ✗ ✓ (when applicable) ✗
+
+Step 2: Retrieval with fallback strategy
+vector_retrieve_memory(query, memory_type="personal", memory_target="jinli")
+ -> Search in memory_type: ["personal", "summary"] # Fallback to summary if personal not found
+```
+
+**设计优势：**
+1. Provides a universal dimension for memory extraction across all memory types
+2. Ensures fallback indexing of original conversations when specific meta memory is not available
+3. Maintains conversation context even when specialized memory extraction fails
+
+### 5.2 Thinking参数实验
+
+探索不同的模型推理能力增强方案，受AgentScope和Claude启发。
+
+#### 5.2.1 Thinking参数设计
+
+```python
+async def record_to_memory(
+ self,
+ thinking: str,
+ content: list[str],
+ **kwargs: Any,
+) -> ToolResponse:
+ """Use this function to record important information that you may
+ need later. The target content should be specific and concise, e.g.
+ who, when, where, do what, why, how, etc.
+
+ Args:
+ thinking (`str`):
+ Your thinking and reasoning about what to record
+ content (`list[str]`):
+ The content to remember, which is a list of strings.
"""
- 被存储到传统的db中的数据结构，一行一条message
- def retrieve_message(key_word=None, time_start=None, limit: int=50):
- 支持关键字检索和时间检索
- content like %{key_word}%
- 暂时不支持向量检索，向量只做抽取的memory的
- """
- name: str = Field(default="", description="The name of the character who actually outputs this message")
- message_id: str = Field(default=..., description="message_id, uuid，每一条message不重样")
- block_id: str = Field(default=..., description="block_id，message分块的id，用户链接summary")
+```
+#### 5.2.2 实验对比方案
-from reme_ai.schema.memory import BaseMemory
+| 方案类型 | 说明 | 灵感来源 |
+|-------------------------------|--------------------------------------------|----------------|
+| Thinking Model | Native reasoning-capable models (e.g., o1) | OpenAI |
+| Instruct Model | Standard instruction-following models | Baseline |
+| Instruct Model + Thinking Params | Add thinking parameter to tool schema | AgentScope |
+| Instruct Model + Thinking Tool | Dedicated thinking tool for explicit reasoning | Claude |
-"""
-基本和原来一致，增加block compress的memory
+### 5.3 多项操作模式实验
-工具记忆放到传统db中
-BlockCompress放到向量库中
-Personal + Task + Self-Cognition 放到向量库中
-"""
+对比单次调用和批量调用的性能与准确性差异。
+**两种模式对比：**
+| 模式 | Tool调用方式 | Model调用次数 | 优势 | 劣势 |
+|--------------|----------------------------|---------------|------------------------------|--------------------------|
+| 单项模式 | Single-item per call | Multiple | Fine-grained control | Higher latency, more tokens |
+| 多项模式 | Batch multiple items | Single | Lower latency, fewer tokens | Potential batch errors |
+**实验目标：**
+- Evaluate accuracy: single vs. batch operations
+- Measure latency and token efficiency
+- Identify optimal use cases for each mode
-class BlockCompressMemory(BaseMemory):
- start_time: str = Field(default=..., description="block的开始时间")
- end_time: str = Field(default=..., description="block的结束时间")
+### 5.4 多版本与扩展性
-class SelfCognitionMemory(BaseMemory):
- target_name: str = Field(default="自我认知")
- memory_type: str = Field(default="self")
+支持从基类继承实现自定义Agent，便于团队协作和功能迭代。
-class ToolMemory(BaseMemory):
- target_name: str = Field(default="", description="工具名称")
- memory_type: str = Field(default="tool")
+**扩展示例：**
-class PersonalMemory(BaseMemory):
- target_name: str = Field(default="", description="这是记录谁的记忆？比如TOM、Jerry？")
- memory_type: str = Field(default="personal")
+```python
+# Version 2 implementations by different team members
+PersonalSummaryAgentV2Op / PersonalRetrieveAgentV2Op # @weikang
+ProceduralSummaryAgentV2Op / ProceduralRetrieveAgentV2Op # @zouyin
-class TaskMemory(BaseMemory):
- target_name: str = Field(default="", description="这是什么任务的记忆、比如bfcl、appworld、金融?")
- memory_type: str = Field(default="task")
+# Inherit from BaseMemoryAgentOp
+class PersonalSummaryAgentV2Op(BaseMemoryAgentOp):
+ """Enhanced personal memory summarization with improved algorithms"""
+ pass
+```
+### 5.5 文件系统集成（未来方向）
+探索将文件操作能力集成到记忆系统中，支持基于文件的记忆管理。
-# Handoffs allow meta agent to delegate tasks to another agent.
-"""
-summarizer
-"""
+**挑战与考虑：**
-def compress_message_block(message_block: List[Message]) -> BlockCompressMemory:
- """
- 压缩message_block信息到BlockCompressMemory
- """
+1. **操作适配性**：Current operations (retrieve/add/update/delete) need adaptation for file-based storage
+2. **工具选择**：Consider file operation tools: `grep`, `glob`, `ls`, `read_file`, `write_file`, `edit_file`
+3. **模型能力**：Base models have limited file operation capabilities; `qwen3-code` shows better performance
-def handoff_to_summary_agent(agent_name: str) -> str:
- """
- 如果agent_name是存在的，那直接使用现成的，否则可以创建新的agent(默认是不开启的)
- """
+**潜在架构：**
-def retrieve_memory_by_query(query: str, limit_size: int = 50, threshold = 0.5) -> List[BaseMemory]:
- """
- 如果之前使用过handoff_to_summary_agent
- 在自我认知
- """
+```python
+# File-based memory operations
+FileMemoryOp(tools=[
+ grep, # Search within files
+ glob, # File pattern matching
+ ls, # List directory contents
+ read_file, # Read file contents
+ write_file, # Write new memory files
+ edit_file, # Update existing memory files
+])
+```
-def datetime_tool(**kwargs)
- """
- 时间工具
- Returns:
+### 5.6 自我修改上下文
- """
+支持Agent动态修改自身的上下文状态，实现自适应记忆管理。
-def delete_memory(memory_id: str):
- ...
+**实现方式：**
-# 做实验
-def update_memory(memory_id: str, content: str, time_stamp: str):
- ...
+1. **Summary Agent 主动修改**：
+ - `add_meta_memory` directly modifies agent context
+ - Updates available memory types and targets during execution
+2. **ReMy Agent 被动修改**：
+ - Retrieves `identity_memory` at each interaction
+ - Dynamically updates self-state based on retrieved identity
+ - Enables adaptive behavior based on accumulated identity knowledge
-def new_memory(content: str, time_stamp: str) -> BaseMemory:
- ...
+## ReMe V2 开发路线图与实施计划
+### 技术改造阶段
+1. **代码整合与兼容**：合并flowllm中reme必要的代码，保留现在server-client的依赖，兼容现在各个仓库的依赖代码
+2. **核心接口重构**：新的ReMe接口设计，支持summary，retrieve，context_offload, context_reload 4个核心接口
+3. **Agentic算法升级**：新的agentic算法方案开发
-"""
-summary reviser / Reward Model
-"""
+### 评估验证阶段
+4. **Benchmark测试**
+ - halumem
+ - locomo
+ - longmemevel
+ - personal-v2 ?
+ - appworld/bfcl-v3
-def revise_memory(messages: List[Message],
- origin_memory: List[BaseMemory],
- output_memory: List[BaseMemory]) -> Tuple["", bool]:
- """
- reasoning and result
- 是否所有信息都加入了？
- 当前结果是不是有冲突和重复。
- 如果有问题，就重新react或者重试
- """
+### 发布推广阶段
+5. **技术报告**撰写与发布
+6. **生态更新**：更新各个仓库的依赖代码
+ - agentscope
+ - agentscope-runtime
+ - evotraders
+ - alias(tool-memory)
+ - agentscope-java
+ - AgentEvolver
+ - cookbook: reme procedural memory paper
+ - tool-memory-upgrade（将要合并）
-"""
-retriever
-"""
+**里程碑目标**：春节前完成小版本发布
-def retrieve_memory_by_query(query: str, limit_size: int = 50, threshold = 0.5) -> List[BaseMemory]:
- """
- 和上面不一致的在于，这里可以搜索到compress的记忆
- """
+---
-def retrieve_message(key_word=None, time_offset=None, time_limit=None):
- """
- 支持关键字检索和时间检索
- content like %{key_word}%
- 暂时不支持向量检索，向量只做抽取的memory的
- """
+## ReMe V2 核心竞争优势
-def retrieve_block_id(block_id):
- """
- 支持关键字检索和时间检索
- content like %{key_word}%
- 暂时不支持向量检索，向量只做抽取的memory的
- """
+### 1. 渐进式 Agentic Memory 架构【核心创新】
+融合了多种记忆的渐进式agentic方案，实现从短期到长期记忆的智能化演进
-```
\ No newline at end of file
+### 2. 全生命周期记忆管理
+同时支持长期记忆（Long-term Memory）和短期记忆（Working Memory），完整覆盖Agent认知周期
+
+### 3. 模型
+提供开源小模型
+
+### 4. 开发者友好生态
+ 1. **简洁接口**：提供简洁的接口设计，全异步接口
+ 2. **即开即用**：提供CLI工具，开箱即用的体验
+ 3. **生态融合**：提供和AgentScope、LangChain无缝集成的方案
+ 4. **高度可扩展**：支持Agentic算法的二次开发与定制
\ No newline at end of file
diff --git a/reme_ai/core/application.py b/reme_ai/core/application.py
new file mode 100644
index 00000000..e15f4fcf
--- /dev/null
+++ b/reme_ai/core/application.py
@@ -0,0 +1,208 @@
+"""Main application module for managing ReMe AI service lifecycle and flow execution."""
+
+import asyncio
+import os
+
+from .context import C
+from .flow import BaseFlow
+from .schema import ServiceConfig, Response
+from .utils import PydanticConfigParser, init_logger, execute_stream_task, run_coro_safely
+
+
+class Application:
+ """
+ Main application class for managing the lifecycle of ReMe AI services.
+
+ Handles initialization, configuration, service management, and flow execution
+ for both synchronous and asynchronous contexts.
+ """
+
+ def __init__(
+ self,
+ *args,
+ llm_api_key: str | None = None,
+ llm_api_base: str | None = None,
+ embedding_api_key: str | None = None,
+ embedding_api_base: str | None = None,
+ service_config: ServiceConfig | None = None,
+ parser: type[PydanticConfigParser] | None = None,
+ config_path: str | None = None,
+ enable_logo: bool = True,
+ llm: dict | None = None,
+ embedding_model: dict | None = None,
+ vector_store: dict | None = None,
+ token_counter: dict | None = None,
+ **kwargs,
+ ):
+ """
+ Initialize the Application with configuration settings.
+
+ Args:
+ llm_api_key: API key for LLM service
+ llm_api_base: Base URL for LLM service
+ embedding_api_key: API key for embedding service
+ embedding_api_base: Base URL for embedding service
+ service_config: Pre-built service configuration object
+ parser: Custom parser class for configuration (defaults to PydanticConfigParser)
+ config_path: Path to configuration file
+ enable_logo: Whether to display the ReMe logo on startup
+ llm: LLM configuration dictionary
+ embedding_model: Embedding model configuration dictionary
+ vector_store: Vector store configuration dictionary
+ token_counter: Token counter configuration dictionary
+ **kwargs: Additional configuration arguments
+ """
+ self._update_env("REME_LLM_API_KEY", llm_api_key)
+ self._update_env("REME_LLM_BASE_URL", llm_api_base)
+ self._update_env("REME_EMBEDDING_API_KEY", embedding_api_key)
+ self._update_env("REME_EMBEDDING_BASE_URL", embedding_api_base)
+
+ init_logger()
+
+ # Use default parser if not provided
+ parser_class = parser if parser is not None else PydanticConfigParser
+ self.parser = parser_class(ServiceConfig)
+
+ if service_config is None:
+ input_args = []
+ if config_path:
+ input_args.append(f"config={config_path}")
+ if args:
+ input_args.extend(args)
+ if kwargs:
+ input_args.extend([f"{k}={v}" for k, v in kwargs.items()])
+ service_config = self.parser.parse_args(*input_args)
+
+ C.service_config = service_config
+
+ if llm:
+ C.update_section_config("llm", **llm)
+ if embedding_model:
+ C.update_section_config("embedding_model", **embedding_model)
+ if vector_store:
+ C.update_section_config("vector_store", **vector_store)
+ if token_counter:
+ C.update_section_config("token_counter", **token_counter)
+ C.service_config.enable_logo = enable_logo
+ C.print_logo()
+
+ @staticmethod
+ def _update_env(key: str, value: str | None) -> None:
+ """Update environment variable if value is provided."""
+ if value:
+ os.environ[key] = value
+
+ @staticmethod
+ async def start() -> None:
+ """Initialize the service context and prepare external MCP servers."""
+ C.initialize_service_context()
+ await C.prepare_mcp_servers()
+
+ @staticmethod
+ def start_sync() -> None:
+ """Synchronous version of start()."""
+ C.initialize_service_context()
+ run_coro_safely(C.prepare_mcp_servers())
+
+ @staticmethod
+ async def stop(wait_thread_pool: bool = True, wait_ray: bool = True) -> None:
+ """
+ Stop the application and cleanup resources.
+
+ Args:
+ wait_thread_pool: Whether to wait for thread pool shutdown
+ wait_ray: Whether to wait for Ray shutdown
+ """
+ await C.close()
+ C.shutdown_thread_pool(wait=wait_thread_pool)
+ C.shutdown_ray(wait=wait_ray)
+
+ @staticmethod
+ def stop_sync(wait_thread_pool: bool = True, wait_ray: bool = True) -> None:
+ """Synchronous version of stop()."""
+ C.close_sync()
+ C.shutdown_thread_pool(wait=wait_thread_pool)
+ C.shutdown_ray(wait=wait_ray)
+
+ async def __aenter__(self):
+ """Async context manager entry."""
+ await self.start()
+ return self
+
+ def __enter__(self):
+ """Context manager entry."""
+ self.start_sync()
+ return self
+
+ async def __aexit__(self, exc_type=None, exc_val=None, exc_tb=None):
+ """Async context manager exit."""
+ await self.stop()
+ return False
+
+ def __exit__(self, exc_type=None, exc_val=None, exc_tb=None):
+ """Context manager exit."""
+ self.stop_sync()
+ return False
+
+ @staticmethod
+ async def execute_flow(name: str, **kwargs) -> Response:
+ """
+ Execute a flow asynchronously.
+
+ Args:
+ name: Name of the flow to execute
+ **kwargs: Arguments to pass to the flow
+
+ Returns:
+ Response object from the flow execution
+ """
+ flow: BaseFlow = C.get_flow(name)
+ return await flow.call(**kwargs)
+
+ @staticmethod
+ def execute_flow_sync(name: str, **kwargs) -> Response:
+ """
+ Execute a flow synchronously.
+
+ Args:
+ name: Name of the flow to execute
+ **kwargs: Arguments to pass to the flow
+
+ Returns:
+ Response object from the flow execution
+ """
+ flow: BaseFlow = C.get_flow(name)
+ return flow.call_sync(**kwargs)
+
+ @staticmethod
+ async def execute_stream_flow(name: str, **kwargs):
+ """
+ Execute a streaming flow asynchronously.
+
+ Args:
+ name: Name of the streaming flow to execute
+ **kwargs: Arguments to pass to the flow
+
+ Yields:
+ Stream chunks from the flow execution
+
+ Raises:
+ AssertionError: If the flow is not configured for streaming
+ """
+ flow: BaseFlow = C.get_flow(name)
+ assert flow.stream is True, "non-stream flow is not supported in execute_stream_flow!"
+ stream_queue = asyncio.Queue()
+ task = asyncio.create_task(flow.call(stream_queue=stream_queue, **kwargs))
+
+ async for chunk in execute_stream_task(
+ queue=stream_queue,
+ task=task,
+ flow_name=name,
+ as_bytes=False,
+ ):
+ yield chunk
+
+ @staticmethod
+ def run_service():
+ """Run the configured service (HTTP, MCP, or CMD)."""
+ C.get_service().run()
diff --git a/reme_ai/core/config/__init__.py b/reme_ai/core/config/__init__.py
new file mode 100644
index 00000000..a8d954d2
--- /dev/null
+++ b/reme_ai/core/config/__init__.py
@@ -0,0 +1,5 @@
+"""config"""
+
+from .reme_config_parser import ReMeConfigParser
+
+__all__ = ["ReMeConfigParser"]
diff --git a/reme_ai/core/config/default.yaml b/reme_ai/core/config/default.yaml
new file mode 100644
index 00000000..cd04da1d
--- /dev/null
+++ b/reme_ai/core/config/default.yaml
@@ -0,0 +1,34 @@
+backend: http
+thread_pool_max_workers: 64
+
+mcp:
+ transport: sse
+ host: "0.0.0.0"
+ port: 8001
+
+http:
+ host: "0.0.0.0"
+ port: 8002
+ timeout_keep_alive: 600
+ limit_concurrency: 64
+
+llm:
+ default:
+ backend: openai
+ model_name: qwen3-30b-a3b-instruct-2507
+ temperature: 0.6
+
+embedding_model:
+ default:
+ backend: openai
+ model_name: text-embedding-v4
+ dimensions: 1024
+
+vector_store:
+ default:
+ backend: local
+ embedding_model: default
+
+token_counter:
+ default:
+ backend: base
diff --git a/reme_ai/core/config/reme_config_parser.py b/reme_ai/core/config/reme_config_parser.py
new file mode 100644
index 00000000..798235b2
--- /dev/null
+++ b/reme_ai/core/config/reme_config_parser.py
@@ -0,0 +1,7 @@
+"""Configuration parser for ReMe framework."""
+
+from ..utils import PydanticConfigParser
+
+
+class ReMeConfigParser(PydanticConfigParser):
+ """Configuration parser for ReMe framework."""
diff --git a/reme_ai/core/context/runtime_context.py b/reme_ai/core/context/runtime_context.py
index ded35ded..d7112e1c 100644
--- a/reme_ai/core/context/runtime_context.py
+++ b/reme_ai/core/context/runtime_context.py
@@ -1,15 +1,14 @@
-"""Module providing a runtime context for managing response states and asynchronous data streaming."""
+"""Runtime context for managing response states and asynchronous data streaming."""
import asyncio
from .base_context import BaseContext
from ..enumeration import ChunkEnum
-from ..schema import Response
-from ..schema import StreamChunk
+from ..schema import Response, StreamChunk
class RuntimeContext(BaseContext):
- """A context class for handling execution state, including response metadata and stream queues."""
+ """Context for execution state, response metadata, and stream queues."""
def __init__(
self,
@@ -17,40 +16,64 @@ class RuntimeContext(BaseContext):
stream_queue: asyncio.Queue | None = None,
**kwargs,
):
- """Initialize the runtime context with optional response objects and message queues."""
+ """Initialize the context with optional response and queue."""
super().__init__(**kwargs)
+ self.response = response or Response()
+ self.stream_queue = stream_queue
- self.response: Response | None = response if response is not None else Response()
- self.stream_queue: asyncio.Queue | None = stream_queue
+ @classmethod
+ def from_context(cls, context: "RuntimeContext | None" = None, **kwargs) -> "RuntimeContext":
+ """Create a new context from an existing instance or keywords."""
+ if context is None:
+ return cls(**kwargs)
- async def add_stream_string_and_type(self, chunk: str, chunk_type: ChunkEnum):
- """Create and enqueue a stream chunk from a raw string and specific type."""
- if self.stream_queue is None:
- return self
+ context.update(kwargs)
+ return context
- # Package raw data into a StreamChunk schema
- stream_chunk = StreamChunk(chunk_type=chunk_type, chunk=chunk)
- await self.stream_queue.put(stream_chunk)
+ async def _enqueue(self, chunk: StreamChunk) -> None:
+ """Internal helper to put a chunk into the queue if it exists."""
+ if self.stream_queue:
+ await self.stream_queue.put(chunk)
+
+ async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext":
+ """Enqueue a stream chunk from a raw string and type."""
+ await self._enqueue(StreamChunk(chunk_type=chunk_type, chunk=chunk))
return self
- async def add_stream_chunk(self, stream_chunk: StreamChunk):
- """Directly enqueue an existing stream chunk into the stream queue."""
- if self.stream_queue is None:
- return self
- await self.stream_queue.put(stream_chunk)
+ async def add_stream_chunk(self, stream_chunk: StreamChunk) -> "RuntimeContext":
+ """Enqueue an existing stream chunk."""
+ await self._enqueue(stream_chunk)
return self
- async def add_stream_done(self):
- """Enqueue a termination chunk to signal the end of the data stream."""
- if self.stream_queue is None:
- return self
-
- # Create a special chunk representing the completion state
- done_chunk = StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True)
- await self.stream_queue.put(done_chunk)
+ async def add_stream_done(self) -> "RuntimeContext":
+ """Enqueue a termination chunk to signal the end of the stream."""
+ await self._enqueue(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True))
return self
- def add_response_error(self, e: Exception):
- """Update the internal response object to reflect a failure state using exception details."""
+ def add_response_error(self, e: Exception) -> "RuntimeContext":
+ """Record an exception into the response object."""
self.response.success = False
- self.response.answer = str(e.args)
+ self.response.answer = str(e)
+ return self
+
+ def apply_mapping(self, mapping: dict[str, str]) -> "RuntimeContext":
+ """Copy internal values based on a source-to-target key map."""
+ if not mapping:
+ return self
+
+ for source, target in mapping.items():
+ if source in self:
+ self[target] = self[source]
+ return self
+
+ def validate_required_keys(self, required_keys: dict[str, bool], context_name: str = "context") -> "RuntimeContext":
+ """Ensure all required keys are present in the context.
+
+ Args:
+ required_keys: Dictionary mapping key names to boolean indicating if required
+ context_name: Name of the context for error messages (e.g., operator name)
+ """
+ for key, is_required in required_keys.items():
+ if is_required and key not in self:
+ raise ValueError(f"{context_name}: missing required input '{key}'")
+ return self
diff --git a/reme_ai/core/context/service_context.py b/reme_ai/core/context/service_context.py
index af99a870..c5bcde0c 100644
--- a/reme_ai/core/context/service_context.py
+++ b/reme_ai/core/context/service_context.py
@@ -1,34 +1,80 @@
"""Module for managing global service configurations and component registries via a singleton context."""
from concurrent.futures import ThreadPoolExecutor
-from typing import Dict
+from typing import TYPE_CHECKING
+
+from loguru import logger
from .base_context import BaseContext
from .registry import Registry
from ..enumeration import RegistryEnum
from ..schema import ServiceConfig
-from ..utils import singleton
+from ..utils import singleton, print_logo
+
+if TYPE_CHECKING:
+ from ..llm import BaseLLM
+ from ..embedding import BaseEmbeddingModel
+ from ..vector_store import BaseVectorStore
+ from ..token_counter import BaseTokenCounter
+ from ..flow import BaseFlow
+ from ..service import BaseService
@singleton
class ServiceContext(BaseContext):
- """A singleton container for global application state, thread pools, and component registries."""
+ """A singleton container for global application state, thread pools, and component registries.
+
+ This class serves as the central management hub for the entire ReMe application, providing:
+ - Service configuration management
+ - Component registration and instantiation (LLMs, embeddings, vector stores, etc.)
+ - Thread pool and Ray distributed computing management
+ - MCP (Model Context Protocol) server integration
+
+ The singleton pattern ensures only one instance exists throughout the application lifecycle,
+ accessible via the global `C` variable exported at the bottom of this module.
+ """
def __init__(self, **kwargs):
- """Initialize the global context with configuration objects and specialized registries."""
+ """Initialize the global context with configuration objects and specialized registries.
+
+ Sets up:
+ - Empty service configuration placeholder
+ - Thread pool for concurrent operations
+ - Registry dictionaries for class registration (templates)
+ - Instance dictionaries for instantiated objects (actual instances)
+ - MCP server mapping for external tool integration
+ """
super().__init__(**kwargs)
+ # Service configuration and runtime settings
self.service_config: ServiceConfig | None = None
self.language: str = ""
self.thread_pool: ThreadPoolExecutor | None = None
- self.vector_store_dict: Dict[str, dict] = {}
- self.external_mcp_tool_call_dict: dict = {}
- # Initialize a registry for every category defined in RegistryEnum
- self.registry_dict: Dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()}
- self.flow_dict: dict = {}
+
+ # Registry system: stores class definitions for different component types
+ self.registry_dict: dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()}
+
+ # Instance system: stores instantiated objects created from registered classes
+ self.instance_dict: dict[RegistryEnum, dict] = {v: {} for v in RegistryEnum.__members__.values()}
+
+ # MCP server mapping: maps server_name -> {tool_name: ToolCall}
+ self.mcp_server_mapping: dict[str, dict] = {}
def register(self, name: str, register_type: RegistryEnum):
- """Return a decorator to register a component within a specific registry category."""
+ """Return a decorator to register a component within a specific registry category.
+
+ Args:
+ name: The registration name for the component (used for lookup)
+ register_type: The type of registry (LLM, EMBEDDING_MODEL, VECTOR_STORE, etc.)
+
+ Returns:
+ A decorator function that registers the decorated class
+
+ Example:
+ @C.register("my_llm", RegistryEnum.LLM)
+ class MyLLM(BaseLLM):
+ pass
+ """
return self.registry_dict[register_type].register(name=name)
def register_llm(self, name: str = ""):
@@ -60,18 +106,29 @@ class ServiceContext(BaseContext):
return self.register(name=name, register_type=RegistryEnum.TOKEN_COUNTER)
def get_model_class(self, name: str, register_type: RegistryEnum):
- """Retrieve a registered class by name from a specific registry category."""
+ """Retrieve a registered class by name from a specific registry category.
+
+ Args:
+ name: The registration name of the class
+ register_type: The type of registry to search in
+
+ Returns:
+ The registered class (not an instance, but the class itself)
+
+ Raises:
+ AssertionError: If the class is not found in the registry
+ """
assert name in self.registry_dict[register_type], f"{name} not in registry_dict[{register_type}]"
return self.registry_dict[register_type][name]
- def get_embedding_model_class(self, name: str):
- """Get the embedding model class registered under the given name."""
- return self.get_model_class(name, RegistryEnum.EMBEDDING_MODEL)
-
def get_llm_class(self, name: str):
"""Get the LLM class registered under the given name."""
return self.get_model_class(name, RegistryEnum.LLM)
+ def get_embedding_model_class(self, name: str):
+ """Get the embedding model class registered under the given name."""
+ return self.get_model_class(name, RegistryEnum.EMBEDDING_MODEL)
+
def get_vector_store_class(self, name: str):
"""Get the vector store class registered under the given name."""
return self.get_model_class(name, RegistryEnum.VECTOR_STORE)
@@ -92,14 +149,378 @@ class ServiceContext(BaseContext):
"""Get the token counter class registered under the given name."""
return self.get_model_class(name, RegistryEnum.TOKEN_COUNTER)
- def get_vector_store(self, name: str):
- """Retrieve a specific vector store instance by name."""
- return self.vector_store_dict[name]
+ def get_llm(self, name: str) -> "BaseLLM":
+ """Retrieve a specific LLM instance by name.
- def get_flow(self, name: str):
- """Retrieve a specific flow instance by name."""
- return self.flow_dict[name]
+ Args:
+ name: The name of the LLM instance (typically 'default' or custom name)
+
+ Returns:
+ The instantiated LLM object
+
+ Raises:
+ KeyError: If no LLM with the given name exists
+ """
+ return self.instance_dict[RegistryEnum.LLM][name]
+
+ def get_embedding_model(self, name: str) -> "BaseEmbeddingModel":
+ """Retrieve a specific embedding model instance by name.
+
+ Args:
+ name: The name of the embedding model instance (typically 'default')
+
+ Returns:
+ The instantiated embedding model object
+
+ Raises:
+ KeyError: If no embedding model with the given name exists
+ """
+ return self.instance_dict[RegistryEnum.EMBEDDING_MODEL][name]
+
+ def get_vector_store(self, name: str) -> "BaseVectorStore":
+ """Retrieve a specific vector store instance by name.
+
+ Args:
+ name: The name of the vector store instance (typically 'default')
+
+ Returns:
+ The instantiated vector store object
+
+ Raises:
+ KeyError: If no vector store with the given name exists
+ """
+ return self.instance_dict[RegistryEnum.VECTOR_STORE][name]
+
+ def get_token_counter(self, name: str) -> "BaseTokenCounter":
+ """Retrieve a specific token counter instance by name.
+
+ Args:
+ name: The name of the token counter instance (typically 'default')
+
+ Returns:
+ The instantiated token counter object
+
+ Raises:
+ KeyError: If no token counter with the given name exists
+ """
+ return self.instance_dict[RegistryEnum.TOKEN_COUNTER][name]
+
+ def get_flow(self, name: str) -> "BaseFlow":
+ """Retrieve a specific flow instance by name.
+
+ Args:
+ name: The name of the flow instance
+
+ Returns:
+ The instantiated flow object
+
+ Raises:
+ KeyError: If no flow with the given name exists
+ """
+ return self.instance_dict[RegistryEnum.FLOW][name]
+
+ def get_service(self) -> "BaseService":
+ """Retrieve the default service instance.
+
+ Returns:
+ The instantiated service backend (HTTP, MCP, or CMD service)
+
+ Raises:
+ KeyError: If the default service was not initialized
+ """
+ return self.instance_dict[RegistryEnum.SERVICE]["default"]
+
+ def update_section_config(self, section_name: str, **kwargs):
+ """Update a specific section of the service config with new values.
+
+ Args:
+ section_name: Name of the config section (e.g., 'llm', 'embedding_model')
+ **kwargs: Key-value pairs to update in the default configuration
+
+ Raises:
+ KeyError: If the default config for the section doesn't exist
+
+ Example:
+ update_section_config('llm', temperature=0.8, max_tokens=1000)
+ """
+ if not hasattr(self.service_config, section_name) or not kwargs:
+ return
+
+ section_dict: dict = getattr(self.service_config, section_name)
+ if "default" not in section_dict:
+ raise KeyError(f"Default `{section_name}` config not found")
+
+ current_config = section_dict["default"]
+ section_dict["default"] = current_config.model_copy(update=kwargs, deep=True)
+
+ def initialize_service_context(self):
+ """Initialize the service context with the configuration.
+
+ This is the main initialization method that sets up all system components in order:
+ 1. Language settings
+ 2. Thread pool for concurrent operations
+ 3. Ray cluster (if configured for distributed computing)
+ 4. LLM instances
+ 5. Embedding model instances
+ 6. Token counter instances
+ 7. Vector store instances (with their embedding models)
+ 8. Flow instances (both registered and configured)
+ 9. Service backend instance
+
+ Note: This method should be called after service_config is set.
+ """
+ self.language = self.service_config.language
+ self.thread_pool = ThreadPoolExecutor(max_workers=self.service_config.thread_pool_max_workers)
+
+ # Initialize Ray for distributed computing if configured
+ if self.service_config.ray_max_workers > 1:
+ import ray
+
+ ray.init(num_cpus=self.service_config.ray_max_workers)
+
+ # Initialize components in dependency order
+ self._initialize_llm()
+ self._initialize_embedding_model()
+ self._initialize_token_counter()
+ self._initialize_vector_store() # Depends on embedding models
+ self._initialize_flow()
+ self._initialize_service()
+
+ def _initialize_llm(self):
+ """Initialize all configured LLM instances.
+
+ For each LLM configuration:
+ - Retrieves the corresponding registered LLM class by backend name
+ - Instantiates it with model_name and additional configuration
+ - Stores the instance in instance_dict for later retrieval
+ """
+ for name, config in self.service_config.llm.items():
+ llm_cls = self.get_llm_class(config.backend)
+ self.instance_dict[RegistryEnum.LLM][name] = llm_cls(model_name=config.model_name, **config.model_extra)
+
+ def _initialize_embedding_model(self):
+ """Initialize all configured embedding model instances.
+
+ For each embedding model configuration:
+ - Retrieves the corresponding registered embedding model class by backend name
+ - Instantiates it with model_name and additional configuration
+ - Stores the instance in instance_dict for later retrieval
+ """
+ for name, config in self.service_config.embedding_model.items():
+ embedding_model_cls = self.get_embedding_model_class(config.backend)
+ self.instance_dict[RegistryEnum.EMBEDDING_MODEL][name] = embedding_model_cls(
+ model_name=config.model_name,
+ **config.model_extra,
+ )
+
+ def _initialize_token_counter(self):
+ """Initialize all configured token counter instances.
+
+ For each token counter configuration:
+ - Retrieves the corresponding registered token counter class by backend name
+ - Instantiates it with model_name and additional configuration
+ - Stores the instance in instance_dict for later retrieval
+ """
+ for name, config in self.service_config.token_counter.items():
+ token_counter_cls = self.get_token_counter_class(config.backend)
+ self.instance_dict[RegistryEnum.TOKEN_COUNTER][name] = token_counter_cls(
+ model_name=config.model_name,
+ **config.model_extra,
+ )
+
+ def _initialize_vector_store(self):
+ """Initialize all configured vector stores with their embedding models.
+
+ For each vector store configuration:
+ - Retrieves the corresponding registered vector store class by backend name
+ - Retrieves the associated embedding model instance by name
+ - Instantiates the vector store with collection name, embedding model, and extra config
+ - Stores the instance in instance_dict for later retrieval
+
+ Note: This must be called after _initialize_embedding_model() since vector stores
+ depend on embedding model instances.
+ """
+ for name, config in self.service_config.vector_store.items():
+ vector_store_cls = self.get_vector_store_class(config.backend)
+ self.instance_dict[RegistryEnum.VECTOR_STORE][name] = vector_store_cls(
+ collection_name=config.collection_name,
+ embedding_model=self.instance_dict[RegistryEnum.EMBEDDING_MODEL][config.embedding_model],
+ **config.model_extra,
+ )
+
+ def _filter_flows(self, name: str) -> bool:
+ """Filter flows based on enabled_flows and disabled_flows configuration.
+
+ The filtering logic follows this priority:
+ 1. If enabled_flows is set: only flows in the list are loaded
+ 2. Else if disabled_flows is set: all flows except those in the list are loaded
+ 3. Otherwise: all flows are loaded
+
+ Args:
+ name: The flow name to check
+
+ Returns:
+ True if the flow should be loaded, False otherwise
+ """
+ if self.service_config.enabled_flows:
+ return name in self.service_config.enabled_flows
+ elif self.service_config.disabled_flows:
+ return name not in self.service_config.disabled_flows
+ else:
+ return True
+
+ def _initialize_flow(self):
+ """Initialize all flows from both registry and configuration.
+
+ Flows can be defined in two ways:
+ 1. Registered flows: Python classes decorated with @register_flow
+ 2. Configuration flows: Defined in config as ExpressionFlow instances
+
+ Process:
+ 1. First, instantiate all registered flow classes (from decorators)
+ - Filter based on enabled_flows/disabled_flows
+ - Create instance with the flow name
+
+ 2. Then, instantiate all configured flows (from config file)
+ - Filter based on enabled_flows/disabled_flows
+ - Create ExpressionFlow instances with flow configuration
+
+ Note: Configuration flows can override registered flows with the same name.
+ """
+
+ # Initialize flows from registry (decorator-based registration)
+ for name, flow_cls in self.registry_dict[RegistryEnum.FLOW].items():
+ if not self._filter_flows(name):
+ continue
+ flow: "BaseFlow" = flow_cls(name=name)
+ self.instance_dict[RegistryEnum.FLOW][flow.name] = flow
+
+ # Initialize flows from configuration (config-based definition)
+ from ..flow import ExpressionFlow
+
+ for name, flow_config in self.service_config.flow.items():
+ if not self._filter_flows(name):
+ continue
+ flow_config.name = name
+ flow: BaseFlow = ExpressionFlow(flow_config=flow_config)
+ self.instance_dict[RegistryEnum.FLOW][name] = flow
+
+ def _initialize_service(self):
+ """Initialize the service backend instance.
+
+ Creates an instance of the configured service backend (e.g., HTTP, MCP, or CMD service)
+ and stores it in the instance dictionary under the 'default' key.
+ """
+ service_cls = self.get_service_class(self.service_config.backend)
+ self.instance_dict[RegistryEnum.SERVICE]["default"] = service_cls()
+
+ async def prepare_mcp_servers(self):
+ """Prepare and initialize MCP (Model Context Protocol) server connections.
+
+ This method:
+ 1. Checks if MCP servers are configured
+ 2. Creates an MCP client instance
+ 3. For each configured server:
+ - Lists available tool calls from the server
+ - Builds a mapping of tool_name -> ToolCall object
+ - Logs available tools for debugging
+
+ The mcp_server_mapping is structured as:
+ {
+ "server_name": {
+ "tool_name": ToolCall(...),
+ ...
+ },
+ ...
+ }
+
+ This allows the application to discover and use external tools provided by MCP servers.
+ """
+ if not self.service_config.mcp_servers:
+ return
+
+ from ..utils import MCPClient
+
+ mcp_client = MCPClient(config={"mcpServers": self.service_config.mcp_servers})
+ for server_name in self.service_config.mcp_servers.keys():
+ try:
+ # Retrieve all available tool calls from this MCP server
+ tool_calls = await mcp_client.list_tool_calls(server_name=server_name, return_dict=False)
+
+ # Build mapping: tool_name -> ToolCall for quick lookup
+ self.mcp_server_mapping[server_name] = {tool_call.name: tool_call for tool_call in tool_calls}
+
+ # Log discovered tools for debugging
+ for tool_call in tool_calls:
+ logger.info(f"list_tool_calls: {server_name}@{tool_call.name} {tool_call.simple_input_dump()}")
+
+ except Exception as e:
+ logger.exception(f"list_tool_calls: {server_name} error: {e}")
+
+ def print_logo(self):
+ """Print the ReMe logo if enabled in configuration."""
+ if self.service_config.enable_logo:
+ print_logo(service_config=self.service_config)
+
+ async def close(self):
+ """Close all service components asynchronously.
+
+ Gracefully closes all instantiated components in order:
+ 1. Vector stores (closes database connections)
+ 2. LLMs (closes API clients and connections)
+ 3. Embedding models (closes API clients and connections)
+
+ This method should be called when shutting down the application
+ to ensure all resources are properly released.
+ """
+ for _, vector_store in self.instance_dict[RegistryEnum.VECTOR_STORE].items():
+ await vector_store.close()
+
+ for _, llm in self.instance_dict[RegistryEnum.LLM].items():
+ await llm.close()
+
+ for _, embedding_model in self.instance_dict[RegistryEnum.EMBEDDING_MODEL].items():
+ await embedding_model.close()
+
+ def close_sync(self):
+ """Close all service components synchronously.
+
+ Synchronous version of close() for non-async contexts.
+ Closes LLMs and embedding models without using async/await.
+
+ Note: Vector stores are not closed here as they typically require async operations.
+ """
+ for _, llm in self.instance_dict[RegistryEnum.LLM].items():
+ llm.close_sync()
+
+ for _, embedding_model in self.instance_dict[RegistryEnum.EMBEDDING_MODEL].items():
+ embedding_model.close_sync()
+
+ def shutdown_thread_pool(self, wait: bool = True):
+ """Shutdown the thread pool executor.
+
+ Args:
+ wait: If True, blocks until all pending futures are executed.
+ If False, returns immediately and pending futures may be cancelled.
+ """
+ if self.thread_pool:
+ self.thread_pool.shutdown(wait=wait)
+
+ def shutdown_ray(self, wait: bool = True):
+ """Shutdown Ray cluster if it was initialized.
+
+ Args:
+ wait: If True, waits for Ray to fully shutdown.
+ If False, returns immediately without waiting.
+
+ Note: Only shuts down Ray if it was configured with ray_max_workers > 1.
+ """
+ if self.service_config and self.service_config.ray_max_workers > 1:
+ import ray
+
+ ray.shutdown(_exiting_interpreter=not wait)
-# Export a global instance for easy access across the application
+# Export a global singleton instance for easy access across the application
+# This is the primary way to access the service context throughout the codebase
C = ServiceContext()
diff --git a/reme_ai/core/embedding/__init__.py b/reme_ai/core/embedding/__init__.py
new file mode 100644
index 00000000..c1d92375
--- /dev/null
+++ b/reme_ai/core/embedding/__init__.py
@@ -0,0 +1,11 @@
+"""embedding"""
+
+from .base_embedding_model import BaseEmbeddingModel
+from .openai_embedding_model import OpenAIEmbeddingModel
+from .openai_embedding_model_sync import OpenAIEmbeddingModelSync
+
+__all__ = [
+ "BaseEmbeddingModel",
+ "OpenAIEmbeddingModel",
+ "OpenAIEmbeddingModelSync",
+]
diff --git a/reme_ai/core/embedding/base_embedding_model.py b/reme_ai/core/embedding/base_embedding_model.py
new file mode 100644
index 00000000..3725f96d
--- /dev/null
+++ b/reme_ai/core/embedding/base_embedding_model.py
@@ -0,0 +1,156 @@
+"""Base embedding model interface for ReMe.
+
+Defines the abstract base class and standard API for all embedding model implementations.
+"""
+
+import asyncio
+import time
+from abc import ABC
+
+from loguru import logger
+
+from ..schema import VectorNode
+
+
+class BaseEmbeddingModel(ABC):
+ """Abstract base class for embedding model implementations.
+
+ Provides a standard interface for text-to-vector generation with
+ built-in batching, retry logic, and error handling.
+ """
+
+ def __init__(
+ self,
+ model_name: str = "",
+ dimensions: int = 1024,
+ max_batch_size: int = 10,
+ max_retries: int = 3,
+ raise_exception: bool = True,
+ **kwargs,
+ ):
+ """Initialize model configuration and parameters."""
+ self.model_name = model_name
+ self.dimensions = dimensions
+ self.max_batch_size = max_batch_size
+ self.max_retries = max_retries
+ self.raise_exception = raise_exception
+ self.kwargs = kwargs
+
+ async def _get_embeddings(self, input_text: list[str]) -> list[list[float]]:
+ """Internal async implementation for calling the embedding API with batch input."""
+
+ def _get_embeddings_sync(self, input_text: list[str]) -> list[list[float]]:
+ """Internal synchronous implementation for calling the embedding API with batch input."""
+
+ async def get_embedding(self, input_text: str) -> list[float]:
+ """Async get embedding for a single text with exponential backoff retries."""
+ for i in range(self.max_retries):
+ try:
+ result = await self._get_embeddings([input_text])
+ return result[0]
+ except Exception as e:
+ logger.error(f"Model {self.model_name} failed: {e}")
+ if i == self.max_retries - 1:
+ if self.raise_exception:
+ raise
+ return []
+ await asyncio.sleep(i + 1)
+ return []
+
+ async def get_embeddings(self, input_text: list[str]) -> list[list[float]]:
+ """Async get embeddings with automatic batching and exponential backoff retries."""
+ # Split into batches and process sequentially to respect rate limits
+ results = []
+ for i in range(0, len(input_text), self.max_batch_size):
+ batch = input_text[i : i + self.max_batch_size]
+ # Process each batch with retry logic
+ for retry in range(self.max_retries):
+ try:
+ batch_res = await self._get_embeddings(batch)
+ if batch_res:
+ results.extend(batch_res)
+ break
+ except Exception as e:
+ logger.error(f"Model {self.model_name} batch failed: {e}")
+ if retry == self.max_retries - 1:
+ if self.raise_exception:
+ raise
+ else:
+ await asyncio.sleep(retry + 1)
+ return results
+
+ def get_embedding_sync(self, input_text: str) -> list[float]:
+ """Synchronous get embedding for a single text with retry logic."""
+ for i in range(self.max_retries):
+ try:
+ result = self._get_embeddings_sync([input_text])
+ return result[0]
+ except Exception as exc:
+ logger.error(f"Model {self.model_name} failed: {exc}")
+ if i == self.max_retries - 1:
+ if self.raise_exception:
+ raise
+ return []
+ time.sleep(i + 1)
+ return []
+
+ def get_embeddings_sync(self, input_text: list[str]) -> list[list[float]]:
+ """Synchronous get embeddings with automatic batching and retry logic."""
+ results = []
+ for i in range(0, len(input_text), self.max_batch_size):
+ batch = input_text[i : i + self.max_batch_size]
+ # Process each batch with retry logic
+ for retry in range(self.max_retries):
+ try:
+ batch_res = self._get_embeddings_sync(batch)
+ if batch_res:
+ results.extend(batch_res)
+ break
+ except Exception as exc:
+ logger.error(f"Model {self.model_name} batch failed: {exc}")
+ if retry == self.max_retries - 1:
+ if self.raise_exception:
+ raise
+ else:
+ time.sleep(retry + 1)
+ return results
+
+ async def get_node_embedding(self, node: VectorNode) -> VectorNode:
+ """Async generate and populate vector field for a single VectorNode object."""
+ node.vector = await self.get_embedding(node.content)
+ return node
+
+ async def get_node_embeddings(self, nodes: list[VectorNode]) -> list[VectorNode]:
+ """Async generate and populate vector fields for a batch of VectorNode objects."""
+ contents = [node.content for node in nodes]
+ embeddings: list[list[float]] = await self.get_embeddings(contents)
+
+ if len(embeddings) == len(nodes):
+ for node, vec in zip(nodes, embeddings):
+ node.vector = vec
+ else:
+ logger.warning(f"Mismatch: got {len(embeddings)} vectors for {len(nodes)} nodes")
+ return nodes
+
+ def get_node_embedding_sync(self, node: VectorNode) -> VectorNode:
+ """Synchronously generate and populate vector field for a single VectorNode object."""
+ node.vector = self.get_embedding_sync(node.content)
+ return node
+
+ def get_node_embeddings_sync(self, nodes: list[VectorNode]) -> list[VectorNode]:
+ """Synchronously generate and populate vector fields for a batch of VectorNode objects."""
+ contents = [node.content for node in nodes]
+ embeddings: list[list[float]] = self.get_embeddings_sync(contents)
+
+ if len(embeddings) == len(nodes):
+ for node, vec in zip(nodes, embeddings):
+ node.vector = vec
+ else:
+ logger.warning(f"Mismatch: got {len(embeddings)} vectors for {len(nodes)} nodes")
+ return nodes
+
+ def close_sync(self):
+ """Synchronously release resources and close connections."""
+
+ async def close(self):
+ """Asynchronously release resources and close connections."""
diff --git a/reme_ai/core/embedding/openai_embedding_model.py b/reme_ai/core/embedding/openai_embedding_model.py
new file mode 100644
index 00000000..e10c7e66
--- /dev/null
+++ b/reme_ai/core/embedding/openai_embedding_model.py
@@ -0,0 +1,52 @@
+"""Asynchronous OpenAI-compatible embedding model implementation for ReMe."""
+
+import os
+from typing import Literal
+
+from openai import AsyncOpenAI
+
+from .base_embedding_model import BaseEmbeddingModel
+from ..context import C
+
+
+@C.register_embedding_model("openai")
+class OpenAIEmbeddingModel(BaseEmbeddingModel):
+ """Asynchronous embedding model implementation compatible with OpenAI-style APIs."""
+
+ def __init__(
+ self,
+ api_key: str | None = None,
+ base_url: str | None = None,
+ encoding_format: Literal["float", "base64"] = "float",
+ **kwargs,
+ ):
+ """Initialize the OpenAI async embedding model with API credentials and configuration."""
+ super().__init__(**kwargs)
+ self.api_key: str = api_key or os.getenv("REME_EMBEDDING_API_KEY", "")
+ self.base_url: str = base_url or os.getenv("REME_EMBEDDING_BASE_URL", "")
+ self.encoding_format: Literal["float", "base64"] = encoding_format
+
+ # Create client using factory method
+ self._client = self._create_client()
+
+ def _create_client(self):
+ """Create and return an internal AsyncOpenAI client instance."""
+ return AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
+
+ async def _get_embeddings(self, input_text: list[str]) -> list[list[float]]:
+ """Fetch embeddings from the API for a batch of strings."""
+ completion = await self._client.embeddings.create(
+ model=self.model_name,
+ input=input_text,
+ dimensions=self.dimensions,
+ encoding_format=self.encoding_format,
+ )
+
+ result_emb = [[] for _ in range(len(input_text))]
+ for emb in completion.data:
+ result_emb[emb.index] = emb.embedding
+ return result_emb
+
+ async def close(self):
+ """Close the asynchronous OpenAI client and release network resources."""
+ await self._client.close()
diff --git a/reme_ai/core/embedding/openai_embedding_model_sync.py b/reme_ai/core/embedding/openai_embedding_model_sync.py
new file mode 100644
index 00000000..b777e6d6
--- /dev/null
+++ b/reme_ai/core/embedding/openai_embedding_model_sync.py
@@ -0,0 +1,33 @@
+"""Synchronous OpenAI-compatible embedding model implementation for ReMe."""
+
+from openai import OpenAI
+
+from .openai_embedding_model import OpenAIEmbeddingModel
+from ..context import C
+
+
+@C.register_embedding_model("openai_sync")
+class OpenAIEmbeddingModelSync(OpenAIEmbeddingModel):
+ """Synchronous embedding model implementation that extends the asynchronous OpenAI model."""
+
+ def _create_client(self):
+ """Create and return an internal synchronous OpenAI client instance."""
+ return OpenAI(api_key=self.api_key, base_url=self.base_url)
+
+ def _get_embeddings_sync(self, input_text: list[str]) -> list[list[float]]:
+ """Fetch embeddings synchronously from the API for a batch of strings."""
+ completion = self._client.embeddings.create(
+ model=self.model_name,
+ input=input_text,
+ dimensions=self.dimensions,
+ encoding_format=self.encoding_format,
+ )
+
+ result_emb = [[] for _ in range(len(input_text))]
+ for emb in completion.data:
+ result_emb[emb.index] = emb.embedding
+ return result_emb
+
+ def close_sync(self):
+ """Close the synchronous OpenAI client and release network resources."""
+ self._client.close()
diff --git a/reme_ai/core/enumeration/json_schema_enum.py b/reme_ai/core/enumeration/json_schema_enum.py
index 17b59380..507645f4 100644
--- a/reme_ai/core/enumeration/json_schema_enum.py
+++ b/reme_ai/core/enumeration/json_schema_enum.py
@@ -3,17 +3,16 @@
from enum import Enum
-class JsonSchemaEnum(str, Enum):
+class JsonSchemaEnum(Enum):
"""Enumeration of valid JSON Schema data types."""
- STRING = "string"
- NUMBER = "number"
- INTEGER = "integer"
- OBJECT = "object"
- ARRAY = "array"
- BOOLEAN = "boolean"
- NULL = "null"
+ STRING = str
+ NUMBER = float
+ INTEGER = int
+ OBJECT = dict
+ ARRAY = list
+ BOOLEAN = bool
def __str__(self) -> str:
"""Returns the string representation of the enum value."""
- return self.value
+ return self.name.lower()
diff --git a/reme_ai/core/flow/__init__.py b/reme_ai/core/flow/__init__.py
new file mode 100644
index 00000000..6d5a053b
--- /dev/null
+++ b/reme_ai/core/flow/__init__.py
@@ -0,0 +1,11 @@
+"""flow"""
+
+from .base_flow import BaseFlow
+from .cmd_flow import CmdFlow
+from .expression_flow import ExpressionFlow
+
+__all__ = [
+ "BaseFlow",
+ "CmdFlow",
+ "ExpressionFlow",
+]
diff --git a/reme_ai/core/flow/base_flow.py b/reme_ai/core/flow/base_flow.py
new file mode 100644
index 00000000..1802af78
--- /dev/null
+++ b/reme_ai/core/flow/base_flow.py
@@ -0,0 +1,214 @@
+"""Base flow module providing abstract flow execution with caching and operation orchestration."""
+
+import asyncio
+import hashlib
+import json
+from abc import ABC, abstractmethod
+
+from loguru import logger
+
+from ..context import C, RuntimeContext
+from ..enumeration import ChunkEnum, RegistryEnum
+from ..op import BaseOp, SequentialOp, ParallelOp
+from ..schema import Response, ToolCall, ToolAttr
+from ..utils import camel_to_snake, CacheHandler
+
+
+class BaseFlow(ABC):
+ """Abstract base class for flow execution with caching, streaming, and operation tree management.
+
+ BaseFlow provides a framework for building complex workflows by composing operations
+ into executable trees. It supports both synchronous and asynchronous execution modes,
+ response caching, streaming outputs, and automatic tool call schema generation.
+ """
+
+ def __init__(
+ self,
+ name: str = "",
+ stream: bool = False,
+ raise_exception: bool = True,
+ enable_cache: bool = False,
+ cache_path: str = "cache/flow",
+ cache_expire_hours: float = 0.1,
+ **kwargs,
+ ):
+ """Initialize flow configuration and execution state."""
+ super().__init__()
+
+ self.name: str = name or camel_to_snake(self.__class__.__name__)
+ self.stream: bool = stream
+ self.raise_exception: bool = raise_exception
+ self.enable_cache: bool = enable_cache
+ self.cache_path: str = cache_path
+ self.cache_expire_hours: float = cache_expire_hours
+ self.flow_params: dict = kwargs
+
+ self._flow_op: BaseOp | None = None
+ self._cache: CacheHandler | None = None
+ self._flow_printed: bool = False
+ self._tool_call: ToolCall | None = None
+
+ def _build_tool_call(self) -> ToolCall | None:
+ """Generate the tool call schema definition for this flow."""
+
+ @abstractmethod
+ def _build_flow(self) -> BaseOp:
+ """Construct the root operation tree for flow execution."""
+
+ def _compute_cache_key(self, params: dict) -> str | None:
+ """Generate a SHA256 hash from input parameters for caching."""
+ try:
+ payload = json.dumps(params, sort_keys=True, ensure_ascii=False, default=str)
+ return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+ except Exception as e:
+ logger.exception(f"{self.name} cache key serialization failed: {e}")
+ return None
+
+ def _maybe_load_cached(self, params: dict) -> Response | None:
+ """Retrieve a cached response if caching is enabled and available."""
+ if not self.enable_cache or self.stream:
+ return None
+
+ if key := self._compute_cache_key(params):
+ if cached := self.cache.load(key):
+ logger.info(f"Loaded {self.name} response from cache.")
+ return Response(**cached)
+ return None
+
+ def _maybe_save_cache(self, params: dict, response: Response):
+ """Persist the execution response to the cache."""
+ if not self.enable_cache or self.stream:
+ return
+
+ if key := self._compute_cache_key(params):
+ self.cache.save(
+ key,
+ response.model_dump(exclude_none=True),
+ expire_hours=self.cache_expire_hours,
+ )
+
+ def _print_operation_tree(self, name: str, op: BaseOp, indent: int):
+ """Recursively log the hierarchy of the flow's operation tree."""
+ prefix = " " * indent
+ op_type = "sequential" if isinstance(op, SequentialOp) else "parallel" if isinstance(op, ParallelOp) else name
+ logger.info(f"{prefix}{op_type} execution")
+
+ for sub_op in op.sub_ops or []:
+ self._print_operation_tree(sub_op.name, sub_op, indent + 2)
+
+ @property
+ def tool_call(self) -> ToolCall | None:
+ """Lazily construct the ToolCall schema describing this flow."""
+ if self.flow_op.tool_call:
+ return self.flow_op.tool_call
+
+ if self._tool_call is None:
+ self._tool_call = self._build_tool_call()
+ if self._tool_call:
+ self._tool_call.name = self._tool_call.name or self.name
+ self._tool_call.output = self._tool_call.output or {
+ f"{self.name}_result": ToolAttr(
+ type="string",
+ description=f"The execution result of the {self.name}",
+ ),
+ }
+ return self._tool_call
+
+ @property
+ def cache(self) -> CacheHandler:
+ """Provide access to the internal CacheHandler instance."""
+ assert self.enable_cache, "Cache usage requested while disabled."
+ if self._cache is None:
+ self._cache = CacheHandler(f"{self.cache_path}/{self.name}")
+ return self._cache
+
+ @property
+ def flow_op(self) -> BaseOp:
+ """Lazily build and retrieve the root operation of the flow."""
+ if self._flow_op is None:
+ self._flow_op = self._build_flow()
+ return self._flow_op
+
+ @property
+ def async_mode(self) -> bool:
+ """Check if the current flow operation tree is asynchronous."""
+ return self.flow_op.async_mode
+
+ @staticmethod
+ def parse_expression(expression: str) -> BaseOp:
+ """Parse a string expression into an executable BaseOp instance."""
+ lines = [x.strip() for x in expression.strip().splitlines() if x.strip()]
+ if not lines:
+ raise ValueError("Expression is empty")
+
+ env: dict = C.registry_dict[RegistryEnum.OP]
+ if len(lines) > 1:
+ exec("\n".join(lines[:-1]), {"__builtins__": {}}, env)
+
+ result = eval(lines[-1], {"__builtins__": {}}, env)
+ if not isinstance(result, BaseOp):
+ raise TypeError(f"Expression evaluated to {type(result)}, expected BaseOp")
+ return result
+
+ def print_flow(self):
+ """Log the visual structure of the flow once."""
+ if not self._flow_printed:
+ logger.info(f"---------- [Flow Structure] {self.name} ----------")
+ self._print_operation_tree(self.name, self.flow_op, 0)
+ logger.info("-" * 50)
+ self._flow_printed = True
+
+ async def call(self, **kwargs) -> Response | asyncio.Queue:
+ """Execute the flow asynchronously with parameter caching."""
+ kwargs["stream"] = self.stream
+ logger.info(f"{self.name} incoming params: {kwargs}")
+ if cached := self._maybe_load_cached(kwargs):
+ return cached
+
+ context = RuntimeContext(**kwargs)
+ try:
+ self.print_flow()
+ flow_op: BaseOp = self._build_flow()
+ assert self.flow_op.async_mode, "Async call requires an async flow operation."
+
+ await flow_op.call(context=context)
+ result = context.stream_queue if self.stream else context.response
+
+ if self.stream:
+ await context.add_stream_done()
+
+ self._maybe_save_cache(kwargs, result)
+ return result
+ except Exception as e:
+ logger.exception(f"{self.name} async call failed: {e}")
+ if self.raise_exception:
+ raise e
+ if self.stream:
+ await context.add_stream_chunk_and_type(str(e), ChunkEnum.ERROR)
+ await context.add_stream_done()
+ return context.stream_queue
+ context.add_response_error(e)
+ return context.response
+
+ def call_sync(self, **kwargs) -> Response:
+ """Execute the flow synchronously with parameter caching."""
+ logger.info(f"{self.name} incoming sync params: {kwargs}")
+ assert not self.stream, "Synchronous call cannot be used in stream mode."
+ if cached := self._maybe_load_cached(kwargs):
+ return cached
+
+ context = RuntimeContext(**kwargs)
+ try:
+ self.print_flow()
+ flow_op: BaseOp = self._build_flow()
+ assert not self.flow_op.async_mode, "Sync call requires a sync flow operation."
+
+ flow_op.call_sync(context=context)
+ self._maybe_save_cache(kwargs, context.response)
+ return context.response
+ except Exception as e:
+ logger.exception(f"{self.name} sync call failed: {e}")
+ if self.raise_exception:
+ raise e
+ context.add_response_error(e)
+ return context.response
diff --git a/reme_ai/core/flow/cmd_flow.py b/reme_ai/core/flow/cmd_flow.py
new file mode 100644
index 00000000..2fe3b97a
--- /dev/null
+++ b/reme_ai/core/flow/cmd_flow.py
@@ -0,0 +1,18 @@
+"""Command-based flow implementation for parsing and executing operation sequences."""
+
+from .base_flow import BaseFlow
+from ..op import BaseOp
+
+
+class CmdFlow(BaseFlow):
+ """A flow class that builds an operation chain from a string expression."""
+
+ def __init__(self, flow: str = "", **kwargs):
+ """Initialize the command flow with a string-based operation definition."""
+ super().__init__(**kwargs)
+ self.flow = flow
+ assert flow, "add `flow=` in cmd!"
+
+ def _build_flow(self) -> BaseOp:
+ """Parse the stored flow expression into a functional operation object."""
+ return self.parse_expression(self.flow)
diff --git a/reme_ai/core/flow/expression_flow.py b/reme_ai/core/flow/expression_flow.py
new file mode 100644
index 00000000..5ac8f0d8
--- /dev/null
+++ b/reme_ai/core/flow/expression_flow.py
@@ -0,0 +1,30 @@
+"""Expression-based flow implementation driven by configuration objects."""
+
+from .base_flow import BaseFlow
+from ..op import BaseOp
+from ..schema import FlowConfig, ToolCall
+
+
+class ExpressionFlow(BaseFlow):
+ """A flow implementation that constructs operations from a FlowConfig definition."""
+
+ def __init__(self, flow_config: FlowConfig):
+ """Initialize the flow using settings and metadata from a FlowConfig instance."""
+ self.flow_config: FlowConfig = flow_config
+ super().__init__(
+ name=flow_config.name,
+ stream=self.flow_config.stream,
+ raise_exception=self.flow_config.raise_exception,
+ enable_cache=self.flow_config.enable_cache,
+ cache_path=self.flow_config.cache_path,
+ cache_expire_hours=self.flow_config.cache_expire_hours,
+ **flow_config.model_extra,
+ )
+
+ def _build_flow(self) -> BaseOp:
+ """Generate the operation chain by parsing the flow content string."""
+ return self.parse_expression(self.flow_config.flow_content)
+
+ def _build_tool_call(self) -> ToolCall:
+ """Construct a tool call representation based on configuration parameters."""
+ return ToolCall(**{"description": self.flow_config.description, "parameters": self.flow_config.parameters})
diff --git a/reme_ai/core/llm/base_llm.py b/reme_ai/core/llm/base_llm.py
index 6f1da121..a7ba0f76 100644
--- a/reme_ai/core/llm/base_llm.py
+++ b/reme_ai/core/llm/base_llm.py
@@ -4,7 +4,7 @@ import asyncio
import json
import time
from abc import ABC
-from typing import List, Callable, Generator, AsyncGenerator, Any, Optional, Dict
+from typing import Callable, Generator, AsyncGenerator, Optional, Any
from loguru import logger
@@ -74,7 +74,7 @@ class BaseLLM(ABC):
@staticmethod
def _accumulate_tool_call_chunk(
tool_call,
- ret_tools: List[ToolCall],
+ ret_tools: list[ToolCall],
) -> None:
"""Assemble incremental tool call fragments into complete ToolCall objects."""
index = tool_call.index
@@ -95,15 +95,15 @@ class BaseLLM(ABC):
@staticmethod
def _validate_and_serialize_tools(
- ret_tools: List[ToolCall],
- tools: Optional[List[ToolCall]],
- ) -> List[Dict]:
+ ret_tools: list[ToolCall],
+ tools: Optional[list[ToolCall]],
+ ) -> list[dict]:
"""Validate tool call integrity and return serialized tool dictionaries."""
if not ret_tools:
return []
# Create lookup dict for tool validation
- tool_dict: Dict[str, ToolCall] = {x.name: x for x in tools} if tools else {}
+ tool_dict: dict[str, ToolCall] = {x.name: x for x in tools} if tools else {}
validated_tools = []
for tool in ret_tools:
@@ -123,8 +123,8 @@ class BaseLLM(ABC):
def _build_stream_kwargs(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
log_params: bool = True,
**kwargs,
) -> dict:
@@ -133,8 +133,8 @@ class BaseLLM(ABC):
async def _stream_chat(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
stream_kwargs: Optional[dict] = None,
) -> AsyncGenerator[StreamChunk, None]:
"""Internal async generator for streaming raw response chunks."""
@@ -142,8 +142,8 @@ class BaseLLM(ABC):
def _stream_chat_sync(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
stream_kwargs: Optional[dict] = None,
) -> Generator[StreamChunk, None, None]:
"""Internal synchronous generator for streaming raw response chunks."""
@@ -152,8 +152,8 @@ class BaseLLM(ABC):
async def _stream_with_retry(
self,
operation_name: str,
- messages: List[Message],
- tools: Optional[List[ToolCall]],
+ messages: list[Message],
+ tools: Optional[list[ToolCall]],
stream_kwargs: dict,
) -> AsyncGenerator[StreamChunk, None]:
"""Execute the async streaming operation with retry logic and error recovery."""
@@ -178,8 +178,8 @@ class BaseLLM(ABC):
def _stream_with_retry_sync(
self,
operation_name: str,
- messages: List[Message],
- tools: Optional[List[ToolCall]],
+ messages: list[Message],
+ tools: Optional[list[ToolCall]],
stream_kwargs: dict,
) -> Generator[StreamChunk, None, None]:
"""Execute the synchronous streaming operation with retry logic and error recovery."""
@@ -202,8 +202,8 @@ class BaseLLM(ABC):
async def stream_chat(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
**kwargs,
) -> AsyncGenerator[StreamChunk, None]:
"""Public async interface for streaming chat completions with retries."""
@@ -213,8 +213,8 @@ class BaseLLM(ABC):
def stream_chat_sync(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
**kwargs,
) -> Generator[StreamChunk, None, None]:
"""Public synchronous interface for streaming chat completions with retries."""
@@ -223,8 +223,8 @@ class BaseLLM(ABC):
async def _chat(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
enable_stream_print: bool = False,
**kwargs,
) -> Message:
@@ -245,8 +245,8 @@ class BaseLLM(ABC):
def _chat_sync(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
enable_stream_print: bool = False,
**kwargs,
) -> Message:
@@ -315,8 +315,8 @@ class BaseLLM(ABC):
async def chat(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
enable_stream_print: bool = False,
callback_fn: Optional[Callable[[Message], Any]] = None,
default_value: Any = None,
@@ -337,8 +337,8 @@ class BaseLLM(ABC):
def chat_sync(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
enable_stream_print: bool = False,
callback_fn: Optional[Callable[[Message], Any]] = None,
default_value: Any = None,
diff --git a/reme_ai/core/llm/lite_llm_sync.py b/reme_ai/core/llm/lite_llm_sync.py
index 5f814a2e..8015a275 100644
--- a/reme_ai/core/llm/lite_llm_sync.py
+++ b/reme_ai/core/llm/lite_llm_sync.py
@@ -1,11 +1,6 @@
-"""Synchronous LiteLLM-based LLM implementation for the ReMe framework.
+"""Synchronous LiteLLM-based LLM implementation for the ReMe framework."""
-This module provides a unified synchronous interface for 100+ LLM providers via LiteLLM,
-supporting streaming completions, tool calling, and reasoning content. For
-asynchronous operations, refer to the LiteLLM class in the lite_llm module.
-"""
-
-from typing import List, Generator, Optional
+from typing import Generator
import litellm
@@ -19,54 +14,21 @@ from ..schema import ToolCall
@C.register_llm("litellm_sync")
class LiteLLMSync(LiteLLM):
- """
- Synchronous LiteLLM client for executing chat completions and streaming responses.
-
- This class extends the base LiteLLM implementation to provide synchronous
- execution of streaming methods, inheriting initialization and configuration
- logic from the parent class.
-
- Example:
- >>> llm = LiteLLMSync(
- ... model_name="qwen3-max",
- ... api_key="sk-...",
- ... temperature=0.7
- ... )
- >>> messages = [Message(role=Role.USER, content="Hello!")]
- >>> for chunk in llm.chat(messages):
- ... print(chunk)
- """
+ """Synchronous LiteLLM client for executing chat completions and streaming responses."""
def _stream_chat_sync(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
- stream_kwargs: Optional[dict] = None,
+ messages: list[Message],
+ tools: list[ToolCall] | None = None,
+ stream_kwargs: dict | None = None,
) -> Generator[StreamChunk, None, None]:
- """
- Internal synchronous generator for processing streaming chat completion chunks.
-
- This method orchestrates the LiteLLM completion lifecycle by categorizing
- raw API chunks into usage data, reasoning content (thinking), regular
- text responses, and aggregated tool calls.
-
- Args:
- messages: List of conversation messages to send to the model.
- tools: Optional list of tool definitions available for the model to call.
- stream_kwargs: Dictionary of pre-built parameters for the LiteLLM API.
-
- Yields:
- StreamChunk: Wrapped response fragments categorized by ChunkEnum.
-
- Raises:
- ValueError: If tool call arguments fail validation or serialization.
- """
+ """Internal synchronous generator for processing streaming chat completion chunks."""
# Create streaming completion request using LiteLLM
stream_kwargs = stream_kwargs or {}
completion = litellm.completion(**stream_kwargs)
# Track accumulated tool calls across chunks
- ret_tools: List[ToolCall] = []
+ ret_tools: list[ToolCall] = []
# Flag to track if we've started receiving answer content
is_answering: bool = False
diff --git a/reme_ai/core/llm/openai_llm.py b/reme_ai/core/llm/openai_llm.py
index 2c593d6f..df368021 100644
--- a/reme_ai/core/llm/openai_llm.py
+++ b/reme_ai/core/llm/openai_llm.py
@@ -1,7 +1,7 @@
"""Asynchronous OpenAI-compatible LLM implementation supporting streaming, tool calls, and reasoning content."""
import os
-from typing import List, AsyncGenerator, Optional
+from typing import AsyncGenerator, Optional
from loguru import logger
from openai import AsyncOpenAI
@@ -38,8 +38,8 @@ class OpenAILLM(BaseLLM):
def _build_stream_kwargs(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
log_params: bool = True,
**kwargs,
) -> dict:
@@ -68,8 +68,8 @@ class OpenAILLM(BaseLLM):
async def _stream_chat(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
stream_kwargs: Optional[dict] = None,
) -> AsyncGenerator[StreamChunk, None]:
"""Generate a stream of chat completion chunks including text, reasoning content, and tool calls."""
@@ -78,7 +78,7 @@ class OpenAILLM(BaseLLM):
completion = await self._client.chat.completions.create(**stream_kwargs)
# Track accumulated tool calls across chunks
- ret_tools: List[ToolCall] = []
+ ret_tools: list[ToolCall] = []
# Flag to track if we've started receiving answer content
is_answering: bool = False
diff --git a/reme_ai/core/llm/openai_llm_sync.py b/reme_ai/core/llm/openai_llm_sync.py
index 3aadab1a..408ce272 100644
--- a/reme_ai/core/llm/openai_llm_sync.py
+++ b/reme_ai/core/llm/openai_llm_sync.py
@@ -1,6 +1,6 @@
"""Synchronous OpenAI-compatible LLM implementation supporting streaming, tool calls, and reasoning content."""
-from typing import List, Generator, Optional
+from typing import Generator, Optional
from openai import OpenAI
@@ -22,8 +22,8 @@ class OpenAILLMSync(OpenAILLM):
def _stream_chat_sync(
self,
- messages: List[Message],
- tools: Optional[List[ToolCall]] = None,
+ messages: list[Message],
+ tools: Optional[list[ToolCall]] = None,
stream_kwargs: Optional[dict] = None,
) -> Generator[StreamChunk, None, None]:
"""Synchronously generate a stream of chat completion chunks including text, reasoning, and tool calls."""
@@ -32,7 +32,7 @@ class OpenAILLMSync(OpenAILLM):
completion = self._client.chat.completions.create(**stream_kwargs)
# Track accumulated tool calls across chunks
- ret_tools: List[ToolCall] = []
+ ret_tools: list[ToolCall] = []
# Flag to track if we've started receiving answer content
is_answering: bool = False
diff --git a/reme_ai/core/op/__init__.py b/reme_ai/core/op/__init__.py
new file mode 100644
index 00000000..c009e89f
--- /dev/null
+++ b/reme_ai/core/op/__init__.py
@@ -0,0 +1,13 @@
+"""op"""
+
+from .base_op import BaseOp
+from .base_ray_op import BaseRayOp
+from .parallel_op import ParallelOp
+from .sequential_op import SequentialOp
+
+__all__ = [
+ "BaseOp",
+ "BaseRayOp",
+ "ParallelOp",
+ "SequentialOp",
+]
diff --git a/reme_ai/core/op/base_op.py b/reme_ai/core/op/base_op.py
new file mode 100644
index 00000000..7eeb5fbd
--- /dev/null
+++ b/reme_ai/core/op/base_op.py
@@ -0,0 +1,363 @@
+"""Base operator class for LLM workflow execution and composition."""
+
+import asyncio
+import copy
+import inspect
+from pathlib import Path
+from typing import Callable, Any, Optional
+
+from loguru import logger
+from tqdm import tqdm
+
+from ..context import RuntimeContext, PromptHandler, C
+from ..embedding import BaseEmbeddingModel
+from ..llm import BaseLLM
+from ..schema import ToolCall, ToolAttr, Response
+from ..token_counter import BaseTokenCounter
+from ..utils import camel_to_snake, CacheHandler, timer
+from ..vector_store import BaseVectorStore
+
+
+class BaseOp:
+ """Base operator class for LLM workflow execution and composition."""
+
+ def __new__(cls, *args, **kwargs):
+ """Capture initialization arguments for object cloning."""
+ instance = super().__new__(cls)
+ instance._init_args = copy.copy(args)
+ instance._init_kwargs = copy.copy(kwargs)
+ return instance
+
+ def __init__(
+ self,
+ name: str = "",
+ async_mode: bool = True,
+ language: str = "",
+ prompt_name: str = "",
+ llm: str | BaseLLM = "default",
+ embedding_model: str | BaseEmbeddingModel = "default",
+ vector_store: str | BaseVectorStore = "default",
+ token_counter: str | BaseTokenCounter = "default",
+ enable_cache: bool = False,
+ cache_path: str = "cache/op",
+ sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"] = None,
+ input_mapping: dict[str, str] | None = None,
+ output_mapping: dict[str, str] | None = None,
+ save_response_result: bool = False,
+ enable_sync_thread_pool: bool = True,
+ max_retries: int = 1,
+ raise_exception: bool = False,
+ **kwargs,
+ ):
+ """Initialize operator configurations and internal state."""
+ self.name = name or camel_to_snake(self.__class__.__name__)
+ self.async_mode = async_mode
+ self.language = language or C.language
+ self.prompt = self._get_prompt_handler(prompt_name)
+
+ self._llm = llm
+ self._embedding_model = embedding_model
+ self._vector_store = vector_store
+ self._token_counter = token_counter
+
+ self.enable_cache = enable_cache
+ self.cache_path = cache_path
+ self.sub_ops: list[BaseOp] = []
+ self.add_sub_ops(sub_ops)
+
+ self.input_mapping = input_mapping
+ self.output_mapping = output_mapping
+ self.save_response_result = save_response_result
+ self.enable_sync_thread_pool = enable_sync_thread_pool
+ self.max_retries = max(1, max_retries)
+ self.raise_exception = raise_exception
+ self.op_params = kwargs
+
+ self._pending_tasks: list = []
+ self.context: RuntimeContext | None = None
+ self._cache: CacheHandler | None = None
+ self._tool_call: ToolCall | None = None
+
+ def _get_prompt_handler(self, prompt_name: str) -> PromptHandler:
+ """Load prompt configuration from the associated YAML file."""
+ path = Path(inspect.getfile(self.__class__))
+ path = path.with_stem(prompt_name) if prompt_name else path
+ return PromptHandler(language=self.language).load_prompt_by_file(path.with_suffix(".yaml"))
+
+ def _build_tool_call(self) -> ToolCall | None:
+ """Build and return the tool call schema; override in subclasses."""
+
+ def _validate_inputs(self):
+ """Ensure all required tool inputs are present in context."""
+ if self.tool_call is not None:
+ parameters = self.tool_call.parameters
+ if parameters.type == "object" and parameters.properties:
+ required_list = parameters.required or []
+ required_keys = {k: (k in required_list) for k in parameters.properties.keys()}
+ self.context.validate_required_keys(required_keys, self.name)
+
+ def _handle_failure(self, e: Exception, attempt: int):
+ """Log failures and handle final retry logic."""
+ message = f"{self.name} failed (attempt {attempt + 1}): {e}"
+ if attempt == self.max_retries - 1:
+ logger.exception(message)
+ if self.raise_exception:
+ raise e
+
+ if self.tool_call is not None:
+ self.output = f"{self.name} failed: {e}"
+ else:
+ logger.warning(message)
+
+ @property
+ def tool_call(self) -> ToolCall | None:
+ """Lazily construct and return the tool call metadata."""
+ if self._tool_call is None:
+ self._tool_call = self._build_tool_call()
+ if self._tool_call is None:
+ return None
+
+ self._tool_call.name = self._tool_call.name or self.name
+ if not self._tool_call.output.properties:
+ self._tool_call.output.properties = {
+ f"{self.name}_result": ToolAttr(type="string", description=f"Execution result of {self.name}"),
+ }
+ return self._tool_call
+
+ def set_tool_call(self, tool_call: ToolCall | dict):
+ """Set the tool call."""
+ if isinstance(tool_call, dict):
+ self._tool_call = ToolCall(**tool_call)
+ elif isinstance(tool_call, ToolCall):
+ self._tool_call = tool_call
+ else:
+ raise ValueError(f"Invalid tool call: {tool_call}")
+
+ self._tool_call.name = self._tool_call.name or self.name
+ if not self._tool_call.output.properties:
+ self._tool_call.output.properties = {
+ f"{self.name}_result": ToolAttr(type="string", description=f"Execution result of {self.name}"),
+ }
+
+ @property
+ def input_dict(self) -> dict:
+ """Extract required and optional inputs from context based on schema."""
+ parameters = self.tool_call.parameters
+ if parameters.type != "object" or not parameters.properties:
+ return {}
+ required_keys = set(parameters.required or [])
+ return {k: self.context[k] for k in parameters.properties.keys() if (k in required_keys or k in self.context)}
+
+ @property
+ def output(self) -> Any:
+ """Get the single output value from context."""
+ output_properties = self.tool_call.output.properties
+ if not output_properties:
+ return None
+
+ keys = list(output_properties.keys())
+ return self.context[keys[0]]
+
+ @output.setter
+ def output(self, value: Any):
+ """Set the single output value into context."""
+ output_properties = self.tool_call.output.properties
+ if not output_properties:
+ return
+
+ keys = list(output_properties.keys())
+ self.context[keys[0]] = value
+
+ @property
+ def cache(self) -> CacheHandler:
+ """Access the operator-specific cache handler."""
+ assert self.enable_cache, "Cache is disabled!"
+ if not self._cache:
+ self._cache = CacheHandler(f"{self.cache_path}/{self.name}")
+ return self._cache
+
+ @property
+ def llm(self) -> BaseLLM:
+ """Get the LLM instance from ServiceContext."""
+ if isinstance(self._llm, str):
+ self._llm = C.get_llm(self._llm)
+ return self._llm
+
+ @property
+ def embedding_model(self) -> BaseEmbeddingModel:
+ """Get the embedding model instance from ServiceContext."""
+ if isinstance(self._embedding_model, str):
+ self._embedding_model = C.get_embedding_model(self._embedding_model)
+ return self._embedding_model
+
+ @property
+ def vector_store(self) -> BaseVectorStore:
+ """Lazily initialize and return the vector store instance."""
+ if isinstance(self._vector_store, str):
+ self._vector_store = C.get_vector_store(self._vector_store)
+ return self._vector_store
+
+ @property
+ def token_counter(self) -> BaseTokenCounter:
+ """Get the token counter instance from ServiceContext."""
+ if isinstance(self._token_counter, str):
+ self._token_counter = C.get_token_counter(self._token_counter)
+ return self._token_counter
+
+ @property
+ def service_metadata(self) -> dict:
+ """Get service configuration metadata."""
+ return C.service_config.model_extra
+
+ @property
+ def response(self) -> Response:
+ """Get the response object."""
+ return self.context.response
+
+ def before_execute_sync(self):
+ """Prepare context and validate before sync execution."""
+ self.context.apply_mapping(self.input_mapping)
+ self._validate_inputs()
+
+ def execute_sync(self):
+ """Define core sync logic in subclasses."""
+
+ def after_execute_sync(self):
+ """Finalize context and mappings after sync execution."""
+ self.context.apply_mapping(self.output_mapping)
+ if self.tool_call is not None and self.save_response_result:
+ self.context.response.answer = self.output
+
+ async def before_execute(self):
+ """Prepare context and validate before async execution."""
+ self.before_execute_sync()
+
+ async def execute(self):
+ """Define core async logic in subclasses."""
+
+ async def after_execute(self):
+ """Finalize context and mappings after async execution."""
+ self.after_execute_sync()
+
+ @timer
+ def call_sync(self, context: RuntimeContext = None, **kwargs):
+ """Execute the operator synchronously with retry logic."""
+ self.context = RuntimeContext.from_context(context, **kwargs)
+ for i in range(self.max_retries):
+ try:
+ self.before_execute_sync()
+ self.execute_sync()
+ self.after_execute_sync()
+ break
+ except Exception as e:
+ self._handle_failure(e, i)
+ return self.output if self.tool_call is not None else None
+
+ async def call(self, context: RuntimeContext = None, **kwargs):
+ """Execute the operator asynchronously with retry logic."""
+ self.context = RuntimeContext.from_context(context, **kwargs)
+ for i in range(self.max_retries):
+ try:
+ await self.before_execute()
+ await self.execute()
+ await self.after_execute()
+ break
+ except Exception as e:
+ self._handle_failure(e, i)
+ return self.output if self.tool_call is not None else None
+
+ def submit_sync_task(self, fn: Callable, *args, **kwargs) -> "BaseOp":
+ """Submit a task to the thread pool or local queue."""
+ task = C.thread_pool.submit(fn, *args, **kwargs) if self.enable_sync_thread_pool else (fn, args, kwargs)
+ self._pending_tasks.append(task)
+ return self
+
+ def submit_async_task(self, coro_fn: Callable, *args, **kwargs) -> "BaseOp":
+ """Submit an async task to the pending tasks queue."""
+ task = coro_fn(*args, **kwargs)
+ self._pending_tasks.append(task)
+ return self
+
+ def join_sync_tasks(self, task_desc: str = None) -> list:
+ """Wait for all pending sync tasks and return flattened results."""
+ results = []
+ for task in tqdm(self._pending_tasks, desc=task_desc or self.name):
+ res = task.result() if self.enable_sync_thread_pool else task[0](*task[1], **task[2])
+ if res:
+ results.extend(res if isinstance(res, list) else [res])
+ self._pending_tasks.clear()
+ return results
+
+ async def join_async_tasks(self, return_exceptions: bool = True) -> list:
+ """Wait for all pending async tasks and aggregate results."""
+ try:
+ raw_results = await asyncio.gather(*self._pending_tasks, return_exceptions=return_exceptions)
+ results = []
+ for res in raw_results:
+ if isinstance(res, Exception):
+ logger.error(f"Async task failed: {res}")
+ continue
+ if res:
+ results.extend(res if isinstance(res, list) else [res])
+ return results
+ finally:
+ self._pending_tasks.clear()
+
+ def add_sub_ops(self, sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"]):
+ """Add child operators to this operator's sub_ops."""
+ if not sub_ops:
+ return
+
+ if isinstance(sub_ops, dict):
+ for name, op in sub_ops.items():
+ assert self.async_mode == op.async_mode, "Async mode mismatch!"
+ op.name = name
+ self.sub_ops.append(op)
+ elif isinstance(sub_ops, list):
+ for op in sub_ops:
+ assert self.async_mode == op.async_mode, "Async mode mismatch!"
+ self.sub_ops.append(op)
+ else:
+ assert self.async_mode == sub_ops.async_mode, "Async mode mismatch!"
+ self.sub_ops.append(sub_ops)
+
+ def add_sub_op(self, sub_op: "BaseOp"):
+ """Add a single child operator to this operator's sub_ops."""
+ self.sub_ops.append(sub_op)
+
+ def __lshift__(self, ops):
+ """Operator overload for adding sub-operators."""
+ self.add_sub_ops(ops)
+ return self
+
+ def __rshift__(self, op: "BaseOp"):
+ """Operator overload for sequential execution composition."""
+ from .sequential_op import SequentialOp
+
+ seq = SequentialOp(sub_ops=[self], async_mode=self.async_mode)
+ seq.add_sub_ops(op.sub_ops if isinstance(op, SequentialOp) else op)
+ return seq
+
+ def __or__(self, op: "BaseOp"):
+ """Operator overload for parallel execution composition."""
+ from .parallel_op import ParallelOp
+
+ par = ParallelOp(sub_ops=[self], async_mode=self.async_mode)
+ par.add_sub_ops(op.sub_ops if isinstance(op, ParallelOp) else op)
+ return par
+
+ def prompt_format(self, prompt_name: str, **kwargs) -> str:
+ """Format a prompt template with provided keyword arguments."""
+ return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs)
+
+ def get_prompt(self, prompt_name: str) -> str:
+ """Get a prompt template by name."""
+ return self.prompt.get_prompt(prompt_name=prompt_name)
+
+ def copy(self, **kwargs):
+ """Create a copy of this operator with optional parameter overrides."""
+ copy_op = self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs})
+ if self.sub_ops:
+ copy_op.sub_ops.clear()
+ copy_op.add_sub_ops(self.sub_ops)
+ return copy_op
diff --git a/reme_ai/core/op/base_ray_op.py b/reme_ai/core/op/base_ray_op.py
new file mode 100644
index 00000000..88cc2985
--- /dev/null
+++ b/reme_ai/core/op/base_ray_op.py
@@ -0,0 +1,124 @@
+"""Base class for Ray-based parallel operations."""
+
+from abc import ABCMeta
+from typing import Callable
+
+import pandas as pd
+from loguru import logger
+from tqdm import tqdm
+
+from .base_op import BaseOp
+from ..context import BaseContext, C
+
+_RAY_IMPORT_ERROR = None
+
+try:
+ import ray
+except ImportError as e:
+ _RAY_IMPORT_ERROR = e
+ ray = None
+
+
+class BaseRayOp(BaseOp, metaclass=ABCMeta):
+ """Base class for Ray-based parallel operations."""
+
+ def __init__(self, **kwargs):
+ if _RAY_IMPORT_ERROR:
+ raise ImportError("Ray requires extra dependencies. Install with `pip install ray`")
+
+ super().__init__(**kwargs)
+ self._ray_task_list: list = []
+
+ def submit_and_join_parallel_op(self, op: BaseOp, **kwargs) -> list:
+ """Submit a BaseOp to be executed in parallel via Ray."""
+ return self.submit_and_join_ray_task(fn=op.call, task_desc=op.name, context=self.context, **kwargs)
+
+ def submit_and_join_ray_task(self, fn: Callable, parallel_key: str = "", task_desc: str = "", **kwargs) -> list:
+ """Divide data into chunks and execute them across Ray workers."""
+ max_workers = C.service_config.ray_max_workers
+ self._ray_task_list.clear()
+
+ # Automatically detect the key containing the list to parallelize
+ if not parallel_key:
+ for key, value in kwargs.items():
+ if isinstance(value, list):
+ parallel_key = key
+ break
+
+ if not parallel_key:
+ raise ValueError("No list found in kwargs to parallelize over.")
+
+ parallel_list = kwargs.pop(parallel_key)
+ logger.info(f"Parallelizing '{parallel_key}' across {max_workers} workers")
+
+ # Put large shared objects into the Ray Object Store once
+ optimized_kwargs = {
+ k: (ray.put(v) if isinstance(v, (pd.DataFrame, pd.Series, dict, list, BaseContext)) else v)
+ for k, v in kwargs.items()
+ }
+
+ # Submit sliced chunks to reduce inter-node data transfer
+ remote_task_loop = ray.remote(self._ray_task_loop)
+ for i in range(max_workers):
+ chunk = parallel_list[i::max_workers]
+ if not chunk:
+ continue
+
+ task = remote_task_loop.remote(
+ fn,
+ parallel_key,
+ chunk,
+ i,
+ **optimized_kwargs,
+ )
+ self._ray_task_list.append(task)
+ logger.info(f"Submitted task {i + 1}/{max_workers} for {task_desc}")
+
+ return self.join_ray_task(task_desc=task_desc)
+
+ @staticmethod
+ def _ray_task_loop(internal_fn: Callable, parallel_key: str, chunk: list, actor_index: int, **kwargs) -> list:
+ """Execute the function over a specific chunk of data on a worker."""
+ results = []
+ for value in chunk:
+ current_kwargs = {**kwargs, "actor_index": actor_index, parallel_key: value}
+ t_result = internal_fn(**current_kwargs)
+
+ if t_result is not None:
+ if isinstance(t_result, list):
+ results.extend(t_result)
+ else:
+ results.append(t_result)
+ return results
+
+ def submit_ray_task(self, fn, *args, **kwargs):
+ """Submit a single Ray task to the task list for later execution."""
+ if not ray.is_initialized():
+ ray.init(num_cpus=C.service_config.ray_max_workers, ignore_reinit_error=True)
+
+ remote_fn = ray.remote(fn)
+ task = remote_fn.remote(*args, **kwargs)
+ self._ray_task_list.append(task)
+ return self
+
+ def join_ray_task(self, task_desc: str | None = None) -> list:
+ """Collect results from Ray workers using a progress bar."""
+ results = []
+ unfinished = list(self._ray_task_list)
+
+ with tqdm(total=len(unfinished), desc=task_desc or f"{self.name}_ray") as pbar:
+ while unfinished:
+ ready, unfinished = ray.wait(unfinished, num_returns=1)
+ for obj_ref in ready:
+ try:
+ t_result = ray.get(obj_ref)
+ if isinstance(t_result, list):
+ results.extend(t_result)
+ elif t_result is not None:
+ results.append(t_result)
+ except Exception as e:
+ logger.error(f"Worker task failed: {e}")
+ pbar.update(1)
+
+ self._ray_task_list.clear()
+ return results
diff --git a/reme_ai/core/op/parallel_op.py b/reme_ai/core/op/parallel_op.py
new file mode 100644
index 00000000..18b84d0c
--- /dev/null
+++ b/reme_ai/core/op/parallel_op.py
@@ -0,0 +1,33 @@
+"""Module providing the ParallelOp class for concurrent operation execution."""
+
+from .base_op import BaseOp
+
+
+class ParallelOp(BaseOp):
+ """Operation class that executes multiple sub-operations in parallel."""
+
+ async def execute(self):
+ """Executes all sub-operations concurrently using asynchronous tasks."""
+ for op in self.sub_ops:
+ assert op.async_mode
+ self.submit_async_task(op.call, context=self.context)
+ await self.join_async_tasks()
+
+ def execute_sync(self):
+ """Executes all sub-operations concurrently using synchronous task management."""
+ for op in self.sub_ops:
+ assert not op.async_mode
+ self.submit_sync_task(op.call_sync, context=self.context)
+ self.join_sync_tasks()
+
+ def __lshift__(self, op: dict[str, BaseOp] | list[BaseOp] | BaseOp):
+ """Raises RuntimeError as the shift operator is not supported for parallel operations."""
+ raise RuntimeError(f"`<<` is not supported in `{self.name}`")
+
+ def __or__(self, op: BaseOp):
+ """Adds sub-operations to the current parallel group using the bitwise OR operator."""
+ if isinstance(op, ParallelOp) and op.sub_ops:
+ self.add_sub_ops(op.sub_ops)
+ else:
+ self.add_sub_op(op)
+ return self
diff --git a/reme_ai/core/op/sequential_op.py b/reme_ai/core/op/sequential_op.py
new file mode 100644
index 00000000..3dabb0c9
--- /dev/null
+++ b/reme_ai/core/op/sequential_op.py
@@ -0,0 +1,31 @@
+"""Module providing the SequentialOp class for serial operation execution."""
+
+from .base_op import BaseOp
+
+
+class SequentialOp(BaseOp):
+ """Operation class that executes sub-operations one after another in order."""
+
+ async def execute(self):
+ """Executes sub-operations sequentially using asynchronous awaits."""
+ for op in self.sub_ops:
+ assert op.async_mode
+ await op.call(context=self.context)
+
+ def execute_sync(self):
+ """Executes sub-operations sequentially in a synchronous blocking manner."""
+ for op in self.sub_ops:
+ assert not op.async_mode
+ op.call_sync(context=self.context)
+
+ def __lshift__(self, op: dict[str, BaseOp] | list[BaseOp] | BaseOp):
+ """Raises RuntimeError as the left shift operator is not supported."""
+ raise RuntimeError(f"`<<` is not supported in `{self.name}`")
+
+ def __rshift__(self, op: BaseOp):
+ """Appends operations to the sequence using the bitwise right shift operator."""
+ if isinstance(op, SequentialOp) and op.sub_ops:
+ self.add_sub_ops(op.sub_ops)
+ else:
+ self.add_sub_op(op)
+ return self
diff --git a/reme_ai/core/reme.py b/reme_ai/core/reme.py
new file mode 100644
index 00000000..d1240cb6
--- /dev/null
+++ b/reme_ai/core/reme.py
@@ -0,0 +1,92 @@
+"""ReMe application classes for simplified configuration and execution."""
+
+import sys
+
+from .application import Application
+from .config import ReMeConfigParser
+from .context import C
+
+
+class ReMe(Application):
+ """Simplified ReMe application that auto-initializes the service context."""
+
+ def __init__(
+ self,
+ *args,
+ llm_api_key: str | None = None,
+ llm_api_base: str | None = None,
+ embedding_api_key: str | None = None,
+ embedding_api_base: str | None = None,
+ enable_logo: bool = True,
+ llm: dict | None = None,
+ embedding_model: dict | None = None,
+ vector_store: dict | None = None,
+ token_counter: dict | None = None,
+ **kwargs,
+ ):
+ super().__init__(
+ *args,
+ llm_api_key=llm_api_key,
+ llm_api_base=llm_api_base,
+ embedding_api_key=embedding_api_key,
+ embedding_api_base=embedding_api_base,
+ service_config=None,
+ parser=ReMeConfigParser,
+ config_path=None,
+ enable_logo=enable_logo,
+ llm=llm,
+ embedding_model=embedding_model,
+ vector_store=vector_store,
+ token_counter=token_counter,
+ **kwargs,
+ )
+
+ C.initialize_service_context()
+
+ async def summary(self):
+ """Execute summary operations."""
+
+ async def retrieve(self):
+ """Execute retrieve operations."""
+
+
+class ReMeApp(Application):
+ """ReMe application with config file support and flow execution methods."""
+
+ def __init__(
+ self,
+ *args,
+ llm_api_key: str | None = None,
+ llm_api_base: str | None = None,
+ embedding_api_key: str | None = None,
+ embedding_api_base: str | None = None,
+ config_path: str | None = None,
+ enable_logo: bool = True,
+ **kwargs,
+ ):
+ super().__init__(
+ *args,
+ llm_api_key=llm_api_key,
+ llm_api_base=llm_api_base,
+ embedding_api_key=embedding_api_key,
+ embedding_api_base=embedding_api_base,
+ service_config=None,
+ parser=ReMeConfigParser,
+ config_path=config_path,
+ enable_logo=enable_logo,
+ **kwargs,
+ )
+
+ async def async_execute(self, name: str, **kwargs) -> dict:
+ """Execute a flow asynchronously and return the result as a dictionary."""
+ return (await self.execute_flow(name=name, **kwargs)).model_dump()
+
+
+def main():
+ """Main entry point for running ReMe application from command line."""
+ with ReMeApp(*sys.argv[1:]) as app:
+ app.run_service()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/reme_ai/core/schema/request.py b/reme_ai/core/schema/request.py
index aa054219..ece942b9 100644
--- a/reme_ai/core/schema/request.py
+++ b/reme_ai/core/schema/request.py
@@ -1,17 +1,11 @@
"""Defines the data structure for processing incoming user requests and message history."""
-from typing import List
-
from pydantic import Field, BaseModel, ConfigDict
-from .message import Message
-
class Request(BaseModel):
"""Represents a structured request payload containing a query, message list, and metadata."""
model_config = ConfigDict(extra="allow")
- query: str = Field(default="")
- messages: List[Message] = Field(default_factory=list)
metadata: dict = Field(default_factory=dict)
diff --git a/reme_ai/core/schema/service_config.py b/reme_ai/core/schema/service_config.py
index 67cac68f..6471e982 100644
--- a/reme_ai/core/schema/service_config.py
+++ b/reme_ai/core/schema/service_config.py
@@ -1,5 +1,6 @@
"""Configuration schemas for service components using Pydantic models."""
+import os
from typing import Dict, List
from pydantic import BaseModel, Field, ConfigDict
@@ -12,7 +13,7 @@ class MCPConfig(BaseModel):
model_config = ConfigDict(extra="allow")
- transport: str = Field(default="")
+ transport: str = Field(default="stdio")
host: str = Field(default="0.0.0.0")
port: int = Field(default=8001)
@@ -92,16 +93,14 @@ class ServiceConfig(BaseModel):
model_config = ConfigDict(extra="allow")
backend: str = Field(default="")
+ app_name: str = Field(default=os.getenv("APP_NAME", "ReMe"))
enable_logo: bool = Field(default=True)
language: str = Field(default="")
thread_pool_max_workers: int = Field(default=16)
ray_max_workers: int = Field(default=-1)
disabled_flows: List[str] = Field(default_factory=list)
enabled_flows: List[str] = Field(default_factory=list)
- external_mcp: Dict[str, dict] = Field(
- default_factory=dict,
- description="External MCP Server configuration",
- )
+ mcp_servers: Dict[str, dict] = Field(default_factory=dict, description="External MCP Server configuration")
mcp: MCPConfig = Field(default_factory=MCPConfig)
http: HttpConfig = Field(default_factory=HttpConfig)
diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py
index a94dfb1d..0a96f8c4 100644
--- a/reme_ai/core/schema/tool_call.py
+++ b/reme_ai/core/schema/tool_call.py
@@ -16,7 +16,7 @@ class ToolAttr(BaseModel):
model_config = ConfigDict(extra="allow")
- type: str = Field(default=JsonSchemaEnum.STRING.value, description="The data type of the attribute")
+ type: str = Field(default=str(JsonSchemaEnum.STRING), description="The data type of the attribute")
description: Optional[str] = Field(default=None, description="Description of the attribute")
required: Optional[List[str]] = Field(default=None, description="Required property names for object types")
properties: Optional[Dict[str, "ToolAttr"]] = Field(default=None, description="Child properties for objects")
@@ -27,7 +27,7 @@ class ToolAttr(BaseModel):
@classmethod
def validate_type_is_valid_enum(cls, v: str) -> str:
"""Validates that the provided type string exists within JsonSchemaEnum values."""
- valid_types = [e.value for e in JsonSchemaEnum]
+ valid_types = [str(e) for e in JsonSchemaEnum]
if v not in valid_types:
raise ValueError(f"Invalid type: '{v}'. Must be one of {valid_types}")
diff --git a/reme_ai/core/service/__init__.py b/reme_ai/core/service/__init__.py
new file mode 100644
index 00000000..e9f00a65
--- /dev/null
+++ b/reme_ai/core/service/__init__.py
@@ -0,0 +1,13 @@
+"""service"""
+
+from .base_service import BaseService
+from .cmd_service import CmdService
+from .http_service import HttpService
+from .mcp_service import MCPService
+
+__all__ = [
+ "BaseService",
+ "CmdService",
+ "HttpService",
+ "MCPService",
+]
diff --git a/reme_ai/core/service/base_service.py b/reme_ai/core/service/base_service.py
new file mode 100644
index 00000000..28c82fb3
--- /dev/null
+++ b/reme_ai/core/service/base_service.py
@@ -0,0 +1,41 @@
+"""Base service definitions for flow management."""
+
+from abc import ABC, abstractmethod
+
+from loguru import logger
+from pydantic import BaseModel
+
+from ..context import C
+from ..flow import BaseFlow
+from ..schema import ToolCall
+from ..utils import create_pydantic_model
+
+
+class BaseService(ABC):
+ """Abstract base class for services that integrate and execute flows."""
+
+ def __init__(self, **kwargs):
+ """Initialize the base service."""
+ self.kwargs = kwargs
+
+ @abstractmethod
+ def integrate_flow(self, flow: BaseFlow) -> str | None:
+ """Integrate a flow into the service and return its name if successful."""
+
+ @staticmethod
+ def _prepare_route(flow: BaseFlow) -> tuple[ToolCall, type[BaseModel]]:
+ """Generate the request model and route name for a flow."""
+ tool_call = flow.tool_call
+ model = create_pydantic_model(tool_call.name, tool_call.parameters)
+ return tool_call, model
+
+ def run(self):
+ """Initialize and integrate all flows registered in the global context."""
+ flow_names: list[str] = []
+ for _, flow in C.flow_dict.items():
+ flow_name = self.integrate_flow(flow)
+ if flow_name:
+ flow_names.append(flow_name)
+
+ if flow_names:
+ logger.info(f"integrate {','.join(flow_names)}")
diff --git a/reme_ai/core/service/cmd_service.py b/reme_ai/core/service/cmd_service.py
new file mode 100644
index 00000000..bed8ec86
--- /dev/null
+++ b/reme_ai/core/service/cmd_service.py
@@ -0,0 +1,36 @@
+"""Command service module for managing and executing command-based workflows."""
+
+from loguru import logger
+
+from .base_service import BaseService
+from ..context import C
+from ..flow import CmdFlow, BaseFlow
+from ..utils.common_utils import run_coro_safely
+
+
+@C.register_service("cmd")
+class CmdService(BaseService):
+ """Service implementation for handling command flow execution logic."""
+
+ def __init__(self, **kwargs):
+ """Initialize the command service instance."""
+ super().__init__(**kwargs)
+ self._cmd_flow: CmdFlow | None = None
+
+ def integrate_flow(self, flow: BaseFlow) -> str | None:
+ """Integrate the workflow configuration into the command service."""
+ self._cmd_flow = CmdFlow(flow=C.service_config.flow)
+
+ def run(self):
+ """Execute the command flow in either asynchronous or synchronous mode."""
+ super().run()
+
+ if self._cmd_flow.async_mode:
+ response = run_coro_safely(
+ self._cmd_flow.call(**C.service_config.cmd.model_extra),
+ )
+ else:
+ response = self._cmd_flow.call_sync(**C.service_config.cmd.model_extra)
+
+ if response.answer:
+ logger.info(f"response.answer={response.answer}")
diff --git a/reme_ai/core/service/http_service.py b/reme_ai/core/service/http_service.py
new file mode 100644
index 00000000..f63d7333
--- /dev/null
+++ b/reme_ai/core/service/http_service.py
@@ -0,0 +1,87 @@
+"""HTTP service implementation using FastAPI."""
+
+import asyncio
+from collections.abc import AsyncGenerator
+
+import uvicorn
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+
+from .base_service import BaseService
+from ..context import C
+from ..flow import BaseFlow
+from ..schema import Response
+from ..utils.common_utils import execute_stream_task
+
+
+@C.register_service("http")
+class HttpService(BaseService):
+ """Expose flows via HTTP REST and SSE endpoints."""
+
+ def __init__(self, **kwargs):
+ """Initialize FastAPI app with CORS and health checks."""
+ super().__init__(**kwargs)
+ self.app = FastAPI(title=C.service_config.app_name)
+ self.app.add_middleware(
+ CORSMiddleware,
+ allow_origins=["*"],
+ allow_credentials=True,
+ allow_methods=["*"],
+ allow_headers=["*"],
+ )
+ self.app.get("/health")(lambda: {"status": "healthy"})
+
+ def _integrate_flow(self, flow: BaseFlow) -> str:
+ """Register a standard flow as a POST endpoint."""
+ tool_call, request_model = self._prepare_route(flow)
+
+ async def execute_endpoint(request: request_model) -> Response:
+ return await flow.call(**request.model_dump(exclude_none=True))
+
+ self.app.post(
+ path=f"/{tool_call.name}",
+ response_model=Response,
+ description=tool_call.description,
+ )(execute_endpoint)
+ return tool_call.name
+
+ def _integrate_stream_flow(self, flow: BaseFlow) -> str:
+ """Register a streaming flow as an SSE endpoint."""
+ tool_call, request_model = self._prepare_route(flow)
+
+ async def execute_stream_endpoint(request: request_model) -> StreamingResponse:
+ queue = asyncio.Queue()
+ # Start flow as a background task
+ task = asyncio.create_task(flow.call(stream_queue=queue, **request.model_dump(exclude_none=True)))
+
+ async def generate_stream() -> AsyncGenerator[bytes, None]:
+ async for chunk in execute_stream_task(
+ queue=queue,
+ task=task,
+ flow_name=tool_call.name,
+ as_bytes=True,
+ ):
+ yield chunk
+
+ return StreamingResponse(generate_stream(), media_type="text/event-stream")
+
+ self.app.post(f"/{tool_call.name}")(execute_stream_endpoint)
+ return tool_call.name
+
+ def integrate_flow(self, flow: BaseFlow) -> str | None:
+ """Register a flow based on its streaming configuration."""
+ return self._integrate_stream_flow(flow) if flow.stream else self._integrate_flow(flow)
+
+ def run(self):
+ """Start the Uvicorn server."""
+ super().run()
+ cfg = C.service_config.http
+ uvicorn.run(
+ self.app,
+ host=cfg.host,
+ port=cfg.port,
+ timeout_keep_alive=cfg.timeout_keep_alive,
+ limit_concurrency=cfg.limit_concurrency,
+ **cfg.model_extra,
+ )
diff --git a/reme_ai/core/service/mcp_service.py b/reme_ai/core/service/mcp_service.py
new file mode 100644
index 00000000..e60aec85
--- /dev/null
+++ b/reme_ai/core/service/mcp_service.py
@@ -0,0 +1,56 @@
+"""Model Context Protocol (MCP) service implementation."""
+
+from typing import Any
+
+from fastmcp import FastMCP
+from fastmcp.tools import FunctionTool
+
+from .base_service import BaseService
+from ..context import C
+from ..flow import BaseFlow
+
+
+@C.register_service("mcp")
+class MCPService(BaseService):
+ """Expose flows as Model Context Protocol (MCP) tools."""
+
+ def __init__(self, **kwargs: Any):
+ """Initialize FastMCP instance with service settings."""
+ super().__init__(**kwargs)
+ self.mcp = FastMCP(name=C.service_config.app_name)
+
+ def integrate_flow(self, flow: BaseFlow) -> str | None:
+ """Register a non-streaming flow as an MCP tool."""
+ if flow.stream:
+ return None
+
+ tool_call, request_model = self._prepare_route(flow)
+
+ async def execute_tool(**kwargs):
+ """Execute flow logic and return the string answer."""
+ request_instance = request_model(**kwargs)
+ response = await flow.call(**request_instance.model_dump(exclude_none=True))
+ return response.answer
+
+ self.mcp.add_tool(
+ FunctionTool(
+ name=tool_call.name,
+ description=tool_call.description,
+ fn=execute_tool,
+ parameters=tool_call.parameters.simple_input_dump(),
+ ),
+ )
+ return tool_call.name
+
+ def run(self):
+ """Run the MCP server with specified transport protocol."""
+ super().run()
+ cfg = C.service_config.mcp
+
+ run_args: dict = {"transport": cfg.transport, "show_banner": False, **cfg.model_extra}
+
+ # Add network settings for non-stdio transports
+ if cfg.transport != "stdio":
+ run_args.update({"host": cfg.host, "port": cfg.port})
+
+ self.mcp.run(**run_args)
diff --git a/reme_ai/core/token_counter/__init__.py b/reme_ai/core/token_counter/__init__.py
new file mode 100644
index 00000000..a9b50826
--- /dev/null
+++ b/reme_ai/core/token_counter/__init__.py
@@ -0,0 +1,11 @@
+"""token counter"""
+
+from .base_token_counter import BaseTokenCounter
+from .hf_token_counter import HFTokenCounter
+from .openai_token_counter import OpenAITokenCounter
+
+__all__ = [
+ "BaseTokenCounter",
+ "HFTokenCounter",
+ "OpenAITokenCounter",
+]
diff --git a/reme_ai/core/token_counter/base_token_counter.py b/reme_ai/core/token_counter/base_token_counter.py
new file mode 100644
index 00000000..76ccdb78
--- /dev/null
+++ b/reme_ai/core/token_counter/base_token_counter.py
@@ -0,0 +1,59 @@
+"""Token counting utility based on character-type rules."""
+
+import math
+import re
+from loguru import logger
+
+from ..context import C
+from ..schema import Message, ToolCall
+
+
+@C.register_token_counter("base")
+class BaseTokenCounter:
+ """A rule-based token counter for Chinese and non-Chinese text."""
+
+ def __init__(self, model_name: str, **kwargs):
+ """Initialize with model name and additional parameters."""
+ self.model_name = model_name
+ self.kwargs = kwargs
+ # Matches Chinese characters including extensions
+ self._cn_regex = re.compile(r"[\u4e00-\u9fff]")
+
+ def _count_chars(self, text: str) -> tuple[int, int]:
+ """Count Chinese and other characters in a string."""
+ if not text:
+ return 0, 0
+ cn_count = len(self._cn_regex.findall(text))
+ return cn_count, len(text) - cn_count
+
+ def count_token(
+ self,
+ messages: list[Message],
+ tools: list[ToolCall] | None = None,
+ **_kwargs,
+ ) -> int:
+ """Calculate total tokens using the 1:2 (CN) and 1:4 (Other) rule."""
+ cn_total = 0
+ ot_total = 0
+ logger.info("Calculating tokens using rule-based estimation.")
+
+ # Extract text from messages
+ segments = []
+ for msg in messages:
+ content = msg.content
+ if isinstance(content, bytes):
+ content = content.decode("utf-8", errors="ignore")
+ segments.extend([content, msg.reasoning_content])
+
+ # Extract text from tools
+ if tools:
+ for tool in tools:
+ segments.extend([tool.name, tool.description, tool.arguments])
+
+ # Process all segments
+ for text in filter(None, segments):
+ cn_chars, ot_chars = self._count_chars(text)
+ cn_total += cn_chars
+ ot_total += ot_chars
+
+ return math.ceil(cn_total / 2) + math.ceil(ot_total / 4)
diff --git a/reme_ai/core/token_counter/hf_token_counter.py b/reme_ai/core/token_counter/hf_token_counter.py
new file mode 100644
index 00000000..0bad4c78
--- /dev/null
+++ b/reme_ai/core/token_counter/hf_token_counter.py
@@ -0,0 +1,82 @@
+"""HuggingFace token counting utilities."""
+
+import os
+
+from loguru import logger
+
+from .base_token_counter import BaseTokenCounter
+from ..context import C
+from ..schema import Message, ToolCall
+
+
+@C.register_token_counter("hf")
+class HFTokenCounter(BaseTokenCounter):
+ """Token counter using transformers.AutoTokenizer.apply_chat_template."""
+
+ def __init__(
+ self,
+ model_name: str,
+ use_fast: bool = False,
+ trust_remote_code: bool = False,
+ use_mirror: bool = True,
+ **kwargs,
+ ):
+ """Initialize the counter with model config and lazy tokenizer loading."""
+ super().__init__(model_name=model_name, **kwargs)
+ self.use_fast = use_fast
+ self.trust_remote_code = trust_remote_code
+ self.use_mirror = use_mirror
+ self._tokenizer = None
+
+ def _ensure_tokenizer(self):
+ """Initialize and cache the HuggingFace tokenizer safely."""
+ if self._tokenizer:
+ return self._tokenizer
+
+ if self.use_mirror:
+ os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
+
+ try:
+ from transformers import AutoTokenizer
+
+ logger.info("Initializing HuggingFace tokenizer for {}", self.model_name)
+
+ tokenizer = AutoTokenizer.from_pretrained(
+ self.model_name,
+ use_fast=self.use_fast,
+ trust_remote_code=self.trust_remote_code,
+ **self.kwargs,
+ )
+
+ if not hasattr(tokenizer, "chat_template") or tokenizer.chat_template is None:
+ raise ValueError(f"Model {self.model_name} lacks a chat template.")
+
+ self._tokenizer = tokenizer
+ return tokenizer
+ except Exception as e:
+ logger.error("Failed to load tokenizer {}: {}", self.model_name, e)
+ raise
+
+ def count_token(
+ self,
+ messages: list[Message],
+ tools: list[ToolCall] | None = None,
+ **kwargs,
+ ) -> int:
+ """Calculate total tokens for messages and tools using the chat template."""
+ tokenizer = self._ensure_tokenizer()
+
+ # Serialize inputs for the template
+ formatted_msgs = [m.simple_dump() for m in messages]
+ formatted_tools = [t.simple_input_dump() for t in tools] if tools else None
+
+ # Setting tokenize=True and leaving return_tensors=None returns a List[int]
+ tokens = tokenizer.apply_chat_template(
+ formatted_msgs,
+ tools=formatted_tools,
+ add_generation_prompt=kwargs.pop("add_generation_prompt", False),
+ tokenize=True,
+ **kwargs,
+ )
+
+ return len(tokens)
diff --git a/reme_ai/core/token_counter/openai_token_counter.py b/reme_ai/core/token_counter/openai_token_counter.py
new file mode 100644
index 00000000..5793e88a
--- /dev/null
+++ b/reme_ai/core/token_counter/openai_token_counter.py
@@ -0,0 +1,58 @@
+"""Token counting implementation for OpenAI-compatible models."""
+
+import json
+from loguru import logger
+from .base_token_counter import BaseTokenCounter
+from ..context import C
+from ..schema import Message, ToolCall
+
+
+@C.register_token_counter("openai")
+class OpenAITokenCounter(BaseTokenCounter):
+ """Token counter for OpenAI models using tiktoken."""
+
+ def __init__(self, model_name: str, **kwargs):
+ super().__init__(model_name, **kwargs)
+ self._encoding = None
+
+ @property
+ def encoding(self):
+ """Get or initialize the tiktoken encoding for the specified model."""
+ if self._encoding is None:
+ import tiktoken
+
+ try:
+ self._encoding = tiktoken.encoding_for_model(self.model_name)
+ except KeyError:
+ logger.warning(f"Model {self.model_name} not found; falling back to o200k_base.")
+ self._encoding = tiktoken.get_encoding("o200k_base")
+ return self._encoding
+
+ def count_token(
+ self,
+ messages: list[Message],
+ tools: list[ToolCall] | None = None,
+ **_kwargs,
+ ) -> int:
+ """Calculate total tokens for a request including messages and tool definitions."""
+ enc = self.encoding
+ total_tokens = 0
+
+ for msg in messages:
+ # Every message has <|start|>{role/name}\n{content}<|end|>\n
+ total_tokens += 3 # Base overhead per message
+ if msg.content:
+ total_tokens += len(enc.encode(msg.content))
+
+ if msg.tool_calls:
+ for tc in msg.tool_calls:
+ dump = json.dumps(tc.simple_output_dump(), ensure_ascii=False)
+ total_tokens += len(enc.encode(dump))
+
+ if tools:
+ # Account for tool/function definitions if provided
+ tool_json = json.dumps([t.simple_input_dump() for t in tools], ensure_ascii=False)
+ total_tokens += len(enc.encode(tool_json))
+
+ total_tokens += 3 # Every reply is primed with <|start|>assistant<|message|>
+ return total_tokens
diff --git a/reme_ai/core/tool/__init__.py b/reme_ai/core/tool/__init__.py
new file mode 100644
index 00000000..02dbd468
--- /dev/null
+++ b/reme_ai/core/tool/__init__.py
@@ -0,0 +1,7 @@
+"""tool"""
+
+from .mcp_tool import MCPTool
+
+__all__ = [
+ "MCPTool",
+]
diff --git a/reme_ai/core/tool/mcp_tool.py b/reme_ai/core/tool/mcp_tool.py
new file mode 100644
index 00000000..802e4138
--- /dev/null
+++ b/reme_ai/core/tool/mcp_tool.py
@@ -0,0 +1,82 @@
+"""MCP (Model Context Protocol) tool integration for remote tool execution."""
+
+from typing import List
+
+from ..context import C
+from ..op import BaseOp
+from ..schema import ToolCall
+from ..utils import MCPClient
+
+
+@C.register_op()
+class MCPTool(BaseOp):
+ """Operator for calling remote MCP (Model Context Protocol) tools.
+
+ This class enables integration with external MCP servers to execute tools
+ and retrieve their results. It supports parameter customization and retry logic.
+ """
+
+ def __init__(
+ self,
+ mcp_server: str = "",
+ tool_name: str = "",
+ save_response_result: bool = True,
+ parameter_required: List[str] | None = None,
+ parameter_optional: List[str] | None = None,
+ parameter_deleted: List[str] | None = None,
+ max_retries: int = 3,
+ timeout: float | None = None,
+ raise_exception: bool = False,
+ **kwargs,
+ ):
+
+ super().__init__(
+ save_response_result=save_response_result,
+ max_retries=max_retries,
+ raise_exception=raise_exception,
+ **kwargs,
+ )
+
+ self.mcp_server: str = mcp_server
+ self.tool_name: str = tool_name
+ self.parameter_required: List[str] | None = parameter_required
+ self.parameter_optional: List[str] | None = parameter_optional
+ self.parameter_deleted: List[str] | None = parameter_deleted
+ self.timeout: float | None = timeout
+ # Example MCP marketplace: https://bailian.console.aliyun.com/?tab=mcp#/mcp-market
+
+ self._client = MCPClient(C.service_config.mcp_servers)
+
+ def _build_tool_call(self) -> ToolCall:
+ tool_call_dict = C.mcp_server_mapping[self.mcp_server]
+ tool_call: ToolCall = tool_call_dict[self.tool_name].model_copy(deep=True)
+
+ # Initialize required list if not exists
+ if tool_call.parameters.required is None:
+ tool_call.parameters.required = []
+
+ if self.parameter_required:
+ for name in self.parameter_required:
+ if name not in tool_call.parameters.required:
+ tool_call.parameters.required.append(name)
+
+ if self.parameter_optional:
+ for name in self.parameter_optional:
+ if name in tool_call.parameters.required:
+ tool_call.parameters.required.remove(name)
+
+ if self.parameter_deleted:
+ for name in self.parameter_deleted:
+ tool_call.parameters.properties.pop(name, None)
+ if tool_call.parameters.required and name in tool_call.parameters.required:
+ tool_call.parameters.required.remove(name)
+
+ return tool_call
+
+ async def execute(self):
+ self.output = await self._client.call_tool(
+ server_name=self.mcp_server,
+ tool_name=self.tool_name,
+ arguments=self.input_dict,
+ parse_text_result=True,
+ )
diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py
index ed1916b2..87e01299 100644
--- a/reme_ai/core/utils/__init__.py
+++ b/reme_ai/core/utils/__init__.py
@@ -1,14 +1,34 @@
"""utils"""
+from .cache_handler import CacheHandler
from .case_converter import snake_to_camel, camel_to_snake
+from .common_utils import run_coro_safely, execute_stream_task
from .env_utils import load_env
+from .http_client import HttpClient
+from .llm_utils import extract_content, format_messages
+from .logger_utils import init_logger
+from .logo_utils import print_logo
+from .mcp_client import MCPClient
+from .pydantic_config_parser import PydanticConfigParser
+from .pydantic_utils import create_pydantic_model
from .singleton import singleton
from .timer import timer
__all__ = [
+ "CacheHandler",
"snake_to_camel",
"camel_to_snake",
+ "run_coro_safely",
+ "execute_stream_task",
"load_env",
+ "HttpClient",
+ "extract_content",
+ "format_messages",
+ "init_logger",
+ "print_logo",
+ "MCPClient",
+ "PydanticConfigParser",
+ "create_pydantic_model",
"singleton",
"timer",
]
diff --git a/reme_ai/core/utils/cache_handler.py b/reme_ai/core/utils/cache_handler.py
new file mode 100644
index 00000000..70c8585a
--- /dev/null
+++ b/reme_ai/core/utils/cache_handler.py
@@ -0,0 +1,184 @@
+"""Local file-based cache utility for DataFrames, lists, dicts, and strings."""
+
+import json
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+from loguru import logger
+
+
+class CacheHandler:
+ """Handles persistent data caching with expiration and type support."""
+
+ _EXTENSIONS = {
+ pd.DataFrame: ".csv",
+ dict: ".json",
+ list: ".json",
+ str: ".txt",
+ }
+
+ _TYPE_NAMES = {
+ "DataFrame": pd.DataFrame,
+ "dict": dict,
+ "list": list,
+ "str": str,
+ }
+
+ def __init__(self, cache_dir: str | Path = "cache"):
+ """Initialize cache directory and load existing metadata."""
+ self.cache_dir = Path(cache_dir)
+ self.cache_dir.mkdir(parents=True, exist_ok=True)
+ self.metadata_file = self.cache_dir / "metadata.json"
+ self.metadata: dict[str, Any] = self._load_metadata()
+
+ def set_cache_dir(self, cache_dir: str | Path) -> None:
+ """Change the cache directory and reload metadata."""
+ self.cache_dir = Path(cache_dir)
+ self.cache_dir.mkdir(parents=True, exist_ok=True)
+ self.metadata_file = self.cache_dir / "metadata.json"
+ self.metadata = self._load_metadata()
+ logger.info(f"Cache directory moved to: {self.cache_dir}")
+
+ def _load_metadata(self) -> dict[str, Any]:
+ """Load metadata from the JSON file."""
+ if self.metadata_file.exists():
+ try:
+ with open(self.metadata_file, "r", encoding="utf-8") as f:
+ return json.load(f)
+ except (json.JSONDecodeError, OSError) as e:
+ logger.warning(f"Metadata load failed: {e}")
+ return {}
+
+ def _save_metadata(self) -> None:
+ """Persist metadata to the disk."""
+ try:
+ with open(self.metadata_file, "w", encoding="utf-8") as f:
+ json.dump(self.metadata, f, ensure_ascii=False, indent=2)
+ except OSError as e:
+ logger.error(f"Metadata save failed: {e}")
+
+ def _get_path(self, key: str, data_type: type | None = None) -> Path:
+ """Resolve the file path based on data type or metadata."""
+ ext = ".dat"
+ if data_type in self._EXTENSIONS:
+ ext = self._EXTENSIONS[data_type]
+ elif key in self.metadata:
+ stored_type = self.metadata[key].get("data_type")
+ ext = self._EXTENSIONS.get(self._TYPE_NAMES.get(stored_type, None), ".dat")
+ return self.cache_dir / f"{key}{ext}"
+
+ @staticmethod
+ def _execute_save(data: Any, path: Path, dtype: type, **kwargs) -> dict:
+ """Execute type-specific save operations."""
+ if dtype is pd.DataFrame:
+ data.to_csv(path, index=kwargs.get("index", False), encoding="utf-8")
+ return {"row_count": len(data), "file_size": path.stat().st_size}
+
+ if dtype in (dict, list):
+ with open(path, "w", encoding="utf-8") as f:
+ json.dump(data, f, ensure_ascii=False, indent=2)
+ return {"item_count": len(data), "file_size": path.stat().st_size}
+
+ if dtype is str:
+ path.write_text(data, encoding=kwargs.get("encoding", "utf-8"))
+ return {"char_count": len(data), "file_size": path.stat().st_size}
+
+ raise ValueError(f"Unsupported type: {dtype}")
+
+ @staticmethod
+ def _execute_load(path: Path, type_name: str, **kwargs) -> Any:
+ """Execute type-specific load operations."""
+ if type_name == "DataFrame":
+ return pd.read_csv(path, encoding=kwargs.get("encoding", "utf-8"))
+ if type_name in ("dict", "list"):
+ with open(path, "r", encoding="utf-8") as f:
+ return json.load(f)
+ if type_name == "str":
+ return path.read_text(encoding=kwargs.get("encoding", "utf-8"))
+ raise ValueError(f"Unknown data type in metadata: {type_name}")
+
+ def save(self, key: str, data: Any, expire_hours: float | None = None, **kwargs) -> bool:
+ """Save data to cache with optional expiration."""
+ try:
+ dtype = type(data)
+ path = self._get_path(key, dtype)
+ stats = self._execute_save(data, path, dtype, **kwargs)
+
+ now = datetime.now()
+ self.metadata[key] = {
+ "created_at": now.isoformat(),
+ "expire_at": (now + timedelta(hours=expire_hours)).isoformat() if expire_hours else None,
+ "data_type": dtype.__name__,
+ **stats,
+ }
+ self._save_metadata()
+ return True
+ except Exception as e:
+ logger.error(f"Save failed for {key}: {e}")
+ return False
+
+ def load(self, key: str, auto_clean: bool = True, **kwargs) -> Any | None:
+ """Load data from cache if not expired."""
+ if self._is_expired(key):
+ if auto_clean:
+ self.delete(key)
+ return None
+
+ path = self._get_path(key)
+ if not path.exists() or key not in self.metadata:
+ return None
+
+ try:
+ return self._execute_load(path, self.metadata[key]["data_type"], **kwargs)
+ except Exception as e:
+ logger.error(f"Load failed for {key}: {e}")
+ return None
+
+ def _is_expired(self, key: str) -> bool:
+ """Check if the cached entry has expired."""
+ entry = self.metadata.get(key)
+ if not entry or not entry.get("expire_at"):
+ return False
+ return datetime.now() > datetime.fromisoformat(entry["expire_at"])
+
+ def delete(self, key: str) -> bool:
+ """Remove a specific cache entry and its file."""
+ try:
+ path = self._get_path(key)
+ if path.exists():
+ path.unlink()
+ if key in self.metadata:
+ del self.metadata[key]
+ self._save_metadata()
+ return True
+ except OSError as e:
+ logger.error(f"Delete failed for {key}: {e}")
+ return False
+
+ def exists(self, key: str) -> bool:
+ """Check if a valid cache entry exists."""
+ return key in self.metadata and not self._is_expired(key)
+
+ def clear_all(self) -> bool:
+ """Purge all cache files and reset metadata."""
+ try:
+ for file in self.cache_dir.iterdir():
+ if file.is_file():
+ file.unlink()
+ self.metadata = {}
+ self._save_metadata()
+ return True
+ except OSError as e:
+ logger.error(f"Clear all failed: {e}")
+ return False
+
+ def get_stats(self) -> dict[str, Any]:
+ """Return cache usage statistics."""
+ total_size = sum(f.stat().st_size for f in self.cache_dir.glob("*") if f.is_file())
+ return {
+ "count": len(self.metadata),
+ "size_mb": round(total_size / (1024 * 1024), 2),
+ "dir": str(self.cache_dir),
+ }
diff --git a/reme_ai/core/utils/common_utils.py b/reme_ai/core/utils/common_utils.py
new file mode 100644
index 00000000..d370064d
--- /dev/null
+++ b/reme_ai/core/utils/common_utils.py
@@ -0,0 +1,83 @@
+"""Common utility functions"""
+
+import asyncio
+from collections.abc import AsyncGenerator, Coroutine
+from typing import Any
+
+from loguru import logger
+
+from ..enumeration import ChunkEnum
+from ..schema import StreamChunk
+
+
+def run_coro_safely(coro: Coroutine[Any, Any, Any]) -> Any | asyncio.Task[Any]:
+ """Run a coroutine in the current event loop or a new one if none exists."""
+ try:
+ # Attempt to retrieve the event loop associated with the current thread
+ loop = asyncio.get_running_loop()
+
+ except RuntimeError:
+ # Start a new event loop to run the coroutine to completion
+ return asyncio.run(coro)
+
+ else:
+ # Schedule the coroutine as a background task in the active loop
+ return loop.create_task(coro)
+
+
+async def execute_stream_task(
+ queue: asyncio.Queue,
+ task: asyncio.Task,
+ flow_name: str | None = None,
+ as_bytes: bool = False,
+) -> AsyncGenerator[str | bytes, None]:
+ """
+ Core stream flow execution logic.
+
+ Handles streaming from a queue while monitoring the task completion.
+ Properly manages errors and resource cleanup.
+
+ Args:
+ queue: Queue to receive StreamChunk objects from
+ task: Background task executing the flow
+ flow_name: Optional flow name for logging purposes
+ as_bytes: If True, yield bytes for HTTP responses; if False, yield strings
+
+ Yields:
+ SSE-formatted data chunks (either str or bytes based on as_bytes)
+ """
+ done_msg = b"data:[DONE]\n\n" if as_bytes else "data:[DONE]\n\n"
+
+ try:
+ while True:
+ # Wait for next chunk or check if task failed
+ get_chunk = asyncio.create_task(queue.get())
+ done, _ = await asyncio.wait({get_chunk, task}, return_when=asyncio.FIRST_COMPLETED)
+
+ if get_chunk in done:
+ chunk: StreamChunk = get_chunk.result()
+ if chunk.done:
+ yield done_msg
+ break
+
+ data = f"data:{chunk.model_dump_json()}\n\n"
+ yield data.encode() if as_bytes else data
+ else:
+ # Task finished unexpectedly or raised exception
+ await task
+ yield done_msg
+ break
+
+ except Exception as e:
+ log_msg = f"Stream error in {flow_name}: {e}" if flow_name else f"Stream error: {e}"
+ logger.exception(log_msg)
+
+ err = StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e), done=True)
+ err_data = f"data:{err.model_dump_json()}\n\n"
+ yield err_data.encode() if as_bytes else err_data
+ yield done_msg
+
+ finally:
+ # Ensure task is cancelled if still running to avoid resource leaks
+ if not task.done():
+ task.cancel()
diff --git a/reme_ai/core/utils/http_client.py b/reme_ai/core/utils/http_client.py
new file mode 100644
index 00000000..8c8e92c3
--- /dev/null
+++ b/reme_ai/core/utils/http_client.py
@@ -0,0 +1,91 @@
+"""Asynchronous HTTP client for executing flows with built-in retry logic."""
+
+import json
+from collections.abc import AsyncIterator
+from typing import Optional
+
+import httpx
+from loguru import logger
+
+from ..schema import Response
+
+
+class HttpClient:
+ """Async client for flow endpoints with automated retries and error handling."""
+
+ def __init__(
+ self,
+ base_url: str = "http://localhost:8001",
+ timeout: float = 3600.0,
+ max_retries: int = 3,
+ raise_exception: bool = True,
+ ):
+ """Initialize the client with base configuration."""
+ self.base_url = base_url.rstrip("/")
+ self.timeout = timeout
+ self.max_retries = max_retries
+ self.raise_exception = raise_exception
+ self.client = httpx.AsyncClient(timeout=timeout)
+
+ async def __aenter__(self):
+ """Enter async context manager."""
+ return self
+
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
+ """Exit async context manager and close connection."""
+ await self.close()
+
+ async def close(self):
+ """Close the underlying HTTP client."""
+ await self.client.aclose()
+
+ async def health_check(self) -> dict[str, str]:
+ """Check the health status of the flow service."""
+ response = await self.client.get(f"{self.base_url}/health")
+ response.raise_for_status()
+ return response.json()
+
+ async def execute_flow(self, flow_name: str, **kwargs) -> Optional[Response]:
+ """Execute a flow with automated retry logic."""
+ endpoint = f"{self.base_url}/{flow_name}"
+
+ for attempt in range(self.max_retries):
+ try:
+ response = await self.client.post(endpoint, json=kwargs)
+ response.raise_for_status()
+ return Response(**response.json())
+
+ except (httpx.HTTPError, Exception) as e:
+ logger.error(f"Flow {flow_name} failed (attempt {attempt + 1}/{self.max_retries}): {e}")
+ if attempt == self.max_retries - 1 and self.raise_exception:
+ raise e
+ return None
+
+ async def list_endpoints(self) -> dict:
+ """Retrieve available endpoints from OpenAPI specification."""
+ response = await self.client.get(f"{self.base_url}/openapi.json")
+ response.raise_for_status()
+ return response.json()
+
+ async def execute_stream_flow(self, flow_name: str, **kwargs) -> AsyncIterator[dict[str, str]]:
+ """Execute a flow and yield parsed SSE stream chunks."""
+ endpoint = f"{self.base_url}/{flow_name}"
+
+ async with self.client.stream("POST", endpoint, json=kwargs) as response:
+ response.raise_for_status()
+ async for line in response.aiter_lines():
+ if not line or not line.startswith("data:"):
+ continue
+
+ content = line.removeprefix("data:").strip()
+ if content == "[DONE]":
+ break
+
+ try:
+ data = json.loads(content)
+ yield {
+ "type": data.get("chunk_type", "answer"),
+ "content": data.get("chunk", ""),
+ }
+ except json.JSONDecodeError:
+ continue
diff --git a/reme_ai/core/utils/llm_utils.py b/reme_ai/core/utils/llm_utils.py
new file mode 100644
index 00000000..ae0e763f
--- /dev/null
+++ b/reme_ai/core/utils/llm_utils.py
@@ -0,0 +1,40 @@
+"""Utility functions for processing and formatting LLM-related message data."""
+
+import json
+import re
+
+from ..enumeration import Role
+from ..schema import Message
+
+
+def format_messages(messages: list[Message | dict], enable_system: bool = False) -> str:
+ """Formats a list of messages into a single string, optionally filtering system roles."""
+ formatted_lines = []
+ for message in messages:
+ if isinstance(message, dict):
+ message = Message(**message)
+ if not enable_system and message.role is Role.SYSTEM:
+ continue
+
+ formatted_lines.append(message.format_message())
+ return "\n".join(formatted_lines)
+
+
+def extract_content(text: str, language_tag: str = "json", greedy: bool = False):
+ """Extracts content from Markdown code blocks and parses it if the tag is JSON."""
+ quantifier = ".*" if greedy else ".*?"
+ pattern = rf"```\s*{re.escape(language_tag)}\s*({quantifier})\s*```"
+ match = re.search(pattern, text, re.DOTALL)
+
+ if match:
+ result = match.group(1).strip()
+ else:
+ result = text
+
+ if language_tag == "json":
+ try:
+ result = json.loads(result)
+ except json.JSONDecodeError:
+ result = None
+
+ return result
diff --git a/reme_ai/core/utils/logger_utils.py b/reme_ai/core/utils/logger_utils.py
new file mode 100644
index 00000000..56478683
--- /dev/null
+++ b/reme_ai/core/utils/logger_utils.py
@@ -0,0 +1,40 @@
+"""Logging configuration module for application-wide tracing."""
+
+import os
+import sys
+from datetime import datetime
+
+
+def init_logger(log_dir: str = "logs", level: str = "INFO") -> None:
+ """Initialize the logger with both file and console handlers."""
+ from loguru import logger
+
+ # Remove default handler to avoid duplicate logs
+ logger.remove()
+
+ # Ensure the logging directory exists
+ os.makedirs(log_dir, exist_ok=True)
+
+ # Generate filename based on the current timestamp
+ current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ log_filename = f"{current_ts}.log"
+ log_filepath = os.path.join(log_dir, log_filename)
+
+ # Configure file-based logging with rotation and compression
+ logger.add(
+ log_filepath,
+ level=level,
+ rotation="00:00",
+ retention="7 days",
+ compression="zip",
+ encoding="utf-8",
+ format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {message}",
+ )
+
+ # Configure colorized standard output logging
+ logger.add(
+ sink=sys.stdout,
+ level=level,
+ format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {message}",
+ colorize=True,
+ )
diff --git a/reme_ai/core/utils/logo_utils.py b/reme_ai/core/utils/logo_utils.py
new file mode 100644
index 00000000..91c46a02
--- /dev/null
+++ b/reme_ai/core/utils/logo_utils.py
@@ -0,0 +1,84 @@
+"""Terminal branding and configuration display utilities."""
+
+import importlib.metadata
+from typing import TYPE_CHECKING
+
+from rich.console import Console, Group
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+if TYPE_CHECKING:
+ from ..schema import ServiceConfig
+
+
+def get_version(package_name: str) -> str:
+ """Return the installed version of a package or 'unknown'."""
+ try:
+ return importlib.metadata.version(package_name)
+ except importlib.metadata.PackageNotFoundError:
+ return ""
+
+
+def print_logo(service_config: "ServiceConfig"):
+ """Print a stylized ASCII logo and service metadata to the console."""
+ ascii_art = [
+ r" ██████╗ ███████╗ ███╗ ███╗ ███████╗ ",
+ r" ██╔══██╗ ██╔════╝ ████╗ ████║ ██╔════╝ ",
+ r" ██████╔╝ █████╗ ██╔████╔██║ █████╗ ",
+ r" ██╔══██╗ ██╔══╝ ██║╚██╔╝██║ ██╔══╝ ",
+ r" ██║ ██║ ███████╗ ██║ ╚═╝ ██║ ███████╗ ",
+ r" ╚═╝ ╚═╝ ╚══════╝ ╚═╝ ╚═╝ ╚══════╝ ",
+ ]
+
+ start_color = (85, 239, 196)
+ end_color = (162, 155, 254)
+
+ logo_text = Text()
+ for line in ascii_art:
+ line_len = max(1, len(line) - 1)
+ for i, char in enumerate(line):
+ # Calculate gradient shift per character
+ ratio = i / line_len
+ rgb = tuple(int(s + (e - s) * ratio) for s, e in zip(start_color, end_color))
+ logo_text.append(char, style=f"bold rgb({rgb[0]},{rgb[1]},{rgb[2]})")
+ logo_text.append("\n")
+
+ # Layout configuration info
+ info_table = Table.grid(padding=(0, 1))
+ info_table.add_column(style="bold", justify="center")
+ info_table.add_column(style="bold cyan", justify="left")
+ info_table.add_column(style="white", justify="left")
+
+ # Add core service info
+ info_table.add_row("📦", "Backend:", service_config.backend)
+
+ match service_config.backend:
+ case "http":
+ host, port = service_config.http.host, service_config.http.port
+ info_table.add_row("🔗", "URL:", f"http://{host}:{port}")
+ info_table.add_row("📚", "FastAPI:", Text(get_version("fastapi"), style="dim"))
+ case "mcp":
+ mcp = service_config.mcp
+ transport = mcp.transport if mcp.transport else "stdio"
+ info_table.add_row("🚌", "Transport:", transport)
+ if transport != "stdio":
+ url = f"http://{mcp.host}:{mcp.port}"
+ if transport == "sse":
+ url += "/sse"
+ info_table.add_row("🔗", "URL:", url)
+ info_table.add_row("📚", "FastMCP:", Text(get_version("fastmcp"), style="dim"))
+
+ info_table.add_row("🚀", "ReMe:", Text(get_version("reme-ai"), style="dim"))
+
+ # Render layout within a panel
+ panel = Panel(
+ Group(logo_text, info_table),
+ title=service_config.app_name,
+ title_align="left",
+ border_style="dim",
+ padding=(1, 4),
+ expand=False,
+ )
+
+ Console().print(Group("\n", panel, "\n"))
diff --git a/reme_ai/core/utils/mcp_client.py b/reme_ai/core/utils/mcp_client.py
new file mode 100644
index 00000000..4c13be9f
--- /dev/null
+++ b/reme_ai/core/utils/mcp_client.py
@@ -0,0 +1,121 @@
+"""Module for managing Model Context Protocol (MCP) server connections."""
+
+import os
+import re
+from contextlib import asynccontextmanager
+from typing import Any
+
+from mcp import ClientSession, StdioServerParameters, Tool
+from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.types import CallToolResult, TextContent
+
+from ..schema import ToolCall
+
+
+class MCPClient:
+ """A client manager for handling multiple MCP transport protocols."""
+
+ def __init__(self, config: dict):
+ """Initialize the client with server configuration."""
+ self.config = config
+
+ @staticmethod
+ def _infer_transport_type(cfg: dict[str, Any]) -> str:
+ """Infer the transport type based on configuration fields."""
+ if "command" in cfg:
+ return "stdio"
+
+ if "url" in cfg:
+ url = cfg["url"].lower()
+ if url.endswith("/sse") or "sse" in url:
+ return "sse"
+ return "streamable-http"
+
+ raise ValueError(f"Could not infer transport type for: {cfg}")
+
+ def _replace_env_vars(self, data: str | dict | list) -> Any:
+ """Replace environment variable placeholders in configuration."""
+ if isinstance(data, str):
+ return re.sub(r"\$\{(\w+)\}", lambda m: os.getenv(m.group(1), m.group(0)), data)
+ if isinstance(data, dict):
+ return {k: self._replace_env_vars(v) for k, v in data.items()}
+ if isinstance(data, list):
+ return [self._replace_env_vars(i) for i in data]
+ return data
+
+ @asynccontextmanager
+ async def _get_transport(self, cfg: dict[str, Any]):
+ """Context manager to yield the appropriate MCP transport."""
+ # Pop 'type' if present, otherwise infer it
+ t_type = cfg.pop("type", None) or self._infer_transport_type(cfg)
+
+ try:
+ if t_type == "stdio":
+ params = StdioServerParameters(
+ command=cfg["command"],
+ args=cfg.get("args", []),
+ env=cfg.get("env", None),
+ )
+ async with stdio_client(params) as transport:
+ yield transport
+ elif t_type == "sse":
+ async with sse_client(**cfg) as transport:
+ yield transport
+ elif t_type == "streamable-http":
+ async with streamablehttp_client(**cfg) as transport:
+ yield transport
+ else:
+ raise NotImplementedError(f"Unsupported transport: {t_type}")
+ finally:
+ pass # Ensure proper cleanup
+
+ @asynccontextmanager
+ async def connect_to_server(self, server_name: str):
+ """Establish a session with the specified MCP server."""
+ server_config = self.config.get("mcpServers", {}).get(server_name)
+ if not server_config:
+ raise ValueError(f"Config for '{server_name}' not found.")
+
+ # Process environment variables and transport selection
+ cfg = self._replace_env_vars(server_config)
+
+ async with self._get_transport(cfg) as (read, write):
+ async with ClientSession(read, write) as session:
+ await session.initialize()
+ yield session
+
+ async def list_tools(self, server_name: str) -> list[Tool]:
+ """Retrieve available tools from a specific server."""
+ async with self.connect_to_server(server_name) as session:
+ result = await session.list_tools()
+ return result.tools
+
+ async def list_tool_calls(self, server_name: str, return_dict: bool = True) -> list[dict | ToolCall]:
+ """Retrieve available tools from a specific server."""
+ tools = await self.list_tools(server_name)
+ tool_calls: list[ToolCall] = [ToolCall.from_mcp_tool(tool) for tool in tools]
+ if return_dict:
+ return [tool_call.simple_input_dump() for tool_call in tool_calls]
+
+ return tool_calls
+
+ async def call_tool(
+ self,
+ server_name: str,
+ tool_name: str,
+ arguments: dict[str, Any],
+ parse_text_result: bool = False,
+ ) -> CallToolResult | str:
+ """Execute a tool on a specific server."""
+ async with self.connect_to_server(server_name) as session:
+ tool_results: CallToolResult = await session.call_tool(tool_name, arguments)
+ if not parse_text_result:
+ return tool_results
+
+ text_result = []
+ for block in tool_results.content:
+ if isinstance(block, TextContent):
+ text_result.append(block.text)
+ return "\n".join(text_result)
diff --git a/reme_ai/core/utils/pydantic_config_parser.py b/reme_ai/core/utils/pydantic_config_parser.py
new file mode 100644
index 00000000..2f4f2ae0
--- /dev/null
+++ b/reme_ai/core/utils/pydantic_config_parser.py
@@ -0,0 +1,206 @@
+"""Parser for Pydantic config models with YAML and CLI argument support."""
+
+import inspect
+import json
+from pathlib import Path
+from typing import Any, TypeVar
+
+import yaml
+from loguru import logger
+from pydantic import BaseModel
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class PydanticConfigParser:
+ """Parser that loads and merges Pydantic configs from YAML files and CLI args."""
+
+ def __init__(self, config_class: type[T], default_config: str = "default"):
+ """Initialize parser with a Pydantic config class.
+
+ Args:
+ config_class: Pydantic BaseModel class to validate configs against.
+ default_config: Default config file name to use if not specified in args.
+ """
+ self.config_class = config_class
+ self.default_config = default_config
+ self.config_dict: dict = {}
+
+ def _deep_merge(self, base_dict: dict, update_dict: dict) -> dict:
+ """Recursively merge two dictionaries."""
+ result = base_dict.copy()
+ for key, value in update_dict.items():
+ if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+ result[key] = self._deep_merge(result[key], value)
+ else:
+ result[key] = value
+ return result
+
+ @staticmethod
+ def _convert_value(value_str: str) -> Any:
+ """Convert string value to appropriate Python type."""
+ value_str = value_str.strip()
+ lower_str = value_str.lower()
+
+ # Boolean and None conversion
+ if lower_str in ("true", "false"):
+ return lower_str == "true"
+ if lower_str in ("none", "null"):
+ return None
+
+ # Numeric conversion
+ if "e" in lower_str or "." in value_str:
+ try:
+ return float(value_str)
+ except ValueError:
+ pass
+ else:
+ try:
+ return int(value_str)
+ except ValueError:
+ pass
+
+ # JSON conversion for complex types
+ try:
+ return json.loads(value_str)
+ except (json.JSONDecodeError, ValueError):
+ return value_str
+
+ @staticmethod
+ def load_from_yaml(yaml_path: str | Path) -> dict:
+ """Load configuration from YAML file.
+
+ Args:
+ yaml_path: Path to YAML configuration file.
+
+ Returns:
+ Dictionary containing configuration data.
+
+ Raises:
+ FileNotFoundError: If YAML file does not exist.
+ """
+ if isinstance(yaml_path, str):
+ yaml_path = Path(yaml_path)
+
+ if not yaml_path.exists():
+ raise FileNotFoundError(f"Configuration file does not exist: {yaml_path}")
+
+ with yaml_path.open(encoding="utf-8") as f:
+ return yaml.safe_load(f) or {}
+
+ def merge_configs(self, *config_dicts: dict) -> dict:
+ """Merge multiple config dictionaries in order.
+
+ Args:
+ *config_dicts: Variable number of config dictionaries to merge.
+
+ Returns:
+ Merged configuration dictionary.
+ """
+ result = {}
+ for config_dict in config_dicts:
+ result = self._deep_merge(result, config_dict)
+ return result
+
+ def parse_dot_notation(self, dot_list: list[str]) -> dict:
+ """Parse dot notation strings into nested dictionary.
+
+ Args:
+ dot_list: List of strings in format "key.subkey=value".
+
+ Returns:
+ Nested dictionary representation of dot notation.
+ """
+ config_dict = {}
+ for item in dot_list:
+ if "=" not in item:
+ continue
+
+ key_path, value_str = item.split("=", 1)
+ keys = key_path.split(".")
+
+ # Build nested dictionary
+ current = config_dict
+ for key in keys[:-1]:
+ current = current.setdefault(key, {})
+ current[keys[-1]] = self._convert_value(value_str)
+
+ return config_dict
+
+ def _find_config_path(self, config_name: str) -> Path:
+ """Find config file path, trying parser directory first then current directory."""
+ if not config_name.endswith(".yaml"):
+ config_name += ".yaml"
+
+ # Try parser class directory first
+ config_path = Path(inspect.getfile(self.__class__)).parent / config_name
+ if config_path.exists():
+ logger.info(f"load config={config_path}")
+ return config_path
+
+ # Try current directory
+ logger.warning(f"config={config_path} not found, try {config_name}")
+ config_path = Path(config_name)
+ if not config_path.exists():
+ raise FileNotFoundError(f"config={config_path} not found")
+ return config_path
+
+ def parse_args(self, *args: str) -> T:
+ """Parse CLI arguments and load configs from YAML files.
+
+ Args:
+ *args: CLI arguments in format "key=value" or "config=file.yaml".
+
+ Returns:
+ Validated Pydantic config instance.
+
+ Raises:
+ ValueError: If no config file is specified.
+ FileNotFoundError: If specified config file does not exist.
+ """
+ configs_to_merge = [self.config_class().model_dump()]
+
+ # Separate config file path from other arguments
+ config = ""
+ filter_args = []
+ for arg in args:
+ if "=" not in arg:
+ continue
+ arg = arg.lstrip("-")
+ if arg.startswith(("c=", "config=")):
+ config = arg.split("=", 1)[1]
+ else:
+ filter_args.append(arg)
+
+ # Use default config if not specified
+ config = config or self.default_config
+
+ # Load each config file
+ for single_config in (c.strip() for c in config.split(",") if c.strip()):
+ config_path = self._find_config_path(single_config)
+ configs_to_merge.append(self.load_from_yaml(config_path))
+
+ # Apply CLI overrides
+ if filter_args:
+ configs_to_merge.append(self.parse_dot_notation(filter_args))
+
+ # Merge all configs and validate
+ self.config_dict = self.merge_configs(*configs_to_merge)
+ return self.config_class.model_validate(self.config_dict)
+
+ def update_config(self, **kwargs) -> T:
+ """Update current config with new values using kwargs.
+
+ Args:
+ **kwargs: Key-value pairs where __ in keys represents nested levels.
+
+ Returns:
+ Updated and validated Pydantic config instance.
+ """
+ # Convert kwargs to dot notation and parse
+ dot_list = [f"{key.replace('__', '.')}={value}" for key, value in kwargs.items()]
+ override_config = self.parse_dot_notation(dot_list)
+
+ # Merge with existing config
+ final_config = self.merge_configs(self.config_dict, override_config)
+ return self.config_class.model_validate(final_config)
diff --git a/reme_ai/core/utils/pydantic_utils.py b/reme_ai/core/utils/pydantic_utils.py
new file mode 100644
index 00000000..06b1f4fe
--- /dev/null
+++ b/reme_ai/core/utils/pydantic_utils.py
@@ -0,0 +1,66 @@
+"""
+Utility module for dynamic Pydantic model generation based on schema definitions.
+"""
+
+from typing import Any, Literal
+
+from pydantic import create_model, Field
+
+from . import snake_to_camel
+from ..enumeration import JsonSchemaEnum
+from ..schema import ToolAttr, Request
+
+TYPE_MAPPING = {str(t): t.value for t in JsonSchemaEnum}
+
+
+def create_pydantic_model(name: str, parameters: ToolAttr | None = None) -> type[Request]:
+ """
+ Recursively generates a Pydantic model from a ToolAttr schema definition.
+ """
+ fields = {}
+
+ if not parameters or not parameters.properties:
+ return create_model(f"{snake_to_camel(name)}Model", __base__=Request)
+
+ for field_name, attr in parameters.properties.items():
+ # 1. Determine the base field type
+ if attr.type == "object" and attr.properties:
+ # Handle nested objects recursively
+ field_type = create_pydantic_model(field_name, attr)
+
+ elif attr.type == "array" and attr.items:
+ # Handle array/list types
+ if isinstance(attr.items, ToolAttr):
+ if attr.items.type == "object":
+ inner_type = create_pydantic_model(f"{field_name}_item", attr.items)
+ else:
+ inner_type = TYPE_MAPPING.get(attr.items.type, Any)
+ field_type = list[inner_type]
+ else:
+ # Fallback for simple dictionary item definitions
+ field_type = list[Any]
+
+ else:
+ # Handle primitive types
+ field_type = TYPE_MAPPING.get(attr.type, Any)
+
+ # 2. Handle enumeration constraints
+ if attr.enum:
+ # Dynamically create a Literal type from the enum list
+ field_type = Literal[tuple(attr.enum)] # type: ignore
+
+ # 3. Determine requirement status and default values
+ is_required = False
+ if parameters.required and field_name in parameters.required:
+ is_required = True
+
+ # 4. Construct Field metadata
+ field_info = Field(default=... if is_required else None, description=attr.description)
+
+ if not is_required:
+ field_type = field_type | None
+
+ fields[field_name] = (field_type, field_info)
+
+ # Dynamically construct the final Pydantic model class
+ return create_model(f"{snake_to_camel(name)}Model", **fields, __base__=Request)
diff --git a/reme_ai/core/vector_store/__init__.py b/reme_ai/core/vector_store/__init__.py
new file mode 100644
index 00000000..79500294
--- /dev/null
+++ b/reme_ai/core/vector_store/__init__.py
@@ -0,0 +1,17 @@
+"""vector store"""
+
+from .base_vector_store import BaseVectorStore
+from .chroma_vector_store import ChromaVectorStore
+from .es_vector_store import ESVectorStore
+from .local_vector_store import LocalVectorStore
+from .pgvector_store import PGVectorStore
+from .qdrant_vector_store import QdrantVectorStore
+
+__all__ = [
+ "BaseVectorStore",
+ "ChromaVectorStore",
+ "ESVectorStore",
+ "LocalVectorStore",
+ "PGVectorStore",
+ "QdrantVectorStore",
+]
diff --git a/reme_ai/core/vector_store/base_vector_store.py b/reme_ai/core/vector_store/base_vector_store.py
new file mode 100644
index 00000000..e64b6dbf
--- /dev/null
+++ b/reme_ai/core/vector_store/base_vector_store.py
@@ -0,0 +1,92 @@
+"""Base vector store interface for managing vector embeddings and similarity search."""
+
+import asyncio
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from functools import partial
+
+from reme_ai.core.context import C
+from reme_ai.core.embedding import BaseEmbeddingModel
+from reme_ai.core.schema import VectorNode
+
+
+class BaseVectorStore(ABC):
+ """Abstract base class defining the interface for vector storage and retrieval."""
+
+ def __init__(
+ self,
+ collection_name: str,
+ embedding_model: BaseEmbeddingModel,
+ **kwargs,
+ ):
+ """Initialize the vector store with a collection name and an embedding model."""
+ if embedding_model is None:
+ raise ValueError("embedding_model is required")
+ self.collection_name: str = collection_name
+ self.embedding_model: BaseEmbeddingModel = embedding_model
+ self.kwargs: dict = kwargs
+
+ @staticmethod
+ async def _run_sync_in_executor(sync_func: Callable, *args, **kwargs):
+ """Run a synchronous function in the context-defined thread pool executor."""
+ loop = asyncio.get_running_loop()
+ return await loop.run_in_executor(C.thread_pool, partial(sync_func, *args, **kwargs))
+
+ async def get_node_embedding(self, node: VectorNode) -> VectorNode:
+ """Generate and assign embedding for a single vector node."""
+ return await self.embedding_model.get_node_embedding(node)
+
+ async def get_node_embeddings(self, nodes: list[VectorNode]) -> list[VectorNode]:
+ """Generate and assign embeddings for multiple vector nodes."""
+ return await self.embedding_model.get_node_embeddings(nodes)
+
+ async def get_embedding(self, query: str) -> list[float]:
+ """Convert a single text query into vector embedding using the configured model."""
+ return await self.embedding_model.get_embedding(query)
+
+ async def get_embeddings(self, queries: list[str]) -> list[list[float]]:
+ """Convert multiple text queries into vector embeddings using the configured model."""
+ return await self.embedding_model.get_embeddings(queries)
+
+ @abstractmethod
+ async def list_collections(self) -> list[str]:
+ """Retrieve a list of all existing collection names in the store."""
+
+ @abstractmethod
+ async def create_collection(self, collection_name: str, **kwargs) -> None:
+ """Create a new vector collection with the specified name and configuration."""
+
+ @abstractmethod
+ async def delete_collection(self, collection_name: str, **kwargs) -> None:
+ """Permanently remove a collection from the vector store."""
+
+ @abstractmethod
+ async def copy_collection(self, collection_name: str, **kwargs) -> None:
+ """Duplicate the current collection to a new one with the given name."""
+
+ @abstractmethod
+ async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs) -> None:
+ """Add one or more vector nodes into the current collection."""
+
+ @abstractmethod
+ async def search(self, query: str, limit: int = 5, filters: dict | None = None, **kwargs) -> list[VectorNode]:
+ """Find the most similar vector nodes based on a text query."""
+
+ @abstractmethod
+ async def delete(self, vector_ids: str | list[str], **kwargs) -> None:
+ """Remove specific vectors from the collection using their identifiers."""
+
+ @abstractmethod
+ async def update(self, nodes: VectorNode | list[VectorNode], **kwargs) -> None:
+ """Update the data or metadata of existing vectors in the collection."""
+
+ @abstractmethod
+ async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
+ """Fetch specific vector nodes from the collection by their IDs."""
+
+ @abstractmethod
+ async def list(self, filters: dict | None = None, limit: int | None = None) -> list[VectorNode]:
+ """Retrieve vectors from the collection that match the given filters."""
+
+ async def close(self) -> None:
+ """Release resources and close active connections to the vector store."""
diff --git a/reme_ai/core/vector_store/chroma_vector_store.py b/reme_ai/core/vector_store/chroma_vector_store.py
new file mode 100644
index 00000000..24b88d36
--- /dev/null
+++ b/reme_ai/core/vector_store/chroma_vector_store.py
@@ -0,0 +1,397 @@
+"""ChromaDB vector store implementation for the ReMe framework."""
+
+from typing import Any
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+_CHROMADB_IMPORT_ERROR = None
+
+try:
+ import chromadb
+ from chromadb.config import Settings
+except ImportError as e:
+ _CHROMADB_IMPORT_ERROR = e
+ chromadb = None
+ Settings = None
+
+
+@C.register_vector_store("chroma")
+class ChromaVectorStore(BaseVectorStore):
+ """ChromaDB-based vector store implementation for local or remote storage."""
+
+ def __init__(
+ self,
+ collection_name: str,
+ embedding_model: BaseEmbeddingModel,
+ client: chromadb.ClientAPI | None = None,
+ host: str | None = None,
+ port: int | None = None,
+ path: str | None = None,
+ api_key: str | None = None,
+ tenant: str | None = None,
+ database: str | None = None,
+ **kwargs,
+ ):
+ """Initialize the ChromaDB vector store with the provided configuration."""
+ if _CHROMADB_IMPORT_ERROR is not None:
+ raise ImportError(
+ "ChromaDB requires extra dependencies. Install with `pip install chromadb`",
+ ) from _CHROMADB_IMPORT_ERROR
+
+ super().__init__(
+ collection_name=collection_name,
+ embedding_model=embedding_model,
+ **kwargs,
+ )
+
+ self.client: chromadb.ClientAPI
+ self.collection: chromadb.Collection
+
+ if client:
+ self.client = client
+ elif api_key and tenant:
+ logger.info("Initializing ChromaDB Cloud client")
+ self.client = chromadb.CloudClient(
+ api_key=api_key,
+ tenant=tenant,
+ database=database or "default",
+ )
+ elif host and port:
+ logger.info(f"Initializing ChromaDB HTTP client at {host}:{port}")
+ self.client = chromadb.HttpClient(host=host, port=port)
+ else:
+ if path is None:
+ path = "./chroma_db"
+ logger.info(f"Initializing local ChromaDB at {path}")
+ self.client = chromadb.PersistentClient(
+ path=path,
+ settings=Settings(anonymized_telemetry=False),
+ )
+
+ self.collection = self.client.get_or_create_collection(
+ name=collection_name,
+ metadata={"hnsw:space": "cosine"},
+ )
+
+ @staticmethod
+ def _parse_results(
+ results: dict,
+ include_score: bool = False,
+ ) -> list[VectorNode]:
+ """Convert ChromaDB query results into a list of VectorNode objects."""
+ nodes = []
+
+ ids = results.get("ids", [])
+ documents = results.get("documents", [])
+ metadatas = results.get("metadatas", [])
+ embeddings = results.get("embeddings") if results.get("embeddings") is not None else []
+ distances = results.get("distances") if results.get("distances") is not None else []
+
+ if ids and isinstance(ids[0], list):
+ ids = ids[0] if ids else []
+ documents = documents[0] if documents else []
+ metadatas = metadatas[0] if metadatas else []
+ embeddings = embeddings[0] if embeddings and len(embeddings) > 0 else []
+ distances = distances[0] if distances and len(distances) > 0 else []
+
+ for i, vector_id in enumerate(ids):
+ metadata = metadatas[i] if i < len(metadatas) and metadatas[i] else {}
+
+ if include_score and distances and i < len(distances):
+ metadata["_score"] = 1.0 - distances[i]
+
+ node = VectorNode(
+ vector_id=vector_id,
+ content=documents[i] if i < len(documents) and documents[i] else "",
+ vector=embeddings[i] if len(embeddings) > i else None,
+ metadata=metadata,
+ )
+ nodes.append(node)
+
+ return nodes
+
+ @staticmethod
+ def _generate_where_clause(filters: dict | None) -> dict | None:
+ """Convert the universal filter format to a ChromaDB-compatible where clause."""
+ if not filters:
+ return None
+
+ def convert_condition(k: str, v: Any) -> dict | None:
+ """Convert a single filter condition to ChromaDB operator format."""
+ if v == "*":
+ return None
+ if isinstance(v, dict):
+ chroma_condition = {}
+ for op, val in v.items():
+ mapping = {
+ "eq": "$eq",
+ "ne": "$ne",
+ "gt": "$gt",
+ "gte": "$gte",
+ "lt": "$lt",
+ "lte": "$lte",
+ "in": "$in",
+ "nin": "$nin",
+ }
+ chroma_op = mapping.get(op, "$eq")
+ chroma_condition[k] = {chroma_op: val}
+ return chroma_condition
+ if isinstance(v, list):
+ return {k: {"$in": v}}
+ return {k: {"$eq": v}}
+
+ processed_filters = []
+
+ for key, value in filters.items():
+ if key == "$or":
+ or_conditions = []
+ for condition in value:
+ or_condition = {}
+ for sub_key, sub_value in condition.items():
+ converted = convert_condition(sub_key, sub_value)
+ if converted:
+ or_condition.update(converted)
+ if or_condition:
+ or_conditions.append(or_condition)
+ if len(or_conditions) > 1:
+ processed_filters.append({"$or": or_conditions})
+ elif len(or_conditions) == 1:
+ processed_filters.append(or_conditions[0])
+
+ elif key == "$and":
+ for condition in value:
+ for sub_key, sub_value in condition.items():
+ converted = convert_condition(sub_key, sub_value)
+ if converted:
+ processed_filters.append(converted)
+ elif key == "$not":
+ continue
+ else:
+ converted = convert_condition(key, value)
+ if converted:
+ processed_filters.append(converted)
+
+ if not processed_filters:
+ return None
+ return processed_filters[0] if len(processed_filters) == 1 else {"$and": processed_filters}
+
+ async def list_collections(self) -> list[str]:
+ """Retrieve a list of all existing collection names."""
+
+ def _list():
+ return [col.name for col in self.client.list_collections()]
+
+ return await self._run_sync_in_executor(_list)
+
+ async def create_collection(self, collection_name: str, **kwargs):
+ """Create a new collection with specified distance metrics and metadata."""
+
+ def _create():
+ distance_metric = kwargs.get("distance_metric", "cosine")
+ metadata = kwargs.get("metadata", {})
+ metadata["hnsw:space"] = distance_metric
+ return self.client.get_or_create_collection(name=collection_name, metadata=metadata)
+
+ new_collection = await self._run_sync_in_executor(_create)
+ if collection_name == self.collection_name:
+ self.collection = new_collection
+ logger.info(f"Created collection {collection_name}")
+
+ async def delete_collection(self, collection_name: str, **kwargs):
+ """Delete a specified collection from the database."""
+
+ def _delete():
+ try:
+ self.client.delete_collection(name=collection_name)
+ return True
+ except Exception as e:
+ logger.warning(f"Failed to delete collection {collection_name}: {e}")
+ return False
+
+ deleted = await self._run_sync_in_executor(_delete)
+ if deleted and collection_name == self.collection_name:
+ self.collection = None
+ logger.info(f"Deleted collection {collection_name}")
+
+ async def copy_collection(self, collection_name: str, **kwargs):
+ """Copy all data from the current collection to a new collection."""
+
+ def _copy():
+ source_data = self.collection.get(include=["documents", "metadatas", "embeddings"])
+ if not source_data["ids"]:
+ logger.warning(f"Source collection {self.collection_name} is empty")
+ return
+
+ target_collection = self.client.get_or_create_collection(
+ name=collection_name,
+ metadata={"hnsw:space": "cosine"},
+ )
+ target_collection.add(
+ ids=source_data["ids"],
+ documents=source_data["documents"],
+ metadatas=source_data["metadatas"],
+ embeddings=source_data["embeddings"],
+ )
+
+ await self._run_sync_in_executor(_copy)
+ logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+ async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
+ """Insert vector nodes into the current collection in batches."""
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+ if not nodes:
+ return
+
+ # Batch generate embeddings for nodes that need them
+ nodes_without_vectors = [node for node in nodes if node.vector is None]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ # Create a mapping for quick lookup
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+ else:
+ nodes_to_insert = nodes
+
+ batch_size = kwargs.get("batch_size", 100)
+
+ def _insert_batch(batch_nodes: list[VectorNode]):
+ self.collection.add(
+ ids=[n.vector_id for n in batch_nodes],
+ documents=[n.content for n in batch_nodes],
+ embeddings=[n.vector for n in batch_nodes],
+ metadatas=[n.metadata for n in batch_nodes],
+ )
+
+ for i in range(0, len(nodes_to_insert), batch_size):
+ await self._run_sync_in_executor(_insert_batch, nodes_to_insert[i : i + batch_size])
+ logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}")
+
+ async def search(
+ self,
+ query: str,
+ limit: int = 5,
+ filters: dict | None = None,
+ **kwargs,
+ ) -> list[VectorNode]:
+ """Search for the most similar vector nodes based on a text query."""
+ query_vector = await self.get_embedding(query)
+ where_clause = self._generate_where_clause(filters)
+ include_embeddings = kwargs.get("include_embeddings", False)
+
+ def _search():
+ include: list = ["documents", "metadatas", "distances"]
+ if include_embeddings:
+ include.append("embeddings")
+ return self.collection.query(
+ query_embeddings=[query_vector],
+ n_results=limit,
+ where=where_clause,
+ include=include,
+ )
+
+ results = await self._run_sync_in_executor(_search)
+ nodes = self._parse_results(results, include_score=True)
+
+ score_threshold = kwargs.get("score_threshold")
+ if score_threshold is not None:
+ nodes = [n for n in nodes if n.metadata.get("_score", 0) >= score_threshold]
+ return nodes
+
+ async def delete(self, vector_ids: str | list[str], **kwargs):
+ """Delete specific vector nodes by their IDs."""
+ if isinstance(vector_ids, str):
+ vector_ids = [vector_ids]
+ if not vector_ids:
+ return
+
+ def _delete():
+ self.collection.delete(ids=vector_ids)
+
+ await self._run_sync_in_executor(_delete)
+ logger.info(f"Deleted {len(vector_ids)} nodes from {self.collection_name}")
+
+ async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
+ """Update existing vector nodes with new content or metadata."""
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+ if not nodes:
+ return
+
+ # Batch generate embeddings for nodes that need them
+ nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ # Create a mapping for quick lookup
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+ else:
+ nodes_to_update = nodes
+
+ def _update():
+ self.collection.upsert(
+ ids=[n.vector_id for n in nodes_to_update],
+ documents=[n.content for n in nodes_to_update],
+ embeddings=[n.vector for n in nodes_to_update if n.vector] or None,
+ metadatas=[n.metadata for n in nodes_to_update],
+ )
+
+ await self._run_sync_in_executor(_update)
+ logger.info(f"Updated {len(nodes_to_update)} nodes in {self.collection_name}")
+
+ async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None:
+ """Fetch vector nodes by their IDs from the collection."""
+ is_single = isinstance(vector_ids, str)
+ ids = [vector_ids] if is_single else vector_ids
+
+ def _get():
+ return self.collection.get(ids=ids, include=["documents", "metadatas", "embeddings"])
+
+ results = await self._run_sync_in_executor(_get)
+ nodes = self._parse_results(results)
+ return nodes[0] if is_single and nodes else (nodes if not is_single else None)
+
+ async def list(
+ self,
+ filters: dict | None = None,
+ limit: int | None = None,
+ ) -> list[VectorNode]:
+ """List vector nodes matching optional metadata filters."""
+ where_clause = self._generate_where_clause(filters)
+
+ def _list():
+ return self.collection.get(
+ where=where_clause,
+ limit=limit,
+ include=["documents", "metadatas", "embeddings"],
+ )
+
+ results = await self._run_sync_in_executor(_list)
+ return self._parse_results(results)
+
+ async def count(self) -> int:
+ """Return the total number of vectors in the current collection."""
+ return await self._run_sync_in_executor(self.collection.count)
+
+ async def reset(self):
+ """Reset the current collection by clearing all its data."""
+ logger.warning(f"Resetting collection {self.collection_name}...")
+ await self.delete_collection(self.collection_name)
+
+ def _recreate():
+ self.collection = self.client.get_or_create_collection(
+ name=self.collection_name,
+ metadata={"hnsw:space": "cosine"},
+ )
+
+ await self._run_sync_in_executor(_recreate)
+ logger.info(f"Collection {self.collection_name} has been reset")
+
+ async def close(self):
+ """Close the vector store and log the shutdown process."""
+ logger.info(f"ChromaDB vector store for collection {self.collection_name} closed")
diff --git a/reme_ai/core/vector_store/es_vector_store.py b/reme_ai/core/vector_store/es_vector_store.py
new file mode 100644
index 00000000..06dcafa5
--- /dev/null
+++ b/reme_ai/core/vector_store/es_vector_store.py
@@ -0,0 +1,458 @@
+"""Elasticsearch vector store implementation for ReMe.
+
+This module provides an Elasticsearch-based vector store that implements the BaseVectorStore
+interface for high-performance dense vector storage and retrieval.
+"""
+
+from typing import Any
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+_ELASTICSEARCH_IMPORT_ERROR = None
+
+try:
+ from elasticsearch import AsyncElasticsearch
+ from elasticsearch.helpers import async_bulk
+except ImportError as e:
+ _ELASTICSEARCH_IMPORT_ERROR = e
+ AsyncElasticsearch = None
+ async_bulk = None
+
+
+@C.register_vector_store("es")
+class ESVectorStore(BaseVectorStore):
+ """Elasticsearch-based vector store for dense vector storage and kNN search."""
+
+ def __init__(
+ self,
+ collection_name: str,
+ embedding_model: BaseEmbeddingModel,
+ hosts: str | list[str] | None = None,
+ basic_auth: tuple[str, str] | None = None,
+ cloud_id: str | None = None,
+ api_key: str | None = None,
+ verify_certs: bool = True,
+ headers: dict[str, str] | None = None,
+ **kwargs,
+ ):
+ """Initialize the Elasticsearch client and vector store configuration.
+
+ Args:
+ collection_name: Name of the Elasticsearch index (converted to lowercase).
+ embedding_model: Model instance used to generate vector embeddings.
+ hosts: Connection host(s) for the Elasticsearch cluster.
+ basic_auth: Credentials for basic authentication.
+ cloud_id: Deployment ID for Elastic Cloud.
+ api_key: API key for authentication.
+ verify_certs: Enable or disable SSL certificate verification.
+ headers: Custom HTTP headers for requests.
+ **kwargs: Additional configuration passed to the base class.
+ """
+ if _ELASTICSEARCH_IMPORT_ERROR is not None:
+ raise ImportError(
+ "Elasticsearch requires extra dependencies. Install with `pip install elasticsearch`",
+ ) from _ELASTICSEARCH_IMPORT_ERROR
+
+ # Elasticsearch requires lowercase index names
+ collection_name = collection_name.lower()
+
+ super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs)
+
+ # Initialize AsyncElasticsearch client
+ self.client = AsyncElasticsearch(
+ hosts=hosts,
+ cloud_id=cloud_id,
+ api_key=api_key,
+ basic_auth=basic_auth,
+ verify_certs=verify_certs,
+ headers=headers or {},
+ )
+
+ async def list_collections(self) -> list[str]:
+ """List all available index names in the Elasticsearch cluster."""
+ aliases = await self.client.indices.get_alias()
+ return list(aliases.keys())
+
+ async def create_collection(self, collection_name: str, **kwargs):
+ """Create a new index with dense vector mappings for kNN search.
+
+ Args:
+ collection_name: Name of the index to create.
+ **kwargs: Settings like dimensions, similarity, shards, and replicas.
+ """
+ collection_name = collection_name.lower()
+
+ if await self.client.indices.exists(index=collection_name):
+ return
+
+ dimensions = kwargs.get("dimensions", self.embedding_model.dimensions)
+ similarity = kwargs.get("similarity", "cosine")
+ number_of_shards = kwargs.get("number_of_shards", 5)
+ number_of_replicas = kwargs.get("number_of_replicas", 1)
+ refresh_interval = kwargs.get("refresh_interval", "1s")
+
+ index_settings = {
+ "settings": {
+ "index": {
+ "number_of_replicas": number_of_replicas,
+ "number_of_shards": number_of_shards,
+ "refresh_interval": refresh_interval,
+ },
+ },
+ "mappings": {
+ "properties": {
+ "vector_id": {"type": "keyword"},
+ "content": {"type": "text"},
+ "vector": {
+ "type": "dense_vector",
+ "dims": dimensions,
+ "index": True,
+ "similarity": similarity,
+ },
+ "metadata": {"type": "object", "enabled": True},
+ },
+ },
+ }
+
+ if not await self.client.indices.exists(index=collection_name):
+ await self.client.indices.create(index=collection_name, body=index_settings)
+ logger.info(f"Created index {collection_name} with dimensions={dimensions}")
+ else:
+ logger.info(f"Index {collection_name} already exists")
+
+ async def delete_collection(self, collection_name: str, **kwargs):
+ """Permanently delete an Elasticsearch index.
+
+ Args:
+ collection_name: Name of the index to delete.
+ **kwargs: Additional parameters for the deletion request.
+ """
+ collection_name = collection_name.lower()
+
+ if await self.client.indices.exists(index=collection_name):
+ await self.client.indices.delete(index=collection_name)
+ logger.info(f"Deleted index {collection_name}")
+ else:
+ logger.warning(f"Index {collection_name} does not exist")
+
+ async def copy_collection(self, collection_name: str, **kwargs):
+ """Reindex the current collection into a new index with identical mappings.
+
+ Args:
+ collection_name: Name of the destination index.
+ **kwargs: Additional parameters for the reindexing process.
+ """
+ collection_name = collection_name.lower()
+
+ current_index = await self.client.indices.get(index=self.collection_name)
+ current_settings = current_index[self.collection_name]
+
+ settings_to_copy = current_settings.get("settings", {}).copy()
+ if "index" in settings_to_copy:
+ index_settings = settings_to_copy["index"].copy()
+ internal_keys = [
+ "uuid",
+ "creation_date",
+ "provided_name",
+ "version",
+ "store",
+ "routing",
+ "replication",
+ ]
+ for key in internal_keys:
+ index_settings.pop(key, None)
+ settings_to_copy["index"] = index_settings
+
+ await self.client.indices.create(
+ index=collection_name,
+ body={
+ "settings": settings_to_copy,
+ "mappings": current_settings.get("mappings", {}),
+ },
+ )
+
+ await self.client.reindex(
+ body={
+ "source": {"index": self.collection_name},
+ "dest": {"index": collection_name},
+ },
+ )
+
+ logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+ async def insert(self, nodes: VectorNode | list[VectorNode], refresh: bool = True, **kwargs):
+ """Insert nodes into the index, generating embeddings if missing.
+
+ Args:
+ nodes: Single or multiple VectorNode objects to index.
+ refresh: If True, makes the operation visible to search immediately.
+ **kwargs: Additional insertion options.
+ """
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+
+ nodes_without_vectors = [node for node in nodes if node.vector is None]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+ else:
+ nodes_to_insert = nodes
+
+ actions = []
+ for node in nodes_to_insert:
+ action = {
+ "_index": self.collection_name,
+ "_id": node.vector_id,
+ "_source": {
+ "vector_id": node.vector_id,
+ "content": node.content,
+ "vector": node.vector,
+ "metadata": node.metadata,
+ },
+ }
+ actions.append(action)
+
+ success, failed = await async_bulk(self.client, actions, raise_on_error=False)
+
+ if failed:
+ logger.warning(f"Failed to insert {len(failed)} documents")
+
+ logger.info(f"Inserted {success} documents into {self.collection_name}")
+
+ if refresh:
+ await self.client.indices.refresh(index=self.collection_name)
+
+ async def search(
+ self,
+ query: str,
+ limit: int = 5,
+ filters: dict | None = None,
+ **kwargs,
+ ) -> list[VectorNode]:
+ """Perform a kNN similarity search based on a text query.
+
+ Args:
+ query: The text to search for.
+ limit: Maximum number of nearest neighbors to return.
+ filters: Metadata filters for exact match or 'IN' operations.
+ **kwargs: Search parameters like num_candidates or score_threshold.
+
+ Returns:
+ List of VectorNode objects ordered by similarity.
+ """
+ query_vector = await self.get_embedding(query)
+ num_candidates = kwargs.get("num_candidates", limit * 2)
+
+ search_query: dict = {
+ "knn": {
+ "field": "vector",
+ "query_vector": query_vector,
+ "k": limit,
+ "num_candidates": num_candidates,
+ },
+ "size": limit,
+ }
+
+ if filters:
+ filter_conditions = []
+ for key, value in filters.items():
+ if isinstance(value, list):
+ filter_conditions.append({"terms": {f"metadata.{key}": value}})
+ else:
+ filter_conditions.append({"term": {f"metadata.{key}": value}})
+ search_query["knn"]["filter"] = {"bool": {"must": filter_conditions}}
+
+ response = await self.client.search(index=self.collection_name, body=search_query)
+
+ results = []
+ for hit in response["hits"]["hits"]:
+ source = hit["_source"]
+ node = VectorNode(
+ vector_id=source.get("vector_id", hit["_id"]),
+ content=source.get("content", ""),
+ vector=source.get("vector"),
+ metadata=source.get("metadata", {}),
+ )
+ node.metadata["_score"] = hit["_score"]
+ results.append(node)
+
+ return results
+
+ async def delete(self, vector_ids: str | list[str], refresh: bool = True, **kwargs):
+ """Delete specific vectors from the index by their IDs.
+
+ Args:
+ vector_ids: Single ID or list of IDs to remove.
+ refresh: If True, refreshes the index after deletion.
+ **kwargs: Additional deletion parameters.
+ """
+ if isinstance(vector_ids, str):
+ vector_ids = [vector_ids]
+
+ actions = []
+ for vector_id in vector_ids:
+ actions.append(
+ {
+ "_op_type": "delete",
+ "_index": self.collection_name,
+ "_id": vector_id,
+ },
+ )
+
+ success, failed = await async_bulk(
+ self.client,
+ actions,
+ raise_on_error=False,
+ raise_on_exception=False,
+ )
+
+ if failed:
+ logger.warning(f"Failed to delete {len(failed)} documents")
+
+ logger.info(f"Deleted {success} documents from {self.collection_name}")
+
+ if refresh:
+ await self.client.indices.refresh(index=self.collection_name)
+
+ async def update(self, nodes: VectorNode | list[VectorNode], refresh: bool = True, **kwargs):
+ """Update existing documents with new content or metadata.
+
+ Args:
+ nodes: Single or multiple VectorNode objects with updated data.
+ refresh: If True, refreshes the index after update.
+ **kwargs: Additional update parameters.
+ """
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+
+ nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+ else:
+ nodes_to_update = nodes
+
+ actions = []
+ for node in nodes_to_update:
+ doc = {
+ "vector_id": node.vector_id,
+ "content": node.content,
+ "metadata": node.metadata,
+ }
+ if node.vector is not None:
+ doc["vector"] = node.vector
+
+ actions.append(
+ {
+ "_op_type": "update",
+ "_index": self.collection_name,
+ "_id": node.vector_id,
+ "doc": doc,
+ },
+ )
+
+ success, failed = await async_bulk(
+ self.client,
+ actions,
+ raise_on_error=False,
+ raise_on_exception=False,
+ )
+
+ if failed:
+ logger.warning(f"Failed to update {len(failed)} documents")
+
+ logger.info(f"Updated {success} documents in {self.collection_name}")
+
+ if refresh:
+ await self.client.indices.refresh(index=self.collection_name)
+
+ async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
+ """Fetch documents by their IDs from the current index.
+
+ Args:
+ vector_ids: Single ID or list of IDs to retrieve.
+
+ Returns:
+ A single VectorNode or a list of VectorNodes.
+ """
+ single_result = isinstance(vector_ids, str)
+ if single_result:
+ vector_ids = [vector_ids]
+
+ response = await self.client.mget(
+ index=self.collection_name,
+ body={"ids": vector_ids},
+ )
+
+ results = []
+ for doc in response["docs"]:
+ if doc.get("found"):
+ source = doc["_source"]
+ node = VectorNode(
+ vector_id=source.get("vector_id", doc["_id"]),
+ content=source.get("content", ""),
+ vector=source.get("vector"),
+ metadata=source.get("metadata", {}),
+ )
+ results.append(node)
+ else:
+ logger.warning(f"Document with ID {doc['_id']} not found")
+
+ return results[0] if single_result and results else results
+
+ async def list(
+ self,
+ filters: dict | None = None,
+ limit: int | None = None,
+ ) -> list[VectorNode]:
+ """Retrieve a list of nodes filtered by metadata or limit.
+
+ Args:
+ filters: Optional metadata filtering criteria.
+ limit: Maximum number of nodes to return.
+
+ Returns:
+ A list of matching VectorNode objects.
+ """
+ query: dict[str, Any] = {"query": {"match_all": {}}}
+
+ if filters:
+ filter_conditions = []
+ for key, value in filters.items():
+ if isinstance(value, list):
+ filter_conditions.append({"terms": {f"metadata.{key}": value}})
+ else:
+ filter_conditions.append({"term": {f"metadata.{key}": value}})
+ query["query"] = {"bool": {"must": filter_conditions}}
+
+ if limit:
+ query["size"] = limit
+ else:
+ query["size"] = 10000
+
+ response = await self.client.search(index=self.collection_name, body=query)
+
+ results = []
+ for hit in response["hits"]["hits"]:
+ source = hit["_source"]
+ node = VectorNode(
+ vector_id=source.get("vector_id", hit["_id"]),
+ content=source.get("content", ""),
+ vector=source.get("vector"),
+ metadata=source.get("metadata", {}),
+ )
+ results.append(node)
+
+ return results
+
+ async def close(self):
+ """Terminate the Elasticsearch client session and release resources."""
+ await self.client.close()
+ logger.info("Elasticsearch client connection closed")
diff --git a/reme_ai/core/vector_store/local_vector_store.py b/reme_ai/core/vector_store/local_vector_store.py
new file mode 100644
index 00000000..25beef8d
--- /dev/null
+++ b/reme_ai/core/vector_store/local_vector_store.py
@@ -0,0 +1,281 @@
+"""Local file system vector store implementation for ReMe."""
+
+import json
+from pathlib import Path
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+
+@C.register_vector_store("local")
+class LocalVectorStore(BaseVectorStore):
+ """Local file system-based vector store using JSON files and manual cosine similarity."""
+
+ def __init__(
+ self,
+ collection_name: str,
+ embedding_model: BaseEmbeddingModel,
+ root_path: str = "./local_vector_store",
+ **kwargs,
+ ):
+ """Initialize the local vector store with a root path and collection name."""
+ super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs)
+ self.root_path = Path(root_path)
+ self.collection_path = self.root_path / collection_name
+ self.root_path.mkdir(parents=True, exist_ok=True)
+
+ def _get_collection_path(self, collection_name: str) -> Path:
+ """Get the file system path for a specific collection."""
+ return self.root_path / collection_name
+
+ def _get_node_file_path(self, vector_id: str, collection_name: str | None = None) -> Path:
+ """Get the JSON file path for a specific vector node."""
+ col_path = self._get_collection_path(collection_name or self.collection_name)
+ return col_path / f"{vector_id}.json"
+
+ def _save_node(self, node: VectorNode, collection_name: str | None = None):
+ """Save a vector node to a JSON file on disk."""
+ file_path = self._get_node_file_path(node.vector_id, collection_name)
+ file_path.parent.mkdir(parents=True, exist_ok=True)
+
+ with open(file_path, "w", encoding="utf-8") as f:
+ json.dump(node.model_dump(), f, ensure_ascii=False, indent=2)
+
+ def _load_node(self, vector_id: str, collection_name: str | None = None) -> VectorNode | None:
+ """Load a vector node from a JSON file."""
+ file_path = self._get_node_file_path(vector_id, collection_name)
+
+ if not file_path.exists():
+ return None
+
+ with open(file_path, "r", encoding="utf-8") as f:
+ data = json.load(f)
+ return VectorNode(**data)
+
+ def _load_all_nodes(self, collection_name: str | None = None) -> list[VectorNode]:
+ """Load all vector nodes existing in a collection."""
+ col_path = self._get_collection_path(collection_name or self.collection_name)
+
+ if not col_path.exists():
+ return []
+
+ nodes = []
+ for file_path in col_path.glob("*.json"):
+ try:
+ with open(file_path, "r", encoding="utf-8") as f:
+ data = json.load(f)
+ nodes.append(VectorNode(**data))
+ except Exception as e:
+ logger.warning(f"Failed to load node from {file_path}: {e}")
+
+ return nodes
+
+ @staticmethod
+ def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
+ """Calculate the cosine similarity between two numeric vectors."""
+ if len(vec1) != len(vec2):
+ raise ValueError(f"Vectors must have same length: {len(vec1)} != {len(vec2)}")
+
+ dot_product = sum(a * b for a, b in zip(vec1, vec2))
+ magnitude1 = sum(a * a for a in vec1) ** 0.5
+ magnitude2 = sum(b * b for b in vec2) ** 0.5
+
+ if magnitude1 == 0 or magnitude2 == 0:
+ return 0.0
+
+ return dot_product / (magnitude1 * magnitude2)
+
+ @staticmethod
+ def _match_filters(node: VectorNode, filters: dict | None) -> bool:
+ """Check if a vector node matches the provided metadata filters."""
+ if not filters:
+ return True
+
+ for key, value in filters.items():
+ node_value = node.metadata.get(key)
+
+ if isinstance(value, list):
+ if node_value not in value:
+ return False
+ else:
+ if node_value != value:
+ return False
+
+ return True
+
+ async def list_collections(self) -> list[str]:
+ """List all collection directories in the root path."""
+ if not self.root_path.exists():
+ return []
+
+ return [d.name for d in self.root_path.iterdir() if d.is_dir() and not d.name.startswith(".")]
+
+ async def create_collection(self, collection_name: str, **kwargs):
+ """Create a new collection directory."""
+ col_path = self._get_collection_path(collection_name)
+ col_path.mkdir(parents=True, exist_ok=True)
+ logger.info(f"Created collection {collection_name} at {col_path}")
+
+ async def delete_collection(self, collection_name: str, **kwargs):
+ """Delete a collection directory and all its JSON files."""
+ col_path = self._get_collection_path(collection_name)
+
+ if not col_path.exists():
+ logger.warning(f"Collection {collection_name} does not exist")
+ return
+
+ for file_path in col_path.glob("*.json"):
+ file_path.unlink()
+
+ col_path.rmdir()
+ logger.info(f"Deleted collection {collection_name}")
+
+ async def copy_collection(self, collection_name: str, **kwargs):
+ """Copy all nodes from the current collection to a new one."""
+ source_path = self._get_collection_path(self.collection_name)
+ target_path = self._get_collection_path(collection_name)
+
+ if not source_path.exists():
+ logger.warning(f"Source collection {self.collection_name} does not exist")
+ return
+
+ target_path.mkdir(parents=True, exist_ok=True)
+
+ for file_path in source_path.glob("*.json"):
+ target_file = target_path / file_path.name
+ target_file.write_text(file_path.read_text(encoding="utf-8"), encoding="utf-8")
+
+ logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+ async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
+ """Insert vector nodes into the local store, generating embeddings if necessary."""
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+
+ nodes_without_vectors = [node for node in nodes if node.vector is None]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+ else:
+ nodes_to_insert = nodes
+
+ for node in nodes_to_insert:
+ self._save_node(node)
+
+ logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}")
+
+ async def search(
+ self,
+ query: str,
+ limit: int = 5,
+ filters: dict | None = None,
+ **kwargs,
+ ) -> list[VectorNode]:
+ """Search for nodes similar to the query using brute-force cosine similarity."""
+ query_vector = await self.get_embedding(query)
+ all_nodes = self._load_all_nodes()
+ filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)]
+
+ scored_nodes = []
+ for node in filtered_nodes:
+ if node.vector is None:
+ logger.warning(f"Node {node.vector_id} has no vector, skipping")
+ continue
+
+ try:
+ score = self._cosine_similarity(query_vector, node.vector)
+ scored_nodes.append((node, score))
+ except ValueError as e:
+ logger.warning(f"Failed to calculate similarity for node {node.vector_id}: {e}")
+
+ scored_nodes.sort(key=lambda x: x[1], reverse=True)
+
+ score_threshold = kwargs.get("score_threshold")
+ if score_threshold is not None:
+ scored_nodes = [(node, score) for node, score in scored_nodes if score >= score_threshold]
+
+ scored_nodes = scored_nodes[:limit]
+ results = []
+ for node, score in scored_nodes:
+ node.metadata["_score"] = score
+ results.append(node)
+
+ return results
+
+ async def delete(self, vector_ids: str | list[str], **kwargs):
+ """Delete specific vector nodes by their IDs."""
+ if isinstance(vector_ids, str):
+ vector_ids = [vector_ids]
+
+ deleted_count = 0
+ for vector_id in vector_ids:
+ file_path = self._get_node_file_path(vector_id)
+ if file_path.exists():
+ file_path.unlink()
+ deleted_count += 1
+ else:
+ logger.warning(f"Node {vector_id} does not exist")
+
+ logger.info(f"Deleted {deleted_count} nodes from {self.collection_name}")
+
+ async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
+ """Update existing vector nodes with new data or embeddings."""
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+
+ nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+ else:
+ nodes_to_update = nodes
+
+ updated_count = 0
+ for node in nodes_to_update:
+ file_path = self._get_node_file_path(node.vector_id)
+ if file_path.exists():
+ self._save_node(node)
+ updated_count += 1
+ else:
+ logger.warning(f"Node {node.vector_id} does not exist, skipping update")
+
+ logger.info(f"Updated {updated_count} nodes in {self.collection_name}")
+
+ async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
+ """Retrieve one or more vector nodes by their unique IDs."""
+ is_single = isinstance(vector_ids, str)
+ ids = [vector_ids] if is_single else vector_ids
+
+ results = []
+ for vector_id in ids:
+ node = self._load_node(vector_id)
+ if node:
+ results.append(node)
+ else:
+ logger.warning(f"Node {vector_id} not found")
+
+ return results[0] if is_single and results else results
+
+ async def list(
+ self,
+ filters: dict | None = None,
+ limit: int | None = None,
+ ) -> list[VectorNode]:
+ """List vector nodes in the collection with optional filtering and limits."""
+ all_nodes = self._load_all_nodes()
+ filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)]
+
+ if limit is not None:
+ filtered_nodes = filtered_nodes[:limit]
+
+ return filtered_nodes
+
+ async def close(self):
+ """Close the vector store (no-op for local file system)."""
+ logger.info("Local vector store closed")
diff --git a/reme_ai/core/vector_store/pgvector_store.py b/reme_ai/core/vector_store/pgvector_store.py
new file mode 100644
index 00000000..38c5667a
--- /dev/null
+++ b/reme_ai/core/vector_store/pgvector_store.py
@@ -0,0 +1,533 @@
+"""PostgreSQL pgvector implementation for vector storage and retrieval."""
+
+import json
+from typing import Any
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+_ASYNCPG_IMPORT_ERROR = None
+
+try:
+ import asyncpg
+ from asyncpg import Pool
+except ImportError as e:
+ _ASYNCPG_IMPORT_ERROR = e
+ asyncpg = None
+ Pool = None
+
+
+@C.register_vector_store("pgvector")
+class PGVectorStore(BaseVectorStore):
+ """Vector store implementation using PostgreSQL and pgvector for efficient similarity search."""
+
+ def __init__(
+ self,
+ collection_name: str,
+ embedding_model: BaseEmbeddingModel,
+ host: str = "localhost",
+ port: int = 5432,
+ database: str = "postgres",
+ user: str = "postgres",
+ password: str = "",
+ min_size: int = 1,
+ max_size: int = 10,
+ dsn: str | None = None,
+ use_hnsw: bool = True,
+ use_diskann: bool = False,
+ **kwargs,
+ ):
+ """Initialize the PGVector store with connection parameters and index settings."""
+ if _ASYNCPG_IMPORT_ERROR is not None:
+ raise ImportError(
+ "PGVector requires extra dependencies. Install with `pip install asyncpg pgvector`",
+ ) from _ASYNCPG_IMPORT_ERROR
+
+ super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs)
+
+ self.dsn = dsn
+ self.host = host
+ self.port = port
+ self.database = database
+ self.user = user
+ self.password = password
+ self.min_size = min_size
+ self.max_size = max_size
+ self.use_hnsw = use_hnsw
+ self.use_diskann = use_diskann
+ self._pool: Pool | None = None
+ self.embedding_model_dims = embedding_model.dimensions
+
+ async def _get_pool(self) -> Pool:
+ """Create or return the existing asyncpg connection pool."""
+ if self._pool is None:
+ if self.dsn:
+ self._pool = await asyncpg.create_pool(
+ dsn=self.dsn,
+ min_size=self.min_size,
+ max_size=self.max_size,
+ )
+ else:
+ self._pool = await asyncpg.create_pool(
+ host=self.host,
+ port=self.port,
+ database=self.database,
+ user=self.user,
+ password=self.password,
+ min_size=self.min_size,
+ max_size=self.max_size,
+ )
+
+ async with self._pool.acquire() as conn:
+ await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+
+ logger.info(f"PGVector connection pool created for database {self.database}")
+
+ return self._pool
+
+ async def _ensure_collection_exists(self):
+ """Check if the collection table exists and create it if missing."""
+ collections = await self.list_collections()
+ if self.collection_name not in collections:
+ await self.create_collection(self.collection_name)
+
+ async def list_collections(self) -> list[str]:
+ """List all available table names in the current database."""
+ pool = await self._get_pool()
+ async with pool.acquire() as conn:
+ rows = await conn.fetch(
+ "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'",
+ )
+ return [row["table_name"] for row in rows]
+
+ async def create_collection(self, collection_name: str, **kwargs):
+ """Create a new PostgreSQL table with vector support and appropriate indexing."""
+ pool = await self._get_pool()
+ dimensions = kwargs.get("dimensions", self.embedding_model_dims)
+
+ async with pool.acquire() as conn:
+ await conn.execute(
+ f"""
+ CREATE TABLE IF NOT EXISTS {collection_name} (
+ id TEXT PRIMARY KEY,
+ content TEXT,
+ vector vector({dimensions}),
+ metadata JSONB
+ )
+ """,
+ )
+
+ if self.use_diskann and dimensions < 2000:
+ result = await conn.fetchval(
+ "SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'",
+ )
+ if result:
+ await conn.execute(
+ f"""
+ CREATE INDEX IF NOT EXISTS {collection_name}_diskann_idx
+ ON {collection_name}
+ USING diskann (vector)
+ """,
+ )
+ logger.info(f"Created DiskANN index for collection {collection_name}")
+ else:
+ logger.warning("vectorscale extension not available, skipping DiskANN index")
+ elif self.use_hnsw:
+ await conn.execute(
+ f"""
+ CREATE INDEX IF NOT EXISTS {collection_name}_hnsw_idx
+ ON {collection_name}
+ USING hnsw (vector vector_cosine_ops)
+ """,
+ )
+ logger.info(f"Created HNSW index for collection {collection_name}")
+
+ logger.info(f"Created collection {collection_name} with dimensions={dimensions}")
+
+ async def delete_collection(self, collection_name: str, **kwargs):
+ """Remove the specified collection table from the database."""
+ pool = await self._get_pool()
+ async with pool.acquire() as conn:
+ await conn.execute(f"DROP TABLE IF EXISTS {collection_name}")
+ logger.info(f"Deleted collection {collection_name}")
+
+ async def copy_collection(self, collection_name: str, **kwargs):
+ """Duplicate the structure and content of the current collection to a new table."""
+ pool = await self._get_pool()
+
+ async with pool.acquire() as conn:
+ columns = await conn.fetch(
+ """
+ SELECT column_name, data_type, udt_name
+ FROM information_schema.columns
+ WHERE table_name = $1 AND table_schema = 'public'
+ """,
+ self.collection_name,
+ )
+
+ if not columns:
+ raise ValueError(f"Source collection {self.collection_name} does not exist")
+
+ await conn.execute(f"CREATE TABLE {collection_name} AS TABLE {self.collection_name}")
+ await conn.execute(f"ALTER TABLE {collection_name} ADD PRIMARY KEY (id)")
+
+ if self.use_hnsw:
+ await conn.execute(
+ f"""
+ CREATE INDEX IF NOT EXISTS {collection_name}_hnsw_idx
+ ON {collection_name}
+ USING hnsw (vector vector_cosine_ops)
+ """,
+ )
+
+ logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+ async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
+ """Insert or upsert vector nodes into the PostgreSQL collection."""
+ await self._ensure_collection_exists()
+
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+
+ if not nodes:
+ return
+
+ nodes_without_vectors = [node for node in nodes if node.vector is None]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+ else:
+ nodes_to_insert = nodes
+
+ pool = await self._get_pool()
+ data = [
+ (
+ node.vector_id,
+ node.content,
+ f"[{','.join(map(str, node.vector))}]",
+ json.dumps(node.metadata),
+ )
+ for node in nodes_to_insert
+ ]
+
+ async with pool.acquire() as conn:
+ on_conflict = kwargs.get("on_conflict", "update")
+
+ if on_conflict == "update":
+ await conn.executemany(
+ f"""
+ INSERT INTO {self.collection_name} (id, content, vector, metadata)
+ VALUES ($1, $2, $3::vector, $4::jsonb)
+ ON CONFLICT (id) DO UPDATE SET
+ content = EXCLUDED.content,
+ vector = EXCLUDED.vector,
+ metadata = EXCLUDED.metadata
+ """,
+ data,
+ )
+ elif on_conflict == "ignore":
+ await conn.executemany(
+ f"""
+ INSERT INTO {self.collection_name} (id, content, vector, metadata)
+ VALUES ($1, $2, $3::vector, $4::jsonb)
+ ON CONFLICT (id) DO NOTHING
+ """,
+ data,
+ )
+ else:
+ await conn.executemany(
+ f"""
+ INSERT INTO {self.collection_name} (id, content, vector, metadata)
+ VALUES ($1, $2, $3::vector, $4::jsonb)
+ """,
+ data,
+ )
+
+ logger.info(f"Inserted {len(nodes_to_insert)} documents into {self.collection_name}")
+
+ @staticmethod
+ def _build_filter_clause(filters: dict | None) -> tuple[str, list]:
+ """Generate an SQL WHERE clause and parameter list from a filter dictionary."""
+ if not filters:
+ return "", []
+
+ conditions = []
+ params = []
+ param_idx = 1
+
+ for key, value in filters.items():
+ if isinstance(value, list):
+ placeholders = ", ".join([f"${param_idx + i}" for i in range(len(value))])
+ conditions.append(f"metadata->>'{key}' IN ({placeholders})")
+ params.extend([str(v) for v in value])
+ param_idx += len(value)
+ else:
+ conditions.append(f"metadata->>'{key}' = ${param_idx}")
+ params.append(str(value))
+ param_idx += 1
+
+ filter_clause = "WHERE " + " AND ".join(conditions) if conditions else ""
+ return filter_clause, params
+
+ async def search(
+ self,
+ query: str,
+ limit: int = 5,
+ filters: dict | None = None,
+ **kwargs,
+ ) -> list[VectorNode]:
+ """Perform vector similarity search with optional metadata filtering."""
+ await self._ensure_collection_exists()
+
+ query_vector = await self.get_embedding(query)
+ vector_str = f"[{','.join(map(str, query_vector))}]"
+ pool = await self._get_pool()
+
+ filter_clause, filter_params = self._build_filter_clause(filters)
+
+ if filter_clause:
+ for i in range(len(filter_params)):
+ old_idx = i + 1
+ new_idx = i + 2
+ filter_clause = filter_clause.replace(f"${old_idx}", f"${new_idx}", 1)
+
+ async with pool.acquire() as conn:
+ sql = f"""
+ SELECT id, content, vector, metadata, vector <=> $1::vector AS distance
+ FROM {self.collection_name}
+ {filter_clause}
+ ORDER BY distance
+ LIMIT ${len(filter_params) + 2}
+ """
+ rows = await conn.fetch(sql, vector_str, *filter_params, limit)
+
+ results = []
+ score_threshold = kwargs.get("score_threshold")
+
+ for row in rows:
+ distance = row["distance"]
+ if score_threshold is not None and distance > score_threshold:
+ continue
+
+ vector_data = None
+ if row["vector"]:
+ vector_str_raw = str(row["vector"])
+ if vector_str_raw.startswith("[") and vector_str_raw.endswith("]"):
+ vector_data = [float(x) for x in vector_str_raw[1:-1].split(",")]
+
+ metadata = row["metadata"] if row["metadata"] else {}
+ if isinstance(metadata, str):
+ metadata = json.loads(metadata)
+
+ metadata["_score"] = 1 - distance
+ metadata["_distance"] = distance
+
+ node = VectorNode(
+ vector_id=row["id"],
+ content=row["content"] or "",
+ vector=vector_data,
+ metadata=metadata,
+ )
+ results.append(node)
+
+ return results
+
+ async def delete(self, vector_ids: str | list[str], **kwargs):
+ """Remove specific vector records from the collection by their IDs."""
+ await self._ensure_collection_exists()
+
+ if isinstance(vector_ids, str):
+ vector_ids = [vector_ids]
+
+ if not vector_ids:
+ return
+
+ pool = await self._get_pool()
+ async with pool.acquire() as conn:
+ placeholders = ", ".join([f"${i + 1}" for i in range(len(vector_ids))])
+ await conn.execute(
+ f"DELETE FROM {self.collection_name} WHERE id IN ({placeholders})",
+ *vector_ids,
+ )
+
+ logger.info(f"Deleted {len(vector_ids)} documents from {self.collection_name}")
+
+ async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
+ """Update existing vector nodes with new content, embeddings, or metadata."""
+ await self._ensure_collection_exists()
+
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+
+ if not nodes:
+ return
+
+ nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+ else:
+ nodes_to_update = nodes
+
+ pool = await self._get_pool()
+ async with pool.acquire() as conn:
+ for node in nodes_to_update:
+ update_fields = []
+ params = []
+ idx = 1
+
+ if node.content:
+ update_fields.append(f"content = ${idx}")
+ params.append(node.content)
+ idx += 1
+
+ if node.vector:
+ vector_str = f"[{','.join(map(str, node.vector))}]"
+ update_fields.append(f"vector = ${idx}::vector")
+ params.append(vector_str)
+ idx += 1
+
+ if node.metadata:
+ update_fields.append(f"metadata = ${idx}::jsonb")
+ params.append(json.dumps(node.metadata))
+ idx += 1
+
+ if update_fields:
+ params.append(node.vector_id)
+ await conn.execute(
+ f"UPDATE {self.collection_name} SET {', '.join(update_fields)} WHERE id = ${idx}",
+ *params,
+ )
+
+ logger.info(f"Updated {len(nodes_to_update)} documents in {self.collection_name}")
+
+ async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None:
+ """Retrieve vector nodes by their unique identifiers."""
+ await self._ensure_collection_exists()
+
+ single_result = isinstance(vector_ids, str)
+ if single_result:
+ vector_ids = [vector_ids]
+
+ if not vector_ids:
+ return [] if not single_result else None
+
+ pool = await self._get_pool()
+ async with pool.acquire() as conn:
+ placeholders = ", ".join([f"${i + 1}" for i in range(len(vector_ids))])
+ rows = await conn.fetch(
+ f"SELECT id, content, vector, metadata FROM {self.collection_name} WHERE id IN ({placeholders})",
+ *vector_ids,
+ )
+
+ results = []
+ for row in rows:
+ vector_data = None
+ if row["vector"]:
+ vector_str_raw = str(row["vector"])
+ if vector_str_raw.startswith("[") and vector_str_raw.endswith("]"):
+ vector_data = [float(x) for x in vector_str_raw[1:-1].split(",")]
+
+ metadata = row["metadata"] if row["metadata"] else {}
+ if isinstance(metadata, str):
+ metadata = json.loads(metadata)
+
+ results.append(
+ VectorNode(
+ vector_id=row["id"],
+ content=row["content"] or "",
+ vector=vector_data,
+ metadata=metadata,
+ ),
+ )
+
+ if single_result:
+ return results[0] if results else None
+ return results
+
+ async def list(
+ self,
+ filters: dict | None = None,
+ limit: int | None = None,
+ ) -> list[VectorNode]:
+ """Return a list of vector nodes matching the provided filters and limit."""
+ await self._ensure_collection_exists()
+
+ pool = await self._get_pool()
+ filter_clause, filter_params = self._build_filter_clause(filters)
+
+ limit_clause = ""
+ if limit:
+ limit_clause = f"LIMIT ${len(filter_params) + 1}"
+ filter_params.append(limit)
+
+ async with pool.acquire() as conn:
+ sql = f"""
+ SELECT id, content, vector, metadata
+ FROM {self.collection_name}
+ {filter_clause}
+ {limit_clause}
+ """
+ rows = await conn.fetch(sql, *filter_params)
+
+ results = []
+ for row in rows:
+ vector_data = None
+ if row["vector"]:
+ vector_str_raw = str(row["vector"])
+ if vector_str_raw.startswith("[") and vector_str_raw.endswith("]"):
+ vector_data = [float(x) for x in vector_str_raw[1:-1].split(",")]
+
+ metadata = row["metadata"] if row["metadata"] else {}
+ if isinstance(metadata, str):
+ metadata = json.loads(metadata)
+
+ results.append(
+ VectorNode(
+ vector_id=row["id"],
+ content=row["content"] or "",
+ vector=vector_data,
+ metadata=metadata,
+ ),
+ )
+
+ return results
+
+ async def collection_info(self) -> dict[str, Any]:
+ """Fetch metadata including record count and disk usage for the collection."""
+ pool = await self._get_pool()
+
+ async with pool.acquire() as conn:
+ row = await conn.fetchrow(
+ f"""
+ SELECT
+ '{self.collection_name}' as name,
+ (SELECT COUNT(*) FROM {self.collection_name}) as row_count,
+ pg_size_pretty(pg_total_relation_size('{self.collection_name}')) as total_size
+ """,
+ )
+
+ return {
+ "name": row["name"],
+ "count": row["row_count"],
+ "size": row["total_size"],
+ }
+
+ async def reset(self):
+ """Purge all data by dropping and recreating the collection table."""
+ logger.warning(f"Resetting collection {self.collection_name}...")
+ await self.delete_collection(self.collection_name)
+ await self.create_collection(self.collection_name)
+
+ async def close(self):
+ """Terminate the database connection pool and release associated resources."""
+ if self._pool is not None:
+ await self._pool.close()
+ self._pool = None
+ logger.info("PGVector connection pool closed")
diff --git a/reme_ai/core/vector_store/qdrant_vector_store.py b/reme_ai/core/vector_store/qdrant_vector_store.py
new file mode 100644
index 00000000..97eb5b61
--- /dev/null
+++ b/reme_ai/core/vector_store/qdrant_vector_store.py
@@ -0,0 +1,444 @@
+"""Qdrant vector store implementation for the ReMe project."""
+
+from typing import Any
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+_QDRANT_IMPORT_ERROR = None
+
+try:
+ from qdrant_client import AsyncQdrantClient
+ from qdrant_client.models import (
+ Distance,
+ FieldCondition,
+ Filter,
+ MatchValue,
+ PointIdsList,
+ PointStruct,
+ Range,
+ VectorParams,
+ )
+except ImportError as e:
+ _QDRANT_IMPORT_ERROR = e
+ AsyncQdrantClient = None
+ Distance = None
+ FieldCondition = None
+ Filter = None
+ MatchValue = None
+ PointIdsList = None
+ PointStruct = None
+ Range = None
+ VectorParams = None
+
+
+@C.register_vector_store("qdrant")
+class QdrantVectorStore(BaseVectorStore):
+ """Vector store implementation using Qdrant for dense vector search."""
+
+ def __init__(
+ self,
+ collection_name: str,
+ embedding_model: BaseEmbeddingModel,
+ host: str | None = None,
+ port: int = 6333,
+ path: str | None = None,
+ url: str | None = None,
+ api_key: str | None = None,
+ https: bool | None = None,
+ grpc_port: int = 6334,
+ prefer_grpc: bool = False,
+ distance: str = "cosine",
+ on_disk: bool = False,
+ **kwargs: Any,
+ ):
+ """Initialize the Qdrant client and collection configuration.
+
+ Args:
+ collection_name: Name of the collection.
+ embedding_model: Model used for generating vector embeddings.
+ host: Server host address.
+ port: HTTP port for the server.
+ path: Local storage path for on-disk/in-memory mode.
+ url: Full connection URL.
+ api_key: Authentication key for Qdrant Cloud.
+ https: Use secure connection if True.
+ grpc_port: gRPC interface port.
+ prefer_grpc: Use gRPC instead of HTTP if True.
+ distance: Metric for similarity (cosine, euclid, dot).
+ on_disk: Enable persistent storage for vectors.
+ **kwargs: Additional client configuration.
+ """
+ if _QDRANT_IMPORT_ERROR is not None:
+ raise ImportError(
+ "Qdrant requires extra dependencies. Install with `pip install qdrant-client`",
+ ) from _QDRANT_IMPORT_ERROR
+
+ super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs)
+
+ self.client = AsyncQdrantClient(
+ host=host,
+ port=port,
+ path=path,
+ url=url,
+ api_key=api_key,
+ https=https,
+ grpc_port=grpc_port,
+ prefer_grpc=prefer_grpc,
+ **kwargs,
+ )
+
+ self.is_local = path is not None
+ distance_map = {
+ "cosine": Distance.COSINE,
+ "euclid": Distance.EUCLID,
+ "dot": Distance.DOT,
+ }
+ self.distance = distance_map.get(distance.lower(), Distance.COSINE)
+ self.on_disk = on_disk
+
+ async def list_collections(self) -> list[str]:
+ """Retrieve names of all existing collections in the Qdrant instance."""
+ collections = await self.client.get_collections()
+ return [collection.name for collection in collections.collections]
+
+ async def create_collection(self, collection_name: str, **kwargs: Any):
+ """Create a new collection with the specified vector configuration.
+
+ Args:
+ collection_name: Name of the collection to create.
+ **kwargs: Overrides for dimensions, distance, or on_disk settings.
+ """
+ collections = await self.list_collections()
+ if collection_name in collections:
+ logger.info(f"Collection {collection_name} already exists")
+ return
+
+ dimensions = kwargs.get("dimensions", self.embedding_model.dimensions)
+ distance = kwargs.get("distance", self.distance)
+ on_disk = kwargs.get("on_disk", self.on_disk)
+
+ await self.client.create_collection(
+ collection_name=collection_name,
+ vectors_config=VectorParams(
+ size=dimensions,
+ distance=distance,
+ on_disk=on_disk,
+ ),
+ )
+
+ logger.info(f"Created collection {collection_name} with dimensions={dimensions}")
+
+ if not self.is_local:
+ await self._create_payload_indexes(collection_name)
+
+ async def _create_payload_indexes(self, collection_name: str):
+ """Create keyword indexes for common metadata fields to optimize filtering."""
+ common_fields = ["user_id", "agent_id", "run_id", "actor_id", "source"]
+
+ for field in common_fields:
+ try:
+ await self.client.create_payload_index(
+ collection_name=collection_name,
+ field_name=field,
+ field_schema="keyword",
+ )
+ logger.debug(f"Created index for {field} in collection {collection_name}")
+ except Exception as e:
+ logger.debug(f"Index for {field} might already exist: {e}")
+
+ async def delete_collection(self, collection_name: str, **kwargs: Any):
+ """Permanently remove a collection from the Qdrant instance."""
+ collections = await self.list_collections()
+ if collection_name in collections:
+ await self.client.delete_collection(collection_name=collection_name)
+ logger.info(f"Deleted collection {collection_name}")
+ else:
+ logger.warning(f"Collection {collection_name} does not exist")
+
+ async def copy_collection(self, collection_name: str, **kwargs: Any):
+ """Duplicate an existing collection to a new one including all data."""
+ collection_info = await self.client.get_collection(collection_name=self.collection_name)
+
+ await self.client.create_collection(
+ collection_name=collection_name,
+ vectors_config=collection_info.config.params.vectors,
+ )
+
+ offset = None
+ batch_size = 100
+
+ while True:
+ records, next_offset = await self.client.scroll(
+ collection_name=self.collection_name,
+ limit=batch_size,
+ offset=offset,
+ with_payload=True,
+ with_vectors=True,
+ )
+
+ if not records:
+ break
+
+ points = [
+ PointStruct(
+ id=record.id,
+ vector=record.vector,
+ payload=record.payload,
+ )
+ for record in records
+ ]
+
+ await self.client.upsert(
+ collection_name=collection_name,
+ points=points,
+ )
+
+ offset = next_offset
+ if offset is None:
+ break
+
+ logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+ async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs: Any):
+ """Insert vector nodes into the collection, generating embeddings as needed."""
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+
+ nodes_without_vectors = [node for node in nodes if node.vector is None]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+ else:
+ nodes_to_insert = nodes
+
+ points = []
+ for node in nodes_to_insert:
+ try:
+ point_id = int(node.vector_id)
+ except ValueError:
+ point_id = abs(hash(node.vector_id)) % (10**18)
+
+ point = PointStruct(
+ id=point_id,
+ vector=node.vector,
+ payload={
+ "vector_id": node.vector_id,
+ "content": node.content,
+ "metadata": node.metadata,
+ },
+ )
+ points.append(point)
+
+ wait = kwargs.get("wait", True)
+ await self.client.upsert(
+ collection_name=self.collection_name,
+ points=points,
+ wait=wait,
+ )
+
+ logger.info(f"Inserted {len(points)} documents into {self.collection_name}")
+
+ @staticmethod
+ def _create_filter(filters: dict) -> Filter | None:
+ """Convert a dictionary of filter conditions into a Qdrant Filter object."""
+ if not filters:
+ return None
+
+ conditions = []
+ for key, value in filters.items():
+ if isinstance(value, dict) and ("gte" in value or "lte" in value):
+ range_params = {}
+ if "gte" in value:
+ range_params["gte"] = value["gte"]
+ if "lte" in value:
+ range_params["lte"] = value["lte"]
+ conditions.append(
+ FieldCondition(
+ key=f"metadata.{key}",
+ range=Range(**range_params),
+ ),
+ )
+ elif isinstance(value, list):
+ conditions.append(
+ FieldCondition(key=f"metadata.{key}", match=MatchValue(value=value[0])),
+ )
+ else:
+ conditions.append(
+ FieldCondition(key=f"metadata.{key}", match=MatchValue(value=value)),
+ )
+
+ return Filter(must=conditions) if conditions else None
+
+ async def search(
+ self,
+ query: str,
+ limit: int = 5,
+ filters: dict | None = None,
+ **kwargs: Any,
+ ) -> list[VectorNode]:
+ """Search for the most similar vectors based on a text query."""
+ query_vector = await self.get_embedding(query)
+ query_filter = self._create_filter(filters) if filters else None
+ score_threshold = kwargs.get("score_threshold", None)
+
+ results = await self.client.query_points(
+ collection_name=self.collection_name,
+ query=query_vector,
+ query_filter=query_filter,
+ limit=limit,
+ score_threshold=score_threshold,
+ )
+
+ nodes = []
+ for point in results.points:
+ payload = point.payload or {}
+ node = VectorNode(
+ vector_id=payload.get("vector_id", str(point.id)),
+ content=payload.get("content", ""),
+ vector=point.vector if hasattr(point, "vector") else None,
+ metadata=payload.get("metadata", {}),
+ )
+ node.metadata["_score"] = point.score
+ nodes.append(node)
+
+ return nodes
+
+ async def delete(self, vector_ids: str | list[str], **kwargs: Any):
+ """Delete specific vectors from the collection using their IDs."""
+ if isinstance(vector_ids, str):
+ vector_ids = [vector_ids]
+
+ point_ids = []
+ for vector_id in vector_ids:
+ try:
+ point_id = int(vector_id)
+ except ValueError:
+ point_id = abs(hash(vector_id)) % (10**18)
+ point_ids.append(point_id)
+
+ wait = kwargs.get("wait", True)
+ await self.client.delete(
+ collection_name=self.collection_name,
+ points_selector=PointIdsList(points=point_ids),
+ wait=wait,
+ )
+
+ logger.info(f"Deleted {len(point_ids)} documents from {self.collection_name}")
+
+ async def update(self, nodes: VectorNode | list[VectorNode], **kwargs: Any):
+ """Update existing vector nodes with new content or metadata."""
+ if isinstance(nodes, VectorNode):
+ nodes = [nodes]
+
+ nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+ if nodes_without_vectors:
+ nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+ vector_map = {n.vector_id: n for n in nodes_with_vectors}
+ nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+ else:
+ nodes_to_update = nodes
+
+ points = []
+ for node in nodes_to_update:
+ try:
+ point_id = int(node.vector_id)
+ except ValueError:
+ point_id = abs(hash(node.vector_id)) % (10**18)
+
+ point = PointStruct(
+ id=point_id,
+ vector=node.vector,
+ payload={
+ "vector_id": node.vector_id,
+ "content": node.content,
+ "metadata": node.metadata,
+ },
+ )
+ points.append(point)
+
+ wait = kwargs.get("wait", True)
+ await self.client.upsert(
+ collection_name=self.collection_name,
+ points=points,
+ wait=wait,
+ )
+
+ logger.info(f"Updated {len(points)} documents in {self.collection_name}")
+
+ async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
+ """Retrieve vector nodes by their IDs from the collection."""
+ single_result = isinstance(vector_ids, str)
+ if single_result:
+ vector_ids = [vector_ids]
+
+ point_ids = []
+ for vector_id in vector_ids:
+ try:
+ point_id = int(vector_id)
+ except ValueError:
+ point_id = abs(hash(vector_id)) % (10**18)
+ point_ids.append(point_id)
+
+ points = await self.client.retrieve(
+ collection_name=self.collection_name,
+ ids=point_ids,
+ with_payload=True,
+ with_vectors=True,
+ )
+
+ results = []
+ for point in points:
+ if point:
+ payload = point.payload or {}
+ node = VectorNode(
+ vector_id=payload.get("vector_id", str(point.id)),
+ content=payload.get("content", ""),
+ vector=point.vector,
+ metadata=payload.get("metadata", {}),
+ )
+ results.append(node)
+ else:
+ logger.warning("Point not found")
+
+ return results[0] if single_result and results else results
+
+ async def list(
+ self,
+ filters: dict | None = None,
+ limit: int | None = None,
+ ) -> list[VectorNode]:
+ """List all vector nodes in the collection matching the filter criteria."""
+ scroll_filter = self._create_filter(filters) if filters else None
+
+ limit = limit or 10000
+ records, _ = await self.client.scroll(
+ collection_name=self.collection_name,
+ scroll_filter=scroll_filter,
+ limit=limit,
+ with_payload=True,
+ with_vectors=True,
+ )
+
+ results = []
+ for record in records:
+ payload = record.payload or {}
+ node = VectorNode(
+ vector_id=payload.get("vector_id", str(record.id)),
+ content=payload.get("content", ""),
+ vector=record.vector,
+ metadata=payload.get("metadata", {}),
+ )
+ results.append(node)
+
+ return results
+
+ async def close(self):
+ """Close the AsyncQdrantClient connection and release resources."""
+ await self.client.close()
+ logger.info("Qdrant client connection closed")
diff --git a/tests/mcp_servers_demo.json b/tests/mcp_servers_demo.json
new file mode 100644
index 00000000..e8f9a858
--- /dev/null
+++ b/tests/mcp_servers_demo.json
@@ -0,0 +1,37 @@
+{
+ "mcpServers": {
+ "sqlite-explorer": {
+ "type": "stdio",
+ "command": "uv",
+ "args": [
+ "run",
+ "--with",
+ "mcp-server-sqlite",
+ "mcp-server-sqlite",
+ "--db-path",
+ "/path/to/your/database.db"
+ ],
+ "env": {
+ "CUSTOM_VAR": "optional_value"
+ }
+ },
+ "remote-fetcher": {
+ "type": "sse",
+ "url": "https://mcp-server.example.com/sse",
+ "headers": {
+ "Authorization": "Bearer {BAILIAN_MCP_API_KEY}",
+ "Content-Type": "application/json"
+ },
+ "timeout": 5,
+ "sse_read_timeout": 300
+ },
+ "my-modern-remote": {
+ "type": "streamable-http",
+ "url": "https://api.example.com/mcp",
+ "headers": {
+ "Authorization": "Bearer {BAILIAN_MCP_API_KEY}"
+ },
+ "timeout": 5
+ }
+ }
+}
\ No newline at end of file
diff --git a/tests/test_base_context.py b/tests/test_base_context.py
new file mode 100644
index 00000000..316a6796
--- /dev/null
+++ b/tests/test_base_context.py
@@ -0,0 +1,72 @@
+"""
+Unit tests for the BaseContext class in reme_ai.core.context.
+Ensures attribute-style and dict-style access work interchangeably.
+"""
+
+import pickle
+from reme_ai.core.context import BaseContext
+
+
+def test_attribute_access():
+ """Test setting values via attributes and retrieving via items."""
+ context = BaseContext()
+ context.xxx = 123
+ assert context.xxx == 123
+ assert context["xxx"] == 123
+
+
+def test_dict_access():
+ """Test setting values via items and retrieving via attributes."""
+ context = BaseContext()
+ context["yyy"] = 456
+ assert context.yyy == 456
+ assert context["yyy"] == 456
+
+
+def test_delete_attribute():
+ """Test that deleting an attribute removes it from the internal state."""
+ context = BaseContext()
+ context.zzz = 789
+ del context.zzz
+ assert "zzz" not in context
+
+
+def test_attribute_error():
+ """Test that accessing non-existent attributes raises the correct error."""
+ context = BaseContext()
+ try:
+ _ = context.nonexistent
+ assert False, "Should raise AttributeError"
+ except AttributeError as error:
+ assert "nonexistent" in str(error)
+
+
+def test_pickling():
+ """Test that BaseContext instances can be serialized and deserialized."""
+ context = BaseContext()
+ context.test_value = "bar"
+ context.num = 42
+
+ pickled = pickle.dumps(context)
+ restored = pickle.loads(pickled)
+
+ assert restored.test_value == "bar"
+ assert restored.num == 42
+ assert isinstance(restored, BaseContext)
+
+
+def test_init_with_data():
+ """Test that the constructor correctly handles initial dictionary data."""
+ context = BaseContext({"a": 1, "b": 2})
+ assert context.a == 1
+ assert context.b == 2
+
+
+if __name__ == "__main__":
+ test_attribute_access()
+ test_dict_access()
+ test_delete_attribute()
+ test_attribute_error()
+ test_pickling()
+ test_init_with_data()
+ print("All tests passed!")
diff --git a/tests/test_cache_handler.py b/tests/test_cache_handler.py
new file mode 100644
index 00000000..ddcac86f
--- /dev/null
+++ b/tests/test_cache_handler.py
@@ -0,0 +1,94 @@
+"""
+Self-contained script for CacheHandler's comprehensive test suite.
+"""
+
+import shutil
+from datetime import datetime, timedelta
+from pathlib import Path
+
+import pandas as pd
+from loguru import logger
+
+from reme_ai.core.utils.cache_handler import CacheHandler
+
+
+def run_tests():
+ """Execute comprehensive tests for CacheHandler."""
+ test_dir = Path("test_cache_system")
+ if test_dir.exists():
+ shutil.rmtree(test_dir)
+
+ handler = CacheHandler(cache_dir=test_dir)
+ logger.info("Starting CacheHandler tests...")
+
+ # 1. Test Data Types
+ logger.info("Testing data types support...")
+
+ # DataFrame
+ df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
+ assert handler.save("df_test", df)
+ assert isinstance(handler.load("df_test"), pd.DataFrame)
+ assert handler.load("df_test").shape == (2, 2)
+
+ # Dict & List
+ d = {"key": "value", "nested": [1, 2]}
+ l_value = [1, "string", {"a": 1}]
+ assert handler.save("dict_test", d)
+ assert handler.save("list_test", l_value)
+ assert handler.load("dict_test")["key"] == "value"
+ assert handler.load("list_test")[1] == "string"
+
+ # String
+ s = "Hello World"
+ assert handler.save("str_test", s)
+ assert handler.load("str_test") == "Hello World"
+
+ # 2. Test Expiration
+ logger.info("Testing expiration logic...")
+ # Save with 1 second expiry (approx 0.00027 hours)
+ handler.save("exp_test", {"data": 1}, expire_hours=0.00001)
+ assert handler.exists("exp_test") is True
+
+ # Manually modify metadata to force expiration for instant test
+ handler.metadata["exp_test"]["expire_at"] = (datetime.now() - timedelta(seconds=1)).isoformat()
+ assert handler.exists("exp_test") is False
+ assert handler.load("exp_test") is None
+ assert "exp_test" not in handler.metadata # Auto-cleaned
+
+ # 3. Test Existence and Deletion
+ logger.info("Testing delete and exists...")
+ handler.save("del_test", "delete me")
+ assert handler.exists("del_test") is True
+ handler.delete("del_test")
+ assert handler.exists("del_test") is False
+ assert not (test_dir / "del_test.txt").exists()
+
+ # 4. Test Persistence (Reload handler)
+ logger.info("Testing persistence...")
+ handler.save("persist_test", [1, 2, 3])
+ new_handler = CacheHandler(cache_dir=test_dir)
+ assert new_handler.exists("persist_test") is True
+ assert new_handler.load("persist_test") == [1, 2, 3]
+
+ # 5. Test Statistics and Clear
+ logger.info("Testing stats and clear...")
+ stats = handler.get_stats()
+ assert stats["count"] > 0
+ handler.clear_all()
+ assert handler.get_stats()["count"] == 0
+ assert len(list(test_dir.glob("*"))) == 1 # Only metadata.json remains
+
+ # 6. Test Error Handling
+ logger.info("Testing error handling...")
+ assert handler.load("non_existent_key") is None
+ # Test unsupported type
+ assert handler.save("invalid", {1, 2}) is False
+
+ logger.success("All tests passed successfully!")
+
+ # Cleanup after tests
+ shutil.rmtree(test_dir)
+
+
+if __name__ == "__main__":
+ run_tests()
diff --git a/tests/test_embedding.py b/tests/test_embedding.py
new file mode 100644
index 00000000..d1c404f5
--- /dev/null
+++ b/tests/test_embedding.py
@@ -0,0 +1,349 @@
+"""
+Async unit tests for Embedding classes (OpenAIEmbeddingModel) covering:
+- Async single text embedding
+- Async batch text embeddings
+- Async large batch with automatic batching
+- Async VectorNode embedding (single and batch)
+- Error handling and retries
+
+Usage:
+ python test_embedding.py --openai # Test OpenAIEmbeddingModel only
+ python test_embedding.py --all # Test all embedding models
+"""
+
+# flake8: noqa: E402
+# pylint: disable=C0413
+
+import asyncio
+import argparse
+from typing import Type, List
+
+from reme_ai.core.utils import load_env
+
+load_env()
+
+from reme_ai.core.embedding import OpenAIEmbeddingModel, BaseEmbeddingModel
+from reme_ai.core.schema import VectorNode
+
+
+def get_embedding_model(model_class: Type[BaseEmbeddingModel]) -> BaseEmbeddingModel:
+ """Create and return an embedding model instance."""
+ return model_class(
+ model_name="text-embedding-v4",
+ dimensions=1024,
+ max_retries=2,
+ raise_exception=True,
+ )
+
+
+def get_test_texts() -> List[str]:
+ """Create test texts for embedding."""
+ return [
+ "The quick brown fox jumps over the lazy dog.",
+ "Machine learning is a subset of artificial intelligence.",
+ "Python is a popular programming language for data science.",
+ "Solar energy is a renewable source of power.",
+ "The capital of France is Paris.",
+ ]
+
+
+def get_large_batch_texts() -> List[str]:
+ """Create a large batch of test texts to test automatic batching."""
+ texts = []
+ topics = [
+ "Climate change and global warming",
+ "Artificial intelligence and machine learning",
+ "Renewable energy sources",
+ "Space exploration and astronomy",
+ "Medical research and healthcare",
+ "Financial markets and economics",
+ "Education and learning systems",
+ "Transportation and urban planning",
+ ]
+
+ for i, topic in enumerate(topics):
+ for j in range(3):
+ texts.append(f"Text {i*3+j+1}: This is a sample text about {topic}.")
+
+ return texts # 24 texts total
+
+
+def get_test_nodes() -> List[VectorNode]:
+ """Create test VectorNodes for embedding."""
+ texts = get_test_texts()
+ return [
+ VectorNode(
+ content=text,
+ metadata={"index": str(i), "category": "test"},
+ )
+ for i, text in enumerate(texts)
+ ]
+
+
+async def test_async_single_embedding(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test asynchronous single text embedding."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Async Single Text Embedding")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ test_text = "Hello, this is a test sentence for embedding."
+
+ print(f"Input text: {test_text}")
+
+ embedding = await model.get_embedding(test_text)
+
+ assert embedding is not None, f"{model_name}: Embedding is None"
+ assert isinstance(embedding, list), f"{model_name}: Embedding is not a list"
+ assert len(embedding) > 0, f"{model_name}: Empty embedding"
+ assert len(embedding) == model.dimensions, f"{model_name}: Embedding dimension mismatch"
+ assert all(isinstance(x, float) for x in embedding), f"{model_name}: Not all elements are floats"
+
+ print("\n✓ Embedding generated successfully")
+ print(f" - Dimension: {len(embedding)}")
+ print(f" - First 5 values: {embedding[:5]}")
+ print(f" - Value range: [{min(embedding):.4f}, {max(embedding):.4f}]")
+
+ await model.close()
+ print(f"✓ PASSED: {model_name} async single embedding")
+
+
+async def test_async_batch_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test asynchronous batch text embeddings."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Async Batch Text Embeddings")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ test_texts = get_test_texts()
+
+ print(f"Input: {len(test_texts)} texts")
+ for i, text in enumerate(test_texts[:3], 1):
+ print(f" {i}. {text[:50]}...")
+
+ embeddings = await model.get_embeddings(test_texts)
+
+ assert embeddings is not None, f"{model_name}: Embeddings is None"
+ assert isinstance(embeddings, list), f"{model_name}: Embeddings is not a list"
+ assert len(embeddings) == len(test_texts), f"{model_name}: Embeddings count mismatch"
+
+ for i, emb in enumerate(embeddings):
+ assert isinstance(emb, list), f"{model_name}: Embedding {i} is not a list"
+ assert len(emb) == model.dimensions, f"{model_name}: Embedding {i} dimension mismatch"
+ assert all(isinstance(x, float) for x in emb), f"{model_name}: Embedding {i} has non-float values"
+
+ print("\n✓ Batch embeddings generated successfully")
+ print(f" - Count: {len(embeddings)}")
+ print(f" - Dimension: {len(embeddings[0])}")
+ print(f" - First embedding preview: {embeddings[0][:3]}...")
+
+ await model.close()
+ print(f"✓ PASSED: {model_name} async batch embeddings")
+
+
+async def test_async_large_batch_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test asynchronous large batch embeddings with automatic batching."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Async Large Batch with Auto-Batching")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ test_texts = get_large_batch_texts()
+
+ print(f"Input: {len(test_texts)} texts")
+ print(f"Max batch size: {model.max_batch_size}")
+ print(f"Expected batches: {(len(test_texts) + model.max_batch_size - 1) // model.max_batch_size}")
+
+ embeddings = await model.get_embeddings(test_texts)
+
+ assert embeddings is not None, f"{model_name}: Embeddings is None"
+ assert isinstance(embeddings, list), f"{model_name}: Embeddings is not a list"
+ assert len(embeddings) == len(test_texts), f"{model_name}: Embeddings count mismatch"
+
+ # Check all embeddings are valid
+ for i, emb in enumerate(embeddings):
+ assert isinstance(emb, list), f"{model_name}: Embedding {i} is not a list"
+ assert len(emb) == model.dimensions, f"{model_name}: Embedding {i} dimension mismatch"
+
+ print("\n✓ Large batch embeddings generated successfully")
+ print(f" - Total texts: {len(test_texts)}")
+ print(f" - Total embeddings: {len(embeddings)}")
+ print(f" - Dimension: {len(embeddings[0])}")
+
+ await model.close()
+ print(f"✓ PASSED: {model_name} async large batch embeddings")
+
+
+async def test_async_single_node_embedding(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test asynchronous single VectorNode embedding."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Async Single VectorNode Embedding")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ node = VectorNode(
+ content="This is a test node for embedding.",
+ metadata={"test": "true"},
+ )
+
+ print(f"Input node content: {node.content}")
+ print(f"Initial vector: {node.vector}")
+
+ result_node = await model.get_node_embedding(node)
+
+ assert result_node is not None, f"{model_name}: Result node is None"
+ assert result_node.vector is not None, f"{model_name}: Node vector is None"
+ assert isinstance(result_node.vector, list), f"{model_name}: Vector is not a list"
+ assert len(result_node.vector) == model.dimensions, f"{model_name}: Vector dimension mismatch"
+
+ print("\n✓ Node embedding generated successfully")
+ print(f" - Vector dimension: {len(result_node.vector)}")
+ print(f" - First 5 values: {result_node.vector[:5]}")
+ print(f" - Metadata preserved: {result_node.metadata}")
+
+ await model.close()
+ print(f"✓ PASSED: {model_name} async single node embedding")
+
+
+async def test_async_batch_node_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test asynchronous batch VectorNode embeddings."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Async Batch VectorNode Embeddings")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ nodes = get_test_nodes()
+
+ print(f"Input: {len(nodes)} nodes")
+ for i, node in enumerate(nodes[:3], 1):
+ print(f" {i}. {node.content[:50]}...")
+
+ result_nodes = await model.get_node_embeddings(nodes)
+
+ assert result_nodes is not None, f"{model_name}: Result nodes is None"
+ assert isinstance(result_nodes, list), f"{model_name}: Result is not a list"
+ assert len(result_nodes) == len(nodes), f"{model_name}: Nodes count mismatch"
+
+ for i, node in enumerate(result_nodes):
+ assert node.vector is not None, f"{model_name}: Node {i} vector is None"
+ assert isinstance(node.vector, list), f"{model_name}: Node {i} vector is not a list"
+ assert len(node.vector) == model.dimensions, f"{model_name}: Node {i} dimension mismatch"
+ assert node.metadata is not None, f"{model_name}: Node {i} metadata is None"
+
+ print("\n✓ Batch node embeddings generated successfully")
+ print(f" - Count: {len(result_nodes)}")
+ print(f" - All vectors populated: {all(n.vector is not None for n in result_nodes)}")
+ print(f" - All metadata preserved: {all(n.metadata is not None for n in result_nodes)}")
+
+ await model.close()
+ print(f"✓ PASSED: {model_name} async batch node embeddings")
+
+
+async def test_async_large_batch_node_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test asynchronous large batch VectorNode embeddings with automatic batching."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Async Large Batch Node Embeddings")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ texts = get_large_batch_texts()
+ nodes = [VectorNode(content=text, metadata={"index": str(i)}) for i, text in enumerate(texts)]
+
+ print(f"Input: {len(nodes)} nodes")
+ print(f"Max batch size: {model.max_batch_size}")
+ print(f"Expected batches: {(len(nodes) + model.max_batch_size - 1) // model.max_batch_size}")
+
+ result_nodes = await model.get_node_embeddings(nodes)
+
+ assert result_nodes is not None, f"{model_name}: Result nodes is None"
+ assert len(result_nodes) == len(nodes), f"{model_name}: Nodes count mismatch"
+
+ # Check all nodes have embeddings
+ nodes_with_vectors = sum(1 for n in result_nodes if n.vector is not None)
+ print("\n✓ Large batch node embeddings generated successfully")
+ print(f" - Total nodes: {len(result_nodes)}")
+ print(f" - Nodes with vectors: {nodes_with_vectors}")
+ print(f" - Success rate: {nodes_with_vectors/len(result_nodes)*100:.1f}%")
+
+ assert nodes_with_vectors == len(nodes), f"{model_name}: Not all nodes have vectors"
+
+ await model.close()
+ print(f"✓ PASSED: {model_name} async large batch node embeddings")
+
+
+async def run_all_tests_for_model(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Run all tests for a specific embedding model class."""
+ print(f"\n\n{'#'*60}")
+ print(f"# Running all tests for: {model_name}")
+ print(f"{'#'*60}")
+
+ await test_async_single_embedding(model_class, model_name)
+ await test_async_batch_embeddings(model_class, model_name)
+ await test_async_large_batch_embeddings(model_class, model_name)
+ await test_async_single_node_embedding(model_class, model_name)
+ await test_async_batch_node_embeddings(model_class, model_name)
+ await test_async_large_batch_node_embeddings(model_class, model_name)
+
+ print(f"\n{'='*60}")
+ print(f"✓ All tests passed for {model_name}!")
+ print(f"{'='*60}")
+
+
+async def main():
+ """Main entry point for running tests."""
+ parser = argparse.ArgumentParser(
+ description="Run async embedding model tests",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ python test_embedding.py --openai # Test OpenAIEmbeddingModel only
+ python test_embedding.py --all # Test all embedding models
+ """,
+ )
+ parser.add_argument(
+ "--openai",
+ action="store_true",
+ help="Test OpenAIEmbeddingModel",
+ )
+ parser.add_argument(
+ "--all",
+ action="store_true",
+ help="Run tests for all available embedding models",
+ )
+
+ args = parser.parse_args()
+
+ # Determine which models to test
+ models_to_test = []
+
+ if args.openai:
+ models_to_test.append((OpenAIEmbeddingModel, "OpenAIEmbeddingModel"))
+ elif args.all:
+ models_to_test.append((OpenAIEmbeddingModel, "OpenAIEmbeddingModel"))
+ else:
+ # Default to all models if no argument provided
+ models_to_test = [(OpenAIEmbeddingModel, "OpenAIEmbeddingModel")]
+ print("No model specified, defaulting to all models")
+ print("Use --openai to test OpenAI specifically\n")
+
+ # Run tests for each model
+ for model_class, model_name in models_to_test:
+ try:
+ await run_all_tests_for_model(model_class, model_name)
+ except Exception as e:
+ print(f"\n✗ FAILED: {model_name} tests failed with error:")
+ print(f" {type(e).__name__}: {e}")
+ raise
+
+ # Final summary
+ print(f"\n\n{'#'*60}")
+ print("# TEST SUMMARY")
+ print(f"{'#'*60}")
+ print(f"✓ All tests passed for {len(models_to_test)} embedding model(s):")
+ for _, model_name in models_to_test:
+ print(f" - {model_name}")
+ print(f"{'#'*60}\n")
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
diff --git a/tests/test_embedding_sync.py b/tests/test_embedding_sync.py
new file mode 100644
index 00000000..361a42b3
--- /dev/null
+++ b/tests/test_embedding_sync.py
@@ -0,0 +1,348 @@
+"""
+Sync unit tests for Embedding classes (OpenAIEmbeddingModelSync) covering:
+- Sync single text embedding
+- Sync batch text embeddings
+- Sync large batch with automatic batching
+- Sync VectorNode embedding (single and batch)
+- Error handling and retries
+
+Usage:
+ python test_embedding_sync.py --openai # Test OpenAIEmbeddingModelSync only
+ python test_embedding_sync.py --all # Test all embedding models
+"""
+
+# flake8: noqa: E402
+# pylint: disable=C0413
+
+import argparse
+from typing import Type, List
+
+from reme_ai.core.utils import load_env
+
+load_env()
+
+from reme_ai.core.embedding import OpenAIEmbeddingModelSync, BaseEmbeddingModel
+from reme_ai.core.schema import VectorNode
+
+
+def get_embedding_model(model_class: Type[BaseEmbeddingModel]) -> BaseEmbeddingModel:
+ """Create and return an embedding model instance."""
+ return model_class(
+ model_name="text-embedding-v4",
+ dimensions=1024,
+ max_retries=2,
+ raise_exception=True,
+ )
+
+
+def get_test_texts() -> List[str]:
+ """Create test texts for embedding."""
+ return [
+ "The quick brown fox jumps over the lazy dog.",
+ "Machine learning is a subset of artificial intelligence.",
+ "Python is a popular programming language for data science.",
+ "Solar energy is a renewable source of power.",
+ "The capital of France is Paris.",
+ ]
+
+
+def get_large_batch_texts() -> List[str]:
+ """Create a large batch of test texts to test automatic batching."""
+ texts = []
+ topics = [
+ "Climate change and global warming",
+ "Artificial intelligence and machine learning",
+ "Renewable energy sources",
+ "Space exploration and astronomy",
+ "Medical research and healthcare",
+ "Financial markets and economics",
+ "Education and learning systems",
+ "Transportation and urban planning",
+ ]
+
+ for i, topic in enumerate(topics):
+ for j in range(3):
+ texts.append(f"Text {i*3+j+1}: This is a sample text about {topic}.")
+
+ return texts # 24 texts total
+
+
+def get_test_nodes() -> List[VectorNode]:
+ """Create test VectorNodes for embedding."""
+ texts = get_test_texts()
+ return [
+ VectorNode(
+ content=text,
+ metadata={"index": str(i), "category": "test"},
+ )
+ for i, text in enumerate(texts)
+ ]
+
+
+def test_sync_single_embedding(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test synchronous single text embedding."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Sync Single Text Embedding")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ test_text = "Hello, this is a test sentence for embedding."
+
+ print(f"Input text: {test_text}")
+
+ embedding = model.get_embedding_sync(test_text)
+
+ assert embedding is not None, f"{model_name}: Embedding is None"
+ assert isinstance(embedding, list), f"{model_name}: Embedding is not a list"
+ assert len(embedding) > 0, f"{model_name}: Empty embedding"
+ assert len(embedding) == model.dimensions, f"{model_name}: Embedding dimension mismatch"
+ assert all(isinstance(x, float) for x in embedding), f"{model_name}: Not all elements are floats"
+
+ print("\n✓ Embedding generated successfully")
+ print(f" - Dimension: {len(embedding)}")
+ print(f" - First 5 values: {embedding[:5]}")
+ print(f" - Value range: [{min(embedding):.4f}, {max(embedding):.4f}]")
+
+ model.close_sync()
+ print(f"✓ PASSED: {model_name} sync single embedding")
+
+
+def test_sync_batch_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test synchronous batch text embeddings."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Sync Batch Text Embeddings")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ test_texts = get_test_texts()
+
+ print(f"Input: {len(test_texts)} texts")
+ for i, text in enumerate(test_texts[:3], 1):
+ print(f" {i}. {text[:50]}...")
+
+ embeddings = model.get_embeddings_sync(test_texts)
+
+ assert embeddings is not None, f"{model_name}: Embeddings is None"
+ assert isinstance(embeddings, list), f"{model_name}: Embeddings is not a list"
+ assert len(embeddings) == len(test_texts), f"{model_name}: Embeddings count mismatch"
+
+ for i, emb in enumerate(embeddings):
+ assert isinstance(emb, list), f"{model_name}: Embedding {i} is not a list"
+ assert len(emb) == model.dimensions, f"{model_name}: Embedding {i} dimension mismatch"
+ assert all(isinstance(x, float) for x in emb), f"{model_name}: Embedding {i} has non-float values"
+
+ print("\n✓ Batch embeddings generated successfully")
+ print(f" - Count: {len(embeddings)}")
+ print(f" - Dimension: {len(embeddings[0])}")
+ print(f" - First embedding preview: {embeddings[0][:3]}...")
+
+ model.close_sync()
+ print(f"✓ PASSED: {model_name} sync batch embeddings")
+
+
+def test_sync_large_batch_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test synchronous large batch embeddings with automatic batching."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Sync Large Batch with Auto-Batching")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ test_texts = get_large_batch_texts()
+
+ print(f"Input: {len(test_texts)} texts")
+ print(f"Max batch size: {model.max_batch_size}")
+ print(f"Expected batches: {(len(test_texts) + model.max_batch_size - 1) // model.max_batch_size}")
+
+ embeddings = model.get_embeddings_sync(test_texts)
+
+ assert embeddings is not None, f"{model_name}: Embeddings is None"
+ assert isinstance(embeddings, list), f"{model_name}: Embeddings is not a list"
+ assert len(embeddings) == len(test_texts), f"{model_name}: Embeddings count mismatch"
+
+ # Check all embeddings are valid
+ for i, emb in enumerate(embeddings):
+ assert isinstance(emb, list), f"{model_name}: Embedding {i} is not a list"
+ assert len(emb) == model.dimensions, f"{model_name}: Embedding {i} dimension mismatch"
+
+ print("\n✓ Large batch embeddings generated successfully")
+ print(f" - Total texts: {len(test_texts)}")
+ print(f" - Total embeddings: {len(embeddings)}")
+ print(f" - Dimension: {len(embeddings[0])}")
+
+ model.close_sync()
+ print(f"✓ PASSED: {model_name} sync large batch embeddings")
+
+
+def test_sync_single_node_embedding(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test synchronous single VectorNode embedding."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Sync Single VectorNode Embedding")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ node = VectorNode(
+ content="This is a test node for embedding.",
+ metadata={"test": "true"},
+ )
+
+ print(f"Input node content: {node.content}")
+ print(f"Initial vector: {node.vector}")
+
+ result_node = model.get_node_embedding_sync(node)
+
+ assert result_node is not None, f"{model_name}: Result node is None"
+ assert result_node.vector is not None, f"{model_name}: Node vector is None"
+ assert isinstance(result_node.vector, list), f"{model_name}: Vector is not a list"
+ assert len(result_node.vector) == model.dimensions, f"{model_name}: Vector dimension mismatch"
+
+ print("\n✓ Node embedding generated successfully")
+ print(f" - Vector dimension: {len(result_node.vector)}")
+ print(f" - First 5 values: {result_node.vector[:5]}")
+ print(f" - Metadata preserved: {result_node.metadata}")
+
+ model.close_sync()
+ print(f"✓ PASSED: {model_name} sync single node embedding")
+
+
+def test_sync_batch_node_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test synchronous batch VectorNode embeddings."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Sync Batch VectorNode Embeddings")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ nodes = get_test_nodes()
+
+ print(f"Input: {len(nodes)} nodes")
+ for i, node in enumerate(nodes[:3], 1):
+ print(f" {i}. {node.content[:50]}...")
+
+ result_nodes = model.get_node_embeddings_sync(nodes)
+
+ assert result_nodes is not None, f"{model_name}: Result nodes is None"
+ assert isinstance(result_nodes, list), f"{model_name}: Result is not a list"
+ assert len(result_nodes) == len(nodes), f"{model_name}: Nodes count mismatch"
+
+ for i, node in enumerate(result_nodes):
+ assert node.vector is not None, f"{model_name}: Node {i} vector is None"
+ assert isinstance(node.vector, list), f"{model_name}: Node {i} vector is not a list"
+ assert len(node.vector) == model.dimensions, f"{model_name}: Node {i} dimension mismatch"
+ assert node.metadata is not None, f"{model_name}: Node {i} metadata is None"
+
+ print("\n✓ Batch node embeddings generated successfully")
+ print(f" - Count: {len(result_nodes)}")
+ print(f" - All vectors populated: {all(n.vector is not None for n in result_nodes)}")
+ print(f" - All metadata preserved: {all(n.metadata is not None for n in result_nodes)}")
+
+ model.close_sync()
+ print(f"✓ PASSED: {model_name} sync batch node embeddings")
+
+
+def test_sync_large_batch_node_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Test synchronous large batch VectorNode embeddings with automatic batching."""
+ print(f"\n{'='*60}")
+ print(f"Testing {model_name}: Sync Large Batch Node Embeddings")
+ print(f"{'='*60}")
+
+ model = get_embedding_model(model_class)
+ texts = get_large_batch_texts()
+ nodes = [VectorNode(content=text, metadata={"index": str(i)}) for i, text in enumerate(texts)]
+
+ print(f"Input: {len(nodes)} nodes")
+ print(f"Max batch size: {model.max_batch_size}")
+ print(f"Expected batches: {(len(nodes) + model.max_batch_size - 1) // model.max_batch_size}")
+
+ result_nodes = model.get_node_embeddings_sync(nodes)
+
+ assert result_nodes is not None, f"{model_name}: Result nodes is None"
+ assert len(result_nodes) == len(nodes), f"{model_name}: Nodes count mismatch"
+
+ # Check all nodes have embeddings
+ nodes_with_vectors = sum(1 for n in result_nodes if n.vector is not None)
+ print("\n✓ Large batch node embeddings generated successfully")
+ print(f" - Total nodes: {len(result_nodes)}")
+ print(f" - Nodes with vectors: {nodes_with_vectors}")
+ print(f" - Success rate: {nodes_with_vectors/len(result_nodes)*100:.1f}%")
+
+ assert nodes_with_vectors == len(nodes), f"{model_name}: Not all nodes have vectors"
+
+ model.close_sync()
+ print(f"✓ PASSED: {model_name} sync large batch node embeddings")
+
+
+def run_all_tests_for_model(model_class: Type[BaseEmbeddingModel], model_name: str):
+ """Run all tests for a specific embedding model class."""
+ print(f"\n\n{'#'*60}")
+ print(f"# Running all tests for: {model_name}")
+ print(f"{'#'*60}")
+
+ test_sync_single_embedding(model_class, model_name)
+ test_sync_batch_embeddings(model_class, model_name)
+ test_sync_large_batch_embeddings(model_class, model_name)
+ test_sync_single_node_embedding(model_class, model_name)
+ test_sync_batch_node_embeddings(model_class, model_name)
+ test_sync_large_batch_node_embeddings(model_class, model_name)
+
+ print(f"\n{'='*60}")
+ print(f"✓ All tests passed for {model_name}!")
+ print(f"{'='*60}")
+
+
+def main():
+ """Main entry point for running tests."""
+ parser = argparse.ArgumentParser(
+ description="Run sync embedding model tests",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ python test_embedding_sync.py --openai # Test OpenAIEmbeddingModelSync only
+ python test_embedding_sync.py --all # Test all embedding models
+ """,
+ )
+ parser.add_argument(
+ "--openai",
+ action="store_true",
+ help="Test OpenAIEmbeddingModelSync",
+ )
+ parser.add_argument(
+ "--all",
+ action="store_true",
+ help="Run tests for all available embedding models",
+ )
+
+ args = parser.parse_args()
+
+ # Determine which models to test
+ models_to_test = []
+
+ if args.openai:
+ models_to_test.append((OpenAIEmbeddingModelSync, "OpenAIEmbeddingModelSync"))
+ elif args.all:
+ models_to_test.append((OpenAIEmbeddingModelSync, "OpenAIEmbeddingModelSync"))
+ else:
+ # Default to all models if no argument provided
+ models_to_test = [(OpenAIEmbeddingModelSync, "OpenAIEmbeddingModelSync")]
+ print("No model specified, defaulting to all models")
+ print("Use --openai to test OpenAI specifically\n")
+
+ # Run tests for each model
+ for model_class, model_name in models_to_test:
+ try:
+ run_all_tests_for_model(model_class, model_name)
+ except Exception as e:
+ print(f"\n✗ FAILED: {model_name} tests failed with error:")
+ print(f" {type(e).__name__}: {e}")
+ raise
+
+ # Final summary
+ print(f"\n\n{'#'*60}")
+ print("# TEST SUMMARY")
+ print(f"{'#'*60}")
+ print(f"✓ All tests passed for {len(models_to_test)} embedding model(s):")
+ for _, model_name in models_to_test:
+ print(f" - {model_name}")
+ print(f"{'#'*60}\n")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/test_logo.py b/tests/test_logo.py
new file mode 100644
index 00000000..eeede81e
--- /dev/null
+++ b/tests/test_logo.py
@@ -0,0 +1,9 @@
+"""test logo"""
+
+from reme_ai.core.schema import ServiceConfig, MCPConfig
+
+if __name__ == "__main__":
+ from reme_ai.core.utils import print_logo
+
+ c = ServiceConfig(app_name="reme", backend="mcp", mcp=MCPConfig(transport="sse"))
+ print_logo(service_config=c)
diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py
new file mode 100644
index 00000000..d2fef40e
--- /dev/null
+++ b/tests/test_mcp_client.py
@@ -0,0 +1,127 @@
+"""Test module for demonstrating MCPClient functionality."""
+
+# pylint: disable=too-many-return-statements,too-many-statements
+
+import asyncio
+import json
+
+from reme_ai.core.utils import MCPClient
+
+
+async def main():
+ """Execute demonstration of the MCPClient."""
+ test_mcp = "test_mcp"
+ config_data = {
+ "mcpServers": {
+ test_mcp: {
+ "url": "http://127.0.0.1:8010/sse",
+ },
+ },
+ }
+
+ client = MCPClient(config_data)
+
+ try:
+ # List all available tools
+ print("=" * 50)
+ print("Listing available tools:")
+ print("=" * 50)
+ t_list = await client.list_tool_calls(test_mcp)
+ for t in t_list:
+ print(json.dumps(t, ensure_ascii=False, indent=2))
+
+ # Helper function to build default values
+ def build_default_value(param_info: dict) -> any:
+ """Build a default value for a parameter based on its schema."""
+ param_type = param_info.get("type", "string")
+
+ # Handle enum types - use the first enum value
+ if "enum" in param_info and param_info["enum"]:
+ return param_info["enum"][0]
+
+ # Handle different types
+ if param_type == "string":
+ return "example_string"
+ elif param_type == "number":
+ return 0.0
+ elif param_type == "integer":
+ return 0
+ elif param_type == "boolean":
+ return False
+ elif param_type == "array":
+ return []
+ elif param_type == "object":
+ # Recursively build nested objects
+ obj = {}
+ nested_properties = param_info.get("properties", {})
+ nested_required = param_info.get("required", [])
+
+ for nested_param_name in nested_required:
+ if nested_param_name in nested_properties:
+ nested_param_info = nested_properties[nested_param_name]
+ obj[nested_param_name] = build_default_value(nested_param_info)
+
+ return obj
+ else:
+ return None
+
+ # Call tools if available
+ if t_list:
+ # Execute the first two tools (or fewer if not enough tools available)
+ tools_to_execute = min(2, len(t_list))
+
+ for idx in range(tools_to_execute):
+ print("\n" + "=" * 50)
+ print(f"Calling tool #{idx + 1}:")
+ print("=" * 50)
+
+ # Get the tool's information
+ current_tool = t_list[idx]
+ tool_type = current_tool.get("type", "function")
+ tool_body = current_tool.get(tool_type, {})
+
+ tool_name = tool_body.get("name")
+ tool_description = tool_body.get("description", "")
+
+ # Prepare arguments based on the tool's input schema
+ tool_arguments = {}
+ parameters = tool_body.get("parameters", {})
+ properties = parameters.get("properties", {})
+ required = parameters.get("required", [])
+
+ # Build minimal arguments for required parameters
+ for param_name in required:
+ if param_name in properties:
+ param_info = properties[param_name]
+ tool_arguments[param_name] = build_default_value(param_info)
+
+ # Validate tool name exists
+ if not tool_name:
+ print("Error: Tool name not found in the tool definition")
+ print(f"Tool structure: {json.dumps(current_tool, ensure_ascii=False, indent=2)}")
+ else:
+ print(f"Tool name: {tool_name}")
+ print(f"Tool description: {tool_description}")
+ print(f"Arguments: {json.dumps(tool_arguments, ensure_ascii=False, indent=2)}")
+
+ # Call the tool
+ result = await client.call_tool(test_mcp, tool_name, tool_arguments, parse_text_result=True)
+
+ print("\n" + "-" * 50)
+ print("Tool call result:")
+ print("-" * 50)
+ print(f"Content: {result}")
+ if hasattr(result, "isError"):
+ print(f"Is Error: {result.isError}")
+ else:
+ print("\nNo tools available to call.")
+
+ except Exception as e:
+ print(f"Error occurred: {e}")
+ import traceback
+
+ traceback.print_exc()
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py
new file mode 100644
index 00000000..67f9c542
--- /dev/null
+++ b/tests/test_mcp_server.py
@@ -0,0 +1,125 @@
+"""Dynamic MCP server implementation with JSON-schema based tool registration."""
+
+from typing import Any
+
+from fastmcp import FastMCP
+from fastmcp.tools import FunctionTool
+
+from reme_ai.core.schema import ToolCall
+from reme_ai.core.utils import create_pydantic_model
+
+mcp = FastMCP("DynamicSchemaServer", port=8010)
+
+# Configuration including enum examples
+MODES_CONFIG = {
+ "register_user": ToolCall(
+ **{
+ "name": "register_user",
+ "description": "Register a new user with metadata, tags, and roles.",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "username": {"type": "string", "description": "Unique username"},
+ "role": {
+ "type": "string",
+ "enum": ["admin", "editor", "viewer"],
+ "description": "User access level",
+ },
+ "metadata": {
+ "type": "object",
+ "description": "User metadata",
+ "properties": {
+ "age": {"type": "integer"},
+ "location": {"type": "string"},
+ },
+ "required": ["age"],
+ },
+ "tags": {
+ "type": "array",
+ "description": "User tags",
+ "items": {
+ "type": "object",
+ "properties": {
+ "tag_id": {"type": "string"},
+ "level": {"type": "number"},
+ },
+ "required": ["tag_id"],
+ },
+ },
+ },
+ "required": ["username", "metadata", "role"],
+ },
+ },
+ ),
+ "create_order": ToolCall(
+ **{
+ "name": "create_order",
+ "description": "创建订单",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "order_id": {"type": "string", "description": "订单ID"},
+ "amount": {"type": "number", "description": "订单金额"},
+ "customer": {
+ "type": "object",
+ "description": "客户信息",
+ "properties": {
+ "name": {"type": "string", "description": "客户姓名"},
+ "email": {"type": "string", "description": "客户邮箱"},
+ "phone": {"type": "string", "description": "联系电话"},
+ },
+ "required": ["name", "email"],
+ },
+ },
+ "required": ["order_id", "customer"],
+ },
+ },
+ ),
+}
+
+
+async def core_handler(mode: str, **kwargs: Any) -> dict[str, Any]:
+ """Process dynamic tool requests and return execution results."""
+ print(f"Executing Mode: {mode}, Parameters: {kwargs}")
+ return {
+ "status": "success",
+ "mode": mode,
+ "received_data": kwargs,
+ }
+
+
+def register_dynamic_tools() -> None:
+ """Iterate over tool configurations and register them to the MCP instance."""
+ for mode_name, tool_call in MODES_CONFIG.items():
+ # Create Pydantic model from tool parameters
+ request_model = create_pydantic_model(tool_call.name, tool_call.parameters)
+
+ # Create execution function with closure to capture current mode and model
+ def create_tool_func(current_mode: str, model: type):
+ async def execute_tool(**kwargs: Any) -> dict[str, Any]:
+ # Validate and normalize input using Pydantic model
+ validated_data = model(**kwargs).model_dump(exclude_none=True)
+ return await core_handler(current_mode, **validated_data)
+
+ return execute_tool
+
+ tool_fn = create_tool_func(mode_name, request_model)
+
+ # Extract parameters schema
+ tool_call_schema = tool_call.simple_input_dump()
+ parameters = tool_call_schema[tool_call_schema["type"]]["parameters"]
+
+ # Create FunctionTool and register
+ tool = FunctionTool(
+ name=tool_call.name,
+ description=tool_call.description,
+ fn=tool_fn,
+ parameters=parameters,
+ )
+
+ mcp.add_tool(tool)
+
+
+if __name__ == "__main__":
+ register_dynamic_tools()
+ mcp.run(transport="sse")
diff --git a/tests/test_op_composition.py b/tests/test_op_composition.py
new file mode 100644
index 00000000..8d32b51a
--- /dev/null
+++ b/tests/test_op_composition.py
@@ -0,0 +1,325 @@
+"""
+Unit tests for BaseOp and operator composition (>>, <<, |).
+Tests asynchronous execution mode.
+"""
+
+import asyncio
+
+from reme_ai.core.op import BaseOp
+from reme_ai.core.schema import ToolCall, ToolAttr
+
+
+class AddOp(BaseOp):
+ """Simple operator that adds a value to a number in context."""
+
+ def __init__(self, value: int = 1, **kwargs):
+ super().__init__(**kwargs)
+ self.value = value
+
+ def _build_tool_call(self) -> ToolCall:
+ return ToolCall(
+ **{
+ "name": self.name,
+ "description": f"Add {self.value} to input",
+ "parameters": ToolAttr(
+ **{
+ "type": "object",
+ "properties": {
+ "number": {"type": "integer", "description": "Input number"},
+ },
+ "required": ["number"],
+ },
+ ),
+ },
+ )
+
+ async def execute(self):
+ """Async execution: add value to input number."""
+ self.context["number"] += self.value
+ self.output = self.context["number"]
+
+
+class MultiplyOp(BaseOp):
+ """Simple operator that multiplies a number in context."""
+
+ def __init__(self, factor: int = 2, **kwargs):
+ super().__init__(**kwargs)
+ self.factor = factor
+
+ def _build_tool_call(self) -> ToolCall:
+ return ToolCall(
+ **{
+ "name": self.name,
+ "description": f"Multiply by {self.factor}",
+ "parameters": ToolAttr(
+ **{
+ "type": "object",
+ "properties": {
+ "number": {"type": "integer", "description": "Input number"},
+ },
+ "required": ["number"],
+ },
+ ),
+ },
+ )
+
+ async def execute(self):
+ """Async execution: multiply input number."""
+ self.context["number"] *= self.factor
+ self.output = self.context["number"]
+
+
+class AppendOp(BaseOp):
+ """Operator that appends a value to a list in context."""
+
+ def __init__(self, value: str = "", **kwargs):
+ super().__init__(**kwargs)
+ self.value = value
+
+ def _build_tool_call(self) -> ToolCall:
+ return ToolCall(
+ **{
+ "name": self.name,
+ "description": f"Append {self.value} to list",
+ "parameters": ToolAttr(
+ **{
+ "type": "object",
+ "properties": {
+ "items": {"type": "array", "description": "List of items"},
+ },
+ "required": ["items"],
+ },
+ ),
+ },
+ )
+
+ async def execute(self):
+ """Async execution: append value to list."""
+ self.context["items"].append(self.value)
+ self.output = self.context["items"]
+
+
+async def test_basic_async_call():
+ """Test basic asynchronous operator execution."""
+ op = AddOp(value=5, name="add_5")
+ await op.call(number=10)
+ number = op.context["number"]
+ assert number == 15, f"Expected context result 15, got {number}"
+ print("✓ test_basic_async_call passed")
+
+
+async def test_sequential_composition_async():
+ """Test >> operator for sequential composition in async mode."""
+ add_op = AddOp(value=5, name="add_5")
+ multiply_op = MultiplyOp(factor=2, name="multiply_2")
+ composed = add_op >> multiply_op
+ await composed.call(number=10)
+
+ # (10 + 5) * 2 = 30
+ assert composed.context["number"] == 30, f"Expected 30, got {composed.context['number']}"
+ print("✓ test_sequential_composition_async passed")
+
+
+async def test_parallel_composition_async():
+ """Test | operator for parallel composition in async mode."""
+ append_a = AppendOp(value="A", name="append_a")
+ append_b = AppendOp(value="B", name="append_b")
+ append_c = AppendOp(value="C", name="append_c")
+
+ composed = append_a | append_b | append_c
+
+ await composed.call(items=[])
+
+ # All should append to the list
+ items = composed.context["items"]
+ assert len(items) == 3, f"Expected 3 items, got {len(items)}"
+ assert set(items) == {"A", "B", "C"}, f"Expected A,B,C, got {items}"
+ print("✓ test_parallel_composition_async passed")
+
+
+async def test_add_sub_ops_async():
+ """Test << operator for adding sub-operations in async mode."""
+ parent_op = BaseOp(name="parent")
+ child1 = AddOp(value=5, name="child1")
+ child2 = MultiplyOp(factor=2, name="child2")
+
+ _ = parent_op << child1
+ _ = parent_op << child2
+
+ assert len(parent_op.sub_ops) == 2, f"Expected 2 sub_ops, got {len(parent_op.sub_ops)}"
+ sub_op_names = [op.name for op in parent_op.sub_ops]
+ assert "child1" in sub_op_names, "child1 not in sub_ops"
+ assert "child2" in sub_op_names, "child2 not in sub_ops"
+ print("✓ test_add_sub_ops_async passed")
+
+
+async def test_add_sub_ops_dict():
+ """Test << operator with dictionary of operations."""
+ parent_op = BaseOp(name="parent")
+ ops_dict = {
+ "add": AddOp(value=5, name="add"),
+ "multiply": MultiplyOp(factor=2, name="multiply"),
+ }
+
+ _ = parent_op << ops_dict
+
+ assert len(parent_op.sub_ops) == 2, f"Expected 2 ops_dict, got {len(parent_op.sub_ops)}"
+ sub_op_names = [op.name for op in parent_op.sub_ops]
+ assert "add" in sub_op_names, "add not in ops_dict"
+ assert "multiply" in sub_op_names, "multiply not in ops_dict"
+ print("✓ test_add_sub_ops_dict passed")
+
+
+async def test_add_sub_ops_list():
+ """Test << operator with list of operations."""
+ parent_op = BaseOp(name="parent")
+ sub_ops = [
+ AddOp(value=5, name="add"),
+ MultiplyOp(factor=2, name="multiply"),
+ ]
+
+ _ = parent_op << sub_ops
+
+ assert len(parent_op.sub_ops) == 2, f"Expected 2 sub_ops, got {len(parent_op.sub_ops)}"
+ sub_op_names = [op.name for op in parent_op.sub_ops]
+ assert "add" in sub_op_names, "add not in sub_ops"
+ assert "multiply" in sub_op_names, "multiply not in sub_ops"
+ print("✓ test_add_sub_ops_list passed")
+
+
+async def test_mixed_composition_async():
+ """Test mixing >> and | operators in async mode."""
+ # (add_5 >> multiply_2) | (add_10 >> multiply_3)
+ seq1 = AddOp(value=5, name="add_5") >> MultiplyOp(factor=2, name="multiply_2")
+ seq2 = AddOp(value=10, name="add_10") >> MultiplyOp(factor=3, name="multiply_3")
+
+ composed = seq1 | seq2
+
+ await composed.call(number=10)
+
+ # Both sequences execute in parallel with shared context
+ # seq1: (10 + 5) * 2 = 30
+ # seq2: (30 + 10) * 3 = 120 (builds on seq1's result due to shared context)
+ # The exact result depends on execution order and timing
+ # With current implementation, result is 120
+ assert composed.context["number"] == 120, f"Expected 120, got {composed.context['number']}"
+ print("✓ test_mixed_composition_async passed")
+
+
+async def test_op_copy():
+ """Test operator copy functionality."""
+ original = AddOp(value=5, name="original")
+ copy_op = original.copy(name="copy")
+
+ assert copy_op.name == "copy", f"Expected name 'copy', got {copy_op.name}"
+ assert copy_op.value == 5, f"Expected value 5, got {copy_op.value}"
+ assert copy_op is not original, "Copy should be a different object"
+ print("✓ test_op_copy passed")
+
+
+async def test_input_mapping():
+ """Test input_mapping parameter."""
+ op = AddOp(
+ value=5,
+ name="add_5",
+ input_mapping={"x": "number"}, # Map x to number
+ )
+
+ await op.call(x=10) # Input is 'x' not 'number'
+
+ assert op.context["number"] == 15, f"Expected number=15, got {op.context['number']}"
+ print("✓ test_input_mapping passed")
+
+
+async def test_output_mapping():
+ """Test output_mapping parameter."""
+ op = AddOp(
+ value=5,
+ name="add_5",
+ output_mapping={"number": "final_result"}, # Map number to final_result
+ )
+
+ await op.call(number=10)
+
+ assert op.context["number"] == 15, f"Expected number=15, got {op.context['number']}"
+ assert op.context["final_result"] == 15, f"Expected final_result=15, got {op.context['final_result']}"
+ print("✓ test_output_mapping passed")
+
+
+async def test_validation_missing_required():
+ """Test that missing required inputs raise an error."""
+ op = AddOp(value=5, name="add_5", raise_exception=True)
+
+ try:
+ await op.call() # Missing 'number' field
+ assert False, "Should have raised ValueError for missing required input"
+ except ValueError as e:
+ assert "number" in str(e), f"Expected error about 'number', got: {e}"
+ print("✓ test_validation_missing_required passed")
+
+
+async def test_max_retries():
+ """Test max_retries parameter with failing operation."""
+
+ class FailingOp(BaseOp):
+ """An operation that always fails."""
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ self.attempt_count = 0
+
+ def _build_tool_call(self) -> ToolCall:
+ return ToolCall(
+ **{
+ "name": self.name,
+ "description": "Always fails",
+ "parameters": ToolAttr(**{"type": "object", "properties": {}}),
+ "output": ToolAttr(
+ **{
+ "type": "object",
+ "properties": {
+ "result": ToolAttr(**{"type": "string", "description": "Result"}),
+ },
+ },
+ ),
+ },
+ )
+
+ async def execute(self):
+ self.attempt_count += 1
+ raise RuntimeError(f"Attempt {self.attempt_count} failed")
+
+ op = FailingOp(max_retries=3, name="failing")
+
+ await op.call()
+
+ assert op.attempt_count == 3, f"Expected 3 attempts, got {op.attempt_count}"
+ print("✓ test_max_retries passed")
+
+
+async def async_main():
+ """Run all async tests."""
+ await test_basic_async_call()
+ await test_sequential_composition_async()
+ await test_parallel_composition_async()
+ await test_add_sub_ops_async()
+ await test_add_sub_ops_dict()
+ await test_add_sub_ops_list()
+ await test_mixed_composition_async()
+ await test_op_copy()
+ await test_input_mapping()
+ await test_output_mapping()
+ await test_validation_missing_required()
+ await test_max_retries()
+
+
+if __name__ == "__main__":
+ print("Running BaseOp composition tests...\n")
+
+ # Async tests
+ print("=== Asynchronous Tests ===")
+ asyncio.run(async_main())
+
+ print("\n" + "=" * 50)
+ print("All tests passed! ✓")
+ print("=" * 50)
diff --git a/tests/test_token_counter.py b/tests/test_token_counter.py
new file mode 100644
index 00000000..3c44a298
--- /dev/null
+++ b/tests/test_token_counter.py
@@ -0,0 +1,511 @@
+"""
+Unit tests for TokenCounter classes covering:
+- BaseTokenCounter (rule-based estimation)
+- OpenAITokenCounter (tiktoken-based)
+- HFTokenCounter (HuggingFace tokenizer-based)
+
+Usage:
+ python test_token_counter.py --base # Test BaseTokenCounter only
+ python test_token_counter.py --openai # Test OpenAITokenCounter only
+ python test_token_counter.py --hf # Test HFTokenCounter only
+ python test_token_counter.py --all # Test all token counters
+"""
+
+import argparse
+from typing import Type, List
+
+from reme_ai.core.enumeration import Role
+from reme_ai.core.schema import Message, ToolCall
+from reme_ai.core.token_counter import BaseTokenCounter, OpenAITokenCounter, HFTokenCounter
+
+
+def get_token_counter(counter_class: Type[BaseTokenCounter], **kwargs) -> BaseTokenCounter:
+ """Create and return a token counter instance."""
+ default_kwargs = {
+ "model_name": "gpt-4o",
+ }
+ default_kwargs.update(kwargs)
+ return counter_class(**default_kwargs)
+
+
+def get_test_messages() -> List[Message]:
+ """Create test messages for token counting."""
+ return [
+ Message(role=Role.SYSTEM, content="You are a helpful assistant."),
+ Message(role=Role.USER, content="Hello, how are you today?"),
+ Message(role=Role.ASSISTANT, content="I'm doing well, thank you for asking! How can I help you?"),
+ Message(role=Role.USER, content="Can you explain what machine learning is?"),
+ Message(
+ role=Role.ASSISTANT,
+ content="Machine learning is a subset of artificial intelligence that enables computers to "
+ "learn from data without being explicitly programmed.",
+ ),
+ ]
+
+
+def get_chinese_messages() -> List[Message]:
+ """Create test messages with Chinese content."""
+ return [
+ Message(role=Role.SYSTEM, content="你是一个有帮助的助手。"),
+ Message(role=Role.USER, content="你好，今天天气怎么样？"),
+ Message(role=Role.ASSISTANT, content="今天天气很好，阳光明媚，适合外出活动。"),
+ Message(role=Role.USER, content="能给我推荐一些好看的电影吗？"),
+ Message(role=Role.ASSISTANT, content="当然可以！我推荐《肖申克的救赎》、《阿甘正传》和《泰坦尼克号》。"),
+ ]
+
+
+def get_mixed_messages() -> List[Message]:
+ """Create test messages with mixed English and Chinese content."""
+ return [
+ Message(role=Role.SYSTEM, content="You are a bilingual assistant. 你是一个双语助手。"),
+ Message(role=Role.USER, content="What is AI? 什么是人工智能？"),
+ Message(
+ role=Role.ASSISTANT,
+ content="AI (Artificial Intelligence) 是人工智能的英文缩写，它是计算机科学的一个分支。",
+ ),
+ ]
+
+
+def get_messages_with_reasoning() -> List[Message]:
+ """Create test messages with reasoning content."""
+ return [
+ Message(role=Role.USER, content="What is 2 + 2?"),
+ Message(
+ role=Role.ASSISTANT,
+ content="The answer is 4.",
+ reasoning_content="Let me think about this step by step. 2 + 2 "
+ "equals 4 because addition combines two quantities.",
+ ),
+ ]
+
+
+def get_test_tools() -> List[ToolCall]:
+ """Create test tool calls for token counting."""
+ return [
+ ToolCall(
+ **{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get the current weather for a specified location.",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {
+ "type": "string",
+ "description": "The city and country, e.g., 'Beijing, China'",
+ },
+ "unit": {
+ "type": "string",
+ "description": "Temperature unit: 'celsius' or 'fahrenheit'",
+ "enum": ["celsius", "fahrenheit"],
+ },
+ },
+ "required": ["location"],
+ },
+ },
+ },
+ ),
+ ToolCall(
+ **{
+ "type": "function",
+ "function": {
+ "name": "search_web",
+ "description": "Search the web for information.",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "query": {
+ "type": "string",
+ "description": "The search query",
+ },
+ "num_results": {
+ "type": "integer",
+ "description": "Number of results to return",
+ },
+ },
+ "required": ["query"],
+ },
+ },
+ },
+ ),
+ ]
+
+
+def get_tool_call_messages() -> List[Message]:
+ """Create messages with tool call responses."""
+ return [
+ Message(role=Role.USER, content="What's the weather in Beijing?"),
+ Message(
+ role=Role.ASSISTANT,
+ content="",
+ tool_calls=[
+ ToolCall(
+ id="call_123",
+ name="get_weather",
+ arguments='{"location": "Beijing, China", "unit": "celsius"}',
+ ),
+ ],
+ ),
+ Message(
+ role=Role.TOOL,
+ content='{"temperature": 25, "condition": "sunny", "humidity": 60}',
+ tool_call_id="call_123",
+ ),
+ Message(
+ role=Role.ASSISTANT,
+ content="The weather in Beijing is sunny with a temperature of 25°C and 60% humidity.",
+ ),
+ ]
+
+
+def test_basic_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Test basic token counting with simple messages."""
+ print(f"\n{'=' * 60}")
+ print(f"Testing {counter_name}: Basic Token Count")
+ print(f"{'=' * 60}")
+
+ counter = get_token_counter(counter_class, **kwargs)
+ messages = get_test_messages()
+
+ print(f"Input: {len(messages)} messages")
+ for i, msg in enumerate(messages, 1):
+ content_preview = msg.content[:50] + "..." if len(msg.content) > 50 else msg.content
+ print(f" {i}. [{msg.role.value}] {content_preview}")
+
+ token_count = counter.count_token(messages)
+
+ assert token_count is not None, f"{counter_name}: Token count is None"
+ assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+ assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+ print(f"\n✓ Token count: {token_count}")
+ print(f"✓ PASSED: {counter_name} basic token count")
+
+
+def test_chinese_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Test token counting with Chinese content."""
+ print(f"\n{'=' * 60}")
+ print(f"Testing {counter_name}: Chinese Token Count")
+ print(f"{'=' * 60}")
+
+ counter = get_token_counter(counter_class, **kwargs)
+ messages = get_chinese_messages()
+
+ print(f"Input: {len(messages)} Chinese messages")
+ for i, msg in enumerate(messages, 1):
+ content_preview = msg.content[:30] + "..." if len(msg.content) > 30 else msg.content
+ print(f" {i}. [{msg.role.value}] {content_preview}")
+
+ token_count = counter.count_token(messages)
+
+ assert token_count is not None, f"{counter_name}: Token count is None"
+ assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+ assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+ print(f"\n✓ Token count: {token_count}")
+ print(f"✓ PASSED: {counter_name} Chinese token count")
+
+
+def test_mixed_language_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Test token counting with mixed English and Chinese content."""
+ print(f"\n{'=' * 60}")
+ print(f"Testing {counter_name}: Mixed Language Token Count")
+ print(f"{'=' * 60}")
+
+ counter = get_token_counter(counter_class, **kwargs)
+ messages = get_mixed_messages()
+
+ print(f"Input: {len(messages)} mixed language messages")
+ for i, msg in enumerate(messages, 1):
+ content_preview = msg.content[:40] + "..." if len(msg.content) > 40 else msg.content
+ print(f" {i}. [{msg.role.value}] {content_preview}")
+
+ token_count = counter.count_token(messages)
+
+ assert token_count is not None, f"{counter_name}: Token count is None"
+ assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+ assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+ print(f"\n✓ Token count: {token_count}")
+ print(f"✓ PASSED: {counter_name} mixed language token count")
+
+
+def test_reasoning_content_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Test token counting with reasoning content."""
+ print(f"\n{'=' * 60}")
+ print(f"Testing {counter_name}: Reasoning Content Token Count")
+ print(f"{'=' * 60}")
+
+ counter = get_token_counter(counter_class, **kwargs)
+ messages = get_messages_with_reasoning()
+
+ print(f"Input: {len(messages)} messages with reasoning content")
+ for i, msg in enumerate(messages, 1):
+ print(f" {i}. [{msg.role.value}] content: {msg.content[:30]}...")
+ if msg.reasoning_content:
+ print(f" reasoning: {msg.reasoning_content[:30]}...")
+
+ token_count = counter.count_token(messages)
+
+ assert token_count is not None, f"{counter_name}: Token count is None"
+ assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+ assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+ print(f"\n✓ Token count: {token_count}")
+ print(f"✓ PASSED: {counter_name} reasoning content token count")
+
+
+def test_token_count_with_tools(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Test token counting with tool definitions."""
+ print(f"\n{'=' * 60}")
+ print(f"Testing {counter_name}: Token Count with Tools")
+ print(f"{'=' * 60}")
+
+ counter = get_token_counter(counter_class, **kwargs)
+ messages = get_test_messages()[:2]
+ tools = get_test_tools()
+
+ print(f"Input: {len(messages)} messages, {len(tools)} tools")
+ for tool in tools:
+ print(f" Tool: {tool.name} - {tool.description[:40]}...")
+
+ token_count = counter.count_token(messages, tools=tools)
+
+ assert token_count is not None, f"{counter_name}: Token count is None"
+ assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+ assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+ # Token count with tools should be higher than without
+ token_count_no_tools = counter.count_token(messages)
+ assert token_count > token_count_no_tools, f"{counter_name}: Token count with tools should be higher"
+
+ print(f"\n✓ Token count without tools: {token_count_no_tools}")
+ print(f"✓ Token count with tools: {token_count}")
+ print(f"✓ Tools added {token_count - token_count_no_tools} tokens")
+ print(f"✓ PASSED: {counter_name} token count with tools")
+
+
+def test_tool_call_messages_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Test token counting with messages containing tool calls."""
+ print(f"\n{'=' * 60}")
+ print(f"Testing {counter_name}: Tool Call Messages Token Count")
+ print(f"{'=' * 60}")
+
+ counter = get_token_counter(counter_class, **kwargs)
+ messages = get_tool_call_messages()
+
+ print(f"Input: {len(messages)} messages with tool calls")
+ for i, msg in enumerate(messages, 1):
+ if msg.tool_calls:
+ print(f" {i}. [{msg.role.value}] tool_calls: {[tc.name for tc in msg.tool_calls]}")
+ else:
+ content_preview = msg.content[:40] + "..." if len(msg.content) > 40 else msg.content
+ print(f" {i}. [{msg.role.value}] {content_preview}")
+
+ token_count = counter.count_token(messages)
+
+ assert token_count is not None, f"{counter_name}: Token count is None"
+ assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+ assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+ print(f"\n✓ Token count: {token_count}")
+ print(f"✓ PASSED: {counter_name} tool call messages token count")
+
+
+def test_empty_messages(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Test token counting with empty message list."""
+ print(f"\n{'=' * 60}")
+ print(f"Testing {counter_name}: Empty Messages")
+ print(f"{'=' * 60}")
+
+ # HFTokenCounter does not support empty message list (apply_chat_template requires at least one message)
+ if counter_class == HFTokenCounter:
+ print("⊘ SKIPPED: HFTokenCounter does not support empty message list")
+ return
+
+ counter = get_token_counter(counter_class, **kwargs)
+ messages: List[Message] = []
+
+ print("Input: 0 messages")
+
+ token_count = counter.count_token(messages)
+
+ assert token_count is not None, f"{counter_name}: Token count is None"
+ assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+
+ # OpenAITokenCounter adds 3 tokens for reply priming even with empty messages
+ if counter_class == OpenAITokenCounter:
+ assert token_count == 3, f"{counter_name}: Empty messages should have 3 tokens (reply priming)"
+ print(f"\n✓ Token count: {token_count} (includes 3 tokens for reply priming)")
+ else:
+ assert token_count == 0, f"{counter_name}: Empty messages should have 0 tokens"
+ print(f"\n✓ Token count: {token_count}")
+
+ print(f"✓ PASSED: {counter_name} empty messages")
+
+
+def test_single_message(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Test token counting with a single message."""
+ print(f"\n{'=' * 60}")
+ print(f"Testing {counter_name}: Single Message")
+ print(f"{'=' * 60}")
+
+ counter = get_token_counter(counter_class, **kwargs)
+ messages = [Message(role=Role.USER, content="Hello!")]
+
+ print(f"Input: 1 message - '{messages[0].content}'")
+
+ token_count = counter.count_token(messages)
+
+ assert token_count is not None, f"{counter_name}: Token count is None"
+ assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+ assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+ print(f"\n✓ Token count: {token_count}")
+ print(f"✓ PASSED: {counter_name} single message")
+
+
+def test_long_content(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Test token counting with long content."""
+ print(f"\n{'=' * 60}")
+ print(f"Testing {counter_name}: Long Content")
+ print(f"{'=' * 60}")
+
+ counter = get_token_counter(counter_class, **kwargs)
+
+ # Create a long message
+ long_text = "This is a test sentence. " * 100
+ messages = [Message(role=Role.USER, content=long_text)]
+
+ print(f"Input: 1 message with {len(long_text)} characters")
+
+ token_count = counter.count_token(messages)
+
+ assert token_count is not None, f"{counter_name}: Token count is None"
+ assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+ assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+ # Long content should have more tokens
+ short_messages = [Message(role=Role.USER, content="This is a test sentence.")]
+ short_token_count = counter.count_token(short_messages)
+ assert token_count > short_token_count, f"{counter_name}: Long content should have more tokens"
+
+ print(f"\n✓ Short content token count: {short_token_count}")
+ print(f"✓ Long content token count: {token_count}")
+ print(f"✓ PASSED: {counter_name} long content")
+
+
+def run_all_tests_for_counter(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+ """Run all tests for a specific token counter class."""
+ print(f"\n\n{'#' * 60}")
+ print(f"# Running all tests for: {counter_name}")
+ print(f"{'#' * 60}")
+
+ test_basic_token_count(counter_class, counter_name, **kwargs)
+ test_chinese_token_count(counter_class, counter_name, **kwargs)
+ test_mixed_language_token_count(counter_class, counter_name, **kwargs)
+ test_reasoning_content_token_count(counter_class, counter_name, **kwargs)
+ test_token_count_with_tools(counter_class, counter_name, **kwargs)
+ test_tool_call_messages_token_count(counter_class, counter_name, **kwargs)
+ test_empty_messages(counter_class, counter_name, **kwargs)
+ test_single_message(counter_class, counter_name, **kwargs)
+ test_long_content(counter_class, counter_name, **kwargs)
+
+ print(f"\n{'=' * 60}")
+ print(f"✓ All tests passed for {counter_name}!")
+ print(f"{'=' * 60}")
+
+
+def main():
+ """Main entry point for running tests."""
+ parser = argparse.ArgumentParser(
+ description="Run token counter tests",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ python test_token_counter.py --base # Test BaseTokenCounter only
+ python test_token_counter.py --openai # Test OpenAITokenCounter only
+ python test_token_counter.py --hf # Test HFTokenCounter only
+ python test_token_counter.py --all # Test all token counters
+ """,
+ )
+ parser.add_argument(
+ "--base",
+ action="store_true",
+ help="Test BaseTokenCounter (rule-based)",
+ )
+ parser.add_argument(
+ "--openai",
+ action="store_true",
+ help="Test OpenAITokenCounter (tiktoken-based)",
+ )
+ parser.add_argument(
+ "--hf",
+ action="store_true",
+ help="Test HFTokenCounter (HuggingFace tokenizer-based)",
+ )
+ parser.add_argument(
+ "--hf-model",
+ type=str,
+ default="Qwen/Qwen2.5-0.5B-Instruct",
+ help="HuggingFace model name for HFTokenCounter (default: Qwen/Qwen2.5-0.5B-Instruct)",
+ )
+ parser.add_argument(
+ "--all",
+ action="store_true",
+ help="Run tests for all available token counters",
+ )
+
+ args = parser.parse_args()
+
+ # Determine which counters to test
+ counters_to_test = []
+
+ if args.all:
+ counters_to_test.append((BaseTokenCounter, "BaseTokenCounter", {}))
+ counters_to_test.append((OpenAITokenCounter, "OpenAITokenCounter", {}))
+ counters_to_test.append(
+ (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}),
+ )
+ else:
+ if args.base:
+ counters_to_test.append((BaseTokenCounter, "BaseTokenCounter", {}))
+ if args.openai:
+ counters_to_test.append((OpenAITokenCounter, "OpenAITokenCounter", {}))
+ if args.hf:
+ counters_to_test.append(
+ (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}),
+ )
+
+ if not counters_to_test:
+ # Default to all counters if no argument provided
+ counters_to_test = [
+ (BaseTokenCounter, "BaseTokenCounter", {}),
+ (OpenAITokenCounter, "OpenAITokenCounter", {}),
+ (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}),
+ ]
+ print("No counter specified, defaulting to test all counters")
+ print("Use --base/--openai/--hf to test specific ones\n")
+
+ # Run tests for each counter
+ for counter_class, counter_name, kwargs in counters_to_test:
+ try:
+ run_all_tests_for_counter(counter_class, counter_name, **kwargs)
+ except Exception as e:
+ print(f"\n✗ FAILED: {counter_name} tests failed with error:")
+ print(f" {type(e).__name__}: {e}")
+ raise
+
+ # Final summary
+ print(f"\n\n{'#' * 60}")
+ print("# TEST SUMMARY")
+ print(f"{'#' * 60}")
+ print(f"✓ All tests passed for {len(counters_to_test)} token counter(s):")
+ for _, counter_name, _ in counters_to_test:
+ print(f" - {counter_name}")
+ print(f"{'#' * 60}\n")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py
new file mode 100644
index 00000000..76de21e1
--- /dev/null
+++ b/tests/test_vector_store.py
@@ -0,0 +1,1417 @@
+# pylint: disable=too-many-lines
+"""Unified test suite for vector store implementations.
+
+This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore,
+PGVectorStore, QdrantVectorStore, and ChromaVectorStore implementations. Tests can be
+run for specific vector stores or all implementations.
+
+Usage:
+ python test_vector_store.py --local # Test LocalVectorStore only
+ python test_vector_store.py --es # Test ESVectorStore only
+ python test_vector_store.py --pgvector # Test PGVectorStore only
+ python test_vector_store.py --qdrant # Test QdrantVectorStore only
+ python test_vector_store.py --chroma # Test ChromaVectorStore only
+ python test_vector_store.py --all # Test all vector stores
+
+"""
+
+import argparse
+import asyncio
+import shutil
+from pathlib import Path
+from typing import List
+
+from loguru import logger
+
+from reme_ai.core.embedding import OpenAIEmbeddingModel
+from reme_ai.core.schema import VectorNode
+from reme_ai.core.vector_store import (
+ BaseVectorStore,
+ ChromaVectorStore,
+ LocalVectorStore,
+ ESVectorStore,
+ PGVectorStore,
+ QdrantVectorStore,
+)
+
+
+# ==================== Configuration ====================
+
+
+class TestConfig:
+ """Configuration for test execution."""
+
+ # LocalVectorStore settings
+ LOCAL_ROOT_PATH = "./test_vector_store_local"
+
+ # ESVectorStore settings
+ ES_HOSTS = "http://11.160.132.46:8200"
+ ES_BASIC_AUTH = None # Set to ("username", "password") if authentication is required
+
+ # QdrantVectorStore settings
+ QDRANT_PATH = None # "./test_vector_store_qdrant" # For local mode
+ QDRANT_HOST = None # Set to host address for remote mode (e.g., "localhost")
+ QDRANT_PORT = None # Set to port for remote mode (e.g., 6333)
+ QDRANT_URL = "http://11.160.132.46:6333" # Alternative to host/port (e.g., http://localhost:6333)
+ QDRANT_API_KEY = None # Set for Qdrant Cloud authentication
+
+ # PGVectorStore settings
+ PG_DSN = "postgresql://localhost/postgres" # PostgreSQL connection string
+ PG_MIN_SIZE = 1 # Minimum connections in pool
+ PG_MAX_SIZE = 5 # Maximum connections in pool
+ PG_USE_HNSW = True # Use HNSW index for faster search
+ PG_USE_DISKANN = False # Use DiskANN index (requires vectorscale extension)
+
+ # ChromaVectorStore settings
+ CHROMA_PATH = "./test_vector_store_chroma" # For local persistent mode
+ CHROMA_HOST = None # Set to host address for remote mode (e.g., "localhost")
+ CHROMA_PORT = None # Set to port for remote mode (e.g., 8000)
+ CHROMA_API_KEY = None # Set for ChromaDB Cloud authentication
+ CHROMA_TENANT = None # Set for ChromaDB Cloud tenant
+ CHROMA_DATABASE = None # Set for ChromaDB Cloud database
+
+ # Embedding model settings
+ EMBEDDING_MODEL_NAME = "text-embedding-v4"
+ EMBEDDING_DIMENSIONS = 64
+
+ # Test collection naming
+ TEST_COLLECTION_PREFIX = "test_vector_store"
+
+
+# ==================== Sample Data Generator ====================
+
+
+class SampleDataGenerator:
+ """Generator for sample test data."""
+
+ @staticmethod
+ def create_sample_nodes(prefix: str = "") -> List[VectorNode]:
+ """Create sample VectorNode instances for testing.
+
+ Args:
+ prefix: Optional prefix for vector_id to avoid conflicts
+
+ Returns:
+ List[VectorNode]: List of sample nodes with diverse metadata
+ """
+ id_prefix = f"{prefix}_" if prefix else ""
+ return [
+ VectorNode(
+ vector_id=f"{id_prefix}node1",
+ content="Artificial intelligence is a technology that simulates human intelligence.",
+ metadata={
+ "node_type": "tech",
+ "category": "AI",
+ "source": "research",
+ "priority": "high",
+ "year": "2023",
+ "department": "engineering",
+ "language": "english",
+ "status": "published",
+ },
+ ),
+ VectorNode(
+ vector_id=f"{id_prefix}node2",
+ content="Machine learning is a subset of artificial intelligence.",
+ metadata={
+ "node_type": "tech",
+ "category": "ML",
+ "source": "research",
+ "priority": "high",
+ "year": "2022",
+ "department": "engineering",
+ "language": "english",
+ "status": "published",
+ },
+ ),
+ VectorNode(
+ vector_id=f"{id_prefix}node3",
+ content="Deep learning uses neural networks with multiple layers.",
+ metadata={
+ "node_type": "tech_new",
+ "category": "DL",
+ "source": "blog",
+ "priority": "medium",
+ "year": "2024",
+ "department": "marketing",
+ "language": "chinese",
+ "status": "draft",
+ },
+ ),
+ VectorNode(
+ vector_id=f"{id_prefix}node4",
+ content="I love eating delicious seafood, especially fresh fish.",
+ metadata={
+ "node_type": "food",
+ "category": "preference",
+ "source": "personal",
+ "priority": "low",
+ "year": "2023",
+ "department": "lifestyle",
+ "language": "english",
+ "status": "published",
+ },
+ ),
+ VectorNode(
+ vector_id=f"{id_prefix}node5",
+ content="Natural language processing enables computers to understand human language.",
+ metadata={
+ "node_type": "tech",
+ "category": "NLP",
+ "source": "research",
+ "priority": "high",
+ "year": "2024",
+ "department": "engineering",
+ "language": "english",
+ "status": "review",
+ },
+ ),
+ ]
+
+
+# ==================== Vector Store Factory ====================
+
+
+def get_store_type(store: BaseVectorStore) -> str:
+ """Get the type identifier of a vector store instance.
+
+ Args:
+ store: Vector store instance
+
+ Returns:
+ str: Type identifier ("local", "es", "pgvector", "qdrant", or "chroma")
+ """
+ if isinstance(store, LocalVectorStore):
+ return "local"
+ elif isinstance(store, QdrantVectorStore):
+ return "qdrant"
+ elif isinstance(store, ESVectorStore):
+ return "es"
+ elif isinstance(store, PGVectorStore):
+ return "pgvector"
+ elif isinstance(store, ChromaVectorStore):
+ return "chroma"
+ else:
+ raise ValueError(f"Unknown vector store type: {type(store)}")
+
+
+def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStore:
+ """Create a vector store instance based on type.
+
+ Args:
+ store_type: Type of vector store ("local", "es", or "qdrant")
+ collection_name: Name of the collection
+
+ Returns:
+ BaseVectorStore: Initialized vector store instance
+ """
+ config = TestConfig()
+
+ # Initialize embedding model
+ embedding_model = OpenAIEmbeddingModel(
+ model_name=config.EMBEDDING_MODEL_NAME,
+ dimensions=config.EMBEDDING_DIMENSIONS,
+ )
+
+ if store_type == "local":
+ return LocalVectorStore(
+ collection_name=collection_name,
+ embedding_model=embedding_model,
+ root_path=config.LOCAL_ROOT_PATH,
+ )
+ elif store_type == "es":
+ return ESVectorStore(
+ collection_name=collection_name,
+ embedding_model=embedding_model,
+ hosts=config.ES_HOSTS,
+ basic_auth=config.ES_BASIC_AUTH,
+ )
+ elif store_type == "qdrant":
+ return QdrantVectorStore(
+ collection_name=collection_name,
+ embedding_model=embedding_model,
+ path=config.QDRANT_PATH,
+ host=config.QDRANT_HOST,
+ port=config.QDRANT_PORT,
+ url=config.QDRANT_URL,
+ api_key=config.QDRANT_API_KEY,
+ distance="cosine",
+ on_disk=False,
+ )
+ elif store_type == "pgvector":
+ return PGVectorStore(
+ collection_name=collection_name,
+ embedding_model=embedding_model,
+ dsn=config.PG_DSN,
+ min_size=config.PG_MIN_SIZE,
+ max_size=config.PG_MAX_SIZE,
+ use_hnsw=config.PG_USE_HNSW,
+ use_diskann=config.PG_USE_DISKANN,
+ )
+ elif store_type == "chroma":
+ return ChromaVectorStore(
+ collection_name=collection_name,
+ embedding_model=embedding_model,
+ path=config.CHROMA_PATH,
+ host=config.CHROMA_HOST,
+ port=config.CHROMA_PORT,
+ api_key=config.CHROMA_API_KEY,
+ tenant=config.CHROMA_TENANT,
+ database=config.CHROMA_DATABASE,
+ )
+ else:
+ raise ValueError(f"Unknown store type: {store_type}")
+
+
+# ==================== Test Functions ====================
+
+
+async def test_create_collection(store: BaseVectorStore, _store_name: str):
+ """Test collection creation."""
+ logger.info("=" * 20 + " CREATE COLLECTION TEST " + "=" * 20)
+
+ # Clean up if exists
+ collections = await store.list_collections()
+ if store.collection_name in collections:
+ await store.delete_collection(store.collection_name)
+ logger.info(f"Cleaned up existing collection: {store.collection_name}")
+
+ # Create collection
+ await store.create_collection(store.collection_name)
+
+ # Verify creation
+ collections = await store.list_collections()
+ assert store.collection_name in collections, "Collection should exist after creation"
+ logger.info(f"✓ Created collection: {store.collection_name}")
+
+
+async def test_insert(store: BaseVectorStore, _store_name: str) -> List[VectorNode]:
+ """Test node insertion (single and batch)."""
+ logger.info("=" * 20 + " INSERT TEST " + "=" * 20)
+
+ # Test single node insertion
+ single_node = VectorNode(
+ vector_id="test_single_insert",
+ content="This is a single node insertion test",
+ metadata={"test_type": "single_insert"},
+ )
+ await store.insert(single_node)
+ logger.info("✓ Inserted single node")
+
+ # Test batch insertion
+ sample_nodes = SampleDataGenerator.create_sample_nodes("test")
+ await store.insert(sample_nodes)
+ logger.info(f"✓ Batch inserted {len(sample_nodes)} nodes")
+
+ # Verify total insertions
+ all_nodes = await store.list(limit=20)
+ assert len(all_nodes) >= len(sample_nodes) + 1, "Should have at least sample nodes + single node"
+ logger.info(f"✓ Total nodes in collection: {len(all_nodes)}")
+
+ return sample_nodes
+
+
+async def test_search(store: BaseVectorStore, _store_name: str):
+ """Test basic vector search."""
+ logger.info("=" * 20 + " SEARCH TEST " + "=" * 20)
+
+ results = await store.search(
+ query="What is artificial intelligence?",
+ limit=3,
+ )
+
+ logger.info(f"Search returned {len(results)} results")
+ for i, r in enumerate(results, 1):
+ score = r.metadata.get("_score", "N/A")
+ logger.info(f" Result {i}: {r.content[:60]}... (score: {score})")
+
+ assert len(results) > 0, "Search should return results"
+ logger.info("✓ Basic search test passed")
+
+
+async def test_search_with_single_filter(store: BaseVectorStore, _store_name: str):
+ """Test vector search with single metadata filter."""
+ logger.info("=" * 20 + " SINGLE FILTER SEARCH TEST " + "=" * 20)
+
+ # Test single value filter
+ filters = {"node_type": "tech"}
+ results = await store.search(
+ query="What is artificial intelligence?",
+ limit=5,
+ filters=filters,
+ )
+
+ logger.info(f"Filtered search (node_type=tech) returned {len(results)} results")
+ for i, r in enumerate(results, 1):
+ node_type = r.metadata.get("node_type")
+ logger.info(f" Result {i}: type={node_type}, content={r.content[:50]}...")
+ assert node_type == "tech", "Result should have node_type='tech'"
+
+ logger.info("✓ Single filter search test passed")
+
+
+async def test_search_with_list_filter(store: BaseVectorStore, _store_name: str):
+ """Test vector search with list filter (IN operation)."""
+ logger.info("=" * 20 + " LIST FILTER SEARCH TEST " + "=" * 20)
+
+ # Test list filter (IN operation)
+ filters = {"node_type": ["tech", "tech_new"]}
+ results = await store.search(
+ query="What is artificial intelligence?",
+ limit=5,
+ filters=filters,
+ )
+
+ logger.info(f"Filtered search (node_type IN [tech, tech_new]) returned {len(results)} results")
+ for i, r in enumerate(results, 1):
+ node_type = r.metadata.get("node_type")
+ logger.info(f" Result {i}: type={node_type}, content={r.content[:50]}...")
+ assert node_type in ["tech", "tech_new"], "Result should have node_type in [tech, tech_new]"
+
+ logger.info("✓ List filter search test passed")
+
+
+async def test_search_with_multiple_filters(store: BaseVectorStore, _store_name: str):
+ """Test vector search with multiple metadata filters (AND operation)."""
+ logger.info("=" * 20 + " MULTIPLE FILTERS SEARCH TEST " + "=" * 20)
+
+ # Test multiple filters (AND operation)
+ filters = {
+ "node_type": ["tech", "tech_new"],
+ "source": "research",
+ }
+ results = await store.search(
+ query="What is artificial intelligence?",
+ limit=5,
+ filters=filters,
+ )
+
+ logger.info(
+ f"Multi-filter search (node_type IN [tech, tech_new] AND source=research) " f"returned {len(results)} results",
+ )
+ for i, r in enumerate(results, 1):
+ node_type = r.metadata.get("node_type")
+ source = r.metadata.get("source")
+ logger.info(f" Result {i}: type={node_type}, source={source}, content={r.content[:40]}...")
+ assert node_type in ["tech", "tech_new"], "Result should have node_type in [tech, tech_new]"
+ assert source == "research", "Result should have source='research'"
+
+ logger.info("✓ Multiple filters search test passed")
+
+
+async def test_get_by_id(store: BaseVectorStore, _store_name: str):
+ """Test retrieving nodes by vector_id (single and batch)."""
+ logger.info("=" * 20 + " GET BY ID TEST " + "=" * 20)
+
+ # Test single ID retrieval
+ target_id = "test_node1"
+ result = await store.get(target_id)
+
+ assert isinstance(result, VectorNode), "Should return a VectorNode for single ID"
+ assert result.vector_id == target_id, f"Result should have vector_id={target_id}"
+ logger.info(f"✓ Retrieved single node: {result.vector_id}")
+
+ # Test batch retrieval (small batch)
+ target_ids = ["test_node1", "test_node2"]
+ results = await store.get(target_ids)
+
+ assert isinstance(results, list), "Should return a list for multiple IDs"
+ assert len(results) == 2, f"Should return 2 results, got {len(results)}"
+ result_ids = {r.vector_id for r in results}
+ assert result_ids == set(target_ids), f"Result IDs should match {target_ids}"
+ logger.info(f"✓ Batch retrieved {len(results)} nodes")
+
+ # Test larger batch retrieval
+ large_batch_ids = ["test_node1", "test_node2", "test_node3", "test_node5"]
+ large_results = await store.get(large_batch_ids)
+ assert isinstance(large_results, list), "Should return a list for batch IDs"
+ assert len(large_results) >= 3, "Should return at least 3 results"
+ logger.info(f"✓ Large batch retrieved {len(large_results)} nodes")
+
+
+async def test_list_all(store: BaseVectorStore, _store_name: str):
+ """Test listing all nodes in collection."""
+ logger.info("=" * 20 + " LIST ALL TEST " + "=" * 20)
+
+ results = await store.list(limit=10)
+
+ logger.info(f"Collection contains {len(results)} nodes")
+ for i, node in enumerate(results, 1):
+ logger.info(f" Node {i}: id={node.vector_id}, content={node.content[:50]}...")
+
+ assert len(results) > 0, "Collection should contain nodes"
+ logger.info("✓ List all nodes test passed")
+
+
+async def test_list_with_filters(store: BaseVectorStore, _store_name: str):
+ """Test listing nodes with metadata filters."""
+ logger.info("=" * 20 + " LIST WITH FILTERS TEST " + "=" * 20)
+
+ filters = {"category": "AI"}
+ results = await store.list(filters=filters, limit=10)
+
+ logger.info(f"Filtered list (category=AI) returned {len(results)} nodes")
+ for i, node in enumerate(results, 1):
+ category = node.metadata.get("category")
+ logger.info(f" Node {i}: category={category}, id={node.vector_id}")
+ assert category == "AI", "All nodes should have category=AI"
+
+ logger.info("✓ List with filters test passed")
+
+
+async def test_update(store: BaseVectorStore, _store_name: str):
+ """Test updating existing nodes (single and batch)."""
+ logger.info("=" * 20 + " UPDATE TEST " + "=" * 20)
+
+ # Test single node update
+ updated_node = VectorNode(
+ vector_id="test_node2",
+ content="Machine learning is a powerful subset of AI that learns from data.",
+ metadata={
+ "node_type": "tech",
+ "category": "ML",
+ "updated": "true",
+ "update_timestamp": "2024-12-26",
+ },
+ )
+
+ await store.update(updated_node)
+
+ # Verify single update
+ result = await store.get("test_node2")
+ assert "updated" in result.metadata, "Updated metadata should be present"
+ logger.info(f"✓ Updated single node: {result.vector_id}")
+ logger.info(f" New content: {result.content[:60]}...")
+
+ # Test batch update (update multiple nodes at once)
+ batch_update_nodes = [
+ VectorNode(
+ vector_id="test_node1",
+ content="Artificial intelligence is evolving rapidly with new breakthroughs.",
+ metadata={
+ "node_type": "tech",
+ "category": "AI",
+ "batch_updated": "true",
+ "update_timestamp": "2024-12-31",
+ },
+ ),
+ VectorNode(
+ vector_id="test_node3",
+ content="Deep learning revolutionizes neural network architectures.",
+ metadata={
+ "node_type": "tech_new",
+ "category": "DL",
+ "batch_updated": "true",
+ "update_timestamp": "2024-12-31",
+ },
+ ),
+ ]
+
+ await store.update(batch_update_nodes)
+ logger.info(f"✓ Batch updated {len(batch_update_nodes)} nodes")
+
+ # Verify batch updates
+ results = await store.get(["test_node1", "test_node3"])
+ for r in results:
+ assert r.metadata.get("batch_updated") == "true", f"Node {r.vector_id} should have batch_updated metadata"
+ logger.info(f"✓ Verified batch update for {len(results)} nodes")
+
+
+async def test_delete(store: BaseVectorStore, _store_name: str):
+ """Test deleting nodes (single and batch)."""
+ logger.info("=" * 20 + " DELETE TEST " + "=" * 20)
+
+ # Test single node deletion
+ node_to_delete = "test_node4"
+ await store.delete(node_to_delete)
+
+ # Verify single deletion - try to get the deleted node
+ try:
+ result = await store.get(node_to_delete)
+ # If result is empty list or None, deletion was successful
+ if isinstance(result, list):
+ assert len(result) == 0, "Deleted node should not be retrievable"
+ else:
+ assert result is None, "Deleted node should not be retrievable"
+ except Exception:
+ pass # Expected if node doesn't exist
+
+ logger.info(f"✓ Deleted single node: {node_to_delete}")
+
+ # Test batch deletion - first insert some nodes to delete
+ batch_delete_nodes = [
+ VectorNode(
+ vector_id=f"delete_test_{i}",
+ content=f"Node to be deleted {i}",
+ metadata={"test_type": "delete_batch"},
+ )
+ for i in range(5)
+ ]
+ await store.insert(batch_delete_nodes)
+ logger.info(f"✓ Inserted {len(batch_delete_nodes)} nodes for batch delete test")
+
+ # Batch delete
+ delete_ids = [f"delete_test_{i}" for i in range(5)]
+ await store.delete(delete_ids)
+ logger.info(f"✓ Batch deleted {len(delete_ids)} nodes")
+
+ # Verify batch deletion
+ try:
+ results = await store.get(delete_ids)
+ if isinstance(results, list):
+ assert len(results) == 0, "All deleted nodes should not be retrievable"
+ except Exception:
+ pass # Expected if nodes don't exist
+ logger.info("✓ Verified batch deletion")
+
+
+async def test_copy_collection(store: BaseVectorStore, store_name: str):
+ """Test copying a collection."""
+ logger.info("=" * 20 + " COPY COLLECTION TEST " + "=" * 20)
+
+ config = TestConfig()
+ copy_collection_name = f"{config.TEST_COLLECTION_PREFIX}_{store_name}_copy"
+
+ # Elasticsearch and PostgreSQL require lowercase table/index names
+ store_type = get_store_type(store)
+ if store_type in ("es", "pgvector"):
+ copy_collection_name = copy_collection_name.lower()
+
+ # Clean up if exists
+ collections = await store.list_collections()
+ if copy_collection_name in collections:
+ await store.delete_collection(copy_collection_name)
+
+ # Copy collection
+ await store.copy_collection(copy_collection_name)
+
+ # Verify copy
+ collections = await store.list_collections()
+ assert copy_collection_name in collections, "Copied collection should exist"
+ logger.info(f"✓ Copied collection to: {copy_collection_name}")
+
+ # Verify content in copied collection
+ copied_store = create_vector_store(store_type, copy_collection_name)
+ copied_nodes = await copied_store.list()
+ logger.info(f"✓ Copied collection has {len(copied_nodes)} nodes")
+ await copied_store.close()
+
+ # Clean up copied collection
+ await store.delete_collection(copy_collection_name)
+ logger.info("✓ Cleaned up copied collection")
+
+
+async def test_list_collections(store: BaseVectorStore, _store_name: str):
+ """Test listing all collections."""
+ logger.info("=" * 20 + " LIST COLLECTIONS TEST " + "=" * 20)
+
+ collections = await store.list_collections()
+
+ logger.info(f"Found {len(collections)} collections")
+ config = TestConfig()
+ test_collections = [c for c in collections if c.startswith(config.TEST_COLLECTION_PREFIX)]
+ logger.info(f" Test collections: {test_collections}")
+
+ assert store.collection_name in collections, "Main test collection should be listed"
+ logger.info("✓ List collections test passed")
+
+
+async def test_delete_collection(store: BaseVectorStore, _store_name: str):
+ """Test deleting a collection."""
+ logger.info("=" * 20 + " DELETE COLLECTION TEST " + "=" * 20)
+
+ await store.delete_collection(store.collection_name)
+
+ # Verify deletion
+ collections = await store.list_collections()
+ assert store.collection_name not in collections, "Collection should not exist after deletion"
+ logger.info(f"✓ Deleted collection: {store.collection_name}")
+
+
+async def test_cosine_similarity(store_name: str):
+ """Test manual cosine similarity calculation (LocalVectorStore only)."""
+ if store_name != "LocalVectorStore":
+ logger.info("=" * 20 + " COSINE SIMILARITY TEST (SKIPPED) " + "=" * 20)
+ logger.info("⊘ Skipped: Only applicable to LocalVectorStore")
+ return
+
+ logger.info("=" * 20 + " COSINE SIMILARITY TEST " + "=" * 20)
+
+ vec1 = [1.0, 0.0, 0.0]
+ vec2 = [0.0, 1.0, 0.0]
+ vec3 = [1.0, 0.0, 0.0]
+
+ # Test perpendicular vectors (similarity = 0)
+ sim1 = LocalVectorStore._cosine_similarity(vec1, vec2) # pylint: disable=protected-access
+ logger.info(f"Similarity between perpendicular vectors: {sim1:.4f}")
+ assert abs(sim1) < 0.0001, "Perpendicular vectors should have similarity close to 0"
+
+ # Test identical vectors (similarity = 1)
+ sim2 = LocalVectorStore._cosine_similarity(vec1, vec3) # pylint: disable=protected-access
+ logger.info(f"Similarity between identical vectors: {sim2:.4f}")
+ assert abs(sim2 - 1.0) < 0.0001, "Identical vectors should have similarity close to 1"
+
+ # Test with real-world like vectors
+ vec4 = [0.5, 0.5, 0.5]
+ vec5 = [0.6, 0.4, 0.5]
+ sim3 = LocalVectorStore._cosine_similarity(vec4, vec5) # pylint: disable=protected-access
+ logger.info(f"Similarity between similar vectors: {sim3:.4f}")
+ assert sim3 > 0.9, "Similar vectors should have high similarity"
+
+ logger.info("✓ Cosine similarity tests passed")
+
+
+async def test_batch_operations(store: BaseVectorStore, _store_name: str):
+ """Test large-scale batch insert, update, and delete operations.
+
+ This test validates the efficiency and correctness of batch operations
+ by processing 100 nodes at once, which is more realistic for production use.
+ """
+ logger.info("=" * 20 + " BATCH OPERATIONS TEST " + "=" * 20)
+
+ # Create a large batch of nodes (100 nodes)
+ batch_nodes = []
+ for i in range(100):
+ batch_nodes.append(
+ VectorNode(
+ vector_id=f"batch_node_{i}",
+ content=f"This is batch test content number {i} about various topics in technology and science.",
+ metadata={
+ "batch_id": str(i // 10), # Group into batches of 10
+ "index": str(i),
+ "category": ["tech", "science", "business"][i % 3],
+ "priority": ["high", "medium", "low"][i % 3],
+ },
+ ),
+ )
+
+ # Batch insert
+ await store.insert(batch_nodes)
+ logger.info(f"✓ Inserted {len(batch_nodes)} nodes in batch")
+
+ # Verify batch insert
+ results = await store.list(limit=150)
+ assert len(results) >= 100, f"Should have at least 100 nodes, got {len(results)}"
+ logger.info(f"✓ Verified batch insert: {len(results)} total nodes")
+
+ # Batch update (update first 20 nodes)
+ update_nodes = []
+ for i in range(20):
+ update_nodes.append(
+ VectorNode(
+ vector_id=f"batch_node_{i}",
+ content=f"UPDATED: This is updated batch content {i}",
+ metadata={
+ "batch_id": str(i // 10),
+ "index": str(i),
+ "updated": "true",
+ "update_timestamp": "2024-12-31",
+ },
+ ),
+ )
+
+ await store.update(update_nodes)
+ logger.info(f"✓ Updated {len(update_nodes)} nodes in batch")
+
+ # Verify updates
+ updated_results = await store.list(filters={"updated": "true"}, limit=50)
+ assert len(updated_results) >= 20, "Should have at least 20 updated nodes"
+ logger.info(f"✓ Verified batch update: {len(updated_results)} updated nodes")
+
+ # Batch delete (delete nodes with batch_id >= 5)
+ delete_ids = [f"batch_node_{i}" for i in range(50, 100)]
+ await store.delete(delete_ids)
+ logger.info(f"✓ Deleted {len(delete_ids)} nodes in batch")
+
+ # Verify deletions
+ remaining = await store.list(limit=150)
+ batch_nodes_remaining = [n for n in remaining if n.vector_id.startswith("batch_node_")]
+ assert len(batch_nodes_remaining) <= 50, "Should have at most 50 batch nodes remaining"
+ logger.info(f"✓ Verified batch delete: {len(batch_nodes_remaining)} nodes remaining")
+
+
+async def test_complex_metadata_queries(store: BaseVectorStore, _store_name: str):
+ """Test complex metadata filtering with nested conditions."""
+ logger.info("=" * 20 + " COMPLEX METADATA QUERIES TEST " + "=" * 20)
+
+ # Insert nodes with rich metadata
+ complex_nodes = [
+ VectorNode(
+ vector_id="complex_1",
+ content="Advanced neural networks for computer vision applications",
+ metadata={
+ "domain": "AI",
+ "subdomain": "computer_vision",
+ "year": "2024",
+ "citations": "150",
+ "impact_factor": "high",
+ "tags": "neural_networks,vision,deep_learning",
+ },
+ ),
+ VectorNode(
+ vector_id="complex_2",
+ content="Natural language processing with transformer models",
+ metadata={
+ "domain": "AI",
+ "subdomain": "nlp",
+ "year": "2023",
+ "citations": "200",
+ "impact_factor": "high",
+ "tags": "transformers,nlp,language_models",
+ },
+ ),
+ VectorNode(
+ vector_id="complex_3",
+ content="Reinforcement learning for robotics control",
+ metadata={
+ "domain": "AI",
+ "subdomain": "robotics",
+ "year": "2024",
+ "citations": "80",
+ "impact_factor": "medium",
+ "tags": "reinforcement_learning,robotics,control",
+ },
+ ),
+ VectorNode(
+ vector_id="complex_4",
+ content="Quantum computing algorithms and applications",
+ metadata={
+ "domain": "quantum",
+ "subdomain": "algorithms",
+ "year": "2024",
+ "citations": "50",
+ "impact_factor": "medium",
+ "tags": "quantum,algorithms,computing",
+ },
+ ),
+ ]
+
+ await store.insert(complex_nodes)
+ logger.info(f"✓ Inserted {len(complex_nodes)} nodes with complex metadata")
+
+ # Test 1: Multiple field filters with list values
+ filters_1 = {
+ "domain": "AI",
+ "year": ["2023", "2024"],
+ "impact_factor": "high",
+ }
+ results_1 = await store.search(
+ query="artificial intelligence research",
+ limit=10,
+ filters=filters_1,
+ )
+ logger.info(f"Test 1 - AI + high impact + recent years: {len(results_1)} results")
+ for r in results_1:
+ assert r.metadata.get("domain") == "AI"
+ assert r.metadata.get("impact_factor") == "high"
+ assert r.metadata.get("year") in ["2023", "2024"]
+
+ # Test 2: List filter with multiple subdomains
+ filters_2 = {
+ "subdomain": ["nlp", "computer_vision"],
+ }
+ results_2 = await store.search(
+ query="deep learning applications",
+ limit=10,
+ filters=filters_2,
+ )
+ logger.info(f"Test 2 - NLP or Computer Vision: {len(results_2)} results")
+ for r in results_2:
+ assert r.metadata.get("subdomain") in ["nlp", "computer_vision"]
+
+ # Test 3: Year-based filtering
+ filters_3 = {
+ "year": "2024",
+ }
+ results_3 = await store.list(filters=filters_3, limit=10)
+ logger.info(f"Test 3 - Year 2024 only: {len(results_3)} results")
+ for r in results_3:
+ assert r.metadata.get("year") == "2024"
+
+ logger.info("✓ Complex metadata queries test passed")
+
+
+async def test_edge_cases(store: BaseVectorStore, _store_name: str):
+ """Test edge cases and boundary conditions."""
+ logger.info("=" * 20 + " EDGE CASES TEST " + "=" * 20)
+
+ # Test 1: Empty content
+ edge_node_1 = VectorNode(
+ vector_id="edge_empty_content",
+ content="",
+ metadata={"type": "empty"},
+ )
+ try:
+ await store.insert([edge_node_1])
+ logger.info("✓ Handled empty content node")
+ except Exception as e:
+ logger.info(f"⊘ Empty content not supported: {e}")
+
+ # Test 2: Very long content
+ edge_node_2 = VectorNode(
+ vector_id="edge_long_content",
+ content="A" * 10000, # 10k characters
+ metadata={"type": "long_content"},
+ )
+ await store.insert([edge_node_2])
+ result = await store.get("edge_long_content")
+ assert len(result.content) == 10000
+ logger.info("✓ Handled very long content (10k chars)")
+
+ # Test 3: Special characters in content
+ edge_node_3 = VectorNode(
+ vector_id="edge_special_chars",
+ content="Special chars: @#$%^&*()[]{}|\\;:'\",.<>?/~`+=−×÷",
+ metadata={"type": "special_chars"},
+ )
+ await store.insert([edge_node_3])
+ result = await store.get("edge_special_chars")
+ assert "@#$%^&*()" in result.content
+ logger.info("✓ Handled special characters in content")
+
+ # Test 4: Unicode and emoji content
+ edge_node_4 = VectorNode(
+ vector_id="edge_unicode",
+ content="Unicode test: 你好世界 🌍 مرحبا العالم Привет мир",
+ metadata={"type": "unicode", "language": "multi"},
+ )
+ await store.insert([edge_node_4])
+ result = await store.get("edge_unicode")
+ assert "你好世界" in result.content
+ assert "🌍" in result.content
+ logger.info("✓ Handled unicode and emoji content")
+
+ # Test 5: Search with empty query
+ try:
+ results = await store.search(query="", limit=5)
+ logger.info(f"✓ Empty query returned {len(results)} results")
+ except Exception as e:
+ logger.info(f"⊘ Empty query not supported: {e}")
+
+ # Test 6: Search with very high limit
+ results = await store.search(query="test", limit=1000)
+ logger.info(f"✓ High limit search returned {len(results)} results")
+
+ # Test 7: Get non-existent ID
+ result = await store.get("non_existent_id_12345")
+ if isinstance(result, list):
+ assert len(result) == 0, "Non-existent ID should return empty list"
+ else:
+ assert result is None, "Non-existent ID should return None"
+ logger.info("✓ Handled non-existent ID gracefully")
+
+ # Test 8: Metadata with empty string values
+ edge_node_5 = VectorNode(
+ vector_id="edge_empty_metadata",
+ content="Testing empty string values in metadata",
+ metadata={"field1": "value1", "field2": "", "field3": "value3"},
+ )
+ await store.insert([edge_node_5])
+ logger.info("✓ Handled empty string values in metadata")
+
+ logger.info("✓ Edge cases test passed")
+
+
+async def test_concurrent_operations(store: BaseVectorStore, _store_name: str):
+ """Test concurrent read/write operations."""
+ logger.info("=" * 20 + " CONCURRENT OPERATIONS TEST " + "=" * 20)
+
+ # Prepare concurrent insert nodes
+ concurrent_nodes = [
+ VectorNode(
+ vector_id=f"concurrent_{i}",
+ content=f"Concurrent test content {i}",
+ metadata={"thread_id": str(i % 5), "index": str(i)},
+ )
+ for i in range(50)
+ ]
+
+ # Test concurrent inserts
+ insert_tasks = []
+ for i in range(0, 50, 10):
+ batch = concurrent_nodes[i : i + 10]
+ insert_tasks.append(store.insert(batch))
+
+ await asyncio.gather(*insert_tasks)
+ logger.info("✓ Completed concurrent inserts")
+
+ # Test concurrent searches
+ search_tasks = [store.search(query=f"concurrent test {i}", limit=5) for i in range(10)]
+ search_results = await asyncio.gather(*search_tasks)
+ logger.info(f"✓ Completed {len(search_results)} concurrent searches")
+
+ # Test concurrent reads
+ get_tasks = [store.get(f"concurrent_{i}") for i in range(0, 50, 5)]
+ get_results = await asyncio.gather(*get_tasks)
+ logger.info(f"✓ Completed {len(get_results)} concurrent reads")
+
+ # Test batch updates (using batch update instead of concurrent individual updates)
+ update_nodes = [
+ VectorNode(
+ vector_id=f"concurrent_{i}",
+ content=f"UPDATED concurrent content {i}",
+ metadata={"thread_id": str(i % 5), "updated": "true"},
+ )
+ for i in range(0, 20, 2)
+ ]
+ await store.update(update_nodes)
+ logger.info(f"✓ Completed batch update of {len(update_nodes)} nodes")
+
+ logger.info("✓ Concurrent operations test passed")
+
+
+async def test_search_relevance_ranking(store: BaseVectorStore, _store_name: str):
+ """Test search result relevance and ranking."""
+ logger.info("=" * 20 + " SEARCH RELEVANCE RANKING TEST " + "=" * 20)
+
+ # Insert nodes with varying relevance
+ relevance_nodes = [
+ VectorNode(
+ vector_id="relevance_exact",
+ content="Machine learning is a subset of artificial intelligence focused on learning from data.",
+ metadata={"relevance": "exact"},
+ ),
+ VectorNode(
+ vector_id="relevance_high",
+ content="Artificial intelligence and machine learning are transforming technology.",
+ metadata={"relevance": "high"},
+ ),
+ VectorNode(
+ vector_id="relevance_medium",
+ content="Deep learning uses neural networks for pattern recognition.",
+ metadata={"relevance": "medium"},
+ ),
+ VectorNode(
+ vector_id="relevance_low",
+ content="Software engineering best practices for code quality.",
+ metadata={"relevance": "low"},
+ ),
+ VectorNode(
+ vector_id="relevance_none",
+ content="Cooking recipes for delicious Italian pasta dishes.",
+ metadata={"relevance": "none"},
+ ),
+ ]
+
+ await store.insert(relevance_nodes)
+ logger.info(f"✓ Inserted {len(relevance_nodes)} nodes with varying relevance")
+
+ # Search with specific query
+ query = "What is machine learning and artificial intelligence?"
+ results = await store.search(query=query, limit=5)
+
+ logger.info(f"Search results for: '{query}'")
+ for i, result in enumerate(results, 1):
+ score = result.metadata.get("_score", "N/A")
+ relevance = result.metadata.get("relevance", "unknown")
+ logger.info(f" {i}. [{relevance}] score={score}: {result.content[:60]}...")
+
+ # Verify that more relevant results appear first
+ if len(results) >= 2:
+ # The exact match should be in top results
+ top_ids = [r.vector_id for r in results[:3]]
+ assert (
+ "relevance_exact" in top_ids or "relevance_high" in top_ids
+ ), "Most relevant results should appear in top 3"
+ logger.info("✓ Relevance ranking verified")
+
+ # Test with different query
+ query2 = "neural networks deep learning"
+ results2 = await store.search(query=query2, limit=5)
+ logger.info(f"\nSearch results for: '{query2}'")
+ for i, result in enumerate(results2, 1):
+ score = result.metadata.get("_score", "N/A")
+ logger.info(f" {i}. score={score}: {result.content[:60]}...")
+
+ logger.info("✓ Search relevance ranking test passed")
+
+
+async def test_metadata_statistics(store: BaseVectorStore, _store_name: str):
+ """Test metadata aggregation and statistics."""
+ logger.info("=" * 20 + " METADATA STATISTICS TEST " + "=" * 20)
+
+ # Get all nodes and analyze metadata
+ all_nodes = await store.list(limit=500)
+ logger.info(f"Total nodes in collection: {len(all_nodes)}")
+
+ # Count by category
+ category_counts = {}
+ for node in all_nodes:
+ category = node.metadata.get("category", "unknown")
+ category_counts[category] = category_counts.get(category, 0) + 1
+
+ logger.info("Category distribution:")
+ for category, count in sorted(category_counts.items()):
+ logger.info(f" {category}: {count}")
+
+ # Count by node_type
+ type_counts = {}
+ for node in all_nodes:
+ node_type = node.metadata.get("node_type", "unknown")
+ type_counts[node_type] = type_counts.get(node_type, 0) + 1
+
+ logger.info("Node type distribution:")
+ for node_type, count in sorted(type_counts.items()):
+ logger.info(f" {node_type}: {count}")
+
+ # Verify we can filter by each category
+ for category in category_counts:
+ if category != "unknown":
+ filtered = await store.list(filters={"category": category}, limit=100)
+ logger.info(f"✓ Filter by category '{category}': {len(filtered)} results")
+
+ logger.info("✓ Metadata statistics test passed")
+
+
+async def test_update_metadata_only(store: BaseVectorStore, _store_name: str):
+ """Test updating only metadata without changing content."""
+ logger.info("=" * 20 + " UPDATE METADATA ONLY TEST " + "=" * 20)
+
+ # Get an existing node
+ original = await store.get("test_node1")
+ original_content = original.content
+
+ # Update with same content but different metadata
+ updated_node = VectorNode(
+ vector_id="test_node1",
+ content=original_content, # Keep same content
+ metadata={
+ **original.metadata,
+ "metadata_updated": "true",
+ "update_count": "1",
+ "last_modified": "2024-12-31",
+ },
+ )
+
+ await store.update(updated_node)
+ logger.info("✓ Updated metadata without changing content")
+
+ # Verify update
+ result = await store.get("test_node1")
+ assert result.content == original_content, "Content should remain unchanged"
+ assert result.metadata.get("metadata_updated") == "true", "Metadata should be updated"
+ logger.info("✓ Verified metadata-only update")
+
+ # Update metadata again
+ updated_node_2 = VectorNode(
+ vector_id="test_node1",
+ content=original_content,
+ metadata={
+ **result.metadata,
+ "update_count": "2",
+ "last_modified": "2024-12-31T12:00:00",
+ },
+ )
+ await store.update(updated_node_2)
+
+ result_2 = await store.get("test_node1")
+ assert result_2.metadata.get("update_count") == "2", "Metadata should be updated again"
+ logger.info("✓ Multiple metadata updates successful")
+
+ logger.info("✓ Update metadata only test passed")
+
+
+async def test_filter_combinations(store: BaseVectorStore, _store_name: str):
+ """Test various filter combinations and edge cases."""
+ logger.info("=" * 20 + " FILTER COMBINATIONS TEST " + "=" * 20)
+
+ # Test 1: Empty filter (should return all results)
+ results_1 = await store.search(query="technology", filters={}, limit=10)
+ logger.info(f"Test 1 - Empty filter: {len(results_1)} results")
+
+ # Test 2: Single value filter
+ results_2 = await store.search(
+ query="technology",
+ filters={"node_type": "tech"},
+ limit=10,
+ )
+ logger.info(f"Test 2 - Single value filter: {len(results_2)} results")
+ for r in results_2:
+ assert r.metadata.get("node_type") == "tech"
+
+ # Test 3: List filter with single item
+ results_3 = await store.search(
+ query="technology",
+ filters={"node_type": ["tech"]},
+ limit=10,
+ )
+ logger.info(f"Test 3 - List filter (single item): {len(results_3)} results")
+
+ # Test 4: List filter with multiple items
+ results_4 = await store.search(
+ query="technology",
+ filters={"category": ["AI", "ML", "DL"]},
+ limit=10,
+ )
+ logger.info(f"Test 4 - List filter (multiple items): {len(results_4)} results")
+ for r in results_4:
+ assert r.metadata.get("category") in ["AI", "ML", "DL"]
+
+ # Test 5: Multiple filters (AND operation)
+ results_5 = await store.search(
+ query="technology",
+ filters={
+ "node_type": ["tech", "tech_new"],
+ "source": "research",
+ "priority": "high",
+ },
+ limit=10,
+ )
+ logger.info(f"Test 5 - Multiple filters (AND): {len(results_5)} results")
+ for r in results_5:
+ assert r.metadata.get("node_type") in ["tech", "tech_new"]
+ assert r.metadata.get("source") == "research"
+ assert r.metadata.get("priority") == "high"
+
+ # Test 6: Filter with non-existent value
+ results_6 = await store.search(
+ query="technology",
+ filters={"category": "NON_EXISTENT_CATEGORY"},
+ limit=10,
+ )
+ logger.info(f"Test 6 - Non-existent filter value: {len(results_6)} results")
+ assert len(results_6) == 0, "Should return no results for non-existent filter value"
+
+ # Test 7: List operation with filters
+ list_results = await store.list(
+ filters={"node_type": "tech", "priority": "high"},
+ limit=20,
+ )
+ logger.info(f"Test 7 - List with filters: {len(list_results)} results")
+ for r in list_results:
+ assert r.metadata.get("node_type") == "tech"
+ assert r.metadata.get("priority") == "high"
+
+ logger.info("✓ Filter combinations test passed")
+
+
+# ==================== Test Runner ====================
+
+
+async def run_all_tests_for_store(store_type: str, store_name: str):
+ """Run all tests for a specific vector store type.
+
+ Args:
+ store_type: Type of vector store ("local" or "es")
+ store_name: Display name for the vector store
+ """
+ logger.info(f"\n\n{'#' * 60}")
+ logger.info(f"# Running all tests for: {store_name}")
+ logger.info(f"{'#' * 60}")
+
+ config = TestConfig()
+ collection_name = f"{config.TEST_COLLECTION_PREFIX}_{store_type}_main"
+
+ # Create vector store instance
+ store = create_vector_store(store_type, collection_name)
+
+ try:
+ # Run cosine similarity test first (only for LocalVectorStore)
+ await test_cosine_similarity(store_name)
+
+ # ========== Basic Tests ==========
+ logger.info(f"\n{'#' * 60}")
+ logger.info("# BASIC FUNCTIONALITY TESTS")
+ logger.info(f"{'#' * 60}")
+
+ await test_create_collection(store, store_name)
+ await test_insert(store, store_name)
+ await test_search(store, store_name)
+ await test_search_with_single_filter(store, store_name)
+ await test_search_with_list_filter(store, store_name)
+ await test_search_with_multiple_filters(store, store_name)
+ await test_get_by_id(store, store_name)
+ await test_list_all(store, store_name)
+ await test_list_with_filters(store, store_name)
+ await test_update(store, store_name)
+ await test_delete(store, store_name)
+
+ # ========== Advanced Tests ==========
+ logger.info(f"\n{'#' * 60}")
+ logger.info("# ADVANCED FUNCTIONALITY TESTS")
+ logger.info(f"{'#' * 60}")
+
+ await test_batch_operations(store, store_name)
+ await test_complex_metadata_queries(store, store_name)
+ await test_edge_cases(store, store_name)
+ await test_concurrent_operations(store, store_name)
+ await test_search_relevance_ranking(store, store_name)
+ await test_metadata_statistics(store, store_name)
+ await test_update_metadata_only(store, store_name)
+ await test_filter_combinations(store, store_name)
+
+ # ========== Collection Management Tests ==========
+ logger.info(f"\n{'#' * 60}")
+ logger.info("# COLLECTION MANAGEMENT TESTS")
+ logger.info(f"{'#' * 60}")
+
+ await test_list_collections(store, store_name)
+ await test_copy_collection(store, store_name)
+ await test_delete_collection(store, store_name)
+
+ logger.info(f"\n{'=' * 60}")
+ logger.info(f"✓ All tests passed for {store_name}!")
+ logger.info(f"{'=' * 60}")
+
+ except Exception as e:
+ logger.error(f"Test failed: {e}")
+ raise
+ finally:
+ # Cleanup
+ await cleanup_store(store, store_type)
+
+
+async def cleanup_store(store: BaseVectorStore, store_type: str):
+ """Clean up test resources for a vector store.
+
+ Args:
+ store: Vector store instance
+ store_type: Type of vector store ("local" or "es")
+ """
+ logger.info("=" * 20 + " CLEANUP " + "=" * 20)
+
+ try:
+ # Clean up test collections
+ config = TestConfig()
+ collections = await store.list_collections()
+ test_collections = [c for c in collections if c.startswith(config.TEST_COLLECTION_PREFIX)]
+
+ for collection in test_collections:
+ try:
+ await store.delete_collection(collection)
+ logger.info(f"Deleted test collection: {collection}")
+ except Exception as e:
+ logger.warning(f"Failed to delete collection {collection}: {e}")
+
+ # Close connections
+ await store.close()
+
+ # Clean up local directory if LocalVectorStore
+ if store_type == "local":
+ test_dir = Path(config.LOCAL_ROOT_PATH)
+ if test_dir.exists():
+ shutil.rmtree(test_dir)
+ logger.info(f"Cleaned up local directory: {config.LOCAL_ROOT_PATH}")
+
+ # Clean up local directory if ChromaVectorStore
+ if store_type == "chroma" and config.CHROMA_PATH:
+ test_dir = Path(config.CHROMA_PATH)
+ if test_dir.exists():
+ shutil.rmtree(test_dir)
+ logger.info(f"Cleaned up chroma directory: {config.CHROMA_PATH}")
+
+ logger.info("✓ Cleanup completed")
+ except Exception as e:
+ logger.error(f"Cleanup error: {e}")
+
+
+# ==================== Main Entry Point ====================
+
+
+async def main():
+ """Main entry point for running tests."""
+ parser = argparse.ArgumentParser(
+ description="Run vector store tests",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ python test_vector_store.py --local # Test LocalVectorStore only
+ python test_vector_store.py --es # Test ESVectorStore only
+ python test_vector_store.py --pgvector # Test PGVectorStore only
+ python test_vector_store.py --qdrant # Test QdrantVectorStore only
+ python test_vector_store.py --chroma # Test ChromaVectorStore only
+ python test_vector_store.py --all # Test all vector stores
+ """,
+ )
+ parser.add_argument(
+ "--local",
+ action="store_true",
+ help="Test LocalVectorStore",
+ )
+ parser.add_argument(
+ "--es",
+ action="store_true",
+ help="Test ESVectorStore",
+ )
+ parser.add_argument(
+ "--qdrant",
+ action="store_true",
+ help="Test QdrantVectorStore",
+ )
+ parser.add_argument(
+ "--pgvector",
+ action="store_true",
+ help="Test PGVectorStore",
+ )
+ parser.add_argument(
+ "--chroma",
+ action="store_true",
+ help="Test ChromaVectorStore",
+ )
+ parser.add_argument(
+ "--all",
+ action="store_true",
+ help="Run tests for all available vector stores",
+ )
+
+ args = parser.parse_args()
+
+ # Determine which vector stores to test
+ stores_to_test = []
+
+ if args.all:
+ stores_to_test = [
+ ("local", "LocalVectorStore"),
+ ("es", "ESVectorStore"),
+ ("pgvector", "PGVectorStore"),
+ ("qdrant", "QdrantVectorStore"),
+ ("chroma", "ChromaVectorStore"),
+ ]
+ else:
+ # Build list based on individual flags
+ if args.local:
+ stores_to_test.append(("local", "LocalVectorStore"))
+ if args.es:
+ stores_to_test.append(("es", "ESVectorStore"))
+ if args.pgvector:
+ stores_to_test.append(("pgvector", "PGVectorStore"))
+ if args.qdrant:
+ stores_to_test.append(("qdrant", "QdrantVectorStore"))
+ if args.chroma:
+ stores_to_test.append(("chroma", "ChromaVectorStore"))
+
+ if not stores_to_test:
+ # Default to all vector stores if no argument provided
+ stores_to_test = [
+ ("local", "LocalVectorStore"),
+ ("es", "ESVectorStore"),
+ ("pgvector", "PGVectorStore"),
+ ("qdrant", "QdrantVectorStore"),
+ ("chroma", "ChromaVectorStore"),
+ ]
+ print("No vector store specified, defaulting to test all vector stores")
+ print(
+ "Use --local/--es/--pgvector/--qdrant/--chroma to test specific ones\n",
+ )
+
+ # Run tests for each vector store
+ for store_type, store_name in stores_to_test:
+ try:
+ await run_all_tests_for_store(store_type, store_name)
+ except Exception as e:
+ logger.error(f"\n✗ FAILED: {store_name} tests failed with error:")
+ logger.error(f" {type(e).__name__}: {e}")
+ raise
+
+ # Final summary
+ print(f"\n\n{'#' * 60}")
+ print("# TEST SUMMARY")
+ print(f"{'#' * 60}")
+ print(f"✓ All tests passed for {len(stores_to_test)} vector store(s):")
+ for _, store_name in stores_to_test:
+ print(f" - {store_name}")
+ print(f"{'#' * 60}\n")
+
+
+if __name__ == "__main__":
+ asyncio.run(main())