Hook up create_collection as a system process

Author: TotallyNotChase

rholang/src/rust/interpreter/rho_runtime.rs

@@ -25,6 +25,7 @@ use rspace_plus_plus::rspace::tuplespace_interface::Tuplespace;
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
+use crate::rust::interpreter::chromadb_service::ChromaDBService;
 use crate::rust::interpreter::openai_service::OpenAIService;
 use crate::rust::interpreter::system_processes::{BodyRefs, FixedChannels};
 
@@ -815,20 +816,44 @@ fn std_rho_ai_processes() -> Vec<Definition> {
     ]
 }
 
+fn std_rho_chroma_processes() -> Vec<Definition> {
+    vec![
+        Definition {
+            urn: "rho:chroma:collection:new".to_string(),
+            fixed_channel: FixedChannels::chroma_create_collection(),
+            // TODO (chase): How to define overloads?
+            // This function can support 3 or 2 arguments (last one is optional).
+            arity: 3,
+            body_ref: BodyRefs::CHROMA_CREATE_COLLECTION,
+            handler: Box::new(|ctx| {
+                Box::new(move |args| {
+                    let ctx = ctx.clone();
+                    Box::pin(
+                        async move { ctx.system_processes.clone().chroma_create_collection(args).await },
+                    )
+                })
+            }),
+            remainder: None,
+        },
+    ]
+}
+
 fn dispatch_table_creator(
     space: RhoISpace,
     dispatcher: RhoDispatch,
     block_data: Arc<tokio::sync::RwLock<BlockData>>,
     invalid_blocks: InvalidBlocks,
     extra_system_processes: &mut Vec<Definition>,
     openai_service: Arc<tokio::sync::Mutex<OpenAIService>>,
+    chromadb_service: Arc<tokio::sync::Mutex<ChromaDBService>>,
 ) -> RhoDispatchMap {
     let mut dispatch_table = HashMap::new();
 
     for def in std_system_processes().iter_mut().chain(
         std_rho_crypto_processes()
             .iter_mut()
             .chain(std_rho_ai_processes().iter_mut())
+            .chain(std_rho_chroma_processes().iter_mut())
             .chain(extra_system_processes.iter_mut()),
     ) {
         // TODO: Remove cloning every time
@@ -838,6 +863,7 @@ fn dispatch_table_creator(
             block_data.clone(),
             invalid_blocks.clone(),
             openai_service.clone(),
+            chromadb_service.clone(),
         ));
 
         dispatch_table.insert(tuple.0, tuple.1);
@@ -888,6 +914,7 @@ async fn setup_reducer(
     merge_chs: Arc<std::sync::RwLock<HashSet<Par>>>,
     mergeable_tag_name: Par,
     openai_service: Arc<tokio::sync::Mutex<OpenAIService>>,
+    chromadb_service: Arc<tokio::sync::Mutex<ChromaDBService>>,
     cost: _cost,
 ) -> DebruijnInterpreter {
     // println!("\nsetup_reducer");
@@ -906,6 +933,7 @@ async fn setup_reducer(
         invalid_blocks,
         extra_system_processes,
         openai_service,
+        chromadb_service,
     );
 
     let dispatcher = Arc::new(RholangAndScalaDispatcher {
@@ -941,10 +969,12 @@ fn setup_maps_and_refs(
     let system_binding = std_system_processes();
     let rho_crypto_binding = std_rho_crypto_processes();
     let rho_ai_binding = std_rho_ai_processes();
+    let rho_chroma_binding = std_rho_chroma_processes();
     let combined_processes = system_binding
         .iter()
         .chain(rho_crypto_binding.iter())
         .chain(rho_ai_binding.iter())
+        .chain(rho_chroma_binding.iter())
         .chain(extra_system_processes.iter())
         .collect::<Vec<&Definition>>();
 
@@ -996,6 +1026,7 @@ where
     )));
 
     let openai_service = Arc::new(tokio::sync::Mutex::new(OpenAIService::new()));
+    let chromadb_service = Arc::new(tokio::sync::Mutex::new(ChromaDBService::new().await));
     let reducer = setup_reducer(
         charging_rspace,
         block_data_ref.clone(),
@@ -1005,6 +1036,7 @@ where
         merge_chs,
         mergeable_tag_name,
         openai_service,
+        chromadb_service,
         cost,
     )
     .await;
@@ -1101,7 +1133,11 @@ where
 /// # Returns
 ///
 /// A configured `RhoRuntimeImpl` instance ready for executing Rholang code.
-#[tracing::instrument(name = "create-play-runtime", target = "f1r3fly.rholang.runtime", skip_all)]
+#[tracing::instrument(
+    name = "create-play-runtime",
+    target = "f1r3fly.rholang.runtime",
+    skip_all
+)]
 pub async fn create_rho_runtime<T>(
     rspace: T,
     mergeable_tag_name: Par,
@@ -1136,7 +1172,11 @@ where
 /// # Returns
 ///
 /// A configured `RhoRuntimeImpl` instance with replay capabilities.
-#[tracing::instrument(name = "create-replay-runtime", target = "f1r3fly.rholang.runtime", skip_all)]
+#[tracing::instrument(
+    name = "create-replay-runtime",
+    target = "f1r3fly.rholang.runtime",
+    skip_all
+)]
 pub async fn create_replay_rho_runtime<T>(
     rspace: T,
     mergeable_tag_name: Par,
@@ -1197,15 +1237,18 @@ where
     (rho_runtime, replay_rho_runtime)
 }
 
