open-standards/eoc/eoc-rs/crates/eoc-agent/src/react.rs

//! ReAct — Thought / Action / Observation.
//!
//! Loop:
//!
//! 1. Ask the LLM for the next *Thought* and *Action*.
//! 2. If the action is `Finish[answer]`, return.
//! 3. Otherwise, look up the named tool, run it, and feed the
//!    *Observation* back as part of the next prompt.
//! 4. Repeat until terminated by [`Budget`](crate::agent::Budget).
//!
//! The provider is expected to emit lines of the form:
//!
//! ```text
//! Thought: I should use the calc tool.
//! Action: calc[2+2]
//! ```
//!
//! Or the final form `Action: Finish[42]`. The format is intentionally
//! permissive — any trailing garbage after `Action:` is treated as the
//! action body.

use std::sync::Arc;

use async_trait::async_trait;

use crate::agent::{Agent, AgentLoop, Budget, LlmProvider, StopReason, ToolBox};
use crate::error::{AgentError, AgentResult};
use crate::trace::{Span, SpanKind};

/// One parsed ReAct step.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReActStep {
    /// The model's chain-of-thought line (free text).
    pub thought: String,
    /// The action name (`"Finish"` if the agent is terminating).
    pub action: String,
    /// The action body — for `Finish[...]` this is the final answer.
    pub argument: String,
}

impl ReActStep {
    /// Parse a model completion into a step. Lenient — missing fields
    /// become empty strings.
    pub fn parse(text: &str) -> AgentResult<Self> {
        let mut thought = String::new();
        let mut action = String::new();
        let mut argument = String::new();
        for line in text.lines() {
            let line = line.trim();
            if let Some(rest) = line.strip_prefix("Thought:") {
                thought = rest.trim().to_string();
            } else if let Some(rest) = line.strip_prefix("Action:") {
                let rest = rest.trim();
                // Parse `name[arg]` or bare `name`.
                if let (Some(lb), Some(rb)) = (rest.find('['), rest.rfind(']')) {
                    if lb < rb {
                        action = rest[..lb].trim().to_string();
                        argument = rest[lb + 1..rb].to_string();
                        continue;
                    }
                }
                action = rest.to_string();
            }
        }
        if action.is_empty() {
            return Err(AgentError::Parse(format!(
                "no Action: line in completion: {text:?}"
            )));
        }
        Ok(Self {
            thought,
            action,
            argument,
        })
    }
}

/// ReAct driver.
pub struct ReActLoop {
    provider: Arc<dyn LlmProvider>,
    tools: ToolBox,
    inner: AgentLoop,
    /// System preamble prepended to each prompt.
    pub system: String,
    transcript: String,
    last_output: String,
}

impl ReActLoop {
    /// Build a ReAct loop. `system` is prepended to every prompt; it
    /// should describe the available tools and the `Thought:` / `Action:`
    /// format the model is expected to emit.
    pub fn new(
        provider: Arc<dyn LlmProvider>,
        tools: ToolBox,
        budget: Budget,
        system: impl Into<String>,
    ) -> Self {
        Self {
            provider,
            tools,
            inner: AgentLoop::new(budget),
            system: system.into(),
            transcript: String::new(),
            last_output: String::new(),
        }
    }

    fn build_prompt(&self, goal: &str) -> String {
        if self.transcript.is_empty() {
            format!("{sys}\n\nGoal: {goal}\n", sys = self.system)
        } else {
            format!(
                "{sys}\n\nGoal: {goal}\n{tr}",
                sys = self.system,
                tr = self.transcript
            )
        }
    }
}

