Skip to main content

lc_agents/plan_execute/
planner.rs

1//! Planner - generates / replans execution plans with the LLM
2
3use lc_core::language_models::BaseChatModel;
4use lc_core::tools::ToolDefinition;
5use lc_providers::ProviderError;
6use lc_schema::Message;
7use serde_json::{json, Value};
8
9use super::plan::Plan;
10use crate::AgentError;
11
12use std::sync::Arc;
13
14/// JSON Schema for the planning tool: forces the LLM to emit a structured steps array (P1-3).
15fn plan_tool() -> ToolDefinition {
16    ToolDefinition::new(
17        "generate_plan",
18        "为给定目标生成执行计划,返回按顺序执行的步骤描述数组",
19    )
20    .with_parameters(json!({
21        "type": "object",
22        "properties": {
23            "steps": {
24                "type": "array",
25                "items": { "type": "string" },
26                "description": "按顺序执行的步骤描述"
27            }
28        },
29        "required": ["steps"]
30    }))
31}
32
33/// Extracts the steps array from the tool_call args, serializing it back to `["a", "b"]` for parse_plan to reuse.
34fn steps_to_json_string(args: &Value) -> String {
35    args.get("steps")
36        .and_then(|v| serde_json::to_string(v).ok())
37        .unwrap_or_default()
38}
39
40/// Planner: calls the LLM to generate a step list
41pub struct Planner {
42    llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
43}
44
45impl Planner {
46    /// Creates a new planner.
47    pub fn new(llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>) -> Self {
48        Self { llm }
49    }
50
51    /// Generates an execution plan
52    pub async fn plan(&self, objective: &str) -> Result<Plan, AgentError> {
53        let prompt = format!(
54            "为以下目标制定执行计划,输出 JSON 字符串数组,每项是一个步骤描述。\n\
55             目标: {}\n\
56             输出格式: [\"步骤1\", \"步骤2\", ...]\n\
57             只输出 JSON,不要任何其他内容。",
58            objective
59        );
60        let messages = vec![
61            Message::system("你是规划助手,只输出 JSON。"),
62            Message::human(prompt),
63        ];
64        let structured = crate::structured::chat_structured(
65            self.llm.as_ref(),
66            Some(plan_tool()),
67            messages,
68            None,
69            &crate::retry::RetryConfig::default(),
70        )
71        .await
72        .map_err(|e| AgentError::Other(format!("LLM error: {:?}", e)))?;
73        let content = match &structured.tool_args {
74            Some(args) => steps_to_json_string(args),
75            None => structured.content,
76        };
77        self.parse_plan(objective, &content)
78    }
79
80    /// Replans (when a step fails).
81    ///
82    /// 0.22.0 C4 fix: `completed` carries the already-finished steps and
83    /// their results; the prompt injects them and explicitly asks for a plan
84    /// of the *remaining* work, so a replan no longer re-executes everything
85    /// (duplicate tool calls / duplicate cost) and the caller can splice the
86    /// completed steps back into the returned plan.
87    pub async fn replan(
88        &self,
89        objective: &str,
90        failed_step: &str,
91        reason: &str,
92        completed: &str,
93    ) -> Result<Plan, AgentError> {
94        let completed_block = if completed.trim().is_empty() {
95            "(none)".to_string()
96        } else {
97            completed.to_string()
98        };
99        let prompt = format!(
100            "原目标: {}\n已完成步骤及结果:\n{}\n失败步骤 '{}' 失败: {}\n请重新制定**剩余工作**的完整计划,不要重复已完成的步骤。输出 JSON 字符串数组 [\"步骤\", ...],只输出 JSON。",
101            objective, completed_block, failed_step, reason
102        );
103        let messages = vec![
104            Message::system("你是规划助手,只输出 JSON。"),
105            Message::human(prompt),
106        ];
107        let structured = crate::structured::chat_structured(
108            self.llm.as_ref(),
109            Some(plan_tool()),
110            messages,
111            None,
112            &crate::retry::RetryConfig::default(),
113        )
114        .await
115        .map_err(|e| AgentError::Other(format!("LLM error: {:?}", e)))?;
116        let content = match &structured.tool_args {
117            Some(args) => steps_to_json_string(args),
118            None => structured.content,
119        };
120        self.parse_plan(objective, &content)
121    }
122
123    fn parse_plan(&self, objective: &str, content: &str) -> Result<Plan, AgentError> {
124        let json_str = extract_json_array(content);
125        let descs: Vec<String> = serde_json::from_str(&json_str).map_err(|e| {
126            AgentError::OutputParsingError(format!(
127                "failed to parse plan: {} | raw: {}",
128                e, content
129            ))
130        })?;
131        Ok(Plan::from_descriptions(objective, descs))
132    }
133}
134
135/// Extracts a JSON array from LLM output (tolerates markdown code fences)
136fn extract_json_array(content: &str) -> String {
137    let trimmed = content.trim();
138    // Strip markdown ```json ... ```
139    let stripped = if trimmed.starts_with("```") {
140        trimmed
141            .strip_prefix("```json")
142            .or_else(|| trimmed.strip_prefix("```"))
143            .unwrap_or(trimmed)
144            .strip_suffix("```")
145            .unwrap_or(trimmed)
146            .trim()
147    } else {
148        trimmed
149    };
150    // Take from the first [ to the last ]
151    if let Some(start) = stripped.find('[') {
152        if let Some(end) = stripped.rfind(']') {
153            if end > start {
154                return stripped[start..=end].to_string();
155            }
156        }
157    }
158    stripped.to_string()
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    #[test]
166    fn test_extract_plain_json() {
167        let s = r#"["步骤1", "步骤2"]"#;
168        assert_eq!(extract_json_array(s), s);
169    }
170
171    #[test]
172    fn test_extract_markdown_json() {
173        let s = "```json\n[\"a\", \"b\"]\n```";
174        assert_eq!(extract_json_array(s), r#"["a", "b"]"#);
175    }
176
177    #[test]
178    fn test_extract_json_with_surrounding_text() {
179        let s = r#"结果如下: ["x", "y"] 完成"#;
180        assert_eq!(extract_json_array(s), r#"["x", "y"]"#);
181    }
182
183    #[test]
184    fn test_parse_plan() {
185        // No LLM needed: test the parse logic directly (through extract)
186        let content = r#"["搜索资料", "总结"]"#;
187        let json = extract_json_array(content);
188        let descs: Vec<String> = serde_json::from_str(&json).unwrap();
189        assert_eq!(descs, vec!["搜索资料", "总结"]);
190    }
191
192    #[test]
193    fn test_steps_to_json_string() {
194        // P1-3: the tool_call's steps array serializes back to a JSON array that parse_plan can consume.
195        let args = serde_json::json!({"steps": ["a", "b"]});
196        assert_eq!(steps_to_json_string(&args), r#"["a","b"]"#);
197    }
198
199    #[test]
200    fn test_steps_to_json_string_missing_steps() {
201        let args = serde_json::json!({"other": 1});
202        assert_eq!(steps_to_json_string(&args), "");
203    }
204
205    #[test]
206    fn test_plan_tool_schema() {
207        let tool = plan_tool();
208        assert_eq!(tool.function.name, "generate_plan");
209        assert!(tool.function.parameters.is_some());
210    }
211}