-#[tracing::instrument(name = "create-play-runtime", target = "f1r3fly.rholang.runtime.create-play", skip_all)]
+#[tracing::instrument(
+    name = "create-play-runtime",
+    target = "f1r3fly.rholang.runtime.create-play",
+    skip_all
+)]
 pub async fn create_runtime_from_kv_store(
     stores: RSpaceStore,
     mergeable_tag_name: Par,
     init_registry: bool,
     additional_system_processes: &mut Vec<Definition>,
     matcher: Arc<Box<dyn Match<BindPattern, ListParWithRandom>>>,
 ) -> RhoRuntimeImpl {
-    
     let space: RSpace<Par, BindPattern, ListParWithRandom, TaggedContinuation> =
         RSpace::create(stores, matcher).unwrap();
 

rholang/src/rust/interpreter/system_processes.rs

@@ -1,3 +1,6 @@
+use crate::rust::interpreter::chromadb_service::{ChromaDBService, CollectionMetadata};
+use crate::rust::interpreter::rho_type::Extractor;
+
 use super::contract_call::ContractCall;
 use super::dispatch::RhoDispatch;
 use super::errors::{illegal_argument_error, InterpreterError};
@@ -37,8 +40,10 @@ use std::sync::Arc;
 // NOTE: Not implementing Logger
 pub type RhoSysFunction = Box<
     dyn Fn(
-        (Vec<ListParWithRandom>, bool, Vec<Par>),
-    ) -> Pin<Box<dyn Future<Output = Result<Vec<Par>, InterpreterError>> + Send>> + Send + Sync,
+            (Vec<ListParWithRandom>, bool, Vec<Par>),
+        ) -> Pin<Box<dyn Future<Output = Result<Vec<Par>, InterpreterError>> + Send>>
+        + Send
+        + Sync,
 >;
 pub type RhoDispatchMap = Arc<tokio::sync::RwLock<HashMap<i64, RhoSysFunction>>>;
 pub type Name = Par;
@@ -170,6 +175,15 @@ impl FixedChannels {
     pub fn dev_null() -> Par {
         byte_name(24)
     }
+
+    // ChromaDB section start
+
+    // these bytes may need to change during finalization.
+    pub fn chroma_create_collection() -> Par {
+        byte_name(25)
+    }
+
+    // ChromaDB section end
 }
 
 pub struct BodyRefs;
