Add fields support

crates/milli/src/prompt/document.rs

@@ -12,6 +12,7 @@ use liquid::{ObjectView, ValueView};
 use rustc_hash::FxBuildHasher;
 use serde_json::value::RawValue;
 
+use crate::constants::{RESERVED_GEO_FIELD_NAME, RESERVED_VECTORS_FIELD_NAME};
 use crate::update::del_add::{DelAdd, KvReaderDelAdd};
 use crate::FieldsIdsMap;
 
@@ -143,6 +144,112 @@ impl ValueView for Document<'_> {
 /// Implementation for any type that implements the Document trait
 use crate::update::new::document::Document as DocumentTrait;
 
+pub struct JsonDocument {
+    object: liquid::Object,
+    cached: BTreeMap<String, Box<RawValue>>,
+}
+
+impl JsonDocument {
+    pub fn new(value: &serde_json::Value) -> Self {
+        let to_string = serde_json::to_string(&value)
+            .expect("JsonDocument should only be created with valid JSON"); // TODO: Remove panic
+        let back_to_value: BTreeMap<String, Box<RawValue>> = serde_json::from_str(&to_string)
+            .expect("JsonDocument should only be created with valid JSON");
+        let object =
+            liquid::to_object(&value).expect("JsonDocument should only be created with valid JSON");
+        Self { object, cached: back_to_value }
+    }
+}
+
+impl std::fmt::Debug for JsonDocument {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.object.fmt(f)
+    }
+}
+
+impl<'a> DocumentTrait<'a> for &'a JsonDocument {
+    fn iter_top_level_fields(
+        &self,
+    ) -> impl Iterator<Item = crate::Result<(&'a str, &'a RawValue)>> {
+        self.cached.iter().filter_map(|(k, v)| {
+            if k == RESERVED_VECTORS_FIELD_NAME || k == RESERVED_GEO_FIELD_NAME {
+                None
+            } else {
+                Some(Ok((k.as_str(), v.as_ref())))
+            }
+        })
+    }
+
+    fn top_level_fields_count(&self) -> usize {
+        self.cached.len()
+            - self.cached.contains_key(RESERVED_VECTORS_FIELD_NAME) as usize
+            - self.cached.contains_key(RESERVED_GEO_FIELD_NAME) as usize
+    }
+
+    fn top_level_field(&self, k: &str) -> crate::Result<Option<&'a RawValue>> {
+        if k == RESERVED_VECTORS_FIELD_NAME || k == RESERVED_GEO_FIELD_NAME {
+            return Ok(None);
+        }
+        Ok(self.cached.get(k).map(|r| r.as_ref()))
+    }
+
+    fn vectors_field(&self) -> crate::Result<Option<&'a RawValue>> {
+        Ok(self.cached.get(RESERVED_VECTORS_FIELD_NAME).map(|r| r.as_ref()))
+    }
+
+    fn geo_field(&self) -> crate::Result<Option<&'a RawValue>> {
+        Ok(self.cached.get(RESERVED_GEO_FIELD_NAME).map(|r| r.as_ref()))
+    }
+}
+
+impl ObjectView for JsonDocument {
+    fn as_value(&self) -> &dyn ValueView {
+        self.object.as_value()
+    }
+    fn size(&self) -> i64 {
+        self.object.size()
+    }
+    fn keys<'k>(&'k self) -> Box<dyn Iterator<Item = KStringCow<'k>> + 'k> {
+        Box::new(self.object.keys().map(|s| s.into()))
+    }
+    fn values<'k>(&'k self) -> Box<dyn Iterator<Item = &'k dyn ValueView> + 'k> {
+        Box::new(self.object.values().map(|v| v.as_view()))
+    }
+    fn iter<'k>(&'k self) -> Box<dyn Iterator<Item = (KStringCow<'k>, &'k dyn ValueView)> + 'k> {
+        Box::new(self.object.iter().map(|(k, v)| (k.into(), v.as_view())))
+    }
+    fn contains_key(&self, index: &str) -> bool {
+        self.object.contains_key(index)
+    }
+    fn get<'s>(&'s self, index: &str) -> Option<&'s dyn ValueView> {
+        self.object.get(index).map(|v| v.as_view())
+    }
+}
+
+impl ValueView for JsonDocument {
+    fn as_debug(&self) -> &dyn fmt::Debug {
+        self.object.as_debug()
+    }
+    fn render(&self) -> DisplayCow<'_> {
+        self.object.render()
+    }
+    fn source(&self) -> DisplayCow<'_> {
+        self.object.source()
+    }
+    fn type_name(&self) -> &'static str {
+        self.object.type_name()
+    }
+    fn query_state(&self, state: State) -> bool {
+        self.object.query_state(state)
+    }
+    fn to_kstr(&self) -> KStringCow<'_> {
+        self.object.to_kstr()
+    }
+    fn to_value(&self) -> LiquidValue {
+        self.object.to_value()
+    }
+}
+
 #[derive(Debug)]
 pub struct ParseableDocument<'a, 'doc, D: DocumentTrait<'a> + Debug> {
     document: D,

