Skip to main content

tiny_agent_rs/services/
planning.rs

1use crate::tools::FunctionFactory;
2
3/// Generate a planning prompt for the agent before tool execution
4pub fn generate_planning_prompt(
5    task: &str,
6    available_tools: &[String],
7    iteration: usize,
8) -> String {
9    let tools_list = if available_tools.is_empty() {
10        "No tools available".to_string()
11    } else {
12        available_tools.join(", ")
13    };
14
15    if iteration == 1 {
16        format!(
17            "Task: {}\n\nAvailable tools: {}\n\nThink step-by-step about how to solve this task. What tools will you need? What's your approach?",
18            task, tools_list
19        )
20    } else {
21        format!(
22            "Continuing task: {}\n\nIteration {}\nAvailable tools: {}\n\nBased on previous observations, what should be the next step?",
23            task, iteration, tools_list
24        )
25    }
26}
27
28/// Extract tool names from FunctionFactory for planning context
29pub fn get_tool_names(factory: &FunctionFactory) -> Vec<String> {
30    factory
31        .get_openai_tools()
32        .iter()
33        .filter_map(|tool| {
34            tool.get("function")
35                .and_then(|f| f.get("name"))
36                .and_then(|n| n.as_str())
37                .map(|s| s.to_string())
38        })
39        .collect()
40}
41
42/// Generate a simplified planning prompt that encourages direct tool use
43pub fn generate_tool_planning_prompt(task: &str, factory: &FunctionFactory) -> String {
44    let tool_descriptions: Vec<String> = factory
45        .get_openai_tools()
46        .iter()
47        .filter_map(|tool| {
48            let function = tool.get("function")?;
49            let name = function.get("name")?.as_str()?;
50            let description = function.get("description")?.as_str()?;
51            Some(format!("- {}: {}", name, description))
52        })
53        .collect();
54
55    if tool_descriptions.is_empty() {
56        format!("Task: {}\n\nHow would you approach solving this?", task)
57    } else {
58        format!(
59            "Task: {}\n\nAvailable tools:\n{}\n\nWhat's your plan to solve this task?",
60            task,
61            tool_descriptions.join("\n")
62        )
63    }
64}
65
66/// Check if a response indicates the agent is planning vs. executing
67pub fn is_planning_response(content: &str) -> bool {
68    let planning_indicators = [
69        "plan", "approach", "strategy", "steps:", "first,", "then,", "finally,", "need to",
70        "should", "will use",
71    ];
72
73    let content_lower = content.to_lowercase();
74    planning_indicators
75        .iter()
76        .any(|indicator| content_lower.contains(indicator))
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn test_generate_planning_prompt() {
85        let tools = vec!["search".to_string(), "calculator".to_string()];
86        let prompt = generate_planning_prompt("Calculate 2+2", &tools, 1);
87
88        assert!(prompt.contains("Calculate 2+2"));
89        assert!(prompt.contains("search, calculator"));
90        assert!(prompt.contains("Think step-by-step"));
91    }
92
93    #[test]
94    fn test_generate_planning_prompt_continuation() {
95        let tools = vec!["search".to_string()];
96        let prompt = generate_planning_prompt("Find info", &tools, 3);
97
98        assert!(prompt.contains("Iteration 3"));
99        assert!(prompt.contains("next step"));
100    }
101
102    #[test]
103    fn test_empty_tools() {
104        let tools: Vec<String> = vec![];
105        let prompt = generate_planning_prompt("Task", &tools, 1);
106
107        assert!(prompt.contains("No tools available"));
108    }
109
110    #[test]
111    fn test_is_planning_response() {
112        assert!(is_planning_response("My plan is to use the search tool"));
113        assert!(is_planning_response(
114            "First, I will search. Then, I will analyze."
115        ));
116        assert!(is_planning_response("I need to calculate the result"));
117        assert!(!is_planning_response("The result is 42"));
118        assert!(!is_planning_response("Done"));
119    }
120
121    #[test]
122    fn test_tool_planning_prompt_empty() {
123        let factory = FunctionFactory::new();
124        let prompt = generate_tool_planning_prompt("Test task", &factory);
125
126        assert!(prompt.contains("Test task"));
127        assert!(prompt.contains("How would you approach"));
128    }
129}