Skip to main content

web_search/providers/
google.rs

1//! Google search provider
2
3use async_trait::async_trait;
4use serde::Deserialize;
5
6use super::base::{SearchOptions, SearchProvider, SearchResult};
7use crate::error::SearchError;
8
9/// Google Custom Search API response
10#[derive(Debug, Deserialize)]
11struct GoogleApiResponse {
12    items: Option<Vec<GoogleApiItem>>,
13}
14
15#[derive(Debug, Deserialize)]
16struct GoogleApiItem {
17    title: String,
18    link: String,
19    snippet: Option<String>,
20}
21
22/// Configuration for Google provider
23#[derive(Debug, Clone, Default)]
24pub struct GoogleConfig {
25    /// Google Custom Search API key
26    pub api_key: Option<String>,
27    /// Google Custom Search Engine ID
28    pub search_engine_id: Option<String>,
29}
30
31/// Google search provider
32pub struct GoogleProvider {
33    name: String,
34    enabled: bool,
35    weight: f64,
36    client: reqwest::Client,
37    config: GoogleConfig,
38    api_url: String,
39}
40
41impl GoogleProvider {
42    /// Create a new Google provider with optional API credentials
43    pub fn new(config: GoogleConfig) -> Self {
44        Self {
45            name: "google".to_string(),
46            enabled: true,
47            weight: 1.0,
48            client: reqwest::Client::builder()
49                .user_agent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36")
50                .build()
51                .expect("Failed to create HTTP client"),
52            config,
53            api_url: "https://www.googleapis.com/customsearch/v1".to_string(),
54        }
55    }
56
57    /// Create a new Google provider from environment variables
58    pub fn from_env() -> Self {
59        Self::new(GoogleConfig {
60            api_key: std::env::var("GOOGLE_API_KEY").ok(),
61            search_engine_id: std::env::var("GOOGLE_CX").ok(),
62        })
63    }
64
65    /// Check if API credentials are configured
66    pub fn has_api_credentials(&self) -> bool {
67        self.config.api_key.is_some() && self.config.search_engine_id.is_some()
68    }
69
70    async fn search_with_api(
71        &self,
72        query: &str,
73        options: &SearchOptions,
74    ) -> Result<Vec<SearchResult>, SearchError> {
75        let api_key = self.config.api_key.as_ref().unwrap();
76        let cx = self.config.search_engine_id.as_ref().unwrap();
77        let limit = options.limit.unwrap_or(10).min(10);
78
79        let mut url = format!(
80            "{}?key={}&cx={}&q={}&num={}",
81            self.api_url,
82            api_key,
83            cx,
84            urlencoding::encode(query),
85            limit
86        );
87
88        if let Some(ref lang) = options.language {
89            url.push_str(&format!("&lr=lang_{}", lang));
90        }
91
92        if let Some(ref region) = options.region {
93            url.push_str(&format!("&gl={}", region));
94        }
95
96        if let Some(safe) = options.safe_search {
97            url.push_str(&format!("&safe={}", if safe { "active" } else { "off" }));
98        }
99
100        let response = self.client.get(&url).send().await?;
101
102        if !response.status().is_success() {
103            let error_text = response.text().await.unwrap_or_default();
104            return Err(SearchError::ApiError {
105                provider: self.name.clone(),
106                message: error_text,
107            });
108        }
109
110        let api_response: GoogleApiResponse = response.json().await?;
111
112        let results = api_response
113            .items
114            .unwrap_or_default()
115            .into_iter()
116            .enumerate()
117            .map(|(i, item)| SearchResult {
118                title: item.title,
119                url: item.link,
120                snippet: item.snippet.unwrap_or_default(),
121                source: self.name.clone(),
122                rank: i + 1,
123                score: None,
124                sources: None,
125            })
126            .collect();
127
128        Ok(results)
129    }
130
131    async fn search_with_scraping(
132        &self,
133        query: &str,
134        options: &SearchOptions,
135    ) -> Result<Vec<SearchResult>, SearchError> {
136        let limit = options.limit.unwrap_or(10);
137        let mut url = format!(
138            "https://www.google.com/search?q={}&num={}",
139            urlencoding::encode(query),
140            limit.min(20)
141        );
142
143        if let Some(ref lang) = options.language {
144            url.push_str(&format!("&hl={}", lang));
145        }
146
147        if let Some(ref region) = options.region {
148            url.push_str(&format!("&gl={}", region));
149        }
150
151        let response = self
152            .client
153            .get(&url)
154            .header(
155                "Accept",
156                "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
157            )
158            .header("Accept-Language", "en-US,en;q=0.5")
159            .send()
160            .await?;
161
162        if !response.status().is_success() {
163            return Err(SearchError::ApiError {
164                provider: self.name.clone(),
165                message: format!("HTTP {}", response.status()),
166            });
167        }
168
169        let html = response.text().await?;
170        Ok(self.parse_scraped_results(&html, limit))
171    }
172
173    fn parse_scraped_results(&self, html: &str, limit: usize) -> Vec<SearchResult> {
174        use scraper::{Html, Selector};
175
176        let document = Html::parse_document(html);
177        let mut results = Vec::new();
178        let mut seen_urls = std::collections::HashSet::new();
179
180        let link_selector = Selector::parse("a").unwrap();
181        let h3_selector = Selector::parse("h3").unwrap();
182
183        for element in document.select(&link_selector) {
184            if results.len() >= limit {
185                break;
186            }
187
188            let href = element.value().attr("href").unwrap_or_default();
189            let url = if href.starts_with("/url?q=") {
190                href.strip_prefix("/url?q=")
191                    .and_then(|u| u.split('&').next())
192                    .map(|u| {
193                        urlencoding::decode(u)
194                            .unwrap_or_else(|_| u.into())
195                            .to_string()
196                    })
197            } else if href.starts_with("http") && !href.contains("google.com") {
198                Some(href.to_string())
199            } else {
200                None
201            };
202
203            if let Some(url) = url {
204                if seen_urls.contains(&url) || url.contains("google.com") {
205                    continue;
206                }
207
208                let title = element
209                    .select(&h3_selector)
210                    .next()
211                    .map(|h3| h3.text().collect::<String>())
212                    .unwrap_or_default()
213                    .trim()
214                    .to_string();
215
216                if title.is_empty() {
217                    continue;
218                }
219
220                seen_urls.insert(url.clone());
221                results.push(SearchResult {
222                    title,
223                    url,
224                    snippet: String::new(),
225                    source: self.name.clone(),
226                    rank: results.len() + 1,
227                    score: None,
228                    sources: None,
229                });
230            }
231        }
232
233        results
234    }
235}
236
237impl Default for GoogleProvider {
238    fn default() -> Self {
239        Self::from_env()
240    }
241}
242
243#[async_trait]
244impl SearchProvider for GoogleProvider {
245    fn name(&self) -> &str {
246        &self.name
247    }
248
249    fn is_available(&self) -> bool {
250        self.enabled
251    }
252
253    fn weight(&self) -> f64 {
254        self.weight
255    }
256
257    fn set_weight(&mut self, weight: f64) {
258        self.weight = weight.clamp(0.0, 1.0);
259    }
260
261    fn set_enabled(&mut self, enabled: bool) {
262        self.enabled = enabled;
263    }
264
265    async fn search(
266        &self,
267        query: &str,
268        options: &SearchOptions,
269    ) -> Result<Vec<SearchResult>, SearchError> {
270        if query.is_empty() {
271            return Ok(Vec::new());
272        }
273
274        if self.has_api_credentials() {
275            match self.search_with_api(query, options).await {
276                Ok(results) => return Ok(results),
277                Err(e) => {
278                    tracing::warn!("Google API search failed, falling back to scraping: {}", e);
279                }
280            }
281        }
282
283        self.search_with_scraping(query, options).await
284    }
285}