mirror of https://github.com/0xplaygrounds/rig
60 lines
1.8 KiB
Rust
60 lines
1.8 KiB
Rust
use rig::completion::Prompt;
|
|
use rig::providers;
|
|
use rig::tool::Tool;
|
|
use rig_macros::rig_tool;
|
|
use std::time::Duration;
|
|
use tracing_subscriber;
|
|
|
|
// Example demonstrating async tool usage
|
|
#[rig_tool(
|
|
description = "A tool that simulates an async operation",
|
|
params(
|
|
input = "Input value to process",
|
|
delay_ms = "Delay in milliseconds before returning result"
|
|
)
|
|
)]
|
|
async fn async_operation(input: String, delay_ms: u64) -> Result<String, rig::tool::ToolError> {
|
|
// Simulate some async work
|
|
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
|
|
|
|
// Process the input
|
|
Ok(format!(
|
|
"Processed after {}ms: {}",
|
|
delay_ms,
|
|
input.to_uppercase()
|
|
))
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
// Initialize tracing
|
|
tracing_subscriber::fmt().pretty().init();
|
|
|
|
// Create an agent with the ASYNCOPERATION tool
|
|
let async_agent = providers::openai::Client::from_env()
|
|
.agent(providers::openai::GPT_4O)
|
|
.preamble("You are an agent with tools access, always use the tools")
|
|
.max_tokens(1024)
|
|
.tool(AsyncOperation)
|
|
.build();
|
|
|
|
// Print out the tool definition to verify
|
|
println!("Tool definition:");
|
|
println!(
|
|
"ASYNCOPERATION: {}",
|
|
serde_json::to_string_pretty(&AsyncOperation.definition(String::default()).await).unwrap()
|
|
);
|
|
|
|
// Test prompts
|
|
for prompt in [
|
|
"What tools do you have?",
|
|
"Process the text 'hello world' with a delay of 1000ms",
|
|
"Process the text 'async operation' with a delay of 500ms",
|
|
"Process the text 'concurrent calls' with a delay of 200ms",
|
|
"Process the text 'error handling' with a delay of 'not a number'",
|
|
] {
|
|
println!("User: {}", prompt);
|
|
println!("Agent: {}", async_agent.prompt(prompt).await.unwrap());
|
|
}
|
|
}
|