web_search/providers/
google.rs

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

web_search/providers/
google.rs
