perf(graph): eliminate heavy allocations by substituting KnowledgeGraph clones with BorrowedGraph for json serialization
This commit is contained in:
1 parent
1d5e4bd00d
commit
8934857635
2 files changed
+38
-18
No files matched your search
@@ -7,6 +7,21 @@ use serde_json::Value;
|
|||||||
use std::collections::HashSet;
|
use std::collections::HashSet;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[derive(serde::Serialize)]
|
||||||
|
struct BorrowedGraph<'a> {
|
||||||
|
entities: std::collections::HashMap<&'a String, &'a crate::models::Entity>,
|
||||||
|
relations: Vec<&'a crate::models::Relation>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a> Default for BorrowedGraph<'a> {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
entities: std::collections::HashMap::new(),
|
||||||
|
relations: Vec::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub struct QueryGraphPathHandler;
|
pub struct QueryGraphPathHandler;
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
@@ -267,15 +282,15 @@ impl McpTool for ReadGraphHandler {
|
|||||||
let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
let data = state.read_graph(|full| {
|
let data = state.read_graph(|full| {
|
||||||
if let Some(ns) = req.namespace {
|
if let Some(ns) = req.namespace {
|
||||||
let mut filtered = KnowledgeGraph::default();
|
let mut filtered = BorrowedGraph::default();
|
||||||
for (k, v) in &full.entities {
|
for (k, v) in &full.entities {
|
||||||
if v.namespace == ns {
|
if v.namespace == ns {
|
||||||
filtered.entities.insert(k.clone(), v.clone());
|
filtered.entities.insert(k, v);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
for r in &full.relations {
|
for r in &full.relations {
|
||||||
if r.namespace == ns {
|
if r.namespace == ns {
|
||||||
filtered.relations.push(r.clone());
|
filtered.relations.push(r);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
serde_json::to_string(&filtered).unwrap_or_default()
|
serde_json::to_string(&filtered).unwrap_or_default()
|
||||||
@@ -308,18 +323,18 @@ impl McpTool for SearchNodesHandler {
|
|||||||
vec![]
|
vec![]
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut result = KnowledgeGraph::default();
|
let data = state.read_graph(|full| {
|
||||||
state.read_graph(|full| {
|
let mut result = BorrowedGraph::default();
|
||||||
for (id, doc_type, _, _, _) in matches {
|
for (id, doc_type, _, _, _) in &matches {
|
||||||
if doc_type == "entity"
|
if doc_type == "entity"
|
||||||
&& let Some(e) = full.entities.get(&id)
|
&& let Some(e) = full.entities.get(id)
|
||||||
{
|
{
|
||||||
result.entities.insert(id, e.clone());
|
result.entities.insert(id, e);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
serde_json::to_string(&result).unwrap_or_default()
|
||||||
});
|
});
|
||||||
let data = serde_json::to_string(&result).unwrap_or_default();
|
Ok(data)
|
||||||
Ok(data.to_string())
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -337,28 +352,27 @@ impl McpTool for OpenNodesHandler {
|
|||||||
|
|
||||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
|
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 req: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||||
let result = state.read_graph(|full| {
|
let data = state.read_graph(|full| {
|
||||||
let targets: HashSet<&str> = req.names.iter().map(|s| s.as_str()).collect();
|
let targets: HashSet<&str> = req.names.iter().map(|s| s.as_str()).collect();
|
||||||
let mut result = KnowledgeGraph::default();
|
let mut result = BorrowedGraph::default();
|
||||||
let mut connected: HashSet<&str> = HashSet::new();
|
let mut connected: HashSet<&str> = HashSet::new();
|
||||||
for r in &full.relations {
|
for r in &full.relations {
|
||||||
if targets.contains(r.from.as_str()) {
|
if targets.contains(r.from.as_str()) {
|
||||||
connected.insert(r.to.as_str());
|
connected.insert(r.to.as_str());
|
||||||
result.relations.push(r.clone());
|
result.relations.push(r);
|
||||||
} else if targets.contains(r.to.as_str()) {
|
} else if targets.contains(r.to.as_str()) {
|
||||||
connected.insert(r.from.as_str());
|
connected.insert(r.from.as_str());
|
||||||
result.relations.push(r.clone());
|
result.relations.push(r);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
for (name, e) in &full.entities {
|
for (name, e) in &full.entities {
|
||||||
if targets.contains(name.as_str()) || connected.contains(name.as_str()) {
|
if targets.contains(name.as_str()) || connected.contains(name.as_str()) {
|
||||||
result.entities.insert(name.clone(), e.clone());
|
result.entities.insert(name, e);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
result
|
serde_json::to_string(&result).unwrap_or_default()
|
||||||
});
|
});
|
||||||
let data = serde_json::to_string(&result).unwrap_or_default();
|
Ok(data)
|
||||||
Ok(data.to_string())
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,6 @@
|
|||||||
|
local buf = vim.fn.bufnr('server/src/handlers_v2/graph.rs')
|
||||||
|
if buf == -1 then
|
||||||
|
vim.cmd('e server/src/handlers_v2/graph.rs')
|
||||||
|
buf = vim.api.nvim_get_current_buf()
|
||||||
|
end
|
||||||
|
return vim.api.nvim_buf_get_name(buf)
|
||||||
Reference in new issue
Block a user