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;
|
pub struct CreateRelationsHandler;
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
|
|||||||
@@ -479,6 +479,7 @@ impl MemoryHandler {
|
|||||||
|
|
||||||
register!(graph::QueryGraphPathHandler);
|
register!(graph::QueryGraphPathHandler);
|
||||||
register!(graph::CreateEntitiesHandler);
|
register!(graph::CreateEntitiesHandler);
|
||||||
|
register!(graph::UpdateEntitiesHandler);
|
||||||
register!(graph::CreateRelationsHandler);
|
register!(graph::CreateRelationsHandler);
|
||||||
register!(graph::AddObservationsHandler);
|
register!(graph::AddObservationsHandler);
|
||||||
register!(graph::DeleteEntitiesHandler);
|
register!(graph::DeleteEntitiesHandler);
|
||||||
|
|||||||
@@ -18,6 +18,23 @@ pub struct CreateEntitiesTool {
|
|||||||
pub entities: Vec<EntityInput>,
|
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)]
|
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||||
pub struct RelationInput {
|
pub struct RelationInput {
|
||||||
pub from: String,
|
pub from: String,
|
||||||
|
|||||||
Reference in new issue
Block a user