Inference Service, query requests (batch processing) (qdrant#5329)

* Inference Service, query requests (batch processing) --------- Co-authored-by: n0x29a <n0x29a@example.com> Co-authored-by: generall <andrey@vasnetsov.com>
AKJUS · Nov 8, 2024 · 097508d · 097508d
1 parent 865268a
commit 097508d
Show file tree

Hide file tree

Showing 11 changed files with 1,046 additions and 295 deletions.
diff --git a/lib/api/src/conversions/inference.rs b/lib/api/src/conversions/inference.rs
@@ -2,14 +2,18 @@ use tonic::Status;
 
 use crate::conversions::json::{dict_to_proto, json_to_proto, proto_dict_to_json, proto_to_json};
 use crate::grpc::qdrant as grpc;
-use crate::rest::schema as rest;
+use crate::rest::{schema as rest, Options};
 
 impl From<rest::Document> for grpc::Document {
     fn from(document: rest::Document) -> Self {
         Self {
             text: document.text,
             model: document.model,
-            options: document.options.map(dict_to_proto).unwrap_or_default(),
+            options: document
+                .options
+                .options
+                .map(dict_to_proto)
+                .unwrap_or_default(),
         }
     }
 }
@@ -21,7 +25,9 @@ impl TryFrom<grpc::Document> for rest::Document {
         Ok(Self {
             text: document.text,
             model: document.model,
-            options: Some(proto_dict_to_json(document.options)?),
+            options: Options {
+                options: Some(proto_dict_to_json(document.options)?),
+            },
         })
     }
 }
@@ -31,7 +37,7 @@ impl From<rest::Image> for grpc::Image {
         Self {
             image: image.image,
             model: image.model,
-            options: image.options.map(dict_to_proto).unwrap_or_default(),
+            options: image.options.options.map(dict_to_proto).unwrap_or_default(),
         }
     }
 }
@@ -43,7 +49,9 @@ impl TryFrom<grpc::Image> for rest::Image {
         Ok(Self {
             image: image.image,
             model: image.model,
-            options: Some(proto_dict_to_json(image.options)?),
+            options: Options {
+                options: Some(proto_dict_to_json(image.options)?),
+            },
         })
     }
 }
@@ -53,7 +61,11 @@ impl From<rest::InferenceObject> for grpc::InferenceObject {
         Self {
             object: Some(json_to_proto(object.object)),
             model: object.model,
-            options: object.options.map(dict_to_proto).unwrap_or_default(),
+            options: object
+                .options
+                .options
+                .map(dict_to_proto)
+                .unwrap_or_default(),
         }
     }
 }
@@ -74,7 +86,9 @@ impl TryFrom<grpc::InferenceObject> for rest::InferenceObject {
         Ok(Self {
             object: proto_to_json(object)?,
             model,
-            options: Some(proto_dict_to_json(options)?),
+            options: Options {
+                options: Some(proto_dict_to_json(options)?),
+            },
         })
     }
 }
diff --git a/lib/api/src/rest/schema.rs b/lib/api/src/rest/schema.rs
@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::hash::{Hash, Hasher};
 
 use common::types::ScoreType;
 use common::validation::validate_multi_vector;
@@ -139,26 +140,46 @@ impl Validate for VectorStruct {
     }
 }
 
+#[derive(Clone, Default, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema)]
+pub struct Options {
+    /// Parameters for the model
+    /// Values of the parameters are model-specific
+    pub options: Option<HashMap<String, Value>>,
+}
+
+impl Hash for Options {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        // Order of keys in the map should not affect the hash
+        if let Some(options) = &self.options {
+            let mut keys: Vec<_> = options.keys().collect();
+            keys.sort();
+            for key in keys {
+                key.hash(state);
+                options.get(key).unwrap().hash(state);
+            }
+        }
+    }
+}
+
 /// WARN: Work-in-progress, unimplemented
 ///
 /// Text document for embedding. Requires inference infrastructure, unimplemented.
-#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema)]
+#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema, Hash)]
 pub struct Document {
     /// Text of the document
     /// This field will be used as input for the embedding model
     pub text: String,
     /// Name of the model used to generate the vector
     /// List of available models depends on a provider
     pub model: Option<String>,
-    /// Parameters for the model
-    /// Values of the parameters are model-specific
-    pub options: Option<HashMap<String, Value>>,
+    #[serde(flatten)]
+    pub options: Options,
 }
 
 /// WARN: Work-in-progress, unimplemented
 ///
 /// Image object for embedding. Requires inference infrastructure, unimplemented.
-#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema)]
+#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema, Hash)]
 pub struct Image {
     /// Image data: base64 encoded image or an URL
     pub image: String,
@@ -167,13 +188,14 @@ pub struct Image {
     pub model: Option<String>,
     /// Parameters for the model
     /// Values of the parameters are model-specific
-    pub options: Option<HashMap<String, Value>>,
+    #[serde(flatten)]
+    pub options: Options,
 }
 
 /// WARN: Work-in-progress, unimplemented
 ///
 /// Custom object for embedding. Requires inference infrastructure, unimplemented.
-#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema)]
+#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema, Hash)]
 pub struct InferenceObject {
     /// Arbitrary data, used as input for the embedding model
     /// Used if the model requires more than one input or a custom input
@@ -183,7 +205,8 @@ pub struct InferenceObject {
     pub model: Option<String>,
     /// Parameters for the model
     /// Values of the parameters are model-specific
-    pub options: Option<HashMap<String, Value>>,
+    #[serde(flatten)]
+    pub options: Options,
 }
 
 #[derive(Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema)]