use crate::llm::structured::{generate_blocking, StructuredMode, StructuredRequest};
use crate::llm::{LlmClient, Message};
use crate::planning::{AgentGoal, Complexity, ExecutionPlan, Task};
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AchievementResult {
pub achieved: bool,
pub progress: f32,
pub remaining_criteria: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct PreAnalysis {
pub intent: crate::prompts::AgentStyle,
pub requires_planning: bool,
pub goal: AgentGoal,
pub execution_plan: ExecutionPlan,
pub optimized_input: String,
}
pub struct LlmPlanner;
#[derive(Debug, Deserialize)]
struct PlanResponse {
goal: String,
complexity: String,
steps: Vec<StepResponse>,
#[serde(default)]
required_tools: Vec<String>,
}
#[derive(Debug, Deserialize)]
struct StepResponse {
id: String,
description: String,
#[serde(default)]
tool: Option<String>,
#[serde(default)]
dependencies: Vec<String>,
#[serde(default)]
success_criteria: Option<String>,
}
#[derive(Debug, Deserialize)]
struct GoalResponse {
description: String,
success_criteria: Vec<String>,
}
#[derive(Debug, Deserialize)]
struct AchievementResponse {
achieved: bool,
progress: f32,
#[serde(default)]
remaining_criteria: Vec<String>,
}
#[derive(Debug, Deserialize)]
struct PreAnalysisResponse {
intent: String,
requires_planning: bool,
goal: GoalResponse,
execution_plan: PreAnalysisPlan,
optimized_input: String,
}
#[derive(Debug, Deserialize)]
struct PreAnalysisPlan {
complexity: String,
steps: Vec<StepResponse>,
#[serde(default)]
required_tools: Vec<String>,
}
impl LlmPlanner {
pub async fn create_plan(llm: &Arc<dyn LlmClient>, prompt: &str) -> Result<ExecutionPlan> {
let system = crate::prompts::LLM_PLAN_SYSTEM;
let messages = vec![Message::user(prompt)];
let response = llm
.complete(&messages, Some(system), &[])
.await
.context("LLM call failed during plan creation")?;
let text = response.text();
Self::parse_plan_response(&text)
}
pub async fn extract_goal(llm: &Arc<dyn LlmClient>, prompt: &str) -> Result<AgentGoal> {
let system = crate::prompts::LLM_GOAL_EXTRACT_SYSTEM;
let messages = vec![Message::user(prompt)];
let response = llm
.complete(&messages, Some(system), &[])
.await
.context("LLM call failed during goal extraction")?;
let text = response.text();
Self::parse_goal_response(&text)
}
pub async fn check_achievement(
llm: &Arc<dyn LlmClient>,
goal: &AgentGoal,
current_state: &str,
) -> Result<AchievementResult> {
let prompt = format!(
"Goal: {}\nSuccess Criteria: {}\nCurrent State: {}",
goal.description,
goal.success_criteria.join("; "),
current_state,
);
let req = StructuredRequest {
prompt,
system: Some(crate::prompts::LLM_GOAL_CHECK_SYSTEM.to_string()),
schema: Self::achievement_schema(),
schema_name: "goal_achievement".to_string(),
schema_description: Some(
"Strict, evidence-backed evaluation of whether every goal criterion is met"
.to_string(),
),
mode: StructuredMode::Auto,
max_repair_attempts: 2,
};
let result = generate_blocking(&**llm, &req)
.await
.context("LLM achievement structured generation failed")?;
Self::achievement_from_value(result.object)
}
pub fn fallback_plan(prompt: &str) -> ExecutionPlan {
let complexity = if prompt.len() < 50 {
Complexity::Simple
} else if prompt.len() < 150 {
Complexity::Medium
} else if prompt.len() < 300 {
Complexity::Complex
} else {
Complexity::VeryComplex
};
let mut plan = ExecutionPlan::new(prompt, complexity);
let step_count = match complexity {
Complexity::Simple => 2,
Complexity::Medium => 4,
Complexity::Complex => 7,
Complexity::VeryComplex => 10,
};
for i in 0..step_count {
let step = Task::new(
format!("step-{}", i + 1),
crate::prompts::render(
crate::prompts::PLAN_FALLBACK_STEP,
&[("step_num", &(i + 1).to_string())],
),
);
plan.add_step(step);
}
plan
}
pub fn fallback_goal(prompt: &str) -> AgentGoal {
AgentGoal::new(prompt).with_criteria(vec![
"Task is completed successfully".to_string(),
"All requirements are met".to_string(),
])
}
pub fn fallback_check_achievement(goal: &AgentGoal, current_state: &str) -> AchievementResult {
let _ = current_state;
AchievementResult {
achieved: false,
progress: goal.progress,
remaining_criteria: goal.success_criteria.clone(),
}
}
pub async fn pre_analyze(llm: &Arc<dyn LlmClient>, prompt: &str) -> Result<PreAnalysis> {
let req = StructuredRequest {
prompt: format!(
"Analyze this user request and return a compact pre-analysis object. \
Use at most 5 execution steps.\n\nUser request:\n{prompt}"
),
system: Some(crate::prompts::PRE_ANALYSIS_SYSTEM.to_string()),
schema: Self::pre_analysis_schema(),
schema_name: "pre_analysis".to_string(),
schema_description: Some(
"Intent, goal, plan, and optimized input for an agent turn".to_string(),
),
mode: StructuredMode::Auto,
max_repair_attempts: 2,
};
let result = generate_blocking(&**llm, &req)
.await
.context("LLM pre-analysis structured generation failed")?;
Self::pre_analysis_from_value(result.object, prompt)
.context("Failed to parse pre-analysis JSON from LLM response")
}
fn pre_analysis_from_value(
value: serde_json::Value,
original_prompt: &str,
) -> Result<PreAnalysis> {
let parsed: PreAnalysisResponse = serde_json::from_value(value)
.context("pre-analysis object did not match the expected response shape")?;
Self::pre_analysis_from_response(parsed, original_prompt)
}
fn pre_analysis_from_response(
parsed: PreAnalysisResponse,
original_prompt: &str,
) -> Result<PreAnalysis> {
let intent = match parsed.intent.to_lowercase().as_str() {
"plan" => crate::prompts::AgentStyle::Plan,
"explore" => crate::prompts::AgentStyle::Explore,
"verification" => crate::prompts::AgentStyle::Verification,
"codereview" | "code review" => crate::prompts::AgentStyle::CodeReview,
_ => crate::prompts::AgentStyle::GeneralPurpose,
};
let goal_description = parsed.goal.description.clone();
let goal =
AgentGoal::new(goal_description.clone()).with_criteria(parsed.goal.success_criteria);
let complexity = match parsed.execution_plan.complexity.as_str() {
"Simple" => Complexity::Simple,
"Medium" => Complexity::Medium,
"Complex" => Complexity::Complex,
"VeryComplex" => Complexity::VeryComplex,
_ => Complexity::Medium,
};
let mut plan = ExecutionPlan::new(goal_description, complexity);
for step_resp in parsed.execution_plan.steps {
let mut task = Task::new(step_resp.id, step_resp.description);
if let Some(tool) = step_resp.tool {
task = task.with_tool(tool);
}
if !step_resp.dependencies.is_empty() {
task = task.with_dependencies(step_resp.dependencies);
}
if let Some(criteria) = step_resp.success_criteria {
task = task.with_success_criteria(criteria);
}
plan.add_step(task);
}
for tool in parsed.execution_plan.required_tools {
plan.add_required_tool(tool);
}
Ok(PreAnalysis {
intent,
requires_planning: parsed.requires_planning,
goal,
execution_plan: plan,
optimized_input: if parsed.optimized_input.is_empty() {
original_prompt.to_string()
} else {
parsed.optimized_input
},
})
}
fn pre_analysis_schema() -> serde_json::Value {
serde_json::json!({
"type": "object",
"required": ["intent", "requires_planning", "goal", "execution_plan", "optimized_input"],
"properties": {
"intent": { "type": "string" },
"requires_planning": { "type": "boolean" },
"goal": {
"type": "object",
"required": ["description", "success_criteria"],
"properties": {
"description": { "type": "string", "minLength": 1 },
"success_criteria": {
"type": "array",
"items": { "type": "string" }
}
}
},
"execution_plan": {
"type": "object",
"required": ["complexity", "steps"],
"properties": {
"complexity": { "type": "string" },
"steps": {
"type": "array",
"items": {
"type": "object",
"required": ["id", "description"],
"properties": {
"id": { "type": "string" },
"description": { "type": "string" },
"tool": { "type": "string" },
"dependencies": {
"type": "array",
"items": { "type": "string" }
},
"success_criteria": { "type": "string" }
}
}
},
"required_tools": {
"type": "array",
"items": { "type": "string" }
}
}
},
"optimized_input": { "type": "string" }
}
})
}
fn achievement_schema() -> serde_json::Value {
serde_json::json!({
"type": "object",
"additionalProperties": false,
"required": ["achieved", "progress", "remaining_criteria"],
"properties": {
"achieved": { "type": "boolean" },
"progress": {
"type": "number",
"minimum": 0.0,
"maximum": 1.0
},
"remaining_criteria": {
"type": "array",
"items": { "type": "string" }
}
}
})
}
fn achievement_from_value(value: serde_json::Value) -> Result<AchievementResult> {
let parsed: AchievementResponse = serde_json::from_value(value)
.context("achievement object did not match the expected response shape")?;
Ok(AchievementResult {
achieved: parsed.achieved,
progress: parsed.progress.clamp(0.0, 1.0),
remaining_criteria: parsed.remaining_criteria,
})
}
fn parse_plan_response(text: &str) -> Result<ExecutionPlan> {
let parsed: PlanResponse = Self::parse_json_lenient(text)
.context("Failed to parse plan JSON from LLM response")?;
let complexity = match parsed.complexity.as_str() {
"Simple" => Complexity::Simple,
"Medium" => Complexity::Medium,
"Complex" => Complexity::Complex,
"VeryComplex" => Complexity::VeryComplex,
_ => Complexity::Medium,
};
let mut plan = ExecutionPlan::new(parsed.goal, complexity);
for step_resp in parsed.steps {
let mut task = Task::new(step_resp.id, step_resp.description);
if let Some(tool) = step_resp.tool {
task = task.with_tool(tool);
}
if !step_resp.dependencies.is_empty() {
task = task.with_dependencies(step_resp.dependencies);
}
if let Some(criteria) = step_resp.success_criteria {
task = task.with_success_criteria(criteria);
}
plan.add_step(task);
}
for tool in parsed.required_tools {
plan.add_required_tool(tool);
}
Ok(plan)
}
fn parse_goal_response(text: &str) -> Result<AgentGoal> {
let parsed: GoalResponse = Self::parse_json_lenient(text)
.context("Failed to parse goal JSON from LLM response")?;
Ok(AgentGoal::new(parsed.description).with_criteria(parsed.success_criteria))
}
#[cfg(test)]
fn parse_achievement_response(text: &str) -> Result<AchievementResult> {
let parsed: AchievementResponse = Self::parse_json_lenient(text)
.context("Failed to parse achievement JSON from LLM response")?;
Ok(AchievementResult {
achieved: parsed.achieved,
progress: parsed.progress.clamp(0.0, 1.0),
remaining_criteria: parsed.remaining_criteria,
})
}
fn parse_json_lenient<T: serde::de::DeserializeOwned>(text: &str) -> Result<T> {
let value = crate::llm::structured::extract_json_value(text)?;
Ok(serde_json::from_value(value)?)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_plan_response() {
let json = r#"{
"goal": "Build a REST API",
"complexity": "Complex",
"steps": [
{
"id": "step-1",
"description": "Set up project structure",
"tool": "bash",
"dependencies": [],
"success_criteria": "Project directory created"
},
{
"id": "step-2",
"description": "Implement endpoints",
"tool": "write",
"dependencies": ["step-1"],
"success_criteria": "Endpoints respond correctly"
}
],
"required_tools": ["bash", "write", "read"]
}"#;
let plan = LlmPlanner::parse_plan_response(json).unwrap();
assert_eq!(plan.goal, "Build a REST API");
assert_eq!(plan.complexity, Complexity::Complex);
assert_eq!(plan.steps.len(), 2);
assert_eq!(plan.steps[0].id, "step-1");
assert_eq!(plan.steps[0].tool, Some("bash".to_string()));
assert_eq!(plan.steps[1].dependencies, vec!["step-1".to_string()]);
assert_eq!(plan.required_tools, vec!["bash", "write", "read"]);
}
#[test]
fn test_parse_plan_response_with_markdown_fences() {
let json = "```json\n{\"goal\": \"Test\", \"complexity\": \"Simple\", \"steps\": [{\"id\": \"step-1\", \"description\": \"Do it\"}], \"required_tools\": []}\n```";
let plan = LlmPlanner::parse_plan_response(json).unwrap();
assert_eq!(plan.goal, "Test");
assert_eq!(plan.complexity, Complexity::Simple);
assert_eq!(plan.steps.len(), 1);
}
#[test]
fn test_parse_plan_response_invalid() {
let bad_json = "This is not JSON at all";
let result = LlmPlanner::parse_plan_response(bad_json);
assert!(result.is_err());
}
#[test]
fn test_parse_plan_response_unknown_complexity() {
let json =
r#"{"goal": "Test", "complexity": "Unknown", "steps": [], "required_tools": []}"#;
let plan = LlmPlanner::parse_plan_response(json).unwrap();
assert_eq!(plan.complexity, Complexity::Medium); }
#[test]
fn test_parse_goal_response() {
let json = r#"{
"description": "Deploy the application to production",
"success_criteria": [
"All tests pass",
"Application is accessible at production URL",
"Health check returns 200"
]
}"#;
let goal = LlmPlanner::parse_goal_response(json).unwrap();
assert_eq!(goal.description, "Deploy the application to production");
assert_eq!(goal.success_criteria.len(), 3);
assert_eq!(goal.success_criteria[0], "All tests pass");
}
#[test]
fn test_parse_goal_response_invalid() {
let result = LlmPlanner::parse_goal_response("not json");
assert!(result.is_err());
}
#[test]
fn test_parse_achievement_response() {
let json = r#"{
"achieved": false,
"progress": 0.65,
"remaining_criteria": ["Health check not verified"]
}"#;
let result = LlmPlanner::parse_achievement_response(json).unwrap();
assert!(!result.achieved);
assert!((result.progress - 0.65).abs() < f32::EPSILON);
assert_eq!(result.remaining_criteria, vec!["Health check not verified"]);
}
#[test]
fn test_parse_achievement_response_achieved() {
let json = r#"{"achieved": true, "progress": 1.0, "remaining_criteria": []}"#;
let result = LlmPlanner::parse_achievement_response(json).unwrap();
assert!(result.achieved);
assert!((result.progress - 1.0).abs() < f32::EPSILON);
assert!(result.remaining_criteria.is_empty());
}
#[test]
fn test_parse_achievement_response_clamps_progress() {
let json = r#"{"achieved": false, "progress": 1.5, "remaining_criteria": []}"#;
let result = LlmPlanner::parse_achievement_response(json).unwrap();
assert!((result.progress - 1.0).abs() < f32::EPSILON);
}
#[test]
fn test_fallback_plan() {
let short_prompt = "Fix bug";
let plan = LlmPlanner::fallback_plan(short_prompt);
assert_eq!(plan.complexity, Complexity::Simple);
assert_eq!(plan.steps.len(), 2);
assert_eq!(plan.goal, short_prompt);
let long_prompt = "Implement a comprehensive authentication system with OAuth2 support, JWT tokens, refresh token rotation, multi-factor authentication, and role-based access control across all API endpoints with proper audit logging and session management capabilities for both web and mobile clients, including password reset flows, account lockout policies, and integration with external identity providers such as Google, GitHub, and SAML-based enterprise SSO systems";
let plan = LlmPlanner::fallback_plan(long_prompt);
assert_eq!(plan.complexity, Complexity::VeryComplex);
assert_eq!(plan.steps.len(), 10);
}
#[test]
fn test_fallback_goal() {
let goal = LlmPlanner::fallback_goal("Fix the login bug");
assert_eq!(goal.description, "Fix the login bug");
assert_eq!(goal.success_criteria.len(), 2);
assert_eq!(goal.success_criteria[0], "Task is completed successfully");
}
#[test]
fn test_fallback_check_achievement_fails_closed_on_done_text() {
let goal = AgentGoal::new("Test task").with_criteria(vec!["Criterion 1".to_string()]);
let result = LlmPlanner::fallback_check_achievement(&goal, "The task is done.");
assert!(!result.achieved);
assert!(result.progress.abs() < f32::EPSILON);
assert_eq!(result.remaining_criteria, vec!["Criterion 1"]);
}
#[test]
fn test_fallback_check_achievement_rejects_negated_completion_text() {
let goal = AgentGoal::new("Test task").with_criteria(vec!["Criterion 1".to_string()]);
for state in [
"The task is not complete.",
"Verification is not done.",
"The implementation is unfinished.",
] {
let result = LlmPlanner::fallback_check_achievement(&goal, state);
assert!(!result.achieved, "must fail closed for {state:?}");
assert_eq!(result.remaining_criteria, vec!["Criterion 1"]);
}
}
#[test]
fn test_fallback_check_achievement_not_done() {
let goal = AgentGoal::new("Test task")
.with_criteria(vec!["Criterion 1".to_string(), "Criterion 2".to_string()]);
let result = LlmPlanner::fallback_check_achievement(&goal, "Work in progress");
assert!(!result.achieved);
assert_eq!(result.remaining_criteria.len(), 2);
}
#[test]
fn test_parse_json_lenient_plain() {
let v: serde_json::Value = LlmPlanner::parse_json_lenient(" {\"a\": 1} ").unwrap();
assert_eq!(v["a"], 1);
}
#[test]
fn test_parse_json_lenient_with_fences() {
let text = "```json\n{\"a\": 1}\n```";
let v: serde_json::Value = LlmPlanner::parse_json_lenient(text).unwrap();
assert_eq!(v["a"], 1);
}
#[test]
fn test_parse_json_lenient_with_surrounding_prose() {
let text = "Here is the plan:\n{\"goal\": \"test\"}\nDone.";
let v: serde_json::Value = LlmPlanner::parse_json_lenient(text).unwrap();
assert_eq!(v["goal"], "test");
}
#[test]
fn test_parse_json_lenient_brace_inside_string_value() {
let text = "Result: {\"note\": \"use a closing brace } here\"} -- end.";
let v: serde_json::Value = LlmPlanner::parse_json_lenient(text).unwrap();
assert_eq!(v["note"], "use a closing brace } here");
}
#[test]
fn test_parse_json_lenient_fenced_with_trailing_prose() {
let text = "```json\n{\"goal\": \"ship\"}\n```\nNote: revisit the `plan` later.";
let v: serde_json::Value = LlmPlanner::parse_json_lenient(text).unwrap();
assert_eq!(v["goal"], "ship");
}
#[test]
fn test_parse_json_lenient_rejects_non_json() {
let err = LlmPlanner::parse_json_lenient::<serde_json::Value>("no json here at all");
assert!(err.is_err());
}
struct ReplayClient {
responses: std::sync::Mutex<Vec<String>>,
}
impl ReplayClient {
fn new(responses: Vec<String>) -> Self {
Self {
responses: std::sync::Mutex::new(responses),
}
}
}
#[async_trait::async_trait]
impl LlmClient for ReplayClient {
async fn complete(
&self,
_messages: &[Message],
_system: Option<&str>,
_tools: &[crate::llm::ToolDefinition],
) -> anyhow::Result<crate::llm::LlmResponse> {
let text = {
let mut r = self.responses.lock().unwrap();
if r.is_empty() {
String::new()
} else {
r.remove(0)
}
};
Ok(crate::llm::LlmResponse {
message: Message {
role: "assistant".to_string(),
content: vec![crate::llm::ContentBlock::Text { text }],
reasoning_content: None,
},
usage: crate::llm::TokenUsage::default(),
stop_reason: None,
token_logprobs: Vec::new(),
meta: None,
})
}
async fn complete_streaming(
&self,
_messages: &[Message],
_system: Option<&str>,
_tools: &[crate::llm::ToolDefinition],
_cancel_token: tokio_util::sync::CancellationToken,
) -> anyhow::Result<tokio::sync::mpsc::Receiver<crate::llm::StreamEvent>> {
anyhow::bail!("streaming not used in planner tests")
}
}
#[tokio::test]
async fn test_pre_analyze_repairs_invalid_json() {
let good = r#"{"intent":"explore","requires_planning":false,"goal":{"description":"Do x","success_criteria":["done"]},"execution_plan":{"complexity":"Simple","steps":[],"required_tools":[]},"optimized_input":"Do x carefully"}"#;
let client: Arc<dyn LlmClient> = Arc::new(ReplayClient::new(vec![
"Sorry — here's the plan, but not as JSON.".to_string(),
good.to_string(),
]));
let pa = LlmPlanner::pre_analyze(&client, "do x").await.unwrap();
assert_eq!(pa.optimized_input, "Do x carefully");
}
#[tokio::test]
async fn test_pre_analyze_first_try_with_fenced_json() {
let good = format!(
"```json\n{}\n```",
r#"{"intent":"plan","requires_planning":true,"goal":{"description":"g","success_criteria":[]},"execution_plan":{"complexity":"Medium","steps":[],"required_tools":[]},"optimized_input":"opt"}"#
);
let client: Arc<dyn LlmClient> = Arc::new(ReplayClient::new(vec![good]));
let pa = LlmPlanner::pre_analyze(&client, "do x").await.unwrap();
assert_eq!(pa.optimized_input, "opt");
}
}