AIProxy Set Command Preprocessing actions for storeinputs (#77)

* Add a new field preprocess_action for aiproxy set command

* Add preprocess_action field to python client

* Add new fields for aiproxy create store commands

* Update python client after changes to create store command

* Remove index and query types from aiproxy create store

* Update python client with changes to aiproxy create store params

* - Pass preproccess actions to set command
- Create base on how to edit inputtypes and convert to storekeys
- Move preprocessing actions
- Tie model info into aimodel
- Fix fmt trait on AIStoreInputTypes

* regenerate types for python client

* Create ai model struct to hold info about supported models, Implement additional functions on aimodel types via traits

* Cleanup and rename aistoreinputtypes

* Improvements to preprocessing logic, remove &mut references

Author: Iamdavidonuh

ahnlich/Makefile

@@ -34,6 +34,9 @@ check: ## cargo check
 test: ## cargo test
 	cargo nextest run --no-capture
 
+generate-specs: ## cargo run --bin typegen generate
+	cargo run --bin typegen generate
+
 
 bump-protocol-version: ## Bump project versions. Rules for bumpversion: patch, minor, major.
 	@echo "Bumping Rust project version to $${RULE:-$(BUMP_RULE)}"

ahnlich/ai/Cargo.toml

@@ -29,6 +29,7 @@ ahnlich_client_rs = { path = "../client", version = "*" }
 ahnlich_similarity = { path = "../similarity", version = "*", features = ["serde"] }
 cap.workspace = true
 deadpool.workspace = true
+nonzero_ext = "0.3.0"
 
 
 [dev-dependencies]

ahnlich/ai/src/engine/ai/mod.rs

@@ -0,0 +1,10 @@
+mod models;
+use ahnlich_types::keyval::{StoreInput, StoreKey};
+use models::ModelInfo;
+use std::num::NonZeroUsize;
+
+pub trait AIModelManager {
+    fn embedding_size(&self) -> NonZeroUsize;
+    fn model_ndarray(&self, storeinput: &StoreInput) -> StoreKey;
+    fn model_info(&self) -> ModelInfo;
+}

ahnlich/ai/src/engine/ai/models.rs

@@ -0,0 +1,46 @@
+use crate::engine::ai::AIModelManager;
+use ahnlich_types::{
+    ai::{AIModel, AIStoreInputType},
+    keyval::{StoreInput, StoreKey},
+};
+use ndarray::Array1;
+use nonzero_ext::nonzero;
+use serde::{Deserialize, Serialize};
+use std::num::NonZeroUsize;
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct ModelInfo {
+    pub name: String,
+    pub embedding_size: NonZeroUsize,
+    pub input_type: AIStoreInputType,
+}
+
+impl AIModelManager for AIModel {
+    fn embedding_size(&self) -> NonZeroUsize {
+        self.model_info().embedding_size
+    }
+    // TODO: model ndarray values is based on length of string or vec, so for now make sure strings
+    // or vecs have different lengths
+    fn model_ndarray(&self, storeinput: &StoreInput) -> StoreKey {
+        let length = storeinput.len() as f32;
+        StoreKey(
+            Array1::from_iter(0..self.model_info().embedding_size.into())
+                .mapv(|v| v as f32 * length),
+        )
+    }
+
+    fn model_info(&self) -> ModelInfo {
+        match self {
+            AIModel::Llama3 => ModelInfo {
+                name: String::from("Llama3"),
+                embedding_size: nonzero!(100usize),
+                input_type: AIStoreInputType::RawString,
+            },
+            AIModel::DALLE3 => ModelInfo {
+                name: String::from("DALL.E 3"),
+                embedding_size: nonzero!(300usize),
+                input_type: AIStoreInputType::Image,
+            },
+        }
+    }
+}

ahnlich/ai/src/engine/mod.rs

@@ -1 +1,2 @@
+pub mod ai;
 pub mod store;

ahnlich/ai/src/engine/store.rs

@@ -1,7 +1,8 @@
-use crate::error::AIProxyError;
 use crate::AHNLICH_AI_RESERVED_META_KEY;
-use ahnlich_types::ai::AIModel;
-use ahnlich_types::ai::AIStoreInfo;
+use crate::{engine::ai::AIModelManager, error::AIProxyError};
+use ahnlich_types::ai::{
+    AIModel, AIStoreInfo, AIStoreInputType, ImageAction, PreprocessAction, StringAction,
+};
 use ahnlich_types::keyval::StoreInput;
 use ahnlich_types::keyval::StoreKey;
 use ahnlich_types::keyval::StoreName;
@@ -57,14 +58,18 @@ impl AIStoreHandler {
     pub(crate) fn create_store(
         &self,
         store_name: StoreName,
-
-        model: AIModel,
+        query_model: AIModel,
+        index_model: AIModel,
     ) -> Result<(), AIProxyError> {
         if self
             .stores
             .try_insert(
                 store_name.clone(),
-                Arc::new(AIStore::create(store_name.clone(), model.clone())),
+                Arc::new(AIStore::create(
+                    store_name.clone(),
+                    query_model,
+                    index_model,
+                )),
                 &self.stores.guard(),
             )
             .is_err()
@@ -82,8 +87,9 @@ impl AIStoreHandler {
             .iter(&self.stores.guard())
             .map(|(store_name, store)| AIStoreInfo {
                 name: store_name.clone(),
-                model: store.model.clone(),
-                embedding_size: store.model.embedding_size().into(),
+                query_model: store.query_model.clone(),
+                index_model: store.index_model.clone(),
+                embedding_size: store.index_model.embedding_size().into(),
             })
             .collect()
     }
@@ -100,28 +106,58 @@ impl AIStoreHandler {
     }
 
     /// Converts storeinput into a tuple of storekey and storevalue.
-    /// Fails if the type of storeinput does not match the store type
+    /// Fails if the store input type does not match the store index_type
     #[tracing::instrument(skip(self))]
     pub(crate) fn store_input_to_store_key_val(
         &self,
         store_name: &StoreName,
         store_input: StoreInput,
-        store_value: &StoreValue,
+        store_value: StoreValue,
+        preprocess_action: &PreprocessAction,
     ) -> Result<(StoreKey, StoreValue), AIProxyError> {
         let metadata_key = &*AHNLICH_AI_RESERVED_META_KEY;
         if store_value.contains_key(metadata_key) {
             return Err(AIProxyError::ReservedError(metadata_key.to_string()));
         }
         let store = self.get(store_name)?;
 
-        let store_key = store.model.model_ndarray(&store_input);
-        let metadata_value: MetadataValue = store_input.into();
+        let store_input_type: AIStoreInputType = (&store_input).into();
+        let store_index_model_info = store.index_model.model_info();
+
+        if store_input_type != store_index_model_info.input_type {
+            return Err(AIProxyError::StoreSetTypeMismatchError {
+                store_index_model_type: store_index_model_info.input_type,
+                storeinput_type: store_input_type,
+            });
+        }
+
+        let metadata_value: MetadataValue = store_input.clone().into();
         let mut final_store_value: StdHashMap<MetadataKey, MetadataValue> =
             store_value.clone().into_iter().collect();
         final_store_value.insert(metadata_key.clone(), metadata_value);
+
+        let store_key =
+            self.create_store_key(store_input, &store.index_model, preprocess_action)?;
         return Ok((store_key, final_store_value));
     }
 
+    /// Converts storeinput into a tuple of storekey and storevalue.
+    /// Fails if the store input type does not match the store index_type
+    #[tracing::instrument(skip(self))]
+    pub(crate) fn create_store_key(
+        &self,
+        store_input: StoreInput,
+        index_store_model: &AIModel,
+        preprocess_action: &PreprocessAction,
+    ) -> Result<StoreKey, AIProxyError> {
+        // Process the inner value of a store input and convert it into a ndarray by passing
+        // it into  index model. Create a storekey from ndarray
+        let processed_input =
+            self.preprocess_store_input(preprocess_action, store_input, index_store_model)?;
+        let store_key = index_store_model.model_ndarray(&processed_input);
+        Ok(store_key)
+    }
+
     /// Converts (storekey, storevalue) into (storeinput, storevalue)
     /// by removing the reserved_key from storevalue
     #[tracing::instrument(skip(self))]
@@ -153,7 +189,7 @@ impl AIStoreHandler {
         store_input: &StoreInput,
     ) -> Result<StoreKey, AIProxyError> {
         let store = self.get(store_name)?;
-        Ok(store.model.model_ndarray(store_input))
+        Ok(store.index_model.model_ndarray(store_input))
     }
 
     /// Matches DROPSTORE - Drops a store if exist, else returns an error
@@ -185,20 +221,105 @@ impl AIStoreHandler {
         self.stores.clear(&guard);
         store_length
     }
+
+    #[tracing::instrument(skip(self))]
+    pub(crate) fn preprocess_store_input(
+        &self,
+        process_action: &PreprocessAction,
+        input: StoreInput,
+        index_model: &AIModel,
+    ) -> Result<StoreInput, AIProxyError> {
+        match (process_action, input) {
+            (PreprocessAction::Image(image_action), StoreInput::Image(image_input)) => {
+                // resize image and edit
+                let output = self.process_image(image_input, index_model, image_action)?;
+                Ok(output)
+            }
+            (PreprocessAction::RawString(string_action), StoreInput::RawString(string_input)) => {
+                let output =
+                    self.preprocess_raw_string(string_input, index_model, string_action)?;
+                Ok(output)
+            }
+            (PreprocessAction::RawString(_), StoreInput::Image(_)) => {
+                Err(AIProxyError::PreprocessingMismatchError {
+                    input_type: AIStoreInputType::Image,
+                    preprocess_action: process_action.clone(),
+                })
+            }
+
+            (PreprocessAction::Image(_), StoreInput::RawString(_)) => {
+                Err(AIProxyError::PreprocessingMismatchError {
+                    input_type: AIStoreInputType::RawString,
+                    preprocess_action: process_action.clone(),
+                })
+            }
+        }
+    }
+    fn preprocess_raw_string(
+        &self,
+        input: String,
+        index_model: &AIModel,
+        string_action: &StringAction,
+    ) -> Result<StoreInput, AIProxyError> {
+        // tokenize string, return error if max token
+        //let tokenized_input;
+        let model_embedding_dim = index_model.model_info().embedding_size;
+        if input.len() > model_embedding_dim.into() {
+            if let StringAction::ErrorIfTokensExceed = string_action {
+                return Err(AIProxyError::TokenExceededError {
+                    input_token_size: input.len(),
+                    model_embedding_size: model_embedding_dim.into(),
+                });
+            } else {
+                // truncate raw string
+                // let tokenized_input;
+                let _input = input.as_str()[..model_embedding_dim.into()].to_string();
+            }
+        }
+        Ok(StoreInput::RawString(input))
+    }
+
+    fn process_image(
+        &self,
+        input: Vec<u8>,
+        index_model: &AIModel,
+        image_action: &ImageAction,
+    ) -> Result<StoreInput, AIProxyError> {
+        // process image, return error if max dimensions exceeded
+        // let image_data;
+        let model_embedding_dim = index_model.embedding_size().into();
+        if input.len() > model_embedding_dim {
+            if let ImageAction::ErrorIfDimensionsMismatch = image_action {
+                return Err(AIProxyError::ImageDimensionsMismatchError {
+                    image_dimensions: input.len(),
+                    max_dimensions: model_embedding_dim,
+                });
+            } else {
+                // resize image
+            }
+        }
+        Ok(StoreInput::Image(input))
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct AIStore {
     name: StoreName,
     /// Making use of a concurrent hashmap, we should be able to create an engine that manages stores
-    model: AIModel,
+    query_model: AIModel,
+    index_model: AIModel,
 }
 
 impl AIStore {
-    pub(super) fn create(store_name: StoreName, model: AIModel) -> Self {
+    pub(super) fn create(
+        store_name: StoreName,
+        query_model: AIModel,
+        index_model: AIModel,
+    ) -> Self {
         Self {
             name: store_name,
-            model,
+            query_model,
+            index_model,
         }
     }
 }

ahnlich/ai/src/error.rs

@@ -1,4 +1,7 @@
-use ahnlich_types::keyval::StoreName;
+use ahnlich_types::{
+    ai::{AIStoreInputType, PreprocessAction},
+    keyval::StoreName,
+};
 use thiserror::Error;
 
 #[derive(Error, Debug, Eq, PartialEq, PartialOrd, Ord)]
@@ -7,12 +10,44 @@ pub enum AIProxyError {
     StoreNotFound(StoreName),
     #[error("Store {0} already exists")]
     StoreAlreadyExists(StoreName),
+
     #[error("Proxy Errored with {0} ")]
     StandardError(String),
+
     #[error("Proxy Errored with {0} ")]
     DatabaseClientError(String),
     #[error("Reserved key {0} used")]
     ReservedError(String),
     #[error("Unexpected DB Response {0} ")]
     UnexpectedDBResponse(String),
+    #[error("Cannot Query Using Input. Store expects [{store_query_model_type}], but input type [{storeinput_type}] was provided")]
+    StoreQueryTypeMismatchError {
+        store_query_model_type: AIStoreInputType,
+        storeinput_type: AIStoreInputType,
+    },
+    #[error("Cannot Set Input. Store expects [{store_index_model_type}], input type [{storeinput_type}] was provided")]
+    StoreSetTypeMismatchError {
+        store_index_model_type: AIStoreInputType,
+        storeinput_type: AIStoreInputType,
+    },
+
+    #[error("Max Token Exceeded. Model Expects [{model_embedding_size}], input type was [{input_token_size}] ")]
+    TokenExceededError {
+        model_embedding_size: usize,
+        input_token_size: usize,
+    },
+
+    #[error(
+        "Image Dimensions [{image_dimensions}] exceeds max model dimensions [{max_dimensions}] "
+    )]
+    ImageDimensionsMismatchError {
+        image_dimensions: usize,
+        max_dimensions: usize,
+    },
+
+    #[error("Used [{preprocess_action}] for [{input_type}] type")]
+    PreprocessingMismatchError {
+        input_type: AIStoreInputType,
+        preprocess_action: PreprocessAction,
+    },
 }