Dark-Alex-17/loki
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implemented a built-in task management system to help smaller LLMs complete larger multistep tasks and minimize context drift

Browse Source
This commit is contained in:
Alex Clarke
2026-02-09 12:49:06 -07:00
parent 8a37a88ffd
commit a935add2a7
13 changed files with 868 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/function/mod.rs
+18 -4
View File
@@ -1,3 +1,5 @@
pub(crate) mod todo;
use crate::{
config::{Agent, Config, GlobalConfig},
utils::*,
@@ -26,6 +28,7 @@ use std::{
process::{Command, Stdio},
};
use strum_macros::AsRefStr;
use todo::TODO_FUNCTION_PREFIX;
#[derive(Embed)]
#[folder = "assets/functions/"]
@@ -262,6 +265,10 @@ impl Functions {
self.declarations.is_empty()
}
pub fn append_todo_functions(&mut self) {
self.declarations.extend(todo::todo_function_declarations());
}
pub fn clear_mcp_meta_functions(&mut self) {
self.declarations.retain(|d| {
!d.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
@@ -850,7 +857,7 @@ impl ToolCall {
_ if cmd_name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) => {
Self::search_mcp_tools(config, &cmd_name, &json_data).unwrap_or_else(|e| {
let error_msg = format!("MCP search failed: {e}");
println!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
json!({"tool_call_error": error_msg})
})
}
@@ -859,7 +866,7 @@ impl ToolCall {
.await
.unwrap_or_else(|e| {
let error_msg = format!("MCP describe failed: {e}");
println!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
json!({"tool_call_error": error_msg})
})
}
@@ -868,10 +875,17 @@ impl ToolCall {
.await
.unwrap_or_else(|e| {
let error_msg = format!("MCP tool invocation failed: {e}");
println!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
json!({"tool_call_error": error_msg})
})
}
_ if cmd_name.starts_with(TODO_FUNCTION_PREFIX) => {
todo::handle_todo_tool(config, &cmd_name, &json_data).unwrap_or_else(|e| {
let error_msg = format!("Todo tool failed: {e}");
eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
json!({"tool_call_error": error_msg})
})
}
_ => match run_llm_function(cmd_name, cmd_args, envs, agent_name) {
Ok(Some(contents)) => serde_json::from_str(&contents)
.ok()
@@ -1052,7 +1066,7 @@ pub fn run_llm_function(
eprintln!("{stderr}");
}
let tool_error_message = format!("Tool call '{command_name}' exited with code {exit_code}");
println!("{}", warning_text(&format!("⚠️ {tool_error_message} ⚠️")));
eprintln!("{}", warning_text(&format!("⚠️ {tool_error_message} ⚠️")));
let mut error_json = json!({"tool_call_error": tool_error_message});
if !stderr.is_empty() {
error_json["stderr"] = json!(stderr);
src/function/todo.rs
+160
View File
@@ -0,0 +1,160 @@
use super::{FunctionDeclaration, JsonSchema};
use crate::config::GlobalConfig;
use anyhow::{Result, bail};
use indexmap::IndexMap;
use serde_json::{Value, json};
pub const TODO_FUNCTION_PREFIX: &str = "todo__";
pub fn todo_function_declarations() -> Vec<FunctionDeclaration> {
vec![
FunctionDeclaration {
name: format!("{TODO_FUNCTION_PREFIX}init"),
description: "Initialize a new todo list with a goal. Clears any existing todos."
.to_string(),
parameters: JsonSchema {
type_value: Some("object".to_string()),
properties: Some(IndexMap::from([(
"goal".to_string(),
JsonSchema {
type_value: Some("string".to_string()),
description: Some(
"The overall goal to achieve when all todos are completed".into(),
),
..Default::default()
},
)])),
required: Some(vec!["goal".to_string()]),
..Default::default()
},
agent: false,
},
FunctionDeclaration {
name: format!("{TODO_FUNCTION_PREFIX}add"),
description: "Add a new todo item to the list.".to_string(),
parameters: JsonSchema {
type_value: Some("object".to_string()),
properties: Some(IndexMap::from([(
"task".to_string(),
JsonSchema {
type_value: Some("string".to_string()),
description: Some("Description of the todo task".into()),
..Default::default()
},
)])),
required: Some(vec!["task".to_string()]),
..Default::default()
},
agent: false,
},
FunctionDeclaration {
name: format!("{TODO_FUNCTION_PREFIX}done"),
description: "Mark a todo item as done by its id.".to_string(),
parameters: JsonSchema {
type_value: Some("object".to_string()),
properties: Some(IndexMap::from([(
"id".to_string(),
JsonSchema {
type_value: Some("integer".to_string()),
description: Some("The id of the todo item to mark as done".into()),
..Default::default()
},
)])),
required: Some(vec!["id".to_string()]),
..Default::default()
},
agent: false,
},
FunctionDeclaration {
name: format!("{TODO_FUNCTION_PREFIX}list"),
description: "Display the current todo list with status of each item.".to_string(),
parameters: JsonSchema {
type_value: Some("object".to_string()),
..Default::default()
},
agent: false,
},
]
}
pub fn handle_todo_tool(config: &GlobalConfig, cmd_name: &str, args: &Value) -> Result<Value> {
let action = cmd_name
.strip_prefix(TODO_FUNCTION_PREFIX)
.unwrap_or(cmd_name);
match action {
"init" => {
let goal = args.get("goal").and_then(Value::as_str).unwrap_or_default();
let mut cfg = config.write();
let agent = cfg.agent.as_mut();
match agent {
Some(agent) => {
agent.init_todo_list(goal);
Ok(json!({"status": "ok", "message": "Initialized new todo list"}))
}
None => bail!("No active agent"),
}
}
"add" => {
let task = args.get("task").and_then(Value::as_str).unwrap_or_default();
if task.is_empty() {
return Ok(json!({"error": "task description is required"}));
}
let mut cfg = config.write();
let agent = cfg.agent.as_mut();
match agent {
Some(agent) => {
let id = agent.add_todo(task);
Ok(json!({"status": "ok", "id": id}))
}
None => bail!("No active agent"),
}
}
"done" => {
let id = args
.get("id")
.and_then(|v| {
v.as_u64()
.or_else(|| v.as_str().and_then(|s| s.parse().ok()))
})
.map(|v| v as usize);
match id {
Some(id) => {
let mut cfg = config.write();
let agent = cfg.agent.as_mut();
match agent {
Some(agent) => {
if agent.mark_todo_done(id) {
Ok(
json!({"status": "ok", "message": format!("Marked todo {id} as done")}),
)
} else {
Ok(json!({"error": format!("Todo {id} not found")}))
}
}
None => bail!("No active agent"),
}
}
None => Ok(json!({"error": "id is required and must be a number"})),
}
}
"list" => {
let cfg = config.read();
let agent = cfg.agent.as_ref();
match agent {
Some(agent) => {
let list = agent.todo_list();
if list.is_empty() {
Ok(json!({"goal": "", "todos": []}))
} else {
Ok(serde_json::to_value(list)
.unwrap_or(json!({"error": "serialization failed"})))
}
}
None => bail!("No active agent"),
}
}
_ => bail!("Unknown todo action: {action}"),
}
}