save temp progress

This commit is contained in:
taggart_comet 2026-01-19 09:33:05 +02:00
parent 15de22fac7
commit bb80abbd41
29 changed files with 523 additions and 232 deletions

10
Cargo.lock generated
View file

@ -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",

View file

@ -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]

View file

@ -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!")

80
src/domain/bt/mod.rs Normal file
View file

@ -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<Box<dyn BTStepNodeInterface>>,
decision: Option<Box<dyn BTDecisionNodeInterface>>,
}
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()
}
}

View file

@ -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;

View file

@ -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
)

View file

@ -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
)
}
}
}

View file

@ -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
)
}

View file

@ -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<Connection>,
}
impl SessionService {
/// Create a new session service with default workflow
pub fn new(engine: Arc<dyn InferenceEngine>, conn: Arc<Connection>) -> Result<Self, String> {
pub fn new(
engine: Arc<dyn InferenceEngine>,
conn: Arc<Connection>,
use_behavior_trees: bool,
) -> Result<Self, String> {
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<Chain, WorkflowError> = 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();

View file

@ -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()
)
}

View file

@ -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]

View file

@ -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 = "<input></input>".to_string();
/// let _result = ToolResult::ok(tool_name, input, xml);
/// ```
pub fn serialize_output<T>(value: &T) -> Result<String, Error>
where
@ -139,7 +144,7 @@ pub fn build_tool_by_name(name: &str) -> Option<Box<dyn Tool>> {
"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,
}

View file

@ -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<Option<PatchFileInput>>,
pub struct PatchFiles {
input: Mutex<Option<PatchFilesInput>>,
}
#[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<PatchFileInput, Error> {
let parsed: PatchFileInputJson =
fn parse_input_json(raw: &str) -> Result<PatchFilesInput, Error> {
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<PatchFileInput, Error> {
fn load_input(&self) -> Result<PatchFilesInput, Error> {
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<Error> {
@ -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<String> = 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);

View file

@ -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<String>,
kind: Option<String>,
query: String,
}
#[derive(Debug, Serialize)]
@ -65,18 +64,6 @@ impl ObjectQuery {
}
}
#[derive(Serialize)]
struct ReadObjectsOutput {
language: String,
object_not_found: bool,
results: HashMap<String, ObjectContent>,
}
#[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<String> {
query
.replace(',', " ")
.split_whitespace()
.map(|part| part.trim().to_string())
.collect()
}
fn format_output(lang: Lang, results: HashMap<String, ObjectContent>) -> 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]

View file

@ -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");

View file

@ -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. <tool_name> + <input>)
let wrapped = format!("<root>{}</root>", 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

View file

@ -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};

View file

@ -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
}
}

View file

@ -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<String, Box<dyn Tool>>,
}
impl GeneralToolset {
impl AllToolset {
pub fn new() -> Self {
let mut tools: HashMap<String, Box<dyn Tool>> = 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<String, Box<dyn Tool>> {
&self.tools
}

View file

@ -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<String, Box<dyn Tool>>,
}
impl EditToolset {
pub fn new() -> Self {
let mut tools: HashMap<String, Box<dyn Tool>> = 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<String, Box<dyn Tool>> {
&self.tools
}
}

View file

@ -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<dyn Toolset> {
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"));

View file

@ -0,0 +1,21 @@
use crate::domain::tools::*;
use crate::domain::workflow::toolset::Toolset;
use std::collections::HashMap;
pub struct NoneToolset {
tools: HashMap<String, Box<dyn Tool>>,
}
impl NoneToolset {
pub fn new() -> Self {
let mut tools: HashMap<String, Box<dyn Tool>> = HashMap::new();
Self { tools }
}
}
impl Toolset for NoneToolset {
fn tools(&self) -> &HashMap<String, Box<dyn Tool>> {
&self.tools
}
}

View file

@ -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<String, Box<dyn Tool>>,
}
impl ReadToolset {
pub fn new() -> Self {
let mut tools: HashMap<String, Box<dyn Tool>> = 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<String, Box<dyn Tool>> {
&self.tools
}
}

View file

@ -1,3 +0,0 @@
mod general;
pub use general::GeneralToolset;

View file

@ -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<dyn Toolset>,
engine: Arc<dyn InferenceEngine>,
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<dyn InferenceEngine>) -> Result<Self, String> {
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<Chain, Error> {
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<String>,
) -> 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<Chain, Error> {
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())
}
}

View file

@ -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<LLMInferenceResult, Box<dyn Error + Send + Sync>> {
// 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::<OpenAIClientError>()
{
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(

View file

@ -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<Self> {
fn from_chain(chain: &Chain) -> Vec<Self> {
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()
}

View file

@ -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::*;

View file

@ -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!("{} {}", "<Ok>".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,