use std::collections::{HashMap, HashSet};
use serde::{Deserialize, Serialize};
use crate::error::MemoryError;
use crate::types::{MemoryPoint, MemoryPointKind, HEADLESS_POINT_ID};
#[derive(Debug, Clone, PartialEq, Eq)]
struct ChildPosition {
parent_id: String,
index: usize,
}
#[derive(Debug, Clone)]
pub struct MemoryTree {
nodes: HashMap<String, MemoryPoint>,
children: HashMap<String, Vec<String>>,
child_positions: HashMap<String, ChildPosition>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemoryTypeSearchOptions {
pub relation_depth: usize,
pub max_parent_types: usize,
pub max_matches_per_type: usize,
pub max_neighborhood_nodes: usize,
}
impl Default for MemoryTypeSearchOptions {
fn default() -> Self {
Self {
relation_depth: 1,
max_parent_types: 5,
max_matches_per_type: 5,
max_neighborhood_nodes: 64,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemoryRecallOptions {
pub max_types: usize,
pub max_matches_per_type: usize,
pub max_returned_memories: usize,
}
impl Default for MemoryRecallOptions {
fn default() -> Self {
Self {
max_types: 5,
max_matches_per_type: 3,
max_returned_memories: 64,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MemoryTypeNeighborhood {
pub query_type: String,
pub matched_point_id: String,
pub points: Vec<MemoryNeighborhoodPoint>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MemoryRecallBranch {
pub query_type: String,
pub matched_point_id: String,
pub memories: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MemoryNeighborhoodPoint {
pub point_id: String,
pub parent_id: Option<String>,
pub kind: MemoryPointKind,
pub relation_depth: i32,
pub storage: Option<String>,
pub types: Option<String>,
}
impl MemoryTree {
pub fn new() -> Self {
let headless = MemoryPoint::headless();
let mut nodes = HashMap::new();
nodes.insert(headless.point_id.clone(), headless);
Self {
nodes,
children: HashMap::new(),
child_positions: HashMap::new(),
}
}
pub fn from_points(points: Vec<MemoryPoint>) -> Result<Self, MemoryError> {
let mut tree = Self::new();
for point in points {
if point.is_headless() {
point.validate()?;
tree.nodes.insert(HEADLESS_POINT_ID.to_string(), point);
continue;
}
if tree.nodes.contains_key(&point.point_id) {
return Err(MemoryError::DuplicatePoint {
point_id: point.point_id,
});
}
point.validate()?;
tree.nodes.insert(point.point_id.clone(), point);
}
tree.rebuild_indexes()?;
tree.validate()?;
Ok(tree)
}
pub fn insert_root(
&mut self,
storage: impl Into<String>,
types: impl Into<String>,
) -> Result<String, MemoryError> {
let point = MemoryPoint::new_root(storage, types)?;
self.insert_point(point)
}
pub fn insert_child(
&mut self,
parent_id: &str,
storage: impl Into<String>,
types: impl Into<String>,
) -> Result<String, MemoryError> {
self.ensure_parent_exists(parent_id)?;
let point = if parent_id == HEADLESS_POINT_ID {
MemoryPoint::new_root(storage, types)?
} else {
MemoryPoint::new_child(parent_id, storage, types)?
};
self.insert_point(point)
}
pub fn insert_point(&mut self, mut point: MemoryPoint) -> Result<String, MemoryError> {
if point.is_headless() {
return Err(MemoryError::CannotInsertHeadless);
}
point.validate()?;
if self.nodes.contains_key(&point.point_id) {
return Err(MemoryError::DuplicatePoint {
point_id: point.point_id,
});
}
let parent_id = point.parent_id.clone().ok_or(MemoryError::MissingParent)?;
self.ensure_parent_exists(&parent_id)?;
point.kind = if parent_id == HEADLESS_POINT_ID {
MemoryPointKind::Root
} else {
MemoryPointKind::Point
};
let point_id = point.point_id.clone();
self.nodes.insert(point_id.clone(), point);
self.attach_child(&parent_id, &point_id);
Ok(point_id)
}
pub fn contains(&self, point_id: &str) -> bool {
self.nodes.contains_key(point_id)
}
pub fn len(&self) -> usize {
self.nodes.len()
}
pub fn memory_len(&self) -> usize {
self.nodes.len().saturating_sub(1)
}
pub fn is_empty(&self) -> bool {
self.memory_len() == 0
}
pub fn get(&self, point_id: &str) -> Option<&MemoryPoint> {
self.nodes.get(point_id)
}
pub fn iter_points(&self) -> impl Iterator<Item = &MemoryPoint> {
self.nodes.values()
}
pub fn to_points(&self) -> Vec<MemoryPoint> {
self.nodes.values().cloned().collect()
}
pub fn root_ids(&self) -> Vec<String> {
self.children
.get(HEADLESS_POINT_ID)
.cloned()
.unwrap_or_default()
}
pub fn child_ids(&self, parent_id: &str) -> Result<Vec<String>, MemoryError> {
self.ensure_point_exists(parent_id)?;
Ok(self.children.get(parent_id).cloned().unwrap_or_default())
}
pub fn children(&self, parent_id: &str) -> Result<Vec<&MemoryPoint>, MemoryError> {
self.ensure_point_exists(parent_id)?;
self.children
.get(parent_id)
.into_iter()
.flatten()
.map(|child_id| self.require_point(child_id))
.collect()
}
pub fn subtree_ids(&self, start_id: &str) -> Result<Vec<String>, MemoryError> {
self.ensure_point_exists(start_id)?;
let mut result = Vec::new();
let mut stack = vec![start_id.to_string()];
while let Some(point_id) = stack.pop() {
result.push(point_id.clone());
if let Some(child_ids) = self.children.get(&point_id) {
for child_id in child_ids.iter().rev() {
stack.push(child_id.clone());
}
}
}
Ok(result)
}
pub fn subtree(&self, start_id: &str) -> Result<Vec<&MemoryPoint>, MemoryError> {
self.subtree_ids(start_id)?
.iter()
.map(|point_id| self.require_point(point_id))
.collect()
}
pub fn path_ids(&self, point_id: &str) -> Result<Vec<String>, MemoryError> {
self.ensure_point_exists(point_id)?;
let mut path = Vec::new();
let mut current_id = Some(point_id.to_string());
let mut visited = HashSet::new();
while let Some(id) = current_id {
if !visited.insert(id.clone()) {
return Err(MemoryError::CycleDetected {
point_id: point_id.to_string(),
parent_id: id,
});
}
let point = self.require_point(&id)?;
path.push(id);
current_id = point.parent_id.clone();
}
path.reverse();
Ok(path)
}
pub fn find_type_neighborhoods(
&self,
possible_parent_types: &[String],
options: &MemoryTypeSearchOptions,
) -> Result<Vec<MemoryTypeNeighborhood>, MemoryError> {
let query_types = normalize_query_types(possible_parent_types, options.max_parent_types);
let mut neighborhoods = Vec::new();
for query_type in query_types {
let mut matches = self.matching_points_by_type(&query_type);
matches.sort_by(|left, right| {
right
.0
.cmp(&left.0)
.then_with(|| left.1.point_id.cmp(&right.1.point_id))
});
for (_, point) in matches.into_iter().take(options.max_matches_per_type) {
let points = self.collect_neighborhood_points(
&point.point_id,
options.relation_depth,
options.max_neighborhood_nodes,
)?;
neighborhoods.push(MemoryTypeNeighborhood {
query_type: query_type.clone(),
matched_point_id: point.point_id.clone(),
points,
});
}
}
if neighborhoods.is_empty() && !self.is_empty() {
let fallback_query_type = possible_parent_types
.iter()
.find_map(|query_type| {
let query_type = query_type.trim();
(!query_type.is_empty()).then(|| query_type.to_string())
})
.unwrap_or_else(|| "root fallback".to_string());
for root_id in self
.root_ids()
.into_iter()
.take(options.max_matches_per_type)
{
let points = self.collect_neighborhood_points(
&root_id,
options.relation_depth,
options.max_neighborhood_nodes,
)?;
neighborhoods.push(MemoryTypeNeighborhood {
query_type: fallback_query_type.clone(),
matched_point_id: root_id,
points,
});
}
}
Ok(neighborhoods)
}
pub fn recall_subtree_memories(
&self,
memory_types: &[String],
options: &MemoryRecallOptions,
) -> Result<Vec<MemoryRecallBranch>, MemoryError> {
let query_types = normalize_query_types(memory_types, options.max_types);
let mut branches = Vec::new();
let mut seen_branch_ids = HashSet::new();
let mut seen_memory_point_ids = HashSet::new();
let mut returned_memories = 0usize;
for query_type in query_types {
if returned_memories >= options.max_returned_memories {
break;
}
let mut matches = self.matching_points_by_type(&query_type);
matches.sort_by(|left, right| {
right
.0
.cmp(&left.0)
.then_with(|| left.1.point_id.cmp(&right.1.point_id))
});
for (_, point) in matches.into_iter().take(options.max_matches_per_type) {
if returned_memories >= options.max_returned_memories {
break;
}
if !seen_branch_ids.insert(point.point_id.clone()) {
continue;
}
let mut memories = Vec::new();
for point_id in self.subtree_ids(&point.point_id)? {
if returned_memories >= options.max_returned_memories {
break;
}
if !seen_memory_point_ids.insert(point_id.clone()) {
continue;
}
let subtree_point = self.require_point(&point_id)?;
if !subtree_point.is_active() {
continue;
}
if let Some(storage) = subtree_point
.storage
.as_deref()
.map(str::trim)
.filter(|storage| !storage.is_empty())
{
memories.push(storage.to_string());
returned_memories += 1;
}
}
if !memories.is_empty() {
branches.push(MemoryRecallBranch {
query_type: query_type.clone(),
matched_point_id: point.point_id.clone(),
memories,
});
}
}
}
Ok(branches)
}
pub fn format_type_neighborhoods_for_model(neighborhoods: &[MemoryTypeNeighborhood]) -> String {
if neighborhoods.is_empty() {
return "<existing-memory-candidates>\n</existing-memory-candidates>".to_string();
}
let mut lines = vec!["<existing-memory-candidates>".to_string()];
for neighborhood in neighborhoods {
lines.push(format!(
" <candidate query_type=\"{}\" matched_point_id=\"{}\">",
escape_model_text(&neighborhood.query_type),
escape_model_text(&neighborhood.matched_point_id)
));
for point in &neighborhood.points {
lines.push(format!(
" - relation_depth={} kind={:?} point_id={} parent_id={} types={} storage={}",
point.relation_depth,
point.kind,
escape_model_text(&point.point_id),
point
.parent_id
.as_deref()
.map(escape_model_text)
.unwrap_or_else(|| "none".to_string()),
point
.types
.as_deref()
.map(escape_model_text)
.unwrap_or_else(|| "none".to_string()),
point
.storage
.as_deref()
.map(escape_model_text)
.unwrap_or_else(|| "none".to_string())
));
}
lines.push(" </candidate>".to_string());
}
lines.push("</existing-memory-candidates>".to_string());
lines.join("\n")
}
pub fn update_types(
&mut self,
point_id: &str,
new_types: impl Into<String>,
) -> Result<(), MemoryError> {
self.require_point_mut(point_id)?.update_types(new_types)
}
pub fn update_storage(
&mut self,
point_id: &str,
new_storage: impl Into<String>,
) -> Result<(), MemoryError> {
self.require_point_mut(point_id)?
.update_storage(new_storage)
}
pub fn deactivate(&mut self, point_id: &str) -> Result<(), MemoryError> {
self.require_point_mut(point_id)?.deactivate()
}
pub fn reactivate(&mut self, point_id: &str) -> Result<(), MemoryError> {
self.require_point_mut(point_id)?.reactivate()
}
pub fn reparent(&mut self, point_id: &str, new_parent_id: &str) -> Result<(), MemoryError> {
if point_id == HEADLESS_POINT_ID {
return Err(MemoryError::HeadlessModification {
operation: "reparent",
});
}
self.ensure_point_exists(point_id)?;
self.ensure_parent_exists(new_parent_id)?;
if point_id == new_parent_id || self.is_ancestor(point_id, new_parent_id)? {
return Err(MemoryError::CycleDetected {
point_id: point_id.to_string(),
parent_id: new_parent_id.to_string(),
});
}
let old_parent_id = self
.require_point(point_id)?
.parent_id
.clone()
.ok_or(MemoryError::MissingParent)?;
if old_parent_id == new_parent_id {
return Ok(());
}
self.detach_child(point_id)?;
self.require_point_mut(point_id)?
.set_parent(new_parent_id.to_string())?;
self.attach_child(new_parent_id, point_id);
Ok(())
}
pub fn validate(&self) -> Result<(), MemoryError> {
let headless = self.require_point(HEADLESS_POINT_ID)?;
headless.validate()?;
for point in self.nodes.values() {
point.validate()?;
if point.is_headless() {
continue;
}
let parent_id = point
.parent_id
.as_deref()
.ok_or(MemoryError::MissingParent)?;
self.ensure_parent_exists(parent_id)?;
if parent_id == HEADLESS_POINT_ID && point.kind != MemoryPointKind::Root {
return Err(MemoryError::InvalidPointKind {
point_id: point.point_id.clone(),
expected: "root",
});
}
if parent_id != HEADLESS_POINT_ID && point.kind != MemoryPointKind::Point {
return Err(MemoryError::InvalidPointKind {
point_id: point.point_id.clone(),
expected: "point",
});
}
self.assert_reaches_headless(&point.point_id)?;
}
for (parent_id, child_ids) in &self.children {
self.ensure_point_exists(parent_id)?;
for (index, child_id) in child_ids.iter().enumerate() {
let child = self.require_point(child_id)?;
if child.parent_id.as_deref() != Some(parent_id.as_str()) {
return Err(MemoryError::TreeInvariant {
message: format!("child '{}' has mismatched parent", child_id),
});
}
let position = self.child_positions.get(child_id).ok_or_else(|| {
MemoryError::TreeInvariant {
message: format!("child '{}' is missing position index", child_id),
}
})?;
if position.parent_id != *parent_id || position.index != index {
return Err(MemoryError::TreeInvariant {
message: format!("child '{}' has stale position index", child_id),
});
}
}
}
for point in self.nodes.values() {
if point.is_headless() {
continue;
}
if !self.child_positions.contains_key(&point.point_id) {
return Err(MemoryError::TreeInvariant {
message: format!("point '{}' is missing from child index", point.point_id),
});
}
}
Ok(())
}
fn require_point(&self, point_id: &str) -> Result<&MemoryPoint, MemoryError> {
self.nodes
.get(point_id)
.ok_or_else(|| MemoryError::PointNotFound {
point_id: point_id.to_string(),
})
}
fn require_point_mut(&mut self, point_id: &str) -> Result<&mut MemoryPoint, MemoryError> {
self.nodes
.get_mut(point_id)
.ok_or_else(|| MemoryError::PointNotFound {
point_id: point_id.to_string(),
})
}
fn ensure_point_exists(&self, point_id: &str) -> Result<(), MemoryError> {
if self.nodes.contains_key(point_id) {
Ok(())
} else {
Err(MemoryError::PointNotFound {
point_id: point_id.to_string(),
})
}
}
fn ensure_parent_exists(&self, parent_id: &str) -> Result<(), MemoryError> {
if self.nodes.contains_key(parent_id) {
Ok(())
} else {
Err(MemoryError::ParentNotFound {
parent_id: parent_id.to_string(),
})
}
}
fn attach_child(&mut self, parent_id: &str, child_id: &str) {
let child_ids = self.children.entry(parent_id.to_string()).or_default();
let index = child_ids.len();
child_ids.push(child_id.to_string());
self.child_positions.insert(
child_id.to_string(),
ChildPosition {
parent_id: parent_id.to_string(),
index,
},
);
}
fn detach_child(&mut self, child_id: &str) -> Result<(), MemoryError> {
let position =
self.child_positions
.remove(child_id)
.ok_or_else(|| MemoryError::TreeInvariant {
message: format!("child '{}' is missing from position index", child_id),
})?;
let remove_parent_entry;
{
let child_ids = self.children.get_mut(&position.parent_id).ok_or_else(|| {
MemoryError::TreeInvariant {
message: format!(
"parent '{}' is missing from children index",
position.parent_id
),
}
})?;
if child_ids.get(position.index).map(String::as_str) != Some(child_id) {
return Err(MemoryError::TreeInvariant {
message: format!("child '{}' has inconsistent position index", child_id),
});
}
child_ids.swap_remove(position.index);
if position.index < child_ids.len() {
let moved_child_id = child_ids[position.index].clone();
if let Some(moved_position) = self.child_positions.get_mut(&moved_child_id) {
moved_position.index = position.index;
}
}
remove_parent_entry = child_ids.is_empty();
}
if remove_parent_entry {
self.children.remove(&position.parent_id);
}
Ok(())
}
fn is_ancestor(&self, ancestor_id: &str, point_id: &str) -> Result<bool, MemoryError> {
let mut current_id = Some(point_id.to_string());
while let Some(id) = current_id {
if id == ancestor_id {
return Ok(true);
}
current_id = self.require_point(&id)?.parent_id.clone();
}
Ok(false)
}
fn assert_reaches_headless(&self, point_id: &str) -> Result<(), MemoryError> {
let mut current_id = Some(point_id.to_string());
let mut steps = 0usize;
let max_steps = self.nodes.len();
while let Some(id) = current_id {
if steps > max_steps {
let parent_id = self
.nodes
.get(&id)
.and_then(|point| point.parent_id.clone())
.unwrap_or_default();
return Err(MemoryError::CycleDetected {
point_id: point_id.to_string(),
parent_id,
});
}
if id == HEADLESS_POINT_ID {
return Ok(());
}
current_id = self.require_point(&id)?.parent_id.clone();
steps += 1;
}
Err(MemoryError::TreeInvariant {
message: format!("point '{}' does not reach headless", point_id),
})
}
fn rebuild_indexes(&mut self) -> Result<(), MemoryError> {
self.children.clear();
self.child_positions.clear();
let links: Vec<(String, String)> = self
.nodes
.values()
.filter(|point| !point.is_headless())
.map(|point| {
let parent_id = point.parent_id.clone().ok_or(MemoryError::MissingParent)?;
Ok((parent_id, point.point_id.clone()))
})
.collect::<Result<_, MemoryError>>()?;
for (parent_id, child_id) in links {
self.ensure_parent_exists(&parent_id)?;
self.attach_child(&parent_id, &child_id);
}
Ok(())
}
fn matching_points_by_type(&self, query_type: &str) -> Vec<(u32, &MemoryPoint)> {
self.nodes
.values()
.filter(|point| !point.is_headless() && point.is_active())
.filter_map(|point| {
let score = type_match_score(point.types.as_deref()?, query_type)?;
Some((score, point))
})
.collect()
}
fn collect_neighborhood_points(
&self,
matched_point_id: &str,
relation_depth: usize,
max_nodes: usize,
) -> Result<Vec<MemoryNeighborhoodPoint>, MemoryError> {
let mut result = Vec::new();
let mut seen = HashSet::new();
self.collect_ancestors(
matched_point_id,
relation_depth,
&mut seen,
&mut result,
max_nodes,
)?;
self.push_neighborhood_point(matched_point_id, 0, &mut seen, &mut result, max_nodes)?;
self.collect_descendants(
matched_point_id,
relation_depth,
&mut seen,
&mut result,
max_nodes,
)?;
Ok(result)
}
fn collect_ancestors(
&self,
matched_point_id: &str,
relation_depth: usize,
seen: &mut HashSet<String>,
result: &mut Vec<MemoryNeighborhoodPoint>,
max_nodes: usize,
) -> Result<(), MemoryError> {
let mut ancestors = Vec::new();
let mut current_parent = self.require_point(matched_point_id)?.parent_id.clone();
for depth in 1..=relation_depth {
let Some(parent_id) = current_parent else {
break;
};
if parent_id == HEADLESS_POINT_ID {
break;
}
let parent = self.require_point(&parent_id)?;
ancestors.push((parent_id.clone(), -(depth as i32)));
current_parent = parent.parent_id.clone();
}
for (point_id, relation_depth) in ancestors.into_iter().rev() {
self.push_neighborhood_point(&point_id, relation_depth, seen, result, max_nodes)?;
}
Ok(())
}
fn collect_descendants(
&self,
matched_point_id: &str,
relation_depth: usize,
seen: &mut HashSet<String>,
result: &mut Vec<MemoryNeighborhoodPoint>,
max_nodes: usize,
) -> Result<(), MemoryError> {
let mut queue = vec![(matched_point_id.to_string(), 0usize)];
while let Some((point_id, depth)) = queue.pop() {
if depth >= relation_depth || result.len() >= max_nodes {
continue;
}
if let Some(child_ids) = self.children.get(&point_id) {
for child_id in child_ids.iter().rev() {
let next_depth = depth + 1;
self.push_neighborhood_point(
child_id,
next_depth as i32,
seen,
result,
max_nodes,
)?;
queue.push((child_id.clone(), next_depth));
if result.len() >= max_nodes {
break;
}
}
}
}
Ok(())
}
fn push_neighborhood_point(
&self,
point_id: &str,
relation_depth: i32,
seen: &mut HashSet<String>,
result: &mut Vec<MemoryNeighborhoodPoint>,
max_nodes: usize,
) -> Result<(), MemoryError> {
if result.len() >= max_nodes || !seen.insert(point_id.to_string()) {
return Ok(());
}
let point = self.require_point(point_id)?;
result.push(MemoryNeighborhoodPoint {
point_id: point.point_id.clone(),
parent_id: point.parent_id.clone(),
kind: point.kind.clone(),
relation_depth,
storage: point.storage.clone(),
types: point.types.clone(),
});
Ok(())
}
}
impl Default for MemoryTree {
fn default() -> Self {
Self::new()
}
}
fn normalize_query_types(possible_parent_types: &[String], max_parent_types: usize) -> Vec<String> {
let mut seen = HashSet::new();
let mut normalized = Vec::new();
for query_type in possible_parent_types {
let query_type = query_type.trim();
if query_type.is_empty() {
continue;
}
let key = normalize_search_term(query_type);
if seen.insert(key) {
normalized.push(query_type.to_string());
}
if normalized.len() >= max_parent_types {
break;
}
}
normalized
}
fn type_match_score(candidate_type: &str, query_type: &str) -> Option<u32> {
let candidate = normalize_search_term(candidate_type);
let query = normalize_search_term(query_type);
if candidate.is_empty() || query.is_empty() {
return None;
}
if candidate == query {
return Some(1000);
}
if candidate.contains(&query) || query.contains(&candidate) {
let shorter = candidate.chars().count().min(query.chars().count()) as u32;
return Some(800 + shorter.min(100));
}
let simplified_candidate = simplify_search_term(&candidate);
let simplified_query = simplify_search_term(&query);
if !simplified_candidate.is_empty() && !simplified_query.is_empty() {
if simplified_candidate == simplified_query {
return Some(760);
}
if simplified_candidate.contains(&simplified_query)
|| simplified_query.contains(&simplified_candidate)
{
let shorter = simplified_candidate
.chars()
.count()
.min(simplified_query.chars().count()) as u32;
return Some(700 + shorter.min(50));
}
}
if let Some(score) = semantic_type_match_score(&candidate, &query) {
return Some(score);
}
let query_chars = query.chars().count();
if query_chars <= 4 {
let candidate_chars: HashSet<char> = candidate.chars().collect();
let query_chars_set: HashSet<char> = query.chars().collect();
let overlap = query_chars_set.intersection(&candidate_chars).count();
let min_len = query_chars_set.len().min(candidate_chars.len());
if min_len > 0 && overlap * 2 > min_len {
return Some(400 + overlap as u32);
}
}
None
}
fn semantic_type_match_score(candidate: &str, query: &str) -> Option<u32> {
let candidate_tokens = semantic_type_tokens(candidate);
let query_tokens = semantic_type_tokens(query);
if candidate_tokens.is_empty() || query_tokens.is_empty() {
return None;
}
let overlap = candidate_tokens.intersection(&query_tokens).count();
if overlap == 0 {
return None;
}
if query_tokens.contains("food") && !candidate_tokens.contains("food") {
return None;
}
Some(560 + (overlap as u32 * 30))
}
fn semantic_type_tokens(value: &str) -> HashSet<&'static str> {
let mut tokens = HashSet::new();
if contains_any(
value,
&[
"饮食", "食物", "食材", "面食", "菜肴", "菜品", "餐饮", "吃", "餐", "口味",
],
) {
tokens.insert("food");
}
if contains_any(value, &["忌口", "过敏", "不喜欢", "避开", "不能吃"]) {
tokens.insert("restriction");
}
if contains_any(
value,
&["前端", "界面", "设计", "ui", "颜色", "主题", "信息密度"],
) {
tokens.insert("frontend_design");
}
if contains_any(value, &["rust", "错误", "thiserror", "anyhow"]) {
tokens.insert("rust_error");
}
if contains_any(value, &["中文", "英文", "语言", "交流"]) {
tokens.insert("language");
}
if contains_any(value, &["名字", "称呼", "叫我"]) {
tokens.insert("identity");
}
tokens
}
fn contains_any(value: &str, terms: &[&str]) -> bool {
terms.iter().any(|term| value.contains(term))
}
fn normalize_search_term(value: &str) -> String {
value
.chars()
.filter(|ch| !ch.is_whitespace())
.flat_map(char::to_lowercase)
.collect()
}
fn simplify_search_term(value: &str) -> String {
let mut simplified = value.to_string();
for generic_term in [
"用户",
"我的",
"喜欢的",
"喜欢",
"偏好",
"类型",
"种类",
"类别",
"信息",
"的",
] {
simplified = simplified.replace(generic_term, "");
}
simplified
}
fn escape_model_text(value: &str) -> String {
value
.replace('&', "&")
.replace('<', "<")
.replace('>', ">")
.replace('"', """)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn new_tree_contains_headless_only() {
let tree = MemoryTree::new();
assert_eq!(tree.len(), 1);
assert_eq!(tree.memory_len(), 0);
assert!(tree.contains(HEADLESS_POINT_ID));
assert!(tree.validate().is_ok());
}
#[test]
fn inserts_root_and_child_points() {
let mut tree = MemoryTree::new();
let root_id = tree
.insert_root("likes noodles", "food preference")
.unwrap();
let child_id = tree
.insert_child(&root_id, "likes zhajiangmian", "noodle preference")
.unwrap();
assert_eq!(tree.root_ids(), vec![root_id.clone()]);
assert_eq!(tree.child_ids(&root_id).unwrap(), vec![child_id]);
assert_eq!(tree.memory_len(), 2);
assert!(tree.validate().is_ok());
}
#[test]
fn updates_point_content_through_tree() {
let mut tree = MemoryTree::new();
let root_id = tree
.insert_root("likes noodles", "noodle preference")
.unwrap();
tree.update_types(&root_id, "food preference").unwrap();
tree.update_storage(&root_id, "likes noodles and ribs")
.unwrap();
let root = tree.get(&root_id).unwrap();
assert_eq!(root.types.as_deref(), Some("food preference"));
assert_eq!(root.storage.as_deref(), Some("likes noodles and ribs"));
}
#[test]
fn reparents_point_with_index_updates() {
let mut tree = MemoryTree::new();
let root_a = tree.insert_root("food", "food preference").unwrap();
let root_b = tree.insert_root("sports", "sports preference").unwrap();
let child = tree
.insert_child(&root_a, "likes tennis", "tennis preference")
.unwrap();
tree.reparent(&child, &root_b).unwrap();
assert!(tree.child_ids(&root_a).unwrap().is_empty());
assert_eq!(tree.child_ids(&root_b).unwrap(), vec![child.clone()]);
assert_eq!(
tree.get(&child).unwrap().parent_id.as_deref(),
Some(root_b.as_str())
);
assert!(tree.validate().is_ok());
}
#[test]
fn returns_subtree_from_start_node() {
let mut tree = MemoryTree::new();
let root = tree.insert_root("food", "food preference").unwrap();
let child_a = tree
.insert_child(&root, "likes noodles", "noodle preference")
.unwrap();
let child_b = tree
.insert_child(&root, "likes ribs", "dish preference")
.unwrap();
let grandchild = tree
.insert_child(&child_a, "likes zhajiangmian", "specific noodle preference")
.unwrap();
let subtree = tree.subtree_ids(&root).unwrap();
assert_eq!(subtree.len(), 4);
assert_eq!(subtree[0], root);
assert!(subtree.contains(&child_a));
assert!(subtree.contains(&child_b));
assert!(subtree.contains(&grandchild));
}
#[test]
fn returns_path_from_headless_to_point() {
let mut tree = MemoryTree::new();
let root = tree.insert_root("food", "food preference").unwrap();
let child = tree
.insert_child(&root, "likes noodles", "noodle preference")
.unwrap();
let path = tree.path_ids(&child).unwrap();
assert_eq!(
path,
vec![HEADLESS_POINT_ID.to_string(), root.clone(), child.clone()]
);
}
#[test]
fn prevents_cycle_when_reparenting() {
let mut tree = MemoryTree::new();
let root = tree.insert_root("food", "food preference").unwrap();
let child = tree
.insert_child(&root, "likes noodles", "noodle preference")
.unwrap();
let grandchild = tree
.insert_child(&child, "likes zhajiangmian", "specific noodle preference")
.unwrap();
let result = tree.reparent(&root, &grandchild);
assert!(matches!(result, Err(MemoryError::CycleDetected { .. })));
}
#[test]
fn rejects_missing_parent() {
let mut tree = MemoryTree::new();
let result = tree.insert_child("missing", "likes noodles", "noodle preference");
assert!(matches!(result, Err(MemoryError::ParentNotFound { .. })));
}
#[test]
fn rebuilds_tree_from_points() {
let mut original = MemoryTree::new();
let root = original.insert_root("food", "food preference").unwrap();
let child = original
.insert_child(&root, "likes noodles", "noodle preference")
.unwrap();
let restored = MemoryTree::from_points(original.to_points()).unwrap();
assert_eq!(restored.memory_len(), 2);
assert_eq!(restored.child_ids(&root).unwrap(), vec![child]);
assert!(restored.validate().is_ok());
}
#[test]
fn preserves_headless_point_when_rebuilding_from_points() {
let mut points = MemoryTree::new().to_points();
let headless = points
.iter_mut()
.find(|point| point.point_id == HEADLESS_POINT_ID)
.unwrap();
headless.created_at = "stable-created-at".to_string();
headless.updated_at = "stable-updated-at".to_string();
let restored = MemoryTree::from_points(points).unwrap();
let restored_headless = restored.get(HEADLESS_POINT_ID).unwrap();
assert_eq!(restored_headless.created_at, "stable-created-at");
assert_eq!(restored_headless.updated_at, "stable-updated-at");
}
#[test]
fn finds_type_neighborhood_with_parent_self_and_children() {
let mut tree = MemoryTree::new();
let food = tree.insert_root("用户喜欢吃面食和菜", "饮食偏好").unwrap();
let noodles = tree
.insert_child(&food, "用户喜欢吃炸酱面", "面食偏好")
.unwrap();
let zhajiangmian = tree
.insert_child(&noodles, "用户特别喜欢炸酱面", "炸酱面偏好")
.unwrap();
let swimming = tree.insert_root("用户喜欢游泳", "运动偏好").unwrap();
let options = MemoryTypeSearchOptions {
relation_depth: 1,
max_parent_types: 3,
max_matches_per_type: 1,
max_neighborhood_nodes: 16,
};
let neighborhoods = tree
.find_type_neighborhoods(&["面食".to_string()], &options)
.unwrap();
assert_eq!(neighborhoods.len(), 1);
assert_eq!(neighborhoods[0].matched_point_id, noodles);
let point_ids: Vec<&str> = neighborhoods[0]
.points
.iter()
.map(|point| point.point_id.as_str())
.collect();
assert!(point_ids.contains(&food.as_str()));
assert!(point_ids.contains(&neighborhoods[0].matched_point_id.as_str()));
assert!(point_ids.contains(&zhajiangmian.as_str()));
assert!(!point_ids.contains(&swimming.as_str()));
assert!(neighborhoods[0]
.points
.iter()
.any(|point| point.point_id == food && point.relation_depth == -1));
assert!(neighborhoods[0]
.points
.iter()
.any(|point| point.point_id == zhajiangmian && point.relation_depth == 1));
}
#[test]
fn limits_possible_parent_types_before_searching() {
let mut tree = MemoryTree::new();
let food = tree.insert_root("用户喜欢吃面食", "饮食偏好").unwrap();
let sport = tree.insert_root("用户喜欢游泳", "运动偏好").unwrap();
let options = MemoryTypeSearchOptions {
relation_depth: 0,
max_parent_types: 1,
max_matches_per_type: 5,
max_neighborhood_nodes: 16,
};
let neighborhoods = tree
.find_type_neighborhoods(&["饮食".to_string(), "运动".to_string()], &options)
.unwrap();
assert_eq!(neighborhoods.len(), 1);
assert_eq!(neighborhoods[0].matched_point_id, food);
assert_ne!(neighborhoods[0].matched_point_id, sport);
}
#[test]
fn formats_type_neighborhoods_for_model() {
let mut tree = MemoryTree::new();
let food = tree.insert_root("用户喜欢吃面食", "饮食偏好").unwrap();
let options = MemoryTypeSearchOptions::default();
let neighborhoods = tree
.find_type_neighborhoods(&["饮食".to_string()], &options)
.unwrap();
let formatted = MemoryTree::format_type_neighborhoods_for_model(&neighborhoods);
assert!(formatted.contains("<existing-memory-candidates>"));
assert!(formatted.contains(&food));
assert!(formatted.contains("饮食偏好"));
assert!(formatted.contains("用户喜欢吃面食"));
}
#[test]
fn falls_back_to_roots_when_type_search_has_no_matches() {
let mut tree = MemoryTree::new();
let noodles = tree
.insert_root("用户喜欢吃炸酱面", "用户喜欢的面食")
.unwrap();
let swimming = tree.insert_root("用户喜欢游泳", "运动偏好").unwrap();
let options = MemoryTypeSearchOptions {
relation_depth: 0,
max_parent_types: 3,
max_matches_per_type: 1,
max_neighborhood_nodes: 16,
};
let neighborhoods = tree
.find_type_neighborhoods(&["菜肴偏好".to_string()], &options)
.unwrap();
assert_eq!(neighborhoods.len(), 1);
assert_eq!(neighborhoods[0].matched_point_id, noodles);
assert_ne!(neighborhoods[0].matched_point_id, swimming);
}
#[test]
fn recalls_storage_from_matched_subtree() {
let mut tree = MemoryTree::new();
let food = tree
.insert_root("用户喜欢吃炸酱面和糖醋排骨等食物", "用户喜欢的食物")
.unwrap();
let ribs = tree
.insert_child(&food, "用户喜欢吃糖醋排骨", "用户喜欢吃的菜")
.unwrap();
let _noodles = tree
.insert_child(&food, "用户喜欢吃炸酱面", "用户喜欢吃的面食")
.unwrap();
let _sport = tree.insert_root("用户周末喜欢跑步", "运动习惯").unwrap();
let branches = tree
.recall_subtree_memories(
&["食物偏好".to_string()],
&MemoryRecallOptions {
max_types: 3,
max_matches_per_type: 2,
max_returned_memories: 16,
},
)
.unwrap();
assert_eq!(branches.len(), 1);
assert_eq!(branches[0].matched_point_id, food);
assert!(branches[0]
.memories
.contains(&"用户喜欢吃炸酱面和糖醋排骨等食物".to_string()));
assert!(branches[0]
.memories
.contains(&"用户喜欢吃糖醋排骨".to_string()));
assert!(branches[0]
.memories
.contains(&"用户喜欢吃炸酱面".to_string()));
assert!(!branches[0]
.memories
.contains(&"用户周末喜欢跑步".to_string()));
tree.deactivate(&ribs).unwrap();
let branches = tree
.recall_subtree_memories(&["食物偏好".to_string()], &MemoryRecallOptions::default())
.unwrap();
assert!(!branches[0]
.memories
.contains(&"用户喜欢吃糖醋排骨".to_string()));
}
#[test]
fn recalls_related_food_roots_by_semantic_type() {
let mut tree = MemoryTree::new();
tree.insert_root("用户喜欢吃炸酱面", "用户喜欢吃的面食")
.unwrap();
tree.insert_root("用户不喜欢吃香菜", "用户不喜欢的食材")
.unwrap();
tree.insert_root("用户对花生过敏,饮食建议必须避开花生", "食物过敏信息")
.unwrap();
tree.insert_root("用户周末喜欢跑步", "运动习惯").unwrap();
let branches = tree
.recall_subtree_memories(
&["饮食偏好".to_string(), "忌口偏好".to_string()],
&MemoryRecallOptions {
max_types: 3,
max_matches_per_type: 8,
max_returned_memories: 16,
},
)
.unwrap();
let memories = branches
.iter()
.flat_map(|branch| branch.memories.iter())
.cloned()
.collect::<Vec<_>>()
.join("\n");
assert!(memories.contains("炸酱面"));
assert!(memories.contains("香菜"));
assert!(memories.contains("花生"));
assert!(!memories.contains("跑步"));
}
}