feat(graph): add ergonomic UpdateEntitiesTool for renaming and type changes
This commit is contained in:
1 parent
cfcd42f41e
commit
eed9a9165c
3 files changed
+120
No files matched your search
@@ -193,6 +193,108 @@ impl McpTool for CreateEntitiesHandler {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct UpdateEntitiesHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for UpdateEntitiesHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"update_entities"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<UpdateEntitiesTool>(
|
||||
"update_entities",
|
||||
"Update existing entities in the knowledge graph (e.g., to rename an entity or change its type).",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: UpdateEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
|
||||
let mut updated_entities = Vec::new();
|
||||
let mut deleted_names = Vec::new();
|
||||
let mut not_found = Vec::new();
|
||||
let mut conflict_names = Vec::new();
|
||||
|
||||
state.modify_graph(|g| {
|
||||
for update in req.updates {
|
||||
if !g.entities.contains_key(&update.name) {
|
||||
not_found.push(update.name.clone());
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(new_name) = &update.new_name {
|
||||
if update.name != *new_name && g.entities.contains_key(new_name) {
|
||||
conflict_names.push(new_name.clone());
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(mut entity) = g.entities.remove(&update.name) {
|
||||
let mut renamed = false;
|
||||
let old_name = entity.name.clone();
|
||||
|
||||
if let Some(new_name) = update.new_name {
|
||||
if old_name != new_name {
|
||||
entity.name = new_name.clone();
|
||||
renamed = true;
|
||||
deleted_names.push(old_name.clone());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(new_type) = update.new_entity_type {
|
||||
entity.entity_type = crate::models::normalize_entity_type(&new_type);
|
||||
}
|
||||
|
||||
let current_name = entity.name.clone();
|
||||
|
||||
if renamed {
|
||||
// update relations
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
g.relations.retain_mut(|r| {
|
||||
if r.from == old_name {
|
||||
r.from = current_name.clone();
|
||||
}
|
||||
if r.to == old_name {
|
||||
r.to = current_name.clone();
|
||||
}
|
||||
// Keep unique relations
|
||||
seen.insert(r.clone())
|
||||
});
|
||||
}
|
||||
|
||||
updated_entities.push(entity.clone());
|
||||
g.entities.insert(current_name, entity);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if !not_found.is_empty() {
|
||||
return Err(crate::error::AppError::Internal(format!(
|
||||
"Error: Entities not found: {}", not_found.join(", ")
|
||||
)));
|
||||
}
|
||||
|
||||
if !conflict_names.is_empty() {
|
||||
return Err(crate::error::AppError::Internal(format!(
|
||||
"Error: Cannot rename to existing entity names: {}", conflict_names.join(", ")
|
||||
)));
|
||||
}
|
||||
|
||||
let idx = state.get_search_index().await;
|
||||
for old_name in deleted_names {
|
||||
drop(idx.delete_document(&old_name));
|
||||
}
|
||||
|
||||
if !updated_entities.is_empty() {
|
||||
let _ = idx.index_entities_batch(&updated_entities).await;
|
||||
}
|
||||
|
||||
let names: Vec<String> = updated_entities.iter().map(|e| e.name.clone()).collect();
|
||||
Ok(format!("Successfully updated {} entity/entities: {}", names.len(), names.join(", ")))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CreateRelationsHandler;
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -479,6 +479,7 @@ impl MemoryHandler {
|
||||
|
||||
register!(graph::QueryGraphPathHandler);
|
||||
register!(graph::CreateEntitiesHandler);
|
||||
register!(graph::UpdateEntitiesHandler);
|
||||
register!(graph::CreateRelationsHandler);
|
||||
register!(graph::AddObservationsHandler);
|
||||
register!(graph::DeleteEntitiesHandler);
|
||||
|
||||
@@ -18,6 +18,23 @@ pub struct CreateEntitiesTool {
|
||||
pub entities: Vec<EntityInput>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct UpdateEntityInput {
|
||||
/// The current name of the entity to update.
|
||||
pub name: String,
|
||||
/// Optional new name to rename the entity to.
|
||||
pub new_name: Option<String>,
|
||||
/// Optional new entity type (e.g. PascalCase).
|
||||
pub new_entity_type: Option<String>,
|
||||
}
|
||||
|
||||
/// Update existing entities in the knowledge graph (e.g., to rename an entity or change its type).
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct UpdateEntitiesTool {
|
||||
/// Array of entity updates to perform.
|
||||
pub updates: Vec<UpdateEntityInput>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct RelationInput {
|
||||
pub from: String,
|
||||
|
||||
Reference in new issue
Block a user