Merge pull request #32 from arcmindai/dev

feat(cycles_battery) enhance arcmindai brain and tools canisters to use cycles battery canister for topup every 3 days

scripts/deploy_brain.sh

@@ -6,12 +6,18 @@ if [[ -z "${OPENAI_API_KEY}" ]]; then
  exit 1
 fi
 
+if [[ -z "${BATTERY_API_KEY}" ]]; then
+ echo "BATTERY_API_KEY is unset."
+ exit 1
+fi
+
 # To deplopy locally, update IC_NETWORK to local. To deploy to ic, update IC_NETWORK to ic.
 IC_NETWORK=${IC_NETWORK:-local}
 GPT_MODEL=gpt-4
 
 # Deploy brain canister 
 CONTROLLER_PRINCIPAL=$(dfx canister --network $IC_NETWORK id arcmindai_controller)
+BATTERY_PRINCIAL=$(dfx canister --network $IC_NETWORK id cycles_battery)
 
-echo Deploying brain canister with owner=$CONTROLLER_PRINCIPAL, GPT model=$GPT_MODEL and OPENAI_API_KEY=$OPENAI_API_KEY on $IC_NETWORK
-dfx deploy --network $IC_NETWORK arcmindai_brain --argument "(opt principal \"$CONTROLLER_PRINCIPAL\", \"$OPENAI_API_KEY\", \"$GPT_MODEL\")"
+echo Deploying brain canister with owner=$CONTROLLER_PRINCIPAL, GPT model=$GPT_MODEL and OPENAI_API_KEY=$OPENAI_API_KEY, BATTERY_PRINCIAL=$BATTERY_PRINCIAL, BATTERY_API_KEY=$BATTERY_API_KEY on $IC_NETWORK
+dfx deploy --network $IC_NETWORK arcmindai_brain --argument "(opt principal \"$CONTROLLER_PRINCIPAL\", \"$OPENAI_API_KEY\", \"$GPT_MODEL\", opt \"$BATTERY_API_KEY\", opt principal \"$BATTERY_PRINCIAL\")"

scripts/deploy_tools.sh

@@ -11,14 +11,20 @@ if [[ -z "${GOOGLE_SEARCH_ENGINE_ID}" ]]; then
  exit 1
 fi
 
+if [[ -z "${BATTERY_API_KEY}" ]]; then
+ echo "BATTERY_API_KEY is unset."
+ exit 1
+fi
+
 # To deplopy locally, update IC_NETWORK to local. To deploy to ic, update IC_NETWORK to ic.
 IC_NETWORK=${IC_NETWORK:-local}
 
 CONTROLLER_PRINCIPAL=$(dfx canister --network $IC_NETWORK id arcmindai_controller)
+BATTERY_PRINCIAL=$(dfx canister --network $IC_NETWORK id cycles_battery)
 
 # Deploy tools canister
-echo Deploying tools canister with owner=$CONTROLLER_PRINCIPAL on $IC_NETWORK, GOOGLE_API_KEY=$GOOGLE_API_KEY, GOOGLE_SEARCH_ENGINE_ID=$GOOGLE_SEARCH_ENGINE_ID
-dfx deploy --network $IC_NETWORK arcmindai_tools --argument "(opt principal \"$CONTROLLER_PRINCIPAL\", \"$GOOGLE_API_KEY\", \"$GOOGLE_SEARCH_ENGINE_ID\")"
+echo Deploying tools canister with owner=$CONTROLLER_PRINCIPAL on $IC_NETWORK, GOOGLE_API_KEY=$GOOGLE_API_KEY, GOOGLE_SEARCH_ENGINE_ID=$GOOGLE_SEARCH_ENGINE_ID, BATTERY_PRINCIAL=$BATTERY_PRINCIAL, BATTERY_API_KEY=$BATTERY_API_KEY on $IC_NETWORK
+dfx deploy --network $IC_NETWORK arcmindai_tools --argument "(opt principal \"$CONTROLLER_PRINCIPAL\", \"$GOOGLE_API_KEY\", \"$GOOGLE_SEARCH_ENGINE_ID\", opt \"$BATTERY_API_KEY\", opt principal \"$BATTERY_PRINCIAL\" )"
 
 echo Tools Owner:
 dfx canister --network $IC_NETWORK call arcmindai_tools get_owner

