docs: 更新 README feature 表 + 升级指南 + 示例注释 + roadmap 同步
- README 添加 feature 组合表 + 模块级 features 清单 + 升级指南 - 18 个 example 顶部添加 Required features 注释 - roadmap.md 和 roadmap-v0.3.2.md 同步 Phase 26-27 完成状态 - cargo fmt 全量格式化(修复预存格式问题,CI format job 可通过)
This commit is contained in:
@@ -38,7 +38,7 @@ use std::sync::Arc;
|
||||
use agcore::agent::{Agent, AgentBuilder, AgentSession};
|
||||
use agcore::llm::hooks::HookExecutor;
|
||||
use agcore::llm::mock::MockProvider;
|
||||
use agcore::llm::provider::LlmProvider;
|
||||
use agcore::llm::LlmProvider;
|
||||
use agcore::llm::types::message::{ContentBlock, Message};
|
||||
use agcore::llm::types::response_v2::{MessageResponse, StopReason};
|
||||
use agcore::llm::types::Usage;
|
||||
@@ -133,6 +133,89 @@ let provider = create_provider(
|
||||
| `bridge_keys_demo` | 桥接键:Agent 间上下文键值透传 |
|
||||
| `agent_switch_demo` | Agent 热切换:会话中动态切换 Agent 角色 |
|
||||
|
||||
## Feature 组合
|
||||
|
||||
AG Core 通过 Cargo features 让下游按需选择模块,跳过不需要的编译单元和重型依赖。`default = ["full"]` 保持向后兼容——不指定 features 时行为与 v0.3.0 一致。
|
||||
|
||||
### 快捷组合
|
||||
|
||||
| 组合 | 场景 | 包含的 features |
|
||||
|------|------|----------------|
|
||||
| `full`(default) | 全栈使用,兼容 v0.3.0 | 全部 16 个 feature |
|
||||
| `light` | 生产常用,跳过 Anthropic/DeepSeek/Qwen/Ollama | llm + provider-openai + tools + tools-mcp + memory + agent + engine + prompt + document |
|
||||
| `chat` | 纯对话(跳过 SQLite 和 MCP) | agent + provider-openai |
|
||||
| `multi` | 多 Agent 复合(chat + subagent + switch + checkpointer) | engine + provider-openai |
|
||||
|
||||
### Cargo.toml 配置示例
|
||||
|
||||
```toml
|
||||
# 默认全栈(兼容 v0.3.0)
|
||||
[dependencies]
|
||||
agcore = "0.3"
|
||||
|
||||
# 纯对话场景:跳过 SQLite 和 MCP,编译更快
|
||||
[dependencies]
|
||||
agcore = { version = "0.3", default-features = false, features = ["chat", "provider-openai"] }
|
||||
|
||||
# 生产常用:OpenAI + 工具 + 记忆 + Agent
|
||||
[dependencies]
|
||||
agcore = { version = "0.3", default-features = false, features = ["light"] }
|
||||
|
||||
# 多 Agent 复合 + MCP 工具
|
||||
[dependencies]
|
||||
agcore = { version = "0.3", default-features = false, features = ["multi", "provider-openai", "tools-mcp"] }
|
||||
```
|
||||
|
||||
### 模块级 features
|
||||
|
||||
如需更细粒度控制,可单独启用模块级 features:
|
||||
|
||||
| Feature | 覆盖内容 | imply |
|
||||
|---------|---------|-------|
|
||||
| `document` | Document + RecursiveCharacterSplitter | — |
|
||||
| `llm-types` | Message / ToolDef / Usage 等 IR 类型 | — |
|
||||
| `prompt` | PromptTemplate + PromptComposer | `llm-types` |
|
||||
| `llm` | Provider trait + LlmCycle + hooks + compact + embedding + mock | `llm-types` |
|
||||
| `tools` | BaseTool + ToolRegistry | `llm-types` |
|
||||
| `tools-mcp` | McpClient(Stdio/StreamableHttp) | `tools` |
|
||||
| `memory` | MemoryStore + Conversation + VectorStore + KnowledgeGraph + Retriever | `document` + `llm` |
|
||||
| `memory-sqlite` | SqliteStore | `memory` |
|
||||
| `agent` | Agent + Builder + Session + ContextSlot + Summary | `llm` + `tools` + `memory` |
|
||||
| `engine` | SessionManager + Checkpointer + SubAgent + Switch | `agent` |
|
||||
| `provider-openai` | OpenAI Provider 实现 | `llm` |
|
||||
| `provider-anthropic` | Anthropic Provider 实现 | `llm` |
|
||||
| `provider-deepseek` | DeepSeek Provider 实现 | `llm` |
|
||||
| `provider-qwen` | Qwen Provider 实现 | `llm` |
|
||||
| `provider-ollama` | Ollama Provider 实现 | `llm` |
|
||||
| `tracing-init` | `init_tracing()` 函数 | — |
|
||||
|
||||
## 升级指南(v0.3.0 → v0.3.2)
|
||||
|
||||
### LlmProvider trait 路径变更
|
||||
|
||||
v0.3.2 起,`LlmProvider` trait 及其关联类型 `ProviderCapabilities` / `ProviderFeatures` 从 `provider` 模块移至 `llm` 模块根级别,归属 `#[cfg(feature = "llm")]` 而非 `any(provider-*)`。纯 Mock 场景不再需要引入任何 provider feature。
|
||||
|
||||
| 旧路径(v0.3.0) | 新路径(v0.3.2) |
|
||||
|------------------|------------------|
|
||||
| `agcore::llm::provider::LlmProvider` | `agcore::llm::LlmProvider` |
|
||||
| `agcore::llm::provider::ProviderCapabilities` | `agcore::llm::ProviderCapabilities` |
|
||||
| `agcore::llm::provider::ProviderFeatures` | `agcore::llm::ProviderFeatures` |
|
||||
|
||||
**向后兼容**:`provider` 模块中保留了 `pub use` 重导出,老路径仍可编译。但推荐迁移至新路径,未来版本可能移除重导出。
|
||||
|
||||
`ProviderConfig` / `ProviderType` / `create_provider()` 等 provider 创建逻辑仍在 `agcore::llm::provider` 下,无需迁移。
|
||||
|
||||
### 迁移步骤
|
||||
|
||||
```bash
|
||||
# 1. 全局替换 use 路径
|
||||
sed -i 's/agcore::llm::provider::LlmProvider/agcore::llm::LlmProvider/g' src/**/*.rs
|
||||
sed -i 's/agcore::llm::provider::{LlmProvider/agcore::llm::{LlmProvider/g' src/**/*.rs
|
||||
|
||||
# 2. 验证编译
|
||||
cargo build --features "full"
|
||||
```
|
||||
|
||||
## 核心模块
|
||||
|
||||
| 模块 | 一句话说明 |
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
# AG Core Roadmap — v0.3.2
|
||||
|
||||
**状态**:✅ Phase 20-25 已交付(Step 1 合并实施);⏳ Phase 26-27 待实施
|
||||
**状态**:✅ Phase 20-27 全部交付(v0.3.2 交付完毕)
|
||||
|
||||
> 本文件聚焦 **v0.3.2 版本** 的规划与交付(Phase 20–27)。
|
||||
> 返回总入口:[`roadmap.md`](./roadmap.md)
|
||||
|
||||
> **进度更新(2026-07-19)**:Step 1 合并完成 Phase 20(Cargo.toml features 定义 + 依赖 optional 化)+ Phase 21-25 的全部 cfg 门控(因编译阻塞提前)。验证矩阵 14 条编译命令全部通过,`cargo test -F full` 427 passed。实施方案见 [`docs/26-step1-phase20-cargo-features-implementation.md`](./26-step1-phase20-cargo-features-implementation.md)。下一步:Phase 26(CI 测试矩阵)+ Phase 27(文档更新)。
|
||||
> **进度更新(2026-07-19)**:v0.3.2 全部交付。Step 1 完成 Phase 20-25(Cargo features 定义 + 依赖 optional 化 + 全模块 cfg 门控);Step 3 完成 Phase 26-27(CI 测试矩阵固化 + LlmProvider trait 归属修正 + examples required-features + 文档更新)。验证矩阵 6 种组合全部通过,clippy 0 警告。实施方案见 [`docs/26-step1-phase20-cargo-features-implementation.md`](./26-step1-phase20-cargo-features-implementation.md) 和 [`docs/27-step3-phase26-ci-verification.md`](./27-step3-phase26-ci-verification.md)。
|
||||
|
||||
## v0.3.2 愿景
|
||||
|
||||
@@ -218,7 +218,7 @@
|
||||
**依赖**:Phase 20-25(所有门控就绪)
|
||||
**优先级**:P0
|
||||
**预估规模**:约 10 行(CI 配置)
|
||||
**状态**:⏳ 待实施
|
||||
**状态**:✅ 已交付(2026-07-19)— 6 种 feature 组合测试矩阵 + clippy + format + examples 验证 job 全部通过;`RUSTFLAGS=-D warnings` 强制零警告;LlmProvider trait 归属修正 + bundle() 门控 + llm feature 补充 imply futures-util
|
||||
|
||||
---
|
||||
|
||||
@@ -235,7 +235,7 @@
|
||||
**依赖**:Phase 20-26
|
||||
**优先级**:P0
|
||||
**预估规模**:约 100 行
|
||||
**状态**:⏳ 待实施
|
||||
**状态**:✅ 已交付(2026-07-19)— README 添加 feature 表格 + 快捷组合 + 模块级 features 清单 + 升级指南(LlmProvider 路径迁移);18 个 example 顶部添加 Required features 注释;roadmap 总入口同步
|
||||
|
||||
---
|
||||
|
||||
@@ -305,8 +305,8 @@ graph TD
|
||||
| **M19** | Phase 23 | tools 不含 mcp 编译通过;加 tools-mcp 引入 McpClient | ✅ 2026-07-19(Step 1 合并) |
|
||||
| **M20** | Phase 24 | memory 不含 llm/sqlite 编译通过;加 memory-sqlite 引入 SqliteStore | ✅ 2026-07-19(Step 1 合并) |
|
||||
| **M21** | Phase 25 | agent + engine 全链路门控编译通过 | ✅ 2026-07-19(Step 1 合并) |
|
||||
| **M22** | Phase 26 | 7 种 CI 组合全部编译 + 测试通过;clippy --all-features 0 警告 | ⏳ |
|
||||
| **M23** | Phase 27 | 文档 review 通过 | ⏳ |
|
||||
| **M22** | Phase 26 | 7 种 CI 组合全部编译 + 测试通过;clippy --all-features 0 警告 | ✅ 2026-07-19 |
|
||||
| **M23** | Phase 27 | 文档 review 通过 | ✅ 2026-07-19 |
|
||||
|
||||
## Cargo.toml [features] 草案
|
||||
|
||||
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
# AG Core Roadmap
|
||||
|
||||
> 拆分式 roadmap:按版本归档 + 未归类内容
|
||||
> 最后更新:2026-07-19(v0.3.2 Step 1 完成 — Phase 20-25 cfg 门控全部交付,427 测试通过)
|
||||
> 最后更新:2026-07-19(v0.3.2 Step 3 完成 — Phase 26-27 CI 固化 + 文档更新交付,427 测试通过)
|
||||
|
||||
## 文件索引
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
| [`roadmap-v0.1.0.md`](./roadmap-v0.1.0.md) | v0.1.0 计划与交付 — Phase 0–4c + v0.1.0 Release | ✅ 已发布 2026-07-04 |
|
||||
| [`roadmap-v0.2.0.md`](./roadmap-v0.2.0.md) | v0.2.0 计划与交付 — Phase 5–12 + v0.2.0-rc.1 | 🟡 Phase 5-11 已完成;Phase 12 P2 锦上添花可选 |
|
||||
| [`roadmap-v0.3.0.md`](./roadmap-v0.3.0.md) | v0.3.0 计划与交付 - Phase 13–19 | ✅ Phase 13-19 全部完成,v0.3.0 交付完毕 |
|
||||
| [`roadmap-v0.3.2.md`](./roadmap-v0.3.2.md) | v0.3.2 计划与交付 — Phase 20–27(Cargo features 拆分) | 🟡 Step 1 完成(Phase 20-25);Phase 26-27 待实施 |
|
||||
| [`roadmap-v0.3.2.md`](./roadmap-v0.3.2.md) | v0.3.2 计划与交付 — Phase 20–27(Cargo features 拆分) | ✅ Phase 20-27 全部完成,v0.3.2 交付完毕 |
|
||||
| [`roadmap-unsorted.md`](./roadmap-unsorted.md) | 未归到任何版本的内容 — 全局愿景、当前状态、模块完整性、v0.4+ 展望、风险与建议、下一步行动、阶段总回顾 | — |
|
||||
|
||||
## 阅读建议
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! agent_session_demo —— Agent 装配 + 会话链路 + SessionMemory 桥接。
|
||||
//! Required features: cargo run --example agent_session_demo --features "agent"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. 实现 `Agent` trait(定义角色 + system prompt)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! agent_switch_demo —— Agent 角色热切换示例。
|
||||
//! Required features: cargo run --example agent_switch_demo --features "engine"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. 创建 session(绑定 Analyst agent)
|
||||
@@ -110,6 +111,9 @@ async fn main() {
|
||||
};
|
||||
println!("\n[verify] turn_index = {turn_index}, agent = {agent_name_owned}");
|
||||
assert_eq!(turn_index, 2, "turn_index should be 2 after 2 turns");
|
||||
assert_eq!(agent_name_owned, "reporter", "current agent should be reporter");
|
||||
assert_eq!(
|
||||
agent_name_owned, "reporter",
|
||||
"current agent should be reporter"
|
||||
);
|
||||
println!("✓ context preserved across agent switch");
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! bridge_keys_demo —— bridge_keys 过滤 + 子↔子共享 namespace 示例。
|
||||
//! Required features: cargo run --example bridge_keys_demo --features "engine"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. 父 session 设置 SessionMemory(key: "project_goal", "constraints", "noise")
|
||||
@@ -103,13 +104,24 @@ async fn main() {
|
||||
.dispatch(&parent_id, worker.clone(), "do work", config)
|
||||
.await
|
||||
.expect("dispatch");
|
||||
println!("[3] dispatched sub-agent (child_id={})\n", &result.child_id[..20]);
|
||||
println!(
|
||||
"[3] dispatched sub-agent (child_id={})\n",
|
||||
&result.child_id[..20]
|
||||
);
|
||||
|
||||
// 验证过滤效果
|
||||
let child_session = sm.get(&result.child_id).await.unwrap();
|
||||
let child_guard = child_session.lock().await;
|
||||
let inherited_goal = child_guard.session_memory().get("project_goal").await.unwrap();
|
||||
let inherited_constraint = child_guard.session_memory().get("constraints").await.unwrap();
|
||||
let inherited_goal = child_guard
|
||||
.session_memory()
|
||||
.get("project_goal")
|
||||
.await
|
||||
.unwrap();
|
||||
let inherited_constraint = child_guard
|
||||
.session_memory()
|
||||
.get("constraints")
|
||||
.await
|
||||
.unwrap();
|
||||
let filtered_noise = child_guard.session_memory().get("noise").await.unwrap();
|
||||
drop(child_guard);
|
||||
|
||||
@@ -172,12 +184,7 @@ async fn main() {
|
||||
|
||||
// dispatch 第二个子 agent
|
||||
let _writer_result = sm
|
||||
.dispatch(
|
||||
&parent_id,
|
||||
worker.clone(),
|
||||
"writing task",
|
||||
shared_ns_config,
|
||||
)
|
||||
.dispatch(&parent_id, worker.clone(), "writing task", shared_ns_config)
|
||||
.await
|
||||
.expect("dispatch writer");
|
||||
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! context_slot_demo —— 多上下文槽位管理示例。
|
||||
//! Required features: cargo run --example context_slot_demo --features "agent"
|
||||
//!
|
||||
//! 场景:法律咨询入口 → 派生两个独立探索方向 → 切换 → 隔离验证 → 删除。
|
||||
//!
|
||||
@@ -13,12 +14,12 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use agcore::agent::{Agent, AgentBuilder, AgentSession};
|
||||
use agcore::llm::LlmProvider;
|
||||
use agcore::llm::hooks::HookExecutor;
|
||||
use agcore::llm::mock::MockProvider;
|
||||
use agcore::llm::LlmProvider;
|
||||
use agcore::llm::types::Usage;
|
||||
use agcore::llm::types::message::{ContentBlock, Message};
|
||||
use agcore::llm::types::response_v2::{MessageResponse, StopReason};
|
||||
use agcore::llm::types::Usage;
|
||||
use agcore::tools::ToolRegistry;
|
||||
|
||||
struct LegalAdvisor;
|
||||
@@ -78,11 +79,19 @@ async fn main() {
|
||||
|
||||
println!("\n=== 3. 派生两个独立探索方向的 slot ===");
|
||||
session
|
||||
.derive_slot("option_jurisdiction", "default", agcore::agent::DeriveStrategy::Full)
|
||||
.derive_slot(
|
||||
"option_jurisdiction",
|
||||
"default",
|
||||
agcore::agent::DeriveStrategy::Full,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
session
|
||||
.derive_slot("option_amendment", "default", agcore::agent::DeriveStrategy::Full)
|
||||
.derive_slot(
|
||||
"option_amendment",
|
||||
"default",
|
||||
agcore::agent::DeriveStrategy::Full,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let slots: Vec<_> = session.list_slots().cloned().collect();
|
||||
@@ -158,4 +167,4 @@ fn message_contains(msg: &Message, needle: &str) -> bool {
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! conversation_memory_demo —— 对话记忆滑动窗口与隔离。
|
||||
//! Required features: cargo run --example conversation_memory_demo --features "memory"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. `ConversationMemoryConfig` 构造(SlidingWindow / Full 策略)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! custom_tool —— 自定义工具注册、单次 / 并行调用、权限检查。
|
||||
//! Required features: cargo run --example custom_tool --features "tools,llm"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. 实现 `BaseTool` trait(WeatherTool + DeleteFileTool)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! dispatch_stream_demo —— 流式子代理调度示例。
|
||||
//! Required features: cargo run --example dispatch_stream_demo --features "engine"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. 创建父 session
|
||||
@@ -92,7 +93,11 @@ async fn main() {
|
||||
saw_stream_count += 1;
|
||||
}
|
||||
SubTaskStreamEvent::Completed(r) => {
|
||||
println!(" → Completed(child_id={}, {} tokens)", &r.child_id[..20], r.usage.total().total_tokens);
|
||||
println!(
|
||||
" → Completed(child_id={}, {} tokens)",
|
||||
&r.child_id[..20],
|
||||
r.usage.total().total_tokens
|
||||
);
|
||||
completed = Some(r);
|
||||
break;
|
||||
}
|
||||
@@ -114,7 +119,10 @@ async fn main() {
|
||||
let child_guard = child_session.lock().await;
|
||||
let child_turn_index = child_guard.turn_index();
|
||||
drop(child_guard);
|
||||
assert_eq!(child_turn_index, 1, "turn_index should increment after finalize");
|
||||
assert_eq!(
|
||||
child_turn_index, 1,
|
||||
"turn_index should increment after finalize"
|
||||
);
|
||||
println!("[4] child session turn_index = {child_turn_index} (finalize works)");
|
||||
|
||||
println!("\n✓ dispatch_stream completed successfully");
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! document_demo —— Document + RecursiveCharacterSplitter + MockEmbedding + RagPipeline 完整衔接示例。
|
||||
//! Required features: cargo run --example document_demo --features "memory,tracing-init"
|
||||
//!
|
||||
//! 演示 RAG 管线:
|
||||
//! 1. 创建多段落 Document
|
||||
@@ -37,11 +38,7 @@ async fn main() {
|
||||
let embedder: Arc<dyn Embedding> = Arc::new(MockEmbedding::new(4));
|
||||
let store: Arc<dyn VectorStore> = Arc::new(InMemoryVectorStore::new());
|
||||
let splitter = RecursiveCharacterSplitter::new(200, 30);
|
||||
let pipeline = RagPipeline::new(
|
||||
Arc::clone(&embedder),
|
||||
Arc::clone(&store),
|
||||
Some(splitter),
|
||||
);
|
||||
let pipeline = RagPipeline::new(Arc::clone(&embedder), Arc::clone(&store), Some(splitter));
|
||||
|
||||
// 3. 一次性 ingest:自动 split → embed → add
|
||||
pipeline.ingest(std::slice::from_ref(&doc)).await.unwrap();
|
||||
@@ -67,4 +64,4 @@ async fn main() {
|
||||
);
|
||||
|
||||
println!("\n✓ document_demo 完成");
|
||||
}
|
||||
}
|
||||
|
||||
+203
-65
@@ -1,4 +1,5 @@
|
||||
//! end_to_end —— 3 工具 + 3 轮对话 + SqliteStore 持久化跨连接验证。
|
||||
//! Required features: cargo run --example end_to_end --features "agent,memory-sqlite,provider-openai"
|
||||
//!
|
||||
//! 运行:`cargo run --example end_to_end`(离线,零配置)
|
||||
//!
|
||||
@@ -16,11 +17,15 @@ use std::env;
|
||||
use std::sync::Arc;
|
||||
|
||||
use agcore::agent::{Agent, AgentBuilder, AgentSession};
|
||||
use agcore::llm::LlmProvider;
|
||||
use agcore::llm::hooks::HookExecutor;
|
||||
use agcore::llm::mock::MockProvider;
|
||||
use agcore::llm::LlmProvider;
|
||||
use agcore::llm::provider::{create_provider, ProviderConfig, ProviderType};
|
||||
use agcore::llm::types::{Usage, message::{ContentBlock, Message}, response_v2::{MessageResponse, StopReason}};
|
||||
use agcore::llm::provider::{ProviderConfig, ProviderType, create_provider};
|
||||
use agcore::llm::types::{
|
||||
Usage,
|
||||
message::{ContentBlock, Message},
|
||||
response_v2::{MessageResponse, StopReason},
|
||||
};
|
||||
use agcore::memory::store::{MemoryStore, SqliteStore};
|
||||
use agcore::memory::types::{MemoryFilter, MemoryItem};
|
||||
use agcore::tools::{BaseTool, ToolContext, ToolError, ToolRegistry};
|
||||
@@ -33,8 +38,12 @@ use time::OffsetDateTime;
|
||||
|
||||
struct AssistantAgent;
|
||||
impl Agent for AssistantAgent {
|
||||
fn name(&self) -> &str { "end-to-end assistant" }
|
||||
fn system_prompt(&self) -> Option<&str> { Some("简洁助手,必要时调用工具完成任务。") }
|
||||
fn name(&self) -> &str {
|
||||
"end-to-end assistant"
|
||||
}
|
||||
fn system_prompt(&self) -> Option<&str> {
|
||||
Some("简洁助手,必要时调用工具完成任务。")
|
||||
}
|
||||
}
|
||||
|
||||
// === Tools ===
|
||||
@@ -42,14 +51,19 @@ impl Agent for AssistantAgent {
|
||||
struct EchoTool;
|
||||
#[async_trait]
|
||||
impl BaseTool for EchoTool {
|
||||
fn name(&self) -> &str { "echo" }
|
||||
fn description(&self) -> &str { "回显输入文本" }
|
||||
fn name(&self) -> &str {
|
||||
"echo"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"回显输入文本"
|
||||
}
|
||||
fn parameters(&self) -> Value {
|
||||
json!({"type":"object","properties":{"text":{"type":"string"}},"required":["text"]})
|
||||
}
|
||||
async fn execute(&self, args: Value, _: &ToolContext<'_>) -> Result<Value, ToolError> {
|
||||
let text = args.get("text").and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("text".into(), "需要 string 类型的 text 参数".into()))?;
|
||||
let text = args.get("text").and_then(|v| v.as_str()).ok_or_else(|| {
|
||||
ToolError::InvalidArguments("text".into(), "需要 string 类型的 text 参数".into())
|
||||
})?;
|
||||
Ok(json!({"echoed": format!("收到: {text}")}))
|
||||
}
|
||||
}
|
||||
@@ -58,8 +72,12 @@ impl BaseTool for EchoTool {
|
||||
struct CalcTool;
|
||||
#[async_trait]
|
||||
impl BaseTool for CalcTool {
|
||||
fn name(&self) -> &str { "calc" }
|
||||
fn description(&self) -> &str { "四则运算:'a op b' 格式,op ∈ {+, -, *, /}" }
|
||||
fn name(&self) -> &str {
|
||||
"calc"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"四则运算:'a op b' 格式,op ∈ {+, -, *, /}"
|
||||
}
|
||||
fn parameters(&self) -> Value {
|
||||
json!({"type":"object","properties":{"expr":{"type":"string"}},"required":["expr"]})
|
||||
}
|
||||
@@ -67,18 +85,30 @@ impl BaseTool for CalcTool {
|
||||
let expr = args["expr"].as_str().unwrap_or("");
|
||||
let parts: Vec<&str> = expr.split_whitespace().collect();
|
||||
if parts.len() != 3 {
|
||||
return Err(ToolError::InvalidArguments("expr".into(), "需要 'a op b' 三段式".into()));
|
||||
return Err(ToolError::InvalidArguments(
|
||||
"expr".into(),
|
||||
"需要 'a op b' 三段式".into(),
|
||||
));
|
||||
}
|
||||
let a: i64 = parts[0].parse().map_err(|_| ToolError::InvalidArguments("expr".into(), format!("无法解析 '{}'", parts[0])))?;
|
||||
let b: i64 = parts[2].parse().map_err(|_| ToolError::InvalidArguments("expr".into(), format!("无法解析 '{}'", parts[2])))?;
|
||||
let a: i64 = parts[0].parse().map_err(|_| {
|
||||
ToolError::InvalidArguments("expr".into(), format!("无法解析 '{}'", parts[0]))
|
||||
})?;
|
||||
let b: i64 = parts[2].parse().map_err(|_| {
|
||||
ToolError::InvalidArguments("expr".into(), format!("无法解析 '{}'", parts[2]))
|
||||
})?;
|
||||
let result = match parts[1] {
|
||||
"+" => a + b,
|
||||
"-" => a - b,
|
||||
"*" => a * b,
|
||||
"/" => a.checked_div(b).ok_or_else(|| {
|
||||
ToolError::InvalidArguments("expr".into(), "除数不能为 0".into())
|
||||
})?,
|
||||
op => return Err(ToolError::InvalidArguments("expr".into(), format!("不支持的运算符: {op}"))),
|
||||
"/" => a
|
||||
.checked_div(b)
|
||||
.ok_or_else(|| ToolError::InvalidArguments("expr".into(), "除数不能为 0".into()))?,
|
||||
op => {
|
||||
return Err(ToolError::InvalidArguments(
|
||||
"expr".into(),
|
||||
format!("不支持的运算符: {op}"),
|
||||
));
|
||||
}
|
||||
};
|
||||
Ok(json!({"result": result}))
|
||||
}
|
||||
@@ -87,13 +117,21 @@ impl BaseTool for CalcTool {
|
||||
/// 通过 MemoryStore trait 读写笔记:直接持有 Arc<dyn MemoryStore>,
|
||||
/// 绕开 AgentSession 封装(NoteTool 在 tool.execute 中直接操作 store)。
|
||||
/// 关键前缀 "note:" 用于 list 过滤。
|
||||
struct NoteTool { store: Arc<dyn MemoryStore> }
|
||||
impl NoteTool { const PREFIX: &'static str = "note:"; }
|
||||
struct NoteTool {
|
||||
store: Arc<dyn MemoryStore>,
|
||||
}
|
||||
impl NoteTool {
|
||||
const PREFIX: &'static str = "note:";
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BaseTool for NoteTool {
|
||||
fn name(&self) -> &str { "note" }
|
||||
fn description(&self) -> &str { "笔记 save/query: save(key, content) / query()" }
|
||||
fn name(&self) -> &str {
|
||||
"note"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"笔记 save/query: save(key, content) / query()"
|
||||
}
|
||||
fn parameters(&self) -> Value {
|
||||
json!({
|
||||
"type":"object",
|
||||
@@ -117,18 +155,29 @@ impl BaseTool for NoteTool {
|
||||
metadata: json!({}),
|
||||
created_at: OffsetDateTime::now_utc(),
|
||||
};
|
||||
self.store.save(item).await
|
||||
self.store
|
||||
.save(item)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionFailed("note".into(), e.to_string()))?;
|
||||
Ok(json!({"saved": key}))
|
||||
}
|
||||
"query" => {
|
||||
let filter = MemoryFilter { prefix: Some(Self::PREFIX.into()), ..Default::default() };
|
||||
let items = self.store.list(&filter).await
|
||||
let filter = MemoryFilter {
|
||||
prefix: Some(Self::PREFIX.into()),
|
||||
..Default::default()
|
||||
};
|
||||
let items = self
|
||||
.store
|
||||
.list(&filter)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionFailed("note".into(), e.to_string()))?;
|
||||
let notes: Vec<String> = items.into_iter().map(|i| i.content).collect();
|
||||
Ok(json!({"notes": notes}))
|
||||
}
|
||||
_ => Err(ToolError::InvalidArguments("action".into(), format!("未知 action: {action}"))),
|
||||
_ => Err(ToolError::InvalidArguments(
|
||||
"action".into(),
|
||||
format!("未知 action: {action}"),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -136,31 +185,91 @@ impl BaseTool for NoteTool {
|
||||
// === Mock response helper ===
|
||||
|
||||
fn resp(content: Vec<ContentBlock>, stop: StopReason, u: (u32, u32)) -> MessageResponse {
|
||||
MessageResponse { id: String::new(), model: "mock".into(),
|
||||
MessageResponse {
|
||||
id: String::new(),
|
||||
model: "mock".into(),
|
||||
message: Message::Assistant { content },
|
||||
usage: Usage::from_input_output(u.0, u.1),
|
||||
stop_reason: stop, extra: Default::default() }
|
||||
stop_reason: stop,
|
||||
extra: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn mock_responses() -> Vec<MessageResponse> {
|
||||
vec![
|
||||
// 第 1 轮:calc(25 * 4) → tool_result(100) → 文本回答
|
||||
resp(vec![ContentBlock::ToolUse { id: "t1".into(), name: "calc".into(),
|
||||
input: json!({"expr": "25 * 4"}) }], StopReason::ToolUse, (5, 8)),
|
||||
resp(vec![ContentBlock::Text { text: "25 * 4 = 100".into() }], StopReason::Stop, (8, 12)),
|
||||
resp(
|
||||
vec![ContentBlock::ToolUse {
|
||||
id: "t1".into(),
|
||||
name: "calc".into(),
|
||||
input: json!({"expr": "25 * 4"}),
|
||||
}],
|
||||
StopReason::ToolUse,
|
||||
(5, 8),
|
||||
),
|
||||
resp(
|
||||
vec![ContentBlock::Text {
|
||||
text: "25 * 4 = 100".into(),
|
||||
}],
|
||||
StopReason::Stop,
|
||||
(8, 12),
|
||||
),
|
||||
// 第 2 轮:note(save, last_calc, "100") → tool_result(saved) → 文本回答
|
||||
resp(vec![ContentBlock::ToolUse { id: "t2".into(), name: "note".into(),
|
||||
input: json!({"action": "save", "key": "last_calc", "content": "100"}) }],
|
||||
StopReason::ToolUse, (10, 14)),
|
||||
resp(vec![ContentBlock::Text { text: "已记录:last_calc = 100".into() }], StopReason::Stop, (12, 16)),
|
||||
resp(
|
||||
vec![ContentBlock::ToolUse {
|
||||
id: "t2".into(),
|
||||
name: "note".into(),
|
||||
input: json!({"action": "save", "key": "last_calc", "content": "100"}),
|
||||
}],
|
||||
StopReason::ToolUse,
|
||||
(10, 14),
|
||||
),
|
||||
resp(
|
||||
vec![ContentBlock::Text {
|
||||
text: "已记录:last_calc = 100".into(),
|
||||
}],
|
||||
StopReason::Stop,
|
||||
(12, 16),
|
||||
),
|
||||
// 第 3 轮:note(query) → tool_result([100]) → 文本回答
|
||||
resp(vec![ContentBlock::ToolUse { id: "t3".into(), name: "note".into(),
|
||||
input: json!({"action": "query"}) }], StopReason::ToolUse, (8, 8)),
|
||||
resp(vec![ContentBlock::Text { text: "您刚才的计算结果是 100".into() }], StopReason::Stop, (10, 14)),
|
||||
resp(
|
||||
vec![ContentBlock::ToolUse {
|
||||
id: "t3".into(),
|
||||
name: "note".into(),
|
||||
input: json!({"action": "query"}),
|
||||
}],
|
||||
StopReason::ToolUse,
|
||||
(8, 8),
|
||||
),
|
||||
resp(
|
||||
vec![ContentBlock::Text {
|
||||
text: "您刚才的计算结果是 100".into(),
|
||||
}],
|
||||
StopReason::Stop,
|
||||
(10, 14),
|
||||
),
|
||||
// 后续冗余响应(防止队列耗尽报错)
|
||||
resp(vec![ContentBlock::Text { text: "done".into() }], StopReason::Stop, (1, 1)),
|
||||
resp(vec![ContentBlock::Text { text: "done".into() }], StopReason::Stop, (1, 1)),
|
||||
resp(vec![ContentBlock::Text { text: "done".into() }], StopReason::Stop, (1, 1)),
|
||||
resp(
|
||||
vec![ContentBlock::Text {
|
||||
text: "done".into(),
|
||||
}],
|
||||
StopReason::Stop,
|
||||
(1, 1),
|
||||
),
|
||||
resp(
|
||||
vec![ContentBlock::Text {
|
||||
text: "done".into(),
|
||||
}],
|
||||
StopReason::Stop,
|
||||
(1, 1),
|
||||
),
|
||||
resp(
|
||||
vec![ContentBlock::Text {
|
||||
text: "done".into(),
|
||||
}],
|
||||
StopReason::Stop,
|
||||
(1, 1),
|
||||
),
|
||||
]
|
||||
}
|
||||
|
||||
@@ -169,14 +278,21 @@ fn mock_responses() -> Vec<MessageResponse> {
|
||||
fn select_provider() -> Arc<dyn LlmProvider> {
|
||||
if env::var("AG_LLM_BASE_URL").is_ok() && env::var("AG_LLM_API_KEY").is_ok() {
|
||||
let cfg = ProviderConfig::from_env("AG_LLM").expect("AG_LLM_* 环境变量解析失败");
|
||||
let provider_type = env::var("AG_LLM_PROVIDER").ok()
|
||||
let provider_type = env::var("AG_LLM_PROVIDER")
|
||||
.ok()
|
||||
.and_then(|s| s.parse::<ProviderType>().ok())
|
||||
.unwrap_or(ProviderType::OpenaiChat);
|
||||
Arc::from(create_provider(provider_type, cfg).expect("Provider 创建失败"))
|
||||
} else {
|
||||
let found: Vec<&str> = ["AG_LLM_BASE_URL", "AG_LLM_API_KEY", "AG_LLM_MODEL"]
|
||||
.iter().filter(|k| env::var(k).is_ok()).copied().collect();
|
||||
eprintln!("AG_LLM_* 环境变量不完整(检测到: {:?}),回退到 MockProvider", found);
|
||||
.iter()
|
||||
.filter(|k| env::var(k).is_ok())
|
||||
.copied()
|
||||
.collect();
|
||||
eprintln!(
|
||||
"AG_LLM_* 环境变量不完整(检测到: {:?}),回退到 MockProvider",
|
||||
found
|
||||
);
|
||||
Arc::new(MockProvider::new(mock_responses()))
|
||||
}
|
||||
}
|
||||
@@ -191,52 +307,74 @@ async fn main() {
|
||||
let backend: Arc<dyn MemoryStore> =
|
||||
Arc::new(SqliteStore::open(&db_path).expect("SqliteStore 打开失败"));
|
||||
println!("💾 SqliteStore: {}", db_path.display());
|
||||
let provider_label = if env::var("AG_LLM_BASE_URL").is_ok() && env::var("AG_LLM_API_KEY").is_ok() {
|
||||
"真实 LLM Provider"
|
||||
} else {
|
||||
"MockProvider (离线回退模式)"
|
||||
};
|
||||
let provider_label =
|
||||
if env::var("AG_LLM_BASE_URL").is_ok() && env::var("AG_LLM_API_KEY").is_ok() {
|
||||
"真实 LLM Provider"
|
||||
} else {
|
||||
"MockProvider (离线回退模式)"
|
||||
};
|
||||
println!("🔄 Provider: {provider_label}");
|
||||
|
||||
let mut registry = ToolRegistry::new();
|
||||
registry.register(Arc::new(EchoTool)).unwrap();
|
||||
registry.register(Arc::new(CalcTool)).unwrap();
|
||||
registry.register(Arc::new(NoteTool { store: backend.clone() })).unwrap();
|
||||
registry
|
||||
.register(Arc::new(NoteTool {
|
||||
store: backend.clone(),
|
||||
}))
|
||||
.unwrap();
|
||||
println!("🔧 注册工具: {:?}", registry.list_tools());
|
||||
|
||||
let bundle = Arc::new(AgentBuilder::new()
|
||||
.provider(select_provider())
|
||||
.tool_registry(Arc::new(registry))
|
||||
.hook_executor(Arc::new(HookExecutor::new()))
|
||||
.build().expect("RuntimeBundle 装配失败"));
|
||||
let bundle = Arc::new(
|
||||
AgentBuilder::new()
|
||||
.provider(select_provider())
|
||||
.tool_registry(Arc::new(registry))
|
||||
.hook_executor(Arc::new(HookExecutor::new()))
|
||||
.build()
|
||||
.expect("RuntimeBundle 装配失败"),
|
||||
);
|
||||
|
||||
let mut session = AgentSession::new(Arc::new(AssistantAgent), "e2e-1", bundle.clone());
|
||||
|
||||
println!("\n第 1 轮 用户: 帮我算 25 * 4");
|
||||
let r1 = session.submit_turn("帮我算 25 * 4").await.expect("turn 1 失败");
|
||||
let r1 = session
|
||||
.submit_turn("帮我算 25 * 4")
|
||||
.await
|
||||
.expect("turn 1 失败");
|
||||
println!(" → 回答: {}", r1.text());
|
||||
|
||||
println!("\n第 2 轮 用户: 记下来:结果是 100");
|
||||
let r2 = session.submit_turn("记下来:结果是 100").await.expect("turn 2 失败");
|
||||
let r2 = session
|
||||
.submit_turn("记下来:结果是 100")
|
||||
.await
|
||||
.expect("turn 2 失败");
|
||||
println!(" → 回答: {}", r2.text());
|
||||
|
||||
println!("\n第 3 轮 用户: 我刚才算了什么?");
|
||||
let r3 = session.submit_turn("我刚才算了什么?").await.expect("turn 3 失败");
|
||||
let r3 = session
|
||||
.submit_turn("我刚才算了什么?")
|
||||
.await
|
||||
.expect("turn 3 失败");
|
||||
println!(" → 回答: {}", r3.text());
|
||||
|
||||
let total = session.usage().total();
|
||||
println!("\n📊 用量: prompt={}, completion={}, total={}",
|
||||
total.prompt_tokens, total.completion_tokens, total.total_tokens);
|
||||
println!(
|
||||
"\n📊 用量: prompt={}, completion={}, total={}",
|
||||
total.prompt_tokens, total.completion_tokens, total.total_tokens
|
||||
);
|
||||
|
||||
println!("\n=== 持久化验证 ===");
|
||||
// 显式释放所有对 backend 的 Arc 引用,确保 SqliteStore Connection 真正关闭。
|
||||
// 释放顺序:session → bundle(间接持有 NoteTool → backend clone)→ backend 局部变量。
|
||||
drop(session); // session.bundle Arc 计数 -1
|
||||
drop(bundle); // bundle Arc 计数归零 → registry → NoteTool → backend clone Arc 计数 2→1
|
||||
drop(backend); // backend 局部变量 Arc 计数 1→0 → SqliteStore::drop → Connection 自动 close
|
||||
drop(session); // session.bundle Arc 计数 -1
|
||||
drop(bundle); // bundle Arc 计数归零 → registry → NoteTool → backend clone Arc 计数 2→1
|
||||
drop(backend); // backend 局部变量 Arc 计数 1→0 → SqliteStore::drop → Connection 自动 close
|
||||
let backend2: Arc<dyn MemoryStore> =
|
||||
Arc::new(SqliteStore::open(&db_path).expect("重开 SqliteStore 失败"));
|
||||
let filter = MemoryFilter { prefix: Some("note:".into()), ..Default::default() };
|
||||
let filter = MemoryFilter {
|
||||
prefix: Some("note:".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let items = backend2.list(&filter).await.expect("list 失败");
|
||||
println!("✓ 跨连接数据存活: 找到 {} 条 note", items.len());
|
||||
assert!(!items.is_empty(), "持久化验证失败:重开后无数据");
|
||||
@@ -245,4 +383,4 @@ async fn main() {
|
||||
}
|
||||
|
||||
println!("\n✓ 端到端演示完成");
|
||||
}
|
||||
}
|
||||
|
||||
+12
-5
@@ -1,4 +1,5 @@
|
||||
//! engine_demo —— SessionManager + Checkpointer 端到端示例。
|
||||
//! Required features: cargo run --example engine_demo --features "engine"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. SessionManager::create 创建 session
|
||||
@@ -190,8 +191,8 @@ async fn main() {
|
||||
snapshot_turn, snapshot_data_count
|
||||
);
|
||||
|
||||
let mut rolled_back =
|
||||
AgentSession::from_snapshot(snapshot, agent.clone(), bundle.clone()).expect("from_snapshot");
|
||||
let mut rolled_back = AgentSession::from_snapshot(snapshot, agent.clone(), bundle.clone())
|
||||
.expect("from_snapshot");
|
||||
rolled_back
|
||||
.restore_memory()
|
||||
.await
|
||||
@@ -216,8 +217,14 @@ async fn main() {
|
||||
after_turn,
|
||||
before_turn
|
||||
);
|
||||
assert_eq!(after_turn, snapshot_turn, "rollback 后 turn_index 应等于 checkpoint 时刻值");
|
||||
assert!(after_cost <= before_cost, "rollback 后 cost 应 ≤ rollback 前");
|
||||
assert_eq!(
|
||||
after_turn, snapshot_turn,
|
||||
"rollback 后 turn_index 应等于 checkpoint 时刻值"
|
||||
);
|
||||
assert!(
|
||||
after_cost <= before_cost,
|
||||
"rollback 后 cost 应 ≤ rollback 前"
|
||||
);
|
||||
println!("✓ rollback + replace 一致性验证通过");
|
||||
|
||||
// 9. destroy 父子 session
|
||||
@@ -249,4 +256,4 @@ async fn main() {
|
||||
sm.destroy(&c_id).await.unwrap();
|
||||
|
||||
println!("\n✓ engine_demo 完成");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! knowledge_graph_demo -- 知识图谱 + 双通道检索演示。
|
||||
//! Required features: cargo run --example knowledge_graph_demo --features "memory"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. 构建 KnowledgeGraph(实体 + 关系)
|
||||
@@ -52,15 +53,30 @@ async fn main() {
|
||||
graph.add_entity(e.clone()).await.unwrap();
|
||||
}
|
||||
graph
|
||||
.add_relation(GraphRelation::new("langchain", "langgraph", "includes", 0.9))
|
||||
.add_relation(GraphRelation::new(
|
||||
"langchain",
|
||||
"langgraph",
|
||||
"includes",
|
||||
0.9,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
graph
|
||||
.add_relation(GraphRelation::new("langchain", "langsmith", "includes", 0.7))
|
||||
.add_relation(GraphRelation::new(
|
||||
"langchain",
|
||||
"langsmith",
|
||||
"includes",
|
||||
0.7,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
graph
|
||||
.add_relation(GraphRelation::new("langchain", "python", "built_with", 0.95))
|
||||
.add_relation(GraphRelation::new(
|
||||
"langchain",
|
||||
"python",
|
||||
"built_with",
|
||||
0.95,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
graph
|
||||
@@ -86,7 +102,10 @@ async fn main() {
|
||||
// ── 3. 标签管理 ──
|
||||
println!("\n=== 3. 标签管理 ===");
|
||||
graph
|
||||
.set_entity_tags("langchain", vec!["ai".into(), "framework".into(), "llm".into()])
|
||||
.set_entity_tags(
|
||||
"langchain",
|
||||
vec!["ai".into(), "framework".into(), "llm".into()],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
graph
|
||||
@@ -118,8 +137,8 @@ async fn main() {
|
||||
.unwrap();
|
||||
|
||||
// Hybrid 策略(默认)
|
||||
let retriever = MemoryRetriever::new(ks, RetrieverConfig::default())
|
||||
.with_knowledge_graph(graph.clone());
|
||||
let retriever =
|
||||
MemoryRetriever::new(ks, RetrieverConfig::default()).with_knowledge_graph(graph.clone());
|
||||
println!("\n--- Hybrid 检索: 'langchain' ---");
|
||||
let result = retriever.retrieve("langchain").await.unwrap();
|
||||
println!("策略: {:?}", result.strategy);
|
||||
@@ -129,7 +148,10 @@ async fn main() {
|
||||
println!(" [Store] {} (score={:.3})", page.title, score);
|
||||
}
|
||||
RetrievalItem::GraphEntity {
|
||||
entity, score, path, ..
|
||||
entity,
|
||||
score,
|
||||
path,
|
||||
..
|
||||
} => {
|
||||
println!(
|
||||
" [Graph] {} (score={:.3}, path={:?})",
|
||||
@@ -172,7 +194,10 @@ async fn main() {
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
result.items.iter().all(|i| matches!(i, RetrievalItem::GraphEntity { .. })),
|
||||
result
|
||||
.items
|
||||
.iter()
|
||||
.all(|i| matches!(i, RetrievalItem::GraphEntity { .. })),
|
||||
"GraphOnly 应只返回 Graph 结果"
|
||||
);
|
||||
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! knowledge_search_demo —— 知识页面存储与关键词检索。
|
||||
//! Required features: cargo run --example knowledge_search_demo --features "memory"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. `KnowledgeStore` 存储多个 `KnowledgePage`
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! prompt_composer —— 提示词模板与组合器离线示例。
|
||||
//! Required features: cargo run --example prompt_composer --features "prompt,llm"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. `PromptTemplate::compile` + `render` 变量插值(`{{var}}` 语法)
|
||||
|
||||
+55
-18
@@ -1,41 +1,61 @@
|
||||
//! quick_start —— 30 行最小可运行示例,展示 Agent / BaseTool / Builder / Session 四层抽象。
|
||||
//! Required features: cargo run --example quick_start --features "agent"
|
||||
//!
|
||||
//! 运行:`cargo run --example quick_start`(离线,零配置)
|
||||
|
||||
use std::sync::Arc;
|
||||
use agcore::agent::{Agent, AgentBuilder, AgentSession};
|
||||
use agcore::llm::LlmProvider;
|
||||
use agcore::llm::hooks::HookExecutor;
|
||||
use agcore::llm::mock::MockProvider;
|
||||
use agcore::llm::LlmProvider;
|
||||
use agcore::llm::types::{Usage, message::{ContentBlock, Message}, response_v2::{MessageResponse, StopReason}};
|
||||
use agcore::llm::types::{
|
||||
Usage,
|
||||
message::{ContentBlock, Message},
|
||||
response_v2::{MessageResponse, StopReason},
|
||||
};
|
||||
use agcore::tools::{BaseTool, ToolContext, ToolError, ToolRegistry};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Arc;
|
||||
|
||||
struct Greeter;
|
||||
impl Agent for Greeter {
|
||||
fn name(&self) -> &str { "greeter" }
|
||||
fn system_prompt(&self) -> Option<&str> { Some("中文助手,先调用 echo 工具,再总结。") }
|
||||
fn name(&self) -> &str {
|
||||
"greeter"
|
||||
}
|
||||
fn system_prompt(&self) -> Option<&str> {
|
||||
Some("中文助手,先调用 echo 工具,再总结。")
|
||||
}
|
||||
}
|
||||
|
||||
struct EchoTool;
|
||||
#[async_trait]
|
||||
impl BaseTool for EchoTool {
|
||||
fn name(&self) -> &str { "echo" }
|
||||
fn description(&self) -> &str { "回显文本" }
|
||||
fn name(&self) -> &str {
|
||||
"echo"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"回显文本"
|
||||
}
|
||||
fn parameters(&self) -> Value {
|
||||
json!({"type":"object","properties":{"text":{"type":"string"}},"required":["text"]})
|
||||
}
|
||||
async fn execute(&self, args: Value, _: &ToolContext<'_>) -> Result<Value, ToolError> {
|
||||
let text = args.get("text").and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("text".into(), "需要 string 类型的 text 参数".into()))?;
|
||||
let text = args.get("text").and_then(|v| v.as_str()).ok_or_else(|| {
|
||||
ToolError::InvalidArguments("text".into(), "需要 string 类型的 text 参数".into())
|
||||
})?;
|
||||
Ok(json!({"echoed": format!("收到: {text}")}))
|
||||
}
|
||||
}
|
||||
|
||||
fn resp(content: Vec<ContentBlock>, stop: StopReason, u: (u32, u32)) -> MessageResponse {
|
||||
MessageResponse { id: String::new(), model: "mock".into(), message: Message::Assistant { content },
|
||||
usage: Usage::from_input_output(u.0, u.1), stop_reason: stop, extra: Default::default() }
|
||||
MessageResponse {
|
||||
id: String::new(),
|
||||
model: "mock".into(),
|
||||
message: Message::Assistant { content },
|
||||
usage: Usage::from_input_output(u.0, u.1),
|
||||
stop_reason: stop,
|
||||
extra: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
@@ -43,14 +63,31 @@ async fn main() {
|
||||
let mut registry = ToolRegistry::new();
|
||||
registry.register(Arc::new(EchoTool)).unwrap();
|
||||
let provider: Arc<dyn LlmProvider> = Arc::new(MockProvider::new(vec![
|
||||
resp(vec![ContentBlock::ToolUse { id: "c1".into(), name: "echo".into(),
|
||||
input: json!({"text": "你好"}) }], StopReason::ToolUse, (5, 8)),
|
||||
resp(vec![ContentBlock::Text { text: "EchoTool 已收到您的消息并完成回传。".into() }],
|
||||
StopReason::Stop, (8, 16)),
|
||||
resp(
|
||||
vec![ContentBlock::ToolUse {
|
||||
id: "c1".into(),
|
||||
name: "echo".into(),
|
||||
input: json!({"text": "你好"}),
|
||||
}],
|
||||
StopReason::ToolUse,
|
||||
(5, 8),
|
||||
),
|
||||
resp(
|
||||
vec![ContentBlock::Text {
|
||||
text: "EchoTool 已收到您的消息并完成回传。".into(),
|
||||
}],
|
||||
StopReason::Stop,
|
||||
(8, 16),
|
||||
),
|
||||
]));
|
||||
let bundle = Arc::new(AgentBuilder::new()
|
||||
.provider(provider).tool_registry(Arc::new(registry))
|
||||
.hook_executor(Arc::new(HookExecutor::new())).build().unwrap());
|
||||
let bundle = Arc::new(
|
||||
AgentBuilder::new()
|
||||
.provider(provider)
|
||||
.tool_registry(Arc::new(registry))
|
||||
.hook_executor(Arc::new(HookExecutor::new()))
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
let mut session = AgentSession::new(Arc::new(Greeter), "qs", bundle);
|
||||
let resp = session.submit_turn("你好").await.unwrap();
|
||||
let text = resp.text();
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
//! Required features: cargo run --example simple_visit --features "llm,provider-openai,tracing-init"
|
||||
|
||||
use std::env;
|
||||
|
||||
use agcore::init_tracing;
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! streaming_events_demo —— LLM 流式响应事件流消费(含错误路径)。
|
||||
//! Required features: cargo run --example streaming_events_demo --features "llm,provider-openai"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. `MockProvider::chat_stream` 输出标准 `StreamEvent` 流(离线可跑)
|
||||
@@ -14,9 +15,9 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use agcore::llm::LlmProvider;
|
||||
use agcore::llm::cycle::{CycleConfig, LlmCycle};
|
||||
use agcore::llm::mock::MockProvider;
|
||||
use agcore::llm::LlmProvider;
|
||||
use agcore::llm::types::Usage;
|
||||
use agcore::llm::types::message::{ContentBlock, Message};
|
||||
use agcore::llm::types::response_v2::{MessageResponse, StopReason, StreamEvent};
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! sub_agent_dispatch_demo —— SubAgent 并行派发示例。
|
||||
//! Required features: cargo run --example sub_agent_dispatch_demo --features "engine"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. 创建父 session("主编" agent)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
//! task_agent_demo —— Plan 解析、Step 状态机、错误路径。
|
||||
//! Required features: cargo run --example task_agent_demo --features "agent"
|
||||
//!
|
||||
//! 演示:
|
||||
//! 1. `JsonPlanParser::parse` 解析合法 JSON 输入
|
||||
@@ -12,9 +13,9 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use agcore::agent::{AgentError, JsonPlanParser, PlanParser, Step, StepStatus};
|
||||
use agcore::llm::types::Usage;
|
||||
use agcore::llm::types::message::Message;
|
||||
use agcore::llm::types::response_v2::{MessageResponse, StopReason};
|
||||
use agcore::llm::types::Usage;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
@@ -111,4 +112,4 @@ async fn main() {
|
||||
assert!(matches!(err, AgentError::PlanParse(_)));
|
||||
|
||||
println!("\n✓ task_agent_demo 完成");
|
||||
}
|
||||
}
|
||||
|
||||
+2
-2
@@ -23,8 +23,8 @@ pub mod task;
|
||||
pub use agent::Agent;
|
||||
pub use builder::AgentBuilder;
|
||||
pub use context::{
|
||||
ContextBudget, ContextSlot, DeriveStrategy, FocusedConfig, MergeStrategy, SlotConfig,
|
||||
SlotMeta, SlotMode, SlotSource,
|
||||
ContextBudget, ContextSlot, DeriveStrategy, FocusedConfig, MergeStrategy, SlotConfig, SlotMeta,
|
||||
SlotMode, SlotSource,
|
||||
};
|
||||
pub use error::AgentError;
|
||||
pub use runtime::{AgentConfig, RuntimeBundle};
|
||||
|
||||
@@ -12,8 +12,8 @@ use std::sync::Arc;
|
||||
use crate::agent::error::AgentError;
|
||||
use crate::agent::runtime::{AgentConfig, RuntimeBundle};
|
||||
use crate::agent::summary::SummaryConfig;
|
||||
use crate::llm::hooks::HookExecutor;
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::llm::hooks::HookExecutor;
|
||||
use crate::memory::retriever::MemoryRetriever;
|
||||
use crate::memory::store::MemoryStore;
|
||||
use crate::tools::ToolRegistry;
|
||||
|
||||
+136
-44
@@ -250,12 +250,12 @@ impl ContextSlot {
|
||||
|
||||
/// 保存 slot 数据到存储后端(全量写入,含 config)。
|
||||
pub async fn save(&self, store: &dyn MemoryStore) -> Result<(), AgentError> {
|
||||
let data = serde_json::to_string(&self.messages)
|
||||
.map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
let meta = serde_json::to_string(&self.meta)
|
||||
.map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
let config = serde_json::to_string(&self.config)
|
||||
.map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
let data =
|
||||
serde_json::to_string(&self.messages).map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
let meta =
|
||||
serde_json::to_string(&self.meta).map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
let config =
|
||||
serde_json::to_string(&self.config).map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
|
||||
store
|
||||
.save(Self::make_item(
|
||||
@@ -551,8 +551,10 @@ mod tests {
|
||||
async fn slot_save_load_roundtrip() {
|
||||
let store = make_store();
|
||||
let mut slot = make_slot("default", "s1");
|
||||
slot.append_messages(vec![Message::user_text("hi")]).unwrap();
|
||||
slot.append_messages(vec![Message::assistant("hello")]).unwrap();
|
||||
slot.append_messages(vec![Message::user_text("hi")])
|
||||
.unwrap();
|
||||
slot.append_messages(vec![Message::assistant("hello")])
|
||||
.unwrap();
|
||||
|
||||
slot.save(&*store).await.unwrap();
|
||||
let loaded = ContextSlot::load("default", "s1", &*store).await.unwrap();
|
||||
@@ -576,8 +578,14 @@ mod tests {
|
||||
.unwrap();
|
||||
b.save(&*store).await.unwrap();
|
||||
|
||||
let loaded_a = ContextSlot::load("main", "sA", &*store).await.unwrap().unwrap();
|
||||
let loaded_b = ContextSlot::load("main", "sB", &*store).await.unwrap().unwrap();
|
||||
let loaded_a = ContextSlot::load("main", "sA", &*store)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let loaded_b = ContextSlot::load("main", "sB", &*store)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(extract_text(&loaded_a.messages[0]), "only in A");
|
||||
assert_eq!(extract_text(&loaded_b.messages[0]), "only in B");
|
||||
}
|
||||
@@ -622,10 +630,13 @@ mod tests {
|
||||
async fn slot_delete_then_load_none() {
|
||||
let store = make_store();
|
||||
let mut slot = make_slot("to_delete", "s1");
|
||||
slot.append_messages(vec![Message::user_text("hi")]).unwrap();
|
||||
slot.append_messages(vec![Message::user_text("hi")])
|
||||
.unwrap();
|
||||
slot.save(&*store).await.unwrap();
|
||||
|
||||
ContextSlot::delete("to_delete", "s1", &*store).await.unwrap();
|
||||
ContextSlot::delete("to_delete", "s1", &*store)
|
||||
.await
|
||||
.unwrap();
|
||||
let loaded = ContextSlot::load("to_delete", "s1", &*store).await.unwrap();
|
||||
assert!(loaded.is_none());
|
||||
}
|
||||
@@ -723,7 +734,10 @@ mod tests {
|
||||
let store = make_store();
|
||||
let slot = make_slot("empty", "s1");
|
||||
slot.save(&*store).await.unwrap();
|
||||
let loaded = ContextSlot::load("empty", "s1", &*store).await.unwrap().unwrap();
|
||||
let loaded = ContextSlot::load("empty", "s1", &*store)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(loaded.messages.is_empty());
|
||||
assert_eq!(loaded.meta.message_count, 0);
|
||||
}
|
||||
@@ -745,7 +759,8 @@ mod tests {
|
||||
async fn derive_full_copies_parent_messages() {
|
||||
let mut parent = make_slot("p", "s1");
|
||||
for i in 0..3 {
|
||||
parent.append_messages(vec![Message::user_text(format!("u{i}"))])
|
||||
parent
|
||||
.append_messages(vec![Message::user_text(format!("u{i}"))])
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
@@ -769,7 +784,10 @@ mod tests {
|
||||
let store = make_store();
|
||||
child.save(&*store).await.unwrap();
|
||||
|
||||
let loaded = ContextSlot::load("c", "s1", &*store).await.unwrap().unwrap();
|
||||
let loaded = ContextSlot::load("c", "s1", &*store)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(loaded.messages.len(), 3);
|
||||
assert!(matches!(loaded.config.source, SlotSource::Derived { .. }));
|
||||
}
|
||||
@@ -777,9 +795,12 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn derive_focused_filters_parent_messages() {
|
||||
let mut parent = make_slot("p", "s1");
|
||||
parent.append_messages(vec![Message::system("sys")]).unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::system("sys")])
|
||||
.unwrap();
|
||||
for i in 0..5 {
|
||||
parent.append_messages(vec![Message::user_text(format!("u{i}"))])
|
||||
parent
|
||||
.append_messages(vec![Message::user_text(format!("u{i}"))])
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
@@ -818,7 +839,9 @@ mod tests {
|
||||
async fn derived_slot_loadable_independently() {
|
||||
let store = make_store();
|
||||
let mut parent = make_slot("p", "s1");
|
||||
parent.append_messages(vec![Message::user_text("u")]).unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::user_text("u")])
|
||||
.unwrap();
|
||||
parent.save(&*store).await.unwrap();
|
||||
|
||||
// 派生 child
|
||||
@@ -835,12 +858,16 @@ mod tests {
|
||||
compact: true,
|
||||
},
|
||||
);
|
||||
child.append_messages(vec![Message::user_text("derived msg")])
|
||||
child
|
||||
.append_messages(vec![Message::user_text("derived msg")])
|
||||
.unwrap();
|
||||
child.save(&*store).await.unwrap();
|
||||
|
||||
// child 可独立加载
|
||||
let loaded = ContextSlot::load("c", "s1", &*store).await.unwrap().unwrap();
|
||||
let loaded = ContextSlot::load("c", "s1", &*store)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(loaded.messages.len(), 1);
|
||||
assert_eq!(extract_text(&loaded.messages[0]), "derived msg");
|
||||
}
|
||||
@@ -867,17 +894,59 @@ mod tests {
|
||||
slot.save(&*store).await.unwrap();
|
||||
|
||||
// 确认所有记录存在
|
||||
assert!(store.get(&ContextSlot::data_key("s1", "x")).await.unwrap().is_some());
|
||||
assert!(store.get(&ContextSlot::meta_key("s1", "x")).await.unwrap().is_some());
|
||||
assert!(store.get(&ContextSlot::config_key("s1", "x")).await.unwrap().is_some());
|
||||
assert!(store.get(&ContextSlot::rel_key("s1", "x")).await.unwrap().is_some());
|
||||
assert!(
|
||||
store
|
||||
.get(&ContextSlot::data_key("s1", "x"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.get(&ContextSlot::meta_key("s1", "x"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.get(&ContextSlot::config_key("s1", "x"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.get(&ContextSlot::rel_key("s1", "x"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
|
||||
ContextSlot::delete("x", "s1", &*store).await.unwrap();
|
||||
|
||||
// data/meta/config 已删
|
||||
assert!(store.get(&ContextSlot::data_key("s1", "x")).await.unwrap().is_none());
|
||||
assert!(store.get(&ContextSlot::meta_key("s1", "x")).await.unwrap().is_none());
|
||||
assert!(store.get(&ContextSlot::config_key("s1", "x")).await.unwrap().is_none());
|
||||
assert!(
|
||||
store
|
||||
.get(&ContextSlot::data_key("s1", "x"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.get(&ContextSlot::meta_key("s1", "x"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.get(&ContextSlot::config_key("s1", "x"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
// ===== 基础类型测试 =====
|
||||
@@ -893,7 +962,10 @@ mod tests {
|
||||
#[test]
|
||||
fn context_budget_default_sum_128k() {
|
||||
let b = ContextBudget::default();
|
||||
assert_eq!(b.system + b.history + b.tools + b.tool_results + b.reserve, 128_000);
|
||||
assert_eq!(
|
||||
b.system + b.history + b.tools + b.tool_results + b.reserve,
|
||||
128_000
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -939,10 +1011,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn filter_focused_injects_summary() {
|
||||
let messages = vec![
|
||||
Message::user_text("u"),
|
||||
Message::assistant("a"),
|
||||
];
|
||||
let messages = vec![Message::user_text("u"), Message::assistant("a")];
|
||||
let cfg = FocusedConfig {
|
||||
keep_system: false,
|
||||
recent_messages: 100,
|
||||
@@ -964,8 +1033,12 @@ mod tests {
|
||||
#[test]
|
||||
fn fork_full_copies_messages() {
|
||||
let mut parent = make_slot("p", "s1");
|
||||
parent.append_messages(vec![Message::user_text("a")]).unwrap();
|
||||
parent.append_messages(vec![Message::assistant("b")]).unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::user_text("a")])
|
||||
.unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::assistant("b")])
|
||||
.unwrap();
|
||||
let child = parent.fork("c".into(), DeriveStrategy::Full);
|
||||
assert_eq!(child.messages.len(), 2);
|
||||
assert!(matches!(child.config.mode, SlotMode::Full));
|
||||
@@ -974,7 +1047,9 @@ mod tests {
|
||||
#[test]
|
||||
fn fork_focused_filters_messages() {
|
||||
let mut parent = make_slot("p", "s1");
|
||||
parent.append_messages(vec![Message::system("sys")]).unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::system("sys")])
|
||||
.unwrap();
|
||||
for i in 0..5 {
|
||||
parent
|
||||
.append_messages(vec![Message::user_text(format!("u{i}"))])
|
||||
@@ -994,7 +1069,9 @@ mod tests {
|
||||
#[test]
|
||||
fn fork_preserves_independence() {
|
||||
let mut parent = make_slot("p", "s1");
|
||||
parent.append_messages(vec![Message::user_text("a")]).unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::user_text("a")])
|
||||
.unwrap();
|
||||
let mut child = parent.fork("c".into(), DeriveStrategy::Full);
|
||||
let child_count_at_fork = child.messages.len();
|
||||
|
||||
@@ -1003,7 +1080,9 @@ mod tests {
|
||||
.append_messages(vec![Message::user_text("b")])
|
||||
.unwrap();
|
||||
// 子 slot 追加
|
||||
child.append_messages(vec![Message::user_text("c")]).unwrap();
|
||||
child
|
||||
.append_messages(vec![Message::user_text("c")])
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(parent.messages.len(), 2);
|
||||
assert_eq!(child.messages.len(), child_count_at_fork + 1);
|
||||
@@ -1013,10 +1092,15 @@ mod tests {
|
||||
#[test]
|
||||
fn fork_sets_derived_source() {
|
||||
let mut parent = make_slot("p", "s1");
|
||||
parent.append_messages(vec![Message::user_text("a")]).unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::user_text("a")])
|
||||
.unwrap();
|
||||
let child = parent.fork("c".into(), DeriveStrategy::Full);
|
||||
match &child.config.source {
|
||||
SlotSource::Derived { parent_id, strategy } => {
|
||||
SlotSource::Derived {
|
||||
parent_id,
|
||||
strategy,
|
||||
} => {
|
||||
assert_eq!(parent_id, "p");
|
||||
assert!(matches!(strategy, DeriveStrategy::Full));
|
||||
}
|
||||
@@ -1030,7 +1114,9 @@ mod tests {
|
||||
#[test]
|
||||
fn merge_append_appends_messages() {
|
||||
let mut parent = make_slot("p", "s1");
|
||||
parent.append_messages(vec![Message::user_text("p1")]).unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::user_text("p1")])
|
||||
.unwrap();
|
||||
let child = {
|
||||
let mut c = parent.fork("c".into(), DeriveStrategy::Full);
|
||||
// fork 时 child 继承父的 "p1";再追加一条 c1
|
||||
@@ -1049,8 +1135,12 @@ mod tests {
|
||||
#[test]
|
||||
fn merge_replace_replaces_messages() {
|
||||
let mut parent = make_slot("p", "s1");
|
||||
parent.append_messages(vec![Message::user_text("p1")]).unwrap();
|
||||
parent.append_messages(vec![Message::user_text("p2")]).unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::user_text("p1")])
|
||||
.unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::user_text("p2")])
|
||||
.unwrap();
|
||||
let child = {
|
||||
let mut c = parent.fork("c".into(), DeriveStrategy::Full);
|
||||
// 清空 child 再追加
|
||||
@@ -1088,7 +1178,9 @@ mod tests {
|
||||
#[test]
|
||||
fn merge_cross_session_rejected() {
|
||||
let mut parent = make_slot("p", "s1");
|
||||
parent.append_messages(vec![Message::user_text("a")]).unwrap();
|
||||
parent
|
||||
.append_messages(vec![Message::user_text("a")])
|
||||
.unwrap();
|
||||
let child = ContextSlot::new("OTHER_SESSION", "c", SlotConfig::default());
|
||||
let err = parent.merge(child, MergeStrategy::Append).unwrap_err();
|
||||
assert!(matches!(err, AgentError::Config(_)));
|
||||
@@ -1121,4 +1213,4 @@ mod tests {
|
||||
.unwrap()
|
||||
.block_on(f)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,9 +16,9 @@ use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::agent::summary::SummaryConfig;
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::llm::compact::CompactConfig;
|
||||
use crate::llm::hooks::HookExecutor;
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::memory::retriever::MemoryRetriever;
|
||||
use crate::memory::store::MemoryStore;
|
||||
use crate::tools::ToolRegistry;
|
||||
|
||||
+74
-65
@@ -15,24 +15,22 @@ use std::sync::Arc;
|
||||
use futures_core::Stream;
|
||||
|
||||
use crate::agent::agent::Agent;
|
||||
use crate::agent::context::{
|
||||
ContextSlot, DeriveStrategy, SlotConfig, SlotMode,
|
||||
};
|
||||
use crate::agent::context::{ContextSlot, DeriveStrategy, SlotConfig, SlotMode};
|
||||
// SlotSource 仅在 `mod tests` 中使用(通过 `use super::*;` 引入),lib 主体保留以避免测试 import 变更。
|
||||
#[allow(unused_imports)]
|
||||
use crate::agent::context::SlotSource;
|
||||
use crate::agent::error::AgentError;
|
||||
use crate::agent::runtime::RuntimeBundle;
|
||||
use crate::agent::session_memory::SessionMemory;
|
||||
use crate::agent::summary::{format_messages_as_text, SummaryConfig};
|
||||
#[cfg(feature = "engine")]
|
||||
use crate::engine::snapshot::{SessionMemoryEntry, SessionSnapshot};
|
||||
use crate::agent::summary::{SummaryConfig, format_messages_as_text};
|
||||
#[cfg(feature = "engine")]
|
||||
use crate::engine::EngineError;
|
||||
#[cfg(feature = "engine")]
|
||||
use crate::engine::snapshot::{SessionMemoryEntry, SessionSnapshot};
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::llm::cycle::{CostTracker, CycleConfig, LlmCycle};
|
||||
use crate::llm::error::LlmError;
|
||||
use crate::llm::hooks::{HookContext, HookEvent};
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::llm::stream::StreamEvent;
|
||||
use crate::llm::types::message::Message;
|
||||
use crate::llm::types::response_v2::MessageResponse;
|
||||
@@ -112,11 +110,7 @@ impl AgentSession {
|
||||
let session_memory = SessionMemory::new(backend, &session_id_str);
|
||||
|
||||
// 自动创建 "default" slot
|
||||
let default_slot = ContextSlot::new(
|
||||
&session_id_str,
|
||||
"default",
|
||||
SlotConfig::default(),
|
||||
);
|
||||
let default_slot = ContextSlot::new(&session_id_str, "default", SlotConfig::default());
|
||||
let mut slots = HashMap::new();
|
||||
slots.insert("default".to_string(), default_slot);
|
||||
|
||||
@@ -162,9 +156,7 @@ impl AgentSession {
|
||||
key: impl Into<String>,
|
||||
value: impl Into<String>,
|
||||
) -> Result<(), AgentError> {
|
||||
self.session_memory
|
||||
.set(&key.into(), &value.into())
|
||||
.await
|
||||
self.session_memory.set(&key.into(), &value.into()).await
|
||||
}
|
||||
|
||||
/// 读取一条会话级数据。
|
||||
@@ -204,11 +196,7 @@ impl AgentSession {
|
||||
if self.slots.contains_key(&id) {
|
||||
return Err(AgentError::SlotAlreadyExists(id));
|
||||
}
|
||||
let slot = ContextSlot::new(
|
||||
&self.session_id,
|
||||
&id,
|
||||
config.unwrap_or_default(),
|
||||
);
|
||||
let slot = ContextSlot::new(&self.session_id, &id, config.unwrap_or_default());
|
||||
slot.save(&*self.resolve_store()).await?;
|
||||
self.slots.insert(id, slot);
|
||||
Ok(())
|
||||
@@ -264,7 +252,9 @@ impl AgentSession {
|
||||
/// - 已删除后再 load 返回 None
|
||||
pub async fn delete_slot(&mut self, id: &str) -> Result<(), AgentError> {
|
||||
if id == "default" {
|
||||
return Err(AgentError::Config("Cannot delete the 'default' slot".into()));
|
||||
return Err(AgentError::Config(
|
||||
"Cannot delete the 'default' slot".into(),
|
||||
));
|
||||
}
|
||||
if self.slots.len() <= 1 {
|
||||
return Err(AgentError::Config("Cannot delete the last slot".into()));
|
||||
@@ -354,12 +344,7 @@ impl AgentSession {
|
||||
// 6. 只将本轮新增消息追加到当前 slot(保留全量历史,确保 Focused 模式的"读时过滤"语义不丢失数据)
|
||||
// cycle.messages() 包含 [system_prompt?, history..., user_input, tool_calls..., final_response]
|
||||
// 新增消息 = cycle.messages()[input_len..](跳过 initial_messages,即跳过已被持久化的内容)
|
||||
let new_messages: Vec<Message> = cycle
|
||||
.messages()
|
||||
.iter()
|
||||
.skip(input_len)
|
||||
.cloned()
|
||||
.collect();
|
||||
let new_messages: Vec<Message> = cycle.messages().iter().skip(input_len).cloned().collect();
|
||||
let store = self.resolve_store();
|
||||
if let Some(slot) = self.slots.get_mut(&self.current_slot_id) {
|
||||
slot.append_messages(new_messages)?;
|
||||
@@ -437,10 +422,7 @@ impl AgentSession {
|
||||
|
||||
// 5. 调用流式工具循环
|
||||
let stream = cycle
|
||||
.submit_with_tools_stream(
|
||||
user_input.into(),
|
||||
Arc::clone(&self.bundle.tool_registry),
|
||||
)
|
||||
.submit_with_tools_stream(user_input.into(), Arc::clone(&self.bundle.tool_registry))
|
||||
.await?;
|
||||
|
||||
// 6. turn_index 递增 —— 配合 finalize_turn 用 (turn_index - 1) 传递正确的 OnTurnEnd 序号
|
||||
@@ -492,7 +474,8 @@ impl AgentSession {
|
||||
.await;
|
||||
|
||||
// Phase 16: 摘要检查点(流式路径 turn_index 已被 submit_turn_stream 提前 ++1)
|
||||
self.maybe_summarize(self.turn_index.saturating_sub(1)).await;
|
||||
self.maybe_summarize(self.turn_index.saturating_sub(1))
|
||||
.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -569,11 +552,7 @@ impl AgentSession {
|
||||
if slots.is_empty() {
|
||||
slots.insert(
|
||||
"default".to_string(),
|
||||
ContextSlot::new(
|
||||
&snapshot.session_id,
|
||||
"default",
|
||||
SlotConfig::default(),
|
||||
),
|
||||
ContextSlot::new(&snapshot.session_id, "default", SlotConfig::default()),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -619,12 +598,7 @@ impl AgentSession {
|
||||
|
||||
for (key, entry) in entries {
|
||||
self.session_memory
|
||||
.set_with_meta(
|
||||
&key,
|
||||
&entry.value,
|
||||
entry.metadata.clone(),
|
||||
entry.created_at,
|
||||
)
|
||||
.set_with_meta(&key, &entry.value, entry.metadata.clone(), entry.created_at)
|
||||
.await
|
||||
.map_err(EngineError::Agent)?;
|
||||
}
|
||||
@@ -686,13 +660,22 @@ impl AgentSession {
|
||||
let model = cfg.summary_model.clone();
|
||||
let prompt = cfg.summary_prompt.clone();
|
||||
|
||||
let result =
|
||||
Self::generate_summary(&provider, &messages, &prompt, model.as_deref(), max_tool_result_chars)
|
||||
.await;
|
||||
let result = Self::generate_summary(
|
||||
&provider,
|
||||
&messages,
|
||||
&prompt,
|
||||
model.as_deref(),
|
||||
max_tool_result_chars,
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(text) => {
|
||||
tracing::info!(turn = current_turn, summary_len = text.len(), "摘要自动生成成功");
|
||||
tracing::info!(
|
||||
turn = current_turn,
|
||||
summary_len = text.len(),
|
||||
"摘要自动生成成功"
|
||||
);
|
||||
// Resolve store first (immutable borrow on self) before mutable borrow on slots.
|
||||
let store = self.resolve_store();
|
||||
if let Some(slot) = self.slots.get_mut(&self.current_slot_id)
|
||||
@@ -755,8 +738,8 @@ impl AgentSession {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agent::builder::AgentBuilder;
|
||||
use crate::agent::FocusedConfig;
|
||||
use crate::agent::builder::AgentBuilder;
|
||||
use crate::llm::hooks::{Hook, HookContext, HookExecutor, HookResult};
|
||||
use crate::llm::mock::MockProvider;
|
||||
use crate::llm::stream::StreamEvent;
|
||||
@@ -815,7 +798,9 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
fn build_session(provider_responses: Vec<MessageResponse>) -> (AgentSession, Arc<CountHook>, Arc<CountHook>) {
|
||||
fn build_session(
|
||||
provider_responses: Vec<MessageResponse>,
|
||||
) -> (AgentSession, Arc<CountHook>, Arc<CountHook>) {
|
||||
let mut hook_executor = HookExecutor::new();
|
||||
let start_count = Arc::new(CountHook(AtomicU32::new(0)));
|
||||
let end_count = Arc::new(CountHook(AtomicU32::new(0)));
|
||||
@@ -849,7 +834,8 @@ mod tests {
|
||||
/// 烟雾测试 1:AgentSession::submit_turn 跑通 mock provider(向后兼容)。
|
||||
#[tokio::test]
|
||||
async fn submit_turn_runs_with_mock_provider() {
|
||||
let (mut session, start_count, end_count) = build_session(vec![assistant_text("hello back")]);
|
||||
let (mut session, start_count, end_count) =
|
||||
build_session(vec![assistant_text("hello back")]);
|
||||
assert_eq!(session.turn_index(), 0);
|
||||
|
||||
let response = session.submit_turn("hi").await.unwrap();
|
||||
@@ -884,10 +870,8 @@ mod tests {
|
||||
/// 烟雾测试 3:submit_turn 触发 OnTurnStart / OnTurnEnd hook。
|
||||
#[tokio::test]
|
||||
async fn submit_turn_triggers_turn_hooks() {
|
||||
let (mut session, start_count, end_count) = build_session(vec![
|
||||
assistant_text("ok"),
|
||||
assistant_text("ok 2"),
|
||||
]);
|
||||
let (mut session, start_count, end_count) =
|
||||
build_session(vec![assistant_text("ok"), assistant_text("ok 2")]);
|
||||
|
||||
session.submit_turn("hi").await.unwrap();
|
||||
assert_eq!(start_count.0.load(Ordering::SeqCst), 1);
|
||||
@@ -940,9 +924,15 @@ mod tests {
|
||||
// 即 [user_input, tool_results?, final_response](不含 system_prompt,system 由 agent 提供)
|
||||
assert!(slot.messages.len() >= 2, "应至少包含 user 和 assistant");
|
||||
// 验证 user 输入和 assistant 响应都已写入
|
||||
let has_user = slot.messages.iter().any(|m| extract_text(m) == "user input");
|
||||
let has_user = slot
|
||||
.messages
|
||||
.iter()
|
||||
.any(|m| extract_text(m) == "user input");
|
||||
let has_resp = slot.messages.iter().any(|m| extract_text(m) == "resp");
|
||||
assert!(has_user && has_resp, "slot 应包含 user input 和 assistant response");
|
||||
assert!(
|
||||
has_user && has_resp,
|
||||
"slot 应包含 user input 和 assistant response"
|
||||
);
|
||||
}
|
||||
|
||||
/// Phase 10: create_slot 创建新 slot。
|
||||
@@ -986,7 +976,11 @@ mod tests {
|
||||
// 3. 检查 slot_a 的消息数
|
||||
let slot_a = session.slots.get("slot_a").unwrap();
|
||||
let slot_a_count = slot_a.messages.len();
|
||||
assert!(slot_a_count >= 2, "slot_a 至少 2 条消息,实际 {}", slot_a_count);
|
||||
assert!(
|
||||
slot_a_count >= 2,
|
||||
"slot_a 至少 2 条消息,实际 {}",
|
||||
slot_a_count
|
||||
);
|
||||
|
||||
// 4. 切回 default,验证 default 不包含 slot_a 的消息
|
||||
session.switch_slot("default").await.unwrap();
|
||||
@@ -1281,7 +1275,10 @@ mod tests {
|
||||
.iter()
|
||||
.any(|m| matches!(m, Message::User { .. }));
|
||||
let has_resp = slot.messages.iter().any(|m| extract_text(m) == "hi back");
|
||||
assert!(has_user && has_resp, "default slot 应包含 user 和 assistant 消息");
|
||||
assert!(
|
||||
has_user && has_resp,
|
||||
"default slot 应包含 user 和 assistant 消息"
|
||||
);
|
||||
}
|
||||
|
||||
/// Phase 9 Step 5.2 — `submit_turn_stream` 触发 OnTurnStart / OnTurnEnd hook。
|
||||
@@ -1380,7 +1377,11 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn summary_not_generated_below_threshold() {
|
||||
let mut session = build_session_with_summary(
|
||||
vec![assistant_text("a"), assistant_text("b"), assistant_text("c")],
|
||||
vec![
|
||||
assistant_text("a"),
|
||||
assistant_text("b"),
|
||||
assistant_text("c"),
|
||||
],
|
||||
vec![assistant_text("should_not_appear")],
|
||||
SummaryConfig::default(),
|
||||
);
|
||||
@@ -1400,16 +1401,20 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn summary_generated_above_threshold() {
|
||||
let mut session = build_session_with_summary(
|
||||
vec![assistant_text("a"), assistant_text("b"), assistant_text("c")],
|
||||
vec![
|
||||
assistant_text("a"),
|
||||
assistant_text("b"),
|
||||
assistant_text("c"),
|
||||
],
|
||||
vec![
|
||||
assistant_text("summary-1"),
|
||||
assistant_text("summary-2"),
|
||||
assistant_text("summary-3"),
|
||||
],
|
||||
SummaryConfig {
|
||||
max_context_tokens: 20, // 阈值 20 * 0.5 = 10
|
||||
max_context_tokens: 20, // 阈值 20 * 0.5 = 10
|
||||
trigger_token_ratio: 0.5,
|
||||
debounce_turns: 0, // 关闭防抖便于测试
|
||||
debounce_turns: 0, // 关闭防抖便于测试
|
||||
..SummaryConfig::default()
|
||||
},
|
||||
);
|
||||
@@ -1459,7 +1464,11 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn summary_written_to_session_memory_but_full_mode_does_not_inject() {
|
||||
let mut session = build_session_with_summary(
|
||||
vec![assistant_text("a"), assistant_text("b"), assistant_text("c")],
|
||||
vec![
|
||||
assistant_text("a"),
|
||||
assistant_text("b"),
|
||||
assistant_text("c"),
|
||||
],
|
||||
vec![assistant_text("captured-summary")],
|
||||
SummaryConfig {
|
||||
max_context_tokens: 20,
|
||||
@@ -1485,7 +1494,7 @@ mod tests {
|
||||
// 故意只提供 1 个对话响应;摘要调用时队列耗尽,MockProvider 返回 LlmError::Other
|
||||
let mut session = build_session_with_summary(
|
||||
vec![assistant_text("only-one")], // 后续摘要会失败
|
||||
vec![], // 无摘要响应
|
||||
vec![], // 无摘要响应
|
||||
SummaryConfig {
|
||||
max_context_tokens: 20,
|
||||
trigger_token_ratio: 0.5,
|
||||
@@ -1639,4 +1648,4 @@ mod tests {
|
||||
let summary = session.get_conversation_summary().await.unwrap();
|
||||
assert!(summary.is_none(), "巨型 max_context_tokens 应永不触发");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -48,7 +48,8 @@ impl SessionMemory {
|
||||
/// **不保留 metadata 和 created_at** —— 写入时 metadata 为空 JSON `{}`,created_at 为 `now_utc()`。
|
||||
/// 若需保留这两个字段(如 checkpoint rollback),使用 [`Self::set_with_meta`]。
|
||||
pub async fn set(&self, key: &str, value: &str) -> Result<(), AgentError> {
|
||||
self.set_with_meta(key, value, serde_json::json!({}), None).await
|
||||
self.set_with_meta(key, value, serde_json::json!({}), None)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 写入一条 key-value 条目(含完整 metadata + created_at)。
|
||||
|
||||
@@ -98,7 +98,11 @@ pub fn format_messages_as_text(messages: &[Message], max_tool_result_chars: usiz
|
||||
content,
|
||||
is_error,
|
||||
} => {
|
||||
let label = if *is_error { "Tool Error" } else { "Tool Result" };
|
||||
let label = if *is_error {
|
||||
"Tool Error"
|
||||
} else {
|
||||
"Tool Result"
|
||||
};
|
||||
if let Some(text) = first_text(content) {
|
||||
let truncated = truncate_chars(text, max_tool_result_chars);
|
||||
lines.push(format!("{} [{}]: {}", label, tool_call_id, truncated));
|
||||
|
||||
+64
-14
@@ -339,7 +339,14 @@ impl RecursiveCharacterSplitter {
|
||||
let take_n = overlap.min(prev_chars_count);
|
||||
|
||||
// 字符级安全地取 prev 末尾 take_n 个字符
|
||||
let tail: String = prev.chars().rev().take(take_n).collect::<Vec<_>>().into_iter().rev().collect();
|
||||
let tail: String = prev
|
||||
.chars()
|
||||
.rev()
|
||||
.take(take_n)
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
.collect();
|
||||
chunks[i] = format!("{}{}", tail, chunks[i]);
|
||||
}
|
||||
|
||||
@@ -408,8 +415,14 @@ mod tests {
|
||||
let chunks = splitter.split(&[doc]);
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0].content, "hello");
|
||||
assert_eq!(chunks[0].metadata.get("chunk_index").map(|s| s.as_str()), Some("0"));
|
||||
assert_eq!(chunks[0].metadata.get("chunk_count").map(|s| s.as_str()), Some("1"));
|
||||
assert_eq!(
|
||||
chunks[0].metadata.get("chunk_index").map(|s| s.as_str()),
|
||||
Some("0")
|
||||
);
|
||||
assert_eq!(
|
||||
chunks[0].metadata.get("chunk_count").map(|s| s.as_str()),
|
||||
Some("1")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -464,7 +477,11 @@ mod tests {
|
||||
// para1 (5 chars) > chunk_size=4 → 递归降级到 char 级拆分
|
||||
// para2 同理
|
||||
// 总共应该产生多个 chunk
|
||||
assert!(chunks.len() >= 2, "expected >= 2 chunks, got {}", chunks.len());
|
||||
assert!(
|
||||
chunks.len() >= 2,
|
||||
"expected >= 2 chunks, got {}",
|
||||
chunks.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -474,10 +491,18 @@ mod tests {
|
||||
let text: String = "a".repeat(200);
|
||||
let doc = Document::from_raw("long", &text);
|
||||
let chunks = splitter.split(&[doc]);
|
||||
assert!(chunks.len() >= 3, "expected >= 3 chunks, got {}", chunks.len());
|
||||
assert!(
|
||||
chunks.len() >= 3,
|
||||
"expected >= 3 chunks, got {}",
|
||||
chunks.len()
|
||||
);
|
||||
for chunk in &chunks {
|
||||
// chunk 内容 = overlap_tail(≤5) + new_content(≤50),故 ≤ 55
|
||||
assert!(chars_len(&chunk.content) <= 55, "chunk too long: {} chars", chars_len(&chunk.content));
|
||||
assert!(
|
||||
chars_len(&chunk.content) <= 55,
|
||||
"chunk too long: {} chars",
|
||||
chars_len(&chunk.content)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -488,7 +513,11 @@ mod tests {
|
||||
let doc = Document::from_raw("g", "aa\n\nbb\n\ncc\n\ndd");
|
||||
let chunks = splitter.split(&[doc]);
|
||||
// 短段应被合并:总共应该少于 4 个 chunk
|
||||
assert!(chunks.len() <= 3, "expected <= 3 chunks after merge, got {}", chunks.len());
|
||||
assert!(
|
||||
chunks.len() <= 3,
|
||||
"expected <= 3 chunks after merge, got {}",
|
||||
chunks.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -538,11 +567,19 @@ mod tests {
|
||||
let doc = Document::from_raw("cjk", &text);
|
||||
let chunks = splitter.split(&[doc]);
|
||||
// 30 字符 / 10 chunk_size = 3 个 chunk
|
||||
assert!(chunks.len() >= 3, "expected >= 3 chunks for 30 chars / chunk_size=10, got {}", chunks.len());
|
||||
assert!(
|
||||
chunks.len() >= 3,
|
||||
"expected >= 3 chunks for 30 chars / chunk_size=10, got {}",
|
||||
chunks.len()
|
||||
);
|
||||
for chunk in &chunks {
|
||||
let char_count = chars_len(&chunk.content);
|
||||
// chunk = overlap_tail(≤2) + new_content(≤10),故 ≤ 12
|
||||
assert!(char_count <= 12, "chunk char count {} exceeds 10+overlap", char_count);
|
||||
assert!(
|
||||
char_count <= 12,
|
||||
"chunk char count {} exceeds 10+overlap",
|
||||
char_count
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -569,12 +606,25 @@ mod tests {
|
||||
fn split_metadata_inheritance() {
|
||||
let splitter = RecursiveCharacterSplitter::new(100, 10);
|
||||
let mut doc = Document::new("m", "short content", "text/plain");
|
||||
doc.metadata.insert("author".to_string(), "alice".to_string());
|
||||
doc.metadata
|
||||
.insert("author".to_string(), "alice".to_string());
|
||||
let chunks = splitter.split(&[doc]);
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0].metadata.get("author").map(|s| s.as_str()), Some("alice"));
|
||||
assert_eq!(chunks[0].metadata.get("source_id").map(|s| s.as_str()), Some("m"));
|
||||
assert_eq!(chunks[0].metadata.get("chunk_index").map(|s| s.as_str()), Some("0"));
|
||||
assert_eq!(chunks[0].metadata.get("chunk_count").map(|s| s.as_str()), Some("1"));
|
||||
assert_eq!(
|
||||
chunks[0].metadata.get("author").map(|s| s.as_str()),
|
||||
Some("alice")
|
||||
);
|
||||
assert_eq!(
|
||||
chunks[0].metadata.get("source_id").map(|s| s.as_str()),
|
||||
Some("m")
|
||||
);
|
||||
assert_eq!(
|
||||
chunks[0].metadata.get("chunk_index").map(|s| s.as_str()),
|
||||
Some("0")
|
||||
);
|
||||
assert_eq!(
|
||||
chunks[0].metadata.get("chunk_count").map(|s| s.as_str()),
|
||||
Some("1")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,8 +16,8 @@ use serde::{Deserialize, Serialize};
|
||||
use time::OffsetDateTime;
|
||||
|
||||
use crate::agent::session::AgentSession;
|
||||
use crate::engine::snapshot::SessionSnapshot;
|
||||
use crate::engine::EngineError;
|
||||
use crate::engine::snapshot::SessionSnapshot;
|
||||
use crate::memory::store::MemoryStore;
|
||||
use crate::memory::types::{MemoryFilter, MemoryItem};
|
||||
|
||||
@@ -80,9 +80,8 @@ impl Checkpointer {
|
||||
let ckpt_id = generate_ckpt_id();
|
||||
let key = ckpt_key(&session.session_id, &ckpt_id);
|
||||
|
||||
let json = serde_json::to_string(&snapshot).map_err(|e| {
|
||||
EngineError::Serialization(format!("snapshot serialize failed: {e}"))
|
||||
})?;
|
||||
let json = serde_json::to_string(&snapshot)
|
||||
.map_err(|e| EngineError::Serialization(format!("snapshot serialize failed: {e}")))?;
|
||||
|
||||
let item = MemoryItem {
|
||||
id: key,
|
||||
@@ -141,10 +140,7 @@ impl Checkpointer {
|
||||
/// prefix 查询 `ckpt:{session_id}:` → 反序列化 `SessionSnapshot` → 提取元数据。
|
||||
/// 不需要 `CkptMeta` 单独存储——`SessionSnapshot` 已含 `turn_index` 字段,
|
||||
/// `created_at` 用 `MemoryItem.created_at` 转换。
|
||||
pub async fn list_checkpoints(
|
||||
&self,
|
||||
session_id: &str,
|
||||
) -> Result<Vec<CkptMeta>, EngineError> {
|
||||
pub async fn list_checkpoints(&self, session_id: &str) -> Result<Vec<CkptMeta>, EngineError> {
|
||||
let prefix = format!("ckpt:{}:", session_id);
|
||||
let filter = MemoryFilter {
|
||||
prefix: Some(prefix),
|
||||
@@ -346,10 +342,7 @@ mod tests {
|
||||
let mut session = new_session_for_test("latest-session");
|
||||
cp.checkpoint(&session).await.unwrap();
|
||||
tokio::time::sleep(std::time::Duration::from_millis(2)).await;
|
||||
session
|
||||
.set_session_data("v", "2")
|
||||
.await
|
||||
.unwrap();
|
||||
session.set_session_data("v", "2").await.unwrap();
|
||||
cp.checkpoint(&session).await.unwrap();
|
||||
|
||||
let latest = cp.latest_snapshot("latest-session").await.unwrap().unwrap();
|
||||
@@ -374,4 +367,4 @@ mod tests {
|
||||
assert_eq!(metas_a[0].session_id, "iso-a");
|
||||
assert_eq!(metas_b[0].session_id, "iso-b");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -48,4 +48,4 @@ pub enum EngineError {
|
||||
/// 调用方收到此错误时,子 session 已通过 `destroy()` 清理(SessionMeta + checkpoint 全部清空)。
|
||||
#[error("Dispatch failed: {0}")]
|
||||
DispatchFailed(String),
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -20,4 +20,4 @@ pub use checkpointer::{Checkpointer, CkptMeta};
|
||||
pub use error::EngineError;
|
||||
pub use session_manager::{SessionManager, SessionManagerConfig};
|
||||
pub use snapshot::{SessionMemoryEntry, SessionSnapshot};
|
||||
pub use sub_agent::{DispatchConfig, SubTaskResult, SubTaskStreamEvent};
|
||||
pub use sub_agent::{DispatchConfig, SubTaskResult, SubTaskStreamEvent};
|
||||
|
||||
@@ -141,11 +141,11 @@ impl SessionManager {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn load_session_meta(&self, session_id: &str) -> Result<Option<SessionMeta>, EngineError> {
|
||||
let item = self
|
||||
.store
|
||||
.get(&SessionMeta::meta_key(session_id))
|
||||
.await?;
|
||||
pub(crate) async fn load_session_meta(
|
||||
&self,
|
||||
session_id: &str,
|
||||
) -> Result<Option<SessionMeta>, EngineError> {
|
||||
let item = self.store.get(&SessionMeta::meta_key(session_id)).await?;
|
||||
match item {
|
||||
Some(item) => {
|
||||
let meta: SessionMeta = serde_json::from_str(&item.content).map_err(|e| {
|
||||
@@ -241,10 +241,7 @@ impl SessionManager {
|
||||
/// 按 ID 获取 session(**仅查内存**,不自动从存储恢复)。
|
||||
///
|
||||
/// 冷启动时 `get()` 未命中返回 `SessionNotFound`。如需从存储恢复,使用 `recover()` 方法。
|
||||
pub async fn get(
|
||||
&self,
|
||||
session_id: &str,
|
||||
) -> Result<Arc<Mutex<AgentSession>>, EngineError> {
|
||||
pub async fn get(&self, session_id: &str) -> Result<Arc<Mutex<AgentSession>>, EngineError> {
|
||||
let sessions = self.sessions.read().await;
|
||||
let result = sessions.get(session_id).cloned();
|
||||
tracing::debug!(
|
||||
@@ -392,12 +389,7 @@ impl SessionManager {
|
||||
session_id: &str,
|
||||
user_input: impl Into<String>,
|
||||
) -> Result<
|
||||
std::pin::Pin<
|
||||
Box<
|
||||
dyn futures_core::Stream<Item = crate::llm::stream::StreamEvent>
|
||||
+ Send,
|
||||
>,
|
||||
>,
|
||||
std::pin::Pin<Box<dyn futures_core::Stream<Item = crate::llm::stream::StreamEvent> + Send>>,
|
||||
EngineError,
|
||||
> {
|
||||
let session = self.get(session_id).await?;
|
||||
@@ -593,8 +585,14 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
let session = sm.get(&id).await.unwrap();
|
||||
sm.checkpointer().checkpoint(&*session.lock().await).await.unwrap();
|
||||
assert_eq!(sm.checkpointer().list_checkpoints(&id).await.unwrap().len(), 1);
|
||||
sm.checkpointer()
|
||||
.checkpoint(&*session.lock().await)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
sm.checkpointer().list_checkpoints(&id).await.unwrap().len(),
|
||||
1
|
||||
);
|
||||
|
||||
sm.destroy(&id).await.unwrap();
|
||||
|
||||
@@ -603,7 +601,10 @@ mod tests {
|
||||
// SessionMeta 已删除
|
||||
assert!(sm.load_session_meta(&id).await.unwrap().is_none());
|
||||
// Checkpoint 已删除
|
||||
assert_eq!(sm.checkpointer().list_checkpoints(&id).await.unwrap().len(), 0);
|
||||
assert_eq!(
|
||||
sm.checkpointer().list_checkpoints(&id).await.unwrap().len(),
|
||||
0
|
||||
);
|
||||
|
||||
// destroy 不存在的 session 不报错
|
||||
sm.destroy(&id).await.unwrap();
|
||||
@@ -682,7 +683,10 @@ mod tests {
|
||||
{
|
||||
let s = sm.get(&id).await.unwrap();
|
||||
s.lock().await.set_session_data("k", "v1").await.unwrap();
|
||||
sm.checkpointer().checkpoint(&*s.lock().await).await.unwrap();
|
||||
sm.checkpointer()
|
||||
.checkpoint(&*s.lock().await)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// 2. 模拟"进程重启"——清空内存但保留 store
|
||||
@@ -710,7 +714,10 @@ mod tests {
|
||||
.unwrap();
|
||||
{
|
||||
let s = sm.get(&id).await.unwrap();
|
||||
sm.checkpointer().checkpoint(&*s.lock().await).await.unwrap();
|
||||
sm.checkpointer()
|
||||
.checkpoint(&*s.lock().await)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let err = sm
|
||||
@@ -750,7 +757,10 @@ mod tests {
|
||||
let err = sm.submit_turn(&id, "hello").await.unwrap_err();
|
||||
match err {
|
||||
EngineError::Agent(crate::agent::error::AgentError::Llm(_)) => {}
|
||||
other => panic!("expected EngineError::Agent(AgentError::Llm), got {:?}", other),
|
||||
other => panic!(
|
||||
"expected EngineError::Agent(AgentError::Llm), got {:?}",
|
||||
other
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -772,8 +782,14 @@ mod tests {
|
||||
let _ = sm.submit_turn(&id, "x").await;
|
||||
|
||||
// 手动 checkpoint 仍可工作
|
||||
sm.checkpointer().checkpoint(&*s.lock().await).await.unwrap();
|
||||
assert_eq!(sm.checkpointer().list_checkpoints(&id).await.unwrap().len(), 1);
|
||||
sm.checkpointer()
|
||||
.checkpoint(&*s.lock().await)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
sm.checkpointer().list_checkpoints(&id).await.unwrap().len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -787,7 +803,8 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
// 构造一个新 session(同 session_id)然后 replace
|
||||
let mut new_session = AgentSession::new(Arc::new(StubAgent("a".into())), &id, make_bundle());
|
||||
let mut new_session =
|
||||
AgentSession::new(Arc::new(StubAgent("a".into())), &id, make_bundle());
|
||||
new_session
|
||||
.set_session_data("replaced", "yes")
|
||||
.await
|
||||
@@ -838,8 +855,7 @@ mod tests {
|
||||
/// restore_memory 幂等性:第二次调用应立即返回 Ok(())(pending 已被清空)。
|
||||
#[tokio::test]
|
||||
async fn restore_memory_is_idempotent() {
|
||||
let store: Arc<dyn MemoryStore> =
|
||||
Arc::new(crate::memory::store::InMemoryStore::new());
|
||||
let store: Arc<dyn MemoryStore> = Arc::new(crate::memory::store::InMemoryStore::new());
|
||||
let sm = SessionManager::new(store.clone());
|
||||
|
||||
let id = sm
|
||||
@@ -848,12 +864,11 @@ mod tests {
|
||||
.unwrap();
|
||||
{
|
||||
let s = sm.get(&id).await.unwrap();
|
||||
s.lock()
|
||||
.await
|
||||
.set_session_data("k", "v")
|
||||
s.lock().await.set_session_data("k", "v").await.unwrap();
|
||||
sm.checkpointer()
|
||||
.checkpoint(&*s.lock().await)
|
||||
.await
|
||||
.unwrap();
|
||||
sm.checkpointer().checkpoint(&*s.lock().await).await.unwrap();
|
||||
}
|
||||
|
||||
// 模拟"进程重启"——新建 SessionManager,复用 store
|
||||
@@ -877,8 +892,7 @@ mod tests {
|
||||
/// 10 并发 session 创建:验证 RwLock 写锁争用下不冲突,所有 ID 唯一。
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn concurrent_create_ten_sessions() {
|
||||
let store: Arc<dyn MemoryStore> =
|
||||
Arc::new(crate::memory::store::InMemoryStore::new());
|
||||
let store: Arc<dyn MemoryStore> = Arc::new(crate::memory::store::InMemoryStore::new());
|
||||
let sm = Arc::new(SessionManager::new(store));
|
||||
|
||||
let mut handles = Vec::with_capacity(10);
|
||||
@@ -905,4 +919,4 @@ mod tests {
|
||||
assert!(sm.get(id).await.is_ok());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -36,4 +36,4 @@ pub struct SessionSnapshot {
|
||||
pub last_summary_turn: Option<u32>,
|
||||
#[serde(default)]
|
||||
pub session_memory_data: std::collections::HashMap<String, SessionMemoryEntry>,
|
||||
}
|
||||
}
|
||||
|
||||
+23
-22
@@ -77,10 +77,7 @@ pub enum SubTaskStreamEvent {
|
||||
/// 执行完成,携带完整结果。
|
||||
Completed(SubTaskResult),
|
||||
/// 流式调度中的错误。
|
||||
Error {
|
||||
child_id: String,
|
||||
error: String,
|
||||
},
|
||||
Error { child_id: String, error: String },
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SubTaskStreamEvent {
|
||||
@@ -150,7 +147,7 @@ impl SessionManager {
|
||||
|
||||
// 3. 过滤(按 bridge_keys)
|
||||
let filtered: Vec<_> = match &config.bridge_keys {
|
||||
None => Vec::new(), // None = 不继承任何
|
||||
None => Vec::new(), // None = 不继承任何
|
||||
Some(keys) if keys.is_empty() => entries, // 空列表 = 全部继承
|
||||
Some(keys) => entries
|
||||
.into_iter()
|
||||
@@ -343,9 +340,7 @@ impl SessionManager {
|
||||
task: impl Into<String>,
|
||||
config: DispatchConfig,
|
||||
) -> Result<
|
||||
std::pin::Pin<
|
||||
Box<dyn futures_core::Stream<Item = SubTaskStreamEvent> + Send>,
|
||||
>,
|
||||
std::pin::Pin<Box<dyn futures_core::Stream<Item = SubTaskStreamEvent> + Send>>,
|
||||
EngineError,
|
||||
> {
|
||||
use futures_util::StreamExt;
|
||||
@@ -473,9 +468,7 @@ impl SessionManager {
|
||||
}
|
||||
};
|
||||
let mut guard = session.lock().await;
|
||||
guard
|
||||
.finalize_turn(response, new_messages)
|
||||
.await
|
||||
guard.finalize_turn(response, new_messages).await
|
||||
};
|
||||
|
||||
if let Err(e) = lock_result {
|
||||
@@ -575,8 +568,7 @@ mod tests {
|
||||
}
|
||||
|
||||
async fn make_manager_and_bundle() -> (Arc<SessionManager>, Arc<RuntimeBundle>) {
|
||||
let store: Arc<dyn crate::memory::store::MemoryStore> =
|
||||
Arc::new(InMemoryStore::new());
|
||||
let store: Arc<dyn crate::memory::store::MemoryStore> = Arc::new(InMemoryStore::new());
|
||||
let provider = Arc::new(MockProvider::new(vec![
|
||||
assistant_text("child response 1"),
|
||||
assistant_text("child response 2"),
|
||||
@@ -702,7 +694,12 @@ mod tests {
|
||||
|
||||
// dispatch 到不存在的 parent_id
|
||||
let result = sm
|
||||
.dispatch("nonexistent_parent", child, "task", DispatchConfig::default())
|
||||
.dispatch(
|
||||
"nonexistent_parent",
|
||||
child,
|
||||
"task",
|
||||
DispatchConfig::default(),
|
||||
)
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
|
||||
@@ -716,7 +713,12 @@ mod tests {
|
||||
let (sm, _bundle) = make_manager_and_bundle().await;
|
||||
let child: Arc<dyn Agent> = Arc::new(MockAgent::new("child"));
|
||||
let result = sm
|
||||
.dispatch("nonexistent_parent", child, "task", DispatchConfig::default())
|
||||
.dispatch(
|
||||
"nonexistent_parent",
|
||||
child,
|
||||
"task",
|
||||
DispatchConfig::default(),
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(result, Err(EngineError::SessionNotFound(_))));
|
||||
}
|
||||
@@ -724,9 +726,10 @@ mod tests {
|
||||
// ====== dispatch_all 测试 ======
|
||||
|
||||
/// 提供充足的 mock response(>= 3)
|
||||
async fn make_manager_and_bundle_for_all(n: usize) -> (Arc<SessionManager>, Arc<RuntimeBundle>) {
|
||||
let store: Arc<dyn crate::memory::store::MemoryStore> =
|
||||
Arc::new(InMemoryStore::new());
|
||||
async fn make_manager_and_bundle_for_all(
|
||||
n: usize,
|
||||
) -> (Arc<SessionManager>, Arc<RuntimeBundle>) {
|
||||
let store: Arc<dyn crate::memory::store::MemoryStore> = Arc::new(InMemoryStore::new());
|
||||
let responses: Vec<_> = (0..n)
|
||||
.map(|i| assistant_text(&format!("response {i}")))
|
||||
.collect();
|
||||
@@ -747,8 +750,7 @@ mod tests {
|
||||
/// 创建空 mock responses 的 manager 和 bundle —— 后续 dispatch 会触发
|
||||
/// "MockProvider: 预设响应已用完" 错误,可用于测试错误传播。
|
||||
async fn make_manager_and_bundle_empty_mock() -> (Arc<SessionManager>, Arc<RuntimeBundle>) {
|
||||
let store: Arc<dyn crate::memory::store::MemoryStore> =
|
||||
Arc::new(InMemoryStore::new());
|
||||
let store: Arc<dyn crate::memory::store::MemoryStore> = Arc::new(InMemoryStore::new());
|
||||
let provider = Arc::new(MockProvider::empty());
|
||||
let bundle = Arc::new(
|
||||
AgentBuilder::new()
|
||||
@@ -1048,8 +1050,7 @@ mod tests {
|
||||
}
|
||||
|
||||
// 验证:收到 Error 事件
|
||||
let (err_child_id, err_message) =
|
||||
error_received.expect("Error event should be received");
|
||||
let (err_child_id, err_message) = error_received.expect("Error event should be received");
|
||||
assert_eq!(
|
||||
Some(err_child_id.as_str()),
|
||||
child_id_from_event.as_deref(),
|
||||
|
||||
@@ -129,8 +129,7 @@ mod tests {
|
||||
}
|
||||
|
||||
async fn make_manager_and_bundle() -> (Arc<SessionManager>, Arc<RuntimeBundle>) {
|
||||
let store: Arc<dyn crate::memory::store::MemoryStore> =
|
||||
Arc::new(InMemoryStore::new());
|
||||
let store: Arc<dyn crate::memory::store::MemoryStore> = Arc::new(InMemoryStore::new());
|
||||
let provider = Arc::new(MockProvider::new(vec![
|
||||
assistant_text("response1"),
|
||||
assistant_text("response2"),
|
||||
@@ -175,10 +174,7 @@ mod tests {
|
||||
{
|
||||
let session = sm.get(&sid).await.unwrap();
|
||||
let mut guard = session.lock().await;
|
||||
guard
|
||||
.set_session_data("key1", "value1")
|
||||
.await
|
||||
.unwrap();
|
||||
guard.set_session_data("key1", "value1").await.unwrap();
|
||||
}
|
||||
sm.submit_turn(&sid, "hello").await.unwrap();
|
||||
|
||||
|
||||
+8
-8
@@ -1,19 +1,19 @@
|
||||
//! agcore —— 智能体(Agent)核心工具箱。
|
||||
|
||||
#[cfg(any(feature = "llm-types", feature = "llm"))]
|
||||
pub mod llm;
|
||||
#[cfg(feature = "agent")]
|
||||
pub mod agent;
|
||||
#[cfg(feature = "document")]
|
||||
pub mod document;
|
||||
#[cfg(feature = "engine")]
|
||||
pub mod engine;
|
||||
#[cfg(any(feature = "llm-types", feature = "llm"))]
|
||||
pub mod llm;
|
||||
#[cfg(feature = "memory")]
|
||||
pub mod memory;
|
||||
#[cfg(feature = "prompt")]
|
||||
pub mod prompt;
|
||||
#[cfg(feature = "tools")]
|
||||
pub mod tools;
|
||||
#[cfg(feature = "memory")]
|
||||
pub mod memory;
|
||||
#[cfg(feature = "agent")]
|
||||
pub mod agent;
|
||||
#[cfg(feature = "engine")]
|
||||
pub mod engine;
|
||||
|
||||
#[cfg(feature = "document")]
|
||||
pub use document::Document;
|
||||
|
||||
+2
-2
@@ -1,7 +1,5 @@
|
||||
//! LLM 调用周期 —— 大模型基础调用周期控制。
|
||||
|
||||
#[cfg(feature = "llm-types")]
|
||||
pub mod types;
|
||||
#[cfg(feature = "llm")]
|
||||
pub mod compact;
|
||||
#[cfg(feature = "llm")]
|
||||
@@ -16,6 +14,8 @@ pub mod error;
|
||||
pub mod hooks;
|
||||
#[cfg(feature = "llm")]
|
||||
pub mod mock;
|
||||
#[cfg(feature = "llm-types")]
|
||||
pub mod types;
|
||||
// provider 模块依赖 reqwest(通过 reqwest::Client),仅在任一 provider feature 启用时编译
|
||||
#[cfg(any(
|
||||
feature = "provider-openai",
|
||||
|
||||
+103
-78
@@ -15,17 +15,17 @@ use serde_json::Value;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_stream::wrappers::UnboundedReceiverStream;
|
||||
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::llm::compact::{CompactConfig, CompactState, microcompact, should_compact};
|
||||
use crate::llm::cycle::retry::should_retry;
|
||||
use crate::llm::error::LlmError;
|
||||
use crate::llm::hooks::{HookContext, HookEvent, HookExecutor};
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::llm::stream::StreamEvent;
|
||||
use crate::llm::types::ToolChoice;
|
||||
use crate::llm::types::message::{ContentBlock, Message};
|
||||
use crate::llm::types::request_v2::MessageRequest;
|
||||
use crate::llm::types::response_v2::{MessageResponse, PartialMessageResponse, StopReason};
|
||||
use crate::llm::types::tool::ToolDef;
|
||||
use crate::llm::types::ToolChoice;
|
||||
#[cfg(feature = "tools")]
|
||||
use crate::tools::ToolRegistry;
|
||||
|
||||
@@ -452,10 +452,7 @@ impl LlmCycle {
|
||||
/// 内部请求方法(与 `submit` 共享重试逻辑,但不 push user message 和 Assistant 响应)。
|
||||
///
|
||||
/// 用于 `submit_with_tools()` 的多轮 tool 循环。
|
||||
async fn submit_request(
|
||||
&mut self,
|
||||
tools: &[ToolDef],
|
||||
) -> Result<MessageResponse, LlmError> {
|
||||
async fn submit_request(&mut self, tools: &[ToolDef]) -> Result<MessageResponse, LlmError> {
|
||||
let mut attempts = 0;
|
||||
|
||||
loop {
|
||||
@@ -913,11 +910,10 @@ async fn run_tool_loop(
|
||||
Ok(v) => {
|
||||
// ponytail: 与 submit_with_tools 行为对齐 —— 用 truncate_tool_result
|
||||
// 截断结果以防止超大工具输出在 tool 循环中膨胀 messages 上下文窗口
|
||||
let serialized =
|
||||
serde_json::to_string(v).unwrap_or_else(|e| {
|
||||
tracing::warn!("工具结果序列化失败: {}", e);
|
||||
"{}".to_string()
|
||||
});
|
||||
let serialized = serde_json::to_string(v).unwrap_or_else(|e| {
|
||||
tracing::warn!("工具结果序列化失败: {}", e);
|
||||
"{}".to_string()
|
||||
});
|
||||
truncate_tool_result(&serialized, max_bytes)
|
||||
}
|
||||
Err(e) if e.is_recoverable() => format!("错误: {}", e),
|
||||
@@ -1204,8 +1200,7 @@ mod tests {
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_submit_with_tools_stream_pure_text() {
|
||||
let provider = Mock::new(vec![assistant_text_response("你好")]);
|
||||
let mut cycle =
|
||||
LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut cycle = LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut registry = ToolRegistry::new();
|
||||
registry.register(std::sync::Arc::new(AddTool)).unwrap();
|
||||
|
||||
@@ -1216,23 +1211,36 @@ mod tests {
|
||||
let events = drain(stream).await;
|
||||
|
||||
// 期望序列:MessageStart → ContentBlockStart(Text) → TextDelta → ContentBlockEnd → CostUpdate → MessageComplete
|
||||
assert!(matches!(events.first(), Some(StreamEvent::MessageStart { .. })));
|
||||
assert!(events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::TextDelta { text } if text == "你好")));
|
||||
assert!(events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::MessageComplete { .. })));
|
||||
assert!(matches!(
|
||||
events.first(),
|
||||
Some(StreamEvent::MessageStart { .. })
|
||||
));
|
||||
assert!(
|
||||
events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::TextDelta { text } if text == "你好"))
|
||||
);
|
||||
assert!(
|
||||
events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::MessageComplete { .. }))
|
||||
);
|
||||
// 纯文本流不应有 ToolExecutionStarted/Completed 事件
|
||||
assert!(!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::ToolExecutionStarted { .. })));
|
||||
assert!(!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::ToolExecutionCompleted { .. })));
|
||||
assert!(!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::Error { .. })));
|
||||
assert!(
|
||||
!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::ToolExecutionStarted { .. }))
|
||||
);
|
||||
assert!(
|
||||
!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::ToolExecutionCompleted { .. }))
|
||||
);
|
||||
assert!(
|
||||
!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::Error { .. }))
|
||||
);
|
||||
}
|
||||
|
||||
/// Phase 9 测试 3.2 — 单轮工具调用
|
||||
@@ -1242,8 +1250,7 @@ mod tests {
|
||||
assistant_tool_call_response(vec![("call_1", "add", r#"{"a":1,"b":2}"#)]),
|
||||
assistant_text_response("答案是 3"),
|
||||
]);
|
||||
let mut cycle =
|
||||
LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut cycle = LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut registry = ToolRegistry::new();
|
||||
registry.register(std::sync::Arc::new(AddTool)).unwrap();
|
||||
|
||||
@@ -1320,8 +1327,7 @@ mod tests {
|
||||
assistant_tool_call_response(vec![("call_3", "add", r#"{"a":5,"b":6}"#)]),
|
||||
assistant_text_response("完成"),
|
||||
]);
|
||||
let mut cycle =
|
||||
LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut cycle = LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut registry = ToolRegistry::new();
|
||||
registry.register(std::sync::Arc::new(AddTool)).unwrap();
|
||||
|
||||
@@ -1381,8 +1387,11 @@ mod tests {
|
||||
})
|
||||
.collect();
|
||||
assert!(
|
||||
error_events.iter().any(|m| m.contains("达到最大工具循环轮次")),
|
||||
"应包含最大轮次超限 Error,实际: {:?}", error_events
|
||||
error_events
|
||||
.iter()
|
||||
.any(|m| m.contains("达到最大工具循环轮次")),
|
||||
"应包含最大轮次超限 Error,实际: {:?}",
|
||||
error_events
|
||||
);
|
||||
|
||||
// 工具调用次数应 ≤ 2
|
||||
@@ -1409,10 +1418,8 @@ mod tests {
|
||||
async fn chat_stream(
|
||||
&self,
|
||||
_r: MessageRequest,
|
||||
) -> Result<
|
||||
Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>,
|
||||
LlmError,
|
||||
> {
|
||||
) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>, LlmError>
|
||||
{
|
||||
Err(LlmError::Other("网络错误".into()))
|
||||
}
|
||||
fn capabilities(&self) -> ProviderCapabilities {
|
||||
@@ -1436,25 +1443,29 @@ mod tests {
|
||||
|
||||
// 第一个事件应是 Error(chat_stream Err 立即事件化)
|
||||
assert!(
|
||||
events.first().map(|e| matches!(e, StreamEvent::Error { .. }))
|
||||
events
|
||||
.first()
|
||||
.map(|e| matches!(e, StreamEvent::Error { .. }))
|
||||
== Some(true),
|
||||
"流中首个事件应是 Error,实际: {:?}", events.first()
|
||||
"流中首个事件应是 Error,实际: {:?}",
|
||||
events.first()
|
||||
);
|
||||
assert!(
|
||||
events
|
||||
.iter()
|
||||
.filter_map(|e| match e {
|
||||
StreamEvent::Error { message } => Some(message.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.any(|m| m.contains("网络错误"))
|
||||
);
|
||||
assert!(events
|
||||
.iter()
|
||||
.filter_map(|e| match e {
|
||||
StreamEvent::Error { message } => Some(message.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.any(|m| m.contains("网络错误")));
|
||||
}
|
||||
|
||||
/// Phase 9 测试 3.6 — 空 tool_registry
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_submit_with_tools_stream_empty_registry() {
|
||||
let provider = Mock::new(vec![assistant_text_response("纯文本回答")]);
|
||||
let mut cycle =
|
||||
LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut cycle = LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let registry = ToolRegistry::new();
|
||||
|
||||
let stream = cycle
|
||||
@@ -1464,18 +1475,26 @@ mod tests {
|
||||
let events = drain(stream).await;
|
||||
|
||||
// 流退化为纯文本流 —— 无 ToolExecution 事件,无 Error
|
||||
assert!(events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::TextDelta { text } if text == "纯文本回答")));
|
||||
assert!(!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::ToolExecutionStarted { .. })));
|
||||
assert!(!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::ToolExecutionCompleted { .. })));
|
||||
assert!(!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::Error { .. })));
|
||||
assert!(
|
||||
events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::TextDelta { text } if text == "纯文本回答"))
|
||||
);
|
||||
assert!(
|
||||
!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::ToolExecutionStarted { .. }))
|
||||
);
|
||||
assert!(
|
||||
!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::ToolExecutionCompleted { .. }))
|
||||
);
|
||||
assert!(
|
||||
!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::Error { .. }))
|
||||
);
|
||||
}
|
||||
|
||||
/// Phase 9 测试 3.7 — 不可恢复工具错误
|
||||
@@ -1507,8 +1526,7 @@ mod tests {
|
||||
assistant_tool_call_response(vec![("call_x", "fail_unrecoverable", "{}")]),
|
||||
assistant_text_response("忽略"),
|
||||
]);
|
||||
let mut cycle =
|
||||
LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut cycle = LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut registry = ToolRegistry::new();
|
||||
registry
|
||||
.register(std::sync::Arc::new(UnrecoverableTool))
|
||||
@@ -1566,8 +1584,7 @@ mod tests {
|
||||
assistant_tool_call_response(vec![("call_y", "fail_recoverable", "{}")]),
|
||||
assistant_text_response("已恢复"),
|
||||
]);
|
||||
let mut cycle =
|
||||
LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut cycle = LlmCycle::new(Box::new(provider), CycleConfig::default());
|
||||
let mut registry = ToolRegistry::new();
|
||||
registry
|
||||
.register(std::sync::Arc::new(RecoverableTool))
|
||||
@@ -1580,9 +1597,11 @@ mod tests {
|
||||
let events = drain(stream).await;
|
||||
|
||||
// 可恢复错误:tool_result 回传 LLM,最终流正常结束
|
||||
assert!(!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::Error { .. })));
|
||||
assert!(
|
||||
!events
|
||||
.iter()
|
||||
.any(|e| matches!(e, StreamEvent::Error { .. }))
|
||||
);
|
||||
// 最终 MessageComplete 应是 Stop(不是 ToolUse)
|
||||
let final_response = events
|
||||
.iter()
|
||||
@@ -1602,7 +1621,10 @@ mod tests {
|
||||
_ => None,
|
||||
})
|
||||
.unwrap();
|
||||
assert!(completed, "可恢复错误的 ToolExecutionCompleted.is_error 应为 true");
|
||||
assert!(
|
||||
completed,
|
||||
"可恢复错误的 ToolExecutionCompleted.is_error 应为 true"
|
||||
);
|
||||
}
|
||||
|
||||
/// Phase 9 测试 3.9 — 工具超时
|
||||
@@ -1651,9 +1673,7 @@ mod tests {
|
||||
]);
|
||||
let mut cycle = LlmCycle::new(Box::new(provider), config);
|
||||
let mut registry = ToolRegistry::new();
|
||||
registry
|
||||
.register(std::sync::Arc::new(SlowTool))
|
||||
.unwrap();
|
||||
registry.register(std::sync::Arc::new(SlowTool)).unwrap();
|
||||
|
||||
let stream = cycle
|
||||
.submit_with_tools_stream("test".to_string(), std::sync::Arc::new(registry))
|
||||
@@ -1673,8 +1693,14 @@ mod tests {
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert!(!tool_completed.is_empty(), "应有 ToolExecutionCompleted 事件");
|
||||
assert!(tool_completed[0].0, "超时后 ToolExecutionCompleted.is_error 应为 true");
|
||||
assert!(
|
||||
!tool_completed.is_empty(),
|
||||
"应有 ToolExecutionCompleted 事件"
|
||||
);
|
||||
assert!(
|
||||
tool_completed[0].0,
|
||||
"超时后 ToolExecutionCompleted.is_error 应为 true"
|
||||
);
|
||||
assert_eq!(tool_completed[0].1, "slow_tool");
|
||||
|
||||
// 2. 流中应有不可恢复错误终止事件(tool_timeout → McpTimeout → 不可恢复 → Error)
|
||||
@@ -1686,10 +1712,9 @@ mod tests {
|
||||
})
|
||||
.collect();
|
||||
assert!(
|
||||
error_events.iter().any(|m| m.contains("不可恢复错误")),
|
||||
"应有不可恢复错误事件终止流,实际事件: {:?}",
|
||||
error_events
|
||||
.iter()
|
||||
.any(|m| m.contains("不可恢复错误")),
|
||||
"应有不可恢复错误事件终止流,实际事件: {:?}", error_events
|
||||
);
|
||||
|
||||
// 3. 第一轮的 MessageComplete { stop_reason: ToolUse } 在 Error 之前已发出
|
||||
|
||||
+11
-6
@@ -151,17 +151,18 @@ mod tests {
|
||||
let result = embedder.embed(&inputs).await.unwrap();
|
||||
for vec in &result {
|
||||
let norm = l2_norm(vec);
|
||||
assert!((norm - 1.0).abs() < 1e-5, "vector norm should be ~1.0, got {}", norm);
|
||||
assert!(
|
||||
(norm - 1.0).abs() < 1e-5,
|
||||
"vector norm should be ~1.0, got {}",
|
||||
norm
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn embed_different_inputs_different_vectors() {
|
||||
let embedder = MockEmbedding::new(16);
|
||||
let r1 = embedder
|
||||
.embed(&["hello world".to_string()])
|
||||
.await
|
||||
.unwrap();
|
||||
let r1 = embedder.embed(&["hello world".to_string()]).await.unwrap();
|
||||
let r2 = embedder
|
||||
.embed(&["completely different".to_string()])
|
||||
.await
|
||||
@@ -178,6 +179,10 @@ mod tests {
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].len(), 4);
|
||||
let norm = l2_norm(&result[0]);
|
||||
assert!((norm - 1.0).abs() < 1e-5, "empty-string vector norm should be ~1.0, got {}", norm);
|
||||
assert!(
|
||||
(norm - 1.0).abs() < 1e-5,
|
||||
"empty-string vector norm should be ~1.0, got {}",
|
||||
norm
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -18,7 +18,6 @@ use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
use crate::llm::{LlmProvider, ProviderCapabilities, ProviderFeatures};
|
||||
use crate::llm::error::LlmError;
|
||||
use crate::llm::types::message::{ContentBlock, ContentBlockType, Message};
|
||||
use crate::llm::types::request_v2::MessageRequest;
|
||||
@@ -26,6 +25,7 @@ use crate::llm::types::response_v2::{
|
||||
MessageResponse, PartialMessageResponse, PartialUsage, StopReason, StreamEvent,
|
||||
};
|
||||
use crate::llm::types::usage::Usage;
|
||||
use crate::llm::{LlmProvider, ProviderCapabilities, ProviderFeatures};
|
||||
|
||||
/// Anthropic Provider 默认 `max_tokens` 兜底值。
|
||||
///
|
||||
|
||||
@@ -11,10 +11,10 @@ use futures_core::Stream;
|
||||
use reqwest::Client;
|
||||
|
||||
use super::openai::GenericOpenaiProvider;
|
||||
use crate::llm::{LlmProvider, ProviderCapabilities};
|
||||
use crate::llm::error::LlmError;
|
||||
use crate::llm::types::request_v2::MessageRequest;
|
||||
use crate::llm::types::response_v2::{MessageResponse, StreamEvent};
|
||||
use crate::llm::{LlmProvider, ProviderCapabilities};
|
||||
|
||||
/// Ollama 本地 Provider —— OpenAI-compatible 协议的 newtype 包装。
|
||||
///
|
||||
|
||||
@@ -21,7 +21,6 @@ use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
use crate::llm::{LlmProvider, ProviderCapabilities, ProviderFeatures};
|
||||
use crate::llm::convert::{from_openai, to_openai};
|
||||
use crate::llm::error::LlmError;
|
||||
use crate::llm::types::message::{ContentBlock, ContentBlockType, Message};
|
||||
@@ -32,6 +31,7 @@ use crate::llm::types::response_v2::{
|
||||
};
|
||||
use crate::llm::types::shared::{FinishReason, ResponseFormat, ServiceTier, StopSequence};
|
||||
use crate::llm::types::tool::{OpenaiToolCall, OpenaiToolDefinition, ToolChoice};
|
||||
use crate::llm::{LlmProvider, ProviderCapabilities, ProviderFeatures};
|
||||
use serde::Deserialize;
|
||||
|
||||
// =============================================================================
|
||||
|
||||
@@ -2,9 +2,9 @@
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::llm::error::LlmError;
|
||||
use crate::llm::provider::{ProviderConfig, ProviderType, create_provider};
|
||||
use crate::llm::LlmProvider;
|
||||
|
||||
/// Provider 注册表 —— 管理多个 LLM Provider 实例。
|
||||
///
|
||||
|
||||
@@ -232,10 +232,16 @@ mod tests {
|
||||
let decoded: MessageRequest = serde_json::from_str(&json).expect("deserialize");
|
||||
assert_eq!(decoded.tools.len(), 1);
|
||||
assert_eq!(decoded.tools[0].name, "add");
|
||||
assert_eq!(decoded.tools[0].description.as_deref(), Some("add two numbers"));
|
||||
assert_eq!(
|
||||
decoded.tools[0].description.as_deref(),
|
||||
Some("add two numbers")
|
||||
);
|
||||
assert_eq!(decoded.tools[0].parameters, params);
|
||||
|
||||
// 验证序列化 JSON 不含 ToolDef 没有的字段(如 strict),保持 wire-format 兼容
|
||||
assert!(!json.contains("strict"), "ToolDef 序列化不应包含 strict 字段");
|
||||
assert!(
|
||||
!json.contains("strict"),
|
||||
"ToolDef 序列化不应包含 strict 字段"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -372,7 +372,8 @@ impl PartialMessageResponse {
|
||||
}
|
||||
// 元事件:不参与内容块累积,不修改 partial 状态
|
||||
//(Phase 9 —— 工具执行透明化,由 run_tool_loop 在工具前后插入)
|
||||
StreamEvent::ToolExecutionStarted { .. } | StreamEvent::ToolExecutionCompleted { .. } => true,
|
||||
StreamEvent::ToolExecutionStarted { .. }
|
||||
| StreamEvent::ToolExecutionCompleted { .. } => true,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -63,9 +63,7 @@ impl CostTracker {
|
||||
|
||||
impl From<Usage> for CostTracker {
|
||||
fn from(usage: Usage) -> Self {
|
||||
CostTracker {
|
||||
accumulated: usage,
|
||||
}
|
||||
CostTracker { accumulated: usage }
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+4
-2
@@ -13,12 +13,14 @@ pub mod vector_store;
|
||||
// 高频类型(大多数下游需要)
|
||||
pub use conversation::{ConversationMemory, ConversationMemoryConfig};
|
||||
pub use error::MemoryError;
|
||||
pub use graph::{GraphEntity, GraphRelation, InMemoryGraph, KnowledgeGraph, RelationDirection, ScoredEntity};
|
||||
pub use graph::{
|
||||
GraphEntity, GraphRelation, InMemoryGraph, KnowledgeGraph, RelationDirection, ScoredEntity,
|
||||
};
|
||||
pub use knowledge::KnowledgeStore;
|
||||
pub use retriever::MemoryRetriever;
|
||||
pub use store::{InMemoryStore, MemoryStore};
|
||||
#[cfg(feature = "memory-sqlite")]
|
||||
pub use store::SqliteStore;
|
||||
pub use store::{InMemoryStore, MemoryStore};
|
||||
#[allow(deprecated)]
|
||||
pub use vector::{InMemoryVectorRetriever, VectorRetriever};
|
||||
pub use vector_store::{InMemoryVectorStore, PersistentVectorStore, RagPipeline, VectorStore};
|
||||
|
||||
+75
-79
@@ -35,7 +35,11 @@ pub struct GraphEntity {
|
||||
|
||||
impl GraphEntity {
|
||||
/// 创建一个最小实体(仅 id + name + type,其余为空)。
|
||||
pub fn new(id: impl Into<String>, name: impl Into<String>, entity_type: impl Into<String>) -> Self {
|
||||
pub fn new(
|
||||
id: impl Into<String>,
|
||||
name: impl Into<String>,
|
||||
entity_type: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
id: id.into(),
|
||||
name: name.into(),
|
||||
@@ -81,7 +85,10 @@ impl GraphRelation {
|
||||
|
||||
/// 复合键:`source_id:target_id:relation_type`,用于去重和查找。
|
||||
pub fn composite_key(&self) -> String {
|
||||
format!("{}:{}:{}", self.source_id, self.target_id, self.relation_type)
|
||||
format!(
|
||||
"{}:{}:{}",
|
||||
self.source_id, self.target_id, self.relation_type
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -253,9 +260,10 @@ impl KnowledgeGraph for InMemoryGraph {
|
||||
if entity.id.is_empty() {
|
||||
return Err(MemoryError::InvalidInput("entity id is empty".into()));
|
||||
}
|
||||
let mut inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
let mut inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
// upsert:若已存在,先收集旧标签用于清理反向引用(避免同时借用 entities 和 tag_index)
|
||||
let old_tags: Vec<String> = inner
|
||||
.entities
|
||||
@@ -283,16 +291,18 @@ impl KnowledgeGraph for InMemoryGraph {
|
||||
}
|
||||
|
||||
async fn get_entity(&self, id: &str) -> Result<Option<GraphEntity>, MemoryError> {
|
||||
let inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
let inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
Ok(inner.entities.get(id).cloned())
|
||||
}
|
||||
|
||||
async fn remove_entity(&self, id: &str) -> Result<(), MemoryError> {
|
||||
let mut inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
let mut inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
// 移除实体并清理其标签反向引用
|
||||
if let Some(entity) = inner.entities.remove(id) {
|
||||
for tag in &entity.tags {
|
||||
@@ -312,9 +322,10 @@ impl KnowledgeGraph for InMemoryGraph {
|
||||
}
|
||||
|
||||
async fn add_relation(&self, relation: GraphRelation) -> Result<(), MemoryError> {
|
||||
let mut inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
let mut inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
// 校验两端实体存在
|
||||
if !inner.entities.contains_key(&relation.source_id) {
|
||||
return Err(MemoryError::InvalidInput(format!(
|
||||
@@ -348,11 +359,14 @@ impl KnowledgeGraph for InMemoryGraph {
|
||||
target_id: &str,
|
||||
relation_type: &str,
|
||||
) -> Result<(), MemoryError> {
|
||||
let mut inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
let mut inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
inner.relations.retain(|r| {
|
||||
!(r.source_id == source_id && r.target_id == target_id && r.relation_type == relation_type)
|
||||
!(r.source_id == source_id
|
||||
&& r.target_id == target_id
|
||||
&& r.relation_type == relation_type)
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
@@ -364,9 +378,10 @@ impl KnowledgeGraph for InMemoryGraph {
|
||||
direction: RelationDirection,
|
||||
relation_types: Option<&[&str]>,
|
||||
) -> Result<Vec<ScoredEntity>, MemoryError> {
|
||||
let inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
let inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
|
||||
// 1. 验证起点存在
|
||||
if !inner.entities.contains_key(entity_id) {
|
||||
@@ -470,9 +485,10 @@ impl KnowledgeGraph for InMemoryGraph {
|
||||
}
|
||||
|
||||
async fn find_by_keywords(&self, keywords: &[String]) -> Result<Vec<GraphEntity>, MemoryError> {
|
||||
let inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
let inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
if keywords.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
@@ -495,9 +511,10 @@ impl KnowledgeGraph for InMemoryGraph {
|
||||
}
|
||||
|
||||
async fn find_tags(&self, prefix: &str) -> Result<Vec<String>, MemoryError> {
|
||||
let inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
let inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
let prefix_l = prefix.to_lowercase();
|
||||
let mut tags: Vec<String> = inner
|
||||
.tag_index
|
||||
@@ -514,9 +531,10 @@ impl KnowledgeGraph for InMemoryGraph {
|
||||
entity_id: &str,
|
||||
tags: Vec<String>,
|
||||
) -> Result<usize, MemoryError> {
|
||||
let mut inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
let mut inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
|
||||
// 先收集旧标签(避免同时借用 entities 和 tag_index)
|
||||
let old_tags: Vec<String> = inner
|
||||
@@ -554,13 +572,11 @@ impl KnowledgeGraph for InMemoryGraph {
|
||||
}
|
||||
|
||||
async fn entity_count_by_tag(&self, tag: &str) -> Result<usize, MemoryError> {
|
||||
let inner = self.inner.lock().map_err(|e| {
|
||||
MemoryError::RetrievalError(format!("lock poisoned: {e}"))
|
||||
})?;
|
||||
Ok(inner
|
||||
.tag_index
|
||||
.get(tag)
|
||||
.map_or(0, |ids| ids.len()))
|
||||
let inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?;
|
||||
Ok(inner.tag_index.get(tag).map_or(0, |ids| ids.len()))
|
||||
}
|
||||
|
||||
fn tag_constraints(&self) -> TagConstraints {
|
||||
@@ -660,10 +676,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn add_relation_validates_entities() {
|
||||
let graph = InMemoryGraph::new();
|
||||
graph
|
||||
.add_entity(make_entity("a", "A", "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
graph.add_entity(make_entity("a", "A", "x")).await.unwrap();
|
||||
// target 不存在
|
||||
let result = graph
|
||||
.add_relation(GraphRelation::new("a", "ghost", "r", 0.5))
|
||||
@@ -674,14 +687,8 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn add_relation_upsert_overrides_weight() {
|
||||
let graph = InMemoryGraph::new();
|
||||
graph
|
||||
.add_entity(make_entity("a", "A", "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
graph
|
||||
.add_entity(make_entity("b", "B", "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
graph.add_entity(make_entity("a", "A", "x")).await.unwrap();
|
||||
graph.add_entity(make_entity("b", "B", "x")).await.unwrap();
|
||||
graph
|
||||
.add_relation(GraphRelation::new("a", "b", "r", 0.5))
|
||||
.await
|
||||
@@ -701,14 +708,8 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn remove_relation() {
|
||||
let graph = InMemoryGraph::new();
|
||||
graph
|
||||
.add_entity(make_entity("a", "A", "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
graph
|
||||
.add_entity(make_entity("b", "B", "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
graph.add_entity(make_entity("a", "A", "x")).await.unwrap();
|
||||
graph.add_entity(make_entity("b", "B", "x")).await.unwrap();
|
||||
graph
|
||||
.add_relation(GraphRelation::new("a", "b", "r", 0.5))
|
||||
.await
|
||||
@@ -726,10 +727,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn get_related_empty_graph() {
|
||||
let graph = InMemoryGraph::new();
|
||||
graph
|
||||
.add_entity(make_entity("a", "A", "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
graph.add_entity(make_entity("a", "A", "x")).await.unwrap();
|
||||
let related = graph
|
||||
.get_related("a", 3, RelationDirection::Outgoing, None)
|
||||
.await
|
||||
@@ -740,10 +738,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn get_related_depth_0() {
|
||||
let graph = InMemoryGraph::new();
|
||||
graph
|
||||
.add_entity(make_entity("a", "A", "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
graph.add_entity(make_entity("a", "A", "x")).await.unwrap();
|
||||
let related = graph
|
||||
.get_related("a", 0, RelationDirection::Outgoing, None)
|
||||
.await
|
||||
@@ -853,10 +848,16 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn get_related_star_topology() {
|
||||
let graph = InMemoryGraph::new();
|
||||
graph.add_entity(make_entity("center", "C", "x")).await.unwrap();
|
||||
graph
|
||||
.add_entity(make_entity("center", "C", "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
for i in 0..5 {
|
||||
let leaf = format!("leaf{i}");
|
||||
graph.add_entity(make_entity(&leaf, &leaf, "x")).await.unwrap();
|
||||
graph
|
||||
.add_entity(make_entity(&leaf, &leaf, "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
graph
|
||||
.add_relation(GraphRelation::new("center", &leaf, "r", 0.8))
|
||||
.await
|
||||
@@ -1016,10 +1017,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn set_and_find_tags() {
|
||||
let graph = InMemoryGraph::new();
|
||||
graph
|
||||
.add_entity(make_entity("a", "A", "x"))
|
||||
.await
|
||||
.unwrap();
|
||||
graph.add_entity(make_entity("a", "A", "x")).await.unwrap();
|
||||
let n = graph
|
||||
.set_entity_tags("a", vec!["rust".into(), "ai".into()])
|
||||
.await
|
||||
@@ -1055,7 +1053,10 @@ mod tests {
|
||||
let graph = InMemoryGraph::with_constraints(constraints);
|
||||
graph.add_entity(make_entity("a", "A", "x")).await.unwrap();
|
||||
let n = graph
|
||||
.set_entity_tags("a", vec!["t1".into(), "t2".into(), "t3".into(), "t4".into()])
|
||||
.set_entity_tags(
|
||||
"a",
|
||||
vec!["t1".into(), "t2".into(), "t3".into(), "t4".into()],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(n, 2, "should truncate to max_tags_per_entity");
|
||||
@@ -1075,10 +1076,7 @@ mod tests {
|
||||
.set_entity_tags("b", vec!["rust".into(), "ai".into()])
|
||||
.await
|
||||
.unwrap();
|
||||
graph
|
||||
.set_entity_tags("c", vec!["ai".into()])
|
||||
.await
|
||||
.unwrap();
|
||||
graph.set_entity_tags("c", vec!["ai".into()]).await.unwrap();
|
||||
assert_eq!(graph.entity_count_by_tag("rust").await.unwrap(), 2);
|
||||
assert_eq!(graph.entity_count_by_tag("ai").await.unwrap(), 2);
|
||||
}
|
||||
@@ -1086,9 +1084,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn set_entity_tags_not_found() {
|
||||
let graph = InMemoryGraph::new();
|
||||
let result = graph
|
||||
.set_entity_tags("ghost", vec!["x".into()])
|
||||
.await;
|
||||
let result = graph.set_entity_tags("ghost", vec!["x".into()]).await;
|
||||
assert!(matches!(result, Err(MemoryError::NotFound(_))));
|
||||
}
|
||||
|
||||
|
||||
+33
-12
@@ -270,7 +270,12 @@ impl MemoryRetriever {
|
||||
}
|
||||
// BFS 找相关实体
|
||||
let related: Vec<ScoredEntity> = graph
|
||||
.get_related(&start.id, depth, crate::memory::graph::RelationDirection::Both, None)
|
||||
.get_related(
|
||||
&start.id,
|
||||
depth,
|
||||
crate::memory::graph::RelationDirection::Both,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
for se in related {
|
||||
if seen.insert(se.entity.id.clone()) {
|
||||
@@ -361,7 +366,7 @@ fn default_stop_words() -> HashSet<String> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::memory::graph::{GraphEntity, InMemoryGraph, GraphRelation};
|
||||
use crate::memory::graph::{GraphEntity, GraphRelation, InMemoryGraph};
|
||||
use crate::memory::knowledge::KnowledgeStore;
|
||||
use crate::memory::{InMemoryStore, MemoryStore};
|
||||
use std::sync::Arc;
|
||||
@@ -502,7 +507,12 @@ mod tests {
|
||||
graph.add_entity(langgraph).await.unwrap();
|
||||
graph.add_entity(python).await.unwrap();
|
||||
graph
|
||||
.add_relation(GraphRelation::new("langchain", "langgraph", "includes", 0.8))
|
||||
.add_relation(GraphRelation::new(
|
||||
"langchain",
|
||||
"langgraph",
|
||||
"includes",
|
||||
0.8,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
graph
|
||||
@@ -523,7 +533,12 @@ mod tests {
|
||||
let result = retriever.retrieve("langchain").await.unwrap();
|
||||
assert_eq!(result.strategy, RetrievalStrategy::GraphOnly);
|
||||
// 应该全部是 GraphEntity 变体
|
||||
assert!(result.items.iter().all(|i| matches!(i, RetrievalItem::GraphEntity { .. })));
|
||||
assert!(
|
||||
result
|
||||
.items
|
||||
.iter()
|
||||
.all(|i| matches!(i, RetrievalItem::GraphEntity { .. }))
|
||||
);
|
||||
// langchain 是起始实体(score=1.0),langgraph 和 python 是 BFS 结果
|
||||
assert!(!result.items.is_empty(), "should find graph entities");
|
||||
}
|
||||
@@ -544,8 +559,8 @@ mod tests {
|
||||
// KnowledgeGraph 也有匹配
|
||||
let graph = make_graph_with_data().await;
|
||||
|
||||
let retriever = MemoryRetriever::new(ks, RetrieverConfig::default())
|
||||
.with_knowledge_graph(graph);
|
||||
let retriever =
|
||||
MemoryRetriever::new(ks, RetrieverConfig::default()).with_knowledge_graph(graph);
|
||||
let result = retriever.retrieve("langchain").await.unwrap();
|
||||
assert_eq!(result.strategy, RetrievalStrategy::Hybrid);
|
||||
// 应该同时包含 KnowledgePage 和 GraphEntity
|
||||
@@ -569,8 +584,8 @@ mod tests {
|
||||
.unwrap();
|
||||
// 空图
|
||||
let graph = Arc::new(InMemoryGraph::new());
|
||||
let retriever = MemoryRetriever::new(ks, RetrieverConfig::default())
|
||||
.with_knowledge_graph(graph);
|
||||
let retriever =
|
||||
MemoryRetriever::new(ks, RetrieverConfig::default()).with_knowledge_graph(graph);
|
||||
let result = retriever.retrieve("langchain").await.unwrap();
|
||||
assert_eq!(result.strategy, RetrievalStrategy::Hybrid);
|
||||
// 图空,只有 Store 结果
|
||||
@@ -587,15 +602,18 @@ mod tests {
|
||||
let ks = KnowledgeStore::new(store);
|
||||
// Store 空,图有数据
|
||||
let graph = make_graph_with_data().await;
|
||||
let retriever = MemoryRetriever::new(ks, RetrieverConfig::default())
|
||||
.with_knowledge_graph(graph);
|
||||
let retriever =
|
||||
MemoryRetriever::new(ks, RetrieverConfig::default()).with_knowledge_graph(graph);
|
||||
let result = retriever.retrieve("langchain").await.unwrap();
|
||||
assert_eq!(result.strategy, RetrievalStrategy::Hybrid);
|
||||
let has_entity = result
|
||||
.items
|
||||
.iter()
|
||||
.any(|i| matches!(i, RetrievalItem::GraphEntity { .. }));
|
||||
assert!(has_entity, "should have graph results even when store is empty");
|
||||
assert!(
|
||||
has_entity,
|
||||
"should have graph results even when store is empty"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -615,7 +633,10 @@ mod tests {
|
||||
.items
|
||||
.iter()
|
||||
.any(|i| matches!(i, RetrievalItem::GraphEntity { .. }));
|
||||
assert!(!has_entity, "KnowledgeOnly should not return graph entities");
|
||||
assert!(
|
||||
!has_entity,
|
||||
"KnowledgeOnly should not return graph entities"
|
||||
);
|
||||
}
|
||||
|
||||
// ── 辅助函数测试 ──
|
||||
|
||||
@@ -6,9 +6,9 @@ use std::path::Path;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use rusqlite::{params, params_from_iter, Connection, ErrorCode};
|
||||
use time::format_description::well_known::Rfc3339;
|
||||
use rusqlite::{Connection, ErrorCode, params, params_from_iter};
|
||||
use time::OffsetDateTime;
|
||||
use time::format_description::well_known::Rfc3339;
|
||||
use tracing::{debug, error, instrument, warn};
|
||||
|
||||
use crate::memory::error::MemoryError;
|
||||
@@ -114,9 +114,7 @@ impl MemoryStore for SqliteStore {
|
||||
.map_err(|e| map_sqlite_error(e, "query get"))?;
|
||||
match rows.next() {
|
||||
None => Ok(None),
|
||||
Some(row) => row
|
||||
.map(Some)
|
||||
.map_err(|e| map_sqlite_error(e, "decode row")),
|
||||
Some(row) => row.map(Some).map_err(|e| map_sqlite_error(e, "decode row")),
|
||||
}
|
||||
})
|
||||
.await
|
||||
@@ -130,11 +128,8 @@ impl MemoryStore for SqliteStore {
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<(), MemoryError> {
|
||||
let conn = conn.lock().unwrap_or_else(|e| e.into_inner());
|
||||
conn.execute(
|
||||
"DELETE FROM memory_items WHERE id = ?1",
|
||||
params![id_owned],
|
||||
)
|
||||
.map_err(|e| map_sqlite_error(e, "delete"))?;
|
||||
conn.execute("DELETE FROM memory_items WHERE id = ?1", params![id_owned])
|
||||
.map_err(|e| map_sqlite_error(e, "delete"))?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
@@ -143,9 +138,8 @@ impl MemoryStore for SqliteStore {
|
||||
|
||||
#[instrument(skip(self, filter))]
|
||||
async fn list(&self, filter: &MemoryFilter) -> Result<Vec<MemoryItem>, MemoryError> {
|
||||
let mut sql = String::from(
|
||||
"SELECT id, content, metadata, created_at FROM memory_items WHERE 1=1",
|
||||
);
|
||||
let mut sql =
|
||||
String::from("SELECT id, content, metadata, created_at FROM memory_items WHERE 1=1");
|
||||
let mut param_values: Vec<String> = Vec::new();
|
||||
let mut ph_idx = 0usize;
|
||||
|
||||
@@ -309,9 +303,7 @@ fn map_sqlite_error(e: rusqlite::Error, ctx: &str) -> MemoryError {
|
||||
rusqlite::Error::InvalidQuery
|
||||
| rusqlite::Error::InvalidParameterName(_)
|
||||
| rusqlite::Error::InvalidColumnIndex(_)
|
||||
| rusqlite::Error::InvalidColumnName(_) => {
|
||||
MemoryError::InvalidInput(format!("{ctx}: {e}"))
|
||||
}
|
||||
| rusqlite::Error::InvalidColumnName(_) => MemoryError::InvalidInput(format!("{ctx}: {e}")),
|
||||
rusqlite::Error::FromSqlConversionFailure(_, _, _)
|
||||
| rusqlite::Error::ToSqlConversionFailure(_) => {
|
||||
MemoryError::Serialization(format!("{ctx}: {e}"))
|
||||
@@ -527,8 +519,7 @@ mod tests {
|
||||
// ponytail: 回归验证 SqliteStore 可作为 Arc<dyn MemoryStore> 与 InMemoryStore 互换
|
||||
// 所有现有消费者(Conversation / Knowledge / Retriever / SessionMemory)均通过 trait object 引用,
|
||||
// 此测试确保 trait 接口契约在 SqliteStore 上同样成立。
|
||||
let sqlite: Arc<dyn MemoryStore> =
|
||||
Arc::new(SqliteStore::open(":memory:").unwrap());
|
||||
let sqlite: Arc<dyn MemoryStore> = Arc::new(SqliteStore::open(":memory:").unwrap());
|
||||
let in_mem: Arc<dyn MemoryStore> = Arc::new(InMemoryStore::new());
|
||||
|
||||
let stores: Vec<Arc<dyn MemoryStore>> = vec![Arc::clone(&sqlite), Arc::clone(&in_mem)];
|
||||
@@ -556,12 +547,7 @@ mod tests {
|
||||
handles.push(tokio::spawn(async move {
|
||||
let id = format!("concurrent_{i}");
|
||||
// 设置每次 save 的 per-call timeout —— busy_timeout=5000ms 应足够
|
||||
match tokio::time::timeout(
|
||||
Duration::from_secs(10),
|
||||
s.save(make_item(&id)),
|
||||
)
|
||||
.await
|
||||
{
|
||||
match tokio::time::timeout(Duration::from_secs(10), s.save(make_item(&id))).await {
|
||||
Ok(res) => res.unwrap(),
|
||||
Err(_) => panic!("save({id}) timed out under 100-way concurrency"),
|
||||
}
|
||||
|
||||
+3
-11
@@ -30,11 +30,7 @@ pub trait VectorRetriever: Send + Sync {
|
||||
///
|
||||
/// 返回 `Vec<(id, score)>`,按 score 降序排列,score ∈ [0.0, 1.0]
|
||||
/// (余弦相似度)。当 `k == 0`、索引为空或 query 为零向量时返回空 Vec。
|
||||
async fn search(
|
||||
&self,
|
||||
query: Vec<f32>,
|
||||
k: usize,
|
||||
) -> Result<Vec<(String, f32)>, MemoryError>;
|
||||
async fn search(&self, query: Vec<f32>, k: usize) -> Result<Vec<(String, f32)>, MemoryError>;
|
||||
}
|
||||
|
||||
/// 进程内向量检索器 —— 基于 HashMap + 全量余弦相似度扫描。
|
||||
@@ -79,11 +75,7 @@ impl VectorRetriever for InMemoryVectorRetriever {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn search(
|
||||
&self,
|
||||
query: Vec<f32>,
|
||||
k: usize,
|
||||
) -> Result<Vec<(String, f32)>, MemoryError> {
|
||||
async fn search(&self, query: Vec<f32>, k: usize) -> Result<Vec<(String, f32)>, MemoryError> {
|
||||
let vectors = self
|
||||
.vectors
|
||||
.lock()
|
||||
@@ -241,4 +233,4 @@ mod tests {
|
||||
h.await.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+26
-46
@@ -11,8 +11,8 @@ use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use time::format_description::well_known::Rfc3339;
|
||||
use time::OffsetDateTime;
|
||||
use time::format_description::well_known::Rfc3339;
|
||||
use tracing::{debug, info};
|
||||
|
||||
use crate::document::{Document, RecursiveCharacterSplitter};
|
||||
@@ -43,11 +43,8 @@ pub trait VectorStore: Send + Sync {
|
||||
/// - 截取 `min(len)` 对处理(部分写入已发生)
|
||||
/// - 返回 `Err(MemoryError::InvalidInput)` 告知截断
|
||||
/// - 调用方可以 `let _ = store.add(...)` 忽略错误
|
||||
async fn add(
|
||||
&self,
|
||||
documents: &[Document],
|
||||
embeddings: &[Vec<f32>],
|
||||
) -> Result<(), MemoryError>;
|
||||
async fn add(&self, documents: &[Document], embeddings: &[Vec<f32>])
|
||||
-> Result<(), MemoryError>;
|
||||
|
||||
/// 检索与 `query` 向量最相似的 `k` 条记录。
|
||||
///
|
||||
@@ -59,11 +56,7 @@ pub trait VectorStore: Send + Sync {
|
||||
/// - 空索引 → 返回 `vec![]`
|
||||
/// - `k == 0` → 返回 `vec![]`
|
||||
/// - 零向量(norm ≈ 0)→ 返回 `vec![]`
|
||||
async fn search(
|
||||
&self,
|
||||
query: &[f32],
|
||||
k: usize,
|
||||
) -> Result<Vec<(Document, f32)>, MemoryError>;
|
||||
async fn search(&self, query: &[f32], k: usize) -> Result<Vec<(Document, f32)>, MemoryError>;
|
||||
|
||||
/// 批量删除文档(幂等)。
|
||||
///
|
||||
@@ -105,9 +98,7 @@ impl InMemoryVectorStore {
|
||||
}
|
||||
|
||||
/// 从预填充的 entries 构造(供 `PersistentVectorStore` 使用)。
|
||||
pub(crate) fn with_entries(
|
||||
entries: HashMap<String, (Document, Vec<f32>)>,
|
||||
) -> Self {
|
||||
pub(crate) fn with_entries(entries: HashMap<String, (Document, Vec<f32>)>) -> Self {
|
||||
Self {
|
||||
entries: Mutex::new(entries),
|
||||
}
|
||||
@@ -142,7 +133,10 @@ impl VectorStore for InMemoryVectorStore {
|
||||
}
|
||||
|
||||
for i in 0..n {
|
||||
entries.insert(documents[i].id.clone(), (documents[i].clone(), embeddings[i].clone()));
|
||||
entries.insert(
|
||||
documents[i].id.clone(),
|
||||
(documents[i].clone(), embeddings[i].clone()),
|
||||
);
|
||||
}
|
||||
|
||||
if documents.len() != embeddings.len() {
|
||||
@@ -156,11 +150,7 @@ impl VectorStore for InMemoryVectorStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn search(
|
||||
&self,
|
||||
query: &[f32],
|
||||
k: usize,
|
||||
) -> Result<Vec<(Document, f32)>, MemoryError> {
|
||||
async fn search(&self, query: &[f32], k: usize) -> Result<Vec<(Document, f32)>, MemoryError> {
|
||||
if k == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
@@ -265,10 +255,7 @@ impl PersistentVectorStore {
|
||||
///
|
||||
/// `MemoryStore::list()` 由 `SqliteStore` 内部使用 `spawn_blocking` 卸载,
|
||||
/// 加载过程本身在 async context 中即可,无需额外 spawn_blocking。
|
||||
pub async fn new(
|
||||
store: Arc<dyn MemoryStore>,
|
||||
namespace: &str,
|
||||
) -> Result<Self, MemoryError> {
|
||||
pub async fn new(store: Arc<dyn MemoryStore>, namespace: &str) -> Result<Self, MemoryError> {
|
||||
let prefix = format!("vec:{namespace}:");
|
||||
let filter = MemoryFilter {
|
||||
prefix: Some(prefix.clone()),
|
||||
@@ -292,7 +279,10 @@ impl PersistentVectorStore {
|
||||
entries.insert(doc.id.clone(), (doc, entry.embedding));
|
||||
}
|
||||
|
||||
info!(entries = entries.len(), "PersistentVectorStore — 内存索引重建完成");
|
||||
info!(
|
||||
entries = entries.len(),
|
||||
"PersistentVectorStore — 内存索引重建完成"
|
||||
);
|
||||
Ok(Self {
|
||||
inner: InMemoryVectorStore::with_entries(entries),
|
||||
store,
|
||||
@@ -340,11 +330,7 @@ impl VectorStore for PersistentVectorStore {
|
||||
self.inner.add(documents, embeddings).await
|
||||
}
|
||||
|
||||
async fn search(
|
||||
&self,
|
||||
query: &[f32],
|
||||
k: usize,
|
||||
) -> Result<Vec<(Document, f32)>, MemoryError> {
|
||||
async fn search(&self, query: &[f32], k: usize) -> Result<Vec<(Document, f32)>, MemoryError> {
|
||||
tracing::trace!(k, "PersistentVectorStore::search");
|
||||
self.inner.search(query, k).await
|
||||
}
|
||||
@@ -575,10 +561,7 @@ mod tests {
|
||||
async fn remove_items() {
|
||||
let store = InMemoryVectorStore::new();
|
||||
let docs = vec![make_doc("a", "alpha"), make_doc("b", "beta")];
|
||||
let embeddings = vec![
|
||||
make_vec(&[1.0, 0.0, 0.0]),
|
||||
make_vec(&[0.0, 1.0, 0.0]),
|
||||
];
|
||||
let embeddings = vec![make_vec(&[1.0, 0.0, 0.0]), make_vec(&[0.0, 1.0, 0.0])];
|
||||
store.add(&docs, &embeddings).await.unwrap();
|
||||
|
||||
store.remove(&["a".to_string()]).await.unwrap();
|
||||
@@ -599,7 +582,9 @@ mod tests {
|
||||
let embeddings = vec![make_vec(&[1.0, 0.0, 0.0])];
|
||||
store.add(&docs, &embeddings).await.unwrap();
|
||||
|
||||
let result = store.remove(&["nonexistent".to_string(), "also_nonexistent".to_string()]).await;
|
||||
let result = store
|
||||
.remove(&["nonexistent".to_string(), "also_nonexistent".to_string()])
|
||||
.await;
|
||||
assert!(result.is_ok(), "批量删除不存在 id 不应报错");
|
||||
|
||||
let results = store.search(&[1.0, 0.0, 0.0], 5).await.unwrap();
|
||||
@@ -657,7 +642,9 @@ mod tests {
|
||||
backend: Arc<dyn MemoryStore>,
|
||||
namespace: &str,
|
||||
) -> PersistentVectorStore {
|
||||
PersistentVectorStore::new(backend, namespace).await.unwrap()
|
||||
PersistentVectorStore::new(backend, namespace)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -755,10 +742,7 @@ mod tests {
|
||||
let store = make_persistent(Arc::clone(&backend), "default").await;
|
||||
|
||||
// 正常写入前 2 条
|
||||
let docs_first = vec![
|
||||
make_doc("doc_0", "first"),
|
||||
make_doc("doc_1", "second"),
|
||||
];
|
||||
let docs_first = vec![make_doc("doc_0", "first"), make_doc("doc_1", "second")];
|
||||
let embeddings_first = vec![make_vec(&[1.0, 0.0, 0.0]), make_vec(&[0.0, 1.0, 0.0])];
|
||||
store.add(&docs_first, &embeddings_first).await.unwrap();
|
||||
|
||||
@@ -797,11 +781,7 @@ mod tests {
|
||||
let store: Arc<dyn VectorStore> = Arc::new(InMemoryVectorStore::new());
|
||||
let splitter = RecursiveCharacterSplitter::new(50, 5);
|
||||
|
||||
let pipeline = RagPipeline::new(
|
||||
Arc::clone(&embedder),
|
||||
Arc::clone(&store),
|
||||
Some(splitter),
|
||||
);
|
||||
let pipeline = RagPipeline::new(Arc::clone(&embedder), Arc::clone(&store), Some(splitter));
|
||||
|
||||
// 创建多段落文档
|
||||
let doc = Document::new(
|
||||
@@ -934,4 +914,4 @@ mod tests {
|
||||
elapsed.as_millis()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user