Merge branch 'master' of https://github.com/stakwork/sphinx-swarm

stakwork · Nov 4, 2024 · 3414293 · 3414293
2 parents a44dd89 + 8df973e
commit 3414293
Show file tree

Hide file tree

Showing 8 changed files with 54 additions and 3 deletions.
diff --git a/app/dist/assets/index-9ca5ffe9.js → app/dist/assets/index-d91d7471.js b/app/dist/assets/index-9ca5ffe9.js → app/dist/assets/index-d91d7471.js
diff --git a/app/dist/index.html b/app/dist/index.html
@@ -6,7 +6,7 @@
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>Sphinx Stack</title>
     <link rel="stylesheet" href="g100.css" />
-    <script type="module" crossorigin src="/assets/index-9ca5ffe9.js"></script>
+    <script type="module" crossorigin src="/assets/index-d91d7471.js"></script>
     <link rel="stylesheet" href="/assets/index-121261a7.css">
   </head>
   <body>

diff --git a/docker-compose.yml b/docker-compose.yml
@@ -91,6 +91,7 @@ services:
       - FEATURE_FLAG_TEXT_EMBEDDINGS=$FEATURE_FLAG_TEXT_EMBEDDINGS
       - SUPER_URL=$SUPER_URL
       - SUPER_TOKEN=$SUPER_TOKEN
+      - QUESTION_AND_ANSWER_WORKFLOW_ID=$QUESTION_AND_ANSWER_WORKFLOW_ID
 
 networks:
   sphinx-swarm:

diff --git a/second-brain.yml b/second-brain.yml
@@ -104,6 +104,7 @@ services:
       - WEBPAGE_TEXT_WORKFLOW_ID=$WEBPAGE_TEXT_WORKFLOW_ID
       - SUPER_URL=$SUPER_URL
       - LOCAL_LLAMA=$LOCAL_LLAMA
+      - QUESTION_AND_ANSWER_WORKFLOW_ID=$QUESTION_AND_ANSWER_WORKFLOW_ID
 
 networks:
   sphinx-swarm:

diff --git a/src/bin/super/state.rs b/src/bin/super/state.rs
@@ -4,14 +4,15 @@ use serde::{Deserialize, Serialize};
 use sphinx_swarm::config::{Role, User};
 use sphinx_swarm::secrets;
 
-use crate::util::get_descriptive_instance_type;
+use crate::util::{get_descriptive_instance_type, get_today_dash_date};
 
 #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
 pub struct Super {
     pub stacks: Vec<RemoteStack>,
     pub users: Vec<User>,
     pub jwt_key: String,
     pub bots: Vec<BotCred>,
+    pub ec2_limit: Ec2Limit,
 }
 
 #[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Default)]
@@ -25,6 +26,12 @@ pub struct RemoteStack {
     pub ec2_instance_id: String,
 }
 
+#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Default)]
+pub struct Ec2Limit {
+    pub count: i32,
+    pub date: String,
+}
+
 #[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Default, Clone)]
 pub struct AwsInstanceType {
     pub name: String,
@@ -52,6 +59,7 @@ impl Default for Super {
             users: vec![default_superuser()],
             jwt_key: secrets::random_word(16),
             bots: Vec::new(),
+            ec2_limit: default_ec2_limit(),
         }
     }
 }
@@ -77,6 +85,14 @@ fn default_superuser() -> User {
     }
 }
 
+fn default_ec2_limit() -> Ec2Limit {
+    let today_dash_date = get_today_dash_date();
+    Ec2Limit {
+        count: 0,
+        date: today_dash_date,
+    }
+}
+
 impl Super {
     pub fn remove_tokens(&self) -> Super {
         let stacks = self
@@ -107,6 +123,10 @@ impl Super {
             users: vec![],
             jwt_key: "".to_string(),
             bots: bots,
+            ec2_limit: Ec2Limit {
+                count: 0,
+                date: "".to_string(),
+            },
         }
     }
 

diff --git a/src/bin/super/util.rs b/src/bin/super/util.rs
@@ -11,6 +11,7 @@ use aws_sdk_ec2::types::{
 };
 use aws_sdk_ec2::Client;
 use aws_smithy_types::retry::RetryConfig;
+use chrono::Local;
 use futures_util::TryFutureExt;
 use reqwest::Response;
 use serde_json::Value;
@@ -605,6 +606,22 @@ pub async fn create_swarm_ec2(
     info: &CreateEc2InstanceInfo,
     state: &mut Super,
 ) -> Result<(), Error> {
+    let daily_limit = getenv("EC2_DAILY_LIMIT")
+        .unwrap_or("5".to_string())
+        .parse()
+        .unwrap_or(5);
+
+    let today_date = get_today_dash_date();
+    if today_date == state.ec2_limit.date {
+        if &state.ec2_limit.count < &daily_limit {
+            state.ec2_limit.count = state.ec2_limit.count + 1;
+        } else {
+            return Err(anyhow!("Daily limit for creating Ec2 Instance exceeded"));
+        }
+    } else {
+        state.ec2_limit.date = today_date;
+        state.ec2_limit.count = 1;
+    }
     let mut actual_vanity_address: Option<String> = None;
 
     let instance_type = get_instance(&info.instance_type);
@@ -902,3 +919,7 @@ pub async fn get_config(state: &mut Super) -> Result<Super, Error> {
     let res = state.remove_tokens();
     Ok(res)
 }
+
+pub fn get_today_dash_date() -> String {
+    Local::now().format("%d-%m-%Y").to_string()
+}
diff --git a/src/images/jarvis.rs b/src/images/jarvis.rs
@@ -209,6 +209,12 @@ fn jarvis(
             webpage_text_workflow_id
         ));
     }
+    if let Ok(question_and_answer_workflow_id) = getenv("QUESTION_AND_ANSWER_WORKFLOW_ID") {
+        env.push(format!(
+            "QUESTION_AND_ANSWER_WORKFLOW_ID={}",
+            question_and_answer_workflow_id
+        ));
+    }
     Config {
         image: Some(format!("{}:{}", img, node.version)),
         hostname: Some(domain(&name)),

diff --git a/superadmin.yml b/superadmin.yml
@@ -93,6 +93,8 @@ services:
       - SWARM_TAG_VALUE=$SWARM_TAG_VALUE
       - SWARM_TAG_KEY=$SWARM_TAG_KEY
       - SWARM_UPDATER_PASSWORD=$SWARM_UPDATER_PASSWORD
+      - EC2_DAILY_LIMIT=$EC2_DAILY_LIMIT
+      - QUESTION_AND_ANSWER_WORKFLOW_ID=$QUESTION_AND_ANSWER_WORKFLOW_ID
 
 networks:
   sphinx-swarm: