chore: store blob file status (#1147)

* chore: store blob file status

* chore: fmt

* chore: fmt

.github/workflows/stress_test.yml

@@ -32,9 +32,26 @@ jobs:
       - name: Install Prerequisites
         run: |
           brew update
-          brew install libpq
+          if ! brew list libpq &>/dev/null; then
-          brew install sqlx-cli
+            echo "Installing libpq..."
-          brew install protobuf
+            brew install libpq
+          else
+            echo "libpq is already installed."
+          fi
+
+          if ! brew list sqlx-cli &>/dev/null; then
+            echo "Installing sqlx-cli..."
+            brew install sqlx-cli
+          else
+            echo "sqlx-cli is already installed."
+          fi
+
+          if ! brew list protobuf &>/dev/null; then
+            echo "Installing protobuf..."
+            brew install protobuf
+          else
+            echo "protobuf is already installed."
+          fi
 
       - name: Replace Values in .env
         run: |
@@ -51,6 +68,7 @@ jobs:
 
       - name: Start Docker Compose Services
         run: |
+          docker compose -f docker-compose-dev.yml down
           docker compose -f docker-compose-dev.yml up -d
           ./script/code_gen.sh
           cargo sqlx database create && cargo sqlx migrate run

.sqlx/query-441316f35ca8c24bf78167f9fec48e28c05969bbbbe3d0e3d9e1569a375de476.json

@@ -27,6 +27,11 @@
         "ordinal": 4,
         "name": "modified_at",
         "type_info": "Timestamptz"
+      },
+      {
+        "ordinal": 5,
+        "name": "status",
+        "type_info": "Int2"
       }
     ],
     "parameters": {
@@ -40,6 +45,7 @@
       false,
       false,
       false,
+      false,
       false
     ]
   },

.sqlx/query-74de473589a405c3ab567e72a881869321095e2de497b2c1866c547f939c359c.json

@@ -27,6 +27,11 @@
         "ordinal": 4,
         "name": "modified_at",
         "type_info": "Timestamptz"
+      },
+      {
+        "ordinal": 5,
+        "name": "status",
+        "type_info": "Int2"
       }
     ],
     "parameters": {
@@ -39,6 +44,7 @@
       false,
       false,
       false,
+      false,
       false
     ]
   },

libs/database/src/pg_row.rs

@@ -198,6 +198,23 @@ pub struct AFCollabMemberRow {
   pub permission_id: i64,
 }
 
+#[derive(Serialize, Deserialize, Eq, PartialEq, Debug, Clone)]
+#[repr(i16)]
+pub enum AFBlobStatus {
+  Ok = 0,
+  DallEContentPolicyViolation = 1,
+}
+
+impl From<i16> for AFBlobStatus {
+  fn from(value: i16) -> Self {
+    match value {
+      0 => AFBlobStatus::Ok,
+      1 => AFBlobStatus::DallEContentPolicyViolation,
+      _ => AFBlobStatus::Ok,
+    }
+  }
+}
+
 #[derive(Debug, FromRow, Serialize, Deserialize)]
 pub struct AFBlobMetadataRow {
   pub workspace_id: Uuid,
@@ -205,6 +222,8 @@ pub struct AFBlobMetadataRow {
   pub file_type: String,
   pub file_size: i64,
   pub modified_at: DateTime<Utc>,
+  #[serde(default)]
+  pub status: i16,
 }
 
 #[derive(Debug, Deserialize, Serialize, Clone)]

migrations/20250109142738_blob_metadata_add_file_status.sql

@@ -0,0 +1,3 @@
+-- Add migration script here
+ALTER TABLE af_blob_metadata
+ADD COLUMN status SMALLINT NOT NULL DEFAULT 0;

src/api/file_storage.rs

@@ -25,6 +25,7 @@ use crate::state::AppState;
 use anyhow::anyhow;
 use aws_sdk_s3::primitives::ByteStream;
 use collab_importer::util::FileId;
+use database::pg_row::AFBlobStatus;
 use serde::Deserialize;
 use shared_entity::dto::file_dto::PutFileResponse;
 use shared_entity::dto::workspace_dto::{BlobMetadata, RepeatedBlobMetaData, WorkspaceSpaceUsage};
@@ -369,6 +370,13 @@ async fn get_blob_by_object_key(
   }
 
   let metadata = result.unwrap();
+  match AFBlobStatus::from(metadata.status) {
+    AFBlobStatus::DallEContentPolicyViolation => {
+      return Ok(HttpResponse::UnprocessableEntity().finish());
+    },
+    AFBlobStatus::Ok => {},
+  };
+
   // Check if the file is modified since the last time
   if let Some(modified_since) = req
     .headers()

tests/ai_test/chat_test.rs

@@ -390,8 +390,8 @@ async fn get_text_with_image_message_test() {
   assert_eq!(workspace_id, workspace_id_url);
   assert_eq!(chat_id, chat_id_url);
 
-  let mut retries = 3;
+  let mut retries = 5;
-  let retry_interval = Duration::from_secs(8);
+  let retry_interval = Duration::from_secs(10);
   let mut last_error = None;
 
   // The image will be generated in the background, so we need to retry until it's available