diff --git a/core/bin/core_api.rs b/core/bin/core_api.rs index 0e920b2ea6b4..572248f2daa5 100644 --- a/core/bin/core_api.rs +++ b/core/bin/core_api.rs @@ -66,7 +66,7 @@ static GLOBAL: Jemalloc = Jemalloc; /// API State struct RunManager { - pending_apps: Vec<(app::App, run::Credentials, run::Secrets)>, + pending_apps: Vec<(app::App, run::Credentials, run::Secrets, bool)>, pending_runs: Vec, } @@ -97,9 +97,17 @@ impl APIState { } } - fn run_app(&self, app: app::App, credentials: run::Credentials, secrets: run::Secrets) { + fn run_app( + &self, + app: app::App, + credentials: run::Credentials, + secrets: run::Secrets, + store_block_result: bool, + ) { let mut run_manager = self.run_manager.lock(); - run_manager.pending_apps.push((app, credentials, secrets)); + run_manager + .pending_apps + .push((app, credentials, secrets, store_block_result)); } async fn stop_loop(&self) { @@ -123,7 +131,7 @@ impl APIState { let mut loop_count = 0; loop { - let apps: Vec<(app::App, run::Credentials, run::Secrets)> = { + let apps: Vec<(app::App, run::Credentials, run::Secrets, bool)> = { let mut manager = self.run_manager.lock(); let apps = manager.pending_apps.drain(..).collect::>(); apps.iter().for_each(|app| { @@ -145,7 +153,15 @@ impl APIState { match app .0 - .run(app.1, app.2, store, databases_store, qdrant_clients, None) + .run( + app.1, + app.2, + store, + databases_store, + qdrant_clients, + None, + app.3, + ) .await { Ok(()) => { @@ -587,6 +603,7 @@ struct RunsCreatePayload { config: run::RunConfig, credentials: run::Credentials, secrets: Vec, + store_block_result: Option, } async fn run_helper( @@ -830,7 +847,12 @@ async fn runs_create( Ok(app) => { // The run is empty for now, we can clone it for the response. let run = app.run_ref().unwrap().clone(); - state.run_app(app, credentials, secrets); + state.run_app( + app, + credentials, + secrets, + payload.store_block_result.unwrap_or(true), + ); ( StatusCode::OK, Json(APIResponse { @@ -915,6 +937,7 @@ async fn runs_create_stream( databases_store, qdrant_clients, Some(tx.clone()), + payload.store_block_result.unwrap_or(true), ) .await { diff --git a/core/src/app.rs b/core/src/app.rs index bf3f07221094..edeb2874c655 100644 --- a/core/src/app.rs +++ b/core/src/app.rs @@ -310,6 +310,7 @@ impl App { databases_store: Box, qdrant_clients: QdrantClients, event_sender: Option>, + store_block_result: bool, ) -> Result<()> { assert!(self.run.is_some()); assert!(self.run_config.is_some()); @@ -731,6 +732,7 @@ impl App { block_idx, &block.block_type(), name, + store_block_result, ) .await?; diff --git a/core/src/stores/postgres.rs b/core/src/stores/postgres.rs index b81f9b38818e..ac55702f8540 100644 --- a/core/src/stores/postgres.rs +++ b/core/src/stores/postgres.rs @@ -776,6 +776,7 @@ impl Store for PostgresStore { block_idx: usize, block_type: &BlockType, block_name: &String, + store_block_result: bool, ) -> Result<()> { let traces = run .traces @@ -796,8 +797,16 @@ impl Store for PostgresStore { .iter() .enumerate() .map(|(map_idx, execution)| { - let execution_json = serde_json::to_string(&execution)?; - + let execution_json = match store_block_result { + true => serde_json::to_string(&execution)?, + false => serde_json::to_string( + &(BlockExecution { + value: None, + error: execution.error.clone(), + meta: execution.meta.clone(), + }), + )?, + }; Ok(( block_idx, block_type.clone(), diff --git a/core/src/stores/store.rs b/core/src/stores/store.rs index a187e6acad63..032840cee90a 100644 --- a/core/src/stores/store.rs +++ b/core/src/stores/store.rs @@ -140,6 +140,7 @@ pub trait Store { block_idx: usize, block_type: &BlockType, block_name: &String, + store_block_result: bool, ) -> Result<()>; async fn load_run( diff --git a/front/pages/api/v1/w/[wId]/spaces/[spaceId]/apps/[aId]/runs/index.ts b/front/pages/api/v1/w/[wId]/spaces/[spaceId]/apps/[aId]/runs/index.ts index eda77ecac621..e2b72f3e9048 100644 --- a/front/pages/api/v1/w/[wId]/spaces/[spaceId]/apps/[aId]/runs/index.ts +++ b/front/pages/api/v1/w/[wId]/spaces/[spaceId]/apps/[aId]/runs/index.ts @@ -22,7 +22,7 @@ import { withPublicAPIAuthentication } from "@app/lib/api/auth_wrappers"; import apiConfig from "@app/lib/api/config"; import { getDustAppSecrets } from "@app/lib/api/dust_app_secrets"; import { withResourceFetchingFromRoute } from "@app/lib/api/resource_wrappers"; -import type { Authenticator } from "@app/lib/auth"; +import { getFeatureFlags, type Authenticator } from "@app/lib/auth"; import { AppResource } from "@app/lib/resources/app_resource"; import type { RunUsageType } from "@app/lib/resources/run_resource"; import { RunResource } from "@app/lib/resources/run_resource"; @@ -276,6 +276,9 @@ async function handler( } } + const flags = await getFeatureFlags(owner); + const storeBlockResult = !flags.includes("disable_run_logs"); + logger.info( { workspace: { @@ -299,6 +302,7 @@ async function handler( credentials, secrets, isSystemKey: auth.isSystemKey(), + storeBlockResult, } ); diff --git a/front/pages/api/w/[wId]/spaces/[spaceId]/apps/[aId]/runs/index.ts b/front/pages/api/w/[wId]/spaces/[spaceId]/apps/[aId]/runs/index.ts index 46cc3972f552..27bc5448fc33 100644 --- a/front/pages/api/w/[wId]/spaces/[spaceId]/apps/[aId]/runs/index.ts +++ b/front/pages/api/w/[wId]/spaces/[spaceId]/apps/[aId]/runs/index.ts @@ -7,7 +7,7 @@ import { withSessionAuthenticationForWorkspace } from "@app/lib/api/auth_wrapper import config from "@app/lib/api/config"; import { getDustAppSecrets } from "@app/lib/api/dust_app_secrets"; import { withResourceFetchingFromRoute } from "@app/lib/api/resource_wrappers"; -import { Authenticator } from "@app/lib/auth"; +import { Authenticator, getFeatureFlags } from "@app/lib/auth"; import type { SessionWithUser } from "@app/lib/iam/provider"; import { AppResource } from "@app/lib/resources/app_resource"; import { RunResource } from "@app/lib/resources/run_resource"; @@ -123,6 +123,9 @@ async function handler( ); const inputDataset = inputConfigEntry ? inputConfigEntry.dataset : null; + const flags = await getFeatureFlags(owner); + const storeBlockResult = !flags.includes("disable_run_logs"); + const dustRun = await coreAPI.createRun(owner, auth.groups(), { projectId: app.dustAPIProjectId, runType: "local", @@ -134,6 +137,7 @@ async function handler( config: { blocks: config }, credentials: credentialsFromProviders(providers), secrets, + storeBlockResult, }); if (dustRun.isErr()) { diff --git a/types/src/front/lib/core_api.ts b/types/src/front/lib/core_api.ts index 74d9e6f895f4..8ced412e6ebb 100644 --- a/types/src/front/lib/core_api.ts +++ b/types/src/front/lib/core_api.ts @@ -103,6 +103,7 @@ type CoreAPICreateRunParams = { credentials: CredentialsType; secrets: DustAppSecretType[]; isSystemKey?: boolean; + storeBlockResult?: boolean; }; type GetDatasetResponse = { @@ -306,6 +307,7 @@ export class CoreAPI { credentials, secrets, isSystemKey, + storeBlockResult = true, }: CoreAPICreateRunParams ): Promise> { const response = await this._fetchWithError( @@ -327,6 +329,7 @@ export class CoreAPI { config: config, credentials: credentials, secrets: secrets, + store_block_result: storeBlockResult, }), } ); @@ -348,6 +351,7 @@ export class CoreAPI { credentials, secrets, isSystemKey, + storeBlockResult = true, }: CoreAPICreateRunParams ): Promise< CoreAPIResponse<{ @@ -374,6 +378,7 @@ export class CoreAPI { config: config, credentials: credentials, secrets: secrets, + store_block_result: storeBlockResult, }), } ); diff --git a/types/src/shared/feature_flags.ts b/types/src/shared/feature_flags.ts index c1e489248e0b..59110a362146 100644 --- a/types/src/shared/feature_flags.ts +++ b/types/src/shared/feature_flags.ts @@ -11,6 +11,7 @@ export const WHITELISTABLE_FEATURES = [ "openai_o1_high_reasoning_feature", "index_private_slack_channel", "conversations_jit_actions", + "disable_run_logs", "labs_trackers", ] as const; export type WhitelistableFeature = (typeof WHITELISTABLE_FEATURES)[number];