llm_coding_tools_rig/
todo.rs1use llm_coding_tools_core::operations::{read_todos, write_todos};
6use llm_coding_tools_core::tool_names;
7use llm_coding_tools_core::{ToolContext, ToolError, ToolOutput};
8use rig::completion::ToolDefinition;
9use rig::tool::Tool;
10use schemars::{schema_for, JsonSchema};
11use serde::Deserialize;
12
13pub use llm_coding_tools_core::{Todo, TodoPriority, TodoState, TodoStatus};
15
16#[derive(Debug, Clone, Deserialize, JsonSchema)]
18pub struct TodoWriteArgs {
19 pub todos: Vec<Todo>,
21}
22
23#[derive(Debug, Clone, Deserialize, JsonSchema)]
25pub struct TodoReadArgs {}
26
27#[derive(Debug, Clone)]
29pub struct TodoWriteTool {
30 state: TodoState,
31}
32
33impl TodoWriteTool {
34 pub fn new(state: TodoState) -> Self {
36 Self { state }
37 }
38}
39
40impl Tool for TodoWriteTool {
41 const NAME: &'static str = tool_names::TODO_WRITE;
42
43 type Error = ToolError;
44 type Args = TodoWriteArgs;
45 type Output = ToolOutput;
46
47 async fn definition(&self, _prompt: String) -> ToolDefinition {
48 ToolDefinition {
49 name: <Self as Tool>::NAME.to_string(),
50 description: "Replace the todo list with new items.".to_string(),
51 parameters: serde_json::to_value(schema_for!(TodoWriteArgs))
52 .expect("schema serialization should never fail"),
53 }
54 }
55
56 async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
57 let message = write_todos(&self.state, args.todos)?;
58 Ok(ToolOutput::new(message))
59 }
60}
61
62#[derive(Debug, Clone)]
64pub struct TodoReadTool {
65 state: TodoState,
66}
67
68impl TodoReadTool {
69 pub fn new(state: TodoState) -> Self {
71 Self { state }
72 }
73}
74
75impl Tool for TodoReadTool {
76 const NAME: &'static str = tool_names::TODO_READ;
77
78 type Error = ToolError;
79 type Args = TodoReadArgs;
80 type Output = ToolOutput;
81
82 async fn definition(&self, _prompt: String) -> ToolDefinition {
83 ToolDefinition {
84 name: <Self as Tool>::NAME.to_string(),
85 description: "Read the current todo list.".to_string(),
86 parameters: serde_json::to_value(schema_for!(TodoReadArgs))
87 .expect("schema serialization should never fail"),
88 }
89 }
90
91 async fn call(&self, _args: Self::Args) -> Result<Self::Output, Self::Error> {
92 let content = read_todos(&self.state);
93 Ok(ToolOutput::new(content))
94 }
95}
96
97impl ToolContext for TodoWriteTool {
98 const NAME: &'static str = tool_names::TODO_WRITE;
99
100 fn context(&self) -> &'static str {
101 llm_coding_tools_core::context::TODO_WRITE
102 }
103}
104
105impl ToolContext for TodoReadTool {
106 const NAME: &'static str = tool_names::TODO_READ;
107
108 fn context(&self) -> &'static str {
109 llm_coding_tools_core::context::TODO_READ
110 }
111}
112
113pub struct TodoTools {
115 pub write: TodoWriteTool,
117 pub read: TodoReadTool,
119}
120
121impl TodoTools {
122 pub fn new() -> Self {
124 let state = TodoState::new();
125 Self {
126 write: TodoWriteTool::new(state.clone()),
127 read: TodoReadTool::new(state),
128 }
129 }
130
131 pub fn with_state(state: TodoState) -> Self {
133 Self {
134 write: TodoWriteTool::new(state.clone()),
135 read: TodoReadTool::new(state),
136 }
137 }
138}
139
140impl Default for TodoTools {
141 fn default() -> Self {
142 Self::new()
143 }
144}
145
146#[cfg(test)]
147mod tests {
148 use super::*;
149
150 fn make_todo(id: &str, status: TodoStatus) -> Todo {
151 Todo {
152 id: id.to_string(),
153 content: format!("Task {id}"),
154 status,
155 priority: TodoPriority::Medium,
156 }
157 }
158
159 #[tokio::test]
160 async fn write_and_read_todos() {
161 let tools = TodoTools::new();
162
163 let write_args = TodoWriteArgs {
164 todos: vec![
165 make_todo("1", TodoStatus::Pending),
166 make_todo("2", TodoStatus::Completed),
167 ],
168 };
169 let write_result = tools.write.call(write_args).await.unwrap();
170 assert!(write_result.content.contains("2 task(s)"));
171
172 let read_result = tools.read.call(TodoReadArgs {}).await.unwrap();
173 assert!(read_result.content.contains("Task 1"));
174 assert!(read_result.content.contains("Task 2"));
175 }
176
177 #[tokio::test]
178 async fn shared_state_works() {
179 let state = TodoState::new();
180 let write_tool = TodoWriteTool::new(state.clone());
181 let read_tool = TodoReadTool::new(state);
182
183 let write_args = TodoWriteArgs {
184 todos: vec![make_todo("shared", TodoStatus::InProgress)],
185 };
186 write_tool.call(write_args).await.unwrap();
187
188 let read_result = read_tool.call(TodoReadArgs {}).await.unwrap();
189 assert!(read_result.content.contains("shared"));
190 }
191
192 #[tokio::test]
193 async fn empty_list_returns_no_tasks() {
194 let tools = TodoTools::new();
195 let result = tools.read.call(TodoReadArgs {}).await.unwrap();
196 assert_eq!(result.content, "No tasks.");
197 }
198}