~ knotless Search/
↑↓ move↵ open⇧↵ all results

johnandersen777.bsky.social/asyncai

0
0
main
144 lines · 3.8 KB Raw History
use std::pin::Pin;
use std::time::Duration;

use async_trait::async_trait;
use futures::Stream;
use futures::StreamExt;
use thiserror::Error;

use crate::ids::AgentId;
use crate::message::{AgentContent, AgentMessage, Role, RunStatus};

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AgentSpec {
    pub agent_id: AgentId,
    pub agent_name: String,
    pub instructions: String,
}

#[derive(Clone, Debug)]
pub struct RunRequest {
    pub agent: AgentSpec,
    /// Conversation thread; stateful tools (code mode REPL) scope by it.
    /// Runs on one thread never overlap (the runtime queues them).
    pub thread_id: Option<String>,
    pub messages: Vec<AgentMessage>,
}

#[derive(Clone, Debug)]
pub struct RunOutcome {
    pub status: RunStatus,
    pub messages: Vec<AgentMessage>,
    pub final_output: Option<AgentContent>,
    pub last_error: Option<String>,
}

#[derive(Clone, Debug)]
pub enum BackendEvent {
    Message(AgentMessage),
    Completed(RunOutcome),
}

#[derive(Clone, Debug, PartialEq, Eq, Error)]
#[error("{0}")]
pub struct BackendError(pub String);

impl BackendError {
    pub fn new(message: impl Into<String>) -> Self {
        Self(message.into())
    }
}

pub type BackendStream = Pin<Box<dyn Stream<Item = Result<BackendEvent, BackendError>> + Send>>;

#[async_trait]
pub trait AgentBackend: Send + Sync {
    async fn start_run(&self, request: RunRequest) -> Result<BackendStream, BackendError>;
}

#[derive(Clone, Debug)]
enum MockReply {
    Fixed(String),
    Echo,
}

/// Deterministic in-process backend used by the CLI and tests.
#[derive(Clone, Debug)]
pub struct MockBackend {
    reply: MockReply,
    failure: Option<String>,
    delay: Option<Duration>,
}

impl MockBackend {
    pub fn fixed(reply: impl Into<String>) -> Self {
        Self {
            reply: MockReply::Fixed(reply.into()),
            failure: None,
            delay: None,
        }
    }

    pub fn echo() -> Self {
        Self {
            reply: MockReply::Echo,
            failure: None,
            delay: None,
        }
    }

    pub fn with_failure(mut self, message: impl Into<String>) -> Self {
        self.failure = Some(message.into());
        self
    }

    pub fn with_delay(mut self, delay: Duration) -> Self {
        self.delay = Some(delay);
        self
    }
}

#[async_trait]
impl AgentBackend for MockBackend {
    async fn start_run(&self, request: RunRequest) -> Result<BackendStream, BackendError> {
        if let Some(failure) = &self.failure {
            return Err(BackendError::new(failure.clone()));
        }

        let reply = match &self.reply {
            MockReply::Fixed(reply) => reply.clone(),
            MockReply::Echo => request
                .messages
                .iter()
                .rev()
                .find_map(|message| message.content.as_text().map(str::to_owned))
                .unwrap_or_else(|| "hello from asyncai".to_string()),
        };

        let assistant = AgentMessage {
            role: Role::Assistant,
            content: AgentContent::Text(reply.clone()),
        };
        let mut messages = request.messages;
        messages.push(assistant.clone());

        let events: Vec<Result<BackendEvent, BackendError>> = vec![
            Ok(BackendEvent::Message(assistant)),
            Ok(BackendEvent::Completed(RunOutcome {
                status: RunStatus::Completed,
                messages,
                final_output: Some(AgentContent::Text(reply)),
                last_error: None,
            })),
        ];

        let stream = futures::stream::iter(events);
        match self.delay {
            Some(delay) => Ok(Box::pin(stream.then(move |event| async move {
                tokio::time::sleep(delay).await;
                event
            }))),
            None => Ok(Box::pin(stream)),
        }
    }
}