langchain_rust/tools/serpapi/
serpapi.rs1use 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}