xt/agcore
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

docs: 更新 README feature 表 + 升级指南 + 示例注释 + roadmap 同步

Browse Source 
- 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:
徐涛
2026-07-19 08:18:04 +08:00
parent bc4eac72e1
commit 5baa170508
53 changed files with 1112 additions and 615 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent.rs
+2 -2
View File
@@ -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};
src/agent/builder.rs
+1 -1
View File
@@ -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;
src/agent/context.rs
+136 -44
View File
@@ -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)
}
}
}
src/agent/runtime.rs
+1 -1
View File
@@ -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;
src/agent/session.rs
+74 -65
View File
@@ -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 应永不触发");
}
}
}
src/agent/session_memory.rs
+2 -1
View File
@@ -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）。
src/agent/summary.rs
+5 -1
View File
@@ -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));
src/document.rs
+64 -14
View File
@@ -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")
);
}
}
src/engine/checkpointer.rs
+6 -13
View File
@@ -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");
}
}
}
src/engine/error.rs
+1 -1
View File
@@ -48,4 +48,4 @@ pub enum EngineError {
/// 调用方收到此错误时，子 session 已通过 `destroy()` 清理（SessionMeta + checkpoint 全部清空）。
#[error("Dispatch failed: {0}")]
DispatchFailed(String),
}
}
src/engine/mod.rs
+1 -1
View File
@@ -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};
src/engine/session_manager.rs
+47 -33
View File
@@ -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());
}
}
}
}
src/engine/snapshot.rs
+1 -1
View File
@@ -36,4 +36,4 @@ pub struct SessionSnapshot {
pub last_summary_turn: Option<u32>,
#[serde(default)]
pub session_memory_data: std::collections::HashMap<String, SessionMemoryEntry>,
}
}
src/engine/sub_agent.rs
+23 -22
View File
@@ -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(),
src/engine/switch.rs
+2 -6
View File
@@ -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();
src/lib.rs
+8 -8
View File
@@ -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;
src/llm.rs
+2 -2
View File
@@ -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",
src/llm/cycle.rs
+103 -78
View File
@@ -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 之前已发出
src/llm/embedding.rs
+11 -6
View File
@@ -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
);
}
}
src/llm/provider/anthropic.rs
+1 -1
View File
@@ -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` 兜底值。
///
src/llm/provider/ollama.rs
+1 -1
View File
@@ -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 包装。
///
src/llm/provider/openai.rs
+1 -1
View File
@@ -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;
// =============================================================================
src/llm/provider/registry.rs
+1 -1
View File
@@ -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 实例。
///
src/llm/types/request_v2.rs
+8 -2
View File
@@ -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 字段"
);
}
}
src/llm/types/response_v2.rs
+2 -1
View File
@@ -372,7 +372,8 @@ impl PartialMessageResponse {
}
// 元事件：不参与内容块累积，不修改 partial 状态
//（Phase 9 —— 工具执行透明化，由 run_tool_loop 在工具前后插入）
StreamEvent::ToolExecutionStarted { .. } | StreamEvent::ToolExecutionCompleted { .. } => true,
StreamEvent::ToolExecutionStarted { .. }
| StreamEvent::ToolExecutionCompleted { .. } => true,
}
}
src/llm/types/usage.rs
+1 -3
View File
@@ -63,9 +63,7 @@ impl CostTracker {
impl From<Usage> for CostTracker {
fn from(usage: Usage) -> Self {
CostTracker {
accumulated: usage,
}
CostTracker { accumulated: usage }
}
}
src/memory.rs
+4 -2
View File
@@ -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};
src/memory/graph.rs
+75 -79
View File
@@ -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(_))));
}
src/memory/retriever.rs
+33 -12
View File
@@ -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"
);
}
// ── 辅助函数测试 ──
src/memory/store/sqlite_store.rs
+10 -24
View File
@@ -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"),
}
src/memory/vector.rs
+3 -11
View File
@@ -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();
}
}
}
}
src/memory/vector_store.rs
+26 -46
View File
@@ -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()
);
}
}
}