use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use agent_base::{
AgentBuilder, AgentError, AgentResult, ChatMessage, Middleware, OpenAiClient, PostLlmCtx,
RuntimeEvent, SessionId, Tool, ToolContext, ToolControlFlow, ToolOutput,
};
use async_trait::async_trait;
use dotenvy::dotenv;
use serde_json::{Value, json};
struct EchoTool;
#[async_trait]
impl Tool for EchoTool {
fn name(&self) -> &'static str {
"echo"
}
fn definition(&self) -> Value {
json!({
"type": "function",
"function": {
"name": "echo",
"description": "echo back the message",
"parameters": {
"type": "object",
"properties": {
"message": { "type": "string" }
},
"required": ["message"]
}
}
})
}
async fn call(&self, args: &Value, _ctx: &ToolContext) -> AgentResult<ToolOutput> {
let msg = args["message"].as_str().unwrap_or("");
Ok(ToolOutput {
summary: format!("echo: {msg}"),
raw: Some(json!({ "echo": msg })),
control_flow: ToolControlFlow::Break,
truncation: None,
})
}
}
struct AddTool;
#[async_trait]
impl Tool for AddTool {
fn name(&self) -> &'static str {
"add"
}
fn definition(&self) -> Value {
json!({
"type": "function",
"function": {
"name": "add",
"description": "Calculate the sum of two integers",
"parameters": {
"type": "object",
"properties": {
"a": { "type": "integer", "description": "First addend" },
"b": { "type": "integer", "description": "Second addend" }
},
"required": ["a", "b"]
}
}
})
}
async fn call(&self, args: &Value, _ctx: &ToolContext) -> AgentResult<ToolOutput> {
let a = args["a"].as_i64().unwrap_or(0);
let b = args["b"].as_i64().unwrap_or(0);
let result = a + b;
Ok(ToolOutput {
summary: format!("{} + {} = {}", a, b, result),
raw: Some(json!({ "result": result })),
control_flow: ToolControlFlow::Break,
truncation: None,
})
}
}
struct AntiHallucinationMiddleware {
max_nudges: usize,
nudge_count: AtomicUsize,
}
impl AntiHallucinationMiddleware {
fn new(max_nudges: usize) -> Self {
Self {
max_nudges,
nudge_count: AtomicUsize::new(0),
}
}
}
#[async_trait]
impl Middleware for AntiHallucinationMiddleware {
async fn on_post_llm(&self, ctx: &mut PostLlmCtx) -> AgentResult<()> {
if ctx.available_tools.is_empty() || ctx.total_tool_calls > 0 {
return Ok(());
}
if ctx.is_tool_call {
return Ok(());
}
if ctx.full_text.is_empty() {
return Ok(());
}
let count = self.nudge_count.fetch_add(1, Ordering::SeqCst);
if count >= self.max_nudges {
return Ok(());
}
println!(
"[Middleware] Nudge #{count}: suppressing hallucination text, injecting follow-up message"
);
ctx.skip_push = true;
ctx.follow_up_message = Some(
"CRITICAL: You have tools available but did not call any. \
Do NOT describe what you would do — actually DO it by calling tools. \
Call the appropriate tool NOW."
.to_string(),
);
Ok(())
}
}
struct SensitiveContentFilter {
blocked_words: Vec<String>,
}
impl SensitiveContentFilter {
fn new(words: Vec<String>) -> Self {
Self {
blocked_words: words,
}
}
}
#[async_trait]
impl Middleware for SensitiveContentFilter {
async fn on_post_llm(&self, ctx: &mut PostLlmCtx) -> AgentResult<()> {
let lower = ctx.full_text.to_lowercase();
for word in &self.blocked_words {
if lower.contains(&word.to_lowercase()) {
println!("[Middleware] Blocked sensitive content containing '{word}', suppressing");
ctx.skip_push = true;
ctx.full_text.clear();
return Ok(());
}
}
Ok(())
}
}
struct EventPrinter;
impl EventPrinter {
fn handle(event: RuntimeEvent) -> AgentResult<()> {
match event {
RuntimeEvent::TextDelta { text, .. } => {
print!("{}", text);
use std::io::{self, Write};
io::stdout().flush().unwrap();
}
RuntimeEvent::ToolCallStarted {
tool_name,
args_json,
..
} => {
println!();
println!("[Tool call] {} (args: {})", tool_name, args_json);
}
RuntimeEvent::ToolCallFinished {
tool_name, summary, ..
} => {
println!("[Tool finished] {} -> {}", tool_name, summary);
}
RuntimeEvent::RunFinished { .. } => {
println!();
println!("[Run finished]");
}
_ => {}
}
Ok(())
}
}
async fn print_session_messages(runtime: &agent_base::AgentRuntime, session_id: &SessionId) {
match runtime.session(session_id).await {
Some(session) => {
println!("\n--- Session Messages ---");
for msg in session.chat_messages() {
match msg {
ChatMessage::System { content, .. } => {
println!("[SYSTEM] {}", content);
}
ChatMessage::User { content, .. } => {
if content.len() > 120 {
println!("[USER] {}...", &content[..120]);
} else {
println!("[USER] {}", content);
}
}
ChatMessage::Assistant {
content,
tool_calls,
..
} => {
if let Some(tc) = tool_calls {
println!(
"[ASSISTANT] tool_calls: {:?}",
tc.iter()
.map(|t| format!("{}({})", t.name, t.arguments))
.collect::<Vec<_>>()
);
} else if let Some(c) = content {
println!("[ASSISTANT] {}", c);
}
}
ChatMessage::Tool {
tool_call_id,
content,
..
} => {
if content.len() > 120 {
println!("[TOOL:{}] {}...", tool_call_id, &content[..120]);
} else {
println!("[TOOL:{}] {}", tool_call_id, content);
}
}
}
}
println!("-------------------------");
}
None => {
println!("Session not found");
}
}
}
const SYSTEM_PROMPT: &str = r#"You are a helpful assistant with calculation tools.
Available tools:
- add: Calculate the sum of two integers
- echo: echo back a message
When the user asks for a calculation, you MUST use the appropriate tool.
Do NOT describe what you would do — actually call the tool."#;
#[tokio::main]
async fn main() -> AgentResult<()> {
dotenv().ok();
let api_key = std::env::var("OPENAI_API_KEY")
.or_else(|_| std::env::var("DASHSCOPE_API_KEY"))
.map_err(|_| {
AgentError::internal(
"Please set OPENAI_API_KEY or DASHSCOPE_API_KEY environment variable",
)
})?;
let model = std::env::var("OPENAI_MODEL")
.or_else(|_| std::env::var("DASHSCOPE_MODEL"))
.unwrap_or_else(|_| "gpt-4o-mini".to_string());
let base_url = std::env::var("OPENAI_BASE_URL")
.or_else(|_| std::env::var("DASHSCOPE_BASE_URL"))
.unwrap_or_else(|_| "https://api.openai.com/v1".to_string());
let llm_client = Arc::new(OpenAiClient::new(api_key, model.clone(), Some(base_url)));
let runtime = AgentBuilder::new(llm_client)
.system_prompt(SYSTEM_PROMPT)
.register_tool(AddTool)
.register_tool(EchoTool)
.middleware(AntiHallucinationMiddleware::new(3))
.middleware(SensitiveContentFilter::new(vec![
"password".to_string(),
"api_key".to_string(),
]))
.build()
.unwrap();
let session_id = runtime.create_session().await;
println!("=== Middleware Demo ===");
println!("Model: {model}");
println!();
println!("This demo shows two middleware use cases:");
println!(
" 1. AntiHallucinationMiddleware — if LLM text-replies without calling tools, suppress + nudge"
);
println!(
" 2. SensitiveContentFilter — if LLM output contains blocked words, suppress the message"
);
println!();
println!("Try asking: 'calculate 3 + 5'");
println!("Type 'exit' to quit, 'session' to view session history");
println!();
loop {
use std::io::{self, Write};
print!("> ");
io::stdout().flush().unwrap();
let mut input = String::new();
io::stdin()
.read_line(&mut input)
.map_err(|e| AgentError::internal(format!("Failed to read input: {e}")))?;
let input = input.trim().to_string();
if input.is_empty() {
continue;
}
if matches!(input.as_str(), "exit" | "quit") {
break;
}
if input == "session" {
print_session_messages(&runtime, &session_id).await;
continue;
}
match runtime
.run_turn(session_id.clone(), &input, |event| {
EventPrinter::handle(event)
})
.await
{
Ok(outcome) => {
println!("\nOutcome: {:?}", outcome);
}
Err(e) => {
if e.is_cancelled() {
println!("Cancelled");
} else {
println!("Error: {}", e);
}
}
}
}
Ok(())
}