use std::fmt;
use serde::{Deserialize, Serialize};
use super::{SearchError, SearchProvider, SearchResult};
pub struct TavilyProvider {
pub(crate) api_key: String,
pub(crate) client: reqwest::Client,
}
impl TavilyProvider {
pub fn new(api_key: String, client: reqwest::Client) -> Self {
Self { api_key, client }
}
}
impl fmt::Debug for TavilyProvider {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TavilyProvider")
.field("api_key", &"[REDACTED]")
.field("client", &self.client)
.finish()
}
}
#[derive(Serialize)]
struct TavilyRequest<'a> {
api_key: &'a str,
query: &'a str,
max_results: usize,
}
#[derive(Deserialize)]
struct TavilyResponse {
results: Vec<TavilyResult>,
}
#[derive(Deserialize)]
struct TavilyResult {
title: String,
url: String,
content: Option<String>,
}
impl SearchProvider for TavilyProvider {
fn name(&self) -> &str {
"tavily"
}
fn search(
&self,
query: &str,
max_results: usize,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<Vec<SearchResult>, SearchError>> + Send + '_>,
> {
let query = query.to_owned();
Box::pin(async move {
if self.api_key.is_empty() {
return Err(SearchError::ApiKeyMissing);
}
let request_body = TavilyRequest {
api_key: &self.api_key,
query: &query,
max_results,
};
let body_json = serde_json::to_string(&request_body)
.map_err(|e| SearchError::ParseError(e.to_string()))?;
let response = self
.client
.post("https://api.tavily.com/search")
.header("Content-Type", "application/json")
.body(body_json)
.send()
.await
.map_err(|e| SearchError::NetworkError(e.to_string()))?;
if response.status() == reqwest::StatusCode::TOO_MANY_REQUESTS {
return Err(SearchError::RateLimited);
}
if response.status() == reqwest::StatusCode::UNAUTHORIZED
|| response.status() == reqwest::StatusCode::FORBIDDEN
{
return Err(SearchError::ApiKeyMissing);
}
if !response.status().is_success() {
return Err(SearchError::ProviderUnavailable(format!(
"Tavily returned status {}",
response.status()
)));
}
let body_text = response
.text()
.await
.map_err(|e| SearchError::ParseError(e.to_string()))?;
let body: TavilyResponse = serde_json::from_str(&body_text)
.map_err(|e| SearchError::ParseError(e.to_string()))?;
let results = body
.results
.into_iter()
.take(max_results)
.map(|r| SearchResult {
title: r.title,
url: r.url,
snippet: r.content.unwrap_or_default(),
})
.collect();
Ok(results)
})
}
}
#[cfg(test)]
mod tests {
use super::TavilyProvider;
#[test]
fn debug_output_redacts_api_key() {
const SECRET: &str = "tavily-super-secret-token";
let provider = TavilyProvider::new(SECRET.to_string(), reqwest::Client::new());
let debug = format!("{provider:?}");
assert!(!debug.contains(SECRET), "api key leaked: {debug}");
assert!(debug.contains("[REDACTED]"), "{debug}");
}
}