fix two bug and add a dump test

This commit is contained in:
Tamo 2024-06-04 16:41:33 +02:00
parent 5d50850e12
commit cc5dca8321
4 changed files with 234 additions and 14 deletions

View file

@ -52,7 +52,7 @@ pub struct SearchQueryGet {
#[deserr(default, error = DeserrQueryParamError<InvalidSearchAttributesToRetrieve>)]
attributes_to_retrieve: Option<CS<String>>,
#[deserr(default, error = DeserrQueryParamError<InvalidSearchRetrieveVectors>)]
retrieve_vectors: bool,
retrieve_vectors: Param<bool>,
#[deserr(default, error = DeserrQueryParamError<InvalidSearchAttributesToCrop>)]
attributes_to_crop: Option<CS<String>>,
#[deserr(default = Param(DEFAULT_CROP_LENGTH()), error = DeserrQueryParamError<InvalidSearchCropLength>)]
@ -155,7 +155,7 @@ impl From<SearchQueryGet> for SearchQuery {
page: other.page.as_deref().copied(),
hits_per_page: other.hits_per_page.as_deref().copied(),
attributes_to_retrieve: other.attributes_to_retrieve.map(|o| o.into_iter().collect()),
retrieve_vectors: other.retrieve_vectors,
retrieve_vectors: other.retrieve_vectors.0,
attributes_to_crop: other.attributes_to_crop.map(|o| o.into_iter().collect()),
crop_length: other.crop_length.0,
attributes_to_highlight: other.attributes_to_highlight.map(|o| o.into_iter().collect()),

View file

@ -1051,6 +1051,7 @@ fn make_hits(
formatter_builder.highlight_prefix(format.highlight_pre_tag);
formatter_builder.highlight_suffix(format.highlight_post_tag);
let mut documents = Vec::new();
let embedding_configs = index.embedding_configs(&rtxn)?;
let documents_iter = index.documents(rtxn, documents_ids)?;
for ((id, obkv), score) in documents_iter.into_iter().zip(document_scores.into_iter()) {
// First generate a document with all the displayed fields
@ -1066,12 +1067,19 @@ fn make_hits(
if retrieve_vectors {
let mut vectors = serde_json::Map::new();
for (name, mut vector) in index.embeddings(&rtxn, id)? {
let user_defined = embedding_configs
.iter()
.find(|conf| conf.name == name)
.is_some_and(|conf| conf.user_defined.contains(id));
let mut embedding = serde_json::Map::new();
embedding.insert("userDefined".to_string(), user_defined.into());
if vector.len() == 1 {
let vector = vector.pop().unwrap();
vectors.insert(name.into(), vector.into());
embedding.insert("embedding".to_string(), vector.into());
} else {
vectors.insert(name.into(), vector.into());
embedding.insert("embedding".to_string(), vector.into());
}
vectors.insert(name.into(), embedding.into());
}
document.insert("_vectors".into(), vectors.into());
}