Vakyartha documentation

1//! Semantic Entity Knowledge Graph (tri-partite memory architecture).
2//!
3//! Stores typed domain entities with attributes and cross-entity relations
4//! per workspace in `<sessions_home>/entities/<hash>/ENTITIES.jsonl`
5//! (and global entities in `<sessions_home>/entities/global/ENTITIES.jsonl`).
6//!
7//! Entities participate in recall via `session_search` and can be inspected
8//! or updated across turns.
9
10use std::collections::BTreeMap;
11use std::fs::{File, OpenOptions};
12use std::io::{BufRead, BufReader, Write};
13use std::path::{Path, PathBuf};
14
15use chrono::{DateTime, Utc};
16use serde::{Deserialize, Serialize};
17
18use crate::memory::hash_cwd;
19
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct EntityRelation {
22 pub relation: String,
23 pub target_entity_id: String,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
27pub struct EntityRecord {
28 pub id: String,
29 pub name: String,
30 pub entity_type: String,
31 pub summary: String,
32 #[serde(default)]
33 pub attributes: BTreeMap<String, String>,
34 #[serde(default)]
35 pub relations: Vec<EntityRelation>,
36 pub updated_at: DateTime<Utc>,
37}
38
39pub fn entities_file(home: &Path, cwd: Option<&Path>) -> PathBuf {
40 let sub = match cwd {
41 Some(dir) => hash_cwd(dir),
42 None => "global".to_string(),
43 };
44 home.join("entities").join(sub).join("ENTITIES.jsonl")
45}
46
47/// List all entities in the target workspace (or global if cwd is None).
48pub fn list_entities(home: &Path, cwd: Option<&Path>) -> Vec<EntityRecord> {
49 let path = entities_file(home, cwd);
50 let Ok(file) = File::open(&path) else {
51 return Vec::new();
52 };
53 let reader = BufReader::new(file);
54 let mut records = Vec::new();
55 for line in reader.lines().map_while(Result::ok) {
56 let trimmed = line.trim();
57 if trimmed.is_empty() {
58 continue;
59 }
60 if let Ok(record) = serde_json::from_str::<EntityRecord>(trimmed) {
61 records.push(record);
62 }
63 }
64 records
65}
66
67/// Retrieve a specific entity by ID.
68pub fn get_entity(home: &Path, cwd: Option<&Path>, id: &str) -> Option<EntityRecord> {
69 list_entities(home, cwd).into_iter().find(|e| e.id == id)
70}
71
72/// Search entities by keyword across name, entity_type, summary, attributes, and relations.
73pub fn search_entities(home: &Path, cwd: Option<&Path>, query: &str) -> Vec<EntityRecord> {
74 let q = query.trim().to_ascii_lowercase();
75 if q.is_empty() {
76 return list_entities(home, cwd);
77 }
78 list_entities(home, cwd)
79 .into_iter()
80 .filter(|e| {
81 e.name.to_ascii_lowercase().contains(&q)
82 || e.entity_type.to_ascii_lowercase().contains(&q)
83 || e.summary.to_ascii_lowercase().contains(&q)
84 || e.attributes.iter().any(|(k, v)| {
85 k.to_ascii_lowercase().contains(&q) || v.to_ascii_lowercase().contains(&q)
86 })
87 || e.relations.iter().any(|r| {
88 r.relation.to_ascii_lowercase().contains(&q)
89 || r.target_entity_id.to_ascii_lowercase().contains(&q)
90 })
91 })
92 .collect()
93}
94
95/// Upsert an entity record: updates in-place if matching ID exists, or appends.
96pub fn upsert_entity(
97 home: &Path,
98 cwd: Option<&Path>,
99 mut record: EntityRecord,
100) -> Result<EntityRecord, std::io::Error> {
101 let path = entities_file(home, cwd);
102 if let Some(parent) = path.parent() {
103 std::fs::create_dir_all(parent)?;
104 }
105 record.updated_at = Utc::now();
106 let mut all = list_entities(home, cwd);
107 if let Some(pos) = all.iter().position(|e| e.id == record.id) {
108 all[pos] = record.clone();
109 } else {
110 all.push(record.clone());
111 }
112
113 // Atomic overwrite via temp file
114 let tmp_path = path.with_extension("tmp");
115 {
116 let mut file = OpenOptions::new()
117 .create(true)
118 .write(true)
119 .truncate(true)
120 .open(&tmp_path)?;
121 for entry in &all {
122 let json = serde_json::to_string(entry)
123 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
124 writeln!(file, "{json}")?;
125 }
126 file.flush()?;
127 }
128 std::fs::rename(&tmp_path, &path)?;
129 Ok(record)
130}
131
132/// Delete an entity by ID. Returns true if removed, false if not found.
133pub fn delete_entity(home: &Path, cwd: Option<&Path>, id: &str) -> Result<bool, std::io::Error> {
134 let path = entities_file(home, cwd);
135 if !path.is_file() {
136 return Ok(false);
137 }
138 let all = list_entities(home, cwd);
139 let orig_len = all.len();
140 let filtered: Vec<_> = all.into_iter().filter(|e| e.id != id).collect();
141 if filtered.len() == orig_len {
142 return Ok(false);
143 }
144
145 let tmp_path = path.with_extension("tmp");
146 {
147 let mut file = OpenOptions::new()
148 .create(true)
149 .write(true)
150 .truncate(true)
151 .open(&tmp_path)?;
152 for entry in &filtered {
153 let json = serde_json::to_string(entry)
154 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
155 writeln!(file, "{json}")?;
156 }
157 file.flush()?;
158 }
159 std::fs::rename(&tmp_path, &path)?;
160 Ok(true)
161}
162
163pub struct EntityRecordTool {
164 pub sessions_home: PathBuf,
165 pub cwd: PathBuf,
166}
167
168#[async_trait::async_trait]
169impl vak_tools::Tool for EntityRecordTool {
170 fn name(&self) -> &str {
171 "entity_record"
172 }
173
174 fn serves(&self) -> &'static [&'static str] {
175 &["memory"]
176 }
177
178 fn description(&self) -> &str {
179 "Record or update a domain entity in the semantic knowledge graph. \
180 Entities represent durable systems, concepts, people, datasets, or components \
181 with typed attributes and directed relations (e.g. depends_on, hosted_on, owns)."
182 }
183
184 fn schema(&self) -> serde_json::Value {
185 serde_json::json!({
186 "type": "object",
187 "properties": {
188 "id": {
189 "type": "string",
190 "description": "Unique slug identifier (e.g. 'db-primary' or 'auth-service'). If omitted, slug is derived from name."
191 },
192 "name": {
193 "type": "string",
194 "description": "Human-readable name of the entity"
195 },
196 "entity_type": {
197 "type": "string",
198 "description": "Category or kind of entity (e.g. 'system', 'service', 'dataset', 'person', 'concept', 'module')"
199 },
200 "summary": {
201 "type": "string",
202 "description": "Concise summary of the entity's purpose, role, or definition"
203 },
204 "attributes": {
205 "type": "object",
206 "description": "Key-value map of attributes or metadata"
207 },
208 "relations": {
209 "type": "array",
210 "items": {
211 "type": "object",
212 "properties": {
213 "relation": { "type": "string", "description": "Relationship verb (e.g. 'depends_on', 'hosted_on', 'owns')" },
214 "target": { "type": "string", "description": "Target entity ID" }
215 },
216 "required": ["relation", "target"]
217 },
218 "description": "Outgoing directed relationships to other entities"
219 }
220 },
221 "required": ["name", "entity_type", "summary"]
222 })
223 }
224
225 async fn execute(
226 &self,
227 args: &serde_json::Value,
228 _ctx: &vak_tools::ToolContext,
229 ) -> vak_tools::ToolOutput {
230 let Some(name) = args.get("name").and_then(|v| v.as_str()).map(str::trim) else {
231 return vak_tools::ToolOutput::error("missing required argument 'name'");
232 };
233 let Some(entity_type) = args
234 .get("entity_type")
235 .and_then(|v| v.as_str())
236 .map(str::trim)
237 else {
238 return vak_tools::ToolOutput::error("missing required argument 'entity_type'");
239 };
240 let Some(summary) = args.get("summary").and_then(|v| v.as_str()).map(str::trim) else {
241 return vak_tools::ToolOutput::error("missing required argument 'summary'");
242 };
243 if name.is_empty() || entity_type.is_empty() || summary.is_empty() {
244 return vak_tools::ToolOutput::error(
245 "'name', 'entity_type', and 'summary' must not be empty",
246 );
247 }
248
249 let id = match args.get("id").and_then(|v| v.as_str()).map(str::trim) {
250 Some(custom) if !custom.is_empty() => custom.to_string(),
251 _ => {
252 let slug: String = name
253 .to_ascii_lowercase()
254 .chars()
255 .map(|c| if c.is_alphanumeric() { c } else { '-' })
256 .collect();
257 let deduped = slug
258 .split('-')
259 .filter(|s| !s.is_empty())
260 .collect::<Vec<_>>()
261 .join("-");
262 if deduped.is_empty() {
263 format!("ent-{}", hash_cwd(Path::new(name)))
264 } else {
265 deduped
266 }
267 }
268 };
269
270 let mut attributes = BTreeMap::new();
271 if let Some(obj) = args.get("attributes").and_then(|v| v.as_object()) {
272 for (k, v) in obj {
273 let val_str = match v {
274 serde_json::Value::String(s) => s.clone(),
275 other => other.to_string(),
276 };
277 attributes.insert(k.clone(), val_str);
278 }
279 }
280
281 let mut relations = Vec::new();
282 if let Some(arr) = args.get("relations").and_then(|v| v.as_array()) {
283 for item in arr {
284 if let (Some(rel), Some(tgt)) = (
285 item.get("relation").and_then(|v| v.as_str()),
286 item.get("target").and_then(|v| v.as_str()),
287 ) {
288 relations.push(EntityRelation {
289 relation: rel.trim().to_string(),
290 target_entity_id: tgt.trim().to_string(),
291 });
292 }
293 }
294 }
295
296 let record = EntityRecord {
297 id: id.clone(),
298 name: name.to_string(),
299 entity_type: entity_type.to_string(),
300 summary: summary.to_string(),
301 attributes,
302 relations,
303 updated_at: Utc::now(),
304 };
305
306 match upsert_entity(&self.sessions_home, Some(&self.cwd), record) {
307 Ok(saved) => vak_tools::ToolOutput::ok(format!(
308 "recorded entity '{id}' ({}) with {} attributes and {} relations",
309 saved.entity_type,
310 saved.attributes.len(),
311 saved.relations.len()
312 )),
313 Err(e) => vak_tools::ToolOutput::error(format!("could not record entity: {e}")),
314 }
315 }
316
317 fn claims(&self, _args: &serde_json::Value) -> vak_tools::ResourceClaims {
318 vak_tools::ResourceClaims {
319 exclusive: true,
320 read_only: false,
321 paths: vec![],
322 }
323 }
324}
325
326pub struct EntityQueryTool {
327 pub sessions_home: PathBuf,
328 pub cwd: PathBuf,
329}
330
331#[async_trait::async_trait]
332impl vak_tools::Tool for EntityQueryTool {
333 fn name(&self) -> &str {
334 "entity_query"
335 }
336
337 fn serves(&self) -> &'static [&'static str] {
338 &["memory"]
339 }
340
341 fn description(&self) -> &str {
342 "Query the semantic entity knowledge graph by keyword or entity type. \
343 Returns matching entities with their attributes and related connections."
344 }
345
346 fn schema(&self) -> serde_json::Value {
347 serde_json::json!({
348 "type": "object",
349 "properties": {
350 "query": {
351 "type": "string",
352 "description": "Keywords to match against entity names, summaries, attributes, or relations"
353 },
354 "entity_type": {
355 "type": "string",
356 "description": "Optional filter by entity category (e.g. 'system', 'dataset')"
357 },
358 "limit": {
359 "type": "integer",
360 "description": "Max results to return (default 10, max 50)"
361 }
362 }
363 })
364 }
365
366 async fn execute(
367 &self,
368 args: &serde_json::Value,
369 _ctx: &vak_tools::ToolContext,
370 ) -> vak_tools::ToolOutput {
371 let query = args.get("query").and_then(|v| v.as_str()).unwrap_or("");
372 let entity_type_filter = args
373 .get("entity_type")
374 .and_then(|v| v.as_str())
375 .map(|s| s.to_ascii_lowercase());
376 let limit = args
377 .get("limit")
378 .and_then(|v| v.as_u64())
379 .map(|l| l as usize)
380 .unwrap_or(10)
381 .min(50);
382
383 let mut results = search_entities(&self.sessions_home, Some(&self.cwd), query);
384 if let Some(ref filter) = entity_type_filter {
385 results.retain(|e| e.entity_type.to_ascii_lowercase() == *filter);
386 }
387 results.truncate(limit);
388
389 if results.is_empty() {
390 return vak_tools::ToolOutput::ok("no matching entities found in knowledge graph");
391 }
392
393 let mut out = format!("Found {} entities:\n\n", results.len());
394 for e in results {
395 out.push_str(&format!(
396 "### [{}] {} (`{}`)\n",
397 e.entity_type, e.name, e.id
398 ));
399 out.push_str(&format!("{}\n", e.summary));
400 if !e.attributes.is_empty() {
401 out.push_str("Attributes:\n");
402 for (k, v) in &e.attributes {
403 out.push_str(&format!("- **{k}**: {v}\n"));
404 }
405 }
406 if !e.relations.is_empty() {
407 out.push_str("Relations:\n");
408 for r in &e.relations {
409 out.push_str(&format!("- {} -> `{}`\n", r.relation, r.target_entity_id));
410 }
411 }
412 out.push('\n');
413 }
414
415 vak_tools::ToolOutput::ok(out.trim_end().to_string())
416 }
417
418 fn claims(&self, _args: &serde_json::Value) -> vak_tools::ResourceClaims {
419 vak_tools::ResourceClaims {
420 exclusive: false,
421 read_only: true,
422 paths: vec![],
423 }
424 }
425}
426
427#[cfg(test)]
428mod tests {
429 #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
430 use super::*;
431
432 #[test]
433 fn entity_crud_and_search_roundtrip() {
434 let temp = tempfile::tempdir().unwrap();
435 let home = temp.path().join("home");
436 let cwd = temp.path().join("cwd");
437 std::fs::create_dir_all(&home).unwrap();
438 std::fs::create_dir_all(&cwd).unwrap();
439
440 let mut attrs = BTreeMap::new();
441 attrs.insert("role".into(), "primary_db".into());
442 attrs.insert("engine".into(), "postgresql".into());
443
444 let entity = EntityRecord {
445 id: "ent-postgres-1".into(),
446 name: "Main PostgreSQL Cluster".into(),
447 entity_type: "database".into(),
448 summary: "Primary transaction database holding customer accounts".into(),
449 attributes: attrs,
450 relations: vec![EntityRelation {
451 relation: "hosted_on".into(),
452 target_entity_id: "srv-aws-us-east-1".into(),
453 }],
454 updated_at: Utc::now(),
455 };
456
457 // 1. Upsert
458 let saved = upsert_entity(&home, Some(&cwd), entity.clone()).unwrap();
459 assert_eq!(saved.id, "ent-postgres-1");
460
461 // 2. Get
462 let fetched = get_entity(&home, Some(&cwd), "ent-postgres-1").unwrap();
463 assert_eq!(fetched.name, "Main PostgreSQL Cluster");
464 assert_eq!(fetched.attributes.get("engine").unwrap(), "postgresql");
465
466 // 3. Search
467 let results = search_entities(&home, Some(&cwd), "transaction");
468 assert_eq!(results.len(), 1);
469 assert_eq!(results[0].id, "ent-postgres-1");
470
471 let rel_results = search_entities(&home, Some(&cwd), "aws-us-east-1");
472 assert_eq!(rel_results.len(), 1);
473
474 let none_results = search_entities(&home, Some(&cwd), "redis");
475 assert!(none_results.is_empty());
476
477 // 4. Update
478 let mut updated = fetched;
479 updated.summary = "Updated summary".into();
480 upsert_entity(&home, Some(&cwd), updated).unwrap();
481 let re_fetched = get_entity(&home, Some(&cwd), "ent-postgres-1").unwrap();
482 assert_eq!(re_fetched.summary, "Updated summary");
483
484 // 5. Delete
485 assert!(delete_entity(&home, Some(&cwd), "ent-postgres-1").unwrap());
486 assert!(get_entity(&home, Some(&cwd), "ent-postgres-1").is_none());
487 assert!(!delete_entity(&home, Some(&cwd), "ent-postgres-1").unwrap());
488 }
489
490 #[tokio::test]
491 async fn entity_record_and_query_tools_execute() {
492 use vak_tools::Tool;
493
494 let temp = tempfile::tempdir().unwrap();
495 let home = temp.path().join("home");
496 let cwd = temp.path().join("cwd");
497 std::fs::create_dir_all(&home).unwrap();
498 std::fs::create_dir_all(&cwd).unwrap();
499
500 let record_tool = EntityRecordTool {
501 sessions_home: home.clone(),
502 cwd: cwd.clone(),
503 };
504 let query_tool = EntityQueryTool {
505 sessions_home: home,
506 cwd: cwd.clone(),
507 };
508 let ctx = vak_tools::ToolContext {
509 cwd,
510 cancel: tokio_util::sync::CancellationToken::new(),
511 sandbox: None,
512 sandbox_sink: None,
513 agent_id: None,
514 new_documents: Vec::new(),
515 };
516
517 // Record entity via tool
518 let record_args = serde_json::json!({
519 "name": "Auth Gateway",
520 "entity_type": "service",
521 "summary": "Handles JWT authentication and token exchange",
522 "attributes": { "port": "8080", "protocol": "https" },
523 "relations": [ { "relation": "depends_on", "target": "redis-session-store" } ]
524 });
525 let rec_out = record_tool.execute(&record_args, &ctx).await;
526 assert!(!rec_out.is_error, "{}", rec_out.content);
527 assert!(
528 rec_out.content.contains("recorded entity 'auth-gateway'"),
529 "{}",
530 rec_out.content
531 );
532
533 // Query entity via tool
534 let query_args = serde_json::json!({ "query": "JWT" });
535 let q_out = query_tool.execute(&query_args, &ctx).await;
536 assert!(!q_out.is_error, "{}", q_out.content);
537 assert!(q_out.content.contains("Auth Gateway"), "{}", q_out.content);
538 assert!(
539 q_out
540 .content
541 .contains("depends_on -> `redis-session-store`"),
542 "{}",
543 q_out.content
544 );
545
546 // Filter query by entity_type
547 let type_args = serde_json::json!({ "entity_type": "service" });
548 let type_out = query_tool.execute(&type_args, &ctx).await;
549 assert!(!type_out.is_error, "{}", type_out.content);
550 assert!(
551 type_out.content.contains("Auth Gateway"),
552 "{}",
553 type_out.content
554 );
555
556 let none_args = serde_json::json!({ "entity_type": "database" });
557 let none_out = query_tool.execute(&none_args, &ctx).await;
558 assert!(!none_out.is_error, "{}", none_out.content);
559 assert!(
560 none_out.content.contains("no matching entities"),
561 "{}",
562 none_out.content
563 );
564 }
565}
566

