agent_base/tool/
mod.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use serde_json::{json, Value};
6use tokio::sync::broadcast;
7
8use crate::llm::LlmClient;
9use crate::types::{AgentResult, AgentError, AgentEvent, SessionId};
10use crate::engine::SessionStore;
11
12pub mod auto_continue;
13pub mod policy;
14pub mod subagent;
15
16pub use auto_continue::AutoContinueTool;
17pub use subagent::{SubAgentSessionPolicy, SubAgentTool};
18
19pub use policy::ToolPolicy;
20
21#[derive(Clone, Debug, Default)]
22pub struct ToolOutput {
23    pub summary: String,
24    pub raw: Option<Value>,
25    pub control_flow: ToolControlFlow,
26    pub truncation: Option<TruncationInfo>,
27}
28
29#[derive(Clone, Debug)]
30pub struct TruncationInfo {
31    pub original_summary_len: usize,
32    pub original_raw_len: Option<usize>,
33    pub max_allowed_chars: usize,
34}
35
36#[derive(Clone, Debug, Default)]
37pub enum ToolControlFlow {
38    #[default]
39    Break,
40    Continue,
41}
42
43#[derive(Clone)]
44pub struct ToolContext {
45    pub session_id: SessionId,
46    pub event_bus: broadcast::Sender<AgentEvent>,
47    pub llm_client: Option<Arc<dyn LlmClient>>,
48    pub session_store: Option<Arc<dyn SessionStore>>,
49    /// Language preference for tool output.
50    /// Defaults to `Language::En` if not set.
51    pub language: crate::types::Language,
52}
53
54#[async_trait]
55pub trait Tool: Send + Sync {
56    fn name(&self) -> &'static str;
57    fn definition(&self) -> Value;
58    async fn call(&self, args: &Value, ctx: &ToolContext) -> AgentResult<ToolOutput>;
59}
60
61#[async_trait]
62pub trait TypedTool: Send + Sync {
63    type Args: serde::de::DeserializeOwned;
64    type Output: serde::Serialize;
65
66    fn name(&self) -> &'static str;
67    fn description(&self) -> &'static str;
68    fn parameters_schema(&self) -> Value;
69    async fn call_typed(&self, args: Self::Args, ctx: &ToolContext) -> AgentResult<Self::Output>;
70
71    fn control_flow() -> ToolControlFlow
72    where
73        Self: Sized,
74    {
75        ToolControlFlow::Break
76    }
77
78    fn format_output(&self, output: Self::Output) -> String {
79        serde_json::to_string(&output).unwrap_or_default()
80    }
81}
82
83#[async_trait]
84impl<T: TypedTool + Send + Sync + 'static> Tool for T {
85    fn name(&self) -> &'static str {
86        TypedTool::name(self)
87    }
88
89    fn definition(&self) -> Value {
90        json!({
91            "type": "function",
92            "function": {
93                "name": self.name(),
94                "description": self.description(),
95                "parameters": self.parameters_schema(),
96            }
97        })
98    }
99
100    async fn call(&self, args: &Value, ctx: &ToolContext) -> AgentResult<ToolOutput> {
101        let typed_args: T::Args = serde_json::from_value(args.clone())
102            .map_err(|_| AgentError::ToolArgsInvalid {
103                name: self.name().to_string(),
104                raw: args.to_string(),
105            })?;
106        let output = self.call_typed(typed_args, ctx).await?;
107        let output_json = serde_json::to_value(&output).ok();
108        let summary = self.format_output(output);
109        Ok(ToolOutput {
110            summary,
111            raw: output_json,
112            control_flow: T::control_flow(),
113            truncation: None,
114        })
115    }
116}
117
118pub(crate) type ToolRef = Arc<dyn Tool>;
119
120#[derive(Clone, Default)]
121pub struct ToolRegistry {
122    tools: HashMap<String, ToolRef>,
123}
124
125impl ToolRegistry {
126    pub fn register(&mut self, tool: impl Tool + 'static) {
127        self.tools.insert(tool.name().to_string(), Arc::new(tool));
128    }
129
130    pub fn register_arc(&mut self, tool: Arc<dyn Tool>) {
131        self.tools.insert(tool.name().to_string(), tool);
132    }
133
134    pub fn update(&mut self, tool: impl Tool + 'static) {
135        self.tools.insert(tool.name().to_string(), Arc::new(tool));
136    }
137
138    pub fn get(&self, name: &str) -> Option<ToolRef> {
139        self.tools.get(name).cloned()
140    }
141
142    pub fn definitions(&self) -> Vec<Value> {
143        self.tools.values().map(|tool| tool.definition()).collect()
144    }
145
146    pub fn len(&self) -> usize {
147        self.tools.len()
148    }
149
150    pub fn is_empty(&self) -> bool {
151        self.tools.is_empty()
152    }
153}
agent_base/tool/mod.rs

agent_base/tool/
mod.rs
