rig_agent/agent/
tool.rs

1use std::sync::Arc;
2
3use crate::{
4    agent::Agent,
5    completion::{CompletionModel, Prompt},
6    tool::{DynamicTool, ToolExecutionError, ToolOutput},
7};
8use schemars::{JsonSchema, schema_for};
9use serde::{Deserialize, Serialize};
10use serde_json::json;
11
12#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
13struct AgentToolArgs {
14    /// The prompt for the agent to call.
15    prompt: String,
16}
17
18const DEFAULT_AGENT_TOOL_NAME: &str = "agent_tool";
19
20impl<M: CompletionModel + 'static> Agent<M> {
21    /// Convert this agent into a runtime-defined tool.
22    ///
23    /// The configured agent name becomes the tool name. Unnamed agents use
24    /// `agent_tool`. This explicit conversion keeps runtime identity out of the
25    /// statically named [`Tool`](crate::tool::Tool) trait.
26    pub fn into_tool(self) -> DynamicTool {
27        let name = self
28            .name
29            .clone()
30            .unwrap_or_else(|| DEFAULT_AGENT_TOOL_NAME.to_string());
31        let description = format!(
32            "
33            Prompt a sub-agent to do a task for you.
34
35            Agent name: {name}
36            Agent description: {description}
37            Agent system prompt: {sysprompt}
38            ",
39            name = name,
40            description = self.description.clone().unwrap_or_default(),
41            sysprompt = self.preamble.clone().unwrap_or_default()
42        );
43        let parameters = json!(schema_for!(AgentToolArgs));
44        let agent = Arc::new(self);
45
46        DynamicTool::new(name, description, parameters, move |context, args| {
47            let agent = Arc::clone(&agent);
48            let inherited_context = context.inbound_only();
49            Box::pin(async move {
50                let args: AgentToolArgs = serde_json::from_value(args).map_err(|error| {
51                    ToolExecutionError::invalid_args(format!(
52                        "failed to parse agent tool arguments: {error}"
53                    ))
54                    .with_source(error)
55                })?;
56                agent
57                    .prompt(args.prompt)
58                    .tool_context(inherited_context)
59                    .await
60                    .map(ToolOutput::text)
61                    .map_err(ToolExecutionError::from_error)
62            })
63        })
64    }
65}
66
67impl<M: CompletionModel + 'static> From<Agent<M>> for DynamicTool {
68    fn from(agent: Agent<M>) -> Self {
69        agent.into_tool()
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use crate::agent::AgentBuilder;
77    use crate::test_utils::{MockCompletionModel, MockContextProbeTool, MockTurn, SessionId};
78    use crate::tool::ToolContext;
79
80    /// A `ToolContext` set on the outer run propagates into a sub-agent
81    /// invoked as a tool, so the inner agent's own tools observe it.
82    #[tokio::test]
83    async fn context_propagates_into_sub_agent() {
84        // Inner agent: calls a context-probing tool, then answers.
85        let probe = MockContextProbeTool::default();
86        let inner_model = MockCompletionModel::new([
87            MockTurn::tool_call("c1", "context_probe", json!({})),
88            MockTurn::text("inner done"),
89        ]);
90        let inner = AgentBuilder::new(inner_model)
91            .name("researcher")
92            .tool(probe.clone())
93            .build();
94
95        // Outer agent: delegates to the inner agent (registered as the
96        // "researcher" tool), then answers.
97        let outer_model = MockCompletionModel::new([
98            MockTurn::tool_call("c2", "researcher", json!({"prompt": "do research"})),
99            MockTurn::text("outer done"),
100        ]);
101        let outer = AgentBuilder::new(outer_model)
102            .dynamic_tool(inner.into_tool())
103            .build();
104
105        let mut context = ToolContext::new();
106        context.insert(SessionId("abc-123".to_string()));
107
108        let out = outer
109            .prompt("start")
110            .tool_context(context)
111            .max_turns(5)
112            .await
113            .expect("run succeeds");
114
115        assert_eq!(out, "outer done");
116        assert_eq!(probe.observed().as_deref(), Some("session:abc-123"));
117    }
118}
rig_agent/agent/tool.rs

rig_agent/agent/
tool.rs
