llm_coding_tools_rig/
webfetch.rs

1//! Web content fetching tool.
2//!
3//! Provides URL fetching with format conversion support.
4
5use 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
14/// Default timeout: 30 seconds.
15const DEFAULT_TIMEOUT_MS: u64 = 30_000;
16
17fn default_timeout_ms() -> u64 {
18    DEFAULT_TIMEOUT_MS
19}
20
21/// Arguments for the webfetch tool.
22#[derive(Debug, Clone, Deserialize, JsonSchema)]
23pub struct WebFetchArgs {
24    /// The URL to fetch.
25    pub url: String,
26    /// Timeout in milliseconds (default: 30000).
27    #[serde(default = "default_timeout_ms")]
28    pub timeout_ms: u64,
29}
30
31/// Tool for fetching web content.
32///
33/// - HTML is converted to markdown
34/// - JSON is pretty-printed
35/// - Other content returned as-is
36#[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    /// Creates a new webfetch tool with default client.
49    pub fn new() -> Self {
50        Self {
51            client: reqwest::Client::new(),
52        }
53    }
54
55    /// Creates a webfetch tool with a custom client.
56    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}