agentlib_reasoning/

reflect.rs

use crate::utils::{call_model, execute_tool_calls, extract_text, parse_json};
use agentlib_core::{
    ModelMessage, ModelRequest, ReasoningContext, ReasoningEngine, ReasoningStep, Role,
};
use anyhow::{Result, anyhow};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
struct CritiqueResult {
    score: u32,
    issues: Vec<String>,
    suggestion: String,
    needs_revision: bool,
}

pub struct ReflectEngine {
    max_reflections: usize,
    acceptance_threshold: u32,
    max_answer_steps: usize,
    critique_prompt: String,
}

impl ReflectEngine {
    pub fn new(
        max_reflections: usize,
        acceptance_threshold: u32,
        max_answer_steps: usize,
        critique_prompt: Option<String>,
    ) -> Self {
        Self {
            max_reflections,
            acceptance_threshold,
            max_answer_steps,
            critique_prompt: critique_prompt.unwrap_or_else(|| {
                "You are a critical evaluator. Review the answer below and assess its quality.\n\nRespond in this exact JSON format (no markdown):\n{\n  \"score\": <0-10>,\n  \"issues\": [\"<issue 1>\", \"<issue 2>\"],\n  \"suggestion\": \"<one-sentence improvement suggestion>\",\n  \"needs_revision\": <true|false>\n}\n\nBe strict. Score 10 only for perfect answers. Score < 8 if the answer is incomplete, incorrect, or could be substantially improved.".to_string()
            }),
        }
    }

    async fn generate_answer(
        &self,
        r_ctx: &mut ReasoningContext<'_>,
        messages: &mut Vec<ModelMessage>,
    ) -> Result<String> {
        let mut steps = 0;
        while steps < self.max_answer_steps {
            let response = call_model(r_ctx, messages.clone()).await?;
            messages.push(response.message.clone());

            let tool_calls = response.message.tool_calls.as_ref();
            if tool_calls.is_none() || tool_calls.unwrap().is_empty() {
                return Ok(extract_text(&response.message.content));
            }

            execute_tool_calls(r_ctx, &response).await?;
            steps += 1;
        }
        Err(anyhow!(
            "[ReflectEngine] Max answer steps ({}) reached.",
            self.max_answer_steps
        ))
    }

    async fn critique(
        &self,
        r_ctx: &mut ReasoningContext<'_>,
        question: &str,
        answer: &str,
    ) -> Result<CritiqueResult> {
        let critique_messages = vec![
            ModelMessage {
                role: Role::System,
                content: self.critique_prompt.clone(),
                tool_call_id: None,
                tool_calls: None,
            },
            ModelMessage {
                role: Role::User,
                content: format!("Question:\n{}\n\nAnswer to evaluate:\n{}", question, answer),
                tool_call_id: None,
                tool_calls: None,
            },
        ];

        let request = ModelRequest {
            messages: critique_messages,
            tools: None, // no tools for critique
        };

        let response = r_ctx.model.complete(request).await?;

        // Accumulate usage
        if let Some(usage) = &response.usage {
            r_ctx.ctx.usage.prompt_tokens += usage.prompt_tokens;
            r_ctx.ctx.usage.completion_tokens += usage.completion_tokens;
            r_ctx.ctx.usage.total_tokens += usage.total_tokens;
        }

        match parse_json::<CritiqueResult>(&response.message.content) {
            Ok(res) => Ok(res),
            Err(_) => Ok(CritiqueResult {
                score: 9,
                issues: vec![],
                suggestion: "".to_string(),
                needs_revision: false,
            }),
        }
    }

    async fn revise(
        &self,
        r_ctx: &mut ReasoningContext<'_>,
        question: &str,
        current_answer: &str,
        critique: &CritiqueResult,
    ) -> Result<String> {
        let revision_messages = vec![
            ModelMessage {
                role: Role::System,
                content: "You are revising your previous answer based on critique. Produce an improved, complete answer.".to_string(),
                tool_call_id: None,
                tool_calls: None,
            },
            ModelMessage {
                role: Role::User,
                content: format!("Original question: {}", question),
                tool_call_id: None,
                tool_calls: None,
            },
            ModelMessage {
                role: Role::Assistant,
                content: current_answer.to_string(),
                tool_call_id: None,
                tool_calls: None,
            },
            ModelMessage {
                role: Role::User,
                content: format!(
                    "Critique of your answer:\n- Score: {}/10\n- Issues: {}\n- Suggestion: {}\n\nPlease revise your answer to address these issues.",
                    critique.score,
                    critique.issues.join(", "),
                    critique.suggestion
                ),
                tool_call_id: None,
                tool_calls: None,
            }
        ];

        let response = call_model(r_ctx, revision_messages).await?;

        if let Some(tool_calls) = &response.message.tool_calls {
            if !tool_calls.is_empty() {
                r_ctx.ctx.messages.push(response.message.clone());
                execute_tool_calls(r_ctx, &response).await?;
                // One more pass after tools
                let request = ModelRequest {
                    messages: r_ctx.ctx.messages.clone(),
                    tools: None,
                };
                let final_res = r_ctx.model.complete(request).await?;
                return Ok(extract_text(&final_res.message.content));
            }
        }

        Ok(extract_text(&response.message.content))
    }
}

impl Default for ReflectEngine {
    fn default() -> Self {
        Self::new(2, 8, 5, None)
    }
}

#[async_trait]
impl ReasoningEngine for ReflectEngine {
    fn name(&self) -> &str {
        "reflect"
    }

    async fn execute(&self, r_ctx: &mut ReasoningContext<'_>) -> Result<String> {
        let mut messages = r_ctx.ctx.messages.clone();
        let mut answer = self.generate_answer(r_ctx, &mut messages).await?;
        r_ctx.push_step(ReasoningStep::Thought {
            content: "Initial answer generated.".to_string(),
            engine: self.name().to_string(),
        });

        for i in 0..self.max_reflections {
            let question = r_ctx.ctx.input.clone();
            let critique = self.critique(r_ctx, &question, &answer).await?;

            r_ctx.push_step(ReasoningStep::Reflection {
                assessment: format!(
                    "Score: {}/10. Issues: {}. {}",
                    critique.score,
                    critique.issues.join("; "),
                    critique.suggestion
                ),
                needs_revision: critique.needs_revision,
                engine: self.name().to_string(),
            });

            if !critique.needs_revision || critique.score >= self.acceptance_threshold {
                break;
            }

            r_ctx.push_step(ReasoningStep::Thought {
                content: format!(
                    "Revising answer (attempt {}/{})...",
                    i + 1,
                    self.max_reflections
                ),
                engine: self.name().to_string(),
            });

            answer = self.revise(r_ctx, &question, &answer, &critique).await?;
        }

        r_ctx.push_step(ReasoningStep::Response {
            content: answer.clone(),
            engine: self.name().to_string(),
        });

        Ok(answer)
    }
}