ive harnessed the harness
refactor tools into separate modules
+794
-530
+23
-21
klbr-core/src/agent.rs
+23
-21
klbr-core/src/agent.rs
···
9
9
interrupt::Interrupt,
10
10
llm::{LlmClient, LlmEvent, Message},
11
11
memory::MemoryStore,
12
-
tools, AgentEvent, AgentMetrics,
12
+
tools::{self, ToolContext},
13
+
AgentEvent, AgentMetrics,
13
14
};
14
15
15
16
pub async fn run(
···
34
35
turn_count = pairs.len();
35
36
}
36
37
37
-
let tool_defs = tools::definitions();
38
+
let tool_ctx = ToolContext::new(memory.clone(), llm.clone());
39
+
let registry = tools::all_tools();
38
40
39
41
while let Some(interrupt) = rx.recv().await {
40
42
match interrupt {
···
46
48
}
47
49
Interrupt::Compact => {
48
50
let _ = output.send(AgentEvent::Status("compacting...".into()));
49
-
if let Err(e) = compact(&llm, &memory, &mut ctx, 0, &output).await {
51
+
if let Err(e) = compact(&tool_ctx, &mut ctx, 0, &output).await {
50
52
let _ = output.send(AgentEvent::Status(format!("compaction failed: {e}")));
51
53
} else {
52
54
turn_count = ctx.turn_count();
···
94
96
let (tok_tx, mut tok_rx) = mpsc::channel(256);
95
97
let llm2 = llm.clone();
96
98
let msgs = ctx.as_messages();
97
-
let defs = tool_defs.clone();
99
+
let defs = registry.definitions();
98
100
tokio::spawn(async move {
99
101
let _ = llm2.stream(&msgs, &defs, tok_tx).await;
100
102
});
···
135
137
args: args.clone(),
136
138
});
137
139
138
-
let result = tools::execute(call, &memory, &llm).await;
140
+
let result = registry.execute(call, &tool_ctx).await;
139
141
140
142
let _ = output.send(AgentEvent::ToolResult {
141
143
name: name.clone(),
···
166
168
167
169
if ctx.total_tokens > config.watermark_tokens {
168
170
let _ = output.send(AgentEvent::Status("compacting...".into()));
169
-
if let Err(e) =
170
-
compact(&llm, &memory, &mut ctx, config.compaction_keep, &output).await
171
+
if let Err(e) = compact(&tool_ctx, &mut ctx, config.compaction_keep, &output).await
171
172
{
172
173
let _ = output.send(AgentEvent::Status(format!("compaction failed: {e}")));
173
174
} else {
···
183
184
}
184
185
185
186
async fn compact(
186
-
llm: &LlmClient,
187
-
memory: &MemoryStore,
187
+
tool_ctx: &ToolContext,
188
188
ctx: &mut Context,
189
189
keep: usize,
190
190
output: &broadcast::Sender<AgentEvent>,
191
191
) -> Result<()> {
192
192
// run reflection before draining so the agent can curate memories
193
193
let _ = output.send(AgentEvent::Status("reflecting...".into()));
194
-
if let Err(e) = reflect(llm, memory, ctx).await {
194
+
if let Err(e) = reflect(tool_ctx, ctx).await {
195
195
tracing::warn!(err = %e, "reflection failed");
196
196
}
197
197
···
217
217
"summarize these conversation turns concisely, preserving key facts, decisions, and topics:\n\n{turns_text}"
218
218
))];
219
219
220
-
let (summary, _) = llm.complete(&prompt).await?;
220
+
let (summary, _) = tool_ctx.llm.complete(&prompt).await?;
221
221
if !summary.is_empty() {
222
-
let emb = llm.embed(&summary).await?;
223
-
memory.store(&summary, &emb, &["compaction_summary".to_string()])?;
222
+
let emb = tool_ctx.llm.embed(&summary).await?;
223
+
tool_ctx
224
+
.memory
225
+
.store(&summary, &emb, &["compaction_summary".to_string()])?;
224
226
}
225
227
226
228
// rebuild context anchor with freshly updated pinned memories
227
-
let pinned = memory.pinned_memories().unwrap_or_default();
229
+
let pinned = tool_ctx.memory.pinned_memories().unwrap_or_default();
228
230
ctx.update_anchor_memories(&pinned);
229
231
230
232
Ok(())
···
232
234
233
235
/// ephemeral reflection loop: let the agent review and curate its memories
234
236
/// without touching the main conversation context
235
-
async fn reflect(llm: &LlmClient, memory: &MemoryStore, ctx: &Context) -> Result<()> {
236
-
let pinned = memory.pinned_memories().unwrap_or_default();
237
-
let unpinned = memory.recent_unpinned(20).unwrap_or_default();
237
+
async fn reflect(tool_ctx: &ToolContext, ctx: &Context) -> Result<()> {
238
+
let pinned = tool_ctx.memory.pinned_memories().unwrap_or_default();
239
+
let unpinned = tool_ctx.memory.recent_unpinned(20).unwrap_or_default();
238
240
239
241
let pinned_text = if pinned.is_empty() {
240
242
"(none yet)".to_string()
···
296
298
be selective — pinned memories appear in every context window."
297
299
);
298
300
299
-
let defs = tools::reflection_definitions();
301
+
let reflect_registry = tools::memory_tools();
300
302
let mut msgs = vec![Message::user(reflection_prompt)];
301
303
302
304
// mini tool loop, max 6 iterations
303
305
for _ in 0..6 {
304
306
let (tok_tx, mut tok_rx) = mpsc::channel(128);
305
-
let llm2 = llm.clone();
307
+
let llm2 = tool_ctx.llm.clone();
306
308
let msgs_snap = msgs.clone();
307
-
let defs_snap = defs.clone();
309
+
let defs_snap = reflect_registry.definitions();
308
310
tokio::spawn(async move {
309
311
let _ = llm2.stream(&msgs_snap, &defs_snap, tok_tx).await;
310
312
});
···
325
327
326
328
msgs.push(Message::with_tool_calls(tool_calls.clone()));
327
329
for call in &tool_calls {
328
-
let result = tools::execute(call, memory, llm).await;
330
+
let result = reflect_registry.execute(call, tool_ctx).await;
329
331
msgs.push(Message::tool_result(&call.id, &result));
330
332
}
331
333
}
-509
klbr-core/src/tools.rs
-509
klbr-core/src/tools.rs
···
1
-
use serde_json::json;
2
-
use tokio::process::Command;
3
-
4
-
use crate::{
5
-
llm::{LlmClient, ToolCall, ToolDef},
6
-
memory::MemoryStore,
7
-
};
8
-
9
-
/// built-in tool definitions sent to the model
10
-
pub fn definitions() -> Vec<ToolDef> {
11
-
vec![
12
-
ToolDef::function(
13
-
"shell",
14
-
"execute a shell command and return its stdout/stderr. use for running code, \
15
-
file operations, system queries, etc.",
16
-
json!({
17
-
"type": "object",
18
-
"properties": {
19
-
"cmd": {
20
-
"type": "string",
21
-
"description": "the shell command to run"
22
-
}
23
-
},
24
-
"required": ["cmd"]
25
-
}),
26
-
),
27
-
ToolDef::function(
28
-
"read_file",
29
-
"read the contents of a file, optionally limited to a line range",
30
-
json!({
31
-
"type": "object",
32
-
"properties": {
33
-
"path": {
34
-
"type": "string",
35
-
"description": "absolute or relative path to the file"
36
-
},
37
-
"start_line": {
38
-
"type": "integer",
39
-
"description": "first line to return (1-based, inclusive). omit for beginning."
40
-
},
41
-
"end_line": {
42
-
"type": "integer",
43
-
"description": "last line to return (1-based, inclusive). omit for end of file."
44
-
}
45
-
},
46
-
"required": ["path"]
47
-
}),
48
-
),
49
-
ToolDef::function(
50
-
"write_file",
51
-
"write text content to a file, creating or overwriting it",
52
-
json!({
53
-
"type": "object",
54
-
"properties": {
55
-
"path": {
56
-
"type": "string",
57
-
"description": "path to the file to write"
58
-
},
59
-
"content": {
60
-
"type": "string",
61
-
"description": "content to write"
62
-
}
63
-
},
64
-
"required": ["path", "content"]
65
-
}),
66
-
),
67
-
ToolDef::function(
68
-
"remember",
69
-
"store something in long-term memory. use whenever you learn something worth keeping \
70
-
across sessions — user preferences, facts about projects, decisions, names, etc. \
71
-
set important=true to pin it so it's always visible at startup.",
72
-
json!({
73
-
"type": "object",
74
-
"properties": {
75
-
"content": {
76
-
"type": "string",
77
-
"description": "the fact or note to remember, written concisely"
78
-
},
79
-
"important": {
80
-
"type": "boolean",
81
-
"description": "if true, pin this memory so it appears at every startup"
82
-
},
83
-
"tags": {
84
-
"type": "array",
85
-
"items": { "type": "string" },
86
-
"description": "optional category tags, e.g. [\"preference\", \"project\", \"person\"]"
87
-
}
88
-
},
89
-
"required": ["content"]
90
-
}),
91
-
),
92
-
ToolDef::function(
93
-
"recall",
94
-
"semantic search over long-term memory. \
95
-
with no tags: searches all memories by meaning. \
96
-
with tags: searches only within memories that match those tags, \
97
-
ranked by semantic similarity (never misses a tag-matched memory due to global ranking). \
98
-
provide at least a query.",
99
-
json!({
100
-
"type": "object",
101
-
"properties": {
102
-
"query": {
103
-
"type": "string",
104
-
"description": "what to search for by meaning"
105
-
},
106
-
"tags": {
107
-
"type": "array",
108
-
"items": { "type": "string" },
109
-
"description": "restrict search to memories with these tags, e.g. [\"person:mayer\"] or [\"preference\"]"
110
-
},
111
-
"tag_mode": {
112
-
"type": "string",
113
-
"enum": ["and", "or"],
114
-
"description": "\"and\" = all tags must match, \"or\" = any tag matches (default: \"or\")"
115
-
},
116
-
"limit": {
117
-
"type": "integer",
118
-
"description": "max results (default 5)"
119
-
}
120
-
},
121
-
"required": ["query"]
122
-
}),
123
-
),
124
-
ToolDef::function(
125
-
"context_for",
126
-
"fetch all memories associated with a tag — a person, project, topic, etc. \
127
-
use this to load everything you know about someone or something before responding. \
128
-
no semantic ranking; returns newest first.",
129
-
json!({
130
-
"type": "object",
131
-
"properties": {
132
-
"tags": {
133
-
"type": "array",
134
-
"items": { "type": "string" },
135
-
"description": "tags to fetch, e.g. [\"person:mayer\"] or [\"project:klbr\", \"preference\"]"
136
-
},
137
-
"tag_mode": {
138
-
"type": "string",
139
-
"enum": ["and", "or"],
140
-
"description": "\"and\" = all tags must match, \"or\" = any tag matches (default: \"or\")"
141
-
},
142
-
"limit": {
143
-
"type": "integer",
144
-
"description": "max results (default 20)"
145
-
}
146
-
},
147
-
"required": ["tags"]
148
-
}),
149
-
),
150
-
ToolDef::function(
151
-
"tag_memory",
152
-
"set or replace the tags on an existing memory",
153
-
json!({
154
-
"type": "object",
155
-
"properties": {
156
-
"id": {
157
-
"type": "integer",
158
-
"description": "memory id"
159
-
},
160
-
"tags": {
161
-
"type": "array",
162
-
"items": { "type": "string" },
163
-
"description": "new tag list (replaces existing tags)"
164
-
}
165
-
},
166
-
"required": ["id", "tags"]
167
-
}),
168
-
),
169
-
ToolDef::function(
170
-
"pin_memory",
171
-
"pin an existing memory so it appears at every startup. use during reflection to \
172
-
promote unpinned memories that turned out to be long-term important.",
173
-
json!({
174
-
"type": "object",
175
-
"properties": {
176
-
"id": {
177
-
"type": "integer",
178
-
"description": "memory id (shown in list_memories or recall results)"
179
-
}
180
-
},
181
-
"required": ["id"]
182
-
}),
183
-
),
184
-
ToolDef::function(
185
-
"unpin_memory",
186
-
"unpin a memory so it's no longer shown at startup (still searchable). use during \
187
-
reflection to demote pinned memories that are no longer important or accurate.",
188
-
json!({
189
-
"type": "object",
190
-
"properties": {
191
-
"id": {
192
-
"type": "integer",
193
-
"description": "memory id"
194
-
}
195
-
},
196
-
"required": ["id"]
197
-
}),
198
-
),
199
-
ToolDef::function(
200
-
"list_memories",
201
-
"list current pinned memories and recent unpinned memories with their ids. \
202
-
useful before a reflection pass to see what's stored.",
203
-
json!({
204
-
"type": "object",
205
-
"properties": {},
206
-
"required": []
207
-
}),
208
-
),
209
-
]
210
-
}
211
-
212
-
/// tool definitions used only during reflection (memory management tools only)
213
-
pub fn reflection_definitions() -> Vec<ToolDef> {
214
-
definitions()
215
-
.into_iter()
216
-
.filter(|d| {
217
-
matches!(
218
-
d.function.name.as_str(),
219
-
"remember"
220
-
| "recall"
221
-
| "context_for"
222
-
| "tag_memory"
223
-
| "pin_memory"
224
-
| "unpin_memory"
225
-
| "list_memories"
226
-
)
227
-
})
228
-
.collect()
229
-
}
230
-
231
-
/// dispatch a tool call and return the result as a string
232
-
pub async fn execute(call: &ToolCall, memory: &MemoryStore, llm: &LlmClient) -> String {
233
-
let args: serde_json::Value =
234
-
serde_json::from_str(&call.function.arguments).unwrap_or_default();
235
-
236
-
match call.function.name.as_str() {
237
-
"shell" => {
238
-
let cmd = match args["cmd"].as_str() {
239
-
Some(c) => c.to_string(),
240
-
None => return "error: missing required arg 'cmd'".into(),
241
-
};
242
-
match Command::new("sh").arg("-c").arg(&cmd).output().await {
243
-
Ok(out) => {
244
-
let stdout = String::from_utf8_lossy(&out.stdout);
245
-
let stderr = String::from_utf8_lossy(&out.stderr);
246
-
let code = out.status.code().unwrap_or(-1);
247
-
let mut result = format!("exit {code}");
248
-
if !stdout.is_empty() {
249
-
result.push('\n');
250
-
if stdout.len() > 20_000 {
251
-
result.push_str(&stdout[..20_000]);
252
-
result.push_str("\n[...truncated]");
253
-
} else {
254
-
result.push_str(&stdout);
255
-
}
256
-
}
257
-
if !stderr.is_empty() {
258
-
result.push_str("\nstderr:\n");
259
-
if stderr.len() > 5_000 {
260
-
result.push_str(&stderr[..5_000]);
261
-
result.push_str("\n[...truncated]");
262
-
} else {
263
-
result.push_str(&stderr);
264
-
}
265
-
}
266
-
result
267
-
}
268
-
Err(e) => format!("error: {e}"),
269
-
}
270
-
}
271
-
272
-
"read_file" => {
273
-
let path = match args["path"].as_str() {
274
-
Some(p) => p.to_string(),
275
-
None => return "error: missing required arg 'path'".into(),
276
-
};
277
-
match tokio::fs::read_to_string(&path).await {
278
-
Ok(content) => {
279
-
let start = args["start_line"]
280
-
.as_u64()
281
-
.map(|n| n.saturating_sub(1) as usize);
282
-
let end = args["end_line"].as_u64().map(|n| n as usize);
283
-
284
-
if start.is_none() && end.is_none() {
285
-
if content.len() > 50_000 {
286
-
format!(
287
-
"{}\n[...truncated, {} total bytes]",
288
-
&content[..50_000],
289
-
content.len()
290
-
)
291
-
} else {
292
-
content
293
-
}
294
-
} else {
295
-
let lines: Vec<&str> = content.lines().collect();
296
-
let s = start.unwrap_or(0).min(lines.len());
297
-
let e = end.unwrap_or(lines.len()).min(lines.len());
298
-
if s >= e {
299
-
return format!("no lines in range {}..{}", s + 1, e);
300
-
}
301
-
lines[s..e].join("\n")
302
-
}
303
-
}
304
-
Err(e) => format!("error: {e}"),
305
-
}
306
-
}
307
-
308
-
"write_file" => {
309
-
let path = match args["path"].as_str() {
310
-
Some(p) => p.to_string(),
311
-
None => return "error: missing required arg 'path'".into(),
312
-
};
313
-
let content = args["content"].as_str().unwrap_or("").to_string();
314
-
match tokio::fs::write(&path, &content).await {
315
-
Ok(_) => format!("wrote {} bytes to {path}", content.len()),
316
-
Err(e) => format!("error: {e}"),
317
-
}
318
-
}
319
-
320
-
"remember" => {
321
-
let content = match args["content"].as_str() {
322
-
Some(c) => c.to_string(),
323
-
None => return "error: missing required arg 'content'".into(),
324
-
};
325
-
let important = args["important"].as_bool().unwrap_or(false);
326
-
let tags: Vec<String> = args["tags"]
327
-
.as_array()
328
-
.map(|arr| {
329
-
arr.iter()
330
-
.filter_map(|v| v.as_str().map(String::from))
331
-
.collect()
332
-
})
333
-
.unwrap_or_default();
334
-
match llm.embed(&content).await {
335
-
Ok(emb) => match memory.store(&content, &emb, &tags) {
336
-
Ok(id) => {
337
-
if important {
338
-
if let Err(e) = memory.set_pinned(id, true) {
339
-
return format!("stored (id:{id}) but pin failed: {e}");
340
-
}
341
-
let tag_info = if tags.is_empty() {
342
-
String::new()
343
-
} else {
344
-
format!(", tags: {}", tags.join(", "))
345
-
};
346
-
format!("stored and pinned (id:{id}{tag_info})")
347
-
} else {
348
-
let tag_info = if tags.is_empty() {
349
-
String::new()
350
-
} else {
351
-
format!(", tags: {}", tags.join(", "))
352
-
};
353
-
format!("stored (id:{id}{tag_info})")
354
-
}
355
-
}
356
-
Err(e) => format!("error storing memory: {e}"),
357
-
},
358
-
Err(e) => format!("error embedding memory: {e}"),
359
-
}
360
-
}
361
-
362
-
"recall" => {
363
-
let query = match args["query"].as_str() {
364
-
Some(q) => q.to_string(),
365
-
None => return "error: missing required arg 'query'".into(),
366
-
};
367
-
let tags: Vec<String> = args["tags"]
368
-
.as_array()
369
-
.map(|arr| {
370
-
arr.iter()
371
-
.filter_map(|v| v.as_str().map(String::from))
372
-
.collect()
373
-
})
374
-
.unwrap_or_default();
375
-
let tag_and = args["tag_mode"].as_str() == Some("and");
376
-
let limit = args["limit"].as_u64().unwrap_or(5) as usize;
377
-
378
-
let emb = match llm.embed(&query).await {
379
-
Ok(e) => e,
380
-
Err(err) => return format!("error embedding query: {err}"),
381
-
};
382
-
383
-
match memory.recall(Some(&emb), &tags, tag_and, limit) {
384
-
Ok(results) if results.is_empty() => "no memories found".into(),
385
-
Ok(results) => results
386
-
.into_iter()
387
-
.map(|e| {
388
-
let tag_str = if e.tags.is_empty() {
389
-
String::new()
390
-
} else {
391
-
format!(" [{}]", e.tags.join(", "))
392
-
};
393
-
format!(
394
-
"[dist:{:.3}][id:{}]{tag_str} {}",
395
-
e.distance.unwrap_or(0.0),
396
-
e.id,
397
-
e.content
398
-
)
399
-
})
400
-
.collect::<Vec<_>>()
401
-
.join("\n"),
402
-
Err(err) => format!("error: {err}"),
403
-
}
404
-
}
405
-
406
-
"context_for" => {
407
-
let tags: Vec<String> = match args["tags"].as_array() {
408
-
Some(arr) => arr
409
-
.iter()
410
-
.filter_map(|v| v.as_str().map(String::from))
411
-
.collect(),
412
-
None => return "error: missing required arg 'tags'".into(),
413
-
};
414
-
let tag_and = args["tag_mode"].as_str() == Some("and");
415
-
let limit = args["limit"].as_u64().unwrap_or(20) as usize;
416
-
417
-
match memory.context_for(&tags, tag_and, limit) {
418
-
Ok(results) if results.is_empty() => {
419
-
format!("no memories found for tags: {}", tags.join(", "))
420
-
}
421
-
Ok(results) => results
422
-
.into_iter()
423
-
.map(|e| {
424
-
let tag_str = if e.tags.is_empty() {
425
-
String::new()
426
-
} else {
427
-
format!(" [{}]", e.tags.join(", "))
428
-
};
429
-
format!("[id:{}]{tag_str} {}", e.id, e.content)
430
-
})
431
-
.collect::<Vec<_>>()
432
-
.join("\n"),
433
-
Err(err) => format!("error: {err}"),
434
-
}
435
-
}
436
-
437
-
"tag_memory" => {
438
-
let id = match args["id"].as_i64() {
439
-
Some(i) => i,
440
-
None => return "error: missing required arg 'id'".into(),
441
-
};
442
-
let tags: Vec<String> = args["tags"]
443
-
.as_array()
444
-
.map(|arr| {
445
-
arr.iter()
446
-
.filter_map(|v| v.as_str().map(String::from))
447
-
.collect()
448
-
})
449
-
.unwrap_or_default();
450
-
match memory.set_tags(id, &tags) {
451
-
Ok(_) => format!("tagged memory {id} with: {}", tags.join(", ")),
452
-
Err(e) => format!("error: {e}"),
453
-
}
454
-
}
455
-
456
-
"pin_memory" => {
457
-
let id = match args["id"].as_i64() {
458
-
Some(i) => i,
459
-
None => return "error: missing required arg 'id'".into(),
460
-
};
461
-
match memory.set_pinned(id, true) {
462
-
Ok(_) => format!("pinned memory {id}"),
463
-
Err(e) => format!("error: {e}"),
464
-
}
465
-
}
466
-
467
-
"unpin_memory" => {
468
-
let id = match args["id"].as_i64() {
469
-
Some(i) => i,
470
-
None => return "error: missing required arg 'id'".into(),
471
-
};
472
-
match memory.set_pinned(id, false) {
473
-
Ok(_) => format!("unpinned memory {id}"),
474
-
Err(e) => format!("error: {e}"),
475
-
}
476
-
}
477
-
478
-
"list_memories" => {
479
-
let pinned = memory.pinned_memories().unwrap_or_default();
480
-
let unpinned = memory.recent_unpinned(10).unwrap_or_default();
481
-
482
-
let mut out = String::new();
483
-
out.push_str("## pinned\n");
484
-
if pinned.is_empty() {
485
-
out.push_str("(none)\n");
486
-
} else {
487
-
for (i, content) in pinned.iter().enumerate() {
488
-
out.push_str(&format!("{}. {content}\n", i + 1));
489
-
}
490
-
}
491
-
out.push_str("\n## recent unpinned\n");
492
-
if unpinned.is_empty() {
493
-
out.push_str("(none)\n");
494
-
} else {
495
-
for (id, content, tags) in &unpinned {
496
-
let tag_str = if tags.is_empty() {
497
-
String::new()
498
-
} else {
499
-
format!(" [{}]", tags.join(", "))
500
-
};
501
-
out.push_str(&format!("[id:{id}]{tag_str} {content}\n"));
502
-
}
503
-
}
504
-
out
505
-
}
506
-
507
-
other => format!("unknown tool: {other}"),
508
-
}
509
-
}
+76
klbr-core/src/tools/context_for.rs
+76
klbr-core/src/tools/context_for.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
6
+
use crate::llm::ToolDef;
7
+
8
+
use super::{Tool, ToolContext};
9
+
10
+
pub fn tool() -> Tool {
11
+
Tool::new(definition(), exec)
12
+
}
13
+
14
+
fn definition() -> ToolDef {
15
+
ToolDef::function(
16
+
"context_for",
17
+
"fetch all memories associated with a tag — a person, project, topic, etc. \
18
+
use this to load everything you know about someone or something before responding. \
19
+
no semantic ranking; returns newest first.",
20
+
json!({
21
+
"type": "object",
22
+
"properties": {
23
+
"tags": {
24
+
"type": "array",
25
+
"items": { "type": "string" },
26
+
"description": "tags to fetch, e.g. [\"person:mayer\"] or [\"project:klbr\", \"preference\"]"
27
+
},
28
+
"tag_mode": {
29
+
"type": "string",
30
+
"enum": ["and", "or"],
31
+
"description": "\"and\" = all tags must match, \"or\" = any tag matches (default: \"or\")"
32
+
},
33
+
"limit": {
34
+
"type": "integer",
35
+
"description": "max results (default 20)"
36
+
}
37
+
},
38
+
"required": ["tags"]
39
+
}),
40
+
)
41
+
}
42
+
43
+
fn exec(args: serde_json::Value, ctx: ToolContext) -> Pin<Box<dyn Future<Output = String> + Send>> {
44
+
Box::pin(execute(args, ctx))
45
+
}
46
+
47
+
async fn execute(args: serde_json::Value, ctx: ToolContext) -> String {
48
+
let tags: Vec<String> = match args["tags"].as_array() {
49
+
Some(arr) => arr
50
+
.iter()
51
+
.filter_map(|v| v.as_str().map(String::from))
52
+
.collect(),
53
+
None => return "error: missing required arg 'tags'".into(),
54
+
};
55
+
let tag_and = args["tag_mode"].as_str() == Some("and");
56
+
let limit = args["limit"].as_u64().unwrap_or(20) as usize;
57
+
58
+
match ctx.memory.context_for(&tags, tag_and, limit) {
59
+
Ok(results) if results.is_empty() => {
60
+
format!("no memories found for tags: {}", tags.join(", "))
61
+
}
62
+
Ok(results) => results
63
+
.into_iter()
64
+
.map(|e| {
65
+
let tag_str = if e.tags.is_empty() {
66
+
String::new()
67
+
} else {
68
+
format!(" [{}]", e.tags.join(", "))
69
+
};
70
+
format!("[id:{}]{tag_str} {}", e.id, e.content)
71
+
})
72
+
.collect::<Vec<_>>()
73
+
.join("\n"),
74
+
Err(err) => format!("error: {err}"),
75
+
}
76
+
}
+58
klbr-core/src/tools/list_memories.rs
+58
klbr-core/src/tools/list_memories.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
6
+
use crate::llm::ToolDef;
7
+
8
+
use super::{Tool, ToolContext};
9
+
10
+
pub fn tool() -> Tool {
11
+
Tool::new(definition(), exec)
12
+
}
13
+
14
+
fn definition() -> ToolDef {
15
+
ToolDef::function(
16
+
"list_memories",
17
+
"list current pinned memories and recent unpinned memories with their ids. \
18
+
useful before a reflection pass to see what's stored.",
19
+
json!({
20
+
"type": "object",
21
+
"properties": {},
22
+
"required": []
23
+
}),
24
+
)
25
+
}
26
+
27
+
fn exec(args: serde_json::Value, ctx: ToolContext) -> Pin<Box<dyn Future<Output = String> + Send>> {
28
+
Box::pin(execute(args, ctx))
29
+
}
30
+
31
+
async fn execute(_args: serde_json::Value, ctx: ToolContext) -> String {
32
+
let pinned = ctx.memory.pinned_memories().unwrap_or_default();
33
+
let unpinned = ctx.memory.recent_unpinned(10).unwrap_or_default();
34
+
35
+
let mut out = String::new();
36
+
out.push_str("## pinned\n");
37
+
if pinned.is_empty() {
38
+
out.push_str("(none)\n");
39
+
} else {
40
+
for (i, content) in pinned.iter().enumerate() {
41
+
out.push_str(&format!("{}. {content}\n", i + 1));
42
+
}
43
+
}
44
+
out.push_str("\n## recent unpinned\n");
45
+
if unpinned.is_empty() {
46
+
out.push_str("(none)\n");
47
+
} else {
48
+
for (id, content, tags) in &unpinned {
49
+
let tag_str = if tags.is_empty() {
50
+
String::new()
51
+
} else {
52
+
format!(" [{}]", tags.join(", "))
53
+
};
54
+
out.push_str(&format!("[id:{id}]{tag_str} {content}\n"));
55
+
}
56
+
}
57
+
out
58
+
}
+104
klbr-core/src/tools/mod.rs
+104
klbr-core/src/tools/mod.rs
···
1
+
mod context_for;
2
+
mod list_memories;
3
+
mod pin_memory;
4
+
mod read_file;
5
+
mod recall;
6
+
mod remember;
7
+
mod shell;
8
+
mod tag_memory;
9
+
mod unpin_memory;
10
+
mod write_file;
11
+
12
+
use std::collections::HashMap;
13
+
use std::future::Future;
14
+
use std::pin::Pin;
15
+
16
+
use crate::llm::{LlmClient, ToolCall, ToolDef};
17
+
use crate::memory::MemoryStore;
18
+
19
+
/// runtime dependencies passed to every tool execution
20
+
#[derive(Clone)]
21
+
pub struct ToolContext {
22
+
pub memory: MemoryStore,
23
+
pub llm: LlmClient,
24
+
}
25
+
26
+
impl ToolContext {
27
+
pub fn new(memory: MemoryStore, llm: LlmClient) -> Self {
28
+
Self { memory, llm }
29
+
}
30
+
}
31
+
32
+
type ToolFn = fn(serde_json::Value, ToolContext) -> Pin<Box<dyn Future<Output = String> + Send>>;
33
+
34
+
pub struct Tool {
35
+
pub definition: ToolDef,
36
+
exec: ToolFn,
37
+
}
38
+
39
+
impl Tool {
40
+
pub fn new(definition: ToolDef, exec: ToolFn) -> Self {
41
+
Self { definition, exec }
42
+
}
43
+
44
+
async fn run(&self, args: serde_json::Value, ctx: ToolContext) -> String {
45
+
(self.exec)(args, ctx).await
46
+
}
47
+
}
48
+
49
+
/// a registered set of tools; call `execute` to dispatch by name
50
+
pub struct Subroutines {
51
+
tools: HashMap<String, Tool>,
52
+
}
53
+
54
+
impl Subroutines {
55
+
fn new(tools: Vec<Tool>) -> Self {
56
+
Self {
57
+
tools: tools
58
+
.into_iter()
59
+
.map(|t| (t.definition.function.name.clone(), t))
60
+
.collect(),
61
+
}
62
+
}
63
+
64
+
pub fn definitions(&self) -> Vec<ToolDef> {
65
+
self.tools.values().map(|t| t.definition.clone()).collect()
66
+
}
67
+
68
+
pub async fn execute(&self, call: &ToolCall, ctx: &ToolContext) -> String {
69
+
let args: serde_json::Value =
70
+
serde_json::from_str(&call.function.arguments).unwrap_or_default();
71
+
match self.tools.get(&call.function.name) {
72
+
Some(tool) => tool.run(args, ctx.clone()).await,
73
+
None => format!("unknown tool: {}", call.function.name),
74
+
}
75
+
}
76
+
}
77
+
78
+
/// tools used only during reflection (memory management)
79
+
pub fn memory_tools() -> Subroutines {
80
+
Subroutines::new(vec![
81
+
remember::tool(),
82
+
recall::tool(),
83
+
context_for::tool(),
84
+
tag_memory::tool(),
85
+
pin_memory::tool(),
86
+
unpin_memory::tool(),
87
+
list_memories::tool(),
88
+
])
89
+
}
90
+
91
+
/// all built-in tools
92
+
pub fn all_tools() -> Subroutines {
93
+
let mut tools = vec![shell::tool(), read_file::tool(), write_file::tool()];
94
+
tools.extend([
95
+
remember::tool(),
96
+
recall::tool(),
97
+
context_for::tool(),
98
+
tag_memory::tool(),
99
+
pin_memory::tool(),
100
+
unpin_memory::tool(),
101
+
list_memories::tool(),
102
+
]);
103
+
Subroutines::new(tools)
104
+
}
+45
klbr-core/src/tools/pin_memory.rs
+45
klbr-core/src/tools/pin_memory.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
6
+
use crate::llm::ToolDef;
7
+
8
+
use super::{Tool, ToolContext};
9
+
10
+
pub fn tool() -> Tool {
11
+
Tool::new(definition(), exec)
12
+
}
13
+
14
+
fn definition() -> ToolDef {
15
+
ToolDef::function(
16
+
"pin_memory",
17
+
"pin an existing memory so it appears at every startup. use during reflection to \
18
+
promote unpinned memories that turned out to be long-term important.",
19
+
json!({
20
+
"type": "object",
21
+
"properties": {
22
+
"id": {
23
+
"type": "integer",
24
+
"description": "memory id (shown in list_memories or recall results)"
25
+
}
26
+
},
27
+
"required": ["id"]
28
+
}),
29
+
)
30
+
}
31
+
32
+
fn exec(args: serde_json::Value, ctx: ToolContext) -> Pin<Box<dyn Future<Output = String> + Send>> {
33
+
Box::pin(execute(args, ctx))
34
+
}
35
+
36
+
async fn execute(args: serde_json::Value, ctx: ToolContext) -> String {
37
+
let id = match args["id"].as_i64() {
38
+
Some(i) => i,
39
+
None => return "error: missing required arg 'id'".into(),
40
+
};
41
+
match ctx.memory.set_pinned(id, true) {
42
+
Ok(_) => format!("pinned memory {id}"),
43
+
Err(e) => format!("error: {e}"),
44
+
}
45
+
}
+80
klbr-core/src/tools/read_file.rs
+80
klbr-core/src/tools/read_file.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
6
+
use crate::llm::ToolDef;
7
+
8
+
use super::{Tool, ToolContext};
9
+
10
+
pub fn tool() -> Tool {
11
+
Tool::new(definition(), exec)
12
+
}
13
+
14
+
fn definition() -> ToolDef {
15
+
ToolDef::function(
16
+
"read_file",
17
+
"read the contents of a file, optionally limited to a line range",
18
+
json!({
19
+
"type": "object",
20
+
"properties": {
21
+
"path": {
22
+
"type": "string",
23
+
"description": "absolute or relative path to the file"
24
+
},
25
+
"start_line": {
26
+
"type": "integer",
27
+
"description": "first line to return (1-based, inclusive). omit for beginning."
28
+
},
29
+
"end_line": {
30
+
"type": "integer",
31
+
"description": "last line to return (1-based, inclusive). omit for end of file."
32
+
}
33
+
},
34
+
"required": ["path"]
35
+
}),
36
+
)
37
+
}
38
+
39
+
fn exec(
40
+
args: serde_json::Value,
41
+
_ctx: ToolContext,
42
+
) -> Pin<Box<dyn Future<Output = String> + Send>> {
43
+
Box::pin(execute(args))
44
+
}
45
+
46
+
async fn execute(args: serde_json::Value) -> String {
47
+
let path = match args["path"].as_str() {
48
+
Some(p) => p.to_string(),
49
+
None => return "error: missing required arg 'path'".into(),
50
+
};
51
+
match tokio::fs::read_to_string(&path).await {
52
+
Ok(content) => {
53
+
let start = args["start_line"]
54
+
.as_u64()
55
+
.map(|n| n.saturating_sub(1) as usize);
56
+
let end = args["end_line"].as_u64().map(|n| n as usize);
57
+
58
+
if start.is_none() && end.is_none() {
59
+
if content.len() > 50_000 {
60
+
format!(
61
+
"{}\n[...truncated, {} total bytes]",
62
+
&content[..50_000],
63
+
content.len()
64
+
)
65
+
} else {
66
+
content
67
+
}
68
+
} else {
69
+
let lines: Vec<&str> = content.lines().collect();
70
+
let s = start.unwrap_or(0).min(lines.len());
71
+
let e = end.unwrap_or(lines.len()).min(lines.len());
72
+
if s >= e {
73
+
return format!("no lines in range {}..{}", s + 1, e);
74
+
}
75
+
lines[s..e].join("\n")
76
+
}
77
+
}
78
+
Err(e) => format!("error: {e}"),
79
+
}
80
+
}
+95
klbr-core/src/tools/recall.rs
+95
klbr-core/src/tools/recall.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
6
+
use crate::llm::ToolDef;
7
+
8
+
use super::{Tool, ToolContext};
9
+
10
+
pub fn tool() -> Tool {
11
+
Tool::new(definition(), exec)
12
+
}
13
+
14
+
fn definition() -> ToolDef {
15
+
ToolDef::function(
16
+
"recall",
17
+
"semantic search over long-term memory. \
18
+
with no tags: searches all memories by meaning. \
19
+
with tags: searches only within memories that match those tags, \
20
+
ranked by semantic similarity (never misses a tag-matched memory due to global ranking). \
21
+
provide at least a query.",
22
+
json!({
23
+
"type": "object",
24
+
"properties": {
25
+
"query": {
26
+
"type": "string",
27
+
"description": "what to search for by meaning"
28
+
},
29
+
"tags": {
30
+
"type": "array",
31
+
"items": { "type": "string" },
32
+
"description": "restrict search to memories with these tags, e.g. [\"person:mayer\"] or [\"preference\"]"
33
+
},
34
+
"tag_mode": {
35
+
"type": "string",
36
+
"enum": ["and", "or"],
37
+
"description": "\"and\" = all tags must match, \"or\" = any tag matches (default: \"or\")"
38
+
},
39
+
"limit": {
40
+
"type": "integer",
41
+
"description": "max results (default 5)"
42
+
}
43
+
},
44
+
"required": ["query"]
45
+
}),
46
+
)
47
+
}
48
+
49
+
fn exec(args: serde_json::Value, ctx: ToolContext) -> Pin<Box<dyn Future<Output = String> + Send>> {
50
+
Box::pin(execute(args, ctx))
51
+
}
52
+
53
+
async fn execute(args: serde_json::Value, ctx: ToolContext) -> String {
54
+
let query = match args["query"].as_str() {
55
+
Some(q) => q.to_string(),
56
+
None => return "error: missing required arg 'query'".into(),
57
+
};
58
+
let tags: Vec<String> = args["tags"]
59
+
.as_array()
60
+
.map(|arr| {
61
+
arr.iter()
62
+
.filter_map(|v| v.as_str().map(String::from))
63
+
.collect()
64
+
})
65
+
.unwrap_or_default();
66
+
let tag_and = args["tag_mode"].as_str() == Some("and");
67
+
let limit = args["limit"].as_u64().unwrap_or(5) as usize;
68
+
69
+
let emb = match ctx.llm.embed(&query).await {
70
+
Ok(e) => e,
71
+
Err(err) => return format!("error embedding query: {err}"),
72
+
};
73
+
74
+
match ctx.memory.recall(Some(&emb), &tags, tag_and, limit) {
75
+
Ok(results) if results.is_empty() => "no memories found".into(),
76
+
Ok(results) => results
77
+
.into_iter()
78
+
.map(|e| {
79
+
let tag_str = if e.tags.is_empty() {
80
+
String::new()
81
+
} else {
82
+
format!(" [{}]", e.tags.join(", "))
83
+
};
84
+
format!(
85
+
"[dist:{:.3}][id:{}]{tag_str} {}",
86
+
e.distance.unwrap_or(0.0),
87
+
e.id,
88
+
e.content
89
+
)
90
+
})
91
+
.collect::<Vec<_>>()
92
+
.join("\n"),
93
+
Err(err) => format!("error: {err}"),
94
+
}
95
+
}
+86
klbr-core/src/tools/remember.rs
+86
klbr-core/src/tools/remember.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
6
+
use crate::llm::ToolDef;
7
+
8
+
use super::{Tool, ToolContext};
9
+
10
+
pub fn tool() -> Tool {
11
+
Tool::new(definition(), exec)
12
+
}
13
+
14
+
fn definition() -> ToolDef {
15
+
ToolDef::function(
16
+
"remember",
17
+
"store something in long-term memory. use whenever you learn something worth keeping \
18
+
across sessions — user preferences, facts about projects, decisions, names, etc. \
19
+
set important=true to pin it so it's always visible at startup.",
20
+
json!({
21
+
"type": "object",
22
+
"properties": {
23
+
"content": {
24
+
"type": "string",
25
+
"description": "the fact or note to remember, written concisely"
26
+
},
27
+
"important": {
28
+
"type": "boolean",
29
+
"description": "if true, pin this memory so it appears at every startup"
30
+
},
31
+
"tags": {
32
+
"type": "array",
33
+
"items": { "type": "string" },
34
+
"description": "optional category tags, e.g. [\"preference\", \"project\", \"person\"]"
35
+
}
36
+
},
37
+
"required": ["content"]
38
+
}),
39
+
)
40
+
}
41
+
42
+
fn exec(args: serde_json::Value, ctx: ToolContext) -> Pin<Box<dyn Future<Output = String> + Send>> {
43
+
Box::pin(execute(args, ctx))
44
+
}
45
+
46
+
async fn execute(args: serde_json::Value, ctx: ToolContext) -> String {
47
+
let content = match args["content"].as_str() {
48
+
Some(c) => c.to_string(),
49
+
None => return "error: missing required arg 'content'".into(),
50
+
};
51
+
let important = args["important"].as_bool().unwrap_or(false);
52
+
let tags: Vec<String> = args["tags"]
53
+
.as_array()
54
+
.map(|arr| {
55
+
arr.iter()
56
+
.filter_map(|v| v.as_str().map(String::from))
57
+
.collect()
58
+
})
59
+
.unwrap_or_default();
60
+
match ctx.llm.embed(&content).await {
61
+
Ok(emb) => match ctx.memory.store(&content, &emb, &tags) {
62
+
Ok(id) => {
63
+
if important {
64
+
if let Err(e) = ctx.memory.set_pinned(id, true) {
65
+
return format!("stored (id:{id}) but pin failed: {e}");
66
+
}
67
+
let tag_info = if tags.is_empty() {
68
+
String::new()
69
+
} else {
70
+
format!(", tags: {}", tags.join(", "))
71
+
};
72
+
format!("stored and pinned (id:{id}{tag_info})")
73
+
} else {
74
+
let tag_info = if tags.is_empty() {
75
+
String::new()
76
+
} else {
77
+
format!(", tags: {}", tags.join(", "))
78
+
};
79
+
format!("stored (id:{id}{tag_info})")
80
+
}
81
+
}
82
+
Err(e) => format!("error storing memory: {e}"),
83
+
},
84
+
Err(e) => format!("error embedding memory: {e}"),
85
+
}
86
+
}
+73
klbr-core/src/tools/shell.rs
+73
klbr-core/src/tools/shell.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
use tokio::process::Command;
6
+
7
+
use crate::llm::ToolDef;
8
+
9
+
use super::{Tool, ToolContext};
10
+
11
+
pub fn tool() -> Tool {
12
+
Tool::new(definition(), exec)
13
+
}
14
+
15
+
fn definition() -> ToolDef {
16
+
ToolDef::function(
17
+
"shell",
18
+
"execute a shell command and return its stdout/stderr. use for running code, \
19
+
file operations, system queries, etc.",
20
+
json!({
21
+
"type": "object",
22
+
"properties": {
23
+
"cmd": {
24
+
"type": "string",
25
+
"description": "the shell command to run"
26
+
}
27
+
},
28
+
"required": ["cmd"]
29
+
}),
30
+
)
31
+
}
32
+
33
+
fn exec(
34
+
args: serde_json::Value,
35
+
_ctx: ToolContext,
36
+
) -> Pin<Box<dyn Future<Output = String> + Send>> {
37
+
Box::pin(execute(args))
38
+
}
39
+
40
+
async fn execute(args: serde_json::Value) -> String {
41
+
let cmd = match args["cmd"].as_str() {
42
+
Some(c) => c.to_string(),
43
+
None => return "error: missing required arg 'cmd'".into(),
44
+
};
45
+
match Command::new("sh").arg("-c").arg(&cmd).output().await {
46
+
Ok(out) => {
47
+
let stdout = String::from_utf8_lossy(&out.stdout);
48
+
let stderr = String::from_utf8_lossy(&out.stderr);
49
+
let code = out.status.code().unwrap_or(-1);
50
+
let mut result = format!("exit {code}");
51
+
if !stdout.is_empty() {
52
+
result.push('\n');
53
+
if stdout.len() > 20_000 {
54
+
result.push_str(&stdout[..20_000]);
55
+
result.push_str("\n[...truncated]");
56
+
} else {
57
+
result.push_str(&stdout);
58
+
}
59
+
}
60
+
if !stderr.is_empty() {
61
+
result.push_str("\nstderr:\n");
62
+
if stderr.len() > 5_000 {
63
+
result.push_str(&stderr[..5_000]);
64
+
result.push_str("\n[...truncated]");
65
+
} else {
66
+
result.push_str(&stderr);
67
+
}
68
+
}
69
+
result
70
+
}
71
+
Err(e) => format!("error: {e}"),
72
+
}
73
+
}
+57
klbr-core/src/tools/tag_memory.rs
+57
klbr-core/src/tools/tag_memory.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
6
+
use crate::llm::ToolDef;
7
+
8
+
use super::{Tool, ToolContext};
9
+
10
+
pub fn tool() -> Tool {
11
+
Tool::new(definition(), exec)
12
+
}
13
+
14
+
fn definition() -> ToolDef {
15
+
ToolDef::function(
16
+
"tag_memory",
17
+
"set or replace the tags on an existing memory",
18
+
json!({
19
+
"type": "object",
20
+
"properties": {
21
+
"id": {
22
+
"type": "integer",
23
+
"description": "memory id"
24
+
},
25
+
"tags": {
26
+
"type": "array",
27
+
"items": { "type": "string" },
28
+
"description": "new tag list (replaces existing tags)"
29
+
}
30
+
},
31
+
"required": ["id", "tags"]
32
+
}),
33
+
)
34
+
}
35
+
36
+
fn exec(args: serde_json::Value, ctx: ToolContext) -> Pin<Box<dyn Future<Output = String> + Send>> {
37
+
Box::pin(execute(args, ctx))
38
+
}
39
+
40
+
async fn execute(args: serde_json::Value, ctx: ToolContext) -> String {
41
+
let id = match args["id"].as_i64() {
42
+
Some(i) => i,
43
+
None => return "error: missing required arg 'id'".into(),
44
+
};
45
+
let tags: Vec<String> = args["tags"]
46
+
.as_array()
47
+
.map(|arr| {
48
+
arr.iter()
49
+
.filter_map(|v| v.as_str().map(String::from))
50
+
.collect()
51
+
})
52
+
.unwrap_or_default();
53
+
match ctx.memory.set_tags(id, &tags) {
54
+
Ok(_) => format!("tagged memory {id} with: {}", tags.join(", ")),
55
+
Err(e) => format!("error: {e}"),
56
+
}
57
+
}
+45
klbr-core/src/tools/unpin_memory.rs
+45
klbr-core/src/tools/unpin_memory.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
6
+
use crate::llm::ToolDef;
7
+
8
+
use super::{Tool, ToolContext};
9
+
10
+
pub fn tool() -> Tool {
11
+
Tool::new(definition(), exec)
12
+
}
13
+
14
+
fn definition() -> ToolDef {
15
+
ToolDef::function(
16
+
"unpin_memory",
17
+
"unpin a memory so it's no longer shown at startup (still searchable). use during \
18
+
reflection to demote pinned memories that are no longer important or accurate.",
19
+
json!({
20
+
"type": "object",
21
+
"properties": {
22
+
"id": {
23
+
"type": "integer",
24
+
"description": "memory id"
25
+
}
26
+
},
27
+
"required": ["id"]
28
+
}),
29
+
)
30
+
}
31
+
32
+
fn exec(args: serde_json::Value, ctx: ToolContext) -> Pin<Box<dyn Future<Output = String> + Send>> {
33
+
Box::pin(execute(args, ctx))
34
+
}
35
+
36
+
async fn execute(args: serde_json::Value, ctx: ToolContext) -> String {
37
+
let id = match args["id"].as_i64() {
38
+
Some(i) => i,
39
+
None => return "error: missing required arg 'id'".into(),
40
+
};
41
+
match ctx.memory.set_pinned(id, false) {
42
+
Ok(_) => format!("unpinned memory {id}"),
43
+
Err(e) => format!("error: {e}"),
44
+
}
45
+
}
+52
klbr-core/src/tools/write_file.rs
+52
klbr-core/src/tools/write_file.rs
···
1
+
use std::future::Future;
2
+
use std::pin::Pin;
3
+
4
+
use serde_json::json;
5
+
6
+
use crate::llm::ToolDef;
7
+
8
+
use super::{Tool, ToolContext};
9
+
10
+
pub fn tool() -> Tool {
11
+
Tool::new(definition(), exec)
12
+
}
13
+
14
+
fn definition() -> ToolDef {
15
+
ToolDef::function(
16
+
"write_file",
17
+
"write text content to a file, creating or overwriting it",
18
+
json!({
19
+
"type": "object",
20
+
"properties": {
21
+
"path": {
22
+
"type": "string",
23
+
"description": "path to the file to write"
24
+
},
25
+
"content": {
26
+
"type": "string",
27
+
"description": "content to write"
28
+
}
29
+
},
30
+
"required": ["path", "content"]
31
+
}),
32
+
)
33
+
}
34
+
35
+
fn exec(
36
+
args: serde_json::Value,
37
+
_ctx: ToolContext,
38
+
) -> Pin<Box<dyn Future<Output = String> + Send>> {
39
+
Box::pin(execute(args))
40
+
}
41
+
42
+
async fn execute(args: serde_json::Value) -> String {
43
+
let path = match args["path"].as_str() {
44
+
Some(p) => p.to_string(),
45
+
None => return "error: missing required arg 'path'".into(),
46
+
};
47
+
let content = args["content"].as_str().unwrap_or("").to_string();
48
+
match tokio::fs::write(&path, &content).await {
49
+
Ok(_) => format!("wrote {} bytes to {path}", content.len()),
50
+
Err(e) => format!("error: {e}"),
51
+
}
52
+
}