apollo/tools/
web_search.rs1use 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}