Skip to main content

langchain_rust/tools/duckduckgo/
duckduckgo_search.rs

1use std::{collections::HashMap, error::Error};
2
3use async_trait::async_trait;
4use reqwest::Client;
5use scraper::{Html, Selector};
6use serde::{Deserialize, Serialize};
7use serde_json::{json, Value};
8use url::Url;
9
10use crate::tools::Tool;
11
12pub struct DuckDuckGoSearchResults {
13    url: String,
14    client: Client,
15    max_results: usize,
16}
17
18impl DuckDuckGoSearchResults {
19    pub fn new() -> Self {
20        Self {
21            client: Client::new(),
22            url: "https://duckduckgo.com/html/".to_string(),
23            max_results: 4,
24        }
25    }
26
27    pub fn with_max_results(mut self, max_results: usize) -> Self {
28        self.max_results = max_results;
29        self
30    }
31
32    pub async fn search(&self, query: &str) -> Result<String, Box<dyn Error>> {
33        let mut url = Url::parse(&self.url)?;
34
35        let mut query_params = HashMap::new();
36        query_params.insert("q", query);
37
38        url.query_pairs_mut().extend_pairs(query_params.iter());
39
40        let response = self.client.get(url).send().await?;
41        let body = response.text().await?;
42        let document = Html::parse_document(&body);
43
44        let result_selector = Selector::parse(".web-result").unwrap();
45        let result_title_selector = Selector::parse(".result__a").unwrap();
46        let result_url_selector = Selector::parse(".result__url").unwrap();
47        let result_snippet_selector = Selector::parse(".result__snippet").unwrap();
48
49        let results = document
50            .select(&result_selector)
51            .map(|result| {
52                let title = result
53                    .select(&result_title_selector)
54                    .next()
55                    .unwrap()
56                    .text()
57                    .collect::<Vec<_>>()
58                    .join("");
59                let link = result
60                    .select(&result_url_selector)
61                    .next()
62                    .unwrap()
63                    .text()
64                    .collect::<Vec<_>>()
65                    .join("")
66                    .trim()
67                    .to_string();
68                let snippet = result
69                    .select(&result_snippet_selector)
70                    .next()
71                    .unwrap()
72                    .text()
73                    .collect::<Vec<_>>()
74                    .join("");
75
76                SearchResult {
77                    title,
78                    link,
79                    snippet,
80                }
81            })
82            .take(self.max_results)
83            .collect::<Vec<_>>();
84
85        Ok(serde_json::to_string(&results)?)
86    }
87}
88
89#[derive(Debug, Clone, Serialize, Deserialize)]
90pub struct SearchResult {
91    title: String,
92    link: String,
93    snippet: String,
94}
95
96#[async_trait]
97impl Tool for DuckDuckGoSearchResults {
98    fn name(&self) -> String {
99        String::from("DuckDuckGoSearch")
100    }
101
102    fn description(&self) -> String {
103        String::from(
104            r#""Wrapper for DuckDuckGo Search API. "
105	"Useful for when you need to answer questions about current events. "
106	"Always one of the first options when you need to find information on internet"
107	"Input should be a search query. Output is a JSON array of the query results."#,
108        )
109    }
110
111    async fn run(&self, input: Value) -> Result<String, Box<dyn Error>> {
112        let input = input.as_str().ok_or("Input should be a string")?;
113        self.search(input).await
114    }
115
116    fn parameters(&self) -> Value {
117        let prompt = r#"A wrapper around DuckDuckGo Search.
118            Useful for when you need to answer questions about current events.
119            Input should be a search query. Output is a JSON array of the query results."#;
120
121        json!({
122            "description": prompt,
123            "type": "object",
124            "properties": {
125                "query": {
126                    "type": "string",
127                    "description": "Search query to look up"
128                }
129            },
130            "required": ["query"]
131        })
132    }
133}
134
135impl Default for DuckDuckGoSearchResults {
136    fn default() -> DuckDuckGoSearchResults {
137        DuckDuckGoSearchResults::new()
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::DuckDuckGoSearchResults;
144
145    #[tokio::test]
146    #[ignore]
147    async fn duckduckgosearch_tool() {
148        let ddg = DuckDuckGoSearchResults::default().with_max_results(5);
149        let s = ddg
150            .search("Who is the current President of Peru?")
151            .await
152            .unwrap();
153
154        println!("{}", s);
155    }
156}