Skip to main content

apollo/tools/
web_search.rs

1//! Web search tool — search the web via Perplexity or Brave API.
2
3use async_trait::async_trait;
4use serde::Deserialize;
5
6use super::traits::*;
7use crate::text::truncate_chars;
8
9pub struct WebSearchTool {
10    api_key: Option<String>,
11}
12
13impl WebSearchTool {
14    pub fn new() -> Self {
15        Self {
16            api_key: std::env::var("PERPLEXITY_API_KEY").ok(),
17        }
18    }
19
20    pub fn with_api_key(mut self, key: String) -> Self {
21        self.api_key = Some(key);
22        self
23    }
24}
25
26impl Default for WebSearchTool {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32#[derive(Deserialize)]
33struct SearchArgs {
34    query: String,
35}
36
37#[async_trait]
38impl Tool for WebSearchTool {
39    fn name(&self) -> &str {
40        "web_search"
41    }
42
43    fn spec(&self) -> ToolSpec {
44        ToolSpec {
45            name: "web_search".to_string(),
46            description: "Search the web for information. Returns relevant results with snippets."
47                .to_string(),
48            parameters: serde_json::json!({
49                "type": "object",
50                "properties": {
51                    "query": {
52                        "type": "string",
53                        "description": "Search query"
54                    }
55                },
56                "required": ["query"]
57            }),
58        }
59    }
60
61    async fn execute(&self, arguments: &str) -> anyhow::Result<ToolResult> {
62        let args: SearchArgs = serde_json::from_str(arguments)?;
63
64        let api_key = match &self.api_key {
65            Some(k) => k.clone(),
66            None => return Ok(ToolResult::error("No PERPLEXITY_API_KEY set")),
67        };
68
69        let client = reqwest::Client::new();
70        let resp = client
71            .post("https://api.perplexity.ai/chat/completions")
72            .header("Authorization", format!("Bearer {}", api_key))
73            .header("Content-Type", "application/json")
74            .json(&serde_json::json!({
75                "model": "sonar",
76                "messages": [
77                    {"role": "user", "content": args.query}
78                ]
79            }))
80            .send()
81            .await?;
82
83        if !resp.status().is_success() {
84            let status = resp.status();
85            let text = resp.text().await.unwrap_or_default();
86            return Ok(ToolResult::error(format!(
87                "Search API error {}: {}",
88                status,
89                truncate_chars(&text, 200)
90            )));
91        }
92
93        let data: serde_json::Value = resp.json().await?;
94        let content = data["choices"][0]["message"]["content"]
95            .as_str()
96            .unwrap_or("No results found")
97            .to_string();
98
99        Ok(ToolResult::success(content))
100    }
101}