#[async_trait]
impl Agent for ReActLoop {
    async fn run(&mut self, goal: &str) -> AgentResult<(String, StopReason)> {
        for iter in 0..self.inner.budget.max_iterations {
            if let Some(stop) = self.inner.iteration_check(iter) {
                return Ok((self.last_output.clone(), stop));
            }

            let prompt = self.build_prompt(goal);
            let reply = self.provider.complete(&prompt).await?;
            self.last_output = reply.text.clone();

            if let Some(stop) =
                self.inner.charge_llm_step(SpanKind::Think, "react.step", &prompt, &reply)
            {
                return Ok((self.last_output.clone(), stop));
            }

            let step = match ReActStep::parse(&reply.text) {
                Ok(s) => s,
                Err(e) => return Ok((self.last_output.clone(), StopReason::Error(e.to_string()))),
            };

            if step.action.eq_ignore_ascii_case("Finish") {
                return Ok((step.argument, StopReason::Done));
            }

            let tool = match self.tools.get(&step.action) {
                Some(t) => t,
                None => {
                    let obs = format!("ERROR: unknown tool `{}`", step.action);
                    self.transcript
                        .push_str(&format!("Observation: {obs}\n"));
                    self.inner.push_span(
                        Span::new(0, SpanKind::Observe, step.action.clone(), step.argument.clone(), obs.clone()),
                    );
                    continue;
                }
            };

            let args = serde_json::json!({ "input": step.argument });
            // Charged what the meter measured, or at least the step's
            // duration estimate: a tool that ran is never 0 µJ. The span
            // says which.
            let (result, cost) = self
                .inner
                .measured_tool(|| async { tool.run(args).await })
                .await;
            let observation = match result {
                Ok(v) => v.to_string(),
                Err(e) => format!("ERROR: {e}"),
            };
            if let Some(stop) = self.inner.budget.charge_tool(cost.microjoules) {
                return Ok((self.last_output.clone(), stop));
            }
            self.inner.push_span(
                Span::new(
                    0,
                    SpanKind::Act,
                    step.action.clone(),
                    step.argument.clone(),
                    observation.clone(),
                )
                .with_joule_cost(cost),
            );
            self.transcript.push_str(&format!(
                "Thought: {th}\nAction: {ac}[{arg}]\nObservation: {obs}\n",
                th = step.thought,
                ac = step.action,
                arg = step.argument,
                obs = observation
            ));
        }
        Ok((self.last_output.clone(), StopReason::MaxIterations))
    }

    fn trace(&self) -> &crate::trace::Trace {
        &self.inner.trace
    }

