mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-10 22:41:06 +00:00
Merge remote-tracking branch 'origin/main'
This commit is contained in:
方应
commit
6ac8912ab5
21 changed files with 1198 additions and 144 deletions
10
README.md
10
README.md
|
|
@ -29,6 +29,14 @@
|
|||
|
||||
---
|
||||
|
||||
## 📰 Latest Articles
|
||||
|
||||
| Date | Title |
|
||||
|------------|-----------------------------------------------------------------|
|
||||
| 2026-03-30 | [CoPaw Context Management Design](docs/copaw_context_design.md) |
|
||||
|
||||
---
|
||||
|
||||
🧠 ReMe is a memory management framework designed for **AI agents**, providing
|
||||
both [file-based](#-file-based-memory-system-remelight) and [vector-based](#-vector-based-memory-system) memory systems.
|
||||
|
||||
|
|
@ -498,7 +506,7 @@ async def main():
|
|||
"dimensions": 1024,
|
||||
},
|
||||
default_vector_store_config={
|
||||
"backend": "local", # Supports local/chroma/qdrant/elasticsearch
|
||||
"backend": "local", # Supports local/chroma/qdrant/elasticsearch/obvec
|
||||
},
|
||||
)
|
||||
await reme.start()
|
||||
|
|
|
|||
19
README_ZH.md
19
README_ZH.md
|
|
@ -29,6 +29,14 @@
|
|||
|
||||
---
|
||||
|
||||
## 📰 最新文章
|
||||
|
||||
| 日期 | 标题 |
|
||||
|------------|----------------------------------------------------|
|
||||
| 2026-03-30 | [CoPaw 上下文管理设计解析](docs/copaw_context_design_zh.md) |
|
||||
|
||||
---
|
||||
|
||||
🧠 ReMe 是一个专为 **AI 智能体** 打造的记忆管理框架，同时提供基于[文件系统](#-基于文件的记忆系统-remelight)
|
||||
和基于[向量库](#-基于向量库的记忆系统)的记忆系统。
|
||||
|
||||
|
|
@ -224,12 +232,13 @@ flowchart TD
|
|||
A --> E[messages: 完整对话历史]
|
||||
A --> F[文件系统缓存]
|
||||
F --> G[dialog/YYYY-MM-DD.jsonl]
|
||||
F --> H[tool_result/uuid.txt N天TTL]
|
||||
F --> H[tool_result/uuid.txt N天TTL]
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
[CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/reme_light_memory_manager.py) 继承
|
||||
[CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/reme_light_memory_manager.py)
|
||||
继承
|
||||
`ReMeLight`，将记忆能力集成到 Agent 推理流程中：
|
||||
|
||||
```mermaid
|
||||
|
|
@ -326,7 +335,8 @@ graph LR
|
|||
|
||||
#### 4. compact_tool_result — 工具结果压缩
|
||||
|
||||
[ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) 解决工具输出过长导致上下文膨胀的问题。根据消息是否在 `recent_n` 范围内，采用不同的截断策略：
|
||||
[ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) 解决工具输出过长导致上下文膨胀的问题。根据消息是否在
|
||||
`recent_n` 范围内，采用不同的截断策略：
|
||||
|
||||
```mermaid
|
||||
graph LR
|
||||
|
|
@ -454,7 +464,6 @@ graph LR
|
|||
|
||||
安装和环境变量配置与 [ReMeLight 一致](#安装)，通过环境变量设置 API 密钥，可写在项目根目录的 `.env` 文件中。
|
||||
|
||||
|
||||
### Python 使用
|
||||
|
||||
```python
|
||||
|
|
@ -477,7 +486,7 @@ async def main():
|
|||
"dimensions": 1024,
|
||||
},
|
||||
default_vector_store_config={
|
||||
"backend": "local", # 支持 local/chroma/qdrant/elasticsearch
|
||||
"backend": "local", # 支持 local/chroma/qdrant/elasticsearch/obvec
|
||||
},
|
||||
)
|
||||
await reme.start()
|
||||
|
|
|
|||
|
|
@ -1,122 +1,278 @@
|
|||
## Copaw Context Management V2
|
||||
# CoPaw Context Management Design
|
||||
|
||||
> 注：不涉及长期记忆
|
||||
> This article focuses on **short-term context management** and does not cover the long-term memory module.
|
||||
|
||||
### 上下文数据结构
|
||||
---
|
||||
|
||||
#### 1. 上下文-内存
|
||||
Sooner or later, every AI Agent hits the same wall: **the context window fills up**.
|
||||
|
||||
- **compact_summary**（可选）：
|
||||
- **历史对话原始数据引导**：存储于 `dialog/YYYY-MM-DD.jsonl`，共 N 行，按时间顺序排列；回顾时建议从后往前读。
|
||||
- **历史对话摘要**：包含 `Goal + Constraints + Progress + KeyDecisions + NextSteps`。
|
||||
- **messages**：当前对话上下文（完整消息列表）。
|
||||
Tool calls return walls of HTML, thousands of lines of logs, or entire file contents — all of which rapidly consume precious token budget. As the conversation grows, early information either gets truncated or blows up the window entirely, and the Agent's performance begins to degrade.
|
||||
|
||||
#### 2. 上下文-缓存到文件系统
|
||||
[CoPaw](https://github.com/agentscope-ai/CoPaw) addresses this problem with a systematic approach to context management. This article provides a complete breakdown of the data structures and runtime mechanics behind **CoPaw Context Management V2**.
|
||||
|
||||
- **历史对话原始数据**：`dialog/YYYY-MM-DD.jsonl`
|
||||
- **工具调用结果原始数据**：`tool_result/{uuid}.txt`（保留 N 天）
|
||||
---
|
||||
|
||||
## What does the context look like?
|
||||
|
||||
Before discussing "how to manage it", let's first understand "what is being managed".
|
||||
|
||||
CoPaw's context is split into two layers: the **in-memory layer** and the **file system layer**.
|
||||
|
||||
### In-Memory Layer
|
||||
|
||||
Two core fields are maintained in memory:
|
||||
|
||||
- **`compact_summary`** (optional): After conversation history has been compacted, this field holds a structured summary covering five dimensions — `Goal`, `Constraints`, `Progress`, `KeyDecisions`, and `NextSteps` — essentially a refined "work memo". It also includes a **path guide to the raw historical dialog**, pointing the Agent to `dialog/YYYY-MM-DD.jsonl` and suggesting reading from the end backwards.
|
||||
- **`messages`**: The complete list of messages for the current conversation — the data actually consumed by the Agent during reasoning.
|
||||
|
||||
### File System Layer (File Cache)
|
||||
|
||||
For content that is too large or too volatile to reside in memory long-term, CoPaw offloads it to the file system:
|
||||
|
||||
- **Raw conversation history**: `dialog/YYYY-MM-DD.jsonl`, stored per day
|
||||
- **Tool call results**: `tool_result/{uuid}.txt`, with an N-day TTL and automatic cleanup on expiry
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A[Context] --> B[compact_summary]
|
||||
B --> C[dialog path guide + Goal/Constraints/Progress/KeyDecisions/NextSteps]
|
||||
A --> E[messages: full dialogue history]
|
||||
A --> F[File System Cache] --> G[dialog/YYYY-MM-DD.jsonl]
|
||||
A --> F[File System Cache]
|
||||
F --> G[dialog/YYYY-MM-DD.jsonl]
|
||||
F --> H[tool_result/uuid.txt N-day TTL]
|
||||
```
|
||||
|
||||
This design allows the Agent to quickly access recent conversations in memory while being able to look up historical context on demand — without stuffing all history into the context window.
|
||||
|
||||
---
|
||||
|
||||
### 上下文机制（Pre-Reasoning Hook）
|
||||
## What happens before reasoning? The Pre-Reasoning Hook
|
||||
|
||||
1. **工具结果 Offload** (`ToolCallResultCompact`)
|
||||
2. **上下文检查** (`ContextChecker`)
|
||||
3. **若 Token 超阈值**：
|
||||
- 保留最近 **X%** 的 Token（保障连贯性）
|
||||
- 其余历史对话生成摘要 (`Compactor`)
|
||||
4. **被摘要的上下文 Offload 到文件系统** (`SaveDialog`)
|
||||
Before each reasoning step begins, CoPaw executes a **Pre-Reasoning Hook** that automatically tidies up the context. The process runs in four steps:
|
||||
|
||||
1. **Tool result compaction** (`ToolCallResultCompact`): Process tool call results first — truncate oversized content and offload it to the file system.
|
||||
2. **Context checking** (`ContextChecker`): Compute the current token usage and determine whether it exceeds the threshold.
|
||||
3. **If the threshold is exceeded**:
|
||||
- Keep the most recent **X%** of tokens (to preserve conversational continuity).
|
||||
- Call the `Compactor` on the earlier history to generate a structured summary.
|
||||
4. **Dialog persistence** (`SaveDialog`): Save the compacted raw conversation to the file system.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[Pre-Reasoning Hook] --> B[ToolCallResultCompact]
|
||||
B --> C[ContextChecker]
|
||||
C --> D{Token > Threshold?}
|
||||
C --> D{Token > threshold?}
|
||||
D -->|Yes| E[Keep recent X% tokens]
|
||||
E --> F[Compact & Summary old context]
|
||||
F --> G[SaveDialog: offload to file]
|
||||
D -->|No| H[Proceed normally]
|
||||
E --> F[Compact & generate summary]
|
||||
F --> G[SaveDialog: persist to file]
|
||||
D -->|No| H[Normal reasoning]
|
||||
```
|
||||
|
||||
This flow ensures that the context is in a "clean" state at the start of every reasoning step.
|
||||
|
||||
---
|
||||
|
||||
### 工具结果 Offload 机制
|
||||
## Tool Result Offload: Unified Two-Phase Truncation
|
||||
|
||||
1. 所有工具调用结果先放入上下文，等待 Pre-Reasoning Hook 处理。
|
||||
2. 根据是否属于 **recent_n** 范围，决定截断策略：
|
||||
- **recent_n 内**：近期内容 → 低截断比例
|
||||
- **recent_n 外**：远期内容 → 高截断比例
|
||||
Tool call results are one of the main causes of context bloat. CoPaw uses a **two-phase truncation** strategy that separates the timing of truncation from its aggressiveness:
|
||||
|
||||
#### 示例：Browser Use 类工具
|
||||
|
||||
| 阶段 | 行为 |
|
||||
|----|------------------------------------------------------------------------------|
|
||||
| 1 | 原始工具调用结果 |
|
||||
| 2 | 保存原始内容到文件：– 若在 recent_n 内：截断较少– 附注：“FullText saved to xxxx”– 提示：“请从第 N 行开始读” |
|
||||
| 3 | 若再次引用且超出 recent_n：– 二次截断（更激进）– 仍指向原文件路径 |
|
||||
- **First truncation**: Triggered immediately when a tool call result is **written into the context**, uniformly applied to all tools (including `read_file`). The full raw content is saved to `tool_result/{uuid}.txt`, and the message is annotated with the file path and a start-line hint.
|
||||
- **Second truncation**: Triggered by the Pre-Reasoning Hook on **messages that have slid out of the `recent_n` window**, applying a more aggressive truncation to further shrink context usage.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[Tool Call Result] --> B{Within recent_n?}
|
||||
B -->|Yes| C[Low truncation<br>Save full text to tool_result/uuid.txt<br>Hint: 'Read from line N']
|
||||
B -->|No| D[High truncation<br>Reference existing file<br>More aggressive truncation]
|
||||
C --> E[Context includes snippet + file ref]
|
||||
A[Tool call completes] --> T[First truncation<br>executed immediately on write]
|
||||
T --> S[Full content written to tool_result/uuid.txt<br>Message annotated with file path + start line]
|
||||
S --> B{Pre-Reasoning Hook<br>Is message within recent_n?}
|
||||
B -->|Yes| C[No action<br>Keep first-truncation result]
|
||||
B -->|No| D[Second truncation<br>More aggressive compression<br>file_path unchanged]
|
||||
```
|
||||
|
||||
The benefit of this design is: the first truncation ensures that no tool result can blow up the context from the moment it is written; the second truncation automatically "fades out" older messages as the conversation progresses, always leaving enough room for recent content.
|
||||
|
||||
### Browser Use Tools as an Example
|
||||
|
||||
| Phase | Behavior |
|
||||
|----------------------------|-----------------------------------------------------------------------------------------------------------------------------|
|
||||
| First truncation | Result is truncated immediately on return; full content written to `tool_result/uuid.txt`; message annotated with "FullText saved to xxxx, please read from line N" |
|
||||
| Within `recent_n` | Pre-Reasoning Hook makes no additional changes; keeps first-truncation result |
|
||||
| Outside `recent_n` (second truncation) | Parses the existing message result, applies more aggressive truncation to the original content, updates meta info (e.g. line-number hint); **`file_path` is unchanged**, still pointing to the original file |
|
||||
|
||||
The key insight of second truncation: **the original full content is always saved under the same file path**. No matter how many rounds of truncation have occurred, the Agent can always retrieve the original content via the file reference. Truncation only affects the message fragment and meta info in the context — never the file itself.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[Browser Use Result] -->|First truncation| B[Context: fragment + file_path + start line]
|
||||
B --> C{Outside recent_n?}
|
||||
C -->|No| D[Unchanged]
|
||||
C -->|Yes| E[Parse existing message<br>Apply second truncation to original content<br>Update meta info]
|
||||
E --> F[Context: shorter fragment + same file_path]
|
||||
```
|
||||
|
||||
### Code Implementation and Examples of Two-Phase Truncation
|
||||
|
||||
The entry point for truncation logic is `truncate_text_output`, which dispatches to two different functions depending on whether the text already contains the `<<<TRUNCATED>>>` marker:
|
||||
|
||||
```python
|
||||
def truncate_text_output(text, start_line=1, total_lines=0,
|
||||
max_bytes=DEFAULT_MAX_BYTES,
|
||||
file_path=None, encoding="utf-8") -> str:
|
||||
if TRUNCATION_NOTICE_MARKER in text:
|
||||
return _retruncate(text, max_bytes=max_bytes, encoding=encoding)
|
||||
else:
|
||||
return _truncate_fresh(text, start_line=start_line,
|
||||
total_lines=total_lines,
|
||||
max_bytes=max_bytes,
|
||||
file_path=file_path, encoding=encoding)
|
||||
```
|
||||
|
||||
#### First Truncation (`_truncate_fresh`)
|
||||
|
||||
**When it fires**: Immediately when the tool call completes and the result is written into the context — the text does not yet contain a truncation marker at this point.
|
||||
|
||||
**Core logic**:
|
||||
|
||||
1. If the text size in bytes does not exceed `max_bytes`, return the original text as-is.
|
||||
2. Otherwise, slice by bytes, keep the last complete line before the cut point, and compute the start line for the next read.
|
||||
3. Append a truncation notice (`<<<TRUNCATED>>>`) at the end, prompting the reader to continue from `start_line=N`.
|
||||
|
||||
**Example**: Suppose a tool returns 3,000 lines of HTML (200 KB in total), and `max_bytes = 50 KB`:
|
||||
|
||||
```
|
||||
# Original tool output (200 KB, 3000 lines)
|
||||
<html>
|
||||
<head>...</head>
|
||||
<body>
|
||||
...(large content)
|
||||
</body>
|
||||
</html>
|
||||
|
||||
# After first truncation, written to context (50 KB, ~750 lines)
|
||||
<html>
|
||||
<head>...</head>
|
||||
<body>
|
||||
...(first 750 lines)
|
||||
<<<TRUNCATED>>>
|
||||
The output above was truncated.
|
||||
The full content is saved to the file and contains 3000 lines in total.
|
||||
This excerpt starts at line 1 and covers the next 51200 bytes.
|
||||
If the current content is not enough, call `read_file` with file_path=tool_result/abc123.txt start_line=751 to read more.
|
||||
```
|
||||
|
||||
The full raw content is simultaneously written to `tool_result/abc123.txt`; only the truncated fragment and the continuation hint are kept in the context.
|
||||
|
||||
#### Second Truncation (`_retruncate`)
|
||||
|
||||
**When it fires**: During Pre-Reasoning Hook processing, applied to messages that have slid out of the `recent_n` window to further shrink context usage.
|
||||
|
||||
**Core logic**:
|
||||
|
||||
1. Split the text into the raw content before `<<<TRUNCATED>>>` and the notice section after it.
|
||||
2. If the raw content still fits within the new `max_bytes` (with a 100-byte slack), return the original text as-is.
|
||||
3. Otherwise, re-slice according to the new, smaller byte limit and use regex to update the **byte count** and **continuation line number** in the notice; `file_path` remains unchanged.
|
||||
|
||||
**Example**: The same tool message from above, after it slides out of `recent_n`. Second truncation reduces `max_bytes` from 50 KB to 10 KB:
|
||||
|
||||
```
|
||||
# Before second truncation (first-truncation result already in context, 50 KB)
|
||||
<html>
|
||||
<head>...</head>
|
||||
<body>
|
||||
...(first 750 lines)
|
||||
<<<TRUNCATED>>>
|
||||
...This excerpt starts at line 1 and covers the next 51200 bytes.
|
||||
...call `read_file` with file_path=tool_result/abc123.txt start_line=751 to read more.
|
||||
|
||||
# After second truncation (further compressed to 10 KB, ~150 lines)
|
||||
<html>
|
||||
<head>...</head>
|
||||
<body>
|
||||
...(first 150 lines)
|
||||
<<<TRUNCATED>>>
|
||||
...This excerpt starts at line 1 and covers the next 10240 bytes.
|
||||
...call `read_file` with file_path=tool_result/abc123.txt start_line=151 to read more.
|
||||
```
|
||||
|
||||
Key point: `file_path` always points to `tool_result/abc123.txt`. The Agent can retrieve the full original content via the file reference at any time; truncation only affects the context fragment and meta info.
|
||||
|
||||
---
|
||||
|
||||
## Special Handling for the ReadFile Tool
|
||||
|
||||
`read_file` shares the same two-phase truncation mechanism as Browser Use tools, with one key difference: **the file it reads already exists on the file system**, so there is no need to save a separate copy during first truncation.
|
||||
|
||||
| Phase | Behavior |
|
||||
|----------------------------|-----------------------------------------------------------------------------------------------|
|
||||
| First truncation | Truncation happens at read time; result is written to context; original file path is already known — no need to save to `tool_result/` |
|
||||
| Within `recent_n` | Pre-Reasoning Hook makes no changes; keeps the read-time truncation result |
|
||||
| Outside `recent_n` (second truncation) | Same as other tools — more aggressive truncation is applied to the message content, meta info is updated |
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[ReadFile call] -->|Truncate at read time| B[Context: truncated content<br>original file path known]
|
||||
B --> C{Outside recent_n?}
|
||||
C -->|No| D[No changes needed]
|
||||
C -->|Yes| E[Second truncation<br>Update meta info<br>Same behavior as other tools]
|
||||
```
|
||||
|
||||
### Special Protection for Markdown Files
|
||||
|
||||
For Markdown files such as `skill.md` and rule files, CoPaw applies a **higher protection threshold** during truncation.
|
||||
|
||||
Markdown files typically carry structured knowledge or instructions; over-aggressive truncation would break their semantic integrity. Therefore, both the first and second truncation thresholds for Markdown files are set higher than those for regular tool outputs, ensuring the Agent can read as complete a structured content as possible.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[Tool result] --> B{Is it a Markdown file?}
|
||||
B -->|Yes| C[Higher truncation threshold<br>Greater protection]
|
||||
B -->|No| D[Standard truncation threshold]
|
||||
C --> E[First / second truncation logic]
|
||||
D --> E
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### ReadFile 工具调用结果变化示例
|
||||
## Long-term Memory Trigger Logic
|
||||
|
||||
| 阶段 | 行为 |
|
||||
|----|---------------------------------------------|
|
||||
| 1 | 原始工具调用结果 |
|
||||
| 2 | 若在 recent_n 内：– 不截断– 不保存文件（因内容已由用户指定） |
|
||||
| 3 | 若超出 recent_n：– 二次截断（更小）– 保存 FullText 到文件并引用 |
|
||||
> This section goes beyond the core scope of context management and briefly introduces CoPaw's long-term memory write mechanism.
|
||||
|
||||
> 注：ReadFile 本身读取的是外部文件，因此首次调用通常无需重复保存。
|
||||
Long-term memory is driven by three trigger paths:
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[ReadFile Result] --> B{Within recent_n?}
|
||||
B -->|Yes| C[No truncation<br>No file save needed]
|
||||
B -->|No| D[Apply secondary truncation<br>Save FullText to tool_result/uuid.txt]
|
||||
C --> E[Include full content in context]
|
||||
D --> F[Include snippet + file ref]
|
||||
```
|
||||
1. **Explicitly written by the Main Agent**:
|
||||
- `Memory.md` (the backbone of long-term memory, recording persistent information such as user preferences)
|
||||
- `YYYY-MM-DD.md` (daily log)
|
||||
|
||||
---
|
||||
2. **Triggered by context compaction**, written by the **Summarizer (ReAct Agent)**:
|
||||
- Personalization information (user preferences, habits, etc.)
|
||||
- Try-error information (failed attempts and corrective lessons)
|
||||
|
||||
## Copaw Memory
|
||||
|
||||
### 触发逻辑
|
||||
|
||||
1. **主 Agent 主动写入**：
|
||||
- `Memory.md`（长期记忆主干）
|
||||
- `YYYY-MM-DD.md`（当日日志）
|
||||
2. **触发阈值时**，由 **Summarizer（React Agent）** 写日志：
|
||||
- 个性化信息（如偏好、习惯）
|
||||
- Try-error 信息（失败尝试与修正）
|
||||
3. **定时任务**（每日 00:00）：
|
||||
- 汇总最近的 `YYYY-MM-DD.md` 文件
|
||||
- 更新 `Memory.md`
|
||||
3. **Scheduled task** (daily at 00:00):
|
||||
- Aggregates recent `YYYY-MM-DD.md` files
|
||||
- Merges and updates the journal into `Memory.md`
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A[Main Agent] --> B[Write Memory.md]
|
||||
A --> C[Write YYYY-MM-DD.md]
|
||||
D[Context Threshold Reached?] -->|Yes| E[Summarizer Agent]
|
||||
E --> F[Log: Personalization]
|
||||
E --> G[Log: Try-Error Info]
|
||||
H[Cron @ 00:00 daily] --> I[Aggregate recent YYYY-MM-DD.md]
|
||||
I --> J[Update Memory.md]
|
||||
D[Context reaches compaction threshold?] -->|Yes| E[Summarizer Agent]
|
||||
E --> F[Record: personalization info]
|
||||
E --> G[Record: try-error experience]
|
||||
H[Scheduled task 00:00] --> I[Aggregate recent YYYY-MM-DD.md]
|
||||
I --> J[Update Memory.md]
|
||||
```
|
||||
|
||||
This mechanism ensures that important information from short-term conversations is distilled into long-term memory and not lost when the session ends.
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
The core design philosophy of CoPaw's context management can be summed up in one sentence:
|
||||
|
||||
**Keep only "what is needed now" in memory; let the file system hold "what might be needed later".**
|
||||
|
||||
Through the four-step Pre-Reasoning Hook flow, a unified two-phase truncation strategy, and persistent file system backing, CoPaw maximizes information availability for the Agent within a limited context window — no matter how long the conversation runs, the Agent can always find the context it needs.
|
||||
|
||||
---
|
||||
|
||||
*The design described in this article is implemented in [CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/reme_light_memory_manager.py) and [ReMe ReMeLight](https://github.com/agentscope-ai/ReMe).*
|
||||
|
|
|
|||
289
docs/copaw_context_design_zh.md
Normal file
289
docs/copaw_context_design_zh.md
Normal file
|
|
@ -0,0 +1,289 @@
|
|||
# CoPaw 上下文管理设计解析
|
||||
|
||||
> 本文聚焦**短期上下文管理**，不涉及长期记忆模块。
|
||||
|
||||
---
|
||||
|
||||
AI Agent 在使用过程中，迟早会遇到一个让人头疼的问题：**上下文窗口被塞满了**。
|
||||
|
||||
工具调用返回了一大段 HTML、几千行日志、或者完整的文件内容——这些都会急剧消耗宝贵的 Token
|
||||
配额。随着对话轮次增加，早期的信息要么被截断，要么把整个窗口撑爆，Agent 的表现开始下滑。
|
||||
|
||||
[CoPaw](https://github.com/agentscope-ai/CoPaw) 在设计上下文管理时，围绕这个问题给出了一套系统性的答案。本文将完整拆解 *
|
||||
*CoPaw Context Management V2** 的数据结构与运行机制。
|
||||
|
||||
---
|
||||
|
||||
## 上下文长什么样？
|
||||
|
||||
在讨论"如何管理"之前，先看清楚"管理的是什么"。
|
||||
|
||||
CoPaw 的上下文分为两层：**内存层**与**文件系统层**。
|
||||
|
||||
### 内存层（In-Memory）
|
||||
|
||||
内存中维护两个核心字段：
|
||||
|
||||
- **`compact_summary`**（可选）：当历史对话被压缩后，这里存放结构化摘要，包含 `Goal`、`Constraints`、`Progress`、`KeyDecisions`、
|
||||
`NextSteps` 五个维度——相当于一份精炼的"工作备忘录"。同时还包含一个**历史对话原始数据的路径引导**，告诉 Agent 去哪里找
|
||||
`dialog/YYYY-MM-DD.jsonl`，以及建议"从后往前读"。
|
||||
- **`messages`**：当前对话的完整消息列表，是 Agent 实际推理时消费的数据。
|
||||
|
||||
### 文件系统层（File Cache）
|
||||
|
||||
对于体积较大、不适合长期驻留内存的内容，CoPaw 将其 offload 到文件系统：
|
||||
|
||||
- **历史对话原始数据**：`dialog/YYYY-MM-DD.jsonl`，按日期分文件存储
|
||||
- **工具调用结果**：`tool_result/{uuid}.txt`，设有 N 天 TTL，过期自动清理
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A[Context] --> B[compact_summary]
|
||||
B --> C[dialog 路径引导 + Goal/Constraints/Progress/KeyDecisions/NextSteps]
|
||||
A --> E[messages: 完整对话历史]
|
||||
A --> F[文件系统缓存]
|
||||
F --> G[dialog/YYYY-MM-DD.jsonl]
|
||||
F --> H[tool_result/uuid.txt N天TTL]
|
||||
```
|
||||
|
||||
这个设计让 Agent 既能在内存中快速访问近期对话，又能在需要时按需回溯历史——而不是把所有历史内容硬塞进上下文。
|
||||
|
||||
---
|
||||
|
||||
## 推理前做什么？Pre-Reasoning Hook
|
||||
|
||||
每轮推理正式开始前，CoPaw 会执行一个 **Pre-Reasoning Hook**，自动完成上下文的整理工作。整个流程分四步：
|
||||
|
||||
1. **工具结果压缩**（`ToolCallResultCompact`）：先处理工具调用结果，将超长内容截断并 offload 到文件系统
|
||||
2. **上下文检查**（`ContextChecker`）：计算当前上下文的 Token 使用量，判断是否超出阈值
|
||||
3. **若超出阈值**：
|
||||
- 保留最近 **X%** 的 Token（保障对话连贯性）
|
||||
- 对更早的历史对话调用 `Compactor` 生成结构化摘要
|
||||
4. **历史对话持久化**（`SaveDialog`）：将被压缩的原始对话保存到文件系统
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[Pre-Reasoning Hook] --> B[ToolCallResultCompact]
|
||||
B --> C[ContextChecker]
|
||||
C --> D{Token > 阈值?}
|
||||
D -->|是| E[保留近期 X% Token]
|
||||
E --> F[Compact & 生成摘要]
|
||||
F --> G[SaveDialog: 持久化到文件]
|
||||
D -->|否| H[正常推理]
|
||||
```
|
||||
|
||||
这个流程确保每次推理开始前，上下文都处于一个"干净"的状态。
|
||||
|
||||
---
|
||||
|
||||
## 工具结果 Offload：统一的两阶段截断
|
||||
|
||||
工具调用结果是上下文膨胀的主要来源之一。CoPaw 采用**两阶段截断**策略，将截断时机与截断力度分离：
|
||||
|
||||
- **一次截断**：在工具调用结果**写入上下文时**立即触发，所有工具（包括 `read_file`）统一适用。截断后将完整原始内容保存到
|
||||
`tool_result/{uuid}.txt`，并在消息中附注文件路径与起始行提示。
|
||||
- **二次截断**：在 Pre-Reasoning Hook 处理时，对**已滑出 `recent_n` 范围**的历史消息触发，截断更为激进，进一步压缩上下文占用。
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[工具调用完成] --> T[一次截断<br>写入上下文时立即执行]
|
||||
T --> S[完整内容写入 tool_result/uuid.txt<br>消息附注文件路径 + 起始行]
|
||||
S --> B{Pre-Reasoning Hook<br>该消息在 recent_n 内?}
|
||||
B -->|是| C[无需处理<br>保持一次截断结果]
|
||||
B -->|否| D[二次截断<br>更激进压缩<br>file_path 不变]
|
||||
```
|
||||
|
||||
这样设计的好处在于：一次截断保证所有工具结果从写入那刻起就不会撑爆上下文；二次截断则随着对话推进自动"淡化"
|
||||
历史信息，始终为近期内容留出充足空间。
|
||||
|
||||
### 以 Browser Use 类工具为例
|
||||
|
||||
| 阶段 | 行为 |
|
||||
|-------------------|------------------------------------------------------------------------------------|
|
||||
| 一次截断 | 工具返回结果后立即截断，完整内容写入 `tool_result/uuid.txt`，消息附注 "FullText saved to xxxx，请从第 N 行开始读" |
|
||||
| 在 recent_n 内 | Pre-Reasoning Hook 不做额外处理，保持一次截断结果 |
|
||||
| 超出 recent_n（二次截断） | 解析已有的消息结果，对原始内容做更激进的截断，同步更新消息中的 meta 信息（如行号提示）；**`file_path` 不变**，仍指向原文件 |
|
||||
|
||||
二次截断的关键在于：**原始完整内容始终保存在同一个文件路径下**，无论经过多少轮截断，Agent 都能通过文件引用找到原始内容；截断只影响上下文中的消息片段和
|
||||
meta 信息，不改变文件。
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[Browser Use Result] -->|一次截断| B[上下文: 片段 + file_path + 起始行]
|
||||
B --> C{超出 recent_n?}
|
||||
C -->|否| D[保持不变]
|
||||
C -->|是| E[解析现有消息<br>对原始内容二次截断<br>更新 meta 信息]
|
||||
E --> F[上下文: 更短片段 + 同一 file_path]
|
||||
```
|
||||
|
||||
### 两阶段截断的代码实现与示例
|
||||
|
||||
截断逻辑的入口是 `truncate_text_output`，它根据文本中是否已包含 `<<<TRUNCATED>>>` 标记来分发到两个不同的函数：
|
||||
|
||||
```python
|
||||
def truncate_text_output(text, start_line=1, total_lines=0,
|
||||
max_bytes=DEFAULT_MAX_BYTES,
|
||||
file_path=None, encoding="utf-8") -> str:
|
||||
if TRUNCATION_NOTICE_MARKER in text:
|
||||
return _retruncate(text, max_bytes=max_bytes, encoding=encoding)
|
||||
else:
|
||||
return _truncate_fresh(text, start_line=start_line,
|
||||
total_lines=total_lines,
|
||||
max_bytes=max_bytes,
|
||||
file_path=file_path, encoding=encoding)
|
||||
```
|
||||
|
||||
#### 一次截断（`_truncate_fresh`）
|
||||
|
||||
**触发时机**：工具调用完成、结果写入上下文时立即执行，此时文本中尚不含截断标记。
|
||||
|
||||
**核心逻辑**：
|
||||
|
||||
1. 若文本字节数未超过 `max_bytes`，直接返回原文；
|
||||
2. 否则按字节切片，保留截断点前最后一个完整行，计算下一段应从哪一行开始；
|
||||
3. 在末尾追加截断通知（`<<<TRUNCATED>>>`），提示后续从 `start_line=N` 继续读取。
|
||||
|
||||
**示例**：假设一个工具返回了 3 000 行的 HTML 内容（共 200 KB），而 `max_bytes = 50 KB`：
|
||||
|
||||
```
|
||||
# 原始工具输出（200 KB，共 3000 行）
|
||||
<html>
|
||||
<head>...</head>
|
||||
<body>
|
||||
...（大量内容）
|
||||
</body>
|
||||
</html>
|
||||
|
||||
# 一次截断后写入上下文（50 KB，约 750 行）
|
||||
<html>
|
||||
<head>...</head>
|
||||
<body>
|
||||
...（前 750 行）
|
||||
<<<TRUNCATED>>>
|
||||
The output above was truncated.
|
||||
The full content is saved to the file and contains 3000 lines in total.
|
||||
This excerpt starts at line 1 and covers the next 51200 bytes.
|
||||
If the current content is not enough, call `read_file` with file_path=tool_result/abc123.txt start_line=751 to read more.
|
||||
```
|
||||
|
||||
完整原始内容同时写入 `tool_result/abc123.txt`，上下文中仅保留截断片段与续读提示。
|
||||
|
||||
#### 二次截断（`_retruncate`）
|
||||
|
||||
**触发时机**：Pre-Reasoning Hook 处理时，对已滑出 `recent_n` 范围的历史消息执行，进一步压缩上下文占用。
|
||||
|
||||
**核心逻辑**：
|
||||
|
||||
1. 从文本中分离出 `<<<TRUNCATED>>>` 前的原始内容与后面的通知部分；
|
||||
2. 若原始内容仍未超出新的 `max_bytes`（带 100 字节宽松量），直接返回原文；
|
||||
3. 否则按新的更小字节限制重新切片，并通过正则替换通知中的 **字节数** 与 **续读行号**，`file_path` 保持不变。
|
||||
|
||||
**示例**：同样是上面那条工具消息，在它滑出 `recent_n` 之后，二次截断将 `max_bytes` 从 50 KB 压缩到 10 KB：
|
||||
|
||||
```
|
||||
# 二次截断前（上下文中已有一次截断结果，50 KB）
|
||||
<html>
|
||||
<head>...</head>
|
||||
<body>
|
||||
...（前 750 行）
|
||||
<<<TRUNCATED>>>
|
||||
...This excerpt starts at line 1 and covers the next 51200 bytes.
|
||||
...call `read_file` with file_path=tool_result/abc123.txt start_line=751 to read more.
|
||||
|
||||
# 二次截断后（进一步压缩至 10 KB，约 150 行）
|
||||
<html>
|
||||
<head>...</head>
|
||||
<body>
|
||||
...（前 150 行）
|
||||
<<<TRUNCATED>>>
|
||||
...This excerpt starts at line 1 and covers the next 10240 bytes.
|
||||
...call `read_file` with file_path=tool_result/abc123.txt start_line=151 to read more.
|
||||
```
|
||||
|
||||
关键点：`file_path` 始终指向 `tool_result/abc123.txt`，Agent 随时可通过文件引用获取完整原始内容；截断只影响上下文片段与 meta 信息。
|
||||
|
||||
---
|
||||
|
||||
## ReadFile 工具的特殊处理
|
||||
|
||||
`read_file` 与 Browser Use 类工具共享同一套两阶段截断机制，但有一个关键区别：**它读取的文件本身已存在于文件系统**
|
||||
，无需在一次截断时另行保存。
|
||||
|
||||
| 阶段 | 行为 |
|
||||
|-------------------|---------------------------------------------------|
|
||||
| 一次截断 | 在读取时即完成截断，结果写入上下文；原始文件路径已知，无需额外保存到 `tool_result/` |
|
||||
| 在 recent_n 内 | Pre-Reasoning Hook 不做任何修改，保持读取时的截断结果 |
|
||||
| 超出 recent_n（二次截断） | 与其他工具相同，对消息内容做更激进的截断，更新 meta 信息 |
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[ReadFile 调用] -->|读取时截断| B[上下文: 截断内容<br>原始文件路径已知]
|
||||
B --> C{超出 recent_n?}
|
||||
C -->|否| D[无需修改]
|
||||
C -->|是| E[二次截断<br>更新 meta 信息<br>与其他工具行为一致]
|
||||
```
|
||||
|
||||
### Markdown 文件的特殊保护
|
||||
|
||||
对于 `skill.md`、规则文件等 Markdown 文件，CoPaw 在截断时给予**更大的保护阈值**。
|
||||
|
||||
Markdown 文件通常承载结构化的知识或指令，过度截断会破坏其完整语义。因此，在一次截断和二次截断时，Markdown
|
||||
文件的截断触发上限均高于普通工具输出，确保 Agent 能读到尽可能完整的结构化内容。
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[工具结果] --> B{是 Markdown 文件?}
|
||||
B -->|是| C[更高截断阈值<br>更大保护]
|
||||
B -->|否| D[标准截断阈值]
|
||||
C --> E[一次 / 二次截断逻辑]
|
||||
D --> E
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 长期记忆的触发逻辑
|
||||
|
||||
> 本节超出上下文管理的核心范畴，简要介绍 CoPaw 的长期记忆写入机制。
|
||||
|
||||
长期记忆由三个触发路径驱动：
|
||||
|
||||
1. **主 Agent 主动写入**：
|
||||
- `Memory.md`（长期记忆主干，记录用户偏好等持久信息）
|
||||
- `YYYY-MM-DD.md`（当日日志）
|
||||
|
||||
2. **上下文压缩触发时**，由 **Summarizer（ReAct Agent）** 写入：
|
||||
- 个性化信息（用户偏好、习惯等）
|
||||
- Try-error 信息（失败尝试与修正经验）
|
||||
|
||||
3. **定时任务**（每日 00:00）：
|
||||
- 汇总最近的 `YYYY-MM-DD.md` 文件
|
||||
- 将日志整合更新到 `Memory.md`
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A[Main Agent] --> B[写入 Memory.md]
|
||||
A --> C[写入 YYYY-MM-DD.md]
|
||||
D[上下文达到压缩阈值?] -->|是| E[Summarizer Agent]
|
||||
E --> F[记录: 个性化信息]
|
||||
E --> G[记录: Try-Error 经验]
|
||||
H[定时任务 00:00] --> I[汇总最近 YYYY-MM-DD.md]
|
||||
I --> J[更新 Memory.md]
|
||||
```
|
||||
|
||||
这套机制确保了短期对话中的重要信息能够沉淀为长期记忆，不因对话结束而丢失。
|
||||
|
||||
---
|
||||
|
||||
## 小结
|
||||
|
||||
CoPaw 上下文管理的核心设计哲学可以用一句话概括：
|
||||
|
||||
**让内存只放"现在需要的"，让文件系统保管"之后可能需要的"。**
|
||||
|
||||
通过 Pre-Reasoning Hook 的四步流程、统一的两阶段截断策略，以及文件系统的持久化支撑，CoPaw 在有限的上下文窗口内为 Agent
|
||||
提供了最大程度的信息可用性——无论对话持续多久，Agent 总能找到它需要的上下文。
|
||||
|
||||
---
|
||||
|
||||
*本文设计对应实现可参考 [CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/reme_light_memory_manager.py)
|
||||
与 [ReMe ReMeLight](https://github.com/agentscope-ai/ReMe)。*
|
||||
|
|
@ -139,7 +139,7 @@ response = requests.post("http://localhost:8002/retrieve_task_memory", json={
|
|||
## 📚 Resources
|
||||
|
||||
- **[Installation Guide](installation.md)**, **[Quick Start](quick_start.md)**: Get started quickly with practical examples
|
||||
- **[Vector Storage Setup](vector_store_api_guide.md)**: Configure local/vector databases and usage
|
||||
- **[Vector Storage Setup](vector_store_api_guide.md)**: Configure local, Elasticsearch, Qdrant, ChromaDB, or ObVec (OceanBase / seekdb via pyobvector) storage and usage
|
||||
- **[MCP Guide](mcp_quick_start.md)**: Create MCP services
|
||||
- **[Personal Memory](personal_memory/personal_memory.md)**, **[Task Memory](task_memory/task_memory.md)** & **[Tool Memory](tool_memory/tool_memory.md)**: Operators used in personal memory, task memory and tool memory. You can modify the config to customize the pipelines.
|
||||
- **[Example Collection](./cookbook/appworld/quickstart.md)**: Real use cases and best practices
|
||||
|
|
|
|||
|
|
@ -33,8 +33,9 @@ FlowLLM provides multiple Vector Store implementations tailored to different use
|
|||
- **QdrantVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/qdrant_vector_store.py)): Built on the Qdrant vector database, supporting high-performance vector search. Recommended for large-scale production environments.
|
||||
- **ChromaVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/chroma_vector_store.py)): Based on ChromaDB, providing persistent storage and metadata filtering capabilities.
|
||||
- **EsVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/es_vector_store.py)): Built on Elasticsearch, enabling powerful combined full-text and vector search functionalities.
|
||||
- **ObVecVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/obvec_vector_store.py)): Uses [pyobvector](https://pypi.org/project/pyobvector/) against **OceanBase** or **seekdb** (MySQL-compatible wire protocol). Suitable when you already run OceanBase/seekdb or need a SQL-native vector table with HNSW-style ANN search and JSON metadata filters.
|
||||
|
||||
All Vector Store implementations inherit from **BaseVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/base_vector_store.py)), ensuring a consistent interface specification.
|
||||
All Vector Store implementations inherit from **BaseVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/base_vector_store.py)) in ReMe, ensuring a consistent interface specification.
|
||||
|
||||
## Core Features
|
||||
|
||||
|
|
@ -108,6 +109,28 @@ The asynchronous interface is particularly useful in the following scenarios:
|
|||
- **hosts**: Elasticsearch host address(es), either a string or a list (default: `http://localhost:9200`).
|
||||
- **basic_auth**: Basic authentication credentials (username and password).
|
||||
|
||||
### ObVecVectorStore Configuration
|
||||
|
||||
- **uri**: Server address as `host:port` (default: `127.0.0.1:2881`).
|
||||
- **user**: MySQL-compatible user. seekdb single-tenant images often use `root`; OceanBase multi-tenant setups typically use `root@<tenant>` (e.g. `root@test`).
|
||||
- **password**: Database password (seekdb Docker images commonly set this via `ROOT_PASSWORD`).
|
||||
- **database**: Logical database name (default: `test`).
|
||||
- **index_metric**: Distance metric for the vector index: `cosine` or `ip` (inner product); default `cosine`.
|
||||
- **index_ef_search**: HNSW `ef_search` parameter passed to pyobvector (default: `100`).
|
||||
- **collection_name**: Table name for the collection (from `VectorStoreConfig`, default `reme`). Use lowercase names if your deployment restricts identifiers.
|
||||
|
||||
**Local seekdb via Docker**
|
||||
|
||||
```text
|
||||
docker run -d --name reme_seekdb -p 2881:2881 -e ROOT_PASSWORD=<your_root_password> quay.io/oceanbase/seekdb:latest
|
||||
```
|
||||
|
||||
**Integration tests** (requires a running server, embedding API credentials in `.env`, and matching DB password):
|
||||
|
||||
```shell
|
||||
OBVEC_PASSWORD=<your_root_password> python tests/test_vector_store.py --obvec
|
||||
```
|
||||
|
||||
## Configuration File Examples
|
||||
|
||||
Configure Vector Store in `flowllm/config/default.yaml` under the `vector_store` section. The basic structure is as follows:
|
||||
|
|
@ -128,7 +151,7 @@ vector_store.default.params.<param_name>=<param_value>
|
|||
|
||||
### Configuration Field Descriptions
|
||||
|
||||
- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`.
|
||||
- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`, `obvec`.
|
||||
- **`embedding_model`** (required): Name of the embedding model configuration, referencing the `embedding_model` section.
|
||||
- **`params`** (optional): Dictionary of backend-specific parameters passed to the vector store constructor.
|
||||
|
||||
|
|
@ -295,6 +318,35 @@ vector_store.default.backend=elasticsearch
|
|||
vector_store.default.params.hosts='["http://es-node1:9200", "http://es-node2:9200", "http://es-node3:9200"]'
|
||||
```
|
||||
|
||||
#### 6. ObVecVectorStore Configuration (OceanBase / seekdb)
|
||||
|
||||
**Implementation**: [`reme/core/vector_store/obvec_vector_store.py`](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/obvec_vector_store.py)
|
||||
|
||||
**Example (seekdb on localhost)**:
|
||||
|
||||
```yaml
|
||||
vector_stores:
|
||||
default:
|
||||
backend: obvec
|
||||
embedding_model: default
|
||||
collection_name: reme
|
||||
uri: "127.0.0.1:2881"
|
||||
user: "root"
|
||||
password: "your-root-password"
|
||||
database: "test"
|
||||
index_metric: "cosine"
|
||||
index_ef_search: 100
|
||||
```
|
||||
|
||||
```shell
|
||||
vector_stores.default.backend=obvec
|
||||
vector_stores.default.uri=127.0.0.1:2881
|
||||
vector_stores.default.user=root
|
||||
vector_stores.default.password=your-root-password
|
||||
```
|
||||
|
||||
ReMe service YAML uses the key `vector_stores` (plural); CLI overrides use the same nested paths.
|
||||
|
||||
### Complete Configuration Example
|
||||
|
||||
Below is a complete `default.yaml` example including both embedding model and vector store configurations:
|
||||
|
|
@ -352,8 +404,9 @@ Two types of metadata filtering are supported:
|
|||
|
||||
- **Development & Testing**: Use MemoryVectorStore or LocalVectorStore—no additional services required.
|
||||
- **Small-Scale Applications**: Use LocalVectorStore or ChromaVectorStore for simplicity and ease of use.
|
||||
- **Production Environments**: Use QdrantVectorStore or EsVectorStore for high performance and scalability.
|
||||
- **Production Environments**: Use QdrantVectorStore, EsVectorStore, or ObVecVectorStore (OceanBase/seekdb) for high performance and scalability, depending on your existing infrastructure.
|
||||
- **Hybrid Search**: Use EsVectorStore to combine vector search with full-text search capabilities.
|
||||
- **OceanBase / seekdb**: Use ObVecVectorStore when you standardize on pyobvector and SQL-accessible vector tables.
|
||||
|
||||
## Important Notes
|
||||
|
||||
|
|
|
|||
|
|
@ -49,6 +49,9 @@ dependencies = [
|
|||
"openai>=2.8.1",
|
||||
"pandas>=2.3.3",
|
||||
"pydantic>=2.12.4",
|
||||
"pyobvector>=0.1.20",
|
||||
# pyobvector imports Expression from sqlglot; removed from sqlglot 30+ top-level API
|
||||
"sqlglot>=25,<30",
|
||||
"qdrant-client>=1.16.0",
|
||||
"tavily-python>=0.7.13",
|
||||
"tiktoken>=0.12.0",
|
||||
|
|
@ -83,7 +86,7 @@ litellm = [
|
|||
]
|
||||
|
||||
light = [
|
||||
"agentscope==1.0.17",
|
||||
"agentscope==1.0.18",
|
||||
"flowllm[reme]>=0.2.0.10",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from . import extension
|
|||
from . import memory
|
||||
from .reme import ReMe
|
||||
|
||||
__version__ = "0.3.1.6"
|
||||
__version__ = "0.3.1.8"
|
||||
|
||||
__all__ = [
|
||||
"config",
|
||||
|
|
|
|||
|
|
@ -173,28 +173,37 @@ class Application:
|
|||
if config.backend not in R.as_llms:
|
||||
logger.warning(f"AS LLM backend {config.backend} is not supported.")
|
||||
else:
|
||||
config_dict = config.model_dump(exclude={"backend"})
|
||||
if not config_dict.get("api_key", ""):
|
||||
config_dict["api_key"] = self.llm_api_key
|
||||
if "client_kwargs" not in config_dict:
|
||||
config_dict["client_kwargs"] = {}
|
||||
if not config_dict["client_kwargs"].get("base_url", ""):
|
||||
config_dict["client_kwargs"]["base_url"] = self.llm_base_url
|
||||
self.service_context.as_llms[name] = R.as_llms[config.backend](**config_dict)
|
||||
try:
|
||||
config_dict = config.model_dump(exclude={"backend"})
|
||||
if not config_dict.get("api_key", ""):
|
||||
config_dict["api_key"] = self.llm_api_key
|
||||
if "client_kwargs" not in config_dict:
|
||||
config_dict["client_kwargs"] = {}
|
||||
if not config_dict["client_kwargs"].get("base_url", ""):
|
||||
config_dict["client_kwargs"]["base_url"] = self.llm_base_url
|
||||
self.service_context.as_llms[name] = R.as_llms[config.backend](**config_dict)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize AS LLM '{name}': {e}")
|
||||
|
||||
for name, config in self.service_config.as_llm_formatters.items():
|
||||
if config.backend not in R.as_llm_formatters:
|
||||
logger.warning(f"AS LLM formatter backend {config.backend} is not supported.")
|
||||
else:
|
||||
config_dict = config.model_dump(exclude={"backend"})
|
||||
self.service_context.as_llm_formatters[name] = R.as_llm_formatters[config.backend](**config_dict)
|
||||
try:
|
||||
config_dict = config.model_dump(exclude={"backend"})
|
||||
self.service_context.as_llm_formatters[name] = R.as_llm_formatters[config.backend](**config_dict)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize AS LLM formatter '{name}': {e}")
|
||||
|
||||
for name, config in self.service_config.as_token_counters.items():
|
||||
if config.backend not in R.as_token_counters:
|
||||
logger.warning(f"Token counter backend {config.backend} is not supported.")
|
||||
else:
|
||||
config_dict = config.model_dump(exclude={"backend"})
|
||||
self.service_context.as_token_counters[name] = R.as_token_counters[config.backend](**config_dict)
|
||||
try:
|
||||
config_dict = config.model_dump(exclude={"backend"})
|
||||
self.service_context.as_token_counters[name] = R.as_token_counters[config.backend](**config_dict)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize AS token counter '{name}': {e}")
|
||||
|
||||
for name, config in self.service_config.llms.items():
|
||||
if config.backend not in R.llms:
|
||||
|
|
@ -287,15 +296,18 @@ class Application:
|
|||
logger.warning(f"AS LLM backend {config.get('backend')} is not supported.")
|
||||
continue
|
||||
|
||||
config_dict = {k: v for k, v in config.items() if k != "backend"}
|
||||
if not config_dict.get("api_key", ""):
|
||||
config_dict["api_key"] = self.llm_api_key
|
||||
if "client_kwargs" not in config_dict:
|
||||
config_dict["client_kwargs"] = {}
|
||||
if not config_dict["client_kwargs"].get("base_url", ""):
|
||||
config_dict["client_kwargs"]["base_url"] = self.llm_base_url
|
||||
self.service_context.as_llms[name] = R.as_llms[config["backend"]](**config_dict)
|
||||
logger.info(f"Restarted AS LLM: {name}")
|
||||
try:
|
||||
config_dict = {k: v for k, v in config.items() if k != "backend"}
|
||||
if not config_dict.get("api_key", ""):
|
||||
config_dict["api_key"] = self.llm_api_key
|
||||
if "client_kwargs" not in config_dict:
|
||||
config_dict["client_kwargs"] = {}
|
||||
if not config_dict["client_kwargs"].get("base_url", ""):
|
||||
config_dict["client_kwargs"]["base_url"] = self.llm_base_url
|
||||
self.service_context.as_llms[name] = R.as_llms[config["backend"]](**config_dict)
|
||||
logger.info(f"Restarted AS LLM: {name}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to restart AS LLM '{name}': {e}")
|
||||
|
||||
# as_llm_formatters
|
||||
if "as_llm_formatters" in restart_config:
|
||||
|
|
@ -308,9 +320,12 @@ class Application:
|
|||
if config.get("backend") not in R.as_llm_formatters:
|
||||
logger.warning(f"AS LLM formatter backend {config.get('backend')} is not supported.")
|
||||
continue
|
||||
config_dict = {k: v for k, v in config.items() if k != "backend"}
|
||||
self.service_context.as_llm_formatters[name] = R.as_llm_formatters[config["backend"]](**config_dict)
|
||||
logger.info(f"Restarted AS LLM formatter: {name}")
|
||||
try:
|
||||
config_dict = {k: v for k, v in config.items() if k != "backend"}
|
||||
self.service_context.as_llm_formatters[name] = R.as_llm_formatters[config["backend"]](**config_dict)
|
||||
logger.info(f"Restarted AS LLM formatter: {name}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to restart AS LLM formatter '{name}': {e}")
|
||||
|
||||
# as_token_counters
|
||||
if "as_token_counters" in restart_config:
|
||||
|
|
@ -323,9 +338,12 @@ class Application:
|
|||
if config.get("backend") not in R.as_token_counters:
|
||||
logger.warning(f"Token counter backend {config.get('backend')} is not supported.")
|
||||
continue
|
||||
config_dict = {k: v for k, v in config.items() if k != "backend"}
|
||||
self.service_context.as_token_counters[name] = R.as_token_counters[config["backend"]](**config_dict)
|
||||
logger.info(f"Restarted AS token counter: {name}")
|
||||
try:
|
||||
config_dict = {k: v for k, v in config.items() if k != "backend"}
|
||||
self.service_context.as_token_counters[name] = R.as_token_counters[config["backend"]](**config_dict)
|
||||
logger.info(f"Restarted AS token counter: {name}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to restart AS token counter '{name}': {e}")
|
||||
|
||||
# llms
|
||||
if "llms" in restart_config:
|
||||
|
|
|
|||
|
|
@ -5,19 +5,20 @@ import random
|
|||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from .base_file_store import BaseFileStore
|
||||
from ..enumeration import MemorySource
|
||||
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
|
||||
from ..utils import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
try:
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
|
||||
CHROMADB_AVAILABLE = True
|
||||
except ImportError:
|
||||
CHROMADB_AVAILABLE = False
|
||||
_CHROMADB_IMPORT_ERROR: Exception | None = None
|
||||
except Exception as e:
|
||||
_CHROMADB_IMPORT_ERROR = e
|
||||
chromadb = None
|
||||
Settings = None
|
||||
|
||||
|
|
@ -39,10 +40,8 @@ class ChromaFileStore(BaseFileStore):
|
|||
self,
|
||||
**kwargs,
|
||||
):
|
||||
if not CHROMADB_AVAILABLE:
|
||||
raise ImportError(
|
||||
"chromadb package is required for ChromaFileStore. Install it with: pip install chromadb",
|
||||
)
|
||||
if _CHROMADB_IMPORT_ERROR is not None:
|
||||
raise _CHROMADB_IMPORT_ERROR
|
||||
|
||||
super().__init__(**kwargs)
|
||||
self.client: "chromadb.ClientAPI | None" = None
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""SQLite storage backend for file store."""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
|
||||
import struct
|
||||
import time
|
||||
|
||||
|
|
@ -29,6 +29,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
def __init__(self, vec_ext_path: str = "", **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.vec_ext_path = vec_ext_path
|
||||
import sqlite3
|
||||
|
||||
self.conn: sqlite3.Connection | None = None
|
||||
|
||||
|
|
@ -61,6 +62,7 @@ class SqliteFileStore(BaseFileStore):
|
|||
"""Initialize database and load extensions."""
|
||||
if self.conn is not None:
|
||||
return
|
||||
import sqlite3
|
||||
|
||||
self.conn = sqlite3.connect(self.db_path / "reme.db", check_same_thread=False)
|
||||
|
||||
|
|
|
|||
|
|
@ -10,11 +10,11 @@ from tqdm import tqdm
|
|||
from .base_op import BaseOp
|
||||
from ..base_dict import BaseDict
|
||||
|
||||
_RAY_IMPORT_ERROR = None
|
||||
_RAY_IMPORT_ERROR: Exception | None = None
|
||||
|
||||
try:
|
||||
import ray
|
||||
except ImportError as _e:
|
||||
except Exception as _e:
|
||||
_RAY_IMPORT_ERROR = _e
|
||||
ray = None
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ 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 .obvec_vector_store import ObVecVectorStore
|
||||
from .pgvector_store import PGVectorStore
|
||||
from .qdrant_vector_store import QdrantVectorStore
|
||||
from ..registry_factory import R
|
||||
|
|
@ -13,6 +14,7 @@ __all__ = [
|
|||
"ChromaVectorStore",
|
||||
"ESVectorStore",
|
||||
"LocalVectorStore",
|
||||
"ObVecVectorStore",
|
||||
"PGVectorStore",
|
||||
"QdrantVectorStore",
|
||||
]
|
||||
|
|
@ -20,5 +22,6 @@ __all__ = [
|
|||
R.vector_stores.register("chroma")(ChromaVectorStore)
|
||||
R.vector_stores.register("es")(ESVectorStore)
|
||||
R.vector_stores.register("local")(LocalVectorStore)
|
||||
R.vector_stores.register("obvec")(ObVecVectorStore)
|
||||
R.vector_stores.register("pgvector")(PGVectorStore)
|
||||
R.vector_stores.register("qdrant")(QdrantVectorStore)
|
||||
|
|
|
|||
|
|
@ -9,12 +9,12 @@ from .base_vector_store import BaseVectorStore
|
|||
from ..embedding import BaseEmbeddingModel
|
||||
from ..schema import VectorNode
|
||||
|
||||
_CHROMADB_IMPORT_ERROR = None
|
||||
_CHROMADB_IMPORT_ERROR: Exception | None = None
|
||||
|
||||
try:
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
except ImportError as e:
|
||||
except Exception as e:
|
||||
_CHROMADB_IMPORT_ERROR = e
|
||||
chromadb = None
|
||||
Settings = None
|
||||
|
|
|
|||
|
|
@ -13,12 +13,12 @@ from .base_vector_store import BaseVectorStore
|
|||
from ..embedding import BaseEmbeddingModel
|
||||
from ..schema import VectorNode
|
||||
|
||||
_ELASTICSEARCH_IMPORT_ERROR = None
|
||||
_ELASTICSEARCH_IMPORT_ERROR: Exception | None = None
|
||||
|
||||
try:
|
||||
from elasticsearch import AsyncElasticsearch
|
||||
from elasticsearch.helpers import async_bulk
|
||||
except ImportError as e:
|
||||
except Exception as e:
|
||||
_ELASTICSEARCH_IMPORT_ERROR = e
|
||||
AsyncElasticsearch = None
|
||||
async_bulk = None
|
||||
|
|
|
|||
453
reme/core/vector_store/obvec_vector_store.py
Normal file
453
reme/core/vector_store/obvec_vector_store.py
Normal file
|
|
@ -0,0 +1,453 @@
|
|||
"""OceanBase / seekdb vector store for ReMe (pyobvector)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from loguru import logger
|
||||
from sqlalchemy import Column, JSON, String, text as sa_text
|
||||
from sqlalchemy.dialects.mysql import LONGTEXT
|
||||
|
||||
from .base_vector_store import BaseVectorStore
|
||||
from ..embedding import BaseEmbeddingModel
|
||||
from ..schema import VectorNode
|
||||
|
||||
_OBVECTOR_IMPORT_ERROR: Exception | None = None
|
||||
|
||||
try:
|
||||
from pyobvector import IndexParams, ObVecClient, VecIndexType, VECTOR
|
||||
from pyobvector import cosine_distance, inner_product
|
||||
except Exception as e:
|
||||
_OBVECTOR_IMPORT_ERROR = e
|
||||
IndexParams = None # type: ignore[misc, assignment]
|
||||
ObVecClient = None # type: ignore[misc, assignment]
|
||||
VecIndexType = None # type: ignore[misc, assignment]
|
||||
VECTOR = None # type: ignore[misc, assignment]
|
||||
|
||||
_COL_SELECT = "id, content, vector, metadata"
|
||||
|
||||
|
||||
def _is_safe_metadata_key(key: str) -> bool:
|
||||
return bool(key.replace("_", "").replace(".", "").isalnum())
|
||||
|
||||
|
||||
def _coerce_db_vector(raw: Any) -> list[float] | None:
|
||||
if raw is None:
|
||||
return None
|
||||
if isinstance(raw, list):
|
||||
return [float(x) for x in raw]
|
||||
if isinstance(raw, str):
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
if isinstance(parsed, list):
|
||||
return [float(x) for x in parsed]
|
||||
except (json.JSONDecodeError, TypeError, ValueError):
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _coerce_db_metadata(raw: Any) -> dict[str, Any]:
|
||||
if raw is None:
|
||||
return {}
|
||||
if isinstance(raw, dict):
|
||||
return raw
|
||||
if isinstance(raw, str):
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
return {}
|
||||
|
||||
|
||||
def _build_metadata_filter_sql(filters: dict[str, Any] | None) -> str:
|
||||
if not filters:
|
||||
return ""
|
||||
|
||||
parts: list[str] = []
|
||||
for key, value in filters.items():
|
||||
if not _is_safe_metadata_key(key):
|
||||
continue
|
||||
|
||||
path = f"$.{key}"
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
lo, hi = value[0], value[1]
|
||||
if isinstance(lo, (int, float)) and isinstance(hi, (int, float)):
|
||||
parts.append(
|
||||
f"(JSON_EXTRACT(metadata, '{path}') >= {lo} AND " f"JSON_EXTRACT(metadata, '{path}') <= {hi})",
|
||||
)
|
||||
else:
|
||||
parts.append(
|
||||
f"(JSON_EXTRACT(metadata, '{path}') >= '{lo}' AND " f"JSON_EXTRACT(metadata, '{path}') <= '{hi}')",
|
||||
)
|
||||
elif isinstance(value, (int, float)):
|
||||
parts.append(f"JSON_EXTRACT(metadata, '{path}') = {value}")
|
||||
else:
|
||||
parts.append(f"JSON_EXTRACT(metadata, '{path}') = '{value}'")
|
||||
|
||||
return " AND ".join(parts)
|
||||
|
||||
|
||||
def _format_vector_sql_literal(vector: list[float]) -> str:
|
||||
return "[" + ",".join(str(float(v)) for v in vector) + "]"
|
||||
|
||||
|
||||
def _normalize_embedding_for_ann(raw: Any) -> list[float]:
|
||||
if hasattr(raw, "tolist"):
|
||||
raw = raw.tolist()
|
||||
return [float(x) for x in raw]
|
||||
|
||||
|
||||
def _vector_node_from_db_row(row: tuple[Any, ...]) -> VectorNode:
|
||||
return VectorNode(
|
||||
vector_id=row[0],
|
||||
content=row[1] or "",
|
||||
vector=_coerce_db_vector(row[2]),
|
||||
metadata=_coerce_db_metadata(row[3]),
|
||||
)
|
||||
|
||||
|
||||
def _normalize_nodes(nodes: VectorNode | list[VectorNode]) -> list[VectorNode]:
|
||||
return [nodes] if isinstance(nodes, VectorNode) else list(nodes)
|
||||
|
||||
|
||||
def _sql_table(name: str) -> str:
|
||||
return f"`{name}`"
|
||||
|
||||
|
||||
class ObVecVectorStore(BaseVectorStore):
|
||||
"""OceanBase or seekdb vector store for dense vectors and kNN search.
|
||||
|
||||
Args:
|
||||
index_metric: ``cosine`` or ``ip`` (inner product). Invalid values raise
|
||||
``ValueError``; unsupported strings are not mapped to another metric.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: str,
|
||||
db_path: str | Path,
|
||||
embedding_model: BaseEmbeddingModel,
|
||||
uri: str = "127.0.0.1:2881",
|
||||
user: str = "root",
|
||||
password: str = "",
|
||||
database: str = "test",
|
||||
index_metric: str = "cosine",
|
||||
index_ef_search: int = 100,
|
||||
**kwargs,
|
||||
):
|
||||
if _OBVECTOR_IMPORT_ERROR is not None:
|
||||
raise ImportError(
|
||||
"ObVecVectorStore requires pyobvector. Install with `pip install pyobvector`",
|
||||
) from _OBVECTOR_IMPORT_ERROR
|
||||
|
||||
super().__init__(
|
||||
collection_name=collection_name,
|
||||
db_path=db_path,
|
||||
embedding_model=embedding_model,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
key = index_metric.strip().lower()
|
||||
if key not in ("cosine", "ip"):
|
||||
raise ValueError(
|
||||
f"ObVecVectorStore index_metric must be 'cosine' or 'ip', got {index_metric!r}",
|
||||
)
|
||||
self.uri = uri
|
||||
self.user = user
|
||||
self.password = password
|
||||
self.database = database
|
||||
self.index_metric = key
|
||||
self.index_ef_search = index_ef_search
|
||||
|
||||
self.client: ObVecClient | None = None
|
||||
self.embedding_model_dims = embedding_model.dimensions
|
||||
|
||||
def _require_client(self) -> ObVecClient:
|
||||
if self.client is None:
|
||||
raise RuntimeError("ObVecVectorStore.start() must be called before this operation")
|
||||
return self.client
|
||||
|
||||
async def list_collections(self) -> list[str]:
|
||||
client = self._require_client()
|
||||
result = client.perform_raw_text_sql(f"SHOW TABLES FROM {_sql_table(self.database)}")
|
||||
rows = result.fetchall()
|
||||
return [row[0] for row in rows if row]
|
||||
|
||||
def _table_columns_for_create(self, dimensions: int) -> list[Any]:
|
||||
return [
|
||||
Column("id", String(255), primary_key=True),
|
||||
Column("content", LONGTEXT),
|
||||
Column("vector", VECTOR(dimensions)),
|
||||
Column("metadata", JSON),
|
||||
]
|
||||
|
||||
def _hnsw_index_params(self, collection_name: str) -> IndexParams:
|
||||
metric = "cosine" if self.index_metric == "cosine" else "inner_product"
|
||||
vidxs = IndexParams()
|
||||
vidxs.add_index(
|
||||
"vector",
|
||||
VecIndexType.HNSW,
|
||||
f"{collection_name}_vidx",
|
||||
metric_type=metric,
|
||||
params={"efSearch": self.index_ef_search},
|
||||
)
|
||||
return vidxs
|
||||
|
||||
async def create_collection(self, collection_name: str, **kwargs):
|
||||
client = self._require_client()
|
||||
dimensions = kwargs.get("dimensions", self.embedding_model_dims)
|
||||
|
||||
if client.check_table_exists(collection_name):
|
||||
logger.info("Collection {} already exists", collection_name)
|
||||
return
|
||||
|
||||
columns = self._table_columns_for_create(dimensions)
|
||||
vidxs = self._hnsw_index_params(collection_name)
|
||||
|
||||
client.create_table_with_index_params(
|
||||
table_name=collection_name,
|
||||
columns=columns,
|
||||
vidxs=vidxs,
|
||||
)
|
||||
logger.info("Created collection {} with dimensions={}", collection_name, dimensions)
|
||||
|
||||
async def delete_collection(self, collection_name: str, **kwargs):
|
||||
client = self._require_client()
|
||||
client.drop_table_if_exist(collection_name)
|
||||
logger.info("Deleted collection {}", collection_name)
|
||||
|
||||
async def copy_collection(self, collection_name: str, **kwargs):
|
||||
client = self._require_client()
|
||||
|
||||
if not client.check_table_exists(self.collection_name):
|
||||
raise ValueError(f"Source collection {self.collection_name} does not exist")
|
||||
|
||||
await self.create_collection(collection_name)
|
||||
|
||||
try:
|
||||
source_data = await self.list(limit=None)
|
||||
if source_data:
|
||||
await self.insert(source_data, collection_name=collection_name)
|
||||
logger.info("Copied collection {} to {}", self.collection_name, collection_name)
|
||||
except Exception:
|
||||
try:
|
||||
client.drop_table_if_exist(collection_name)
|
||||
except Exception as cleanup_err:
|
||||
logger.warning("Cleanup after failed copy failed: {}", cleanup_err)
|
||||
raise
|
||||
|
||||
async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
|
||||
nodes = _normalize_nodes(nodes)
|
||||
if not nodes:
|
||||
return
|
||||
|
||||
client = self._require_client()
|
||||
|
||||
need_emb = [n for n in nodes if n.vector is None]
|
||||
if need_emb:
|
||||
filled = await self.get_node_embeddings(need_emb)
|
||||
by_id = {n.vector_id: n for n in filled}
|
||||
nodes_to_insert = [by_id.get(n.vector_id, n) for n in nodes]
|
||||
else:
|
||||
nodes_to_insert = nodes
|
||||
|
||||
data = [
|
||||
{
|
||||
"id": node.vector_id,
|
||||
"content": node.content,
|
||||
"vector": node.vector if node.vector is not None else [],
|
||||
"metadata": node.metadata if node.metadata else {},
|
||||
}
|
||||
for node in nodes_to_insert
|
||||
]
|
||||
target = kwargs.get("collection_name", self.collection_name)
|
||||
client.insert(table_name=target, data=data)
|
||||
logger.info("Inserted {} documents into {}", len(nodes_to_insert), target)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query: str,
|
||||
limit: int = 5,
|
||||
filters: dict | None = None,
|
||||
**kwargs,
|
||||
) -> list[VectorNode]:
|
||||
client = self._require_client()
|
||||
raw_vec = await self.get_embedding(query)
|
||||
query_vector = _normalize_embedding_for_ann(raw_vec)
|
||||
dist_fn = cosine_distance if self.index_metric == "cosine" else inner_product
|
||||
|
||||
filter_sql = _build_metadata_filter_sql(filters)
|
||||
where_parts = [sa_text(filter_sql)] if filter_sql else None
|
||||
|
||||
results = client.ann_search(
|
||||
table_name=self.collection_name,
|
||||
vec_data=query_vector,
|
||||
vec_column_name="vector",
|
||||
distance_func=dist_fn,
|
||||
with_dist=True,
|
||||
topk=limit,
|
||||
output_column_names=["id", "content", "metadata"],
|
||||
where_clause=where_parts,
|
||||
)
|
||||
|
||||
score_threshold = kwargs.get("score_threshold")
|
||||
out: list[VectorNode] = []
|
||||
for row in results:
|
||||
if len(row) < 4:
|
||||
raise RuntimeError(
|
||||
"ann_search row must have id, content, metadata, distance " f"(got {len(row)} columns)",
|
||||
)
|
||||
vid, content, metadata_raw, distance = row[0], row[1], row[2], row[3]
|
||||
dist_f = float(distance)
|
||||
if self.index_metric == "cosine":
|
||||
score = max(0.0, 1.0 - dist_f / 2.0)
|
||||
else:
|
||||
score = max(0.0, dist_f)
|
||||
if score_threshold is not None and score < score_threshold:
|
||||
continue
|
||||
meta = _coerce_db_metadata(metadata_raw) if metadata_raw is not None else {}
|
||||
meta["score"] = score
|
||||
meta["_distance"] = dist_f
|
||||
out.append(
|
||||
VectorNode(
|
||||
vector_id=vid,
|
||||
content=content or "",
|
||||
vector=None,
|
||||
metadata=meta,
|
||||
),
|
||||
)
|
||||
return out
|
||||
|
||||
async def delete(self, vector_ids: str | list[str], **kwargs):
|
||||
if isinstance(vector_ids, str):
|
||||
vector_ids = [vector_ids]
|
||||
if not vector_ids:
|
||||
return
|
||||
client = self._require_client()
|
||||
client.delete(self.collection_name, ids=vector_ids)
|
||||
logger.info("Deleted {} documents from {}", len(vector_ids), self.collection_name)
|
||||
|
||||
async def delete_all(self, **kwargs):
|
||||
client = self._require_client()
|
||||
client.delete(self.collection_name)
|
||||
logger.info("Deleted all documents from {}", self.collection_name)
|
||||
|
||||
async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
|
||||
nodes = _normalize_nodes(nodes)
|
||||
if not nodes:
|
||||
return
|
||||
|
||||
client = self._require_client()
|
||||
need_emb = [n for n in nodes if n.vector is None and bool(n.content)]
|
||||
if need_emb:
|
||||
filled = await self.get_node_embeddings(need_emb)
|
||||
by_id = {n.vector_id: n for n in filled}
|
||||
nodes_to_update = [
|
||||
by_id.get(n.vector_id, n) if (n.vector is None and bool(n.content)) else n for n in nodes
|
||||
]
|
||||
else:
|
||||
nodes_to_update = nodes
|
||||
|
||||
for node in nodes_to_update:
|
||||
updates: list[str] = []
|
||||
params: dict[str, Any] = {}
|
||||
|
||||
if node.content is not None:
|
||||
updates.append("content = :content")
|
||||
params["content"] = node.content
|
||||
|
||||
if node.vector is not None:
|
||||
updates.append("vector = :vector")
|
||||
params["vector"] = _format_vector_sql_literal(node.vector)
|
||||
|
||||
if node.metadata is not None:
|
||||
updates.append("metadata = :metadata")
|
||||
params["metadata"] = json.dumps(node.metadata)
|
||||
|
||||
if not updates:
|
||||
continue
|
||||
|
||||
params["vid"] = node.vector_id
|
||||
update_sql = f"UPDATE {_sql_table(self.collection_name)} SET {', '.join(updates)} WHERE id = :vid"
|
||||
with client.engine.connect() as conn:
|
||||
with conn.begin():
|
||||
conn.execute(sa_text(update_sql), params)
|
||||
|
||||
logger.info("Updated {} documents in {}", len(nodes_to_update), self.collection_name)
|
||||
|
||||
async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None:
|
||||
single = isinstance(vector_ids, str)
|
||||
if single:
|
||||
vector_ids = [vector_ids]
|
||||
if not vector_ids:
|
||||
return [] if not single else None
|
||||
|
||||
client = self._require_client()
|
||||
ids_str = "', '".join(vector_ids)
|
||||
select_sql = f"SELECT {_COL_SELECT} FROM {_sql_table(self.collection_name)} " f"WHERE id IN ('{ids_str}')"
|
||||
result = client.perform_raw_text_sql(select_sql)
|
||||
rows = result.fetchall()
|
||||
parsed = [_vector_node_from_db_row(row) for row in rows if row]
|
||||
if single:
|
||||
return parsed[0] if parsed else None
|
||||
return parsed
|
||||
|
||||
async def list(
|
||||
self,
|
||||
filters: dict | None = None,
|
||||
limit: int | None = None,
|
||||
sort_key: str | None = None,
|
||||
reverse: bool = False,
|
||||
) -> list[VectorNode]:
|
||||
client = self._require_client()
|
||||
select_sql = f"SELECT {_COL_SELECT} FROM {_sql_table(self.collection_name)}"
|
||||
where_clause = _build_metadata_filter_sql(filters)
|
||||
if where_clause:
|
||||
select_sql += f" WHERE {where_clause}"
|
||||
if sort_key and _is_safe_metadata_key(sort_key):
|
||||
order = "DESC" if reverse else "ASC"
|
||||
select_sql += f" ORDER BY JSON_EXTRACT(metadata, '$.{sort_key}') {order}"
|
||||
if limit is not None:
|
||||
select_sql += f" LIMIT {limit}"
|
||||
result = client.perform_raw_text_sql(select_sql)
|
||||
rows = result.fetchall()
|
||||
return [_vector_node_from_db_row(row) for row in rows if row]
|
||||
|
||||
async def collection_info(self) -> dict[str, Any]:
|
||||
"""Return collection name and row count."""
|
||||
client = self._require_client()
|
||||
count_sql = f"SELECT COUNT(*) FROM {_sql_table(self.collection_name)}"
|
||||
result = client.perform_raw_text_sql(count_sql)
|
||||
row = result.fetchone()
|
||||
count = row[0] if row else 0
|
||||
return {"name": self.collection_name, "count": count}
|
||||
|
||||
async def reset(self):
|
||||
"""Drop and recreate the current collection table."""
|
||||
logger.warning("Resetting collection {}...", self.collection_name)
|
||||
await self.delete_collection(self.collection_name)
|
||||
await self.create_collection(self.collection_name)
|
||||
|
||||
async def reset_collection(self, collection_name: str):
|
||||
self.collection_name = collection_name
|
||||
await self.create_collection(collection_name)
|
||||
logger.info("Collection reset to {}", collection_name)
|
||||
|
||||
async def start(self) -> None:
|
||||
self.client = ObVecClient(
|
||||
uri=self.uri,
|
||||
user=self.user,
|
||||
password=self.password,
|
||||
db_name=self.database,
|
||||
)
|
||||
|
||||
await super().start()
|
||||
logger.info("seekdb / OceanBase vector table {} ready", self.collection_name)
|
||||
|
||||
async def close(self):
|
||||
self.client = None
|
||||
logger.info("ObVec client connection closed")
|
||||
|
|
@ -11,12 +11,12 @@ from .base_vector_store import BaseVectorStore
|
|||
from ..embedding import BaseEmbeddingModel
|
||||
from ..schema import VectorNode
|
||||
|
||||
_ASYNCPG_IMPORT_ERROR = None
|
||||
_ASYNCPG_IMPORT_ERROR: Exception | None = None
|
||||
|
||||
try:
|
||||
import asyncpg
|
||||
from asyncpg import Pool
|
||||
except ImportError as e:
|
||||
except Exception as e:
|
||||
_ASYNCPG_IMPORT_ERROR = e
|
||||
asyncpg = None
|
||||
Pool = None
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from .base_vector_store import BaseVectorStore
|
|||
from ..embedding import BaseEmbeddingModel
|
||||
from ..schema import VectorNode
|
||||
|
||||
_QDRANT_IMPORT_ERROR = None
|
||||
_QDRANT_IMPORT_ERROR: Exception | None = None
|
||||
|
||||
try:
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
|
|
@ -23,7 +23,7 @@ try:
|
|||
Range,
|
||||
VectorParams,
|
||||
)
|
||||
except ImportError as e:
|
||||
except Exception as e:
|
||||
_QDRANT_IMPORT_ERROR = e
|
||||
AsyncQdrantClient = None
|
||||
Distance = None
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ class Compactor(BaseOp):
|
|||
console_enabled: bool = False,
|
||||
return_dict: bool = False,
|
||||
add_thinking_block: bool = True,
|
||||
extra_instruction: str = "",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
|
|
@ -42,6 +43,7 @@ class Compactor(BaseOp):
|
|||
self.console_enabled: bool = console_enabled
|
||||
self.return_dict: bool = return_dict
|
||||
self.add_thinking_block: bool = add_thinking_block
|
||||
self.extra_instruction: str = extra_instruction
|
||||
|
||||
# pylint: disable=too-many-return-statements
|
||||
async def execute(self):
|
||||
|
|
@ -84,6 +86,9 @@ class Compactor(BaseOp):
|
|||
)
|
||||
else:
|
||||
user_message: str = f"# conversation\n{history_formatted_str}\n\n" + self.get_prompt("initial_user_message")
|
||||
|
||||
if self.extra_instruction:
|
||||
user_message += f"\n\n# extra-instruction\n{self.extra_instruction}"
|
||||
logger.info(f"Compactor sys_prompt={agent.sys_prompt} user_message={user_message}")
|
||||
|
||||
compact_msg: Msg = await agent.reply(
|
||||
|
|
|
|||
|
|
@ -369,6 +369,7 @@ class ReMeLight(Application):
|
|||
previous_summary: str = "",
|
||||
return_dict: bool = False,
|
||||
add_thinking_block: bool = True,
|
||||
extra_instruction: str = "",
|
||||
) -> str | dict:
|
||||
"""
|
||||
Compact a list of messages into a condensed summary.
|
||||
|
|
@ -395,6 +396,12 @@ class ReMeLight(Application):
|
|||
summary for continuity. Defaults to empty string.
|
||||
return_dict (bool): If True, returns a dict with user_message,
|
||||
history_compact, and is_valid. Defaults to False.
|
||||
add_thinking_block (bool): If True, adds a thinking block to the summary.
|
||||
extra_instruction (str): Optional additional instruction appended to the
|
||||
compaction prompt. Use this to guide what information to keep or
|
||||
remove. For example: "Remove debug logs and tool-call details. Keep
|
||||
requirements, decisions, and pending tasks." Defaults to empty string
|
||||
(no extra instruction, preserving default behavior).
|
||||
|
||||
Returns:
|
||||
str | dict: The condensed summary string, or a dict containing
|
||||
|
|
@ -410,6 +417,7 @@ class ReMeLight(Application):
|
|||
language=language if language == "zh" else "",
|
||||
return_dict=return_dict,
|
||||
add_thinking_block=add_thinking_block,
|
||||
extra_instruction=extra_instruction,
|
||||
)
|
||||
|
||||
return await compactor.call(
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@
|
|||
"""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.
|
||||
PGVectorStore, QdrantVectorStore, ChromaVectorStore, and ObVecVectorStore implementations.
|
||||
Tests can be run for specific vector stores or all implementations.
|
||||
|
||||
Usage:
|
||||
python test_vector_store.py --local # Test LocalVectorStore only
|
||||
|
|
@ -11,12 +11,13 @@ Usage:
|
|||
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 --obvec # Test ObVecVectorStore only (needs seekdb / OceanBase)
|
||||
python test_vector_store.py --all # Test all vector stores
|
||||
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
|
@ -32,12 +33,19 @@ from reme.core.vector_store import (
|
|||
ChromaVectorStore,
|
||||
LocalVectorStore,
|
||||
ESVectorStore,
|
||||
ObVecVectorStore,
|
||||
PGVectorStore,
|
||||
QdrantVectorStore,
|
||||
)
|
||||
|
||||
load_env()
|
||||
|
||||
|
||||
def _search_score_for_log(metadata: dict) -> object:
|
||||
"""Similarity score for log lines (implementations use ``metadata['score']``)."""
|
||||
return metadata.get("score", metadata.get("_score", "N/A"))
|
||||
|
||||
|
||||
# ==================== Configuration ====================
|
||||
|
||||
|
||||
|
|
@ -73,6 +81,15 @@ class TestConfig:
|
|||
CHROMA_TENANT = None # Set for ChromaDB Cloud tenant
|
||||
CHROMA_DATABASE = None # Set for ChromaDB Cloud database
|
||||
|
||||
# ObVecVectorStore: seekdb docker often uses user `root` + ROOT_PASSWORD; OceanBase
|
||||
# multi-tenant commonly uses `root@<tenant>` (see pyobvector defaults).
|
||||
# OBVEC_PASSWORD default `root` matches docker-compose.obvec.yml only—override if your
|
||||
# seekdb uses another ROOT_PASSWORD (e.g. another compose stack on the same port).
|
||||
OBVEC_URI = os.environ.get("OBVEC_URI", "127.0.0.1:2881")
|
||||
OBVEC_USER = os.environ.get("OBVEC_USER", "root")
|
||||
OBVEC_PASSWORD = os.environ.get("OBVEC_PASSWORD", "root")
|
||||
OBVEC_DATABASE = os.environ.get("OBVEC_DATABASE", "test")
|
||||
|
||||
# Embedding model settings
|
||||
EMBEDDING_MODEL_NAME = "text-embedding-v4"
|
||||
EMBEDDING_DIMENSIONS = 64
|
||||
|
|
@ -182,7 +199,7 @@ def get_store_type(store: BaseVectorStore) -> str:
|
|||
store: Vector store instance
|
||||
|
||||
Returns:
|
||||
str: Type identifier ("local", "es", "pgvector", "qdrant", or "chroma")
|
||||
str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", or "obvec")
|
||||
"""
|
||||
if isinstance(store, LocalVectorStore):
|
||||
return "local"
|
||||
|
|
@ -194,6 +211,8 @@ def get_store_type(store: BaseVectorStore) -> str:
|
|||
return "pgvector"
|
||||
elif isinstance(store, ChromaVectorStore):
|
||||
return "chroma"
|
||||
elif isinstance(store, ObVecVectorStore):
|
||||
return "obvec"
|
||||
else:
|
||||
raise ValueError(f"Unknown vector store type: {type(store)}")
|
||||
|
||||
|
|
@ -202,7 +221,7 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor
|
|||
"""Create a vector store instance based on type.
|
||||
|
||||
Args:
|
||||
store_type: Type of vector store ("local", "es", "pgvector", "qdrant", or "chroma")
|
||||
store_type: Type of vector store ("local", "es", "pgvector", "qdrant", "chroma", or "obvec")
|
||||
collection_name: Name of the collection
|
||||
|
||||
Returns:
|
||||
|
|
@ -264,6 +283,18 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor
|
|||
tenant=config.CHROMA_TENANT,
|
||||
database=config.CHROMA_DATABASE,
|
||||
)
|
||||
elif store_type == "obvec":
|
||||
return ObVecVectorStore(
|
||||
collection_name=collection_name,
|
||||
embedding_model=embedding_model,
|
||||
db_path=tempfile.mkdtemp(prefix="test_obvec_"),
|
||||
uri=config.OBVEC_URI,
|
||||
user=config.OBVEC_USER,
|
||||
password=config.OBVEC_PASSWORD,
|
||||
database=config.OBVEC_DATABASE,
|
||||
index_metric="cosine",
|
||||
index_ef_search=100,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown store type: {store_type}")
|
||||
|
||||
|
|
@ -327,7 +358,7 @@ async def test_search(store: BaseVectorStore, _store_name: str):
|
|||
|
||||
logger.info(f"Search returned {len(results)} results")
|
||||
for i, r in enumerate(results, 1):
|
||||
score = r.metadata.get("_score", "N/A")
|
||||
score = _search_score_for_log(r.metadata)
|
||||
logger.info(f" Result {i}: {r.content[:60]}... (score: {score})")
|
||||
|
||||
assert len(results) > 0, "Search should return results"
|
||||
|
|
@ -581,9 +612,9 @@ async def test_copy_collection(store: BaseVectorStore, store_name: str):
|
|||
config = TestConfig()
|
||||
copy_collection_name = f"{config.TEST_COLLECTION_PREFIX}_{store_name}_copy"
|
||||
|
||||
# Elasticsearch and PostgreSQL require lowercase table/index names
|
||||
# Elasticsearch, PostgreSQL and OceanBase require lowercase table/index names
|
||||
store_type = get_store_type(store)
|
||||
if store_type in ("es", "pgvector"):
|
||||
if store_type in ("es", "pgvector", "obvec"):
|
||||
copy_collection_name = copy_collection_name.lower()
|
||||
|
||||
# Clean up if exists
|
||||
|
|
@ -1010,7 +1041,7 @@ async def test_search_relevance_ranking(store: BaseVectorStore, _store_name: str
|
|||
|
||||
logger.info(f"Search results for: '{query}'")
|
||||
for i, result in enumerate(results, 1):
|
||||
score = result.metadata.get("_score", "N/A")
|
||||
score = _search_score_for_log(result.metadata)
|
||||
relevance = result.metadata.get("relevance", "unknown")
|
||||
logger.info(f" {i}. [{relevance}] score={score}: {result.content[:60]}...")
|
||||
|
||||
|
|
@ -1028,7 +1059,7 @@ async def test_search_relevance_ranking(store: BaseVectorStore, _store_name: str
|
|||
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")
|
||||
score = _search_score_for_log(result.metadata)
|
||||
logger.info(f" {i}. score={score}: {result.content[:60]}...")
|
||||
|
||||
logger.info("✓ Search relevance ranking test passed")
|
||||
|
|
@ -1718,7 +1749,7 @@ async def cleanup_store(store: BaseVectorStore, store_type: str):
|
|||
|
||||
Args:
|
||||
store: Vector store instance
|
||||
store_type: Type of vector store ("local" or "es")
|
||||
store_type: Backend key (e.g. ``"local"``, ``"obvec"``)
|
||||
"""
|
||||
logger.info("=" * 20 + " CLEANUP " + "=" * 20)
|
||||
|
||||
|
|
@ -1752,6 +1783,13 @@ async def cleanup_store(store: BaseVectorStore, store_type: str):
|
|||
shutil.rmtree(test_dir)
|
||||
logger.info(f"Cleaned up chroma directory: {config.CHROMA_PATH}")
|
||||
|
||||
# ObVecVectorStore uses a temp db_path per run (reserved for local sidecar files).
|
||||
if store_type == "obvec":
|
||||
obvec_dir = getattr(store, "db_path", None)
|
||||
if obvec_dir and Path(obvec_dir).exists():
|
||||
shutil.rmtree(obvec_dir, ignore_errors=True)
|
||||
logger.info(f"Cleaned up obvec temp directory: {obvec_dir}")
|
||||
|
||||
logger.info("✓ Cleanup completed")
|
||||
except Exception as e:
|
||||
logger.error(f"Cleanup error: {e}")
|
||||
|
|
@ -1772,6 +1810,7 @@ Examples:
|
|||
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 --obvec # Test ObVecVectorStore (seekdb / OceanBase)
|
||||
python test_vector_store.py --all # Test all vector stores
|
||||
""",
|
||||
)
|
||||
|
|
@ -1800,6 +1839,11 @@ Examples:
|
|||
action="store_true",
|
||||
help="Test ChromaVectorStore",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--obvec",
|
||||
action="store_true",
|
||||
help="Test ObVecVectorStore",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--all",
|
||||
action="store_true",
|
||||
|
|
@ -1818,6 +1862,7 @@ Examples:
|
|||
("pgvector", "PGVectorStore"),
|
||||
("qdrant", "QdrantVectorStore"),
|
||||
("chroma", "ChromaVectorStore"),
|
||||
("obvec", "ObVecVectorStore"),
|
||||
]
|
||||
else:
|
||||
# Build list based on individual flags
|
||||
|
|
@ -1831,6 +1876,8 @@ Examples:
|
|||
stores_to_test.append(("qdrant", "QdrantVectorStore"))
|
||||
if args.chroma:
|
||||
stores_to_test.append(("chroma", "ChromaVectorStore"))
|
||||
if args.obvec:
|
||||
stores_to_test.append(("obvec", "ObVecVectorStore"))
|
||||
|
||||
if not stores_to_test:
|
||||
# Default to all vector stores if no argument provided
|
||||
|
|
@ -1840,10 +1887,11 @@ Examples:
|
|||
("pgvector", "PGVectorStore"),
|
||||
("qdrant", "QdrantVectorStore"),
|
||||
("chroma", "ChromaVectorStore"),
|
||||
("obvec", "ObVecVectorStore"),
|
||||
]
|
||||
print("No vector store specified, defaulting to test all vector stores")
|
||||
print(
|
||||
"Use --local/--es/--pgvector/--qdrant/--chroma to test specific ones\n",
|
||||
"Use --local/--es/--pgvector/--qdrant/--chroma/--obvec to test specific ones\n",
|
||||
)
|
||||
|
||||
# Run tests for each vector store
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue