Skip to content

Commit e4925ee

Browse files
committed
fix(engine): apply web search domain filters
1 parent 7954d02 commit e4925ee

1 file changed

Lines changed: 86 additions & 2 deletions

File tree

src/cortex-engine/src/tools/handlers/web_search.rs

Lines changed: 86 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,10 +110,12 @@ impl ToolHandler for WebSearchHandler {
110110
));
111111
}
112112

113+
let query = build_search_query(&args);
114+
113115
// Build DuckDuckGo Instant Answer API URL with optional parameters
114116
let mut url = format!(
115117
"https://api.duckduckgo.com/?q={}&format=json&no_redirect=1&no_html=1&skip_disambig=1",
116-
urlencoding::encode(&args.query)
118+
urlencoding::encode(&query)
117119
);
118120

119121
// Add include_text parameter if requested
@@ -185,7 +187,7 @@ impl ToolHandler for WebSearchHandler {
185187
Ok(ToolResult::success(format!(
186188
"No immediate results found for '{}'. You may want to try a more specific query or search directly at https://duckduckgo.com/?q={}",
187189
args.query,
188-
urlencoding::encode(&args.query)
190+
urlencoding::encode(&query)
189191
)))
190192
} else {
191193
Ok(ToolResult::success(format!(
@@ -197,6 +199,33 @@ impl ToolHandler for WebSearchHandler {
197199
}
198200
}
199201

202+
fn build_search_query(args: &WebQueryArgs) -> String {
203+
let mut query = args.query.clone();
204+
205+
if let Some(domains) = &args.include_domains {
206+
append_domain_filters(&mut query, "site:", domains);
207+
}
208+
209+
if let Some(domains) = &args.exclude_domains {
210+
append_domain_filters(&mut query, "-site:", domains);
211+
}
212+
213+
query
214+
}
215+
216+
fn append_domain_filters(query: &mut String, prefix: &str, domains: &[String]) {
217+
for domain in domains {
218+
let domain = domain.trim();
219+
if domain.is_empty() {
220+
continue;
221+
}
222+
223+
query.push(' ');
224+
query.push_str(prefix);
225+
query.push_str(domain);
226+
}
227+
}
228+
200229
mod urlencoding {
201230
pub fn encode(s: &str) -> String {
202231
let mut result = String::new();
@@ -259,4 +288,59 @@ mod tests {
259288
let tool_result = result.unwrap();
260289
assert!(!tool_result.success);
261290
}
291+
292+
#[test]
293+
fn test_build_search_query_with_include_domains() {
294+
let args = WebQueryArgs {
295+
query: "rust async".to_string(),
296+
search_type: default_search_type(),
297+
category: None,
298+
num_results: default_num_results(),
299+
include_domains: Some(vec!["docs.rs".to_string(), "rust-lang.org".to_string()]),
300+
exclude_domains: None,
301+
include_text: false,
302+
};
303+
304+
assert_eq!(
305+
build_search_query(&args),
306+
"rust async site:docs.rs site:rust-lang.org"
307+
);
308+
}
309+
310+
#[test]
311+
fn test_build_search_query_with_exclude_domains() {
312+
let args = WebQueryArgs {
313+
query: "rust async".to_string(),
314+
search_type: default_search_type(),
315+
category: None,
316+
num_results: default_num_results(),
317+
include_domains: None,
318+
exclude_domains: Some(vec!["example.com".to_string(), "spam.test".to_string()]),
319+
include_text: false,
320+
};
321+
322+
assert_eq!(
323+
build_search_query(&args),
324+
"rust async -site:example.com -site:spam.test"
325+
);
326+
}
327+
328+
#[test]
329+
fn test_build_search_query_ignores_blank_domains() {
330+
let args = WebQueryArgs {
331+
query: "rust async".to_string(),
332+
search_type: default_search_type(),
333+
category: None,
334+
num_results: default_num_results(),
335+
include_domains: Some(vec![
336+
" docs.rs ".to_string(),
337+
"".to_string(),
338+
" ".to_string(),
339+
]),
340+
exclude_domains: None,
341+
include_text: false,
342+
};
343+
344+
assert_eq!(build_search_query(&args), "rust async site:docs.rs");
345+
}
262346
}

0 commit comments

Comments
 (0)