Riz Ashraf
2024d9521b
- Replaced unconstrained JSON serialization (serde_json) with dense YAML (serde_yaml) to prevent context window bloat (ADR-0122). - Replaced 404/empty array dead ends with token-efficient LLM fallback hints in router and meta/graph handlers. - Ensured graceful error fallbacks instead of hard crashes.
1828 lines
69 KiB
Rust
1828 lines
69 KiB
Rust
use crate::router::McpTool;
|
|
use crate::state::MemoryState;
|
|
use crate::tools::*;
|
|
use async_trait::async_trait;
|
|
use serde_json::Value;
|
|
use std::collections::HashSet;
|
|
use std::sync::Arc;
|
|
|
|
#[derive(serde::Serialize, Default)]
|
|
struct BorrowedGraph<'a> {
|
|
entities: std::collections::HashMap<&'a String, &'a crate::models::Entity>,
|
|
relations: Vec<&'a crate::models::Relation>,
|
|
}
|
|
|
|
pub struct GraphQueryBuilder<'a> {
|
|
graph: &'a crate::models::KnowledgeGraph,
|
|
max_depth: usize,
|
|
relation_filter: Option<&'a str>,
|
|
}
|
|
|
|
impl<'a> GraphQueryBuilder<'a> {
|
|
pub fn new(graph: &'a crate::models::KnowledgeGraph) -> Self {
|
|
Self {
|
|
graph,
|
|
max_depth: 5,
|
|
relation_filter: None,
|
|
}
|
|
}
|
|
|
|
pub fn max_depth(mut self, depth: usize) -> Self {
|
|
self.max_depth = depth;
|
|
self
|
|
}
|
|
|
|
pub fn relation_filter(mut self, filter: Option<&'a str>) -> Self {
|
|
self.relation_filter = filter;
|
|
self
|
|
}
|
|
|
|
pub fn find_shortest_path(&self, start: &str, end: &str) -> Option<Vec<String>> {
|
|
let adj = self.graph.build_adjacency_map(self.relation_filter);
|
|
|
|
let mut queue = std::collections::VecDeque::new();
|
|
let mut visited = std::collections::HashSet::new();
|
|
let mut parents = std::collections::HashMap::new();
|
|
|
|
queue.push_back(start);
|
|
visited.insert(start);
|
|
|
|
let mut found = false;
|
|
let mut current_depth = 0;
|
|
let mut nodes_at_current_depth = 1;
|
|
let mut nodes_at_next_depth = 0;
|
|
|
|
while let Some(current) = queue.pop_front() {
|
|
if current == end {
|
|
found = true;
|
|
break;
|
|
}
|
|
// Visited node upper-bound cap to guarantee deterministic BFS bounds on dense graphs
|
|
if visited.len() > 10_000 {
|
|
break;
|
|
}
|
|
nodes_at_current_depth -= 1;
|
|
if current_depth < self.max_depth
|
|
&& let Some(neighbors) = adj.get(current)
|
|
{
|
|
for &(neighbor, rel_type, is_inverse) in neighbors {
|
|
if !visited.contains(neighbor) {
|
|
visited.insert(neighbor);
|
|
parents.insert(neighbor, (current, rel_type, is_inverse));
|
|
queue.push_back(neighbor);
|
|
nodes_at_next_depth += 1;
|
|
}
|
|
}
|
|
}
|
|
if nodes_at_current_depth == 0 {
|
|
current_depth += 1;
|
|
nodes_at_current_depth = nodes_at_next_depth;
|
|
nodes_at_next_depth = 0;
|
|
}
|
|
}
|
|
|
|
if found {
|
|
let mut path = Vec::new();
|
|
let mut curr = end;
|
|
while curr != start {
|
|
if let Some((parent, rel_type, is_inverse)) = parents.get(&curr) {
|
|
if *is_inverse {
|
|
path.push(format!("{} -[inverse({})]-> {}", parent, rel_type, curr));
|
|
} else {
|
|
path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
|
|
}
|
|
curr = parent;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
path.reverse();
|
|
Some(path)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct QueryGraphPathHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for QueryGraphPathHandler {
|
|
fn name(&self) -> &'static str {
|
|
"query_graph_path"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<QueryGraphPathTool>(
|
|
"query_graph_path",
|
|
"Find the shortest relationship path between two entities in the knowledge graph within a maximum depth.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: crate::tools::QueryGraphPathTool =
|
|
serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
state.read_graph(|graph| {
|
|
let max_depth = req.max_depth.unwrap_or(5);
|
|
let builder = GraphQueryBuilder::new(graph).max_depth(max_depth as usize);
|
|
if let Some(path) = builder.find_shortest_path(&req.start_node, &req.end_node) {
|
|
Ok(format!("Path found:\n{}", path.join("\n")))
|
|
} else {
|
|
Ok(format!(
|
|
"No path found between {} and {} within depth {}",
|
|
req.start_node, req.end_node, max_depth
|
|
))
|
|
}
|
|
})
|
|
})
|
|
.await
|
|
.map_err(|e| {
|
|
crate::error::AppError::Internal(format!("Graph traversal task failed: {}", e))
|
|
})?
|
|
}
|
|
}
|
|
|
|
pub struct CreateEntitiesHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for CreateEntitiesHandler {
|
|
fn name(&self) -> &'static str {
|
|
"create_entities"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<CreateEntitiesTool>(
|
|
"create_entities",
|
|
"Create new entities in the knowledge graph with normalized entity types.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: CreateEntitiesTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let mut inserted = Vec::new();
|
|
state.modify_graph(|g| {
|
|
for entity_input in req.entities {
|
|
if !entity_input.name.is_empty() {
|
|
let 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);
|
|
}
|
|
}
|
|
});
|
|
let names: Vec<String> = inserted
|
|
.iter()
|
|
.map(|e| format!("{} ({})", e.name, e.entity_type))
|
|
.collect();
|
|
if !inserted.is_empty() {
|
|
let idx = state.get_search_index().await;
|
|
let _ = idx.index_entities_batch(&inserted).await;
|
|
}
|
|
Ok(format!(
|
|
"Successfully created {} entity/entities: {}",
|
|
names.len(),
|
|
names.join(", ")
|
|
))
|
|
}
|
|
}
|
|
|
|
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| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
|
|
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 {
|
|
let mut target_name = update.name.clone();
|
|
if !g.entities.contains_key(&target_name) {
|
|
let lower_target = target_name.to_lowercase();
|
|
if let Some(matched_key) = g.entities.keys().find(|k| k.to_lowercase() == lower_target).cloned() {
|
|
target_name = matched_key;
|
|
} else {
|
|
not_found.push(update.name.clone());
|
|
continue;
|
|
}
|
|
}
|
|
|
|
if let Some(new_name) = &update.new_name {
|
|
if target_name != *new_name && g.entities.contains_key(new_name) {
|
|
conflict_names.push(new_name.clone());
|
|
continue;
|
|
}
|
|
}
|
|
|
|
if let Some(mut entity) = g.entities.remove(&target_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);
|
|
}
|
|
}
|
|
});
|
|
|
|
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();
|
|
let mut msg = format!("Successfully updated {} entity/entities: {}", names.len(), names.join(", "));
|
|
|
|
if !not_found.is_empty() {
|
|
msg.push_str(&format!("\nNote: {} entities were not found and skipped: {}", not_found.len(), not_found.join(", ")));
|
|
}
|
|
if !conflict_names.is_empty() {
|
|
msg.push_str(&format!("\nNote: {} entity renames were skipped due to name conflicts: {}", conflict_names.len(), conflict_names.join(", ")));
|
|
}
|
|
|
|
Ok(msg)
|
|
}
|
|
}
|
|
|
|
pub struct CreateRelationsHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for CreateRelationsHandler {
|
|
fn name(&self) -> &'static str {
|
|
"create_relations"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<CreateRelationsTool>(
|
|
"create_relations",
|
|
"Create directed relationships between entities in the knowledge graph. Requires 'from', 'to', and 'relation_type'.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: CreateRelationsTool = match serde_json::from_value(args.clone()) {
|
|
Ok(r) => r,
|
|
Err(e) => {
|
|
let err_msg = e.to_string();
|
|
if err_msg.contains("missing field `from`")
|
|
|| err_msg.contains("missing field `to`")
|
|
|| err_msg.contains("missing field `relation_type`")
|
|
{
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Schema error: {}. Note that the relation schema strictly uses 'from', 'to', and 'relation_type' (not 'source', 'target', or 'relationType'). Please correct your tool call arguments.",
|
|
err_msg
|
|
)));
|
|
}
|
|
return Err(err_msg.into());
|
|
}
|
|
};
|
|
let mut missing_nodes = std::collections::HashSet::new();
|
|
state.read_graph(|g| {
|
|
for relation in &req.relations {
|
|
if !relation.from.is_empty() && !relation.to.is_empty() {
|
|
let from_exists = g.entities.contains_key(&relation.from);
|
|
let to_exists = g.entities.contains_key(&relation.to);
|
|
if !from_exists {
|
|
missing_nodes.insert(relation.from.clone());
|
|
}
|
|
if !to_exists {
|
|
missing_nodes.insert(relation.to.clone());
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
let mut auto_created = Vec::new();
|
|
let mut added_relations = Vec::new();
|
|
state.modify_graph(|g| {
|
|
for node_name in missing_nodes {
|
|
if !g.entities.contains_key(&node_name) {
|
|
g.entities.insert(
|
|
node_name.clone(),
|
|
crate::models::Entity {
|
|
name: node_name.clone(),
|
|
entity_type: "Entity".to_string(),
|
|
observations: vec![
|
|
"Auto-created stub entity for relation endpoint".to_string(),
|
|
],
|
|
namespace: crate::models::default_namespace(),
|
|
git_branch: None,
|
|
..Default::default()
|
|
},
|
|
);
|
|
auto_created.push(node_name);
|
|
}
|
|
}
|
|
for relation_input in req.relations {
|
|
if !relation_input.from.is_empty() && !relation_input.to.is_empty() {
|
|
let 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
|
|
));
|
|
g.relations.push(relation);
|
|
}
|
|
}
|
|
});
|
|
|
|
let mut msg = format!(
|
|
"Successfully created {} relation(s):\n{}",
|
|
added_relations.len(),
|
|
added_relations.join("\n")
|
|
);
|
|
if !auto_created.is_empty() {
|
|
msg.push_str(&format!(
|
|
"\nNote: Auto-created {} missing stub entity/entities: {}",
|
|
auto_created.len(),
|
|
auto_created.join(", ")
|
|
));
|
|
}
|
|
Ok(msg)
|
|
}
|
|
}
|
|
|
|
pub struct AddObservationsHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for AddObservationsHandler {
|
|
fn name(&self) -> &'static str {
|
|
"add_observations"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<AddObservationsTool>(
|
|
"add_observations",
|
|
"Add new observations and factual statements to existing entities in the knowledge graph.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: AddObservationsTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let mut missing_entities = Vec::new();
|
|
state.read_graph(|g| {
|
|
for o in &req.observations {
|
|
if !g.entities.contains_key(&o.entity_name) {
|
|
missing_entities.push(o.entity_name.clone());
|
|
}
|
|
}
|
|
});
|
|
|
|
if !missing_entities.is_empty() {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Observations dropped for missing entities: {}",
|
|
missing_entities.join(", ")
|
|
)));
|
|
}
|
|
|
|
state.modify_graph(|g| {
|
|
for o in req.observations {
|
|
if let Some(e) = g.entities.get_mut(&o.entity_name) {
|
|
e.observations.extend(o.contents);
|
|
}
|
|
}
|
|
});
|
|
Ok("Observations added".to_string())
|
|
}
|
|
}
|
|
|
|
pub struct DeleteEntitiesHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for DeleteEntitiesHandler {
|
|
fn name(&self) -> &'static str {
|
|
"delete_entities"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<DeleteEntitiesTool>(
|
|
"delete_entities",
|
|
"Delete specified entities and all their associated relations and observations from the knowledge graph.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: DeleteEntitiesTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
|
|
let mut actual_deletes = Vec::new();
|
|
let mut not_found = Vec::new();
|
|
|
|
state.modify_graph(|master| {
|
|
for target in req.entity_names {
|
|
if master.entities.contains_key(&target) {
|
|
actual_deletes.push(target);
|
|
} else {
|
|
let lower_target = target.to_lowercase();
|
|
if let Some(matched_key) = master.entities.keys().find(|k| k.to_lowercase() == lower_target).cloned() {
|
|
actual_deletes.push(matched_key);
|
|
} else {
|
|
not_found.push(target);
|
|
}
|
|
}
|
|
}
|
|
|
|
for name in &actual_deletes {
|
|
master.entities.remove(name);
|
|
}
|
|
let delete_set: std::collections::HashSet<_> = actual_deletes.iter().cloned().collect();
|
|
master
|
|
.relations
|
|
.retain(|r| !delete_set.contains(&r.from) && !delete_set.contains(&r.to));
|
|
});
|
|
|
|
let idx = state.get_search_index().await;
|
|
for name in &actual_deletes {
|
|
drop(idx.delete_document(name));
|
|
}
|
|
|
|
let mut msg = format!("Successfully deleted {} entities.", actual_deletes.len());
|
|
if !not_found.is_empty() {
|
|
msg.push_str(&format!(" Note: {} entities were not found and skipped: {}", not_found.len(), not_found.join(", ")));
|
|
}
|
|
Ok(msg)
|
|
}
|
|
}
|
|
|
|
pub struct DeleteObservationsHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for DeleteObservationsHandler {
|
|
fn name(&self) -> &'static str {
|
|
"delete_observations"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<DeleteObservationsTool>(
|
|
"delete_observations",
|
|
"Delete specific observations from existing entities in the knowledge graph without removing the entity itself.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: DeleteObservationsTool =
|
|
serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
|
|
let mut missing_entities = Vec::new();
|
|
let mut missing_observations = Vec::new();
|
|
|
|
state.modify_graph(|master| {
|
|
for d in req.deletions {
|
|
if let Some(e) = master.entities.get_mut(&d.entity_name) {
|
|
for target in d.observations {
|
|
let target_clean = target.trim();
|
|
if target_clean.is_empty() { continue; }
|
|
|
|
let initial_len = e.observations.len();
|
|
e.observations.retain(|o| {
|
|
let o_clean = o.trim();
|
|
// Exact match OR substantial substring match (to forgive LLM whitespace/newline errors)
|
|
!(o_clean == target_clean || (target_clean.len() >= 10 && o_clean.contains(target_clean)))
|
|
});
|
|
|
|
if e.observations.len() == initial_len {
|
|
missing_observations.push(target);
|
|
}
|
|
}
|
|
} else {
|
|
missing_entities.push(d.entity_name);
|
|
}
|
|
}
|
|
});
|
|
|
|
if !missing_entities.is_empty() {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Entities not found: {}. Please use search_nodes to verify the exact entity names.",
|
|
missing_entities.join(", ")
|
|
)));
|
|
}
|
|
|
|
if !missing_observations.is_empty() {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Failed to find {} observation(s) to delete. Make sure you provide a sufficiently long exact substring match.",
|
|
missing_observations.len()
|
|
)));
|
|
}
|
|
|
|
Ok("Observations deleted".to_string())
|
|
}
|
|
}
|
|
|
|
pub struct DeleteRelationsHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for DeleteRelationsHandler {
|
|
fn name(&self) -> &'static str {
|
|
"delete_relations"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<DeleteRelationsTool>(
|
|
"delete_relations",
|
|
"Delete specific relationships between entities in the knowledge graph by specifying source 'from', target 'to', and 'relation_type'.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: DeleteRelationsTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let mut deleted_count = 0;
|
|
let requested_count = req.relations.len();
|
|
|
|
state.modify_graph(|master| {
|
|
let initial_len = master.relations.len();
|
|
master.relations.retain(|r| {
|
|
let should_delete = req.relations.iter().any(|target| {
|
|
if target.from == r.from && target.to == r.to && target.relation_type == r.relation_type {
|
|
return true;
|
|
}
|
|
target.from.eq_ignore_ascii_case(&r.from)
|
|
&& target.to.eq_ignore_ascii_case(&r.to)
|
|
&& target.relation_type.eq_ignore_ascii_case(&r.relation_type)
|
|
});
|
|
!should_delete
|
|
});
|
|
deleted_count = initial_len - master.relations.len();
|
|
});
|
|
|
|
let missing_count = requested_count.saturating_sub(deleted_count);
|
|
let mut msg = format!("Successfully deleted {} relations.", deleted_count);
|
|
if missing_count > 0 {
|
|
msg.push_str(&format!(" Note: {} relations were not found and skipped.", missing_count));
|
|
}
|
|
Ok(msg)
|
|
}
|
|
}
|
|
|
|
pub struct ReadGraphHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for ReadGraphHandler {
|
|
fn name(&self) -> &'static str {
|
|
"read_graph"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<ReadGraphTool>(
|
|
"read_graph",
|
|
"Read entities and relations from the knowledge graph with optional namespace filtering and token truncation. For large graphs, specify 'namespace' or use 'search_nodes' or 'get_subgraph' for targeted discovery.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let max_tokens = req.max_tokens;
|
|
let limit = req.limit;
|
|
let offset = req.offset.unwrap_or(0);
|
|
|
|
let data = state.read_graph(|full| -> crate::error::Result<String> {
|
|
let filtered_entities: Vec<(&String, &crate::models::Entity)> = full
|
|
.entities
|
|
.iter()
|
|
.filter(|(_, v)| req.namespace.as_ref().is_none_or(|ns| &v.namespace == ns))
|
|
.filter(|(_, v)| req.entity_type.as_ref().is_none_or(|et| &v.entity_type == et))
|
|
.collect();
|
|
|
|
let total_count = filtered_entities.len();
|
|
|
|
let (entities_map, relations_vec, meta_json) = if limit.is_some() || req.offset.is_some() {
|
|
let effective_limit = limit.unwrap_or(50);
|
|
let paged_entities: std::collections::HashMap<&String, &crate::models::Entity> = filtered_entities
|
|
.into_iter()
|
|
.skip(offset)
|
|
.take(effective_limit)
|
|
.collect();
|
|
|
|
let has_more = offset + paged_entities.len() < total_count;
|
|
|
|
let relevant_relations: Vec<&crate::models::Relation> = full
|
|
.relations
|
|
.iter()
|
|
.filter(|r| {
|
|
let ns_match = req.namespace.as_ref().is_none_or(|ns| &r.namespace == ns);
|
|
ns_match && (paged_entities.contains_key(&r.from) || paged_entities.contains_key(&r.to))
|
|
})
|
|
.collect();
|
|
|
|
let meta = serde_json::json!({
|
|
"total_count": total_count,
|
|
"offset": offset,
|
|
"limit": effective_limit,
|
|
"has_more": has_more,
|
|
});
|
|
(paged_entities, relevant_relations, Some(meta))
|
|
} else {
|
|
let ent_map: std::collections::HashMap<&String, &crate::models::Entity> =
|
|
filtered_entities.into_iter().collect();
|
|
let rel_vec: Vec<&crate::models::Relation> = full
|
|
.relations
|
|
.iter()
|
|
.filter(|r| req.namespace.as_ref().is_none_or(|ns| &r.namespace == ns))
|
|
.collect();
|
|
(ent_map, rel_vec, None)
|
|
};
|
|
|
|
let serialized = if let Some(meta) = &meta_json {
|
|
serde_json::json!({
|
|
"entities": entities_map,
|
|
"relations": relations_vec,
|
|
"total_count": meta["total_count"],
|
|
"offset": meta["offset"],
|
|
"limit": meta["limit"],
|
|
"has_more": meta["has_more"],
|
|
})
|
|
} else {
|
|
serde_json::json!({
|
|
"entities": entities_map,
|
|
"relations": relations_vec,
|
|
})
|
|
};
|
|
|
|
let mut result_json = serde_yaml::to_string(&serialized)?;
|
|
|
|
if let Some(max_tok) = max_tokens {
|
|
let max_chars = max_tok * 4;
|
|
if result_json.len() > max_chars {
|
|
let total_ent = entities_map.len();
|
|
let total_rel = relations_vec.len();
|
|
let frac = (max_chars as f64) / (result_json.len() as f64);
|
|
let mut keep_count = (total_ent as f64 * frac * 0.85).max(1.0) as usize;
|
|
|
|
let ent_vec: Vec<_> = entities_map.into_iter().collect();
|
|
|
|
while keep_count > 0 {
|
|
let sliced_ent: std::collections::HashMap<_, _> = ent_vec.iter().take(keep_count).cloned().collect();
|
|
let sliced_rel: Vec<_> = relations_vec.iter().filter(|r| sliced_ent.contains_key(&r.from) && sliced_ent.contains_key(&r.to)).copied().collect();
|
|
|
|
let mut truncated_val = serde_json::json!({
|
|
"entities": sliced_ent,
|
|
"relations": sliced_rel,
|
|
"_meta": {
|
|
"truncated": true,
|
|
"total_entities": total_ent,
|
|
"returned_entities": sliced_ent.len(),
|
|
"total_relations": total_rel,
|
|
"returned_relations": sliced_rel.len(),
|
|
"note": "Output structurally truncated to fit within max_tokens. Use pagination (limit, offset), search_nodes, or get_subgraph for targeted discovery."
|
|
}
|
|
});
|
|
|
|
if let Some(meta) = &meta_json
|
|
&& let Some(obj) = truncated_val.as_object_mut() {
|
|
obj.insert("total_count".to_string(), meta["total_count"].clone());
|
|
obj.insert("offset".to_string(), meta["offset"].clone());
|
|
obj.insert("limit".to_string(), meta["limit"].clone());
|
|
obj.insert("has_more".to_string(), serde_json::Value::Bool(true));
|
|
}
|
|
|
|
if let Ok(trunc_str) = serde_yaml::to_string(&truncated_val)
|
|
&& (trunc_str.len() <= max_chars || keep_count == 1) {
|
|
result_json = trunc_str;
|
|
break;
|
|
}
|
|
keep_count = keep_count.saturating_sub(keep_count / 3 + 1);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(result_json)
|
|
})?;
|
|
Ok(data)
|
|
}
|
|
}
|
|
pub struct SearchNodesHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for SearchNodesHandler {
|
|
fn name(&self) -> &'static str {
|
|
"search_nodes"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<SearchNodesTool>(
|
|
"search_nodes",
|
|
"Search knowledge graph entities and nodes by name or observation query. Returns matching entity names, types, and observations.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: SearchNodesTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let limit = req.limit.unwrap_or(10);
|
|
let include_body = req.include_body.unwrap_or(false);
|
|
|
|
let idx = state.get_search_index().await;
|
|
let matches = idx
|
|
.search(&req.query, req.namespace.as_deref())
|
|
.unwrap_or_default();
|
|
|
|
let data = state.read_graph(|full| -> crate::error::Result<String> {
|
|
let mut matched_entities = Vec::new();
|
|
for (id, doc_type, _, _, _) in matches.iter().take(limit) {
|
|
if doc_type == "entity"
|
|
&& let Some(e) = full.entities.get(id)
|
|
{
|
|
if include_body {
|
|
matched_entities.push(serde_json::to_value(e)?);
|
|
} else {
|
|
matched_entities.push(serde_json::json!({
|
|
"name": e.name,
|
|
"entity_type": e.entity_type,
|
|
"namespace": e.namespace,
|
|
"git_branch": e.git_branch,
|
|
"observations_count": e.observations.len()
|
|
}));
|
|
}
|
|
}
|
|
}
|
|
if matched_entities.is_empty() { Ok::<String, crate::error::AppError>(format!("No entities found matching query. Try using a different keyword or relaxing the search query.")) } else { Ok::<String, crate::error::AppError>(serde_yaml::to_string(&matched_entities)?) }
|
|
})?;
|
|
Ok(data)
|
|
}
|
|
}
|
|
|
|
pub struct OpenNodesHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for OpenNodesHandler {
|
|
fn name(&self) -> &'static str {
|
|
"open_nodes"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<OpenNodesTool>(
|
|
"open_nodes",
|
|
"Retrieve complete node details, observations, and connected relations for specific entities in the knowledge graph.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: OpenNodesTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let data = state.read_graph(|full| -> crate::error::Result<String> {
|
|
let targets: HashSet<&str> = req.names.iter().map(|s| s.as_str()).collect();
|
|
let mut result = BorrowedGraph::default();
|
|
let mut connected: HashSet<&str> = HashSet::new();
|
|
for r in &full.relations {
|
|
if targets.contains(r.from.as_str()) {
|
|
connected.insert(r.to.as_str());
|
|
result.relations.push(r);
|
|
} else if targets.contains(r.to.as_str()) {
|
|
connected.insert(r.from.as_str());
|
|
result.relations.push(r);
|
|
}
|
|
}
|
|
for (name, e) in &full.entities {
|
|
if targets.contains(name.as_str()) || connected.contains(name.as_str()) {
|
|
result.entities.insert(name, e);
|
|
}
|
|
}
|
|
if result.entities.is_empty() { Ok::<String, crate::error::AppError>(format!("No matching nodes found.")) } else { Ok::<String, crate::error::AppError>(serde_yaml::to_string(&result)?) }
|
|
})?;
|
|
Ok(data)
|
|
}
|
|
}
|
|
|
|
pub struct VisualizeGraphHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for VisualizeGraphHandler {
|
|
fn name(&self) -> &'static str {
|
|
"visualize_graph"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<VisualizeGraphTool>(
|
|
"visualize_graph",
|
|
"Generate a visual Mermaid representation of the knowledge graph topology or a filtered query subset.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: VisualizeGraphTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let query = req.query.unwrap_or_default();
|
|
let mut output = state.read_graph(|full| {
|
|
let mut included: HashSet<&str> = HashSet::new();
|
|
let mut to_draw = Vec::new();
|
|
|
|
for (name, e) in &full.entities {
|
|
if let Some(ns) = &req.namespace
|
|
&& e.namespace != *ns
|
|
{
|
|
continue;
|
|
}
|
|
if query.is_empty()
|
|
|| contains_ignore_ascii_case(name, &query)
|
|
|| contains_ignore_ascii_case(&e.entity_type, &query)
|
|
{
|
|
included.insert(name.as_str());
|
|
}
|
|
}
|
|
|
|
for r in &full.relations {
|
|
if let Some(ns) = &req.namespace
|
|
&& r.namespace != *ns
|
|
{
|
|
continue;
|
|
}
|
|
if query.is_empty()
|
|
|| included.contains(r.from.as_str())
|
|
|| included.contains(r.to.as_str())
|
|
{
|
|
included.insert(r.from.as_str());
|
|
included.insert(r.to.as_str());
|
|
to_draw.push(r);
|
|
}
|
|
}
|
|
|
|
let mut out = String::with_capacity(included.len() * 40 + to_draw.len() * 60);
|
|
out.push_str("graph TD;\n");
|
|
|
|
let sanitize_to = |out_str: &mut String, s: &str, id_mode: bool| {
|
|
for c in s.chars() {
|
|
if c != '"' && c != '(' && c != ')' {
|
|
if id_mode && (c == ' ' || c == '-' || c == '.') {
|
|
out_str.push('_');
|
|
} else {
|
|
out_str.push(c);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
for name in &included {
|
|
out.push_str(" id_");
|
|
sanitize_to(&mut out, name, true);
|
|
out.push_str("[\"");
|
|
sanitize_to(&mut out, name, false);
|
|
out.push_str("\"];\n");
|
|
}
|
|
for r in to_draw {
|
|
out.push_str(" id_");
|
|
sanitize_to(&mut out, &r.from, true);
|
|
out.push_str("-->|\"");
|
|
out.push_str(&r.relation_type.replace("\"", ""));
|
|
out.push_str("\"|id_");
|
|
sanitize_to(&mut out, &r.to, true);
|
|
out.push_str(";\n");
|
|
}
|
|
out
|
|
});
|
|
if output == "graph TD;\n" {
|
|
output = "No nodes found to visualize.".to_string();
|
|
}
|
|
Ok(output.to_string())
|
|
}
|
|
}
|
|
|
|
pub struct CondenseEntityHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for CondenseEntityHandler {
|
|
fn name(&self) -> &'static str {
|
|
"condense_entity"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<CondenseEntityTool>(
|
|
"condense_entity",
|
|
"Condense, summarize, and deduplicate redundant observations on an entity to optimize context size while preserving key facts.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: CondenseEntityTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let mut missing = false;
|
|
|
|
if req.summarized_observations.is_empty() {
|
|
let mut obs_to_condense = Vec::new();
|
|
state.read_graph(|g| {
|
|
if let Some(e) = g.entities.get(&req.entity_name) {
|
|
obs_to_condense = e.observations.clone();
|
|
} else {
|
|
missing = true;
|
|
}
|
|
});
|
|
if missing {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Entity '{}' not found. Please verify the exact entity name using search_nodes.",
|
|
req.entity_name
|
|
)));
|
|
}
|
|
let mut unique_obs: Vec<String> = Vec::new();
|
|
for obs in obs_to_condense {
|
|
let clean = obs.trim().to_string();
|
|
if !clean.is_empty() && !unique_obs.contains(&clean) {
|
|
unique_obs.push(clean);
|
|
}
|
|
}
|
|
if unique_obs.len() > 3 && state.ollama.is_available().await {
|
|
let prompt = format!(
|
|
"Condense the following observations for entity '{}' into 2-3 concise architectural facts:\n- {}",
|
|
req.entity_name,
|
|
unique_obs.join("\n- ")
|
|
);
|
|
if let Ok(summary) = state.ollama.generate(&prompt, None, None, None).await {
|
|
let lines: Vec<String> = summary
|
|
.lines()
|
|
.map(|l| l.trim().trim_start_matches('-').trim().to_string())
|
|
.filter(|l| !l.is_empty())
|
|
.collect();
|
|
if !lines.is_empty() {
|
|
unique_obs = lines;
|
|
}
|
|
}
|
|
}
|
|
state.modify_graph(|master| {
|
|
if let Some(e) = master.entities.get_mut(&req.entity_name) {
|
|
e.observations = unique_obs;
|
|
}
|
|
});
|
|
return Ok(format!(
|
|
"Entity '{}' observations condensed with deduplication/AI summary.",
|
|
req.entity_name
|
|
));
|
|
}
|
|
|
|
state.modify_graph(|master| {
|
|
if let Some(e) = master.entities.get_mut(&req.entity_name) {
|
|
e.observations = req.summarized_observations;
|
|
} else {
|
|
missing = true;
|
|
}
|
|
});
|
|
if missing {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Entity '{}' not found. Please verify the exact entity name using search_nodes.",
|
|
req.entity_name
|
|
)));
|
|
}
|
|
Ok("Entity condensed".to_string())
|
|
}
|
|
}
|
|
|
|
pub struct MergeEntitiesHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for MergeEntitiesHandler {
|
|
fn name(&self) -> &'static str {
|
|
"merge_entities"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<MergeEntitiesTool>(
|
|
"merge_entities",
|
|
"Merge two duplicate entities into a primary target entity, consolidating all observations and re-pointing relations.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: MergeEntitiesTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let mut missing = false;
|
|
state.modify_graph(|master| {
|
|
if let Some(src) = master.entities.remove(&req.source_entity) {
|
|
if let Some(tgt) = master.entities.get_mut(&req.target_entity) {
|
|
tgt.observations.extend(src.observations);
|
|
tgt.observations.sort();
|
|
tgt.observations.dedup();
|
|
} else {
|
|
let mut new_tgt = src.clone();
|
|
new_tgt.name = req.target_entity.clone();
|
|
master.entities.insert(req.target_entity.clone(), new_tgt);
|
|
}
|
|
} else {
|
|
missing = true;
|
|
}
|
|
let mut seen = std::collections::HashSet::new();
|
|
master.relations.retain_mut(|r| {
|
|
if r.from == req.source_entity {
|
|
r.from = req.target_entity.clone();
|
|
}
|
|
if r.to == req.source_entity {
|
|
r.to = req.target_entity.clone();
|
|
}
|
|
|
|
// Filter out self-loops
|
|
if r.from == r.to {
|
|
return false;
|
|
}
|
|
|
|
if r.from == req.target_entity || r.to == req.target_entity {
|
|
seen.insert(r.clone())
|
|
} else {
|
|
true
|
|
}
|
|
});
|
|
});
|
|
if missing {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Source entity '{}' not found. Please verify the exact entity name using search_nodes.",
|
|
req.source_entity
|
|
)));
|
|
}
|
|
Ok("Entities merged".to_string())
|
|
}
|
|
}
|
|
|
|
pub struct FindOrphansHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for FindOrphansHandler {
|
|
fn name(&self) -> &'static str {
|
|
"find_orphans"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<FindOrphansTool>(
|
|
"find_orphans",
|
|
"Find orphan nodes in the knowledge graph that have no incoming or outgoing relationships.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let orphans = state.read_graph(|full| {
|
|
let mut connected = std::collections::HashSet::new();
|
|
for r in &full.relations {
|
|
connected.insert(r.from.as_str());
|
|
connected.insert(r.to.as_str());
|
|
}
|
|
full.entities
|
|
.keys()
|
|
.filter(|k| !connected.contains(k.as_str()))
|
|
.cloned()
|
|
.collect::<Vec<String>>()
|
|
});
|
|
let data = Ok::<String, crate::error::AppError>(serde_yaml::to_string(&orphans)?)?;
|
|
Ok(data)
|
|
}
|
|
}
|
|
|
|
pub struct GetSubgraphHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for GetSubgraphHandler {
|
|
fn name(&self) -> &'static str {
|
|
"get_subgraph"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<GetSubgraphTool>(
|
|
"get_subgraph",
|
|
"Extract a connected neighborhood subgraph centered around specified seed nodes up to a given hop depth.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: GetSubgraphTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let root = req.root_entity.ok_or_else(|| {
|
|
crate::error::AppError::Internal("root_entity is required".to_string())
|
|
})?;
|
|
let depth = req.depth.unwrap_or(2);
|
|
let format = req.format.unwrap_or(SubgraphFormat::MarkdownTree);
|
|
|
|
let (matched_entities, matched_relations) = state.read_graph(|full| {
|
|
let mut adj: std::collections::HashMap<&str, Vec<&str>> =
|
|
std::collections::HashMap::with_capacity(full.relations.len() * 2);
|
|
for r in &full.relations {
|
|
adj.entry(r.from.as_str()).or_default().push(r.to.as_str());
|
|
adj.entry(r.to.as_str()).or_default().push(r.from.as_str());
|
|
}
|
|
|
|
let mut visited = std::collections::HashSet::new();
|
|
let mut queue = std::collections::VecDeque::new();
|
|
queue.push_back((root.as_str(), 0));
|
|
visited.insert(root.clone());
|
|
|
|
while let Some((curr, d)) = queue.pop_front() {
|
|
if d < depth
|
|
&& let Some(neighbors) = adj.get(curr)
|
|
{
|
|
for &next in neighbors {
|
|
if visited.insert(next.to_string()) {
|
|
queue.push_back((next, d + 1));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let ent: std::collections::HashMap<String, crate::models::Entity> = visited
|
|
.iter()
|
|
.filter_map(|name| full.entities.get(name).map(|e| (name.clone(), e.clone())))
|
|
.collect();
|
|
let rel: Vec<crate::models::Relation> = full
|
|
.relations
|
|
.iter()
|
|
.filter(|r| visited.contains(&r.from) && visited.contains(&r.to))
|
|
.cloned()
|
|
.collect();
|
|
|
|
(ent, rel)
|
|
});
|
|
|
|
if format == SubgraphFormat::MarkdownTree {
|
|
let mut markdown = format!("# Subgraph Topology for `{}`\n\n", root);
|
|
markdown.push_str("## Entities\n");
|
|
for e in matched_entities.values() {
|
|
markdown.push_str(&format!("- **{}** (`{}`)\n", e.name, e.entity_type));
|
|
for obs in &e.observations {
|
|
markdown.push_str(&format!(" - {}\n", obs));
|
|
}
|
|
}
|
|
|
|
markdown.push_str("\n## Relations\n");
|
|
for r in &matched_relations {
|
|
markdown.push_str(&format!(
|
|
"- `{}` --[{}]--> `{}`\n",
|
|
r.from, r.relation_type, r.to
|
|
));
|
|
}
|
|
return Ok(markdown);
|
|
}
|
|
|
|
let result = serde_json::json!({
|
|
"root_entity": root,
|
|
"depth": depth,
|
|
"entities": matched_entities,
|
|
"relations": matched_relations,
|
|
});
|
|
|
|
Ok(serde_yaml::to_string(&result)?)
|
|
}
|
|
}
|
|
|
|
use crate::handlers::utils::*;
|
|
|
|
pub struct SweepGraphHealthHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for SweepGraphHealthHandler {
|
|
fn name(&self) -> &'static str {
|
|
"sweep_graph_health"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<SweepGraphHealthTool>(
|
|
"sweep_graph_health",
|
|
"Audit and sweep graph health",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: SweepGraphHealthTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let auto_prune = req.auto_prune_orphans.unwrap_or(false);
|
|
let auto_prune_stale = req.auto_prune_stale_files.unwrap_or(false);
|
|
|
|
let mut orphans = Vec::new();
|
|
let mut duplicates = Vec::new();
|
|
let mut stale_entities = Vec::new();
|
|
|
|
state.modify_graph(|g| {
|
|
// 1. Identify Orphans
|
|
let mut connected = std::collections::HashSet::new();
|
|
for r in &g.relations {
|
|
connected.insert(r.from.clone());
|
|
connected.insert(r.to.clone());
|
|
}
|
|
|
|
for name in g.entities.keys() {
|
|
if !connected.contains(name) {
|
|
orphans.push(name.clone());
|
|
}
|
|
}
|
|
|
|
// ADR-0111: Automated Stale Symbol Pruning & Graph Tombstoning
|
|
for (name, entity) in g.entities.iter() {
|
|
if entity.entity_type == "File" || name.contains(".rs") || name.contains(".ts") || name.contains(".js") || name.contains(".py") {
|
|
let check_path = entity.file_path.as_deref().unwrap_or(name.as_str());
|
|
// Strip any symbol qualifiers like File::symbol
|
|
let base_file = check_path.split("::").next().unwrap_or(check_path);
|
|
if (base_file.contains('/') || base_file.contains('\\') || base_file.ends_with(".rs") || base_file.ends_with(".ts")) && !std::path::Path::new(base_file).exists() {
|
|
stale_entities.push(name.clone());
|
|
}
|
|
}
|
|
}
|
|
|
|
if auto_prune {
|
|
for orphan in &orphans {
|
|
g.entities.remove(orphan);
|
|
}
|
|
}
|
|
|
|
if auto_prune_stale {
|
|
for stale in &stale_entities {
|
|
g.entities.remove(stale);
|
|
g.relations.retain(|r| &r.from != stale && &r.to != stale);
|
|
}
|
|
}
|
|
|
|
// 2. Compute similarity pairs for duplicate detection using pre-computed lowercase names
|
|
let names: Vec<_> = g.entities.keys().cloned().collect();
|
|
let lower_names: Vec<String> = names.iter().map(|n| n.to_lowercase()).collect();
|
|
for i in 0..names.len() {
|
|
for j in (i + 1)..names.len() {
|
|
let n1 = &names[i];
|
|
let n2 = &names[j];
|
|
let l1 = &lower_names[i];
|
|
let l2 = &lower_names[j];
|
|
|
|
if l1 == l2 || ((l1.contains(l2.as_str()) || l2.contains(l1.as_str())) && l1.len().min(l2.len()) > 3) {
|
|
duplicates.push(serde_json::json!({
|
|
"entity_a": n1,
|
|
"entity_b": n2,
|
|
"suggested_action": format!("merge_entities(source: '{}', target: '{}')", n1, n2)
|
|
}));
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
let mut out = String::new();
|
|
out.push_str("## Graph Health Report\n\n");
|
|
let score = if orphans.is_empty() && duplicates.is_empty() && stale_entities.is_empty() { "100% (Healthy)" } else { "Needs Maintenance" };
|
|
out.push_str(&format!("**Health Status:** {}\n\n", score));
|
|
|
|
out.push_str(&format!("### Orphaned Entities ({})\n", orphans.len()));
|
|
if auto_prune {
|
|
out.push_str("*These orphans were automatically pruned.*\n");
|
|
}
|
|
for o in orphans {
|
|
out.push_str(&format!("- `{}`\n", o));
|
|
}
|
|
out.push_str("\n");
|
|
|
|
out.push_str(&format!("### Stale Entities ({})\n", stale_entities.len()));
|
|
if auto_prune_stale {
|
|
out.push_str("*These stale entities were automatically pruned.*\n");
|
|
}
|
|
for s in stale_entities {
|
|
out.push_str(&format!("- `{}`\n", s));
|
|
}
|
|
out.push_str("\n");
|
|
|
|
out.push_str(&format!("### Potential Duplicates ({})\n", duplicates.len()));
|
|
for d in duplicates {
|
|
let a = d.get("entity_a").and_then(|x| x.as_str()).unwrap_or("?");
|
|
let b = d.get("entity_b").and_then(|x| x.as_str()).unwrap_or("?");
|
|
let act = d.get("suggested_action").and_then(|x| x.as_str()).unwrap_or("");
|
|
out.push_str(&format!("- `{}` and `{}`\n - Suggested: `{}`\n", a, b, act));
|
|
}
|
|
|
|
Ok(out)
|
|
}
|
|
}
|
|
|
|
pub struct ResolveStaleSymbolsHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for ResolveStaleSymbolsHandler {
|
|
fn name(&self) -> &'static str {
|
|
"resolve_stale_symbols"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<ResolveStaleSymbolsTool>(
|
|
"resolve_stale_symbols",
|
|
"Inspect Knowledge Graph entities and tech debt symbol/line references against files on disk, healing stale pointers.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: ResolveStaleSymbolsTool =
|
|
serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let auto_heal = req.auto_heal.unwrap_or(true);
|
|
|
|
let mut checked = 0;
|
|
let mut stale_entities = Vec::new();
|
|
|
|
state.modify_graph(|g| {
|
|
for name in g.entities.keys() {
|
|
if let Some(target) = &req.target
|
|
&& name != target
|
|
&& !name.contains(target)
|
|
{
|
|
continue;
|
|
}
|
|
checked += 1;
|
|
let path = std::path::Path::new(name);
|
|
if (name.contains('/') || name.contains('\\') || name.contains('.'))
|
|
&& !path.exists()
|
|
{
|
|
stale_entities.push(name.clone());
|
|
}
|
|
}
|
|
|
|
if auto_heal {
|
|
for stale in &stale_entities {
|
|
g.entities.remove(stale);
|
|
g.relations.retain(|r| &r.from != stale && &r.to != stale);
|
|
}
|
|
}
|
|
});
|
|
|
|
let report = serde_json::json!({
|
|
"checked_count": checked,
|
|
"stale_entities_found": stale_entities,
|
|
"auto_healed": auto_heal,
|
|
"status": "success"
|
|
});
|
|
|
|
Ok(serde_yaml::to_string(&report)?)
|
|
}
|
|
}
|
|
|
|
pub struct SummarizeSubgraphHandler;
|
|
|
|
#[async_trait]
|
|
impl McpTool for SummarizeSubgraphHandler {
|
|
fn name(&self) -> &'static str {
|
|
"summarize_subgraph"
|
|
}
|
|
|
|
fn schema(&self) -> Value {
|
|
crate::mcp::tool_def::<SummarizeSubgraphTool>(
|
|
"summarize_subgraph",
|
|
"Generate compact, LLM-optimized Markdown topology of a graph component capped within a token budget.",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: SummarizeSubgraphTool = serde_json::from_value(args).map_err(|e| crate::error::AppError::BadRequest(format!("Schema validation failed. Your JSON arguments do not match the expected tool schema: {}", e)))?;
|
|
let depth = req.depth.unwrap_or(2);
|
|
|
|
let (sub_entities, sub_relations) = state.read_graph(|g| {
|
|
let mut visited = std::collections::HashSet::new();
|
|
let mut queue = std::collections::VecDeque::new();
|
|
queue.push_back((req.root_entity.clone(), 0));
|
|
visited.insert(req.root_entity.clone());
|
|
|
|
while let Some((curr, d)) = queue.pop_front() {
|
|
if d >= depth {
|
|
continue;
|
|
}
|
|
for r in &g.relations {
|
|
if r.from == curr && !visited.contains(&r.to) {
|
|
visited.insert(r.to.clone());
|
|
queue.push_back((r.to.clone(), d + 1));
|
|
} else if r.to == curr && !visited.contains(&r.from) {
|
|
visited.insert(r.from.clone());
|
|
queue.push_back((r.from.clone(), d + 1));
|
|
}
|
|
}
|
|
}
|
|
|
|
let ent: Vec<_> = g
|
|
.entities
|
|
.values()
|
|
.filter(|e| visited.contains(&e.name))
|
|
.cloned()
|
|
.collect();
|
|
let rel: Vec<_> = g
|
|
.relations
|
|
.iter()
|
|
.filter(|r| visited.contains(&r.from) && visited.contains(&r.to))
|
|
.cloned()
|
|
.collect();
|
|
(ent, rel)
|
|
});
|
|
|
|
let mut markdown = format!("# Subgraph Topology for `{}`\n\n", req.root_entity);
|
|
markdown.push_str("## Entities\n");
|
|
for e in &sub_entities {
|
|
markdown.push_str(&format!("- **{}** (`{}`)\n", e.name, e.entity_type));
|
|
for obs in &e.observations {
|
|
markdown.push_str(&format!(" - {}\n", obs));
|
|
}
|
|
}
|
|
|
|
markdown.push_str("\n## Relations\n");
|
|
for r in &sub_relations {
|
|
markdown.push_str(&format!(
|
|
"- `{}` --[{}]--> `{}`\n",
|
|
r.from, r.relation_type, r.to
|
|
));
|
|
}
|
|
|
|
let max_tokens = req.max_tokens.unwrap_or(1000);
|
|
let max_chars = max_tokens * 4;
|
|
if markdown.len() > max_chars {
|
|
let valid_boundary = markdown.floor_char_boundary(max_chars);
|
|
markdown.truncate(valid_boundary);
|
|
markdown.push_str("\n... [Truncated to fit token budget]");
|
|
}
|
|
|
|
Ok(markdown)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::handlers::meta::AgentSignalsHandler;
|
|
use serde_json::json;
|
|
|
|
#[tokio::test]
|
|
async fn test_create_and_read_entities() {
|
|
let state = Arc::new(MemoryState::new_in_memory());
|
|
|
|
let create_handler = CreateEntitiesHandler;
|
|
let args = json!({
|
|
"entities": [
|
|
{"name": "Alice", "entity_type": "Person", "observations": ["Likes Bob"]}
|
|
]
|
|
});
|
|
|
|
let res = create_handler
|
|
.execute(args, state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert!(res.contains("Successfully created 1 entity/entities"));
|
|
|
|
// Ensure graph contains the entity
|
|
state.graph.read_with(|g| {
|
|
assert!(g.entities.contains_key("Alice"));
|
|
assert_eq!(g.entities["Alice"].entity_type, "Person");
|
|
});
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_create_relations() {
|
|
let state = Arc::new(MemoryState::new_in_memory());
|
|
|
|
// Needs entities first
|
|
state.graph.modify(|g| {
|
|
g.entities.insert(
|
|
"A".to_string(),
|
|
crate::models::Entity {
|
|
name: "A".to_string(),
|
|
entity_type: "Node".to_string(),
|
|
observations: vec![],
|
|
namespace: "global".to_string(),
|
|
git_branch: None,
|
|
..Default::default()
|
|
},
|
|
);
|
|
g.entities.insert(
|
|
"B".to_string(),
|
|
crate::models::Entity {
|
|
name: "B".to_string(),
|
|
entity_type: "Node".to_string(),
|
|
observations: vec![],
|
|
namespace: "global".to_string(),
|
|
git_branch: None,
|
|
..Default::default()
|
|
},
|
|
);
|
|
});
|
|
|
|
let handler = CreateRelationsHandler;
|
|
let args = json!({
|
|
"relations": [
|
|
{"from": "A", "to": "B", "relation_type": "knows"}
|
|
]
|
|
});
|
|
let res = handler
|
|
.execute(args, state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert!(res.contains("Successfully created 1 relation(s)"));
|
|
|
|
let alias_args = json!({
|
|
"relations": [
|
|
{"from": "A", "to": "B", "relation_type": "knows"}
|
|
]
|
|
});
|
|
let alias_res = handler.execute(alias_args, state.clone()).await.unwrap();
|
|
assert!(alias_res.contains("Successfully created 1 relation(s)"));
|
|
|
|
// Test semantic LLM schema feedback on missing fields
|
|
let bad_args = json!({
|
|
"relations": [
|
|
{"invalid_field": "X"}
|
|
]
|
|
});
|
|
let err_res = handler
|
|
.execute(bad_args, state.clone())
|
|
.await
|
|
.unwrap_err()
|
|
.to_string();
|
|
assert!(err_res.contains("Schema error:"));
|
|
assert!(err_res.contains("strictly uses 'from', 'to', and 'relation_type'"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_observations_and_reads() {
|
|
let state = Arc::new(MemoryState::new_in_memory());
|
|
|
|
// Need entity first
|
|
state.graph.modify(|g| {
|
|
g.entities.insert(
|
|
"A".to_string(),
|
|
crate::models::Entity {
|
|
name: "A".to_string(),
|
|
entity_type: "Node".to_string(),
|
|
observations: vec![],
|
|
namespace: "global".to_string(),
|
|
git_branch: None,
|
|
..Default::default()
|
|
},
|
|
);
|
|
});
|
|
|
|
let add_obs = AddObservationsHandler;
|
|
let args_obs = json!({
|
|
"observations": [
|
|
{"entity_name": "A", "contents": ["Obs 1", "Obs 2"]}
|
|
]
|
|
});
|
|
let res1 = add_obs
|
|
.execute(args_obs, state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert_eq!(res1, "Observations added");
|
|
|
|
let read_graph = ReadGraphHandler;
|
|
let res2 = read_graph
|
|
.execute(json!({"namespace": "global"}), state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert!(res2.contains("Obs 1"));
|
|
assert!(res2.contains("Obs 2"));
|
|
|
|
let del_entity = DeleteEntitiesHandler;
|
|
let res4 = del_entity
|
|
.execute(json!({"entity_names": ["A"]}), state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert_eq!(res4, "Successfully deleted 1 entities.");
|
|
|
|
let res5 = read_graph
|
|
.execute(json!({"namespace": "global"}), state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert!(!res5.contains("A"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_advanced_graph_operations() {
|
|
let state = Arc::new(MemoryState::new_in_memory());
|
|
|
|
let create_handler = CreateEntitiesHandler;
|
|
let args_ent = json!({
|
|
"entities": [
|
|
{"name": "X", "entity_type": "File", "observations": ["Obs X"], "namespace": "global"},
|
|
{"name": "Y", "entity_type": "File", "observations": ["Obs Y"], "namespace": "global"}
|
|
]
|
|
});
|
|
create_handler
|
|
.execute(args_ent, state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
|
|
let rel_handler = CreateRelationsHandler;
|
|
let args_rel = json!({
|
|
"relations": [
|
|
{"from": "X", "to": "Y", "relation_type": "depends_on", "namespace": "global"}
|
|
]
|
|
});
|
|
rel_handler
|
|
.execute(args_rel, state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
|
|
let read_handler = ReadGraphHandler;
|
|
let res_read = read_handler
|
|
.execute(json!({"namespace": "global"}), state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert!(res_read.contains("X"));
|
|
assert!(res_read.contains("depends_on"));
|
|
|
|
let open_handler = OpenNodesHandler;
|
|
let res_open = open_handler
|
|
.execute(json!({"names": ["X"]}), state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert!(res_open.contains("Y"));
|
|
|
|
let viz_handler = VisualizeGraphHandler;
|
|
let res_viz = viz_handler
|
|
.execute(json!({"query": "X"}), state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert!(!res_viz.is_empty());
|
|
|
|
let condense = CondenseEntityHandler;
|
|
let res_cond = condense
|
|
.execute(
|
|
json!({"entity_name": "X", "summarized_observations": ["X condensed"]}),
|
|
state.clone(),
|
|
)
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert_eq!(res_cond, "Entity condensed");
|
|
|
|
let merge = MergeEntitiesHandler;
|
|
let res_merge = merge
|
|
.execute(
|
|
json!({"source_entity": "X", "target_entity": "Y"}),
|
|
state.clone(),
|
|
)
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
assert_eq!(res_merge, "Entities merged");
|
|
|
|
let orphans = FindOrphansHandler;
|
|
let res_orphans = orphans
|
|
.execute(json!({}), state.clone())
|
|
.await
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
|
.unwrap();
|
|
// After merging X into Y and purging self-loops, Y is the sole node and becomes an orphan
|
|
assert!(res_orphans.contains("Y"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_more_graph_handlers() {
|
|
let state = Arc::new(MemoryState::new_in_memory());
|
|
|
|
let create_handler = CreateEntitiesHandler;
|
|
let args_ent = json!({
|
|
"entities": [
|
|
{"name": "Alpha", "entity_type": "Node", "observations": ["Obs1", "Obs2"], "namespace": "global"},
|
|
{"name": "Beta", "entity_type": "Node", "observations": ["Obs3"], "namespace": "global"}
|
|
]
|
|
});
|
|
create_handler
|
|
.execute(args_ent, state.clone())
|
|
.await
|
|
.unwrap();
|
|
|
|
let rel_handler = CreateRelationsHandler;
|
|
let args_rel = json!({
|
|
"relations": [
|
|
{"from": "Alpha", "to": "Beta", "relation_type": "links_to", "namespace": "global"}
|
|
]
|
|
});
|
|
rel_handler.execute(args_rel, state.clone()).await.unwrap();
|
|
|
|
let path_handler = QueryGraphPathHandler;
|
|
let path_res = path_handler
|
|
.execute(
|
|
json!({"start_node": "Alpha", "end_node": "Beta"}),
|
|
state.clone(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert!(path_res.contains("Alpha"));
|
|
|
|
let sub_handler = GetSubgraphHandler;
|
|
let sub_res = sub_handler
|
|
.execute(json!({"root_entity": "Alpha"}), state.clone())
|
|
.await
|
|
.unwrap();
|
|
assert!(sub_res.contains("Alpha"));
|
|
|
|
let sum_handler = SummarizeSubgraphHandler;
|
|
let sum_res = sum_handler
|
|
.execute(json!({"root_entity": "Alpha"}), state.clone())
|
|
.await
|
|
.unwrap();
|
|
assert!(!sum_res.is_empty());
|
|
|
|
let sweep_handler = SweepGraphHealthHandler;
|
|
let sweep_res = sweep_handler
|
|
.execute(json!({}), state.clone())
|
|
.await
|
|
.unwrap();
|
|
assert!(!sweep_res.is_empty());
|
|
|
|
let del_obs_handler = DeleteObservationsHandler;
|
|
let del_obs_res = del_obs_handler
|
|
.execute(
|
|
json!({
|
|
"deletions": [{"entity_name": "Alpha", "observations": ["Obs1"]}]
|
|
}),
|
|
state.clone(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(del_obs_res, "Observations deleted");
|
|
|
|
let del_rel_handler = DeleteRelationsHandler;
|
|
let del_rel_res = del_rel_handler
|
|
.execute(
|
|
json!({
|
|
"relations": [{"from": "Alpha", "to": "Beta", "relation_type": "links_to"}]
|
|
}),
|
|
state.clone(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(del_rel_res, "Successfully deleted 1 relations.");
|
|
|
|
let bcast_handler = AgentSignalsHandler;
|
|
let bcast_res = bcast_handler
|
|
.execute(
|
|
json!({
|
|
"action": "broadcast",
|
|
"sender": "agent1",
|
|
"signal_type": "task_completed",
|
|
"payload": "fix_bug"
|
|
}),
|
|
state.clone(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert!(bcast_res.contains("Broadcasted signal"));
|
|
|
|
let qsignal_handler = AgentSignalsHandler;
|
|
let qsignal_res = qsignal_handler
|
|
.execute(
|
|
json!({
|
|
"action": "query",
|
|
"sender": "agent1"
|
|
}),
|
|
state.clone(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert!(qsignal_res.contains("task_completed"));
|
|
let read_paged_handler = ReadGraphHandler;
|
|
let paged_res = read_paged_handler
|
|
.execute(json!({"limit": 1, "offset": 0}), state.clone())
|
|
.await
|
|
.unwrap();
|
|
let parsed_paged: serde_json::Value = serde_yaml::from_str(&paged_res).unwrap();
|
|
assert_eq!(parsed_paged["total_count"], 2);
|
|
assert_eq!(parsed_paged["has_more"], true);
|
|
}
|
|
}
|