91 lines
2.7 KiB
Rust
91 lines
2.7 KiB
Rust
use serde::{Deserialize, Serialize};
|
|
|
|
use super::{brave, tavily};
|
|
|
|
/// Which search backend the AI tool layer should call.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum SearchProvider {
|
|
Tavily,
|
|
Brave,
|
|
}
|
|
|
|
impl SearchProvider {
|
|
/// Parse the `SEARCH_PROVIDER` environment variable into a supported variant.
|
|
pub fn from_env(value: Option<&str>) -> Result<Self, String> {
|
|
match value
|
|
.unwrap_or("tavily")
|
|
.trim()
|
|
.to_ascii_lowercase()
|
|
.as_str()
|
|
{
|
|
"tavily" => Ok(Self::Tavily),
|
|
"brave" => Ok(Self::Brave),
|
|
other => Err(format!(
|
|
"Unsupported SEARCH_PROVIDER '{}'. Expected 'tavily' or 'brave'.",
|
|
other
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Return the environment variable name required by the selected provider.
|
|
pub fn required_key_name(self) -> &'static str {
|
|
match self {
|
|
Self::Tavily => "TAVILY_API_KEY",
|
|
Self::Brave => "BRAVE_API_KEY",
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Normalized search result shape shared across providers.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SearchResult {
|
|
pub title: String,
|
|
pub url: String,
|
|
pub description: String,
|
|
pub age: Option<String>,
|
|
}
|
|
|
|
/// Dispatch a search request to whichever provider the server is configured to use.
|
|
pub async fn search(
|
|
provider: SearchProvider,
|
|
query: &str,
|
|
tavily_api_key: Option<&str>,
|
|
brave_api_key: Option<&str>,
|
|
count: u8,
|
|
) -> Result<Vec<SearchResult>, String> {
|
|
match provider {
|
|
SearchProvider::Tavily => {
|
|
let api_key = tavily_api_key
|
|
.filter(|key| !key.is_empty())
|
|
.ok_or_else(|| "TAVILY_API_KEY is not configured".to_string())?;
|
|
tavily::search(query, api_key, count).await
|
|
}
|
|
SearchProvider::Brave => {
|
|
let api_key = brave_api_key
|
|
.filter(|key| !key.is_empty())
|
|
.ok_or_else(|| "BRAVE_API_KEY is not configured".to_string())?;
|
|
brave::search(query, api_key, count).await
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Turn search results into plain text the AI model can read as tool output.
|
|
pub fn format_results(results: &[SearchResult]) -> String {
|
|
if results.is_empty() {
|
|
return "No search results found.".to_string();
|
|
}
|
|
|
|
let mut output = String::new();
|
|
for (i, r) in results.iter().enumerate() {
|
|
output.push_str(&format!("{}. {}\n", i + 1, r.title));
|
|
if !r.url.is_empty() {
|
|
output.push_str(&format!(" URL: {}\n", r.url));
|
|
}
|
|
if let Some(age) = &r.age {
|
|
output.push_str(&format!(" Age: {}\n", age));
|
|
}
|
|
output.push_str(&format!(" {}\n\n", r.description));
|
|
}
|
|
output
|
|
}
|