feat: release/v0.1.13 (#496)

* feat: add usage per day/user/model (#494)

* feat: add usage per day/user/model

* update db to date

* rename timestamp to date

* fix

* add user id to request headers (#495)

---------

Co-authored-by: Martin Stefcek <35243812+Cifko@users.noreply.github.com>

Author: jorgeantonio21

Cargo.lock

@@ -186,8 +186,10 @@ pub fn update_state_manager_fiat(
             user_id,
             estimated_input_amount,
             input_amount,
+            input_tokens,
             estimated_output_amount,
             output_amount,
+            output_tokens,
             model_name,
         })
         .map_err(|e| AtomaProxyError::InternalError {

atoma-proxy/src/server/handlers/mod.rs

@@ -607,6 +607,10 @@ pub async fn confidential_compute_middleware(
         req_parts
             .headers
             .insert(constants::SIGNATURE, signature_header);
+        req_parts.headers.insert(
+            constants::USER_ID,
+            HeaderValue::from_str(&user_id.to_string()).unwrap(),
+        );
         let request_metadata = req_parts
             .extensions
             .get::<RequestMetadataExtension>()
@@ -2514,6 +2518,10 @@ pub mod utils {
         req_parts
             .headers
             .insert(CONTENT_LENGTH, content_length_header);
+        req_parts.headers.insert(
+            constants::USER_ID,
+            HeaderValue::from_str(&user_id.to_string()).unwrap(),
+        );
         let request_model = body_json
             .get(MODEL)
             .and_then(|m| m.as_str())

atoma-proxy/src/server/middleware.rs

@@ -1435,8 +1435,10 @@ pub async fn handle_state_manager_event(
             model_name,
             estimated_input_amount,
             input_amount,
+            input_tokens,
             estimated_output_amount,
             output_amount,
+            output_tokens,
         } => {
             state_manager
                 .state
@@ -1449,6 +1451,17 @@ pub async fn handle_state_manager_event(
                 )
                 .await?;
             if output_amount > 0 {
+                state_manager
+                    .state
+                    .update_per_day_table(
+                        user_id,
+                        model_name.clone(),
+                        input_amount,
+                        input_tokens,
+                        output_amount,
+                        output_tokens,
+                    )
+                    .await?;
                 state_manager
                     .state
                     .update_usage_per_model(user_id, model_name, input_amount, output_amount)

atoma-state/src/handlers.rs

@@ -0,0 +1,11 @@
+CREATE TABLE
+  IF NOT EXISTS usage_per_day (
+    user_id BIGINT NOT NULL,
+    date DATE DEFAULT CURRENT_DATE NOT NULL,
+    model TEXT NOT NULL,
+    input_amount BIGINT NOT NULL DEFAULT 0,
+    input_tokens BIGINT NOT NULL DEFAULT 0,
+    output_amount BIGINT NOT NULL DEFAULT 0,
+    output_tokens BIGINT NOT NULL DEFAULT 0,
+    UNIQUE (user_id, date, model)
+  );

atoma-state/src/migrations/20250528070551_usage-per-day.sql

@@ -4499,6 +4499,67 @@ impl AtomaState {
         Ok(())
     }
 
+    /// Updates the usage per day for a user and model.
+    ///
+    /// This method updates the `usage_per_day` table for the specified user and model.
+    ///
+    /// # Arguments
+    ///
+    /// * `user_id` - The unique identifier of the user.
+    /// * `model` - The name of the model.
+    /// * `input_amount` - The input amount for the model.
+    /// * `input_tokens` - The input tokens for the model.
+    /// * `output_amount` - The output amount for the model.
+    /// * `output_tokens` - The output tokens for the model.
+    ///
+    /// # Returns
+    ///
+    /// - `Result<()>`: A result indicating success (Ok(())) or failure (Err(AtomaStateManagerError)).
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if:
+    /// - The database query fails to execute.
+    ///
+    /// # Example
+    ///
+    /// ```rust,ignore
+    /// use atoma_node::atoma_state::AtomaStateManager;
+    /// use chrono::Utc;
+    /// async fn update_usage_per_day(state_manager: &AtomaStateManager, user_id: i64, model: String, input_amount: i64, input_tokens: i64, output_amount: i64, output_tokens: i64) -> Result<(), AtomaStateManagerError> {
+    ///     state_manager.update_per_day_table(user_id, model, input_amount, input_tokens, output_amount, output_tokens).await
+    /// }
+    /// ```
+    #[instrument(level = "trace", skip(self))]
+    pub async fn update_per_day_table(
+        &self,
+        user_id: i64,
+        model: String,
+        input_amount: i64,
+        input_tokens: i64,
+        output_amount: i64,
+        output_tokens: i64,
+    ) -> Result<()> {
+        sqlx::query(
+            "INSERT INTO usage_per_day (user_id, model, input_amount, input_tokens, output_amount, output_tokens)
+                VALUES ($1, $2, $3, $4, $5, $6)
+                ON CONFLICT (user_id, model, date) DO UPDATE SET
+                    input_amount = usage_per_day.input_amount + EXCLUDED.input_amount,
+                    input_tokens = usage_per_day.input_tokens + EXCLUDED.input_tokens,
+                    output_amount = usage_per_day.output_amount + EXCLUDED.output_amount,
+                    output_tokens = usage_per_day.output_tokens + EXCLUDED.output_tokens",
+        )
+        .bind(user_id)
+        .bind(model)
+        .bind(input_amount)
+        .bind(input_tokens)
+        .bind(output_amount)
+        .bind(output_tokens)
+        .execute(&self.db)
+        .await?;
+        Ok(())
+    }
+
     /// Updates the usage per model for a user.
     ///
     /// This method updates the `usage_per_model` table for the specified user and model.

atoma-state/src/state_manager.rs

@@ -918,9 +918,13 @@ pub enum AtomaAtomaStateManagerEvent {
         estimated_input_amount: i64,
         /// The actual input amount
         input_amount: i64,
+        /// Number of input tokens,
+        input_tokens: i64,
         /// The original estimated output amount
         estimated_output_amount: i64,
         /// The actual output amount
         output_amount: i64,
+        /// Number of output tokens
+        output_tokens: i64,
     },
 }