langchain_rust/tools/serpapi/
serpapi.rs

1use std::error::Error;
2
3use async_trait::async_trait;
4use serde_json::Value;
5
6use crate::tools::Tool;
7
8pub struct SerpApi {
9    api_key: String,
10    location: Option<String>,
11    hl: Option<String>,
12    gl: Option<String>,
13    google_domain: Option<String>,
14}
15
16impl SerpApi {
17    pub fn new(api_key: String) -> Self {
18        Self {
19            api_key,
20            location: None,
21            hl: None,
22            gl: None,
23            google_domain: None,
24        }
25    }
26    pub fn with_location<S: Into<String>>(mut self, location: S) -> Self {
27        self.location = Some(location.into());
28        self
29    }
30    pub fn with_hl<S: Into<String>>(mut self, hl: S) -> Self {
31        self.hl = Some(hl.into());
32        self
33    }
34    pub fn with_gl(mut self, gl: String) -> Self {
35        self.gl = Some(gl);
36        self
37    }
38    pub fn with_google_domain<S: Into<String>>(mut self, google_domain: S) -> Self {
39        self.google_domain = Some(google_domain.into());
40        self
41    }
42
43    pub fn with_api_key<S: Into<String>>(mut self, api_key: S) -> Self {
44        self.api_key = api_key.into();
45        self
46    }
47
48    pub async fn simple_search(&self, query: &str) -> Result<String, Box<dyn Error>> {
49        let mut url = format!(
50            "https://serpapi.com/search.json?q={}&api_key={}",
51            query, self.api_key
52        );
53        if let Some(location) = &self.location {
54            url.push_str(&format!("&location={}", location));
55        }
56        if let Some(hl) = &self.hl {
57            url.push_str(&format!("&hl={}", hl));
58        }
59        if let Some(gl) = &self.gl {
60            url.push_str(&format!("&gl={}", gl));
61        }
62        if let Some(google_domain) = &self.google_domain {
63            url.push_str(&format!("&google_domain={}", google_domain));
64        }
65        let results: Value = reqwest::get(&url).await?.json().await?;
66
67        let res = process_response(&results)?;
68
69        Ok(res)
70    }
71}
72
73fn get_answer_box(result: &Value) -> String {
74    if let Some(map) = result["answer_box"].as_object() {
75        if let Some(answer) = map.get("answer").and_then(|v| v.as_str()) {
76            return answer.to_string();
77        }
78
79        if let Some(snippet) = map.get("snippet").and_then(|v| v.as_str()) {
80            return snippet.to_string();
81        }
82
83        if let Some(snippet) = map
84            .get("snippet_highlighted_words")
85            .and_then(|v| v.as_array())
86        {
87            if !snippet.is_empty() {
88                if let Some(first) = snippet.first().and_then(|v| v.as_str()) {
89                    return first.to_string();
90                }
91            }
92        }
93    }
94
95    "".to_string()
96}
97
98fn process_response(res: &Value) -> Result<String, Box<dyn Error>> {
99    if !get_answer_box(res).is_empty() {
100        return Ok(get_answer_box(res));
101    }
102    if !get_sport_result(res).is_empty() {
103        return Ok(get_sport_result(res));
104    }
105    if !get_knowledge_graph(res).is_empty() {
106        return Ok(get_knowledge_graph(res));
107    }
108    if !get_organic_result(res).is_empty() {
109        return Ok(get_organic_result(res));
110    }
111    Err("No good result".into())
112}
113
114fn get_sport_result(result: &Value) -> String {
115    if let Some(map) = result["sports_results"].as_object() {
116        if let Some(game_spotlight) = map.get("game_spotlight").and_then(|v| v.as_str()) {
117            return game_spotlight.to_string();
118        }
119    }
120
121    "".to_string()
122}
123
124fn get_knowledge_graph(result: &Value) -> String {
125    if let Some(map) = result["knowledge_graph"].as_object() {
126        if let Some(description) = map.get("description").and_then(|v| v.as_str()) {
127            return description.to_string();
128        }
129    }
130
131    "".to_string()
132}
133
134fn get_organic_result(result: &Value) -> String {
135    if let Some(array) = result["organic_results"].as_array() {
136        if !array.is_empty() {
137            if let Some(first) = array.first() {
138                if let Some(first_map) = first.as_object() {
139                    if let Some(snippet) = first_map.get("snippet").and_then(|v| v.as_str()) {
140                        return snippet.to_string();
141                    }
142                }
143            }
144        }
145    }
146
147    "".to_string()
148}
149
150#[async_trait]
151impl Tool for SerpApi {
152    fn name(&self) -> String {
153        String::from("GoogleSearch")
154    }
155    fn description(&self) -> String {
156        String::from(
157            r#""A wrapper around Google Search. "
158	"Useful for when you need to answer questions about current events. "
159	"Always one of the first options when you need to find information on internet"
160	"Input should be a search query."#,
161        )
162    }
163
164    async fn run(&self, input: Value) -> Result<String, Box<dyn Error>> {
165        let input = input.as_str().ok_or("Input should be a string")?;
166        self.simple_search(input).await
167    }
168}
169
170impl Default for SerpApi {
171    fn default() -> SerpApi {
172        SerpApi {
173            api_key: std::env::var("SERPAPI_API_KEY").unwrap_or_default(),
174            location: None,
175            hl: None,
176            gl: None,
177            google_domain: None,
178        }
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::SerpApi;
185
186    #[tokio::test]
187    #[ignore]
188    async fn serpapi_tool() {
189        let serpapi = SerpApi::default();
190        let s = serpapi
191            .simple_search("Who is the President of Peru")
192            .await
193            .unwrap();
194        println!("{}", s);
195    }
196}