Skip to main content

nodedb_types/protocol/
batch.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Batch operation payload types.
4
5use serde::{Deserialize, Serialize};
6
7use crate::json_msgpack::JsonValue;
8
9/// A single vector in a batch insert.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct BatchVector {
12    pub id: String,
13    pub embedding: Vec<f32>,
14    #[serde(skip_serializing_if = "Option::is_none")]
15    pub metadata: Option<serde_json::Value>,
16}
17
18impl zerompk::ToMessagePack for BatchVector {
19    fn write<W: zerompk::Write>(&self, writer: &mut W) -> zerompk::Result<()> {
20        writer.write_array_len(3)?;
21        writer.write_string(&self.id)?;
22        self.embedding.write(writer)?;
23        self.metadata
24            .as_ref()
25            .map(|v| JsonValue(v.clone()))
26            .write(writer)
27    }
28}
29
30impl<'a> zerompk::FromMessagePack<'a> for BatchVector {
31    fn read<R: zerompk::Read<'a>>(reader: &mut R) -> zerompk::Result<Self> {
32        let len = reader.read_array_len()?;
33        if len != 3 {
34            return Err(zerompk::Error::ArrayLengthMismatch {
35                expected: 3,
36                actual: len,
37            });
38        }
39        let id = reader.read_string()?.into_owned();
40        let embedding = Vec::<f32>::read(reader)?;
41        let metadata = Option::<JsonValue>::read(reader)?.map(|v| v.0);
42        Ok(Self {
43            id,
44            embedding,
45            metadata,
46        })
47    }
48}
49
50/// A single document in a batch insert.
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct BatchDocument {
53    pub id: String,
54    pub fields: serde_json::Value,
55}
56
57impl zerompk::ToMessagePack for BatchDocument {
58    fn write<W: zerompk::Write>(&self, writer: &mut W) -> zerompk::Result<()> {
59        writer.write_array_len(2)?;
60        writer.write_string(&self.id)?;
61        JsonValue(self.fields.clone()).write(writer)
62    }
63}
64
65impl<'a> zerompk::FromMessagePack<'a> for BatchDocument {
66    fn read<R: zerompk::Read<'a>>(reader: &mut R) -> zerompk::Result<Self> {
67        let len = reader.read_array_len()?;
68        if len != 2 {
69            return Err(zerompk::Error::ArrayLengthMismatch {
70                expected: 2,
71                actual: len,
72            });
73        }
74        let id = reader.read_string()?.into_owned();
75        let fields = JsonValue::read(reader)?.0;
76        Ok(Self { id, fields })
77    }
78}