crates/milli/src/prompt/mod.rs

@@ -12,11 +12,16 @@ use bumpalo::Bump;
 pub(crate) use document::{Document, ParseableDocument};
 use error::{NewPromptError, RenderPromptError};
 pub use fields::{BorrowedFields, OwnedFields};
+use heed::RoTxn;
+use liquid::model::Value as LiquidValue;
+use liquid::ValueView;
 
 pub use self::context::Context;
 use crate::fields_ids_map::metadata::FieldIdMapWithMetadata;
+use crate::prompt::document::JsonDocument;
 use crate::update::del_add::DelAdd;
-use crate::GlobalFieldsIdsMap;
+use crate::update::new::document::DocumentFromDb;
+use crate::{GlobalFieldsIdsMap, Index, MetadataBuilder};
 
 pub struct Prompt {
     template: liquid::Template,
@@ -164,6 +169,47 @@ fn truncate(s: &mut String, max_bytes: usize) {
     }
 }
 
+pub fn get_inline_document_fields(
+    index: &Index,
+    rtxn: &RoTxn<'_>,
+    inline_doc: &serde_json::Value,
+) -> Result<LiquidValue, crate::Error> {
+    let fid_map_with_meta = index.fields_ids_map_with_metadata(rtxn)?;
+    let inline_doc = JsonDocument::new(inline_doc);
+    let fields = OwnedFields::new(&inline_doc, &fid_map_with_meta);
+
+    Ok(fields.to_value())
+}
+
+pub fn get_document(
+    index: &Index,
+    rtxn: &RoTxn<'_>,
+    external_id: &str,
+    with_fields: bool,
+) -> Result<Option<(LiquidValue, Option<LiquidValue>)>, crate::Error> {
+    let Some(internal_id) = index.external_documents_ids().get(rtxn, external_id)? else {
+        return Ok(None);
+    };
+
+    let fid_map = index.fields_ids_map(rtxn)?;
+    let Some(document_from_db) = DocumentFromDb::new(internal_id, rtxn, index, &fid_map)? else {
+        return Ok(None);
+    };
+
+    let doc_alloc = Bump::new();
+    let parseable_document = ParseableDocument::new(document_from_db, &doc_alloc);
+
+    if with_fields {
+        let metadata_builder = MetadataBuilder::from_index(index, rtxn)?;
+        let fid_map_with_meta = FieldIdMapWithMetadata::new(fid_map.clone(), metadata_builder);
+        let fields = OwnedFields::new(&parseable_document, &fid_map_with_meta);
+
+        Ok(Some((parseable_document.to_value(), Some(fields.to_value()))))
+    } else {
+        Ok(Some((parseable_document.to_value(), None)))
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::Prompt;

crates/milli/src/vector/json_template/mod.rs

@@ -143,14 +143,6 @@ impl JsonTemplate {
         self.render(&search_data)
     }
 
-    /// Renders any serializable value by converting it to a liquid object and rendering it with the template.
-    /// If its a map, values inside can be accessed directly by their keys.
-    pub fn render_serializable<T: Serialize>(&self, object: &T) -> Result<Value, Error> {
-        let object = liquid::to_object(object)
-            .map_err(|err| Error { template_error: err, path: ValuePath::new() })?;
-        self.render(&object)
-    }
-
     /// The JSON value representing the underlying template
     pub fn template(&self) -> &Value {
         &self.value