Skip to main content

golem_ai_web_search_serper/
lib.rs

1mod client;
2pub mod config;
3mod conversions;
4
5use std::cell::RefCell;
6
7use crate::client::SerperSearchApi;
8use crate::conversions::{params_to_request, response_to_results, validate_search_params};
9use golem_ai_web_search::durability::DurableWebSearch;
10use golem_ai_web_search::durability::ExtendedWebSearchProvider;
11use golem_ai_web_search::model::web_search::{
12    SearchError, SearchMetadata, SearchParams, SearchResult, SearchSession,
13};
14use golem_ai_web_search::{SearchSessionInterface, WebSearchProvider};
15
16pub use config::SerperConfig;
17#[cfg(feature = "golem")]
18pub use config::SerperHostConfig;
19
20#[cfg(feature = "golem")]
21#[derive(Debug, Clone, PartialEq, golem_rust::FromValueAndType, golem_rust::IntoValue)]
22pub struct SerperReplayState {
23    pub current_page: u32,
24    pub metadata: Option<SearchMetadata>,
25    pub finished: bool,
26}
27
28struct SerperSearchSessionImpl {
29    client: SerperSearchApi,
30    params: SearchParams,
31    metadata: Option<SearchMetadata>,
32    current_page: u32, // 1-based
33    finished: bool,
34}
35
36impl SerperSearchSessionImpl {
37    fn new(client: SerperSearchApi, params: SearchParams) -> Self {
38        Self {
39            client,
40            params,
41            metadata: None,
42            current_page: 1, // 1-based
43            finished: false,
44        }
45    }
46
47    fn next_page(&mut self) -> Result<Vec<SearchResult>, SearchError> {
48        if self.finished {
49            return Ok(Vec::new());
50        }
51
52        let request = params_to_request(self.params.clone(), self.current_page)?;
53        let num_results = request.num.unwrap_or(10);
54        let response = self.client.search(request)?;
55        let (results, metadata) = response_to_results(response, &self.params, self.current_page);
56
57        self.finished = results.len() < (num_results as usize);
58        self.current_page += 1;
59        self.metadata = Some(metadata);
60
61        Ok(results)
62    }
63
64    fn get_metadata(&self) -> Option<SearchMetadata> {
65        self.metadata.clone()
66    }
67}
68
69// Create a wrapper that implements GuestSearchSession properly
70pub struct SerperSearchSession(RefCell<SerperSearchSessionImpl>);
71
72impl SerperSearchSession {
73    fn new(search: SerperSearchSessionImpl) -> Self {
74        Self(RefCell::new(search))
75    }
76}
77
78impl SearchSessionInterface for SerperSearchSession {
79    fn as_any(&self) -> &dyn std::any::Any {
80        self
81    }
82
83    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
84        self
85    }
86
87    fn next_page(&self) -> Result<Vec<SearchResult>, SearchError> {
88        let mut search = self.0.borrow_mut();
89        search.next_page()
90    }
91
92    fn get_metadata(&self) -> Option<SearchMetadata> {
93        let search = self.0.borrow();
94        search.get_metadata()
95    }
96}
97
98pub struct SerperSearch;
99
100impl SerperSearch {
101    fn execute_search(
102        provider_config: &SerperConfig,
103        params: SearchParams,
104    ) -> Result<(Vec<SearchResult>, SearchMetadata), SearchError> {
105        validate_search_params(&params)?;
106
107        let client = SerperSearchApi::new(provider_config);
108        let request = params_to_request(params.clone(), 1)?;
109
110        let response = client.search(request)?;
111        let (results, metadata) = response_to_results(response, &params, 1);
112
113        Ok((results, metadata))
114    }
115
116    fn start_search_session(
117        provider_config: &SerperConfig,
118        params: SearchParams,
119    ) -> Result<SerperSearchSession, SearchError> {
120        validate_search_params(&params)?;
121
122        let client = SerperSearchApi::new(provider_config);
123        let search = SerperSearchSessionImpl::new(client, params);
124        Ok(SerperSearchSession::new(search))
125    }
126}
127
128impl WebSearchProvider for SerperSearch {
129    type SearchSession = SerperSearchSession;
130    type ProviderConfig = SerperConfig;
131
132    fn start_search(
133        provider_config: Self::ProviderConfig,
134        params: SearchParams,
135    ) -> Result<SearchSession, SearchError> {
136        match Self::start_search_session(&provider_config, params) {
137            Ok(session) => Ok(SearchSession::new(session)),
138            Err(err) => Err(err),
139        }
140    }
141
142    fn search_once(
143        provider_config: Self::ProviderConfig,
144        params: SearchParams,
145    ) -> Result<(Vec<SearchResult>, Option<SearchMetadata>), SearchError> {
146        let (results, metadata) = Self::execute_search(&provider_config, params)?;
147        Ok((results, Some(metadata)))
148    }
149}
150
151#[cfg(feature = "golem")]
152impl ExtendedWebSearchProvider for SerperSearch {
153    type ReplayState = SerperReplayState;
154
155    fn unwrapped_search_session(
156        provider_config: Self::ProviderConfig,
157        params: SearchParams,
158    ) -> Result<Self::SearchSession, SearchError> {
159        let client = SerperSearchApi::new(&provider_config);
160        let search = SerperSearchSessionImpl::new(client, params);
161        Ok(SerperSearchSession::new(search))
162    }
163
164    fn session_to_state(session: &Self::SearchSession) -> Self::ReplayState {
165        let search = session.0.borrow_mut();
166        SerperReplayState {
167            current_page: search.current_page,
168            metadata: search.metadata.clone(),
169            finished: search.finished,
170        }
171    }
172
173    fn session_from_state(
174        provider_config: Self::ProviderConfig,
175        state: &Self::ReplayState,
176        params: SearchParams,
177    ) -> Result<Self::SearchSession, SearchError> {
178        let client = SerperSearchApi::new(&provider_config);
179        let mut search = SerperSearchSessionImpl::new(client, params);
180        search.current_page = state.current_page;
181        search.metadata = state.metadata.clone();
182        search.finished = state.finished;
183        Ok(SerperSearchSession::new(search))
184    }
185}
186
187#[cfg(not(feature = "golem"))]
188impl ExtendedWebSearchProvider for SerperSearch {}
189
190pub type DurableSerperSearch = DurableWebSearch<SerperSearch>;