1321 lines
48 KiB
Rust
1321 lines
48 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 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", "Execute query_graph_path")
|
|
}
|
|
|
|
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| e.to_string())?;
|
|
state.read_graph(|graph| {
|
|
let max_depth = req.max_depth.unwrap_or(5);
|
|
let mut queue: std::collections::VecDeque<&str> = std::collections::VecDeque::new();
|
|
let mut visited: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
|
let mut parents: std::collections::HashMap<&str, (&str, &str, bool)> =
|
|
std::collections::HashMap::new();
|
|
|
|
queue.push_back(req.start_node.as_str());
|
|
visited.insert(req.start_node.as_str());
|
|
|
|
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 == req.end_node {
|
|
found = true;
|
|
break;
|
|
}
|
|
nodes_at_current_depth -= 1;
|
|
if current_depth < max_depth {
|
|
for rel in &graph.relations {
|
|
if rel.from == current && !visited.contains(rel.to.as_str()) {
|
|
visited.insert(rel.to.as_str());
|
|
parents.insert(
|
|
rel.to.as_str(),
|
|
(current, rel.relation_type.as_str(), false),
|
|
);
|
|
queue.push_back(rel.to.as_str());
|
|
nodes_at_next_depth += 1;
|
|
} else if rel.to == current && !visited.contains(rel.from.as_str()) {
|
|
visited.insert(rel.from.as_str());
|
|
parents.insert(
|
|
rel.from.as_str(),
|
|
(current, rel.relation_type.as_str(), true),
|
|
);
|
|
queue.push_back(rel.from.as_str());
|
|
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 = req.end_node.as_str();
|
|
while curr != req.start_node {
|
|
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();
|
|
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
|
|
))
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
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", "Execute create_entities")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
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);
|
|
inserted.push(entity.clone());
|
|
g.entities.insert(entity.name.clone(), entity);
|
|
}
|
|
}
|
|
});
|
|
let idx = state.get_search_index();
|
|
for entity in inserted {
|
|
drop(idx.index_entity(&entity));
|
|
}
|
|
Ok("Entities created".to_string())
|
|
}
|
|
}
|
|
|
|
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", "Execute create_relations")
|
|
}
|
|
|
|
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.modify_graph(|g| {
|
|
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);
|
|
let from_exists = g.entities.contains_key(&relation.from);
|
|
let to_exists = g.entities.contains_key(&relation.to);
|
|
if from_exists && to_exists {
|
|
g.relations.push(relation);
|
|
} else {
|
|
if !from_exists {
|
|
missing_nodes.insert(relation.from);
|
|
}
|
|
if !to_exists {
|
|
missing_nodes.insert(relation.to);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
if !missing_nodes.is_empty() {
|
|
let missing: Vec<_> = missing_nodes.into_iter().collect();
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Relations dropped due to missing entities: {}",
|
|
missing.join(", ")
|
|
)));
|
|
}
|
|
Ok("Relations created".to_string())
|
|
}
|
|
}
|
|
|
|
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", "Execute add_observations")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: AddObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
let mut missing_entities = Vec::new();
|
|
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);
|
|
} else {
|
|
missing_entities.push(o.entity_name);
|
|
}
|
|
}
|
|
});
|
|
if !missing_entities.is_empty() {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Observations dropped for missing entities: {}",
|
|
missing_entities.join(", ")
|
|
)));
|
|
}
|
|
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", "Execute delete_entities")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: DeleteEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
let to_delete: std::collections::HashSet<_> = req.entity_names.into_iter().collect();
|
|
let mut missing = Vec::new();
|
|
state.modify_graph(|master| {
|
|
for name in &to_delete {
|
|
if master.entities.remove(name).is_none() {
|
|
missing.push(name.clone());
|
|
}
|
|
}
|
|
master
|
|
.relations
|
|
.retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
|
|
});
|
|
|
|
if !missing.is_empty() {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
|
|
missing.join(", ")
|
|
)));
|
|
}
|
|
|
|
let idx = state.get_search_index();
|
|
for name in to_delete {
|
|
drop(idx.delete_document(&name));
|
|
}
|
|
Ok("Entities deleted".to_string())
|
|
}
|
|
}
|
|
|
|
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",
|
|
"Execute delete_observations",
|
|
)
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: DeleteObservationsTool =
|
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
let mut missing = Vec::new();
|
|
state.modify_graph(|master| {
|
|
for d in req.deletions {
|
|
if let Some(e) = master.entities.get_mut(&d.entity_name) {
|
|
let to_rem: HashSet<_> = d.observations.into_iter().collect();
|
|
e.observations.retain(|o| !to_rem.contains(o));
|
|
} else {
|
|
missing.push(d.entity_name);
|
|
}
|
|
}
|
|
});
|
|
if !missing.is_empty() {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
|
|
missing.join(", ")
|
|
)));
|
|
}
|
|
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", "Execute delete_relations")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: DeleteRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
let mut missing_count = 0;
|
|
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());
|
|
});
|
|
if missing_count > 0 {
|
|
return Err(crate::error::AppError::Internal(format!(
|
|
"Error: {} relation(s) not found in graph. Please verify exact relation properties using read_graph.",
|
|
missing_count
|
|
)));
|
|
}
|
|
Ok("Relations deleted".to_string())
|
|
}
|
|
}
|
|
|
|
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", "Execute read_graph")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
let max_tokens = req.max_tokens;
|
|
let data = state.read_graph(|full| -> crate::error::Result<String> {
|
|
let mut result_json = if let Some(ns) = req.namespace {
|
|
let mut filtered = BorrowedGraph::default();
|
|
for (k, v) in &full.entities {
|
|
if v.namespace == ns {
|
|
filtered.entities.insert(k, v);
|
|
}
|
|
}
|
|
for r in &full.relations {
|
|
if r.namespace == ns {
|
|
filtered.relations.push(r);
|
|
}
|
|
}
|
|
serde_json::to_string(&filtered)?
|
|
} else {
|
|
serde_json::to_string(full)?
|
|
};
|
|
|
|
if let Some(max_tok) = max_tokens {
|
|
let max_chars = max_tok * 4;
|
|
if result_json.len() > max_chars {
|
|
result_json.truncate(max_chars);
|
|
result_json.push_str("... [TRUNCATED_TO_MAX_TOKENS]");
|
|
}
|
|
}
|
|
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", "Execute search_nodes")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: SearchNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
let matches = if let Ok(idx) = state.search_index.read() {
|
|
idx.search(&req.query, req.namespace.as_deref())
|
|
.unwrap_or_default()
|
|
} else {
|
|
vec![]
|
|
};
|
|
|
|
let data = state.read_graph(|full| -> crate::error::Result<String> {
|
|
let mut result = BorrowedGraph::default();
|
|
for (id, doc_type, _, _, _) in &matches {
|
|
if doc_type == "entity"
|
|
&& let Some(e) = full.entities.get(id)
|
|
{
|
|
result.entities.insert(id, e);
|
|
}
|
|
}
|
|
Ok::<String, crate::error::AppError>(serde_json::to_string(&result)?)
|
|
})?;
|
|
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", "Execute open_nodes")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
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);
|
|
}
|
|
}
|
|
Ok::<String, crate::error::AppError>(serde_json::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", "Execute visualize_graph")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: VisualizeGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
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", "Execute condense_entity")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: CondenseEntityTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
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).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", "Execute merge_entities")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: MergeEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
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();
|
|
}
|
|
|
|
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", "Execute find_orphans")
|
|
}
|
|
|
|
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_json::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", "Execute get_subgraph")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let req: GetSubgraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
let root = req
|
|
.root_entity
|
|
.or(req.root_node)
|
|
.ok_or_else(|| crate::error::AppError::Internal("root_entity or root_node is required".to_string()))?;
|
|
let depth = req.depth.unwrap_or(2);
|
|
let format = req.format.unwrap_or(SubgraphFormat::Json);
|
|
|
|
if format == SubgraphFormat::MarkdownTree {
|
|
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((root.clone(), 0));
|
|
visited.insert(root.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", root);
|
|
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));
|
|
}
|
|
return Ok(markdown);
|
|
}
|
|
|
|
let data = state.read_graph(|full| {
|
|
let mut visited_nodes: std::collections::HashSet<String> = std::collections::HashSet::new();
|
|
let mut matched_entities = std::collections::HashMap::new();
|
|
let mut matched_relations = Vec::new();
|
|
|
|
let mut queue = std::collections::VecDeque::new();
|
|
queue.push_back((root.clone(), 0));
|
|
visited_nodes.insert(root.clone());
|
|
|
|
while let Some((curr, current_depth)) = queue.pop_front() {
|
|
if let Some(entity) = full.entities.get(&curr) {
|
|
matched_entities.insert(curr.clone(), entity.clone());
|
|
}
|
|
|
|
if current_depth < depth {
|
|
for rel in &full.relations {
|
|
if rel.from == curr {
|
|
matched_relations.push(rel.clone());
|
|
if !visited_nodes.contains(&rel.to) {
|
|
visited_nodes.insert(rel.to.clone());
|
|
queue.push_back((rel.to.clone(), current_depth + 1));
|
|
}
|
|
} else if rel.to == curr {
|
|
matched_relations.push(rel.clone());
|
|
if !visited_nodes.contains(&rel.from) {
|
|
visited_nodes.insert(rel.from.clone());
|
|
queue.push_back((rel.from.clone(), current_depth + 1));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let result = serde_json::json!({
|
|
"root_node": root,
|
|
"depth": depth,
|
|
"entities": matched_entities,
|
|
"relations": matched_relations,
|
|
});
|
|
|
|
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&result)?)
|
|
})?;
|
|
|
|
Ok(data)
|
|
}
|
|
}
|
|
|
|
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| e.to_string())?;
|
|
let auto_prune = req.auto_prune_orphans.unwrap_or(false);
|
|
|
|
let mut orphans = Vec::new();
|
|
let mut duplicates = 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());
|
|
}
|
|
}
|
|
|
|
if auto_prune {
|
|
for orphan in &orphans {
|
|
g.entities.remove(orphan);
|
|
}
|
|
}
|
|
|
|
// 2. Compute similarity pairs for duplicate detection
|
|
let names: Vec<_> = g.entities.keys().cloned().collect();
|
|
for i in 0..names.len() {
|
|
for j in (i + 1)..names.len() {
|
|
let n1 = &names[i];
|
|
let n2 = &names[j];
|
|
|
|
let l1 = n1.to_lowercase();
|
|
let l2 = n2.to_lowercase();
|
|
if l1 == l2 || ((l1.contains(&l2) || l2.contains(&l1)) && 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 report = serde_json::json!({
|
|
"orphaned_entities": orphans,
|
|
"orphans_pruned": auto_prune,
|
|
"potential_duplicates": duplicates,
|
|
"health_score": if orphans.is_empty() && duplicates.is_empty() { "100%" } else { "Needs Maintenance" }
|
|
});
|
|
|
|
Ok(serde_json::to_string_pretty(&report)?)
|
|
}
|
|
}
|
|
|
|
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| e.to_string())?;
|
|
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, _entity) in g.entities.iter_mut() {
|
|
if let Some(target) = &req.target {
|
|
if 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_json::to_string_pretty(&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| e.to_string())?;
|
|
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 {
|
|
markdown.truncate(max_chars);
|
|
markdown.push_str("\n... [Truncated to fit token budget]");
|
|
}
|
|
|
|
Ok(markdown)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::handlers::meta::{BroadcastAgentSignalHandler, QueryAgentSignalsHandler};
|
|
use serde_json::json;
|
|
use tempfile::tempdir;
|
|
|
|
|
|
#[tokio::test]
|
|
async fn test_create_and_read_entities() {
|
|
let dir = tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
|
|
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_eq!(res, "Entities created");
|
|
|
|
// 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 dir = tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
|
|
// 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,
|
|
},
|
|
);
|
|
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,
|
|
},
|
|
);
|
|
});
|
|
|
|
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_eq!(res, "Relations created");
|
|
|
|
// Test semantic LLM schema feedback (User request)
|
|
let bad_args = json!({
|
|
"relations": [
|
|
{"source": "A", "target": "B", "relationType": "knows"}
|
|
]
|
|
});
|
|
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 dir = tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
|
|
// 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,
|
|
},
|
|
);
|
|
});
|
|
|
|
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, "Entities deleted");
|
|
|
|
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 dir = tempfile::tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
|
|
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();
|
|
assert!(!res_orphans.contains("Y"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_more_graph_handlers() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
|
|
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, "Relations deleted");
|
|
|
|
let bcast_handler = BroadcastAgentSignalHandler;
|
|
let bcast_res = bcast_handler.execute(json!({
|
|
"sender": "agent1",
|
|
"signal_type": "task_completed",
|
|
"payload": "fix_bug"
|
|
}), state.clone()).await.unwrap();
|
|
assert!(bcast_res.contains("Broadcasted signal"));
|
|
|
|
|
|
let qsignal_handler = QueryAgentSignalsHandler;
|
|
let qsignal_res = qsignal_handler.execute(json!({"sender": "agent1"}), state.clone()).await.unwrap();
|
|
assert!(qsignal_res.contains("task_completed"));
|
|
|
|
}
|
|
}
|
|
|