@@ -196,6 +210,7 @@ impl BodyRefs {
     pub const RANDOM: i64 = 20;
     pub const GRPC_TELL: i64 = 21;
     pub const DEV_NULL: i64 = 22;
+    pub const CHROMA_CREATE_COLLECTION: i64 = 25;
 }
 
 pub fn non_deterministic_ops() -> HashSet<i64> {
@@ -204,6 +219,7 @@ pub fn non_deterministic_ops() -> HashSet<i64> {
         BodyRefs::DALLE3,
         BodyRefs::TEXT_TO_AUDIO,
         BodyRefs::RANDOM,
+        BodyRefs::CHROMA_CREATE_COLLECTION,
     ])
 }
 
@@ -223,6 +239,7 @@ impl ProcessContext {
         block_data: Arc<tokio::sync::RwLock<BlockData>>,
         invalid_blocks: InvalidBlocks,
         openai_service: Arc<tokio::sync::Mutex<OpenAIService>>,
+        chromadb_service: Arc<tokio::sync::Mutex<ChromaDBService>>,
     ) -> Self {
         ProcessContext {
             space: space.clone(),
@@ -234,6 +251,7 @@ impl ProcessContext {
                 space,
                 block_data,
                 openai_service,
+                chromadb_service,
             ),
         }
     }
@@ -246,13 +264,15 @@ pub struct Definition {
     pub body_ref: BodyRef,
     pub handler: Box<
         dyn FnMut(
-            ProcessContext,
-        ) -> Box<
-            dyn Fn(
-                (Vec<ListParWithRandom>, bool, Vec<Par>),
-            )
-                -> Pin<Box<dyn Future<Output = Result<Vec<Par>, InterpreterError>> + Send>> + Send + Sync,
-        > + Send,
+                ProcessContext,
+            ) -> Box<
+                dyn Fn(
+                        (Vec<ListParWithRandom>, bool, Vec<Par>),
+                    )
+                        -> Pin<Box<dyn Future<Output = Result<Vec<Par>, InterpreterError>> + Send>>
+                    + Send
+                    + Sync,
+            > + Send,
     >,
     pub remainder: Remainder,
 }
@@ -265,13 +285,15 @@ impl Definition {
         body_ref: BodyRef,
         handler: Box<
             dyn FnMut(
-                ProcessContext,
-            ) -> Box<
-                dyn Fn(
-                    (Vec<ListParWithRandom>, bool, Vec<Par>),
-                )
-                    -> Pin<Box<dyn Future<Output = Result<Vec<Par>, InterpreterError>> + Send>> + Send + Sync,
-            > + Send,
+                    ProcessContext,
+                ) -> Box<
+                    dyn Fn(
+                            (Vec<ListParWithRandom>, bool, Vec<Par>),
+                        ) -> Pin<
+                            Box<dyn Future<Output = Result<Vec<Par>, InterpreterError>> + Send>,
+                        > + Send
+                        + Sync,
+                > + Send,
         >,
         remainder: Remainder,
     ) -> Self {
@@ -292,9 +314,11 @@ impl Definition {
         BodyRef,
         Box<
             dyn Fn(
-                (Vec<ListParWithRandom>, bool, Vec<Par>),
-            )
-                -> Pin<Box<dyn Future<Output = Result<Vec<Par>, InterpreterError>> + Send>> + Send + Sync,
+                    (Vec<ListParWithRandom>, bool, Vec<Par>),
+                )
+                    -> Pin<Box<dyn Future<Output = Result<Vec<Par>, InterpreterError>> + Send>>
+                + Send
+                + Sync,
         >,
     ) {
         (self.body_ref, (self.handler)(context))
@@ -355,6 +379,7 @@ pub struct SystemProcesses {
     pub space: RhoISpace,
     pub block_data: Arc<tokio::sync::RwLock<BlockData>>,
     openai_service: Arc<tokio::sync::Mutex<OpenAIService>>,
+    chromadb_service: Arc<tokio::sync::Mutex<ChromaDBService>>,
     pretty_printer: PrettyPrinter,
 }
 
@@ -364,12 +389,14 @@ impl SystemProcesses {
         space: RhoISpace,
         block_data: Arc<tokio::sync::RwLock<BlockData>>,
         openai_service: Arc<tokio::sync::Mutex<OpenAIService>>,
+        chromadb_service: Arc<tokio::sync::Mutex<ChromaDBService>>,
     ) -> Self {
         SystemProcesses {
             dispatcher,
             space,
             block_data,
             openai_service,
+            chromadb_service,
             pretty_printer: PrettyPrinter::new(),
         }
     }
@@ -1281,6 +1308,67 @@ impl SystemProcesses {
             Err(illegal_argument_error("casper_invalid_blocks_set"))
         }
     }
