llm_coding_tools_rig/
webfetch.rs1use llm_coding_tools_core::operations::fetch_url;
6use llm_coding_tools_core::tool_names;
7use llm_coding_tools_core::{ToolContext, ToolError, ToolOutput};
8use rig::completion::ToolDefinition;
9use rig::tool::Tool;
10use schemars::{schema_for, JsonSchema};
11use serde::Deserialize;
12use std::time::Duration;
13
14const DEFAULT_TIMEOUT_MS: u64 = 30_000;
16
17fn default_timeout_ms() -> u64 {
18 DEFAULT_TIMEOUT_MS
19}
20
21#[derive(Debug, Clone, Deserialize, JsonSchema)]
23pub struct WebFetchArgs {
24 pub url: String,
26 #[serde(default = "default_timeout_ms")]
28 pub timeout_ms: u64,
29}
30
31#[derive(Debug, Clone)]
37pub struct WebFetchTool {
38 client: reqwest::Client,
39}
40
41impl Default for WebFetchTool {
42 fn default() -> Self {
43 Self::new()
44 }
45}
46
47impl WebFetchTool {
48 pub fn new() -> Self {
50 Self {
51 client: reqwest::Client::new(),
52 }
53 }
54
55 pub fn with_client(client: reqwest::Client) -> Self {
57 Self { client }
58 }
59}
60
61impl Tool for WebFetchTool {
62 const NAME: &'static str = tool_names::WEBFETCH;
63
64 type Error = ToolError;
65 type Args = WebFetchArgs;
66 type Output = ToolOutput;
67
68 async fn definition(&self, _prompt: String) -> ToolDefinition {
69 ToolDefinition {
70 name: <Self as Tool>::NAME.to_string(),
71 description:
72 "Fetch content from a URL. HTML is converted to markdown, JSON is prettified."
73 .to_string(),
74 parameters: serde_json::to_value(schema_for!(WebFetchArgs))
75 .expect("schema serialization should never fail"),
76 }
77 }
78
79 async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
80 let timeout = Duration::from_millis(args.timeout_ms);
81 let result = fetch_url(&self.client, &args.url, timeout).await?;
82 Ok(result.into())
83 }
84}
85
86impl ToolContext for WebFetchTool {
87 const NAME: &'static str = tool_names::WEBFETCH;
88
89 fn context(&self) -> &'static str {
90 llm_coding_tools_core::context::WEBFETCH
91 }
92}
93
94#[cfg(test)]
95mod tests {
96 use super::*;
97
98 #[test]
99 fn creates_with_default_client() {
100 let _tool = WebFetchTool::new();
101 }
102
103 #[test]
104 fn creates_with_custom_client() {
105 let client = reqwest::Client::builder()
106 .user_agent("test")
107 .build()
108 .unwrap();
109 let _tool = WebFetchTool::with_client(client);
110 }
111}