rig_agent/agent/
tool.rs

1use std::sync::Arc;
2
3use crate::{
4    agent::Agent,
5    completion::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 Agent {
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            .config
29            .name
30            .clone()
31            .unwrap_or_else(|| DEFAULT_AGENT_TOOL_NAME.to_string());
32        let description = format!(
33            "
34            Prompt a sub-agent to do a task for you.
35
36            Agent name: {name}
37            Agent description: {description}
38            Agent system prompt: {sysprompt}
39            ",
40            name = name,
41            description = self.config.description.clone().unwrap_or_default(),
42            sysprompt = self.config.preamble.clone().unwrap_or_default()
43        );
44        let parameters = json!(schema_for!(AgentToolArgs));
45        let agent = Arc::new(self);
46
47        DynamicTool::new(name, description, parameters, move |context, args| {
48            let agent = Arc::clone(&agent);
49            let inherited_context = context.for_dispatch();
50            Box::pin(async move {
51                let args: AgentToolArgs = serde_json::from_value(args).map_err(|error| {
52                    ToolExecutionError::invalid_args(format!(
53                        "failed to parse agent tool arguments: {error}"
54                    ))
55                    .with_source(error)
56                })?;
57                agent
58                    .prompt(args.prompt)
59                    .tool_context(inherited_context)
60                    .await
61                    .map(ToolOutput::text)
62                    .map_err(ToolExecutionError::from_error)
63            })
64        })
65    }
66}
67
68impl From<Agent> for DynamicTool {
69    fn from(agent: Agent) -> Self {
70        agent.into_tool()
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77    use crate::agent::AgentBuilder;
78    use crate::test_utils::{MockCompletionModel, MockContextProbeTool, MockTurn, SessionId};
79    use crate::tool::ToolContext;
80
81    /// A `ToolContext` set on the outer run propagates into a sub-agent
82    /// invoked as a tool, so the inner agent's own tools observe it.
83    #[tokio::test]
84    async fn context_propagates_into_sub_agent() {
85        // Inner agent: calls a context-probing tool, then answers.
86        let probe = MockContextProbeTool::default();
87        let inner_model = MockCompletionModel::new([
88            MockTurn::tool_call("c1", "context_probe", json!({})),
89            MockTurn::text("inner done"),
90        ]);
91        let inner = AgentBuilder::new(inner_model)
92            .name("researcher")
93            .tool(probe.clone())
94            .build();
95
96        // Outer agent: delegates to the inner agent (registered as the
97        // "researcher" tool), then answers.
98        let outer_model = MockCompletionModel::new([
99            MockTurn::tool_call("c2", "researcher", json!({"prompt": "do research"})),
100            MockTurn::text("outer done"),
101        ]);
102        let outer = AgentBuilder::new(outer_model)
103            .dynamic_tool(inner.into_tool())
104            .build();
105
106        let mut context = ToolContext::new();
107        context.insert(SessionId("abc-123".to_string()));
108
109        let out = outer
110            .prompt("start")
111            .tool_context(context)
112            .max_turns(5)
113            .await
114            .expect("run succeeds");
115
116        assert_eq!(out, "outer done");
117        assert_eq!(probe.observed().as_deref(), Some("session:abc-123"));
118    }
119}
rig_agent/agent/tool.rs

rig_agent/agent/
tool.rs
