diff --git a/.github/workflows/stress_test.yml b/.github/workflows/stress_test.yml index f5575737..ba60ab00 100644 --- a/.github/workflows/stress_test.yml +++ b/.github/workflows/stress_test.yml @@ -32,9 +32,26 @@ jobs: - name: Install Prerequisites run: | brew update - brew install libpq - brew install sqlx-cli - brew install protobuf + if ! brew list libpq &>/dev/null; then + echo "Installing libpq..." + 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 diff --git a/.sqlx/query-441316f35ca8c24bf78167f9fec48e28c05969bbbbe3d0e3d9e1569a375de476.json b/.sqlx/query-441316f35ca8c24bf78167f9fec48e28c05969bbbbe3d0e3d9e1569a375de476.json index 60a66dee..aad756ac 100644 --- a/.sqlx/query-441316f35ca8c24bf78167f9fec48e28c05969bbbbe3d0e3d9e1569a375de476.json +++ b/.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 ] }, diff --git a/.sqlx/query-74de473589a405c3ab567e72a881869321095e2de497b2c1866c547f939c359c.json b/.sqlx/query-74de473589a405c3ab567e72a881869321095e2de497b2c1866c547f939c359c.json index 8e843a10..9abff493 100644 --- a/.sqlx/query-74de473589a405c3ab567e72a881869321095e2de497b2c1866c547f939c359c.json +++ b/.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 ] }, diff --git a/libs/database/src/pg_row.rs b/libs/database/src/pg_row.rs index 84307498..fece72b6 100644 --- a/libs/database/src/pg_row.rs +++ b/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 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, + #[serde(default)] + pub status: i16, } #[derive(Debug, Deserialize, Serialize, Clone)] diff --git a/migrations/20250109142738_blob_metadata_add_file_status.sql b/migrations/20250109142738_blob_metadata_add_file_status.sql new file mode 100644 index 00000000..eb5aa44d --- /dev/null +++ b/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; diff --git a/src/api/file_storage.rs b/src/api/file_storage.rs index 62ea16e4..03fc9b08 100644 --- a/src/api/file_storage.rs +++ b/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() diff --git a/tests/ai_test/chat_test.rs b/tests/ai_test/chat_test.rs index e67ea834..9c92bea8 100644 --- a/tests/ai_test/chat_test.rs +++ b/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 retry_interval = Duration::from_secs(8); + let mut retries = 5; + 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