rs_agent/
utcp.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use rs_utcp::tools::Tool as UtcpTool;
6use rs_utcp::UtcpClientInterface;
7
8use crate::error::{AgentError, Result};
9use crate::tools::Tool;
10use crate::types::{ToolRequest, ToolResponse, ToolSpec};
11
12/// Adapter that exposes a UTCP tool through the rs-agent `Tool` trait.
13pub struct UtcpToolAdapter {
14    client: Arc<dyn UtcpClientInterface>,
15    tool: UtcpTool,
16}
17
18impl UtcpToolAdapter {
19    pub fn new(client: Arc<dyn UtcpClientInterface>, tool: UtcpTool) -> Self {
20        Self { client, tool }
21    }
22
23    fn tool_spec(&self) -> ToolSpec {
24        let input_schema = serde_json::to_value(&self.tool.inputs)
25            .unwrap_or_else(|_| serde_json::json!({"type": "object"}));
26
27        ToolSpec {
28            name: self.tool.name.clone(),
29            description: self.tool.description.clone(),
30            input_schema,
31            examples: None,
32        }
33    }
34}
35
36#[async_trait]
37impl Tool for UtcpToolAdapter {
38    fn spec(&self) -> ToolSpec {
39        self.tool_spec()
40    }
41
42    async fn invoke(&self, req: ToolRequest) -> Result<ToolResponse> {
43        // Forward invocation through the UTCP client
44        let result = self
45            .client
46            .call_tool(&self.tool.name, req.arguments)
47            .await
48            .map_err(|e| AgentError::UtcpError(e.to_string()))?;
49
50        // Preserve string outputs as-is; serialize other payloads to JSON text
51        let content = match result {
52            serde_json::Value::String(s) => s,
53            other => serde_json::to_string(&other).unwrap_or_else(|_| format!("{other:?}")),
54        };
55
56        Ok(ToolResponse {
57            content,
58            metadata: Some(HashMap::from([(
59                "provider".to_string(),
60                "utcp".to_string(),
61            )])),
62        })
63    }
64}
65
66/// Registers UTCP tools into the agent's tool catalog.
67pub fn register_utcp_tools(
68    catalog: &crate::tools::ToolCatalog,
69    client: Arc<dyn UtcpClientInterface>,
70    tools: Vec<UtcpTool>,
71) -> Result<()> {
72    for tool in tools {
73        let adapter = UtcpToolAdapter::new(client.clone(), tool);
74        catalog.register(Box::new(adapter))?;
75    }
76    Ok(())
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use crate::agent::Agent;
83    use crate::memory::{InMemoryStore, SessionMemory};
84    use crate::models::LLM;
85    use crate::types::{AgentOptions, File, GenerationResponse, Message};
86    use anyhow::anyhow;
87    use rs_utcp::providers::base::Provider;
88    use rs_utcp::tools::ToolInputOutputSchema;
89    use rs_utcp::transports::stream::StreamResult;
90    use rs_utcp::transports::CommunicationProtocol;
91    use std::sync::Mutex;
92
93    struct MockUtcpClient {
94        calls: Mutex<Vec<(String, HashMap<String, serde_json::Value>)>>,
95    }
96
97    impl MockUtcpClient {
98        fn new() -> Self {
99            Self {
100                calls: Mutex::new(Vec::new()),
101            }
102        }
103    }
104
105    #[async_trait]
106    impl UtcpClientInterface for MockUtcpClient {
107        async fn register_tool_provider(
108            &self,
109            _prov: Arc<dyn Provider>,
110        ) -> anyhow::Result<Vec<UtcpTool>> {
111            Ok(vec![])
112        }
113
114        async fn register_tool_provider_with_tools(
115            &self,
116            _prov: Arc<dyn Provider>,
117            tools: Vec<UtcpTool>,
118        ) -> anyhow::Result<Vec<UtcpTool>> {
119            Ok(tools)
120        }
121
122        async fn deregister_tool_provider(&self, _provider_name: &str) -> anyhow::Result<()> {
123            Ok(())
124        }
125
126        async fn call_tool(
127            &self,
128            tool_name: &str,
129            args: HashMap<String, serde_json::Value>,
130        ) -> anyhow::Result<serde_json::Value> {
131            self.calls
132                .lock()
133                .unwrap()
134                .push((tool_name.to_string(), args));
135            Ok(serde_json::json!({"ok": true}))
136        }
137
138        async fn search_tools(&self, _query: &str, _limit: usize) -> anyhow::Result<Vec<UtcpTool>> {
139            Ok(vec![])
140        }
141
142        fn get_transports(&self) -> HashMap<String, Arc<dyn CommunicationProtocol>> {
143            HashMap::new()
144        }
145
146        async fn call_tool_stream(
147            &self,
148            _tool_name: &str,
149            _args: HashMap<String, serde_json::Value>,
150        ) -> anyhow::Result<Box<dyn StreamResult>> {
151            Err(anyhow!("not implemented"))
152        }
153    }
154
155    struct MockLLM;
156
157    #[async_trait]
158    impl LLM for MockLLM {
159        async fn generate(
160            &self,
161            messages: Vec<Message>,
162            _files: Option<Vec<File>>,
163        ) -> Result<GenerationResponse> {
164            let last = messages.last().unwrap();
165            Ok(GenerationResponse {
166                content: format!("Echo: {}", last.content),
167                metadata: None,
168            })
169        }
170
171        fn model_name(&self) -> &str {
172            "mock"
173        }
174    }
175
176    #[tokio::test]
177    async fn registers_and_invokes_utcp_tool() {
178        let client = Arc::new(MockUtcpClient::new());
179        let memory = Arc::new(SessionMemory::new(Box::new(InMemoryStore::new()), 4));
180        let agent = Agent::new(Arc::new(MockLLM), memory, AgentOptions::default());
181
182        let tool = UtcpTool {
183            name: "dummy.echo".to_string(),
184            description: "Echo via UTCP".to_string(),
185            inputs: ToolInputOutputSchema {
186                type_: "object".to_string(),
187                properties: Some(HashMap::from([(
188                    "text".to_string(),
189                    serde_json::json!({"type": "string"}),
190                )])),
191                required: Some(vec!["text".to_string()]),
192                description: None,
193                title: None,
194                items: None,
195                enum_: None,
196                minimum: None,
197                maximum: None,
198                format: None,
199            },
200            outputs: ToolInputOutputSchema {
201                type_: "object".to_string(),
202                properties: None,
203                required: None,
204                description: None,
205                title: None,
206                items: None,
207                enum_: None,
208                minimum: None,
209                maximum: None,
210                format: None,
211            },
212            tags: vec![],
213            average_response_size: None,
214            provider: None,
215        };
216
217        register_utcp_tools(agent.tools().as_ref(), client.clone(), vec![tool]).unwrap();
218
219        let mut args = HashMap::new();
220        args.insert("text".to_string(), serde_json::json!("hello"));
221
222        let result = agent.invoke_tool("s", "dummy.echo", args).await.unwrap();
223
224        assert_eq!(result, r#"{"ok":true}"#);
225        let calls = client.calls.lock().unwrap();
226        assert_eq!(calls.len(), 1);
227        assert_eq!(calls[0].0, "dummy.echo");
228    }
229}
rs_agent/utcp.rs

rs_agent/
utcp.rs
