diff --git a/server/src/handlers/graph.rs b/server/src/handlers/graph.rs index 8dbd8ee..568c865 100644 --- a/server/src/handlers/graph.rs +++ b/server/src/handlers/graph.rs @@ -163,9 +163,15 @@ impl McpTool for CreateEntitiesHandler { let req: CreateEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let mut inserted = Vec::new(); state.modify_graph(|g| { - for mut entity in req.entities { - if !entity.name.is_empty() { - entity.entity_type = crate::models::normalize_entity_type(&entity.entity_type); + for entity_input in req.entities { + if !entity_input.name.is_empty() { + let mut entity = crate::models::Entity { + name: entity_input.name.clone(), + entity_type: crate::models::normalize_entity_type(&entity_input.entity_type), + observations: entity_input.observations, + namespace: entity_input.namespace.unwrap_or_else(crate::models::default_namespace), + ..Default::default() + }; inserted.push(entity.clone()); g.entities.insert(entity.name.clone(), entity); } @@ -256,10 +262,15 @@ impl McpTool for CreateRelationsHandler { auto_created.push(node_name); } } - for mut relation in req.relations { - if !relation.from.is_empty() && !relation.to.is_empty() { - relation.relation_type = - crate::models::normalize_relation_type(&relation.relation_type); + for relation_input in req.relations { + if !relation_input.from.is_empty() && !relation_input.to.is_empty() { + let mut relation = crate::models::Relation { + from: relation_input.from, + to: relation_input.to, + relation_type: crate::models::normalize_relation_type(&relation_input.relation_type), + namespace: relation_input.namespace.unwrap_or_else(crate::models::default_namespace), + ..Default::default() + }; added_relations.push(format!( "{} -[{}]-> {}", relation.from, relation.relation_type, relation.to @@ -443,13 +454,21 @@ impl McpTool for DeleteRelationsHandler { async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let req: DeleteRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?; - let mut missing_count = 0; + let mut deleted_count = 0; + let requested_count = req.relations.len(); + state.modify_graph(|master| { - let to_rem: HashSet<_> = req.relations.into_iter().collect(); let initial_len = master.relations.len(); - master.relations.retain(|r| !to_rem.contains(r)); - missing_count = to_rem.len() - (initial_len - master.relations.len()); + master.relations.retain(|r| { + let should_delete = req.relations.iter().any(|target| { + target.from == r.from && target.to == r.to && target.relation_type == r.relation_type + }); + !should_delete + }); + deleted_count = initial_len - master.relations.len(); }); + + let missing_count = requested_count.saturating_sub(deleted_count); if missing_count > 0 { return Err(crate::error::AppError::Internal(format!( "Error: {} relation(s) not found in graph. Please verify exact relation properties (from, to, relation_type) using read_graph or get_subgraph.", diff --git a/server/src/tools.rs b/server/src/tools.rs index e451f45..4b3ed97 100644 --- a/server/src/tools.rs +++ b/server/src/tools.rs @@ -2,18 +2,36 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; /// Create new entities in the knowledge graph. Note: entity_type MUST ALWAYS be PascalCase (e.g. DatabaseTable, McpTool, File). +#[derive(Debug, Deserialize, Serialize, JsonSchema)] +pub struct EntityInput { + pub name: String, + pub entity_type: String, + #[serde(default)] + pub observations: Vec, + #[serde(default)] + pub namespace: Option, +} + #[derive(Debug, Deserialize, Serialize, JsonSchema)] pub struct CreateEntitiesTool { /// Array of entities to create. - pub entities: Vec, + pub entities: Vec, +} + +#[derive(Debug, Deserialize, Serialize, JsonSchema)] +pub struct RelationInput { + pub from: String, + pub to: String, + pub relation_type: String, + #[serde(default)] + pub namespace: Option, } -/// Create new relations between entities in the knowledge graph. /// Create new relations between entities in the knowledge graph. #[derive(Debug, Deserialize, Serialize, JsonSchema)] pub struct CreateRelationsTool { - /// Array of relations to create. Each relation item MUST use 'from', 'to', and 'relation_type'. - pub relations: Vec, + /// Array of relations to create. + pub relations: Vec, } #[derive(Debug, Deserialize, Serialize, JsonSchema)] @@ -57,7 +75,7 @@ pub struct DeleteObservationsTool { #[derive(Debug, Deserialize, Serialize, JsonSchema)] pub struct DeleteRelationsTool { /// Array of relations to delete. - pub relations: Vec, + pub relations: Vec, } /// Read the entire knowledge graph. WARNING: For large graphs, use search_nodes or pagination (limit, offset) to avoid context limits.