From bb80abbd41b09571f2bb9e3aa1ecbc67cbfbcc01 Mon Sep 17 00:00:00 2001 From: taggart_comet Date: Mon, 19 Jan 2026 09:33:05 +0200 Subject: [PATCH] save temp progress --- Cargo.lock | 10 ++ Cargo.toml | 1 + cats/test_cat.py | 21 +++- src/domain/bt/mod.rs | 80 ++++++++++++++ src/domain/mod.rs | 1 + src/domain/prompting/general.rs | 25 ++++- src/domain/prompting/session.rs | 6 +- src/domain/prompting/tools.rs | 7 +- src/domain/session/service.rs | 24 ++++- src/domain/tools/discover_objects.rs | 4 +- src/domain/tools/find_files.rs | 9 +- src/domain/tools/mod.rs | 15 ++- .../tools/{patch_file.rs => patch_files.rs} | 82 ++++++-------- src/domain/tools/read_objects.rs | 87 +++++---------- src/domain/tools/shell_exec.rs | 4 +- src/domain/workflow/chain.rs | 5 +- src/domain/workflow/mod.rs | 3 +- src/domain/workflow/step.rs | 17 +-- .../{toolsets/general.rs => toolset/all.rs} | 16 +-- src/domain/workflow/toolset/edit.rs | 25 +++++ .../workflow/{toolset.rs => toolset/mod.rs} | 36 ++++++- src/domain/workflow/toolset/none.rs | 21 ++++ src/domain/workflow/toolset/read.rs | 34 ++++++ src/domain/workflow/toolsets/mod.rs | 3 - src/domain/workflow/workflow.rs | 100 +++++++++++++----- .../api_clients/openai/client.rs | 67 ++++++++---- .../api_clients/openai/dto/request.rs | 14 +-- src/lib.rs | 2 + src/main.rs | 36 +++++-- 29 files changed, 523 insertions(+), 232 deletions(-) create mode 100644 src/domain/bt/mod.rs rename src/domain/tools/{patch_file.rs => patch_files.rs} (81%) rename src/domain/workflow/{toolsets/general.rs => toolset/all.rs} (78%) create mode 100644 src/domain/workflow/toolset/edit.rs rename src/domain/workflow/{toolset.rs => toolset/mod.rs} (70%) create mode 100644 src/domain/workflow/toolset/none.rs create mode 100644 src/domain/workflow/toolset/read.rs delete mode 100644 src/domain/workflow/toolsets/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 423d955..20de478 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -252,6 +252,15 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "colored" +version = "3.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -337,6 +346,7 @@ name = "drastis" version = "0.1.0" dependencies = [ "clap", + "colored", "ctrlc", "directories", "env_logger", diff --git a/Cargo.toml b/Cargo.toml index bec9d8c..ea3cf09 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,6 +51,7 @@ reqwest = { version = "0.12", features = ["json", "blocking"] } tokio = { version = "1", features = ["rt-multi-thread", "macros"] } tempfile = "3.0" zenpatch = "0.1.0" +colored = "3.1.1" [dev-dependencies] diff --git a/cats/test_cat.py b/cats/test_cat.py index 6a81274..1b9a8b9 100644 --- a/cats/test_cat.py +++ b/cats/test_cat.py @@ -1,4 +1,3 @@ -import pytest from cat import Cat, create_cat @@ -46,3 +45,23 @@ def test_create_cat_helper(): cat = create_cat("Simba", 2) assert isinstance(cat, Cat) assert cat.name == "Simba" + + +def run_all_tests() -> None: + tests = [ + test_cat_creation, + test_meow, + test_feed_changes_mood, + test_pet_when_neutral, + test_pet_when_happy, + test_describe, + test_create_cat_helper, + ] + + for test in tests: + test() + + +if __name__ == "__main__": + run_all_tests() + print("All tests passed!") \ No newline at end of file diff --git a/src/domain/bt/mod.rs b/src/domain/bt/mod.rs new file mode 100644 index 0000000..65faf59 --- /dev/null +++ b/src/domain/bt/mod.rs @@ -0,0 +1,80 @@ +// module for behaviour trees definitions + +mod general; + +pub use general::GeneralTree; + +use crate::domain::workflow::toolset::ToolsetType; + +pub enum BTNodeStatus { + Pending, + Running, + Success, + Failure, +} + +pub trait BehaviorTree { + fn get_next_step(&self) -> Option<&BTStepNode>; +} + +pub struct BTStepNode { + prompt: String, + status: BTNodeStatus, + toolset: ToolsetType, + max_tools_calls: u32, + max_retries: u32, + next_step: Option>, + decision: Option>, +} + +pub trait BTStepNodeInterface { + fn prompt(&self) -> String; + fn toolset(&self) -> ToolsetType; + fn max_tools_calls(&self) -> u32; + fn max_retries(&self) -> u32; + fn next_step(&self) -> Option<&dyn BTStepNodeInterface>; + fn has_decision(&self) -> bool; + fn get_decision(&self) -> Option<&dyn BTDecisionNodeInterface>; +} + +pub struct BTDecisionNode { + prompt: String, + status: BTNodeStatus, + option_a: BTStepNode, + option_b: BTStepNode, +} + +pub trait BTDecisionNodeInterface { + fn prompt(&self) -> String; + fn get_next(&self, llm_output: &str) -> Option<&BTStepNode>; +} + +impl BTStepNodeInterface for BTStepNode { + fn prompt(&self) -> String { + self.prompt.clone() + } + + fn toolset(&self) -> ToolsetType { + self.toolset + } + + fn max_tools_calls(&self) -> u32 { + self.max_tools_calls + } + + fn max_retries(&self) -> u32 { + self.max_retries + } + + fn next_step(&self) -> Option<&dyn BTStepNodeInterface> { + self.next_step.as_deref() + } + + fn has_decision(&self) -> bool { + self.decision.is_some() + } + + fn get_decision(&self) -> Option<&dyn BTDecisionNodeInterface> { + self.decision.as_deref() + } +} diff --git a/src/domain/mod.rs b/src/domain/mod.rs index a112efc..98d4a51 100644 --- a/src/domain/mod.rs +++ b/src/domain/mod.rs @@ -4,6 +4,7 @@ pub mod session; pub mod startup; pub mod tools; pub mod workflow; +mod bt; pub use project::Project; pub use session::service::SessionService; diff --git a/src/domain/prompting/general.rs b/src/domain/prompting/general.rs index 9a5c215..16f06ff 100644 --- a/src/domain/prompting/general.rs +++ b/src/domain/prompting/general.rs @@ -1,8 +1,10 @@ -use crate::domain::prompting; +use crate::domain::{prompting, SessionRequest}; +use crate::domain::bt::BTStepNodeInterface; use crate::domain::session::Request; -use crate::domain::workflow::Chain; +use crate::domain::workflow::{Chain, ChainStep}; use crate::domain::workflow::Toolset; use crate::domain::ModelType; +use crate::domain::prompting::get_user_prompt; /// LLM prompt templates for the coding assistant /// @@ -64,14 +66,27 @@ OUTPUT (XML ONLY):\n", } } +pub fn get_bt_tree_step_prompt(model_type: ModelType, step: &dyn BTStepNodeInterface, request: &dyn Request) -> String { + if model_type == ModelType::OpenAI { + format!( + "Objective:\n{}\n\ + Current action:\n{}\n.", + request.current_request(), step.prompt() + ) + } else { + format!( + "{}\n{}", get_user_prompt(model_type, request), step.prompt() + ) + } +} + pub fn get_system_prompt(model_type: ModelType) -> String { let (os_name, shell_name) = get_runtime_environment(); if model_type == ModelType::OpenAI { format!( - "You are Drastis, a coding agent. \ + "You are Drastis, a coding agent. \n\ Use the available tools to gather context and make changes. \ -When using tools, pass JSON arguments that match their parameters. \ -If you call a tool, you MUST also include an output_text message explaining why. \ +When using tools, pass JSON arguments that match their parameters. \n\ Runtime: os={}, shell={}.", os_name, shell_name ) diff --git a/src/domain/prompting/session.rs b/src/domain/prompting/session.rs index 58422b0..cb6b528 100644 --- a/src/domain/prompting/session.rs +++ b/src/domain/prompting/session.rs @@ -32,15 +32,15 @@ pub fn session_naming_prompt(model_type: ModelType, prompt_preview: &str) -> Str if model_type == ModelType::OpenAI { format!( "Generate a short title (3-5 words) for the following request: {}\n\ - The title should be concise and easy to remember.", + Respond with ONLY the title itself. Do not add explanations, punctuation, or extra text.", prompt_preview ) } else { format!( - "<|im_start|>system\nYou generate very short titles (3-5 words). Respond with only the title.<|im_end|>\n\ + "<|im_start|>system\nYou generate very short titles (3-5 words). Respond with ONLY the title itself, without any extra text.<|im_end|>\n\ <|im_start|>user\nTitle for: {}<|im_end|>\n\ <|im_start|>assistant\n", prompt_preview ) } -} +} \ No newline at end of file diff --git a/src/domain/prompting/tools.rs b/src/domain/prompting/tools.rs index c1403ab..77753c5 100644 --- a/src/domain/prompting/tools.rs +++ b/src/domain/prompting/tools.rs @@ -21,7 +21,10 @@ pub fn get_tool_result( let tool_name = chain_step.tool_name.unwrap_or_else(|| "unspecified".to_string()); let output = chain_step.tool_output.unwrap_or_default(); format!( - "Tool `{}` execution output is: \n{}\n", - tool_name, output + "Tool `{}` execution output is: \n{}\n\ +---\n\ +Tool input was: \n\ +{}\n", + tool_name, output, chain_step.input_payload ) } diff --git a/src/domain/session/service.rs b/src/domain/session/service.rs index 6b66a54..0186d23 100644 --- a/src/domain/session/service.rs +++ b/src/domain/session/service.rs @@ -8,14 +8,20 @@ use std::sync::Arc; /// Service for running workflows on sessions pub struct SessionService { workflow: Workflow, + use_behavior_trees: bool, conn: Arc, } impl SessionService { /// Create a new session service with default workflow - pub fn new(engine: Arc, conn: Arc) -> Result { + pub fn new( + engine: Arc, + conn: Arc, + use_behavior_trees: bool, + ) -> Result { Ok(Self { workflow: Workflow::new(engine)?, + use_behavior_trees, conn, }) } @@ -25,7 +31,7 @@ impl SessionService { /// Creates a new session_request, runs the workflow, and updates the request with the result /// Returns the execution chain pub fn run( - &self, + &mut self, session: &Session, prompt: &str, cancel: &CancellationToken, @@ -45,7 +51,19 @@ impl SessionService { session_with_request.set_current_request(prompt.to_string()); // Run the workflow - let result = match self.workflow.run(&mut session_with_request, cancel) { + let result: Result = if self.use_behavior_trees { + self.workflow.reset_chain(); + self.workflow + .run_using_bt(&mut session_with_request, cancel) + } else { + self.workflow.reset_chain(); + self.workflow + .run(&mut session_with_request, cancel, 128, None) + .map_err(ServiceError::Workflow)?; + Ok(self.workflow.chain().clone()) + }; + + let result = match result { Ok(chain) => { // Get summary and log from chain let summary = chain.get_summary(); diff --git a/src/domain/tools/discover_objects.rs b/src/domain/tools/discover_objects.rs index 6ee5c5c..74e9e07 100644 --- a/src/domain/tools/discover_objects.rs +++ b/src/domain/tools/discover_objects.rs @@ -133,7 +133,7 @@ impl Tool for DiscoverObjects { "properties": { "full_path_to_file": { "type": "string", - "description": "path to the source file" + "description": "path to the source file, use `find_files` tool to determine the correct path to the file" } }, "required": ["full_path_to_file"], @@ -143,7 +143,7 @@ impl Tool for DiscoverObjects { fn desc(&self) -> String { format!( - "Use the `{}` tool to discover functions, classes, structs in a source file.", + "Use the `{}` tool to discover exact names for functions, classes, structs in a source file.", self.name() ) } diff --git a/src/domain/tools/find_files.rs b/src/domain/tools/find_files.rs index f46eaf6..050f40a 100644 --- a/src/domain/tools/find_files.rs +++ b/src/domain/tools/find_files.rs @@ -262,9 +262,14 @@ mod tests { #[test] fn test_resolve_search_root_with_input() { // When input root is provided, validate it's within project root - let result = FindFiles::resolve_search_root(Some("/tmp/subdir"), Path::new("/tmp")); + let tmp_dir = tempfile::tempdir().unwrap(); + let subdir = tmp_dir.path().join("subdir"); + std::fs::create_dir_all(&subdir).unwrap(); + + let result = + FindFiles::resolve_search_root(subdir.to_str(), tmp_dir.path()); assert!(result.is_ok()); - assert_eq!(result.unwrap(), PathBuf::from("/tmp/subdir")); + assert_eq!(result.unwrap(), subdir); } #[test] diff --git a/src/domain/tools/mod.rs b/src/domain/tools/mod.rs index e576222..e2398b9 100644 --- a/src/domain/tools/mod.rs +++ b/src/domain/tools/mod.rs @@ -1,13 +1,13 @@ mod discover_objects; mod find_files; -mod patch_file; +mod patch_files; mod read_objects; mod shell_exec; mod structure; pub use discover_objects::DiscoverObjects; pub use find_files::FindFiles; -pub use patch_file::PatchFile; +pub use patch_files::PatchFiles; pub use read_objects::ReadObjects; pub use shell_exec::ShellExec; pub use structure::Structure; @@ -99,14 +99,19 @@ impl ToolResult { /// /// Example: /// ```rust +/// use serde::Serialize; +/// use drastis::domain::tools::{serialize_output, ToolResult}; +/// /// #[derive(Serialize)] /// struct MyOutput { /// result: String, /// } /// /// let output = MyOutput { result: "success".to_string() }; -/// let xml = serialize_output(&output)?; -/// ToolResult::ok(tool_name, input, xml) +/// let xml = serialize_output(&output).unwrap(); +/// let tool_name = "my_tool".to_string(); +/// let input = "".to_string(); +/// let _result = ToolResult::ok(tool_name, input, xml); /// ``` pub fn serialize_output(value: &T) -> Result where @@ -139,7 +144,7 @@ pub fn build_tool_by_name(name: &str) -> Option> { "read_objects" => Some(Box::new(ReadObjects::new())), "find_files" => Some(Box::new(FindFiles::new())), "structure" => Some(Box::new(Structure::new())), - "patch_file" => Some(Box::new(PatchFile::new())), + "patch_files" => Some(Box::new(PatchFiles::new())), "shell_exec" => Some(Box::new(ShellExec::new())), _ => None, } diff --git a/src/domain/tools/patch_file.rs b/src/domain/tools/patch_files.rs similarity index 81% rename from src/domain/tools/patch_file.rs rename to src/domain/tools/patch_files.rs index 71d3921..3f85c54 100644 --- a/src/domain/tools/patch_file.rs +++ b/src/domain/tools/patch_files.rs @@ -9,47 +9,41 @@ use zenpatch::apply as apply_patch; use zenpatch::parser::text_to_patch::text_to_patch; use zenpatch::Vfs; -pub struct PatchFile { - input: Mutex>, +pub struct PatchFiles { + input: Mutex>, } #[derive(Debug, Clone)] -struct PatchFileInput { +struct PatchFilesInput { raw: String, - file_path: String, patch: String, } #[derive(Debug, Deserialize)] -struct PatchFileInputJson { - file_path: String, +struct PatchFilesInputJson { patch: String, } -impl PatchFile { +impl PatchFiles { pub fn new() -> Self { Self { input: Mutex::new(None), } } - fn parse_input_json(raw: &str) -> Result { - let parsed: PatchFileInputJson = + fn parse_input_json(raw: &str) -> Result { + let parsed: PatchFilesInputJson = serde_json::from_str(raw).map_err(|e| Error::Parse(e.to_string()))?; - if parsed.file_path.is_empty() { - return Err(Error::Parse("file_path is required".into())); - } if parsed.patch.trim().is_empty() { return Err(Error::Parse("patch is required".into())); } - Ok(PatchFileInput { + Ok(PatchFilesInput { raw: raw.to_string(), - file_path: parsed.file_path, patch: parsed.patch, }) } - fn load_input(&self) -> Result { + fn load_input(&self) -> Result { self.input .lock() .unwrap() @@ -58,9 +52,9 @@ impl PatchFile { } } -impl Tool for PatchFile { +impl Tool for PatchFiles { fn name(&self) -> &'static str { - "patch_file" + "patch_files" } fn parse_input(&self, input: String) -> Option { @@ -99,24 +93,6 @@ impl Tool for PatchFile { } }; - let input_path_match = actions.iter().any(|action| { - action.path == input.file_path - || action - .new_path - .as_deref() - .is_some_and(|path| path == input.file_path.as_str()) - }); - if !input_path_match { - return ToolResult::error( - self.name().to_string(), - input.raw, - format!( - "Patch does not reference file_path '{}'", - input.file_path - ), - ); - } - let mut vfs = Vfs::new(); let mut touched_paths: Vec = Vec::new(); for action in &actions { @@ -217,30 +193,36 @@ impl Tool for PatchFile { json!({ "type": "object", "properties": { - "file_path": { - "type": "string", - "description": "path to the file being patched" - }, "patch": { "type": "string", - "description": "Begin Patch / Update File patch format content" + "description": "Begin Patch / Update File patch content; can include multiple file updates. \ +Example: \ +*** Begin Patch\\n\ +*** Update File: foo.php\\n\ +@@\\n\ +-old\\n\ ++new\\n\ +*** Update File: bar.md\\n\ +@@\\n\ +-Old title\\n\ ++New title\\n\ +*** End Patch" } }, - "required": ["file_path", "patch"], - "additionalProperties": false + "required": ["patch"] }) } fn desc(&self) -> String { format!( - "Use the `{name}` tool to edit files using the Begin Patch / Update File patch format.", + "Use the `{name}` tool to edit one or more files using the Begin Patch / Update File patch format.", name = self.name() ) } } -impl Default for PatchFile { +impl Default for PatchFiles { fn default() -> Self { Self::new() } @@ -261,7 +243,7 @@ mod tests { let request = VirtualRequest::new("test", temp.path()); - let tool = PatchFile::new(); + let tool = PatchFiles::new(); let patch = "*** Begin Patch\n\ *** Update File: test.txt\n\ @@\n\ @@ -269,7 +251,7 @@ mod tests { +line2_modified\n\ *** End Patch"; let input = format!( - "{{\"file_path\":\"test.txt\",\"patch\":\"{}\"}}", + "{{\"patch\":\"{}\"}}", patch.replace('\n', "\\n").replace('"', "\\\"") ); assert!(tool.parse_input(input).is_none()); @@ -298,7 +280,7 @@ mod tests { let request = VirtualRequest::new("test", temp.path()); - let tool = PatchFile::new(); + let tool = PatchFiles::new(); let patch = "*** Begin Patch\n\ *** Update File: code.py\n\ @@\n\ @@ -309,7 +291,7 @@ mod tests { + pass\n\ *** End Patch"; let input = format!( - "{{\"file_path\":\"code.py\",\"patch\":\"{}\"}}", + "{{\"patch\":\"{}\"}}", patch.replace('\n', "\\n").replace('"', "\\\"") ); assert!(tool.parse_input(input).is_none()); @@ -331,8 +313,8 @@ mod tests { let request = VirtualRequest::new("test", temp.path()); - let tool = PatchFile::new(); - let input = "{\"file_path\":\"test.txt\",\"patch\":\"not a patch\"}".to_string(); + let tool = PatchFiles::new(); + let input = "{\"patch\":\"not a patch\"}".to_string(); assert!(tool.parse_input(input).is_none()); let result = tool.work(&request); diff --git a/src/domain/tools/read_objects.rs b/src/domain/tools/read_objects.rs index d23024b..1fc12e4 100644 --- a/src/domain/tools/read_objects.rs +++ b/src/domain/tools/read_objects.rs @@ -1,4 +1,4 @@ -use super::{Error, Tool, ToolResult}; +use super::{DiscoverObjects, Error, Tool, ToolResult}; use crate::domain::session::Request; use crate::utils::{Lang, ObjectKind, ParsedObject, UniversalParser}; use serde::Deserialize; @@ -22,8 +22,7 @@ struct ReadObjectsInput { #[derive(Debug, Deserialize)] struct ReadObjectsInputJson { path: String, - names: Vec, - kind: Option, + query: String, } #[derive(Debug, Serialize)] @@ -65,18 +64,6 @@ impl ObjectQuery { } } -#[derive(Serialize)] -struct ReadObjectsOutput { - language: String, - object_not_found: bool, - results: HashMap, -} - -#[derive(Serialize)] -struct ErrorOutput { - error: String, -} - impl ReadObjects { pub fn new() -> Self { Self { @@ -130,25 +117,18 @@ impl ReadObjects { return Err(Error::Parse("path is required".into())); } - if parsed.names.is_empty() { - return Err(Error::Parse("names is required".into())); - } - - let kind = parsed.kind.as_deref().and_then(ObjectKind::from_str); - if parsed.kind.is_some() && kind.is_none() { - return Err(Error::Parse("invalid kind".into())); + let mut names = Self::parse_query(&parsed.query); + names.retain(|name| !name.trim().is_empty()); + if names.is_empty() { + return Err(Error::Parse("query is required".into())); } let mut queries = Vec::new(); - for name in parsed.names { - if name.is_empty() { + for name in names { + if name.trim().is_empty() { return Err(Error::Parse("names cannot include empty strings".into())); } - let query = match kind { - Some(kind) => ObjectQuery::with_kind(kind, &name), - None => ObjectQuery::new(&name), - }; - queries.push(query); + queries.push(ObjectQuery::new(&name)); } Ok(ReadObjectsInput { @@ -166,6 +146,14 @@ impl ReadObjects { .ok_or_else(|| Error::Parse("input not parsed".into())) } + fn parse_query(query: &str) -> Vec { + query + .replace(',', " ") + .split_whitespace() + .map(|part| part.trim().to_string()) + .collect() + } + fn format_output(lang: Lang, results: HashMap) -> String { // Format as simple text output instead of YAML if results.is_empty() { @@ -234,27 +222,20 @@ impl Tool for ReadObjects { "type": "string", "description": "path to the source file" }, - "names": { - "type": "array", - "description": "object names to find", - "items": { - "type": "string" - } - }, - "kind": { + "query": { "type": "string", - "description": "optional: apply one kind to all names (function, class, struct, etc.)" + "description": "comma- or space-separated object names. Example: \"main, Config, Parser\"" } }, - "required": ["path", "names"], + "required": ["path", "query"], "additionalProperties": false }) } fn desc(&self) -> String { format!( - "Use the `{}` tool to read source code of specific objects from a file.", - self.name() + "Use the `{}` tool to read source code of specific objects from a file. To determine correct properties to use for `{}`, use the `discover_objects` tool first.", + self.name(), self.name() ) } @@ -354,7 +335,7 @@ pub struct Parser {} std::fs::write(&file_path, source).unwrap(); let input_json = format!( - r#"{{"path":"{}","names":["main","Config","Parser"]}}"#, + r#"{{"path":"{}","query":"main, Config, Parser"}}"#, file_path.display() ); @@ -372,32 +353,16 @@ pub struct Parser {} } #[test] - fn test_parse_input_invalid_kind() { + fn test_parse_input_missing_query() { let temp_dir = tempfile::tempdir().unwrap(); let file_path = temp_dir.path().join("test.rs"); std::fs::write(&file_path, "pub fn main() {}").unwrap(); - let input_json = format!( - r#"{{"path":"{}","names":["main"],"kind":"not_a_kind"}}"#, - file_path.display() - ); + let input_json = format!(r#"{{"path":"{}","query":""}}"#, file_path.display()); let tool = ReadObjects::new(); let err = tool.parse_input(input_json).unwrap(); - assert!(err.to_string().contains("invalid kind")); - } - - #[test] - fn test_parse_input_missing_names() { - let temp_dir = tempfile::tempdir().unwrap(); - let file_path = temp_dir.path().join("test.rs"); - std::fs::write(&file_path, "pub fn main() {}").unwrap(); - - let input_json = format!(r#"{{"path":"{}","names":[]}}"#, file_path.display()); - - let tool = ReadObjects::new(); - let err = tool.parse_input(input_json).unwrap(); - assert!(err.to_string().contains("names is required")); + assert!(err.to_string().contains("query is required")); } #[test] diff --git a/src/domain/tools/shell_exec.rs b/src/domain/tools/shell_exec.rs index f90d912..18f7f9a 100644 --- a/src/domain/tools/shell_exec.rs +++ b/src/domain/tools/shell_exec.rs @@ -163,7 +163,7 @@ impl Tool for ShellExec { fn desc(&self) -> String { format!( r#"Use the `{}` tool to execute shell commands. -Please DO NOT use it to read the full content of a file, this is not efficient, use other tools for this."#, +Please DO NOT use it to read the full content of a file, this is not efficient, use `read_objects` tool for this."#, self.name() ) } @@ -188,7 +188,7 @@ impl ShellExec { /// Display warning and countdown before executing command #[cfg(not(test))] fn warn_and_wait(command: &str, work_dir: &std::path::Path) { - println!("\n\x1b[33m┌─ shell_exec COMMAND ─────────────────────────────\x1b[0m"); + println!("\n\x1b[33m┌─ Action Approval: shell_exec ─────────────────────────────\x1b[0m"); println!("\x1b[33m│\x1b[0m Command: \x1b[1m{}\x1b[0m", command); println!("\x1b[33m│\x1b[0m Workdir: {}", work_dir.display()); println!("\x1b[33m│\x1b[0m"); diff --git a/src/domain/workflow/chain.rs b/src/domain/workflow/chain.rs index 26358f6..f1abb33 100644 --- a/src/domain/workflow/chain.rs +++ b/src/domain/workflow/chain.rs @@ -120,8 +120,11 @@ pub fn parse_tool_choice(llm_output: &str) -> Result<(String, String), Error> { let content = content.trim().strip_suffix("```").unwrap_or(content).trim(); + // Wrap to allow multiple top-level nodes (e.g. + ) + let wrapped = format!("{}", content); + // Parse the XML - let doc = roxmltree::Document::parse(content) + let doc = roxmltree::Document::parse(&wrapped) .map_err(|e| Error::Parse(format!("invalid xml: {}", e)))?; // Extract tool_name diff --git a/src/domain/workflow/mod.rs b/src/domain/workflow/mod.rs index 1190771..945adc8 100644 --- a/src/domain/workflow/mod.rs +++ b/src/domain/workflow/mod.rs @@ -1,12 +1,11 @@ pub mod chain; pub mod step; pub mod toolset; -pub mod toolsets; pub mod workflow; pub use chain::Chain; pub use step::{ChainStep, StepType}; -pub use toolset::{GeneralToolset, Toolset}; +pub use toolset::{AllToolset, EditToolset, ReadToolset, Toolset}; pub use workflow::Workflow; use std::sync::atomic::{AtomicBool, Ordering}; diff --git a/src/domain/workflow/step.rs b/src/domain/workflow/step.rs index 34ad5e1..d840825 100644 --- a/src/domain/workflow/step.rs +++ b/src/domain/workflow/step.rs @@ -9,10 +9,8 @@ pub enum StepType { ToolCall, /// User interrupted the workflow (Ctrl+C) UserInterruption, - /// A todo item was created - TodoCreation, - /// A todo item was updated - TodoUpdate, + /// Signifies that a step in the behavior tree was completed + BehaviorTreeStepPassed } impl StepType { @@ -20,8 +18,7 @@ impl StepType { match self { StepType::ToolCall => "tool_call", StepType::UserInterruption => "user_interruption", - StepType::TodoCreation => "todo_creation", - StepType::TodoUpdate => "todo_update", + StepType::BehaviorTreeStepPassed => "behavior_tree_step_passed", } } } @@ -73,6 +70,12 @@ impl ChainStep { } pub fn get_output(&self, model_type: ModelType) -> String { - prompting::get_tool_result(model_type, self.clone()) + if self.step_type == StepType::ToolCall.as_str() { + return prompting::get_tool_result(model_type, self.clone()); + } + + let mut output = format!("Previous step `{}`: {}", self.step_type, self.input_payload); + output.push_str(&format!("\nStep output: {}", self.summary)); + output } } diff --git a/src/domain/workflow/toolsets/general.rs b/src/domain/workflow/toolset/all.rs similarity index 78% rename from src/domain/workflow/toolsets/general.rs rename to src/domain/workflow/toolset/all.rs index 9ee665c..ea52dd8 100644 --- a/src/domain/workflow/toolsets/general.rs +++ b/src/domain/workflow/toolset/all.rs @@ -3,11 +3,11 @@ use crate::domain::workflow::toolset::Toolset; use std::collections::HashMap; /// General toolset containing read-only and utility tools -pub struct GeneralToolset { +pub struct AllToolset { tools: HashMap>, } -impl GeneralToolset { +impl AllToolset { pub fn new() -> Self { let mut tools: HashMap> = HashMap::new(); @@ -23,8 +23,8 @@ impl GeneralToolset { let structure = Box::new(Structure::new()); tools.insert(structure.name().to_string(), structure); - let patch_file = Box::new(PatchFile::new()); - tools.insert(patch_file.name().to_string(), patch_file); + let patch_files = Box::new(PatchFiles::new()); + tools.insert(patch_files.name().to_string(), patch_files); let shell_exec = Box::new(crate::domain::tools::ShellExec::new()); tools.insert(shell_exec.name().to_string(), shell_exec); @@ -33,13 +33,7 @@ impl GeneralToolset { } } -impl Default for GeneralToolset { - fn default() -> Self { - Self::new() - } -} - -impl Toolset for GeneralToolset { +impl Toolset for AllToolset { fn tools(&self) -> &HashMap> { &self.tools } diff --git a/src/domain/workflow/toolset/edit.rs b/src/domain/workflow/toolset/edit.rs new file mode 100644 index 0000000..7c200cd --- /dev/null +++ b/src/domain/workflow/toolset/edit.rs @@ -0,0 +1,25 @@ +use crate::domain::tools::*; +use crate::domain::workflow::toolset::Toolset; +use std::collections::HashMap; + +/// General toolset containing read-only and utility tools +pub struct EditToolset { + tools: HashMap>, +} + +impl EditToolset { + pub fn new() -> Self { + let mut tools: HashMap> = HashMap::new(); + + let patch_files = Box::new(PatchFiles::new()); + tools.insert(patch_files.name().to_string(), patch_files); + + Self { tools } + } +} + +impl Toolset for EditToolset { + fn tools(&self) -> &HashMap> { + &self.tools + } +} diff --git a/src/domain/workflow/toolset.rs b/src/domain/workflow/toolset/mod.rs similarity index 70% rename from src/domain/workflow/toolset.rs rename to src/domain/workflow/toolset/mod.rs index f62dac0..194e624 100644 --- a/src/domain/workflow/toolset.rs +++ b/src/domain/workflow/toolset/mod.rs @@ -1,11 +1,37 @@ +mod all; +mod edit; +mod read; +mod none; + +pub use all::AllToolset; +pub use edit::EditToolset; +pub use read::ReadToolset; +pub use none::NoneToolset; +#[derive(Clone, Copy)] +pub enum ToolsetType { + None, + Read, + Edit, + All, +} + +impl ToolsetType { + pub fn build(self) -> Box { + match self { + ToolsetType::Read => Box::new(ReadToolset::new()), + ToolsetType::Edit => Box::new(EditToolset::new()), + ToolsetType::All => Box::new(AllToolset::new()), + ToolsetType::None => Box::new(NoneToolset::new()), + } + } +} + use super::Error; use crate::domain::prompting::format_tools_description; use crate::domain::session::Request; use crate::domain::tools::Tool; use std::collections::HashMap; -pub use super::toolsets::GeneralToolset; - /// Trait for toolset implementations that provide a set of tools /// /// Implementations should create their tools in the `new()` constructor @@ -43,16 +69,16 @@ mod tests { #[test] fn test_general_toolset() { - let toolset = GeneralToolset::new(); + let toolset = AllToolset::new(); assert!(toolset.get_tool("read_objects").is_some()); assert!(toolset.get_tool("find_files").is_some()); - assert!(toolset.get_tool("patch_file").is_some()); + assert!(toolset.get_tool("patch_files").is_some()); assert!(toolset.get_tool("nonexistent").is_none()); } #[test] fn test_tools_description() { - let toolset = GeneralToolset::new(); + let toolset = AllToolset::new(); let description = toolset.get_tools_description(); assert!(description.contains("Available Tools:")); assert!(description.contains("read_objects")); diff --git a/src/domain/workflow/toolset/none.rs b/src/domain/workflow/toolset/none.rs new file mode 100644 index 0000000..362698b --- /dev/null +++ b/src/domain/workflow/toolset/none.rs @@ -0,0 +1,21 @@ +use crate::domain::tools::*; +use crate::domain::workflow::toolset::Toolset; +use std::collections::HashMap; + +pub struct NoneToolset { + tools: HashMap>, +} + +impl NoneToolset { + pub fn new() -> Self { + let mut tools: HashMap> = HashMap::new(); + + Self { tools } + } +} + +impl Toolset for NoneToolset { + fn tools(&self) -> &HashMap> { + &self.tools + } +} diff --git a/src/domain/workflow/toolset/read.rs b/src/domain/workflow/toolset/read.rs new file mode 100644 index 0000000..eed3178 --- /dev/null +++ b/src/domain/workflow/toolset/read.rs @@ -0,0 +1,34 @@ +use crate::domain::tools::*; +use crate::domain::workflow::toolset::Toolset; +use std::collections::HashMap; + +/// General toolset containing read-only and utility tools +pub struct ReadToolset { + tools: HashMap>, +} + +impl ReadToolset { + pub fn new() -> Self { + let mut tools: HashMap> = HashMap::new(); + + let discover_objects = Box::new(DiscoverObjects::new()); + tools.insert(discover_objects.name().to_string(), discover_objects); + + let read_objects = Box::new(ReadObjects::new()); + tools.insert(read_objects.name().to_string(), read_objects); + + let find_files = Box::new(FindFiles::new()); + tools.insert(find_files.name().to_string(), find_files); + + let structure = Box::new(Structure::new()); + tools.insert(structure.name().to_string(), structure); + + Self { tools } + } +} + +impl Toolset for ReadToolset { + fn tools(&self) -> &HashMap> { + &self.tools + } +} diff --git a/src/domain/workflow/toolsets/mod.rs b/src/domain/workflow/toolsets/mod.rs deleted file mode 100644 index c0b359f..0000000 --- a/src/domain/workflow/toolsets/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -mod general; - -pub use general::GeneralToolset; diff --git a/src/domain/workflow/workflow.rs b/src/domain/workflow/workflow.rs index 3f10a6c..7476264 100644 --- a/src/domain/workflow/workflow.rs +++ b/src/domain/workflow/workflow.rs @@ -1,8 +1,9 @@ -use super::{chain::Chain, toolset::Toolset, CancellationToken, Error}; +use super::{chain::Chain, step::ChainStep, step::StepType, toolset::Toolset, CancellationToken, Error}; +use crate::domain::bt::GeneralTree; use crate::domain::prompting; use crate::domain::session::Request; use crate::domain::tools::Tool; -use crate::domain::workflow::GeneralToolset; +use crate::domain::workflow::AllToolset; use crate::infrastructure::inference::{local::LocalEngine, InferenceEngine}; use std::sync::Arc; @@ -15,6 +16,7 @@ use std::sync::Arc; pub struct Workflow { toolset: Box, engine: Arc, + chain: Chain, } impl Workflow { @@ -22,11 +24,20 @@ impl Workflow { /// The engine is created internally by scanning and selecting a model pub fn new(engine: Arc) -> Result { Ok(Self { - toolset: Box::new(GeneralToolset::new()), + toolset: Box::new(AllToolset::new()), engine, + chain: Chain::new(), }) } + pub fn reset_chain(&mut self) { + self.chain = Chain::new(); + } + + pub fn chain(&self) -> &Chain { + &self.chain + } + /// Run the workflow for a given request /// Implements the eternal agent loop: /// - Asks LLM to choose next tool @@ -37,26 +48,27 @@ impl Workflow { /// The workflow checks the cancellation token before each iteration /// and returns Error::Cancelled if cancellation was requested. pub fn run( - &self, + &mut self, request: &mut dyn Request, cancel: &CancellationToken, - ) -> Result { - const MAX_ITERATIONS: usize = 128; - - // Initialize the chain to track executed steps - let mut chain = Chain::new(); - + max_tool_calls: usize, + user_prompt_override: Option, + ) -> Result<(), Error> { // Eternal agent loop - for _iteration in 1..=MAX_ITERATIONS { + for _iteration in 1..=max_tool_calls { // Check for cancellation at the start of each iteration if cancel.is_cancelled() { log::info!("Workflow cancelled by user"); - chain.add_interruption(); - return Ok(chain); + self.chain.add_interruption(); + return Ok(()); } let system_prompt = prompting::get_system_prompt(self.engine.get_type()); - let user_prompt = prompting::get_user_prompt(self.engine.get_type(), request); + let base_user_prompt = prompting::get_user_prompt(self.engine.get_type(), request); + let user_prompt = user_prompt_override + .as_deref() + .unwrap_or(&base_user_prompt) + .to_string(); // Ask LLM to choose next tool let llm_output = self @@ -66,40 +78,74 @@ impl Workflow { &user_prompt, 1024, &self.toolset.tool_refs(), - &chain, + &self.chain, ) .map_err(|e| Error::Inference(e.to_string()))?; // Check for cancellation after inference if cancel.is_cancelled() { log::info!("Workflow cancelled by user"); - chain.add_interruption(); - return Ok(chain); + self.chain.add_interruption(); + return Ok(()); } if llm_output.chosen_tool.is_none() { let final_message = llm_output.summary; - chain.set_final_message(final_message.clone()); + self.chain.set_final_message(final_message.clone()); request.set_final_message(final_message); - return Ok(chain); + return Ok(()); } let tool_result = llm_output.chosen_tool.unwrap().work(request); // Add step to chain - chain.add_step(tool_result); + self.chain.add_step(tool_result); } - if chain.len() >= MAX_ITERATIONS { - log::warn!("Workflow reached maximum iterations ({})", MAX_ITERATIONS); - chain.mark_failed("Maximum iterations reached".to_string()); + if max_tool_calls > 0 { + log::warn!("Workflow reached maximum iterations ({})", max_tool_calls); } - Ok(chain) + Ok(()) } - /// Get the toolset (for testing or inspection) - pub fn toolset(&self) -> &dyn Toolset { - self.toolset.as_ref() + pub fn run_using_bt( + &mut self, + request: &mut dyn Request, + cancel: &CancellationToken, + ) -> Result { + self.reset_chain(); + let tree = GeneralTree::new(); + let mut current_step = Some(tree.head()); + + while let Some(step) = current_step { + let step_user_prompt = + prompting::get_bt_tree_step_prompt(self.engine.get_type(), step, request); + + self.toolset = step.toolset().build(); + + self.run( + request, + cancel, + step.max_tools_calls() as usize, + Some(step_user_prompt.clone()), + )?; + self.chain.steps.push(ChainStep { + step_type: StepType::BehaviorTreeStepPassed.as_str().to_string(), + summary: self.chain.final_message().unwrap_or("").to_string(), + context_payload: String::new(), + input_payload: step_user_prompt.to_string(), + tool_name: None, + tool_output: None, + is_successful: Some(true), + }); + if self.chain.is_failed { + return Ok(self.chain.clone()); + } + + current_step = step.next_step(); + } + + Ok(self.chain.clone()) } } diff --git a/src/infrastructure/api_clients/openai/client.rs b/src/infrastructure/api_clients/openai/client.rs index f4cec3a..6776989 100644 --- a/src/infrastructure/api_clients/openai/client.rs +++ b/src/infrastructure/api_clients/openai/client.rs @@ -21,14 +21,18 @@ pub enum OpenAIClientError { impl std::fmt::Display for OpenAIClientError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - OpenAIClientError::Api { status, .. } => { - write!(f, "OpenAI API error (status={})", status) + OpenAIClientError::Api { status, body } => { + write!(f, "OpenAI API error (status={}, body={})", status, body) } - OpenAIClientError::Deserialize { source, .. } => { - write!(f, "Failed to deserialize OpenAI response: {}", source) + OpenAIClientError::Deserialize { source, body } => { + write!( + f, + "Failed to deserialize OpenAI response: {} (body={})", + source, body + ) } - OpenAIClientError::NoText { .. } => { - write!(f, "No output_text found in OpenAI response") + OpenAIClientError::NoText { body } => { + write!(f, "No output_text found in OpenAI response (body={})", body) } } } @@ -65,24 +69,45 @@ impl OpenAIClient { tools: &[&dyn crate::domain::tools::Tool], chain: &crate::domain::workflow::Chain, ) -> Result> { - // Try once, retry on transient errors - match self.call_responses_api_inner(system_prompt, user_prompt, tools, chain) { - Ok(result) => Ok(result), - Err(e) => { - // Check if it's a transient error worth retrying - let error_str = e.to_string(); - if error_str.contains("error sending request") - || error_str.contains("connection") - || error_str.contains("timeout") - { - log::warn!("OpenAI API request failed, retrying once: {}", error_str); - std::thread::sleep(std::time::Duration::from_secs(2)); - self.call_responses_api_inner(system_prompt, user_prompt, tools, chain) - } else { - Err(e) + let max_attempts = 3; + for attempt in 1..=max_attempts { + match self.call_responses_api_inner(system_prompt, user_prompt, tools, chain) { + Ok(result) => return Ok(result), + Err(e) => { + let error_str = e.to_string(); + let mut should_retry = error_str.contains("error sending request") + || error_str.contains("connection") + || error_str.contains("timeout"); + + if let Some(OpenAIClientError::Api { status, .. }) = + e.downcast_ref::() + { + let status_code = status.as_u16(); + should_retry = should_retry + || status.is_server_error() + || status_code == 429 + || status_code == 408 + || status_code == 409; + } + + if should_retry && attempt < max_attempts { + let backoff_secs = 2u64.saturating_mul(attempt as u64); + log::warn!( + "OpenAI API request failed, retrying (attempt {}/{}): {}", + attempt, + max_attempts, + error_str + ); + std::thread::sleep(std::time::Duration::from_secs(backoff_secs)); + continue; + } + + return Err(e); } } } + + unreachable!("retry loop exits via return"); } fn call_responses_api_inner( diff --git a/src/infrastructure/api_clients/openai/dto/request.rs b/src/infrastructure/api_clients/openai/dto/request.rs index e908078..d6451f4 100644 --- a/src/infrastructure/api_clients/openai/dto/request.rs +++ b/src/infrastructure/api_clients/openai/dto/request.rs @@ -1,7 +1,7 @@ use crate::domain::tools::Tool; use serde::Serialize; use serde_json::Value; -use crate::domain::ModelType; +use crate::domain::{Chain, ModelType}; #[derive(Debug, Serialize)] pub struct RequestDTO { @@ -93,21 +93,17 @@ impl InputDto { } } - fn from_chain(chain: &crate::domain::workflow::Chain) -> Vec { + fn from_chain(chain: &Chain) -> Vec { chain .steps .iter() - .filter(|step| { - step.step_type == crate::domain::workflow::step::StepType::ToolCall.as_str() - }) - .filter_map(|step| { - let tool_name = step.tool_name.as_ref()?; + .map(|step| { let status = if step.is_successful.unwrap_or(false) { "completed" } else { "failed" }; - Some(Self { + Self { content: vec![InputContent { kind: "input_text".to_string(), text: step.get_output(ModelType::OpenAI), @@ -115,7 +111,7 @@ impl InputDto { role: ROLE_SYSTEM.to_string(), kind: "message".to_string(), status: status.to_string(), - }) + } }) .collect() } diff --git a/src/lib.rs b/src/lib.rs index fc12cfc..bbc501f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,8 +2,10 @@ pub mod domain; pub mod infrastructure; pub mod repository; pub mod utils; +pub mod config; pub use domain::*; pub use infrastructure::*; pub use repository::*; pub use utils::*; +pub use config::*; diff --git a/src/main.rs b/src/main.rs index c9dd83b..d4761dc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,8 +2,10 @@ mod domain; mod infrastructure; mod repository; mod utils; +mod config; use clap::Parser; +use config::AppConfig; use domain::{CancellationToken, Session, SessionService, StartupService}; use infrastructure::{ change_model, get_current_model_name, InfrastructureComponents, InfrastructureInitializer, @@ -12,7 +14,7 @@ use rustyline::Cmd; use rustyline::Editor; use rustyline::KeyEvent; use std::env; -use std::sync::Arc; +use colored::Colorize; use utils::{handle_readline_error, Action, StatusBarHelper}; #[derive(Parser)] @@ -27,19 +29,21 @@ struct Args { fn main() { let args = Args::parse(); - if let Err(e) = run(args.debug) { + let config = AppConfig::new(args.debug); + + if let Err(e) = run(config) { eprintln!("Error: {}", e); std::process::exit(1); } } -fn run(debug: bool) -> Result<(), String> { +fn run(config: AppConfig) -> Result<(), String> { // Initialize logging { let mut builder = env_logger::Builder::from_default_env(); // Our crate logs: respect debug flag - if debug { + if config.debug { builder.filter_level(log::LevelFilter::Debug); } else { builder.filter_level(log::LevelFilter::Info); @@ -52,16 +56,20 @@ fn run(debug: bool) -> Result<(), String> { } // Initialize infrastructure components - let infrastructure_initializer = InfrastructureInitializer::with_debug(debug); + let infrastructure_initializer = InfrastructureInitializer::with_debug(config.debug); let infra = infrastructure_initializer.initialize()?; println!("Application initialized. Type :q or press Ctrl+Q to exit.\n"); // Start the REPL with infrastructure components - repl(infra, debug) + repl(infra, config.debug, config.use_behavior_trees) } -fn repl(infra: InfrastructureComponents, debug: bool) -> Result<(), String> { +fn repl( + infra: InfrastructureComponents, + debug: bool, + use_behavior_trees: bool, +) -> Result<(), String> { // Get current model name for status bar let current_model_name = get_current_model_name(&infra.connection).unwrap_or_else(|_| "unknown".to_string()); @@ -74,7 +82,11 @@ fn repl(infra: InfrastructureComponents, debug: bool) -> Result<(), String> { // Bind Ctrl+Q to EOF (quit) rl.bind_sequence(KeyEvent::ctrl('q'), Cmd::EndOfFile); - let session_service = SessionService::new(infra.engine.clone(), infra.connection.clone()) + let mut session_service = SessionService::new( + infra.engine.clone(), + infra.connection.clone(), + use_behavior_trees, + ) .map_err(|e| format!("Failed to create session service: {}", e))?; let startup_service = StartupService::with_debug(debug, infra.engine.clone(), infra.connection.clone()); @@ -152,8 +164,12 @@ fn repl(infra: InfrastructureComponents, debug: bool) -> Result<(), String> { // Run the session service (creates request, runs workflow, updates result) match session_service.run(session, trimmed, &cancel_token) { Ok(chain) => { - println!("Ok> {}", chain.get_summary()); - println!("{}", chain.final_message().unwrap().to_string()); + println!("{} {}", "".green(), chain.get_summary()); + if let Some(message) = chain.final_message() { + println!("{}", message); + } else if chain.is_failed { + println!("{}", chain.fail_reason); + } } Err(domain::session::ServiceError::Workflow( domain::workflow::Error::Cancelled,