+
+    // ChromaDB section start
+
+    /// This supports two overloads:
+    /// - (collection_name: &str, ignore_if_exists: bool)
+    /// - (collection_name: &str, update_if_exists: bool, metadata: CollectionMetadata)
+    pub async fn chroma_create_collection(
+        &self,
+        contract_args: (Vec<ListParWithRandom>, bool, Vec<Par>),
+    ) -> Result<Vec<Par>, InterpreterError> {
+        let Some((produce, is_replay, previous_output, args)) =
+            self.is_contract_call().unapply(contract_args)
+        else {
+            return Err(illegal_argument_error("chroma_create_collection"));
+        };
+
+        let (collection_name, ignore_or_update_if_exists, metadata, ack) = match args.as_slice() {
+            [collection_name_par, update_if_exists_par, metadata_par, ack] => {
+                let (Some(collection_name), Some(update_if_exists), Some(metadata)) = (
+                    RhoString::unapply(collection_name_par),
+                    RhoBoolean::unapply(update_if_exists_par),
+                    <CollectionMetadata as Extractor>::unapply(metadata_par),
+                ) else {
+                    return Err(illegal_argument_error("chroma_create_collection"));
+                };
+                Ok((collection_name, update_if_exists, Some(metadata), ack))
+            }
+            [collection_name_par, ignore_if_exists_par, ack] => {
+                let (Some(collection_name), Some(ignore_if_exists)) = (
+                    RhoString::unapply(collection_name_par),
+                    RhoBoolean::unapply(ignore_if_exists_par),
+                ) else {
+                    return Err(illegal_argument_error("chroma_create_collection"));
+                };
+                Ok((collection_name, ignore_if_exists, None, ack))
+            }
+            _ => Err(illegal_argument_error("chroma_create_collection")),
+        }?;
+
+        // Common piece of code.
+        if is_replay {
+            produce(&previous_output, ack).await?;
+            return Ok(previous_output);
+        }
+
+        let chromadb_service = self.chromadb_service.lock().await;
+        match chromadb_service
+            .create_collection(&collection_name, ignore_or_update_if_exists, metadata)
+            .await
+        {
+            Ok(_) => Ok(vec![]),
+            Err(e) => {
+                // TODO (chase): Is this right? It seems like other service methods do something similar.
+                let p = RhoString::create_par(collection_name);
+                produce(&[p], ack).await?;
+                return Err(e);
+            }
+        }
+    }
+
+    // ChromaDB section end
 }
 
 // See casper/src/test/scala/coop/rchain/casper/helper/RhoSpec.scala
@@ -1398,7 +1486,9 @@ pub fn test_framework_contracts() -> Vec<Definition> {
                 Box::new(move |args| {
                     let sp = sp.clone();
                     let invalid_blocks = invalid_blocks.clone();
-                    Box::pin(async move { sp.casper_invalid_blocks_set(args, &invalid_blocks).await })
+                    Box::pin(
+                        async move { sp.casper_invalid_blocks_set(args, &invalid_blocks).await },
+                    )
                 })
             }),
             remainder: None,