Skip to main content

ras_agent/application/
run_step.rs

1use std::sync::Arc;
2use std::time::Instant;
3
4use chrono::Utc;
5use ras_errors::AppError;
6use ras_llm::{ChatMessage, ChatResponse, InvokeOptions, LlmClient};
7use ras_types::{ActionResult, StepId};
8
9use crate::application::compute_action_hash::compute_action_hash;
10use crate::application::detect_loop::{build_budget_warning, build_loop_nudge};
11use crate::application::fallback_llm::should_switch_to_fallback;
12use crate::domain::agent_history::StepRecord;
13use crate::domain::agent_output::{ActionInvocation, AgentBrain, AgentOutput};
14use crate::domain::loop_detector::ActionLoopDetector;
15use crate::domain::step_metadata::StepMetadata;
16
17pub struct RunStep {
18    primary_llm: Arc<dyn LlmClient>,
19    fallback_llm: Option<Arc<dyn LlmClient>>,
20}
21
22impl RunStep {
23    #[must_use]
24    pub fn new(primary: Arc<dyn LlmClient>, fallback: Option<Arc<dyn LlmClient>>) -> Self {
25        Self {
26            primary_llm: primary,
27            fallback_llm: fallback,
28        }
29    }
30
31    pub async fn execute(
32        &self,
33        step: StepId,
34        max_steps: u32,
35        prompt: Vec<ChatMessage>,
36        detector: &mut ActionLoopDetector,
37    ) -> Result<StepRecord, AppError> {
38        let started = Instant::now();
39        let mut messages = prompt;
40        if let Some(nudge) = build_loop_nudge(detector) {
41            messages.push(nudge);
42        }
43        if let Some(warn) = build_budget_warning(step.0, max_steps) {
44            messages.push(warn);
45        }
46        let response = self.invoke_with_fallback(messages).await?;
47        let output = parse_agent_output(&response)?;
48        for action in &output.action {
49            detector.record_action(compute_action_hash(action));
50        }
51        let metadata = StepMetadata {
52            duration_ms: started.elapsed().as_millis() as u64,
53            step_interval_ms: None,
54            usage: response.usage,
55            model: Some(response.model.clone()),
56            fallback_used: false,
57        };
58        Ok(StepRecord {
59            step,
60            started_at: Utc::now(),
61            url: None,
62            output,
63            results: Vec::new(),
64            metadata,
65        })
66    }
67
68    async fn invoke_with_fallback(
69        &self,
70        messages: Vec<ChatMessage>,
71    ) -> Result<ChatResponse, AppError> {
72        let opts = InvokeOptions::default();
73        match self
74            .primary_llm
75            .ainvoke(messages.clone(), opts.clone())
76            .await
77        {
78            Ok(r) => Ok(r),
79            Err(e) if should_switch_to_fallback(&e) => match &self.fallback_llm {
80                Some(fb) => fb.ainvoke(messages, opts).await,
81                None => Err(e),
82            },
83            Err(e) => Err(e),
84        }
85    }
86}
87
88fn parse_agent_output(response: &ChatResponse) -> Result<AgentOutput, AppError> {
89    if let Some(content) = &response.content {
90        if let Ok(parsed) = serde_json::from_str::<AgentOutput>(content) {
91            return Ok(parsed);
92        }
93    }
94    Ok(AgentOutput {
95        current_state: AgentBrain {
96            evaluation_previous_goal: String::new(),
97            memory: String::new(),
98            next_goal: response.content.clone().unwrap_or_default(),
99        },
100        action: tool_calls_to_actions(&response.tool_calls),
101        plan: None,
102        current_plan_item: None,
103    })
104}
105
106fn tool_calls_to_actions(calls: &[ras_llm::ToolCall]) -> Vec<ActionInvocation> {
107    calls
108        .iter()
109        .map(|c| ActionInvocation {
110            name: ras_types::ActionName(c.name.clone().into()),
111            parameters: c.arguments.clone(),
112        })
113        .collect()
114}
115
116#[must_use]
117pub fn done_result(text: impl Into<String>) -> ActionResult {
118    ActionResult::done(text)
119}