use async_trait::async_trait;
use oris_runtime::agent::{create_agent, Middleware, MiddlewareContext, MiddlewareError};
use oris_runtime::schemas::agent::AgentAction;
use oris_runtime::schemas::Message;
use std::sync::Arc;
struct PrefixMiddleware {
prefix: String,
}
impl PrefixMiddleware {
fn new(prefix: String) -> Self {
Self { prefix }
}
}
#[async_trait]
impl Middleware for PrefixMiddleware {
async fn after_tool_call(
&self,
_action: &AgentAction,
observation: &str,
_context: &mut MiddlewareContext,
) -> Result<Option<String>, MiddlewareError> {
Ok(Some(format!("{} {}", self.prefix, observation)))
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
let prefix_middleware = PrefixMiddleware::new("[PREFIX]".to_string());
let agent = create_agent(
"gpt-4o-mini",
&[],
Some("You are a helpful assistant."),
Some(vec![Arc::new(prefix_middleware)]),
)?;
let result = agent
.invoke_messages(vec![Message::new_human_message("Hello, how are you?")])
.await?;
println!("Agent response: {}", result);
Ok(())
}