1//! Semantic Entity Knowledge Graph (tri-partite memory architecture).
2//!
3//! Stores typed domain entities with attributes and cross-entity relations
4//! per workspace in `<sessions_home>/entities/<hash>/ENTITIES.jsonl`
5//! (and global entities in `<sessions_home>/entities/global/ENTITIES.jsonl`).
6//!
7//! Entities participate in recall via `session_search` and can be inspected
8//! or updated across turns.
9
10use std::collections::BTreeMap;
11use std::fs::{File, OpenOptions};
12use std::io::{BufRead, BufReader, Write};
13use std::path::{Path, PathBuf};
14
15use chrono::{DateTime, Utc};
16use serde::{Deserialize, Serialize};
17
18use crate::memory::hash_cwd;
19
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct EntityRelation {
22 pub relation: String,
23 pub target_entity_id: String,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
27pub struct EntityRecord {
28 pub id: String,
29 pub name: String,
30 pub entity_type: String,
31 pub summary: String,
32 #[serde(default)]
33 pub attributes: BTreeMap<String, String>,
34 #[serde(default)]
35 pub relations: Vec<EntityRelation>,
36 pub updated_at: DateTime<Utc>,
37}
38
39pub fn entities_file(home: &Path, cwd: Option<&Path>) -> PathBuf {
40 let sub = match cwd {
41 Some(dir) => hash_cwd(dir),
42 None => "global".to_string(),
43 };
44 home.join("entities").join(sub).join("ENTITIES.jsonl")
45}
46
47/// List all entities in the target workspace (or global if cwd is None).
48pub fn list_entities(home: &Path, cwd: Option<&Path>) -> Vec<EntityRecord> {
49 let path = entities_file(home, cwd);
50 let Ok(file) = File::open(&path) else {
51 return Vec::new();
52 };
53 let reader = BufReader::new(file);
54 let mut records = Vec::new();
55 for line in reader.lines().map_while(Result::ok) {
56 let trimmed = line.trim();
57 if trimmed.is_empty() {
58 continue;
59 }
60 if let Ok(record) = serde_json::from_str::<EntityRecord>(trimmed) {
61 records.push(record);
62 }
63 }
64 records
65}
66
67/// Retrieve a specific entity by ID.
68pub fn get_entity(home: &Path, cwd: Option<&Path>, id: &str) -> Option<EntityRecord> {
69 list_entities(home, cwd).into_iter().find(|e| e.id == id)
70}
71
72/// Search entities by keyword across name, entity_type, summary, attributes, and relations.
73pub fn search_entities(home: &Path, cwd: Option<&Path>, query: &str) -> Vec<EntityRecord> {
74 let q = query.trim().to_ascii_lowercase();
75 if q.is_empty() {
76 return list_entities(home, cwd);
77 }
78 list_entities(home, cwd)
79 .into_iter()
80 .filter(|e| {
81 e.name.to_ascii_lowercase().contains(&q)
82 || e.entity_type.to_ascii_lowercase().contains(&q)
83 || e.summary.to_ascii_lowercase().contains(&q)
84 || e.attributes.iter().any(|(k, v)| {
85 k.to_ascii_lowercase().contains(&q) || v.to_ascii_lowercase().contains(&q)
86 })
87 || e.relations.iter().any(|r| {
88 r.relation.to_ascii_lowercase().contains(&q)
89 || r.target_entity_id.to_ascii_lowercase().contains(&q)
90 })
91 })
92 .collect()
93}
94
95/// Upsert an entity record: updates in-place if matching ID exists, or appends.
96pub fn upsert_entity(
97 home: &Path,
98 cwd: Option<&Path>,
99 mut record: EntityRecord,
100) -> Result<EntityRecord, std::io::Error> {
101 let path = entities_file(home, cwd);
102 if let Some(parent) = path.parent() {
103 std::fs::create_dir_all(parent)?;
104 }
105 record.updated_at = Utc::now();
106 let mut all = list_entities(home, cwd);
107 if let Some(pos) = all.iter().position(|e| e.id == record.id) {
108 all[pos] = record.clone();
109 } else {
110 all.push(record.clone());
111 }
112
113 // Atomic overwrite via temp file
114 let tmp_path = path.with_extension("tmp");
115 {
116 let mut file = OpenOptions::new()
117 .create(true)
118 .write(true)
119 .truncate(true)
120 .open(&tmp_path)?;
121 for entry in &all {
122 let json = serde_json::to_string(entry)
123 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
124 writeln!(file, "{json}")?;
125 }
126 file.flush()?;
127 }
128 std::fs::rename(&tmp_path, &path)?;
129 Ok(record)
130}
131
132/// Delete an entity by ID. Returns true if removed, false if not found.
133pub fn delete_entity(home: &Path, cwd: Option<&Path>, id: &str) -> Result<bool, std::io::Error> {
134 let path = entities_file(home, cwd);
135 if !path.is_file() {
136 return Ok(false);
137 }
138 let all = list_entities(home, cwd);
139 let orig_len = all.len();
140 let filtered: Vec<_> = all.into_iter().filter(|e| e.id != id).collect();
141 if filtered.len() == orig_len {
142 return Ok(false);
143 }
144
145 let tmp_path = path.with_extension("tmp");
146 {
147 let mut file = OpenOptions::new()
148 .create(true)
149 .write(true)
150 .truncate(true)
151 .open(&tmp_path)?;
152 for entry in &filtered {
153 let json = serde_json::to_string(entry)
154 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
155 writeln!(file, "{json}")?;
156 }
157 file.flush()?;
158 }
159 std::fs::rename(&tmp_path, &path)?;
160 Ok(true)
161}
162
163pub struct EntityRecordTool {
164 pub sessions_home: PathBuf,
165 pub cwd: PathBuf,
166}
167
168#[async_trait::async_trait]
169impl vak_tools::Tool for EntityRecordTool {
170 fn name(&self) -> &str {
171 "entity_record"
172 }
173
174 fn serves(&self) -> &'static [&'static str] {
175 &["memory"]
176 }
177
178 fn description(&self) -> &str {
179 "Record or update a domain entity in the semantic knowledge graph. \
180 Entities represent durable systems, concepts, people, datasets, or components \
181 with typed attributes and directed relations (e.g. depends_on, hosted_on, owns)."
182 }
183
184 fn schema(&self) -> serde_json::Value {
185 serde_json::json!({
186 "type": "object",
187 "properties": {
188 "id": {
189 "type": "string",
190 "description": "Unique slug identifier (e.g. 'db-primary' or 'auth-service'). If omitted, slug is derived from name."
191 },
192 "name": {
193 "type": "string",
194 "description": "Human-readable name of the entity"
195 },
196 "entity_type": {
197 "type": "string",
198 "description": "Category or kind of entity (e.g. 'system', 'service', 'dataset', 'person', 'concept', 'module')"
199 },
200 "summary": {
201 "type": "string",
202 "description": "Concise summary of the entity's purpose, role, or definition"
203 },
204 "attributes": {
205 "type": "object",
206 "description": "Key-value map of attributes or metadata"
207 },
208 "relations": {
209 "type": "array",
210 "items": {
211 "type": "object",
212 "properties": {
213 "relation": { "type": "string", "description": "Relationship verb (e.g. 'depends_on', 'hosted_on', 'owns')" },
214 "target": { "type": "string", "description": "Target entity ID" }
215 },
216 "required": ["relation", "target"]
217 },
218 "description": "Outgoing directed relationships to other entities"
219 }
220 },
221 "required": ["name", "entity_type", "summary"]
222 })
223 }
224
225 async fn execute(
226 &self,
227 args: &serde_json::Value,
228 _ctx: &vak_tools::ToolContext,
229 ) -> vak_tools::ToolOutput {
230 let Some(name) = args.get("name").and_then(|v| v.as_str()).map(str::trim) else {
231 return vak_tools::ToolOutput::error("missing required argument 'name'");
232 };
233 let Some(entity_type) = args
234 .get("entity_type")
235 .and_then(|v| v.as_str())
236 .map(str::trim)
237 else {
238 return vak_tools::ToolOutput::error("missing required argument 'entity_type'");
239 };
240 let Some(summary) = args.get("summary").and_then(|v| v.as_str()).map(str::trim) else {
241 return vak_tools::ToolOutput::error("missing required argument 'summary'");
242 };
243 if name.is_empty() || entity_type.is_empty() || summary.is_empty() {
244 return vak_tools::ToolOutput::error(
245 "'name', 'entity_type', and 'summary' must not be empty",
246 );
247 }
248
249 let id = match args.get("id").and_then(|v| v.as_str()).map(str::trim) {
250 Some(custom) if !custom.is_empty() => custom.to_string(),
251 _ => {
252 let slug: String = name
253 .to_ascii_lowercase()
254 .chars()
255 .map(|c| if c.is_alphanumeric() { c } else { '-' })
256 .collect();
257 let deduped = slug
258 .split('-')
259 .filter(|s| !s.is_empty())
260 .collect::<Vec<_>>()
261 .join("-");
262 if deduped.is_empty() {
263 format!("ent-{}", hash_cwd(Path::new(name)))
264 } else {
265 deduped
266 }
267 }
268 };
269
270 let mut attributes = BTreeMap::new();
271 if let Some(obj) = args.get("attributes").and_then(|v| v.as_object()) {
272 for (k, v) in obj {
273 let val_str = match v {
274 serde_json::Value::String(s) => s.clone(),
275 other => other.to_string(),
276 };
277 attributes.insert(k.clone(), val_str);
278 }
279 }
280
281 let mut relations = Vec::new();
282 if let Some(arr) = args.get("relations").and_then(|v| v.as_array()) {
283 for item in arr {
284 if let (Some(rel), Some(tgt)) = (
285 item.get("relation").and_then(|v| v.as_str()),
286 item.get("target").and_then(|v| v.as_str()),
287 ) {
288 relations.push(EntityRelation {
289 relation: rel.trim().to_string(),
290 target_entity_id: tgt.trim().to_string(),
291 });
292 }
293 }
294 }
295
296 let record = EntityRecord {
297 id: id.clone(),
298 name: name.to_string(),
299 entity_type: entity_type.to_string(),
300 summary: summary.to_string(),
301 attributes,
302 relations,
303 updated_at: Utc::now(),
304 };
305
306 match upsert_entity(&self.sessions_home, Some(&self.cwd), record) {
307 Ok(saved) => vak_tools::ToolOutput::ok(format!(
308 "recorded entity '{id}' ({}) with {} attributes and {} relations",
309 saved.entity_type,
310 saved.attributes.len(),
311 saved.relations.len()
312 )),
313 Err(e) => vak_tools::ToolOutput::error(format!("could not record entity: {e}")),
314 }
315 }
316
317 fn claims(&self, _args: &serde_json::Value) -> vak_tools::ResourceClaims {
318 vak_tools::ResourceClaims {
319 exclusive: true,
320 read_only: false,
321 paths: vec![],
322 }
323 }
324}
325
326pub struct EntityQueryTool {
327 pub sessions_home: PathBuf,
328 pub cwd: PathBuf,
329}
330
331#[async_trait::async_trait]
332impl vak_tools::Tool for EntityQueryTool {
333 fn name(&self) -> &str {
334 "entity_query"
335 }
336
337 fn serves(&self) -> &'static [&'static str] {
338 &["memory"]
339 }
340
341 fn description(&self) -> &str {
342 "Query the semantic entity knowledge graph by keyword or entity type. \
343 Returns matching entities with their attributes and related connections."
344 }
345
346 fn schema(&self) -> serde_json::Value {
347 serde_json::json!({
348 "type": "object",
349 "properties": {
350 "query": {
351 "type": "string",
352 "description": "Keywords to match against entity names, summaries, attributes, or relations"
353 },
354 "entity_type": {
355 "type": "string",
356 "description": "Optional filter by entity category (e.g. 'system', 'dataset')"
357 },
358 "limit": {
359 "type": "integer",
360 "description": "Max results to return (default 10, max 50)"
361 }
362 }
363 })
364 }
365
366 async fn execute(
367 &self,
368 args: &serde_json::Value,
369 _ctx: &vak_tools::ToolContext,
370 ) -> vak_tools::ToolOutput {
371 let query = args.get("query").and_then(|v| v.as_str()).unwrap_or("");
372 let entity_type_filter = args
373 .get("entity_type")
374 .and_then(|v| v.as_str())
375 .map(|s| s.to_ascii_lowercase());
376 let limit = args
377 .get("limit")
378 .and_then(|v| v.as_u64())
379 .map(|l| l as usize)
380 .unwrap_or(10)
381 .min(50);
382
383 let mut results = search_entities(&self.sessions_home, Some(&self.cwd), query);
384 if let Some(ref filter) = entity_type_filter {
385 results.retain(|e| e.entity_type.to_ascii_lowercase() == *filter);
386 }
387 results.truncate(limit);
388
389 if results.is_empty() {
390 return vak_tools::ToolOutput::ok("no matching entities found in knowledge graph");
391 }
392
393 let mut out = format!("Found {} entities:\n\n", results.len());
394 for e in results {
395 out.push_str(&format!(
396 "### [{}] {} (`{}`)\n",
397 e.entity_type, e.name, e.id
398 ));
399 out.push_str(&format!("{}\n", e.summary));
400 if !e.attributes.is_empty() {
401 out.push_str("Attributes:\n");
402 for (k, v) in &e.attributes {
403 out.push_str(&format!("- **{k}**: {v}\n"));
404 }
405 }
406 if !e.relations.is_empty() {
407 out.push_str("Relations:\n");
408 for r in &e.relations {
409 out.push_str(&format!("- {} -> `{}`\n", r.relation, r.target_entity_id));
410 }
411 }
412 out.push('\n');
413 }
414
415 vak_tools::ToolOutput::ok(out.trim_end().to_string())
416 }
417
418 fn claims(&self, _args: &serde_json::Value) -> vak_tools::ResourceClaims {
419 vak_tools::ResourceClaims {
420 exclusive: false,
421 read_only: true,
422 paths: vec![],
423 }
424 }
425}
426
427#[cfg(test)]
428mod tests {
429 #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
430 use super::*;
431
432 #[test]
433 fn entity_crud_and_search_roundtrip() {
434 let temp = tempfile::tempdir().unwrap();
435 let home = temp.path().join("home");
436 let cwd = temp.path().join("cwd");
437 std::fs::create_dir_all(&home).unwrap();
438 std::fs::create_dir_all(&cwd).unwrap();
439
440 let mut attrs = BTreeMap::new();
441 attrs.insert("role".into(), "primary_db".into());
442 attrs.insert("engine".into(), "postgresql".into());
443
444 let entity = EntityRecord {
445 id: "ent-postgres-1".into(),
446 name: "Main PostgreSQL Cluster".into(),
447 entity_type: "database".into(),
448 summary: "Primary transaction database holding customer accounts".into(),
449 attributes: attrs,
450 relations: vec![EntityRelation {
451 relation: "hosted_on".into(),
452 target_entity_id: "srv-aws-us-east-1".into(),
453 }],
454 updated_at: Utc::now(),
455 };
456
457 // 1. Upsert
458 let saved = upsert_entity(&home, Some(&cwd), entity.clone()).unwrap();
459 assert_eq!(saved.id, "ent-postgres-1");
460
461 // 2. Get
462 let fetched = get_entity(&home, Some(&cwd), "ent-postgres-1").unwrap();
463 assert_eq!(fetched.name, "Main PostgreSQL Cluster");
464 assert_eq!(fetched.attributes.get("engine").unwrap(), "postgresql");
465
466 // 3. Search
467 let results = search_entities(&home, Some(&cwd), "transaction");
468 assert_eq!(results.len(), 1);
469 assert_eq!(results[0].id, "ent-postgres-1");
470
471 let rel_results = search_entities(&home, Some(&cwd), "aws-us-east-1");
472 assert_eq!(rel_results.len(), 1);
473
474 let none_results = search_entities(&home, Some(&cwd), "redis");
475 assert!(none_results.is_empty());
476
477 // 4. Update
478 let mut updated = fetched;
479 updated.summary = "Updated summary".into();
480 upsert_entity(&home, Some(&cwd), updated).unwrap();
481 let re_fetched = get_entity(&home, Some(&cwd), "ent-postgres-1").unwrap();
482 assert_eq!(re_fetched.summary, "Updated summary");
483
484 // 5. Delete
485 assert!(delete_entity(&home, Some(&cwd), "ent-postgres-1").unwrap());
486 assert!(get_entity(&home, Some(&cwd), "ent-postgres-1").is_none());
487 assert!(!delete_entity(&home, Some(&cwd), "ent-postgres-1").unwrap());
488 }
489
490 #[tokio::test]
491 async fn entity_record_and_query_tools_execute() {
492 use vak_tools::Tool;
493
494 let temp = tempfile::tempdir().unwrap();
495 let home = temp.path().join("home");
496 let cwd = temp.path().join("cwd");
497 std::fs::create_dir_all(&home).unwrap();
498 std::fs::create_dir_all(&cwd).unwrap();
499
500 let record_tool = EntityRecordTool {
501 sessions_home: home.clone(),
502 cwd: cwd.clone(),
503 };
504 let query_tool = EntityQueryTool {
505 sessions_home: home,
506 cwd: cwd.clone(),
507 };
508 let ctx = vak_tools::ToolContext {
509 cwd,
510 cancel: tokio_util::sync::CancellationToken::new(),
511 sandbox: None,
512 sandbox_sink: None,
513 agent_id: None,
514 new_documents: Vec::new(),
515 };
516
517 // Record entity via tool
518 let record_args = serde_json::json!({
519 "name": "Auth Gateway",
520 "entity_type": "service",
521 "summary": "Handles JWT authentication and token exchange",
522 "attributes": { "port": "8080", "protocol": "https" },
523 "relations": [ { "relation": "depends_on", "target": "redis-session-store" } ]
524 });
525 let rec_out = record_tool.execute(&record_args, &ctx).await;
526 assert!(!rec_out.is_error, "{}", rec_out.content);
527 assert!(
528 rec_out.content.contains("recorded entity 'auth-gateway'"),
529 "{}",
530 rec_out.content
531 );
532
533 // Query entity via tool
534 let query_args = serde_json::json!({ "query": "JWT" });
535 let q_out = query_tool.execute(&query_args, &ctx).await;
536 assert!(!q_out.is_error, "{}", q_out.content);
537 assert!(q_out.content.contains("Auth Gateway"), "{}", q_out.content);
538 assert!(
539 q_out
540 .content
541 .contains("depends_on -> `redis-session-store`"),
542 "{}",
543 q_out.content
544 );
545
546 // Filter query by entity_type
547 let type_args = serde_json::json!({ "entity_type": "service" });
548 let type_out = query_tool.execute(&type_args, &ctx).await;
549 assert!(!type_out.is_error, "{}", type_out.content);
550 assert!(
551 type_out.content.contains("Auth Gateway"),
552 "{}",
553 type_out.content
554 );
555
556 let none_args = serde_json::json!({ "entity_type": "database" });
557 let none_out = query_tool.execute(&none_args, &ctx).await;
558 assert!(!none_out.is_error, "{}", none_out.content);
559 assert!(
560 none_out.content.contains("no matching entities"),
561 "{}",
562 none_out.content
563 );
564 }
565}
566

Indexing the workspace…