scripts/provision-instance.sh

@@ -64,17 +64,17 @@ cd ../
 pwd
 
 # Deploy brain canister 
-echo Deploying brain canister with owner $CONTROLLER_PRINCIPAL, GPT model $GPT_MODEL and openai_api_key $OPENAI_API_KEY
-dfx deploy --network $IC_NETWORK arcmindai_brain --argument "(opt principal \"$CONTROLLER_PRINCIPAL\", \"$OPENAI_API_KEY\", \"$GPT_MODEL\")"
+echo Deploying brain canister with owner $CONTROLLER_PRINCIPAL, GPT model $GPT_MODEL, openai_api_key $OPENAI_API_KEY, BATTERY_API_KEY=$BATTERY_API_KEY, BATTERY_PRINCIAL=$BATTERY_PRINCIAL on $IC_NETWORK
+dfx deploy --network $IC_NETWORK arcmindai_brain --argument "(opt principal \"$CONTROLLER_PRINCIPAL\", \"$OPENAI_API_KEY\", \"$GPT_MODEL\", opt \"$BATTERY_API_KEY\", opt principal \"$BATTERY_PRINCIAL\")"
 
 echo Brain Owner:
 dfx canister --network $IC_NETWORK call arcmindai_brain get_owner
 
 BRAIN_PRINCIPAL=$(dfx canister --network $IC_NETWORK id arcmindai_brain)
 
 # Deploy tools canister
-echo Deploying tools canister with owner $CONTROLLER_PRINCIPAL on $IC_NETWORK with GOOGLE_API_KEY=$GOOGLE_API_KEY, GOOGLE_SEARCH_ENGINE_ID=$GOOGLE_SEARCH_ENGINE_ID
-dfx deploy --network $IC_NETWORK arcmindai_tools --argument "(opt principal \"$CONTROLLER_PRINCIPAL\", \"$GOOGLE_API_KEY\", \"$GOOGLE_SEARCH_ENGINE_ID\")"
+echo Deploying tools canister with owner $CONTROLLER_PRINCIPAL on $IC_NETWORK with GOOGLE_API_KEY=$GOOGLE_API_KEY, GOOGLE_SEARCH_ENGINE_ID=$GOOGLE_SEARCH_ENGINE_ID, BATTERY_API_KEY=$BATTERY_API_KEY, BATTERY_PRINCIAL=$BATTERY_PRINCIAL on $IC_NETWORK
+dfx deploy --network $IC_NETWORK arcmindai_tools --argument "(opt principal \"$CONTROLLER_PRINCIPAL\", \"$GOOGLE_API_KEY\", \"$GOOGLE_SEARCH_ENGINE_ID\", opt \"$BATTERY_API_KEY\", opt principal \"$BATTERY_PRINCIAL\")"
 
 echo Tools Owner:
 dfx canister --network $IC_NETWORK call arcmindai_tools get_owner

src/arcmindai_brain/Cargo.toml

@@ -12,12 +12,14 @@ crate-type = ["cdylib"]
 candid = "0.8"
 ic-cdk = "0.7"
 ic-cdk-macros = "0.6.0"
+ic-cdk-timers = "0.1.0"
 serde = "1.0.152"
 serde_json = "1.0"
 ciborium = "0.2"
 tiktoken-rs = "0.5.5"
 async-recursion = { version = "1.0.5"}
 
+
 [build-dependencies]
 candid = "0.8"
 ic-cdk = "0.7"

src/arcmindai_brain/arcmindai_brain.did

@@ -1,7 +1,9 @@
 type Result = variant { Ok : vec float32; Err : text };
