use crate::error::Result;
use crate::framework::ChaoticSemanticFramework;
use crate::hyperdim::HVec10240;
use crate::retrieval::rerank::{RerankCandidate, Reranker};
use std::sync::Arc;
use tracing::instrument;
impl ChaoticSemanticFramework {
#[instrument(err, skip(self, query, rerankers))]
pub async fn probe_with_rerankers(
&self,
query: HVec10240,
initial_top_k: usize,
rerankers: &[Box<dyn Reranker>],
final_top_k: usize,
) -> Result<Vec<(String, f32)>> {
self.validate_top_k(initial_top_k)?;
self.validate_top_k(final_top_k)?;
let actual_initial_k = initial_top_k.max(final_top_k);
let initial_results = self.probe(query, actual_initial_k).await?;
if rerankers.is_empty() {
let mut results = initial_results;
results.truncate(final_top_k);
return Ok(results);
}
let mut candidates = Vec::with_capacity(initial_results.len());
{
let sing = self.singularity.read().await;
let ns = self.namespace.read().await;
for (id, score) in initial_results {
if let Some(concept) = sing.get(&ns, &id) {
candidates.push(RerankCandidate {
id: concept.id.clone(),
vector: Arc::new(concept.vector),
metadata: concept.metadata.clone(),
score,
created_at_unix: concept.created_at,
});
}
}
}
for reranker in rerankers {
candidates = reranker.rerank(&query, candidates, actual_initial_k);
}
candidates.truncate(final_top_k);
Ok(candidates.into_iter().map(|c| (c.id, c.score)).collect())
}
}