ooodc/PicoBot
1
0
Fork 0
Code Pull Requests Activity

feat: 更新内存维护逻辑，调整返回类型为Option以处理无结果情况

Browse Source
This commit is contained in:
ooodc 2026-05-12 21:04:17 +08:00
parent 03c95e6b8f
commit 5a0c018ee7
5 changed files with 95 additions and 86 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/gateway/agent_task_executor.rs
View File

@ -61,7 +61,7 @@ impl SchedulerMaintenanceService {
self.session_manager.cleanup_expired_sessions().await
}
async fn run_memory_maintenance(&self) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> {
async fn run_memory_maintenance(&self) -> Result<Option<MemoryMaintenanceScopeResult>, AgentError> {
self.session_manager.run_memory_maintenance_for_all_scopes().await
}
}

143
src/gateway/memory_maintenance.rs
View File

@ -64,7 +64,6 @@ pub(crate) struct MemoryOrganizationOutput {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct MemorySummaryInput {
pub(crate) scope_key: String,
pub(crate) organized_memories: Vec<OrganizedMemory>,
}
@ -217,6 +216,7 @@ impl MemoryMaintenanceService {
Ok(output)
}
#[cfg_attr(not(test), allow(dead_code))]
pub(crate) async fn run_for_scope(
&self,
scope_key: &str,
@ -264,7 +264,6 @@ impl MemoryMaintenanceService {
})?;
let input = MemorySummaryInput {
scope_key: scope_key.to_string(),
organized_memories: remaining_memories
.iter()
.map(|m| OrganizedMemory {
@ -353,15 +352,20 @@ impl MemoryMaintenanceService {
pub(crate) async fn run_for_all_scopes(
&self,
) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> {
) -> Result<Option<MemoryMaintenanceScopeResult>, AgentError> {
let scope_keys = self.store.list_memory_scope_keys("user").map_err(|err| {
AgentError::Other(format!("list memory scope keys error: {}", err))
})?;
let mut results = Vec::new();
for scope_key in scope_keys {
let result = match self.run_for_scope(&scope_key).await {
Ok(Some(result)) => result,
if scope_keys.is_empty() {
return Ok(None);
}
// 步骤1：逐个 scope 整理记忆（merge/delete），但不生成摘要
let mut all_outputs = Vec::new();
for scope_key in &scope_keys {
match self.run_organize_for_scope(scope_key).await {
Ok(Some(output)) => all_outputs.push((scope_key.clone(), output)),
Ok(None) => continue,
Err(error) if is_recoverable_maintenance_scope_error(&error) => {
tracing::warn!(
@ -372,23 +376,78 @@ impl MemoryMaintenanceService {
continue;
}
Err(error) => return Err(error),
};
results.push(result);
}
}
let combined_markdown = combine_managed_memory_markdown(
&results
if all_outputs.is_empty() {
return Ok(None);
}
// 步骤2：收集所有 scope 整理后的剩余记忆
let mut all_remaining_memories = Vec::new();
for (scope_key, _) in &all_outputs {
let memories = self
.store
.list_memories_for_scope("user", scope_key)
.map_err(|err| {
AgentError::Other(format!(
"list remaining memories for scope {} error: {}",
scope_key, err
))
})?;
all_remaining_memories.extend(memories);
}
// 步骤3：统一生成一个摘要
let managed_markdown = if all_remaining_memories.is_empty() {
String::new()
} else {
self.generate_summary("all", &all_remaining_memories).await?
};
if !managed_markdown.is_empty() {
upsert_managed_agent_memory_summary(&managed_markdown)?;
}
// 合并所有输出用于返回
let combined_output = MemoryOrganizationOutput {
merges: all_outputs
.iter()
.map(|result| result.managed_markdown.clone())
.collect::<Vec<_>>(),
);
.flat_map(|(_, o)| o.merges.clone())
.collect(),
conflicts: all_outputs
.iter()
.flat_map(|(_, o)| o.conflicts.clone())
.collect(),
low_value_ids: all_outputs
.iter()
.flat_map(|(_, o)| o.low_value_ids.clone())
.collect(),
};
if !combined_markdown.is_empty() {
upsert_managed_agent_memory_summary(&combined_markdown)?;
}
Ok(Some(MemoryMaintenanceScopeResult {
scope_key: "all".to_string(),
output: combined_output,
managed_markdown,
}))
}
Ok(results)
/// 仅执行整理步骤（organize + apply），不生成摘要
async fn run_organize_for_scope(
&self,
scope_key: &str,
) -> Result<Option<MemoryOrganizationOutput>, AgentError> {
let Some(plan) = self.build_plan_for_scope(scope_key)? else {
return Ok(None);
};
// 步骤1：整理记忆
let organize_output = self.organize_plan(scope_key, &plan).await?;
// 应用整理结果（merge和delete）
apply_memory_maintenance_output(self.store.as_ref(), scope_key, &plan, &organize_output)?;
Ok(Some(organize_output))
}
}
@ -512,47 +571,6 @@ pub(crate) fn extract_json_object(content: &str) -> Option<&str> {
None
}
pub(crate) fn combine_managed_memory_markdown(chunks: &[String]) -> String {
let normalized_chunks = chunks
.iter()
.map(|chunk| chunk.trim())
.filter(|chunk| !chunk.is_empty())
.collect::<Vec<_>>();
let mut combined = Vec::new();
for (index, chunk) in normalized_chunks.iter().enumerate() {
let chunk_lines = chunk
.lines()
.map(str::trim)
.filter(|line| !line.is_empty())
.collect::<HashSet<_>>();
let is_subset_of_other =
normalized_chunks
.iter()
.enumerate()
.any(|(other_index, other)| {
if index == other_index {
return false;
}
let other_lines = other
.lines()
.map(str::trim)
.filter(|line| !line.is_empty())
.collect::<HashSet<_>>();
chunk_lines.len() < other_lines.len() && chunk_lines.is_subset(&other_lines)
});
if !is_subset_of_other && !combined.iter().any(|existing: &String| existing == chunk) {
combined.push((*chunk).to_string());
}
}
combined.join("\n\n")
}
pub(crate) fn apply_memory_maintenance_output(
store: &SessionStore,
scope_key: &str,
@ -651,3 +669,8 @@ fn preview_text(content: &str, max_chars: usize) -> String {
}
preview.replace('\n', "\\n")
}
#[cfg(test)]
mod tests {
use super::*;
}

2
src/gateway/memory_maintenance_coordinator.rs
View File

@ -32,7 +32,7 @@ impl MemoryMaintenanceCoordinator {
pub(crate) async fn run_for_all_scopes(
&self,
) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> {
) -> Result<Option<MemoryMaintenanceScopeResult>, AgentError> {
self.service()?.run_for_all_scopes().await
}

1
src/gateway/memory_maintenance_step2_system_prompt.md
View File

@ -3,7 +3,6 @@
你的任务是基于整理后的用户记忆生成结构化的 Markdown 摘要。
输入格式：
- scope_key: 用户的唯一标识
- organized_memories: 整理后的记忆列表，每个包含 namespace、memory_key、content
输出要求：

33
src/gateway/session.rs
View File

@ -20,7 +20,7 @@ use super::execution::should_display_message_to_user;
#[cfg(test)]
use super::memory_maintenance::{
MemoryMaintenanceMerge, apply_memory_maintenance_output, build_memory_maintenance_plan,
combine_managed_memory_markdown, extract_json_object, is_recoverable_maintenance_llm_error,
extract_json_object, is_recoverable_maintenance_llm_error,
strip_json_code_fence,
};
use super::memory_maintenance::{MemoryMaintenanceScopeResult, MemoryOrganizationOutput};
@ -427,7 +427,7 @@ impl SessionManager {
pub(crate) async fn run_memory_maintenance_for_all_scopes(
&self,
) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> {
) -> Result<Option<MemoryMaintenanceScopeResult>, AgentError> {
self.memory_maintenance.run_for_all_scopes().await
}
@ -1317,15 +1317,16 @@ mod tests {
})
.unwrap();
let results = session_manager
let result = session_manager
.run_memory_maintenance_for_all_scopes()
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].scope_key, "feishu:user-1");
assert!(result.is_some());
let result = result.unwrap();
assert_eq!(result.scope_key, "all");
// 由于步骤2需要新的提示词和输入格式，这里只验证基本功能
assert!(!results[0].managed_markdown.is_empty());
assert!(!result.managed_markdown.is_empty());
}
#[tokio::test]
@ -1378,12 +1379,13 @@ mod tests {
.unwrap();
}
let results = session_manager
let result = session_manager
.run_memory_maintenance_for_all_scopes()
.await
.unwrap();
assert!(results.is_empty());
// 当遇到可恢复错误时，没有整理任何记忆，返回 None
assert!(result.is_none());
}
#[test]
@ -1460,21 +1462,6 @@ mod tests {
assert_eq!(all_memories[0].content, "用户主要在做AI产品设计与实现");
}
#[test]
fn test_combine_managed_memory_markdown_prefers_richer_summary_over_subset() {
let combined = combine_managed_memory_markdown(&[
"### 用户事实\n- 用户在做AI产品\n\n### 用户偏好\n- 偏好简洁表达".to_string(),
"- 用户在做AI产品".to_string(),
"### 用户事实\n- 用户名为区德成，昵称DC。".to_string(),
]);
assert!(
combined.contains("### 用户事实\n- 用户在做AI产品\n\n### 用户偏好\n- 偏好简洁表达")
);
assert!(combined.contains("### 用户事实\n- 用户名为区德成，昵称DC。"));
assert_eq!(combined.matches("- 用户在做AI产品").count(), 1);
}
#[test]
fn test_should_display_message_to_user_hides_completed_tool_results_by_default() {
let completed = ChatMessage::tool("call-1", "calculator", "2");