Chore/refractor (#4454)

* markdown files and env examples cleanup

* components update

* update jsonlines description

* server refractor

* update telemetry

* add execute custom node

* add ui refractor

* add username and password authenticate

* correctly retrieve past images in agentflowv2

* disable e2e temporarily

* add existing username and password authenticate

* update migration to default workspace

* update todo

* blob storage migrating

* throw error on agent tool call error

* add missing execution import

* add referral

* chore: add error message when importData is undefined

* migrate api keys to db

* fix: data too long for column executionData

* migrate api keys from json to db at init

* add info on account setup

* update docstore missing fields

---------

Co-authored-by: chungyau97 <chungyau97@gmail.com>

packages/components/src/Interface.Evaluation.ts

@@ -0,0 +1,43 @@
+// Evaluation Related Interfaces
+export interface IDataset {
+    id: string
+    name: string
+    createdDate: Date
+    updatedDate: Date
+}
+export interface IDatasetRow {
+    id: string
+    datasetId: string
+    input: string
+    output: string
+    updatedDate: Date
+}
+
+export enum EvaluationStatus {
+    PENDING = 'pending',
+    COMPLETED = 'completed'
+}
+export interface IEvaluation {
+    id: string
+    name: string
+    chatflowId: string
+    chatflowName: string
+    datasetId: string
+    datasetName: string
+    evaluationType: string
+    average_metrics: string
+    status: string
+    runDate: Date
+}
+
+export interface IEvaluationRun {
+    id: string
+    evaluationId: string
+    input: string
+    expectedOutput: string
+    actualOutput: string
+    metrics: string
+    runDate: Date
+    reasoning: string
+    score: number
+}

packages/components/src/Interface.ts

@@ -414,11 +414,14 @@ export interface IVisionChatModal {
     revertToOriginalModel(): void
     setMultiModalOption(multiModalOption: IMultiModalOption): void
 }
+
 export interface IStateWithMessages extends ICommonObject {
     messages: BaseMessage[]
     [key: string]: any
 }
 
+export * from './Interface.Evaluation'
+
 export interface IServerSideEventStreamer {
     streamStartEvent(chatId: string, data: any): void
     streamTokenEvent(chatId: string, data: string): void

packages/components/src/MetricsLogger.ts

@@ -0,0 +1,135 @@
+import { BaseTracer, Run } from '@langchain/core/tracers/base'
+import { Logger } from 'winston'
+import { AgentRun, elapsed, tryJsonStringify } from './handler'
+
+export class MetricsLogger extends BaseTracer {
+    name = 'console_callback_handler' as const
+    logger: Logger
+    orgId?: string
+
+    protected persistRun(_run: Run) {
+        return Promise.resolve()
+    }
+
+    constructor(logger: Logger, orgId?: string) {
+        super()
+        this.logger = logger
+        this.orgId = orgId
+    }
+
+    // utility methods
+
+    getParents(run: Run) {
+        const parents: Run[] = []
+        let currentRun = run
+        while (currentRun.parent_run_id) {
+            const parent = this.runMap.get(currentRun.parent_run_id)
+            if (parent) {
+                parents.push(parent)
+                currentRun = parent
+            } else {
+                break
+            }
+        }
+        return parents
+    }
+
+    getBreadcrumbs(run: Run) {
+        const parents = this.getParents(run).reverse()
+        const string = [...parents, run]
+            .map((parent) => {
+                const name = `${parent.execution_order}:${parent.run_type}:${parent.name}`
+                return name
+            })
+            .join(' > ')
+        return string
+    }
+
+    // logging methods
+
+    onChainStart(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[${this.orgId}]: [chain/start] [${crumbs}] Entering Chain run with input: ${tryJsonStringify(run.inputs, '[inputs]')}`
+        )
+    }
+
+    onChainEnd(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[${this.orgId}]: [chain/end] [${crumbs}] [${elapsed(run)}] Exiting Chain run with output: ${tryJsonStringify(
+                run.outputs,
+                '[outputs]'
+            )}`
+        )
+    }
+
+    onChainError(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[${this.orgId}]: [chain/error] [${crumbs}] [${elapsed(run)}] Chain run errored with error: ${tryJsonStringify(
+                run.error,
+                '[error]'
+            )}`
+        )
+    }
+
+    onLLMStart(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        const inputs = 'prompts' in run.inputs ? { prompts: (run.inputs.prompts as string[]).map((p) => p.trim()) } : run.inputs
+        this.logger.verbose(`[${this.orgId}]: [llm/start] [${crumbs}] Entering LLM run with input: ${tryJsonStringify(inputs, '[inputs]')}`)
+    }
+
+    onLLMEnd(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[${this.orgId}]: [llm/end] [${crumbs}] [${elapsed(run)}] Exiting LLM run with output: ${tryJsonStringify(
+                run.outputs,
+                '[response]'
+            )}`
+        )
+    }
+
+    onLLMError(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[${this.orgId}]: [llm/error] [${crumbs}] [${elapsed(run)}] LLM run errored with error: ${tryJsonStringify(
+                run.error,
+                '[error]'
+            )}`
+        )
+    }
+
+    onToolStart(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(`[${this.orgId}]: [tool/start] [${crumbs}] Entering Tool run with input: "${run.inputs.input?.trim()}"`)
+    }
+
+    onToolEnd(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[${this.orgId}]: [tool/end] [${crumbs}] [${elapsed(run)}] Exiting Tool run with output: "${run.outputs?.output?.trim()}"`
+        )
+    }
+
+    onToolError(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[${this.orgId}]: [tool/error] [${crumbs}] [${elapsed(run)}] Tool run errored with error: ${tryJsonStringify(
+                run.error,
+                '[error]'
+            )}`
+        )
+    }
+
+    onAgentAction(run: Run) {
+        const agentRun = run as AgentRun
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[${this.orgId}]: [agent/action] [${crumbs}] Agent selected action: ${tryJsonStringify(
+                agentRun.actions[agentRun.actions.length - 1],
+                '[action]'
+            )}`
+        )
+    }
+}

packages/components/src/followUpPrompts.ts

@@ -36,6 +36,7 @@ export const generateFollowUpPrompts = async (
                     model: providerConfig.modelName,
                     temperature: parseFloat(`${providerConfig.temperature}`)
                 })
+                // @ts-ignore
                 const structuredLLM = llm.withStructuredOutput(FollowUpPromptType)
                 const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
                 return structuredResponse

packages/components/src/handler.ts

@@ -25,6 +25,8 @@ import { AgentAction } from '@langchain/core/agents'
 import { LunaryHandler } from '@langchain/community/callbacks/handlers/lunary'
 
 import { getCredentialData, getCredentialParam, getEnvironmentVariable } from './utils'
+import { EvaluationRunTracer } from '../evaluation/EvaluationRunTracer'
+import { EvaluationRunTracerLlama } from '../evaluation/EvaluationRunTracerLlama'
 import { ICommonObject, IDatabaseEntity, INodeData, IServerSideEventStreamer } from './Interface'
 import { LangWatch, LangWatchSpan, LangWatchTrace, autoconvertTypedValues } from 'langwatch'
 import { DataSource } from 'typeorm'
