use async_trait::async_trait;
use futures_util::Stream;
use lc_core::runnables::{LcelError, Runnable, RunnableConfig};
use std::pin::Pin;
use std::sync::Arc;
use crate::base::AgentExecutor;
use crate::orchestration::{Orchestrator, RunContext};
use crate::streaming::AgentStreamEvent;
pub struct AgentRunnable {
executor: Arc<AgentExecutor>,
}
impl AgentRunnable {
pub fn new(executor: Arc<AgentExecutor>) -> Self {
Self { executor }
}
}
#[async_trait]
impl Runnable<String, String> for AgentRunnable {
type Error = LcelError;
async fn invoke(
&self,
input: String,
config: Option<RunnableConfig>,
) -> Result<String, LcelError> {
self.executor
.invoke_with_config(input, config)
.await
.map_err(|e| LcelError::Agent(e.to_string()))
}
async fn stream(
&self,
input: String,
_config: Option<RunnableConfig>,
) -> Result<Pin<Box<dyn Stream<Item = Result<String, LcelError>> + Send>>, LcelError> {
use futures_util::StreamExt;
let event_stream = self.executor.stream(input);
let mapped = event_stream.filter_map(|event_result| async move {
match event_result {
Ok(event) => match event {
crate::streaming::AgentStreamEvent::FinalAnswer { content } => {
Some(Ok(content))
}
crate::streaming::AgentStreamEvent::Error { message } => {
Some(Err(LcelError::Agent(message)))
}
_ => None,
},
Err(e) => Some(Err(LcelError::Agent(e.to_string()))),
}
});
Ok(Box::pin(mapped))
}
}
pub struct AgentEventRunnable {
executor: Arc<AgentExecutor>,
}
impl AgentEventRunnable {
pub fn new(executor: Arc<AgentExecutor>) -> Self {
Self { executor }
}
}
#[async_trait]
impl Runnable<String, AgentStreamEvent> for AgentEventRunnable {
type Error = LcelError;
async fn invoke(
&self,
input: String,
config: Option<RunnableConfig>,
) -> Result<AgentStreamEvent, LcelError> {
let output = self
.executor
.invoke_with_config(input, config)
.await
.map_err(|e| LcelError::Agent(e.to_string()))?;
Ok(AgentStreamEvent::FinalAnswer { content: output })
}
async fn stream(
&self,
input: String,
_config: Option<RunnableConfig>,
) -> Result<Pin<Box<dyn Stream<Item = Result<AgentStreamEvent, LcelError>> + Send>>, LcelError>
{
use futures_util::StreamExt;
let event_stream = self.executor.stream(input);
let mapped = event_stream
.map(|event_result| event_result.map_err(|e| LcelError::Agent(e.to_string())));
Ok(Box::pin(mapped))
}
}
pub struct OrchestratorRunnable<O: Orchestrator> {
orchestrator: O,
}
impl<O: Orchestrator> OrchestratorRunnable<O> {
pub fn new(orchestrator: O) -> Self {
Self { orchestrator }
}
}
#[async_trait]
impl<O> Runnable<O::Input, O::Output> for OrchestratorRunnable<O>
where
O: Orchestrator,
O::Input: Send + Sync + 'static,
O::Output: Send + Sync + 'static,
{
type Error = LcelError;
async fn invoke(
&self,
input: O::Input,
config: Option<RunnableConfig>,
) -> Result<O::Output, LcelError> {
let ctx = match &config {
Some(cfg) => RunContext::from_config(cfg),
None => RunContext::new_random(),
};
self.orchestrator
.run_with_context(input, &ctx)
.await
.map_err(|e| LcelError::Agent(e.to_string()))
}
}
impl From<crate::base::AgentError> for LcelError {
fn from(err: crate::base::AgentError) -> Self {
LcelError::Agent(err.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::streaming::AgentStreamEvent;
use std::collections::HashMap;
#[test]
fn agent_error_into_lcel_error() {
let agent_err = crate::base::AgentError::MaxIterationsReached;
let lcel_err: LcelError = agent_err.into();
assert!(matches!(lcel_err, LcelError::Agent(_)));
assert!(lcel_err.to_string().contains("Max iterations"));
}
struct TestFinishAgent;
#[async_trait]
impl crate::BaseAgent for TestFinishAgent {
async fn plan(
&self,
_intermediate_steps: &[crate::types::AgentStep],
_inputs: &HashMap<String, String>,
) -> Result<crate::types::AgentOutput, crate::base::AgentError> {
Ok(crate::types::AgentOutput::Finish(
crate::types::AgentFinish::new("answer".to_string(), String::new()),
))
}
}
#[tokio::test]
async fn agent_event_runnable_preserves_all_events() {
use futures_util::StreamExt;
let executor = Arc::new(crate::base::AgentExecutor::new(
Arc::new(TestFinishAgent),
vec![],
));
let runnable = AgentEventRunnable::new(executor);
let mut stream = runnable.stream("hi".to_string(), None).await.unwrap();
let mut events = Vec::new();
while let Some(item) = stream.next().await {
events.push(item.unwrap());
}
assert_eq!(events.len(), 2);
assert!(matches!(events[0], AgentStreamEvent::Text { .. }));
assert!(matches!(events[1], AgentStreamEvent::FinalAnswer { .. }));
}
#[tokio::test]
async fn agent_event_runnable_invoke_returns_final_answer() {
let executor = Arc::new(crate::base::AgentExecutor::new(
Arc::new(TestFinishAgent),
vec![],
));
let runnable = AgentEventRunnable::new(executor);
let event = runnable.invoke("hi".to_string(), None).await.unwrap();
match event {
AgentStreamEvent::FinalAnswer { content } => assert_eq!(content, "answer"),
other => panic!("expected FinalAnswer, got {:?}", other),
}
}
}