pub use super::health::{EmbedderState, HealthResponse};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
#[derive(Debug, thiserror::Error)]
pub enum SearchClientError {
#[error("trusty-search transport error: {0}")]
Transport(String),
#[error("trusty-search API returned {status}: {body}")]
Api {
status: u16,
body: String,
},
#[error("trusty-search response parse error: {0}")]
Parse(String),
#[error("trusty-search is unavailable: {0}")]
Unavailable(String),
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct IndexInfo {
pub id: String,
#[serde(default)]
pub name: Option<String>,
#[serde(default)]
pub root_path: Option<String>,
}
#[derive(Debug, Deserialize)]
pub(crate) struct ListIndexesResponse {
pub(crate) indexes: Vec<IndexInfo>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct SearchResult {
pub file: String,
#[serde(default)]
pub snippet: Option<String>,
#[serde(default)]
pub score: f32,
#[serde(default)]
pub start_line: Option<u32>,
#[serde(default)]
pub end_line: Option<u32>,
}
#[derive(Debug, Clone, Serialize)]
pub struct SearchRequest {
pub text: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<u32>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct SearchResponse {
#[serde(default)]
pub results: Vec<SearchResult>,
}
#[async_trait]
pub trait SearchClient: Send + Sync {
async fn health(&self) -> Result<HealthResponse, SearchClientError>;
async fn list_indexes(&self) -> Result<Vec<IndexInfo>, SearchClientError>;
async fn search(
&self,
index_id: &str,
query: &str,
top_k: Option<u32>,
) -> Result<Vec<SearchResult>, SearchClientError>;
}
pub struct HttpSearchClient {
base_url: String,
http: reqwest::Client,
}
impl HttpSearchClient {
pub fn new(base_url: impl Into<String>) -> Self {
let raw = base_url.into();
let base_url = raw.trim_end_matches('/').to_string();
let http = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.build()
.expect("reqwest::Client::build failed");
Self { base_url, http }
}
pub fn from_config(config: &crate::config::ReviewConfig) -> Self {
Self::new(config.search_url.clone())
}
pub fn base_url(&self) -> &str {
&self.base_url
}
}
#[async_trait]
impl SearchClient for HttpSearchClient {
async fn health(&self) -> Result<HealthResponse, SearchClientError> {
let url = format!("{}/health", self.base_url);
let resp = self
.http
.get(&url)
.send()
.await
.map_err(|e| SearchClientError::Unavailable(format!("GET {url}: {e}")))?;
let status = resp.status();
let body = resp
.text()
.await
.map_err(|e| SearchClientError::Transport(format!("read body of {url}: {e}")))?;
if !status.is_success() {
return Err(SearchClientError::Unavailable(format!(
"GET {url} returned {status}: {body}"
)));
}
serde_json::from_str(&body)
.map_err(|e| SearchClientError::Parse(format!("health response: {e}")))
}
async fn list_indexes(&self) -> Result<Vec<IndexInfo>, SearchClientError> {
let url = format!("{}/indexes?details=true", self.base_url);
let resp = self
.http
.get(&url)
.send()
.await
.map_err(|e| SearchClientError::Transport(format!("GET {url}: {e}")))?;
let status = resp.status();
let body = resp
.text()
.await
.map_err(|e| SearchClientError::Transport(format!("read body of {url}: {e}")))?;
if !status.is_success() {
return Err(SearchClientError::Api {
status: status.as_u16(),
body,
});
}
let envelope: ListIndexesResponse = serde_json::from_str(&body)
.map_err(|e| SearchClientError::Parse(format!("list indexes response: {e}")))?;
Ok(envelope.indexes)
}
async fn search(
&self,
index_id: &str,
query: &str,
top_k: Option<u32>,
) -> Result<Vec<SearchResult>, SearchClientError> {
let url = format!("{}/indexes/{index_id}/search", self.base_url);
let request_body = SearchRequest {
text: query.to_string(),
top_k,
};
let resp = self
.http
.post(&url)
.json(&request_body)
.send()
.await
.map_err(|e| SearchClientError::Transport(format!("POST {url}: {e}")))?;
let status = resp.status();
let body = resp
.text()
.await
.map_err(|e| SearchClientError::Transport(format!("read body of {url}: {e}")))?;
if !status.is_success() {
return Err(SearchClientError::Api {
status: status.as_u16(),
body,
});
}
let search_resp: SearchResponse = serde_json::from_str(&body)
.map_err(|e| SearchClientError::Parse(format!("search response: {e}")))?;
Ok(search_resp.results)
}
}
#[cfg(test)]
#[path = "search_client_tests.rs"]
mod tests;