Riz Ashraf
37003be620
- Optimized memory allocation in router.rs by offloading JSON serialization to spawn_blocking and using references. - Prevented full graph duplication on startup in state.rs index rebuild. - Eliminated massive String allocations in QueryGraphPathHandler BFS loops. - Avoided temporary Strings in VisualizeGraphHandler via inline writing. - Fixed O(N) full-graph deduplication in MergeEntitiesHandler to scale efficiently.
767 lines
26 KiB
Rust
767 lines
26 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>) -> Result<String, 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>) -> Result<String, 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 entity in req.entities {
|
|
if !entity.name.is_empty() {
|
|
inserted.push(entity.clone());
|
|
g.entities.insert(entity.name.clone(), entity);
|
|
}
|
|
}
|
|
});
|
|
let idx = state
|
|
.search_index
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.clone();
|
|
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>) -> Result<String, String> {
|
|
let req: CreateRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
state.modify_graph(|g| {
|
|
for relation in req.relations {
|
|
if !relation.from.is_empty() && !relation.to.is_empty() {
|
|
g.relations.push(relation);
|
|
}
|
|
}
|
|
});
|
|
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>) -> Result<String, String> {
|
|
let req: AddObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
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", "Execute delete_entities")
|
|
}
|
|
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, 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();
|
|
state.modify_graph(|master| {
|
|
for name in &to_delete {
|
|
master.entities.remove(name);
|
|
}
|
|
master
|
|
.relations
|
|
.retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
|
|
});
|
|
|
|
let idx = state
|
|
.search_index
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.clone();
|
|
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>) -> Result<String, String> {
|
|
let req: DeleteObservationsTool =
|
|
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
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));
|
|
}
|
|
}
|
|
});
|
|
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>) -> Result<String, String> {
|
|
let req: DeleteRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
state.modify_graph(|master| {
|
|
let to_rem: HashSet<_> = req.relations.into_iter().collect();
|
|
master.relations.retain(|r| !to_rem.contains(r));
|
|
});
|
|
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>) -> Result<String, String> {
|
|
let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
let data = state.read_graph(|full| {
|
|
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).unwrap_or_default()
|
|
} else {
|
|
serde_json::to_string(full).unwrap_or_default()
|
|
}
|
|
});
|
|
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>) -> Result<String, 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| {
|
|
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);
|
|
}
|
|
}
|
|
serde_json::to_string(&result).unwrap_or_default()
|
|
});
|
|
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>) -> Result<String, String> {
|
|
let req: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
let data = state.read_graph(|full| {
|
|
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);
|
|
}
|
|
}
|
|
serde_json::to_string(&result).unwrap_or_default()
|
|
});
|
|
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>) -> Result<String, 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>) -> Result<String, String> {
|
|
let req: CondenseEntityTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
state.modify_graph(|master| {
|
|
if let Some(e) = master.entities.get_mut(&req.entity_name) {
|
|
e.observations = req.summarized_observations;
|
|
}
|
|
});
|
|
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>) -> Result<String, String> {
|
|
let req: MergeEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
|
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);
|
|
MemoryState::deduplicate(&mut tgt.observations);
|
|
} else {
|
|
let mut new_tgt = src.clone();
|
|
new_tgt.name = req.target_entity.clone();
|
|
master.entities.insert(req.target_entity.clone(), new_tgt);
|
|
}
|
|
}
|
|
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
|
|
}
|
|
});
|
|
});
|
|
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>) -> Result<String, 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 = serde_json::to_string(&orphans).unwrap_or_default();
|
|
Ok(data.to_string())
|
|
}
|
|
}
|
|
|
|
use crate::handlers::utils::*;
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
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", "entityType": "Person", "observations": ["Likes Bob"]}
|
|
]
|
|
});
|
|
|
|
let res = create_handler.execute(args, state.clone()).await.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", "relationType": "knows"}
|
|
]
|
|
});
|
|
let res = handler.execute(args, state.clone()).await.unwrap();
|
|
assert_eq!(res, "Relations created");
|
|
}
|
|
|
|
#[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": [
|
|
{"entityName": "A", "contents": ["Obs 1", "Obs 2"]}
|
|
]
|
|
});
|
|
let res1 = add_obs.execute(args_obs, state.clone()).await.unwrap();
|
|
assert_eq!(res1, "Observations added");
|
|
|
|
let read_graph = ReadGraphHandler;
|
|
let res2 = read_graph
|
|
.execute(json!({"namespace": "global"}), state.clone())
|
|
.await
|
|
.unwrap();
|
|
assert!(res2.contains("Obs 1"));
|
|
assert!(res2.contains("Obs 2"));
|
|
|
|
let del_entity = DeleteEntitiesHandler;
|
|
let res4 = del_entity
|
|
.execute(json!({"entityNames": ["A"]}), state.clone())
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(res4, "Entities deleted");
|
|
|
|
let res5 = read_graph
|
|
.execute(json!({"namespace": "global"}), state.clone())
|
|
.await
|
|
.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", "entityType": "File", "observations": ["Obs X"], "namespace": "global"},
|
|
{"name": "Y", "entityType": "File", "observations": ["Obs Y"], "namespace": "global"}
|
|
]
|
|
});
|
|
create_handler
|
|
.execute(args_ent, state.clone())
|
|
.await
|
|
.unwrap();
|
|
|
|
let rel_handler = CreateRelationsHandler;
|
|
let args_rel = json!({
|
|
"relations": [
|
|
{"from": "X", "to": "Y", "relationType": "depends_on", "namespace": "global"}
|
|
]
|
|
});
|
|
rel_handler.execute(args_rel, state.clone()).await.unwrap();
|
|
|
|
let read_handler = ReadGraphHandler;
|
|
let res_read = read_handler
|
|
.execute(json!({"namespace": "global"}), state.clone())
|
|
.await
|
|
.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
|
|
.unwrap();
|
|
assert!(res_open.contains("Y"));
|
|
|
|
let viz_handler = VisualizeGraphHandler;
|
|
let res_viz = viz_handler
|
|
.execute(json!({"query": "X"}), state.clone())
|
|
.await
|
|
.unwrap();
|
|
assert!(!res_viz.is_empty());
|
|
|
|
let condense = CondenseEntityHandler;
|
|
let res_cond = condense
|
|
.execute(
|
|
json!({"entityName": "X", "summarized_observations": ["X condensed"]}),
|
|
state.clone(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(res_cond, "Entity condensed");
|
|
|
|
let merge = MergeEntitiesHandler;
|
|
let res_merge = merge
|
|
.execute(
|
|
json!({"sourceEntity": "X", "targetEntity": "Y"}),
|
|
state.clone(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(res_merge, "Entities merged");
|
|
|
|
let orphans = FindOrphansHandler;
|
|
let res_orphans = orphans.execute(json!({}), state.clone()).await.unwrap();
|
|
assert!(!res_orphans.contains("Y"));
|
|
}
|
|
}
|