use super::VectorizerClient;
use crate::error::{Result, VectorizerError};
use crate::models::{
AnswerPlan, AnswerPlanRequest, BroadDiscoveryRequest, BroadDiscoveryResponse,
CompressEvidenceRequest, CompressEvidenceResponse, LlmPrompt, PromoteReadmeRequest,
PromoteReadmeResponse, RenderPromptRequest, SemanticFocusRequest, SemanticFocusResponse,
};
impl VectorizerClient {
#[allow(clippy::too_many_arguments)]
pub async fn discover(
&self,
query: &str,
include_collections: Option<Vec<String>>,
exclude_collections: Option<Vec<String>>,
max_bullets: Option<usize>,
broad_k: Option<usize>,
focus_k: Option<usize>,
) -> Result<serde_json::Value> {
if query.trim().is_empty() {
return Err(VectorizerError::validation("Query cannot be empty"));
}
if let Some(max) = max_bullets
&& max == 0
{
return Err(VectorizerError::validation(
"max_bullets must be greater than 0",
));
}
let mut payload = serde_json::Map::new();
payload.insert(
"query".to_string(),
serde_json::Value::String(query.to_string()),
);
if let Some(inc) = include_collections {
payload.insert(
"include_collections".to_string(),
serde_json::to_value(inc).unwrap(),
);
}
if let Some(exc) = exclude_collections {
payload.insert(
"exclude_collections".to_string(),
serde_json::to_value(exc).unwrap(),
);
}
if let Some(max) = max_bullets {
payload.insert(
"max_bullets".to_string(),
serde_json::Value::Number(max.into()),
);
}
if let Some(k) = broad_k {
payload.insert("broad_k".to_string(), serde_json::Value::Number(k.into()));
}
if let Some(k) = focus_k {
payload.insert("focus_k".to_string(), serde_json::Value::Number(k.into()));
}
let response = self
.make_request(
"POST",
"/discover",
Some(serde_json::Value::Object(payload)),
)
.await?;
serde_json::from_str(&response)
.map_err(|e| VectorizerError::server(format!("Failed to parse discover response: {e}")))
}
pub async fn filter_collections(
&self,
query: &str,
include: Option<Vec<String>>,
exclude: Option<Vec<String>>,
) -> Result<serde_json::Value> {
if query.trim().is_empty() {
return Err(VectorizerError::validation("Query cannot be empty"));
}
let mut payload = serde_json::Map::new();
payload.insert(
"query".to_string(),
serde_json::Value::String(query.to_string()),
);
if let Some(inc) = include {
payload.insert("include".to_string(), serde_json::to_value(inc).unwrap());
}
if let Some(exc) = exclude {
payload.insert("exclude".to_string(), serde_json::to_value(exc).unwrap());
}
let response = self
.make_request(
"POST",
"/discovery/filter_collections",
Some(serde_json::Value::Object(payload)),
)
.await?;
serde_json::from_str(&response)
.map_err(|e| VectorizerError::server(format!("Failed to parse filter response: {e}")))
}
pub async fn score_collections(
&self,
query: &str,
name_match_weight: Option<f32>,
term_boost_weight: Option<f32>,
signal_boost_weight: Option<f32>,
) -> Result<serde_json::Value> {
if let Some(w) = name_match_weight
&& !(0.0..=1.0).contains(&w)
{
return Err(VectorizerError::validation(
"name_match_weight must be between 0.0 and 1.0",
));
}
if let Some(w) = term_boost_weight
&& !(0.0..=1.0).contains(&w)
{
return Err(VectorizerError::validation(
"term_boost_weight must be between 0.0 and 1.0",
));
}
if let Some(w) = signal_boost_weight
&& !(0.0..=1.0).contains(&w)
{
return Err(VectorizerError::validation(
"signal_boost_weight must be between 0.0 and 1.0",
));
}
let mut payload = serde_json::Map::new();
payload.insert(
"query".to_string(),
serde_json::Value::String(query.to_string()),
);
if let Some(w) = name_match_weight {
payload.insert("name_match_weight".to_string(), serde_json::json!(w));
}
if let Some(w) = term_boost_weight {
payload.insert("term_boost_weight".to_string(), serde_json::json!(w));
}
if let Some(w) = signal_boost_weight {
payload.insert("signal_boost_weight".to_string(), serde_json::json!(w));
}
let response = self
.make_request(
"POST",
"/discovery/score_collections",
Some(serde_json::Value::Object(payload)),
)
.await?;
serde_json::from_str(&response)
.map_err(|e| VectorizerError::server(format!("Failed to parse score response: {e}")))
}
pub async fn expand_queries(
&self,
query: &str,
max_expansions: Option<usize>,
include_definition: Option<bool>,
include_features: Option<bool>,
include_architecture: Option<bool>,
) -> Result<serde_json::Value> {
let mut payload = serde_json::Map::new();
payload.insert(
"query".to_string(),
serde_json::Value::String(query.to_string()),
);
if let Some(max) = max_expansions {
payload.insert(
"max_expansions".to_string(),
serde_json::Value::Number(max.into()),
);
}
if let Some(def) = include_definition {
payload.insert(
"include_definition".to_string(),
serde_json::Value::Bool(def),
);
}
if let Some(feat) = include_features {
payload.insert(
"include_features".to_string(),
serde_json::Value::Bool(feat),
);
}
if let Some(arch) = include_architecture {
payload.insert(
"include_architecture".to_string(),
serde_json::Value::Bool(arch),
);
}
let response = self
.make_request(
"POST",
"/discovery/expand_queries",
Some(serde_json::Value::Object(payload)),
)
.await?;
serde_json::from_str(&response)
.map_err(|e| VectorizerError::server(format!("Failed to parse expand response: {e}")))
}
pub async fn broad_discovery(
&self,
request: BroadDiscoveryRequest,
) -> Result<BroadDiscoveryResponse> {
let payload = serde_json::json!({
"queries": request.queries,
"k": request.k.unwrap_or(50),
});
let response = self
.make_request("POST", "/discovery/broad_discovery", Some(payload))
.await?;
serde_json::from_str(&response).map_err(|e| {
VectorizerError::server(format!("Failed to parse broad_discovery response: {e}"))
})
}
pub async fn semantic_focus(
&self,
request: SemanticFocusRequest,
) -> Result<SemanticFocusResponse> {
let payload = serde_json::json!({
"collection": request.collection,
"queries": request.queries,
"k": request.k.unwrap_or(15),
});
let response = self
.make_request("POST", "/discovery/semantic_focus", Some(payload))
.await?;
serde_json::from_str(&response).map_err(|e| {
VectorizerError::server(format!("Failed to parse semantic_focus response: {e}"))
})
}
pub async fn promote_readme(
&self,
request: PromoteReadmeRequest,
) -> Result<PromoteReadmeResponse> {
let payload = serde_json::json!({ "chunks": request.chunks });
let response = self
.make_request("POST", "/discovery/promote_readme", Some(payload))
.await?;
serde_json::from_str(&response).map_err(|e| {
VectorizerError::server(format!("Failed to parse promote_readme response: {e}"))
})
}
pub async fn compress_evidence(
&self,
request: CompressEvidenceRequest,
) -> Result<CompressEvidenceResponse> {
let mut payload = serde_json::json!({ "chunks": request.chunks });
if let Some(mb) = request.max_bullets {
payload["max_bullets"] = serde_json::json!(mb);
}
if let Some(mpd) = request.max_per_doc {
payload["max_per_doc"] = serde_json::json!(mpd);
}
let response = self
.make_request("POST", "/discovery/compress_evidence", Some(payload))
.await?;
serde_json::from_str(&response).map_err(|e| {
VectorizerError::server(format!("Failed to parse compress_evidence response: {e}"))
})
}
pub async fn build_answer_plan(&self, request: AnswerPlanRequest) -> Result<AnswerPlan> {
let payload = serde_json::json!({ "bullets": request.bullets });
let response = self
.make_request("POST", "/discovery/build_answer_plan", Some(payload))
.await?;
serde_json::from_str(&response).map_err(|e| {
VectorizerError::server(format!("Failed to parse build_answer_plan response: {e}"))
})
}
pub async fn render_llm_prompt(&self, request: RenderPromptRequest) -> Result<LlmPrompt> {
let payload = serde_json::json!({ "plan": request.plan });
let response = self
.make_request("POST", "/discovery/render_llm_prompt", Some(payload))
.await?;
serde_json::from_str(&response).map_err(|e| {
VectorizerError::server(format!("Failed to parse render_llm_prompt response: {e}"))
})
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use serde_json::json;
use crate::models::{
AnswerPlan, AnswerPlanRequest, BroadDiscoveryRequest, BroadDiscoveryResponse,
CompressEvidenceRequest, CompressEvidenceResponse, LlmPrompt, PromoteReadmeRequest,
PromoteReadmeResponse, RenderPromptRequest, SemanticFocusRequest, SemanticFocusResponse,
};
#[test]
fn broad_discovery_request_serializes() {
let req = BroadDiscoveryRequest {
queries: vec!["HNSW index".into(), "embedding model".into()],
k: Some(30),
};
let v = serde_json::to_value(&req).unwrap();
assert_eq!(v["queries"][0], "HNSW index");
assert_eq!(v["k"], 30);
}
#[test]
fn broad_discovery_response_deserializes() {
let raw = json!({
"chunks": [{"collection": "docs", "score": 0.9, "content_preview": "test"}],
"count": 1
});
let resp: BroadDiscoveryResponse = serde_json::from_value(raw).unwrap();
assert_eq!(resp.count, 1);
assert_eq!(resp.chunks.len(), 1);
}
#[test]
fn semantic_focus_request_serializes() {
let req = SemanticFocusRequest {
collection: "code".into(),
queries: vec!["async runtime".into()],
k: None,
};
let v = serde_json::to_value(&req).unwrap();
assert_eq!(v["collection"], "code");
assert_eq!(v["queries"][0], "async runtime");
}
#[test]
fn semantic_focus_response_deserializes() {
let raw = json!({ "chunks": [], "count": 0 });
let resp: SemanticFocusResponse = serde_json::from_value(raw).unwrap();
assert_eq!(resp.count, 0);
}
#[test]
fn promote_readme_request_serializes() {
let req = PromoteReadmeRequest {
chunks: vec![json!({"collection": "docs", "score": 0.8, "content": "README text"})],
};
let v = serde_json::to_value(&req).unwrap();
assert!(v["chunks"].is_array());
}
#[test]
fn promote_readme_response_deserializes() {
let raw = json!({ "promoted_chunks": [], "count": 0 });
let resp: PromoteReadmeResponse = serde_json::from_value(raw).unwrap();
assert_eq!(resp.count, 0);
}
#[test]
fn compress_evidence_round_trip() {
let req = CompressEvidenceRequest {
chunks: vec![json!({"collection": "c", "score": 1.0, "content": "x"})],
max_bullets: Some(5),
max_per_doc: Some(2),
};
let v = serde_json::to_value(&req).unwrap();
assert_eq!(v["max_bullets"], 5);
let raw = json!({ "bullets": [{"text": "b", "source_id": "s", "category": "Feature", "score": 0.9}], "count": 1 });
let resp: CompressEvidenceResponse = serde_json::from_value(raw).unwrap();
assert_eq!(resp.count, 1);
}
#[test]
fn answer_plan_round_trip() {
let plan = AnswerPlan {
sections: vec![json!({"title": "Intro", "bullets_count": 1, "bullets": []})],
total_bullets: 1,
sources: vec!["docs".into()],
};
let serialized = serde_json::to_value(&plan).unwrap();
let parsed: AnswerPlan = serde_json::from_value(serialized).unwrap();
assert_eq!(parsed.total_bullets, 1);
assert_eq!(parsed.sources[0], "docs");
}
#[test]
fn llm_prompt_deserializes() {
let raw = json!({ "prompt": "Answer: ...", "length": 10, "estimated_tokens": 2 });
let lp: LlmPrompt = serde_json::from_value(raw).unwrap();
assert_eq!(lp.prompt, "Answer: ...");
assert_eq!(lp.estimated_tokens, 2);
}
#[test]
fn render_prompt_request_serializes() {
let req = RenderPromptRequest {
plan: AnswerPlan {
sections: vec![],
total_bullets: 0,
sources: vec![],
},
};
let v = serde_json::to_value(&req).unwrap();
assert!(v["plan"].is_object());
}
}