kalosm_language/context/search/
mod.rs

1#![allow(missing_docs)]
2
3use rand::seq::SliceRandom;
4use url::Url;
5
6use super::{
7    document::{Document, IntoDocuments},
8    get_article, ExtractDocumentError,
9};
10
11/// A search query that can be used to search for documents on the web.
12///
13/// # Example
14/// ```rust, no_run
15/// // You must have the SERPER_API_KEY environment variable set to run this example.
16/// use kalosm_language::prelude::*;
17///
18/// #[tokio::main]
19/// async fn main() {
20///     let query = "What is the best way to learn a language?";
21///     let api_key = std::env::var("SERPER_API_KEY").unwrap();
22///     let search_query = SearchQuery::new(query, &api_key, 5);
23///     let documents = search_query.into_documents().await.unwrap();
24///     let mut text = String::new();
25///     for document in documents {
26///         for word in document.body().split(' ').take(300) {
27///             text.push_str(word);
28///             text.push(' ');
29///         }
30///         text.push('\n');
31///     }
32///     println!("{}", text);
33/// }
34/// ```
35pub struct SearchQuery<'a> {
36    query: &'a str,
37    api_key: &'a str,
38    top: usize,
39}
40
41impl<'a> SearchQuery<'a> {
42    /// Create a new search query.
43    pub fn new(query: &'a str, api_key: &'a str, top_n: usize) -> Self {
44        Self {
45            query,
46            api_key,
47            top: top_n,
48        }
49    }
50}
51
52impl IntoDocuments for SearchQuery<'_> {
53    type Error = ExtractDocumentError;
54
55    async fn into_documents(self) -> Result<Vec<Document>, Self::Error> {
56        let mut search_results = search(self.api_key, self.query).await?;
57
58        let mut documents = vec![];
59        search_results.organic.shuffle(&mut rand::thread_rng());
60
61        for result in search_results.organic.into_iter().take(self.top) {
62            if let Some(link) = &result.link {
63                documents.push(get_article(Url::parse(link)?).await?);
64            }
65        }
66
67        Ok(documents)
68    }
69}
70
71#[derive(serde::Serialize, serde::Deserialize, Debug)]
72pub struct SearchResult {
73    pub knowledge_graph: Option<KnowledgeGraph>,
74    #[serde(default)]
75    pub organic: Vec<Organic>,
76    #[serde(default)]
77    pub people_also_ask: Vec<PeopleAlsoAsk>,
78    #[serde(default)]
79    pub related_searches: Vec<RelatedSearches>,
80}
81
82#[derive(serde::Serialize, serde::Deserialize, Debug)]
83pub struct KnowledgeGraph {
84    pub title: String,
85    pub type_: String,
86    pub website: String,
87    pub image_url: String,
88    pub description: String,
89    pub description_source: String,
90    pub description_link: String,
91    #[serde(default)]
92    pub attributes: Vec<Attributes>,
93}
94
95#[derive(serde::Serialize, serde::Deserialize, Debug)]
96pub struct Attributes {
97    pub key: String,
98    pub value: String,
99}
100
101#[derive(serde::Serialize, serde::Deserialize, Debug)]
102pub struct Organic {
103    pub title: Option<String>,
104    pub link: Option<String>,
105    #[serde(default)]
106    pub snippet: String,
107    #[serde(default)]
108    pub sitelinks: Vec<Sitelinks>,
109    pub position: u32,
110}
111
112#[derive(serde::Serialize, serde::Deserialize, Debug)]
113pub struct Sitelinks {
114    pub title: String,
115    pub link: String,
116}
117
118#[derive(serde::Serialize, serde::Deserialize, Debug)]
119pub struct PeopleAlsoAsk {
120    pub question: String,
121    pub snippet: String,
122    pub title: String,
123    pub link: String,
124}
125
126#[derive(serde::Serialize, serde::Deserialize, Debug)]
127pub struct RelatedSearches {
128    pub query: String,
129}
130
131pub async fn search(api_key: &str, query: &str) -> Result<SearchResult, reqwest::Error> {
132    let url = Url::parse("https://google.serper.dev/search").unwrap();
133    let client = reqwest::Client::new();
134    let res = client
135        .post(url)
136        .header("X-API-KEY", api_key)
137        .json(&serde_json::json!({
138            "q": query
139        }))
140        .send()
141        .await
142        .unwrap();
143    res.json().await
144}
145
146#[tokio::test]
147async fn search_result() {
148    if let Some(key) = option_env!("SERPER_API_KEY") {
149        let result = search(key, "apple inc").await.unwrap();
150        println!("{:#?}", result);
151    }
152}
kalosm_language/context/search/mod.rs

kalosm_language/context/search/
mod.rs
