Skip to main content

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};
8use serde_json::json;
9
10#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
11pub struct AgentToolArgs {
12    /// The prompt for the agent to call.
13    prompt: String,
14}
15
16impl<M: CompletionModel + 'static> Tool for Agent<M> {
17    const NAME: &'static str = "agent_tool";
18
19    type Error = PromptError;
20    type Args = AgentToolArgs;
21    type Output = String;
22
23    async fn definition(&self, _prompt: String) -> ToolDefinition {
24        let description = format!(
25            "
26            Prompt a sub-agent to do a task for you.
27
28            Agent name: {name}
29            Agent description: {description}
30            Agent system prompt: {sysprompt}
31            ",
32            name = self.name(),
33            description = self.description.clone().unwrap_or_default(),
34            sysprompt = self.preamble.clone().unwrap_or_default()
35        );
36        ToolDefinition {
37            name: <Self as Tool>::name(self),
38            description,
39            parameters: json!(schema_for!(AgentToolArgs)),
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}