    fn budget(&self) -> &Budget {
        &self.inner.budget
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use eoc_core::{JouleCost, JouleSource};

    #[test]
    fn parse_thought_action() {
        let s = ReActStep::parse("Thought: hi\nAction: calc[1+1]").unwrap();
        assert_eq!(s.thought, "hi");
        assert_eq!(s.action, "calc");
        assert_eq!(s.argument, "1+1");
    }

    #[test]
    fn parse_finish() {
        let s = ReActStep::parse("Action: Finish[42]").unwrap();
        assert_eq!(s.action, "Finish");
        assert_eq!(s.argument, "42");
    }

    #[test]
    fn parse_missing_action_errors() {
        let err = ReActStep::parse("Thought: nope\n").unwrap_err();
        assert!(matches!(err, AgentError::Parse(_)));
    }

    /// A provider that replies with a script, 1,000 µJ a reply.
    struct Replies(std::sync::Mutex<Vec<&'static str>>);

    #[async_trait]
    impl LlmProvider for Replies {
        async fn complete(&self, _prompt: &str) -> AgentResult<crate::agent::LlmReply> {
            let text = self.0.lock().unwrap().remove(0);
            Ok(crate::agent::LlmReply::new(text, 5, 1_000))
        }
    }

    struct Echo;

    #[async_trait]
    impl crate::agent::Tool for Echo {
        fn name(&self) -> &str {
            "echo"
        }
        async fn run(&self, args: serde_json::Value) -> AgentResult<serde_json::Value> {
            Ok(args)
        }
    }

    /// A tool that takes at least 20 ms and echoes its arguments.
    struct Slow;

    #[async_trait]
    impl crate::agent::Tool for Slow {
        fn name(&self) -> &str {
            "slow"
        }
        async fn run(&self, args: serde_json::Value) -> AgentResult<serde_json::Value> {
            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
            Ok(args)
        }
    }

    /// A tool step whose first meter read fails is charged its own estimate
    /// (its duration at the default power), never the counter's since-boot
    /// total: with a 5 J budget and a counter that reads 9 J once it reads at
    /// all, the agent finishes, having charged its two replies and the tool
    /// step's estimate, which is not 0.
    #[tokio::test]
    async fn a_failed_first_read_never_exhausts_the_joule_budget() {
        let meter = Arc::new(crate::agent::scripted::Scripted::new(&[None, Some(9_000_000)]));
        let provider = Arc::new(Replies(std::sync::Mutex::new(vec!["Action: echo[hi]", "Action: Finish[done]"])));
        let mut tools = ToolBox::new();
        tools.insert(Echo);
        let budget = Budget::new(1_000, 5.0, 10, 10);
        let mut agent = ReActLoop::new(provider, tools, budget, "system");
        agent.inner = AgentLoop::with_meter(budget, meter.clone());
        let (out, stop) = agent.run("go").await.unwrap();
        assert_eq!((out.as_str(), stop), ("done", StopReason::Done));
        assert_eq!(meter.reads(), 1, "the counter was read again after a failed first read");
        let act = agent.trace().spans.iter().find(|s| s.kind == SpanKind::Act).expect("an act span");
        assert!(act.microjoules > 0 && act.microjoules < 9_000_000, "{:?}", act);
        assert_eq!(agent.budget().joules.used_microjoules, 2_000 + act.microjoules, "{:?}", agent.trace().spans);
    }

    /// A tool step the meter does not measure (the host default, the stub) is
    /// charged its duration estimate: at least 20 ms at the default power for
    /// a tool that runs at least 20 ms, on the span and against the budget.
    #[tokio::test]
    async fn an_unmetered_tool_step_is_charged_its_duration_estimate() {
        let provider = Arc::new(Replies(std::sync::Mutex::new(vec!["Action: slow[hi]", "Action: Finish[done]"])));
        let mut tools = ToolBox::new();
        tools.insert(Slow);
        let mut agent = ReActLoop::new(provider, tools, Budget::unlimited(10), "system");
        let (out, stop) = agent.run("go").await.unwrap();
        assert_eq!((out.as_str(), stop), ("done", StopReason::Done));
        let act = agent.trace().spans.iter().find(|s| s.kind == SpanKind::Act).expect("an act span");
        let floor = crate::agent::duration_estimate(std::time::Duration::from_millis(20), crate::agent::DEFAULT_TOOL_WATTS);
        assert!(act.microjoules >= floor, "a 20 ms tool step charged {} µJ", act.microjoules);
        assert_eq!(agent.budget().joules.used_microjoules, 2_000 + act.microjoules);
        for span in &agent.trace().spans {
            assert_eq!(span.joule_source, JouleSource::Estimated, "{span:?}");
        }
    }

    /// Each span says whether its figure was measured: a tool step over which
    /// the counter advanced 50 J is measured, and charged that; a model step
    /// (the provider's own figure) and an unknown tool's observation (nothing
    /// ran, 0 µJ) are estimates.
    #[tokio::test]
    async fn a_span_is_measured_only_when_a_live_counter_read_it() {
        let meter = Arc::new(crate::agent::scripted::Scripted::new(&[Some(0), Some(50_000_000)]));
        let provider = Arc::new(Replies(std::sync::Mutex::new(vec![
            "Action: nope[x]",
            "Action: echo[hi]",
            "Action: Finish[done]",
        ])));
        let mut tools = ToolBox::new();
        tools.insert(Echo);
        let budget = Budget::unlimited(10);
        let mut agent = ReActLoop::new(provider, tools, budget, "system");
        agent.inner = AgentLoop::with_meter(budget, meter.clone());
        let (out, stop) = agent.run("go").await.unwrap();
        assert_eq!((out.as_str(), stop), ("done", StopReason::Done));
        let costs: Vec<_> = agent.trace().spans.iter().map(|s| (s.kind, s.joule_cost())).collect();
        assert_eq!(
            costs,
            [
                (SpanKind::Think, JouleCost::estimated(1_000)),
                (SpanKind::Observe, JouleCost::estimated(0)),
                (SpanKind::Think, JouleCost::estimated(1_000)),
                (SpanKind::Act, JouleCost::measured(50_000_000)),
                (SpanKind::Think, JouleCost::estimated(1_000)),
            ]
        );
        assert_eq!(agent.budget().joules.used_microjoules, 3_000 + 50_000_000);
    }
}