@@ -32,7 +34,7 @@ import { ChatGenerationChunk } from '@langchain/core/outputs'
 import { AIMessageChunk, BaseMessageLike } from '@langchain/core/messages'
 import { Serialized } from '@langchain/core/load/serializable'
 
-interface AgentRun extends Run {
+export interface AgentRun extends Run {
     actions: AgentAction[]
 }
 
@@ -173,7 +175,7 @@ function tryGetJsonSpaces() {
     }
 }
 
-function tryJsonStringify(obj: unknown, fallback: string) {
+export function tryJsonStringify(obj: unknown, fallback: string) {
     try {
         return JSON.stringify(obj, null, tryGetJsonSpaces())
     } catch (err) {
@@ -181,7 +183,7 @@ function tryJsonStringify(obj: unknown, fallback: string) {
     }
 }
 
-function elapsed(run: Run): string {
+export function elapsed(run: Run): string {
     if (!run.end_time) return ''
     const elapsed = run.end_time - run.start_time
     if (elapsed < 1000) {
@@ -193,14 +195,16 @@ function elapsed(run: Run): string {
 export class ConsoleCallbackHandler extends BaseTracer {
     name = 'console_callback_handler' as const
     logger: Logger
+    orgId?: string
 
     protected persistRun(_run: Run) {
         return Promise.resolve()
     }
 
-    constructor(logger: Logger) {
+    constructor(logger: Logger, orgId?: string) {
         super()
         this.logger = logger
+        this.orgId = orgId
         if (getEnvironmentVariable('DEBUG') === 'true') {
             logger.level = getEnvironmentVariable('LOG_LEVEL') ?? 'info'
         }
@@ -235,57 +239,76 @@ export class ConsoleCallbackHandler extends BaseTracer {
     onChainStart(run: Run) {
         const crumbs = this.getBreadcrumbs(run)
 
-        this.logger.verbose(`[chain/start] [${crumbs}] Entering Chain run with input: ${tryJsonStringify(run.inputs, '[inputs]')}`)
+        this.logger.verbose(
+            `[${this.orgId}]: [chain/start] [${crumbs}] Entering Chain run with input: ${tryJsonStringify(run.inputs, '[inputs]')}`
+        )
     }
 
     onChainEnd(run: Run) {
         const crumbs = this.getBreadcrumbs(run)
         this.logger.verbose(
-            `[chain/end] [${crumbs}] [${elapsed(run)}] Exiting Chain run with output: ${tryJsonStringify(run.outputs, '[outputs]')}`
+            `[${this.orgId}]: [chain/end] [${crumbs}] [${elapsed(run)}] Exiting Chain run with output: ${tryJsonStringify(
+                run.outputs,
+                '[outputs]'
+            )}`
         )
     }
 
     onChainError(run: Run) {
         const crumbs = this.getBreadcrumbs(run)
         this.logger.verbose(
-            `[chain/error] [${crumbs}] [${elapsed(run)}] Chain run errored with error: ${tryJsonStringify(run.error, '[error]')}`
+            `[${this.orgId}]: [chain/error] [${crumbs}] [${elapsed(run)}] Chain run errored with error: ${tryJsonStringify(
+                run.error,
+                '[error]'
+            )}`
         )
     }
 
     onLLMStart(run: Run) {
         const crumbs = this.getBreadcrumbs(run)
         const inputs = 'prompts' in run.inputs ? { prompts: (run.inputs.prompts as string[]).map((p) => p.trim()) } : run.inputs
-        this.logger.verbose(`[llm/start] [${crumbs}] Entering LLM run with input: ${tryJsonStringify(inputs, '[inputs]')}`)
+        this.logger.verbose(`[${this.orgId}]: [llm/start] [${crumbs}] Entering LLM run with input: ${tryJsonStringify(inputs, '[inputs]')}`)
     }
 
     onLLMEnd(run: Run) {
         const crumbs = this.getBreadcrumbs(run)
         this.logger.verbose(
-            `[llm/end] [${crumbs}] [${elapsed(run)}] Exiting LLM run with output: ${tryJsonStringify(run.outputs, '[response]')}`
+            `[${this.orgId}]: [llm/end] [${crumbs}] [${elapsed(run)}] Exiting LLM run with output: ${tryJsonStringify(
+                run.outputs,
+                '[response]'
+            )}`
         )
     }
 
     onLLMError(run: Run) {
         const crumbs = this.getBreadcrumbs(run)
         this.logger.verbose(
-            `[llm/error] [${crumbs}] [${elapsed(run)}] LLM run errored with error: ${tryJsonStringify(run.error, '[error]')}`
+            `[${this.orgId}]: [llm/error] [${crumbs}] [${elapsed(run)}] LLM run errored with error: ${tryJsonStringify(
+                run.error,
+                '[error]'
+            )}`
         )
     }
 
     onToolStart(run: Run) {
         const crumbs = this.getBreadcrumbs(run)
-        this.logger.verbose(`[tool/start] [${crumbs}] Entering Tool run with input: "${run.inputs.input?.trim()}"`)
+        this.logger.verbose(`[${this.orgId}]: [tool/start] [${crumbs}] Entering Tool run with input: "${run.inputs.input?.trim()}"`)
     }
 
     onToolEnd(run: Run) {
         const crumbs = this.getBreadcrumbs(run)
-        this.logger.verbose(`[tool/end] [${crumbs}] [${elapsed(run)}] Exiting Tool run with output: "${run.outputs?.output?.trim()}"`)
+        this.logger.verbose(
+            `[${this.orgId}]: [tool/end] [${crumbs}] [${elapsed(run)}] Exiting Tool run with output: "${run.outputs?.output?.trim()}"`
+        )
     }
 
     onToolError(run: Run) {
         const crumbs = this.getBreadcrumbs(run)
         this.logger.verbose(
-            `[tool/error] [${crumbs}] [${elapsed(run)}] Tool run errored with error: ${tryJsonStringify(run.error, '[error]')}`
+            `[${this.orgId}]: [tool/error] [${crumbs}] [${elapsed(run)}] Tool run errored with error: ${tryJsonStringify(
+                run.error,
+                '[error]'
+            )}`
         )
     }
 
@@ -293,7 +316,7 @@ export class ConsoleCallbackHandler extends BaseTracer {
         const agentRun = run as AgentRun
         const crumbs = this.getBreadcrumbs(run)
         this.logger.verbose(
-            `[agent/action] [${crumbs}] Agent selected action: ${tryJsonStringify(
+            `[${this.orgId}]: [agent/action] [${crumbs}] Agent selected action: ${tryJsonStringify(
                 agentRun.actions[agentRun.actions.length - 1],
                 '[action]'
             )}`
@@ -396,6 +419,7 @@ export class CustomChainHandler extends BaseCallbackHandler {
     }
 }
 
+/*TODO - Add llamaIndex tracer to non evaluation runs*/
 class ExtendedLunaryHandler extends LunaryHandler {
     chatId: string
     appDataSource: DataSource
@@ -550,6 +574,13 @@ export const additionalCallbacks = async (nodeData: INodeData, options: ICommonO
                     const handler = new ExtendedLunaryHandler(lunaryFields)
 
                     callbacks.push(handler)
+                } else if (provider === 'evaluation') {
+                    if (options.llamaIndex) {
+                        new EvaluationRunTracerLlama(options.evaluationRunId)
+                    } else {
+                        const evaluationHandler = new EvaluationRunTracer(options.evaluationRunId)
+                        callbacks.push(evaluationHandler)
+                    }
                 } else if (provider === 'langWatch') {
                     const langWatchApiKey = getCredentialParam('langWatchApiKey', credentialData, nodeData)
                     const langWatchEndpoint = getCredentialParam('langWatchEndpoint', credentialData, nodeData)

packages/components/src/index.ts

@@ -9,6 +9,7 @@ export * from './utils'
 export * from './speechToText'
 export * from './storageUtils'
 export * from './handler'
+export * from '../evaluation/EvaluationRunner'
 export * from './followUpPrompts'
 export * from './validator'
 export * from './agentflowv2Generator'

packages/components/src/modelLoader.ts

@@ -76,6 +76,66 @@ const getModelConfig = async (category: MODEL_TYPE, name: string) => {
     }
 }
 
+export const getModelConfigByModelName = async (category: MODEL_TYPE, provider: string | undefined, name: string | undefined) => {
+    const modelFile = process.env.MODEL_LIST_CONFIG_JSON || MASTER_MODEL_LIST
+
+    if (!modelFile) {
+        throw new Error('MODEL_LIST_CONFIG_JSON not set')
+    }
+    if (isValidUrl(modelFile)) {
+        try {
+            const resp = await axios.get(modelFile)
+            if (resp.status === 200 && resp.data) {
+                const models = resp.data
+                const categoryModels = models[category]
+                // each element of categoryModels is an object, with an array of models (models) and regions (regions)
+                // check if the name is in models
+                return getSpecificModelFromCategory(categoryModels, provider, name)
+            } else {
+                throw new Error('Error fetching model list')
+            }
+        } catch (e) {
+            const models = await fs.promises.readFile(getModelsJSONPath(), 'utf8')
+            if (models) {
+                const categoryModels = JSON.parse(models)[category]
+                return getSpecificModelFromCategory(categoryModels, provider, name)
+            }
+            return {}
+        }
+    } else {
+        try {
+            if (fs.existsSync(modelFile)) {
+                const models = await fs.promises.readFile(modelFile, 'utf8')
+                if (models) {
+                    const categoryModels = JSON.parse(models)[category]
+                    return getSpecificModelFromCategory(categoryModels, provider, name)
+                }
+            }
+            return {}
+        } catch (e) {
+            const models = await fs.promises.readFile(getModelsJSONPath(), 'utf8')
+            if (models) {
+                const categoryModels = JSON.parse(models)[category]
+                return getSpecificModelFromCategory(categoryModels, provider, name)
+            }
+            return {}
+        }
+    }
+}
+
+const getSpecificModelFromCategory = (categoryModels: any, provider: string | undefined, name: string | undefined) => {
+    for (const cm of categoryModels) {
+        if (cm.models && cm.name.toLowerCase() === provider?.toLowerCase()) {
+            for (const m of cm.models) {
+                if (m.name === name) {
+                    return m
+                }
+            }
+        }
+    }
+    return undefined
+}
+
 export const getModels = async (category: MODEL_TYPE, name: string) => {
     const returnData: INodeOptionsValue[] = []
     try {

packages/components/src/multiModalUtils.ts

@@ -16,7 +16,7 @@ export const addImagesToMessages = async (
             for (const upload of imageUploads) {
                 let bf = upload.data
                 if (upload.type == 'stored-file') {
-                    const contents = await getFileFromStorage(upload.name, options.chatflowid, options.chatId)
+                    const contents = await getFileFromStorage(upload.name, options.orgId, options.chatflowid, options.chatId)
                     // as the image is stored in the server, read the file and convert it to base64
                     bf = 'data:' + upload.mime + ';base64,' + contents.toString('base64')
 

packages/components/src/speechToText.ts

@@ -18,7 +18,7 @@ export const convertSpeechToText = async (upload: IFileUpload, speechToTextConfi
     if (speechToTextConfig) {
         const credentialId = speechToTextConfig.credentialId as string
         const credentialData = await getCredentialData(credentialId ?? '', options)
-        const audio_file = await getFileFromStorage(upload.name, options.chatflowid, options.chatId)
+        const audio_file = await getFileFromStorage(upload.name, options.orgId, options.chatflowid, options.chatId)
 
         switch (speechToTextConfig.name) {
             case SpeechToTextType.OPENAI_WHISPER: {

packages/components/src/storageUtils.ts

@@ -5,6 +5,7 @@ import {
     GetObjectCommand,
     ListObjectsV2Command,
     PutObjectCommand,
+    ListObjectsCommand,
     S3Client,
     S3ClientConfig
 } from '@aws-sdk/client-s3'
@@ -13,7 +14,32 @@ import { Readable } from 'node:stream'
 import { getUserHome } from './utils'
 import sanitize from 'sanitize-filename'
 
-export const addBase64FilesToStorage = async (fileBase64: string, chatflowid: string, fileNames: string[]) => {
+const dirSize = async (directoryPath: string) => {
+    let totalSize = 0
+
+    async function calculateSize(itemPath: string) {
+        const stats = await fs.promises.stat(itemPath)
+
+        if (stats.isFile()) {
+            totalSize += stats.size
+        } else if (stats.isDirectory()) {
+            const files = await fs.promises.readdir(itemPath)
+            for (const file of files) {
+                await calculateSize(path.join(itemPath, file))
+            }
+        }
+    }
+
+    await calculateSize(directoryPath)
+    return totalSize
+}
+
+export const addBase64FilesToStorage = async (
+    fileBase64: string,
+    chatflowid: string,
+    fileNames: string[],
+    orgId: string
+): Promise<{ path: string; totalSize: number }> => {
     const storageType = getStorageType()
     if (storageType === 's3') {
         const { s3Client, Bucket } = getS3Config()
@@ -24,8 +50,8 @@ export const addBase64FilesToStorage = async (fileBase64: string, chatflowid: st
         const mime = splitDataURI[0].split(':')[1].split(';')[0]
 
         const sanitizedFilename = _sanitizeFilename(filename)
+        const Key = orgId + '/' + chatflowid + '/' + sanitizedFilename
 
-        const Key = chatflowid + '/' + sanitizedFilename
         const putObjCmd = new PutObjectCommand({
             Bucket,
             Key,
@@ -36,7 +62,9 @@ export const addBase64FilesToStorage = async (fileBase64: string, chatflowid: st
         await s3Client.send(putObjCmd)
 
         fileNames.push(sanitizedFilename)
-        return 'FILE-STORAGE::' + JSON.stringify(fileNames)
+        const totalSize = await getS3StorageSize(orgId)
+
+        return { path: 'FILE-STORAGE::' + JSON.stringify(fileNames), totalSize: totalSize / 1024 / 1024 }
     } else if (storageType === 'gcs') {
         const { bucket } = getGcsClient()
         const splitDataURI = fileBase64.split(',')
@@ -55,9 +83,11 @@ export const addBase64FilesToStorage = async (fileBase64: string, chatflowid: st
                 .end(bf)
         })
         fileNames.push(sanitizedFilename)
-        return 'FILE-STORAGE::' + JSON.stringify(fileNames)
+        const totalSize = await getGCSStorageSize(orgId)
+
+        return { path: 'FILE-STORAGE::' + JSON.stringify(fileNames), totalSize: totalSize / 1024 / 1024 }
     } else {
-        const dir = path.join(getStoragePath(), chatflowid)
+        const dir = path.join(getStoragePath(), orgId, chatflowid)
         if (!fs.existsSync(dir)) {
             fs.mkdirSync(dir, { recursive: true })
         }
@@ -68,13 +98,22 @@ export const addBase64FilesToStorage = async (fileBase64: string, chatflowid: st
         const sanitizedFilename = _sanitizeFilename(filename)
 
         const filePath = path.join(dir, sanitizedFilename)
+
         fs.writeFileSync(filePath, bf)
         fileNames.push(sanitizedFilename)
-        return 'FILE-STORAGE::' + JSON.stringify(fileNames)
+
+        const totalSize = await dirSize(path.join(getStoragePath(), orgId))
+        return { path: 'FILE-STORAGE::' + JSON.stringify(fileNames), totalSize: totalSize / 1024 / 1024 }
     }
 }
 
-export const addArrayFilesToStorage = async (mime: string, bf: Buffer, fileName: string, fileNames: string[], ...paths: string[]) => {
+export const addArrayFilesToStorage = async (
+    mime: string,
+    bf: Buffer,
+    fileName: string,
+    fileNames: string[],
+    ...paths: string[]
+): Promise<{ path: string; totalSize: number }> => {
     const storageType = getStorageType()
 
     const sanitizedFilename = _sanitizeFilename(fileName)
@@ -95,7 +134,10 @@ export const addArrayFilesToStorage = async (mime: string, bf: Buffer, fileName:
         })
         await s3Client.send(putObjCmd)
         fileNames.push(sanitizedFilename)
-        return 'FILE-STORAGE::' + JSON.stringify(fileNames)
+
+        const totalSize = await getS3StorageSize(paths[0])
+
+        return { path: 'FILE-STORAGE::' + JSON.stringify(fileNames), totalSize: totalSize / 1024 / 1024 }
     } else if (storageType === 'gcs') {
         const { bucket } = getGcsClient()
         const normalizedPaths = paths.map((p) => p.replace(/\\/g, '/'))
@@ -109,7 +151,10 @@ export const addArrayFilesToStorage = async (mime: string, bf: Buffer, fileName:
                 .end(bf)
         })
         fileNames.push(sanitizedFilename)
-        return 'FILE-STORAGE::' + JSON.stringify(fileNames)
+
+        const totalSize = await getGCSStorageSize(paths[0])
+
+        return { path: 'FILE-STORAGE::' + JSON.stringify(fileNames), totalSize: totalSize / 1024 / 1024 }
     } else {
         const dir = path.join(getStoragePath(), ...paths.map(_sanitizeFilename))
         if (!fs.existsSync(dir)) {
@@ -118,11 +163,19 @@ export const addArrayFilesToStorage = async (mime: string, bf: Buffer, fileName:
         const filePath = path.join(dir, sanitizedFilename)
         fs.writeFileSync(filePath, bf)
         fileNames.push(sanitizedFilename)
-        return 'FILE-STORAGE::' + JSON.stringify(fileNames)
+
+        const totalSize = await dirSize(path.join(getStoragePath(), paths[0]))
+
+        return { path: 'FILE-STORAGE::' + JSON.stringify(fileNames), totalSize: totalSize / 1024 / 1024 }
     }
 }
 
-export const addSingleFileToStorage = async (mime: string, bf: Buffer, fileName: string, ...paths: string[]) => {
+export const addSingleFileToStorage = async (
+    mime: string,
+    bf: Buffer,
+    fileName: string,
+    ...paths: string[]
+): Promise<{ path: string; totalSize: number }> => {
     const storageType = getStorageType()
     const sanitizedFilename = _sanitizeFilename(fileName)
 
@@ -142,7 +195,10 @@ export const addSingleFileToStorage = async (mime: string, bf: Buffer, fileName:
             Body: bf
         })
         await s3Client.send(putObjCmd)
-        return 'FILE-STORAGE::' + sanitizedFilename
+
+        const totalSize = await getS3StorageSize(paths[0])
+
+        return { path: 'FILE-STORAGE::' + sanitizedFilename, totalSize: totalSize / 1024 / 1024 }
     } else if (storageType === 'gcs') {
         const { bucket } = getGcsClient()
         const normalizedPaths = paths.map((p) => p.replace(/\\/g, '/'))
@@ -155,7 +211,10 @@ export const addSingleFileToStorage = async (mime: string, bf: Buffer, fileName:
                 .on('finish', () => resolve())
                 .end(bf)
         })
-        return 'FILE-STORAGE::' + sanitizedFilename
+
+        const totalSize = await getGCSStorageSize(paths[0])
+
+        return { path: 'FILE-STORAGE::' + sanitizedFilename, totalSize: totalSize / 1024 / 1024 }
     } else {
         const dir = path.join(getStoragePath(), ...paths.map(_sanitizeFilename))
         if (!fs.existsSync(dir)) {
@@ -163,7 +222,9 @@ export const addSingleFileToStorage = async (mime: string, bf: Buffer, fileName:
         }
         const filePath = path.join(dir, sanitizedFilename)
         fs.writeFileSync(filePath, bf)
-        return 'FILE-STORAGE::' + sanitizedFilename
+
+        const totalSize = await dirSize(path.join(getStoragePath(), paths[0]))
+        return { path: 'FILE-STORAGE::' + sanitizedFilename, totalSize: totalSize / 1024 / 1024 }
     }
 }
 
@@ -215,36 +276,246 @@ export const getFileFromStorage = async (file: string, ...paths: string[]): Prom
             Key = Key.substring(1)
         }
 
-        const getParams = {
-            Bucket,
-            Key
-        }
+        try {
+            const getParams = {
+                Bucket,
+                Key
+            }
 
-        const response = await s3Client.send(new GetObjectCommand(getParams))
-        const body = response.Body
-        if (body instanceof Readable) {
-            const streamToString = await body.transformToString('base64')
-            if (streamToString) {
-                return Buffer.from(streamToString, 'base64')
+            const response = await s3Client.send(new GetObjectCommand(getParams))
+            const body = response.Body
+            if (body instanceof Readable) {
+                const streamToString = await body.transformToString('base64')
+                if (streamToString) {
+                    return Buffer.from(streamToString, 'base64')
+                }
+            }
+            // @ts-ignore
+            const buffer = Buffer.concat(response.Body.toArray())
+            return buffer
+        } catch (error) {
+            // Fallback: Check if file exists without the first path element (likely orgId)
+            if (paths.length > 1) {
+                const fallbackPaths = paths.slice(1)
+                let fallbackKey = fallbackPaths.reduce((acc, cur) => acc + '/' + cur, '') + '/' + sanitizedFilename
+                if (fallbackKey.startsWith('/')) {
+                    fallbackKey = fallbackKey.substring(1)
+                }
+
+                try {
+                    const fallbackParams = {
+                        Bucket,
+                        Key: fallbackKey
+                    }
+                    const fallbackResponse = await s3Client.send(new GetObjectCommand(fallbackParams))
+                    const fallbackBody = fallbackResponse.Body
+
+                    // Get the file content
+                    let fileContent: Buffer
+                    if (fallbackBody instanceof Readable) {
+                        const streamToString = await fallbackBody.transformToString('base64')
+                        if (streamToString) {
+                            fileContent = Buffer.from(streamToString, 'base64')
+                        } else {
+                            // @ts-ignore
+                            fileContent = Buffer.concat(fallbackBody.toArray())
+                        }
+                    } else {
+                        // @ts-ignore
+                        fileContent = Buffer.concat(fallbackBody.toArray())
+                    }
+
+                    // Move to correct location with orgId
+                    const putObjCmd = new PutObjectCommand({
+                        Bucket,
+                        Key,
+                        Body: fileContent
+                    })
+                    await s3Client.send(putObjCmd)
+
+                    // Delete the old file
+                    await s3Client.send(
+                        new DeleteObjectsCommand({
+                            Bucket,
+                            Delete: {
+                                Objects: [{ Key: fallbackKey }],
+                                Quiet: false
+                            }
+                        })
+                    )
+
+                    // Check if the directory is empty and delete recursively if needed
+                    if (fallbackPaths.length > 0) {
+                        await _cleanEmptyS3Folders(s3Client, Bucket, fallbackPaths[0])
+                    }
+
+                    return fileContent
+                } catch (fallbackError) {
+                    // Throw the original error since the fallback also failed
+                    throw error
+                }
+            } else {
+                throw error
             }
         }
-        // @ts-ignore
-        const buffer = Buffer.concat(response.Body.toArray())
-        return buffer
     } else if (storageType === 'gcs') {
         const { bucket } = getGcsClient()
         const normalizedPaths = paths.map((p) => p.replace(/\\/g, '/'))
         const normalizedFilename = sanitizedFilename.replace(/\\/g, '/')
         const filePath = [...normalizedPaths, normalizedFilename].join('/')
-        const file = bucket.file(filePath)
-        const [buffer] = await file.download()
-        return buffer
+
+        try {
+            const file = bucket.file(filePath)
+            const [buffer] = await file.download()
+            return buffer
+        } catch (error) {
+            // Fallback: Check if file exists without the first path element (likely orgId)
+            if (normalizedPaths.length > 1) {
+                const fallbackPaths = normalizedPaths.slice(1)
+                const fallbackPath = [...fallbackPaths, normalizedFilename].join('/')
+
+                try {
+                    const fallbackFile = bucket.file(fallbackPath)
+                    const [buffer] = await fallbackFile.download()
+
+                    // Move to correct location with orgId
+                    const file = bucket.file(filePath)
+                    await new Promise<void>((resolve, reject) => {
+                        file.createWriteStream()
+                            .on('error', (err) => reject(err))
+                            .on('finish', () => resolve())
+                            .end(buffer)
+                    })
+
+                    // Delete the old file
+                    await fallbackFile.delete()
+
+                    // Check if the directory is empty and delete recursively if needed
+                    if (fallbackPaths.length > 0) {
+                        await _cleanEmptyGCSFolders(bucket, fallbackPaths[0])
+                    }
+
+                    return buffer
+                } catch (fallbackError) {
+                    // Throw the original error since the fallback also failed
+                    throw error
+                }
+            } else {
+                throw error
+            }
+        }
     } else {
-        const fileInStorage = path.join(getStoragePath(), ...paths.map(_sanitizeFilename), sanitizedFilename)
-        return fs.readFileSync(fileInStorage)
+        try {
+            const fileInStorage = path.join(getStoragePath(), ...paths.map(_sanitizeFilename), sanitizedFilename)
+            return fs.readFileSync(fileInStorage)
+        } catch (error) {
+            // Fallback: Check if file exists without the first path element (likely orgId)
+            if (paths.length > 1) {
+                const fallbackPaths = paths.slice(1)
+                const fallbackPath = path.join(getStoragePath(), ...fallbackPaths.map(_sanitizeFilename), sanitizedFilename)
+
+                if (fs.existsSync(fallbackPath)) {
+                    // Create directory if it doesn't exist
+                    const targetPath = path.join(getStoragePath(), ...paths.map(_sanitizeFilename), sanitizedFilename)
+                    const dir = path.dirname(targetPath)
+                    if (!fs.existsSync(dir)) {
+                        fs.mkdirSync(dir, { recursive: true })
+                    }
+
+                    // Copy file to correct location with orgId
+                    fs.copyFileSync(fallbackPath, targetPath)
+
+                    // Delete the old file
+                    fs.unlinkSync(fallbackPath)
+
+                    // Clean up empty directories recursively
+                    if (fallbackPaths.length > 0) {
+                        _cleanEmptyLocalFolders(path.join(getStoragePath(), ...fallbackPaths.map(_sanitizeFilename).slice(0, -1)))
+                    }
+
+                    return fs.readFileSync(targetPath)
+                } else {
+                    throw error
+                }
+            } else {
+                throw error
+            }
+        }
     }
 }
 
+export const getFilesListFromStorage = async (...paths: string[]): Promise<Array<{ name: string; path: string; size: number }>> => {
+    const storageType = getStorageType()
+    if (storageType === 's3') {
+        const { s3Client, Bucket } = getS3Config()
+
+        let Key = paths.reduce((acc, cur) => acc + '/' + cur, '')
+        if (Key.startsWith('/')) {
+            Key = Key.substring(1)
+        }
+
+        const listCommand = new ListObjectsV2Command({
+            Bucket,
+            Prefix: Key
+        })
+        const list = await s3Client.send(listCommand)
+
+        if (list.Contents && list.Contents.length > 0) {
+            return list.Contents.map((item) => ({
+                name: item.Key?.split('/').pop() || '',
+                path: item.Key ?? '',
+                size: item.Size || 0
+            }))
+        } else {
+            return []
+        }
+    } else {
+        const directory = path.join(getStoragePath(), ...paths)
+        const filesList = getFilePaths(directory)
+        return filesList
+    }
+}
+
+interface FileInfo {
+    name: string
+    path: string
+    size: number
+}
+
+function getFilePaths(dir: string): FileInfo[] {
+    let results: FileInfo[] = []
+
+    function readDirectory(directory: string) {
+        try {
+            if (!fs.existsSync(directory)) {
+                console.warn(`Directory does not exist: ${directory}`)
+                return
+            }
+
+            const list = fs.readdirSync(directory)
+            list.forEach((file) => {
+                const filePath = path.join(directory, file)
+                try {
+                    const stat = fs.statSync(filePath)
+                    if (stat && stat.isDirectory()) {
+                        readDirectory(filePath)
+                    } else {
+                        const sizeInMB = stat.size / (1024 * 1024)
+                        results.push({ name: file, path: filePath, size: sizeInMB })
+                    }
+                } catch (error) {
+                    console.error(`Error processing file ${filePath}:`, error)
+                }
+            })
+        } catch (error) {
+            console.error(`Error reading directory ${directory}:`, error)
+        }
+    }
+
+    readDirectory(dir)
+    return results
+}
+
 /**
  * Prepare storage path
  */
@@ -267,14 +538,26 @@ export const removeFilesFromStorage = async (...paths: string[]) => {
         if (Key.startsWith('/')) {
             Key = Key.substring(1)
         }
+
         await _deleteS3Folder(Key)
+
+        // check folder size after deleting all the files
+        const totalSize = await getS3StorageSize(paths[0])
+        return { totalSize: totalSize / 1024 / 1024 }
     } else if (storageType === 'gcs') {
         const { bucket } = getGcsClient()
         const normalizedPath = paths.map((p) => p.replace(/\\/g, '/')).join('/')
         await bucket.deleteFiles({ prefix: `${normalizedPath}/` })
+
+        const totalSize = await getGCSStorageSize(paths[0])
+        return { totalSize: totalSize / 1024 / 1024 }
     } else {
         const directory = path.join(getStoragePath(), ...paths.map(_sanitizeFilename))
-        _deleteLocalFolderRecursive(directory)
+        await _deleteLocalFolderRecursive(directory)
+
+        const totalSize = await dirSize(path.join(getStoragePath(), paths[0]))
+
+        return { totalSize: totalSize / 1024 / 1024 }
     }
 }
 
@@ -304,6 +587,10 @@ export const removeSpecificFileFromStorage = async (...paths: string[]) => {
             Key = Key.substring(1)
         }
         await _deleteS3Folder(Key)
+
+        // check folder size after deleting all the files
+        const totalSize = await getS3StorageSize(paths[0])
+        return { totalSize: totalSize / 1024 / 1024 }
     } else if (storageType === 'gcs') {
         const { bucket } = getGcsClient()
         const fileName = paths.pop()
@@ -313,6 +600,9 @@ export const removeSpecificFileFromStorage = async (...paths: string[]) => {
         }
         const normalizedPath = paths.map((p) => p.replace(/\\/g, '/')).join('/')
         await bucket.file(normalizedPath).delete()
+
+        const totalSize = await getGCSStorageSize(paths[0])
+        return { totalSize: totalSize / 1024 / 1024 }
     } else {
         const fileName = paths.pop()
         if (fileName) {
@@ -320,7 +610,15 @@ export const removeSpecificFileFromStorage = async (...paths: string[]) => {
             paths.push(sanitizedFilename)
         }
         const file = path.join(getStoragePath(), ...paths.map(_sanitizeFilename))
-        fs.unlinkSync(file)
+        // check if file exists, if not skip delete
+        // this might happen when user tries to delete a document loader but the attached file is already deleted
+        const stat = fs.statSync(file, { throwIfNoEntry: false })
+        if (stat && stat.isFile()) {
+            fs.unlinkSync(file)
+        }
+
+        const totalSize = await dirSize(path.join(getStoragePath(), paths[0]))
+        return { totalSize: totalSize / 1024 / 1024 }
     }
 }
 
@@ -333,52 +631,63 @@ export const removeFolderFromStorage = async (...paths: string[]) => {
             Key = Key.substring(1)
         }
         await _deleteS3Folder(Key)
+
+        // check folder size after deleting all the files
+        const totalSize = await getS3StorageSize(paths[0])
+        return { totalSize: totalSize / 1024 / 1024 }
     } else if (storageType === 'gcs') {
         const { bucket } = getGcsClient()
         const normalizedPath = paths.map((p) => p.replace(/\\/g, '/')).join('/')
         await bucket.deleteFiles({ prefix: `${normalizedPath}/` })
+
+        const totalSize = await getGCSStorageSize(paths[0])
+        return { totalSize: totalSize / 1024 / 1024 }
     } else {
         const directory = path.join(getStoragePath(), ...paths.map(_sanitizeFilename))
-        _deleteLocalFolderRecursive(directory, true)
+        await _deleteLocalFolderRecursive(directory, true)
+
+        const totalSize = await dirSize(path.join(getStoragePath(), paths[0]))
+        return { totalSize: totalSize / 1024 / 1024 }
     }
 }
 
-const _deleteLocalFolderRecursive = (directory: string, deleteParentChatflowFolder?: boolean) => {
-    // Console error here as failing is not destructive operation
-    if (fs.existsSync(directory)) {
+const _deleteLocalFolderRecursive = async (directory: string, deleteParentChatflowFolder?: boolean) => {
+    try {
+        // Check if the path exists
+        await fs.promises.access(directory)
+
         if (deleteParentChatflowFolder) {
-            fs.rmSync(directory, { recursive: true, force: true })
-        } else {
-            fs.readdir(directory, (error, files) => {
-                if (error) console.error('Could not read directory')
-
-                for (let i = 0; i < files.length; i++) {
-                    const file = files[i]
-                    const file_path = path.join(directory, file)
-
-                    fs.stat(file_path, (error, stat) => {
-                        if (error) console.error('File do not exist')
-
-                        if (!stat.isDirectory()) {
-                            fs.unlink(file_path, (error) => {
-                                if (error) console.error('Could not delete file')
-                            })
-                            if (i === files.length - 1) {
-                                fs.rmSync(directory, { recursive: true, force: true })
-                            }
-                        } else {
-                            _deleteLocalFolderRecursive(file_path)
-                        }
-                    })
-                }
-            })
+            await fs.promises.rmdir(directory, { recursive: true })
         }
+
+        // Get stats of the path to determine if it's a file or directory
+        const stats = await fs.promises.stat(directory)
+
+        if (stats.isDirectory()) {
+            // Read all directory contents
+            const files = await fs.promises.readdir(directory)
+
+            // Recursively delete all contents
+            for (const file of files) {
+                const currentPath = path.join(directory, file)
+                await _deleteLocalFolderRecursive(currentPath) // Recursive call
+            }
+
+            // Delete the directory itself after emptying it
+            await fs.promises.rmdir(directory, { recursive: true })
+        } else {
+            // If it's a file, delete it directly
+            await fs.promises.unlink(directory)
+        }
+    } catch (error) {
+        // Error handling
     }
 }
 
 const _deleteS3Folder = async (location: string) => {
     let count = 0 // number of files deleted
     const { s3Client, Bucket } = getS3Config()
+
     async function recursiveS3Delete(token?: any) {
         // get the files
         const listCommand = new ListObjectsV2Command({
@@ -410,6 +719,7 @@ const _deleteS3Folder = async (location: string) => {
         // return total deleted count when finished
         return `${count} files deleted from S3`
     }
+
     // start the recursive function
     return recursiveS3Delete()
 }
@@ -417,34 +727,120 @@ const _deleteS3Folder = async (location: string) => {
 export const streamStorageFile = async (
     chatflowId: string,
     chatId: string,
-    fileName: string
+    fileName: string,
+    orgId: string
 ): Promise<fs.ReadStream | Buffer | undefined> => {
     const storageType = getStorageType()
     const sanitizedFilename = sanitize(fileName)
     if (storageType === 's3') {
         const { s3Client, Bucket } = getS3Config()
 
-        const Key = chatflowId + '/' + chatId + '/' + sanitizedFilename
+        const Key = orgId + '/' + chatflowId + '/' + chatId + '/' + sanitizedFilename
         const getParams = {
             Bucket,
             Key
         }
-        const response = await s3Client.send(new GetObjectCommand(getParams))
-        const body = response.Body
-        if (body instanceof Readable) {
-            const blob = await body.transformToByteArray()
-            return Buffer.from(blob)
+        try {
+            const response = await s3Client.send(new GetObjectCommand(getParams))
+            const body = response.Body
+            if (body instanceof Readable) {
+                const blob = await body.transformToByteArray()
+                return Buffer.from(blob)
+            }
+        } catch (error) {
+            // Fallback: Check if file exists without orgId
+            const fallbackKey = chatflowId + '/' + chatId + '/' + sanitizedFilename
+            try {
+                const fallbackParams = {
+                    Bucket,
+                    Key: fallbackKey
+                }
+                const fallbackResponse = await s3Client.send(new GetObjectCommand(fallbackParams))
+                const fallbackBody = fallbackResponse.Body
+
+                // If found, copy to correct location with orgId
+                if (fallbackBody) {
+                    // Get the file content
+                    let fileContent: Buffer
+                    if (fallbackBody instanceof Readable) {
+                        const blob = await fallbackBody.transformToByteArray()
+                        fileContent = Buffer.from(blob)
+                    } else {
+                        // @ts-ignore
+                        fileContent = Buffer.concat(fallbackBody.toArray())
+                    }
+
+                    // Move to correct location with orgId
+                    const putObjCmd = new PutObjectCommand({
+                        Bucket,
+                        Key,
+                        Body: fileContent
+                    })
+                    await s3Client.send(putObjCmd)
+
+                    // Delete the old file
+                    await s3Client.send(
+                        new DeleteObjectsCommand({
+                            Bucket,
+                            Delete: {
+                                Objects: [{ Key: fallbackKey }],
+                                Quiet: false
+                            }
+                        })
+                    )
+
+                    // Check if the directory is empty and delete recursively if needed
+                    await _cleanEmptyS3Folders(s3Client, Bucket, chatflowId)
+
+                    return fileContent
+                }
+            } catch (fallbackError) {
+                // File not found in fallback location either
+                throw new Error(`File ${fileName} not found`)
+            }
         }
     } else if (storageType === 'gcs') {
         const { bucket } = getGcsClient()
         const normalizedChatflowId = chatflowId.replace(/\\/g, '/')
         const normalizedChatId = chatId.replace(/\\/g, '/')
         const normalizedFilename = sanitizedFilename.replace(/\\/g, '/')
-        const filePath = `${normalizedChatflowId}/${normalizedChatId}/${normalizedFilename}`
-        const [buffer] = await bucket.file(filePath).download()
-        return buffer
+        const filePath = `${orgId}/${normalizedChatflowId}/${normalizedChatId}/${normalizedFilename}`
+
+        try {
+            const [buffer] = await bucket.file(filePath).download()
+            return buffer
+        } catch (error) {
+            // Fallback: Check if file exists without orgId
+            const fallbackPath = `${normalizedChatflowId}/${normalizedChatId}/${normalizedFilename}`
+            try {
+                const fallbackFile = bucket.file(fallbackPath)
+                const [buffer] = await fallbackFile.download()
+
+                // If found, copy to correct location with orgId
+                if (buffer) {
+                    const file = bucket.file(filePath)
+                    await new Promise<void>((resolve, reject) => {
+                        file.createWriteStream()
+                            .on('error', (err) => reject(err))
+                            .on('finish', () => resolve())
+                            .end(buffer)
+                    })
+
+                    // Delete the old file
+                    await fallbackFile.delete()
+
+                    // Check if the directory is empty and delete recursively if needed
+                    await _cleanEmptyGCSFolders(bucket, normalizedChatflowId)
+
+                    return buffer
+                }
+            } catch (fallbackError) {
+                // File not found in fallback location either
+                throw new Error(`File ${fileName} not found`)
+            }
+        }
     } else {
-        const filePath = path.join(getStoragePath(), chatflowId, chatId, sanitizedFilename)
+        const filePath = path.join(getStoragePath(), orgId, chatflowId, chatId, sanitizedFilename)
         //raise error if file path is not absolute
         if (!path.isAbsolute(filePath)) throw new Error(`Invalid file path`)
         //raise error if file path contains '..'
@@ -455,11 +851,159 @@ export const streamStorageFile = async (
         if (fs.existsSync(filePath)) {
             return fs.createReadStream(filePath)
         } else {
-            throw new Error(`File ${fileName} not found`)
+            // Fallback: Check if file exists without orgId
+            const fallbackPath = path.join(getStoragePath(), chatflowId, chatId, sanitizedFilename)
+
+            if (fs.existsSync(fallbackPath)) {
+                // Create directory if it doesn't exist
+                const dir = path.dirname(filePath)
+                if (!fs.existsSync(dir)) {
+                    fs.mkdirSync(dir, { recursive: true })
+                }
+
+                // Copy file to correct location with orgId
+                fs.copyFileSync(fallbackPath, filePath)
+
+                // Delete the old file
+                fs.unlinkSync(fallbackPath)
+
+                // Clean up empty directories recursively
+                _cleanEmptyLocalFolders(path.join(getStoragePath(), chatflowId, chatId))
+
+                return fs.createReadStream(filePath)
+            } else {
+                throw new Error(`File ${fileName} not found`)
+            }
         }
     }
 }
 
+/**
+ * Check if a local directory is empty and delete it if so,
+ * then check parent directories recursively
+ */
+const _cleanEmptyLocalFolders = (dirPath: string) => {
+    try {
+        // Stop if we reach the storage root
+        if (dirPath === getStoragePath()) return
+
+        // Check if directory exists
+        if (!fs.existsSync(dirPath)) return
+
+        // Read directory contents
+        const files = fs.readdirSync(dirPath)
+
+        // If directory is empty, delete it and check parent
+        if (files.length === 0) {
+            fs.rmdirSync(dirPath)
+            // Recursively check parent directory
+            _cleanEmptyLocalFolders(path.dirname(dirPath))
+        }
+    } catch (error) {
+        // Ignore errors during cleanup
+        console.error('Error cleaning empty folders:', error)
+    }
+}
+
+/**
+ * Check if an S3 "folder" is empty and delete it recursively
+ */
+const _cleanEmptyS3Folders = async (s3Client: S3Client, Bucket: string, prefix: string) => {
+    try {
+        // Skip if prefix is empty
+        if (!prefix) return
+
+        // List objects in this "folder"
+        const listCmd = new ListObjectsV2Command({
+            Bucket,
+            Prefix: prefix + '/',
+            Delimiter: '/'
+        })
+
+        const response = await s3Client.send(listCmd)
+
+        // If folder is empty (only contains common prefixes but no files)
+        if (
+            (response.Contents?.length === 0 || !response.Contents) &&
+            (response.CommonPrefixes?.length === 0 || !response.CommonPrefixes)
+        ) {
+            // Delete the folder marker if it exists
+            await s3Client.send(
+                new DeleteObjectsCommand({
+                    Bucket,
+                    Delete: {
+                        Objects: [{ Key: prefix + '/' }],
+                        Quiet: true
+                    }
+                })
+            )
+
+            // Recursively check parent folder
+            const parentPrefix = prefix.substring(0, prefix.lastIndexOf('/'))
+            if (parentPrefix) {
+                await _cleanEmptyS3Folders(s3Client, Bucket, parentPrefix)
+            }
+        }
+    } catch (error) {
+        // Ignore errors during cleanup
+        console.error('Error cleaning empty S3 folders:', error)
+    }
+}
+
+/**
+ * Check if a GCS "folder" is empty and delete recursively if so
+ */
+const _cleanEmptyGCSFolders = async (bucket: any, prefix: string) => {
+    try {
+        // Skip if prefix is empty
+        if (!prefix) return
+
+        // List files with this prefix
+        const [files] = await bucket.getFiles({
+            prefix: prefix + '/',
+            delimiter: '/'
+        })
+
+        // If folder is empty (no files)
+        if (files.length === 0) {
+            // Delete the folder marker if it exists
+            try {
+                await bucket.file(prefix + '/').delete()
+            } catch (err) {
+                // Folder marker might not exist, ignore
+            }
+
+            // Recursively check parent folder
+            const parentPrefix = prefix.substring(0, prefix.lastIndexOf('/'))
+            if (parentPrefix) {
+                await _cleanEmptyGCSFolders(bucket, parentPrefix)
+            }
+        }
+    } catch (error) {
+        // Ignore errors during cleanup
+        console.error('Error cleaning empty GCS folders:', error)
+    }
+}
+
+export const getGCSStorageSize = async (orgId: string): Promise<number> => {
+    const { bucket } = getGcsClient()
+    let totalSize = 0
+
+    const [files] = await bucket.getFiles({ prefix: orgId })
+
+    for (const file of files) {
+        const size = file.metadata.size
+        // Handle different types that size could be
+        if (typeof size === 'string') {
+            totalSize += parseInt(size, 10) || 0
+        } else if (typeof size === 'number') {
+            totalSize += size
+        }
+    }
+
+    return totalSize
+}
+
 export const getGcsClient = () => {
     const pathToGcsCredential = process.env.GOOGLE_CLOUD_STORAGE_CREDENTIAL
     const projectId = process.env.GOOGLE_CLOUD_STORAGE_PROJ_ID
@@ -482,6 +1026,20 @@ export const getGcsClient = () => {
     return { storage, bucket }
 }
 
+export const getS3StorageSize = async (orgId: string): Promise<number> => {
+    const { s3Client, Bucket } = getS3Config()
+    const getCmd = new ListObjectsCommand({
+        Bucket,
+        Prefix: orgId
+    })
+    const headObj = await s3Client.send(getCmd)
+    let totalSize = 0
+    for (const obj of headObj.Contents || []) {
+        totalSize += obj.Size || 0
+    }
+    return totalSize
+}
+
 export const getS3Config = () => {
     const accessKeyId = process.env.S3_STORAGE_ACCESS_KEY_ID
     const secretAccessKey = process.env.S3_STORAGE_SECRET_ACCESS_KEY

packages/components/src/utils.ts

@@ -4,7 +4,7 @@ import * as fs from 'fs'
 import * as path from 'path'
 import { JSDOM } from 'jsdom'
 import { z } from 'zod'
-import { DataSource } from 'typeorm'
+import { DataSource, Equal } from 'typeorm'
 import { ICommonObject, IDatabaseEntity, IFileUpload, IMessage, INodeData, IVariable, MessageContentImageUrl } from './Interface'
 import { AES, enc } from 'crypto-js'
 import { omit } from 'lodash'
@@ -706,7 +706,7 @@ export const getUserHome = (): string => {
  * @param {IChatMessage[]} chatmessages
  * @returns {BaseMessage[]}
  */
-export const mapChatMessageToBaseMessage = async (chatmessages: any[] = []): Promise<BaseMessage[]> => {
+export const mapChatMessageToBaseMessage = async (chatmessages: any[] = [], orgId: string): Promise<BaseMessage[]> => {
     const chatHistory = []
 
     for (const message of chatmessages) {
@@ -722,7 +722,7 @@ export const mapChatMessageToBaseMessage = async (chatmessages: any[] = []): Pro
                     const imageContents: MessageContentImageUrl[] = []
                     for (const upload of uploads) {
                         if (upload.type === 'stored-file' && upload.mime.startsWith('image/')) {
-                            const fileData = await getFileFromStorage(upload.name, message.chatflowid, message.chatId)
+                            const fileData = await getFileFromStorage(upload.name, orgId, message.chatflowid, message.chatId)
                             // as the image is stored in the server, read the file and convert it to base64
                             const bf = 'data:' + upload.mime + ';base64,' + fileData.toString('base64')
 
@@ -746,7 +746,8 @@ export const mapChatMessageToBaseMessage = async (chatmessages: any[] = []): Pro
                             const options = {
                                 retrieveAttachmentChatId: true,
                                 chatflowid: message.chatflowid,
-                                chatId: message.chatId
+                                chatId: message.chatId,
+                                orgId
                             }
                             let fileInputFieldFromMimeType = 'txtFile'
                             fileInputFieldFromMimeType = mapMimeTypeToInputField(upload.mime)
@@ -935,8 +936,16 @@ export const convertMultiOptionsToStringArray = (inputString: string): string[]
  * @param {IDatabaseEntity} databaseEntities
  * @param {INodeData} nodeData
  */
-export const getVars = async (appDataSource: DataSource, databaseEntities: IDatabaseEntity, nodeData: INodeData) => {
-    const variables = ((await appDataSource.getRepository(databaseEntities['Variable']).find()) as IVariable[]) ?? []
+export const getVars = async (
+    appDataSource: DataSource,
+    databaseEntities: IDatabaseEntity,
+    nodeData: INodeData,
+    options: ICommonObject
+) => {
+    const variables =
+        ((await appDataSource
+            .getRepository(databaseEntities['Variable'])
+            .findBy(options.workspaceId ? { workspaceId: Equal(options.workspaceId) } : {})) as IVariable[]) ?? []
 
     // override variables defined in overrideConfig
     // nodeData.inputs.vars is an Object, check each property and override the variable