refactor(personalization): improve Cohere reranking logic and error handling

- Replace and_then() with early return for missing personalization
- Simplify reranking by building new hits vector instead of swapping
- Add debug logging for reranked indices
- Fix potential index out-of-bounds issues in reranking
This commit is contained in:
ManyTheFish
2025-07-23 12:13:53 +02:00
parent c56be3d820
commit e3062c4070

View File

@ -23,7 +23,8 @@ impl CohereService {
query: Option<&str>, query: Option<&str>,
) -> Result<SearchResult, ResponseError> { ) -> Result<SearchResult, ResponseError> {
// Extract user context from personalization // Extract user context from personalization
let user_context = personalize.and_then(|p| p.user_context.as_deref()); let Some(personalize) = personalize else { return Ok(search_result) };
let user_context = personalize.user_context.as_deref();
// Build the prompt by merging query and user context // Build the prompt by merging query and user context
let prompt = match (query, user_context) { let prompt = match (query, user_context) {
@ -64,13 +65,12 @@ impl CohereService {
// Create a mapping from original index to new rank // Create a mapping from original index to new rank
let reranked_indices: Vec<usize> = let reranked_indices: Vec<usize> =
rerank_response.iter().map(|result| result.index as usize).collect(); rerank_response.iter().map(|result| result.index as usize).collect();
debug!("Reranked indices: {:?}", reranked_indices);
// Reorder the hits based on Cohere's reranking // Reorder the hits based on Cohere's reranking
let mut reranked_hits = search_result.hits.clone(); let mut reranked_hits = Vec::new();
for (new_index, original_index) in reranked_indices.iter().enumerate() { for index in reranked_indices.iter() {
if *original_index < reranked_hits.len() { reranked_hits.push(search_result.hits[*index].clone());
reranked_hits.swap(new_index, *original_index);
}
} }
Ok(SearchResult { hits: reranked_hits, ..search_result }) Ok(SearchResult { hits: reranked_hits, ..search_result })