-service : (opt principal, text, text) -> {
+service : (opt principal, text, text, opt text, opt principal) -> {
  ask : (text, opt text, int8, opt text) -> (text);
+ check_cycles_and_topup : () -> ();
  generate_embeddings : (text, int8, opt text) -> (Result);
+ get_battery_canister : () -> (opt principal) query;
  get_owner : () -> (opt principal) query;
  update_owner : (principal) -> ();
 }

src/arcmindai_brain/src/lib.rs

@@ -1,4 +1,4 @@
-use std::{cell::RefCell, ops::Deref};
+use std::{cell::RefCell, ops::Deref, time::Duration};
 
 use ic_cdk::api::management_canister::http_request::{
  http_request, CanisterHttpRequestArgument, HttpMethod, HttpResponse, TransformArgs,
@@ -37,6 +37,8 @@ pub struct State {
  pub owner: Option<Principal>,
  pub openai_api_key: String,
  pub gpt_model: String,
+ pub battery_api_key: Option<String>,
+ pub battery_canister: Option<Principal>,
 }
 
 // Mutable global state
@@ -49,6 +51,15 @@ const MAX_DEFAULT_TOKENS: usize = 8000;
 const MAX_NUM_RETIRES: i8 = 2;
 const GPT_TEMPERATURE: f32 = 0.5;
 
+// 3 days
+const CYCLES_BALANCE_CHECK_MIN_INTERVAL_SECS: u64 = 60 * 60 * 24 * 3;
+// Cycle usage threshold
+const CYCLES_ONE_TC: u64 = 1_000_000_000_000;
+const CYCLES_THRESHOLD: u64 = 3 * CYCLES_ONE_TC;
+const CYCLES_TOPUP_AMT: u64 = 4 * CYCLES_ONE_TC;
+
+const CYCLES_TOPUP_GROUP: &str = "arcmindai_brain";
+
 // entry function for user to ask questions
 #[update(guard = "assert_owner")]
 #[candid_method(update)]
@@ -255,15 +266,25 @@ fn transform_openai_embeddings(args: TransformArgs) -> HttpResponse {
 // Controller canister must be created with principal
 #[init]
 #[candid_method(init)]
-fn init(owner: Option<Principal>, openai_api_key: String, gpt_model: String) {
+fn init(
+ owner: Option<Principal>,
+ openai_api_key: String,
+ gpt_model: String,
+ battery_api_key: Option<String>,
+ battery_canister: Option<Principal>,
+) {
  let my_owner: Principal = owner.unwrap_or_else(|| api::caller());
  STATE.with(|state| {
  *state.borrow_mut() = State {
  owner: Some(my_owner),
  openai_api_key: openai_api_key,
  gpt_model: gpt_model,
+ battery_api_key: battery_api_key,
+ battery_canister: battery_canister,
  };
  });
+
+ start_cycles_check_timer(CYCLES_BALANCE_CHECK_MIN_INTERVAL_SECS);
 }
 
 #[query]
@@ -278,14 +299,98 @@ pub fn update_owner(new_owner: Principal) {
  STATE.with(|state| {
  let open_api_key = state.borrow().openai_api_key.clone();
  let gpt_model = state.borrow().gpt_model.clone();
+ let battery_api_key = state.borrow().battery_api_key.clone();
+ let battery_canister = state.borrow().battery_canister.clone();
+
  *state.borrow_mut() = State {
  owner: Some(new_owner),
  openai_api_key: open_api_key,
  gpt_model: gpt_model,
+ battery_api_key: battery_api_key,
+ battery_canister: battery_canister,
  };
  });
 }
 
+#[update]
+fn start_cycles_check_timer(secs: u64) {
+ let secs = Duration::from_secs(secs);
+ ic_cdk::println!(
+ "Controller canister: checking its cycles balance and request topup with {secs:?} interval..."
+ );
+
+ ic_cdk_timers::set_timer_interval(secs, || ic_cdk::spawn(check_cycles_and_topup()));
+}
+
+// Check if the cycles balance is below the threshold, and topup from Cycles Battery canister if necessary
+#[update]
+#[candid_method(update)]
+async fn check_cycles_and_topup() {
+ // Get the cycles balance
+ let current_canister_balance = ic_cdk::api::canister_balance();
+
+ // log the cycles balance
+ ic_cdk::println!("Current canister balance: {}", current_canister_balance);
+
+ let battery_api_key: Option<String> =
+ STATE.with(|state| (*state.borrow()).battery_api_key.clone());
+ let battery_canister = STATE.with(|state| (*state.borrow()).battery_canister.unwrap());
+
+ // Make Topup request if the balance is below the threshold
+ if current_canister_balance < CYCLES_THRESHOLD {
+ ic_cdk::println!("Cycles balance is below the threshold");
+
+ let cycles_topup: u64 = CYCLES_TOPUP_AMT;
+ // convert cycles_topup to u128
+ let cycles_topup_input: u128 = cycles_topup as u128;
+
+ let (result,): (Result<(), String>,) = ic_cdk::api::call::call(
+ battery_canister.clone(),
+ "topup_cycles",
+ (
+ CYCLES_TOPUP_GROUP,
+ battery_api_key.unwrap(),
+ cycles_topup_input,
+ current_canister_balance,
+ ),
+ )
+ .await
+ .expect("call to ask failed");
+
+ if result.is_ok() {
+ ic_cdk::println!("Cycles balance topped up by {}", cycles_topup);
+ } else {
+ ic_cdk::println!("Cycles balance topup failed: {}", result.unwrap_err());
+ }
+ } else {
+ ic_cdk::println!("Cycles balance is above the threshold");
+
+ let (result,): (Result<(), String>,) = ic_cdk::api::call::call(
+ battery_canister.clone(),
+ "log_cycles",
+ (
+ CYCLES_TOPUP_GROUP,
+ battery_api_key.unwrap(),
+ current_canister_balance,
+ ),
+ )
+ .await
+ .expect("call to ask failed");
+
+ if result.is_ok() {
+ ic_cdk::println!("Cycles balance logged: {}", current_canister_balance);
+ } else {
+ ic_cdk::println!("Cycles balance log failed: {}", result.unwrap_err());
+ }
+ }
+}
+
+#[query]
+#[candid_method(query)]
+pub fn get_battery_canister() -> Option<Principal> {
+ STATE.with(|state| (*state.borrow()).battery_canister)
+}
+
 #[pre_upgrade]
 fn pre_upgrade() {
  STATE.with(|cell| {
@@ -295,11 +400,23 @@ fn pre_upgrade() {
 }
 
 #[post_upgrade]
-fn post_upgrade() {
+fn post_upgrade(
+ _owner: Option<Principal>,
+ _openai_api_key: String,
+ _gpt_model: String,
+ battery_api_key: Option<String>,
+ battery_canister: Option<Principal>,
+) {
  STATE.with(|cell| {
  *cell.borrow_mut() =
  ciborium::de::from_reader(StableReader::default()).expect("failed to decode state");
- })
+ });
+
+ // Update newly added state in the latest version state using argument
+ STATE.with(|s| {
+ s.borrow_mut().battery_canister = battery_canister;
+ s.borrow_mut().battery_api_key = battery_api_key.clone();
+ });
 }
 
 // ---------------------- Candid declarations did file generator ----------------------

src/arcmindai_tools/Cargo.toml

@@ -12,6 +12,7 @@ crate-type = ["cdylib"]
 candid = "0.8"
 ic-cdk = "0.7"
 ic-cdk-macros = "0.6.0"
+ic-cdk-timers = "0.1.0"
 serde = { version = "1.0.152", features = ["derive"]}
 serde_json = "1.0"
 ciborium = "0.2"

src/arcmindai_tools/arcmindai_tools.did

@@ -1,5 +1,7 @@
-service : (opt principal, text, text) -> {
+service : (opt principal, text, text, opt text, opt principal) -> {
  browse_website : (text) -> (text);
+ check_cycles_and_topup : () -> ();
+ get_battery_canister : () -> (opt principal) query;
  get_owner : () -> (opt principal) query;
  google : (text) -> (text);
  update_owner : (principal) -> ();