rig/agent/
tool.rs

1use crate::{
2    agent::Agent,
3    completion::{CompletionModel, Prompt, PromptError, ToolDefinition},
4    tool::Tool,
5};
6use schemars::{JsonSchema, schema_for};
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
10pub struct AgentToolArgs {
11    /// The prompt for the agent to call.
12    prompt: String,
13}
14
15impl<M: CompletionModel> Tool for Agent<M> {
16    const NAME: &'static str = "agent_tool";
17
18    type Error = PromptError;
19    type Args = AgentToolArgs;
20    type Output = String;
21
22    async fn definition(&self, _prompt: String) -> ToolDefinition {
23        let description = format!(
24            "
25            Prompt a sub-agent to do a task for you.
26
27            Agent name: {name}
28            Agent description: {description}
29            Agent system prompt: {sysprompt}
30            ",
31            name = self.name(),
32            description = self.description.clone().unwrap_or_default(),
33            sysprompt = self.preamble.clone().unwrap_or_default()
34        );
35        ToolDefinition {
36            name: <Self as Tool>::name(self),
37            description,
38            parameters: serde_json::to_value(schema_for!(AgentToolArgs))
39                .expect("converting JSON schema to JSON value should never fail"),
40        }
41    }
42
43    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
44        self.prompt(args.prompt).await
45    }
46
47    fn name(&self) -> String {
48        self.name.clone().unwrap_or_else(|| Self::NAME.to_string())
49    }
50}