Refactor Autonomous Agent Execution to Asynchronous Task Management with Executor instead of loop

[hghalebi]

i replicate agent_autonomous.rs

With the Tokio executor, as I feel its final abstract is more intuitive. What do you think?

use rig::extractor::Extractor;
use rig::providers::openai::{CompletionModel, client::Client};
use schemars::JsonSchema;
use std::{
    env,
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};
use tokio::sync::oneshot;
use tracing::{error, info, instrument};

#[derive(serde::Deserialize, serde::Serialize, JsonSchema, Debug)]
pub struct AgentResponse {
    action: u32,
    goal: u32,
}

pub struct AutonomousAgent {
    llm: Arc<Extractor<CompletionModel, AgentResponse>>,
    progress: u32,
    goal: u32,
    response_channel: Option<oneshot::Receiver<Result<AgentResponse, String>>>,
}

impl AutonomousAgent {
    pub fn new(llm: Extractor<CompletionModel, AgentResponse>) -> Self {
        Self {
            llm: Arc::new(llm),
            progress: 0,
            goal: 0,
            response_channel: None,
        }
    }
}

impl Future for AutonomousAgent {
    type Output = u32;

    #[instrument(skip(self, cx))]
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.progress >= self.goal && self.goal > 0 {
            info!(progress = self.progress, goal = self.goal, "Goal achieved!");
            return Poll::Ready(self.progress);
        }

        if let Some(mut rx) = self.response_channel.take() {
            match Pin::new(&mut rx).poll(cx) {
                Poll::Pending => {
                    self.response_channel = Some(rx);
                    return Poll::Pending;
                }
                Poll::Ready(Ok(Ok(response))) => {
                    if self.goal == 0 {
                        self.goal = response.goal;
                        info!(goal = self.goal, "Goal set by agent");
                    }
                    self.progress += response.action;
                    info!(
                        action = response.action,
                        progress = self.progress,
                        goal = self.goal
                    );
                    cx.waker().wake_by_ref();
                }
                Poll::Ready(Ok(Err(e))) => {
                    error!(%e, "Agent failed to respond");
                    cx.waker().wake_by_ref();
                }
                Poll::Ready(Err(_)) => {
                    error!("Communication channel closed unexpectedly");
                    cx.waker().wake_by_ref();
                }
            }
        }

        if self.progress >= self.goal && self.goal > 0 {
            info!(
                progress = self.progress,
                goal = self.goal,
                "Goal achieved after response!"
            );
            return Poll::Ready(self.progress);
        }

        let (tx, rx) = oneshot::channel();
        let llm = Arc::clone(&self.llm);
        let context = self.progress.to_string();
        let waker = cx.waker().clone();

        tokio::spawn(async move {
            let result = llm.extract(&context).await;
            let _ = tx.send(result.map_err(|e| e.to_string()));
            waker.wake();
        });

        self.response_channel = Some(rx);
        Poll::Pending
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    let api_key = env::var("OPENAI_API_KEY")?;
    let client = Client::new(&api_key);

    let llm = client
        .extractor::<AgentResponse>("gpt-4")
        .preamble("You are an autonomous agent. Return JSON {\"action\": <random 1-500>, \"goal\": <random 1000-5000>}. On first call, set both action and goal. On subsequent calls, only set action.")
        .build();

    let final_progress = AutonomousAgent::new(llm).await;
    println!("Agent completed with progress: {final_progress}");
    Ok(())
}

[joshua-mo-143]

Hi, thanks for posting!

I'm not quite sure I agree with this being more intuitive. impl'ing Future when there is no real reason to just makes things overly complicated for no reason and while rig is being used in production, there are also additionally a number of people who are not so experienced at Rust who want to use rig. We should therefore leverage the principle of least astonishment (ie, consistency) by not randomly implementing Future on user-facing code unless there's a good reason to do so.

edit: Of course, there may be instances in the future where there may be a good reason to impl Future on things. We'll need to make sure to thoroughly document it and why it's been done that way, though.

[hghalebi]

Thanks for your feedback!

I implemented the Future trait mainly to unify error handling within the execution loop. Sometimes, even if the loop condition (like progress < goal) hasn't been met, we may need to retry due to transient issues, such as network glitches, API unavailability, or IO/system errors.

This implementation ensures that these errors don’t break the loop. Instead, they are absorbed and retired without crashing the flow. By encapsulating this logic in a Future, we can effectively leverage Tokio’s async runtime while maintaining continuity and responsiveness in the agent’s behavior.

I agree that documentation will be key if we adopt this pattern, primarily to support contributors who are less familiar with Rust’s async model (Like myself, I am in the process of learning)