From 3f81c7c231558e777d5aabcd93d8a5a455dd9c4b Mon Sep 17 00:00:00 2001 From: Henry Fontanier Date: Tue, 21 Jan 2025 14:32:32 +0100 Subject: [PATCH] overwrite Devin's mess --- .../lib/api/assistant/actions/dust_app_run.ts | 147 ++++++++++++++++-- .../assistant/actions/result_file_helpers.ts | 47 ++++++ .../lib/api/assistant/actions/tables_query.ts | 46 +----- .../models/assistant/actions/dust_app_run.ts | 26 +++- front/migrations/db/migration_145.sql | 7 + .../front/assistant/actions/dust_app_run.ts | 2 + 6 files changed, 218 insertions(+), 57 deletions(-) create mode 100644 front/lib/api/assistant/actions/result_file_helpers.ts create mode 100644 front/migrations/db/migration_145.sql diff --git a/front/lib/api/assistant/actions/dust_app_run.ts b/front/lib/api/assistant/actions/dust_app_run.ts index 5b77b0971d16..657b4a1e3596 100644 --- a/front/lib/api/assistant/actions/dust_app_run.ts +++ b/front/lib/api/assistant/actions/dust_app_run.ts @@ -1,5 +1,10 @@ import { DustAPI } from "@dust-tt/client"; import type { + ActionGeneratedFileType, + AgentActionSpecification, + DatasetSchema, + DustAppParameters, + DustAppRunActionType, DustAppRunBlockEvent, DustAppRunConfigurationType, DustAppRunErrorEvent, @@ -8,20 +13,19 @@ import type { FunctionCallType, FunctionMessageTypeModel, ModelId, + Result, + SpecificationType, } from "@dust-tt/types"; -import type { DustAppParameters, DustAppRunActionType } from "@dust-tt/types"; -import type { AgentActionSpecification } from "@dust-tt/types"; -import type { SpecificationType } from "@dust-tt/types"; -import type { DatasetSchema } from "@dust-tt/types"; -import type { Result } from "@dust-tt/types"; import { BaseAction, + Err, getHeaderFromGroupIds, + Ok, SUPPORTED_MODEL_CONFIGS, } from "@dust-tt/types"; -import { Err, Ok } from "@dust-tt/types"; import { DUST_CONVERSATION_HISTORY_MAGIC_INPUT_KEY } from "@app/lib/api/assistant/actions/constants"; +import { getToolResultOutputCsvFileAndSnippet } from "@app/lib/api/assistant/actions/result_file_helpers"; import type { BaseActionRunParams } from "@app/lib/api/assistant/actions/types"; import { BaseActionConfigurationServerRunner } from "@app/lib/api/assistant/actions/types"; import { renderConversationForModel } from "@app/lib/api/assistant/generation"; @@ -32,6 +36,8 @@ import { prodAPICredentialsForOwner } from "@app/lib/auth"; import { extractConfig } from "@app/lib/config"; import { AgentDustAppRunAction } from "@app/lib/models/assistant/actions/dust_app_run"; import { AppResource } from "@app/lib/resources/app_resource"; +import { FileResource } from "@app/lib/resources/file_resource"; +import { FileModel } from "@app/lib/resources/storage/models/files"; import { sanitizeJSONOutput } from "@app/lib/utils"; import logger from "@app/logger/logger"; @@ -51,6 +57,9 @@ interface DustAppRunActionBlob { functionCallId: string | null; functionCallName: string | null; step: number; + resultsFileId: string | null; + resultsFileSnippet: string | null; + generatedFiles: ActionGeneratedFileType[]; } export class DustAppRunAction extends BaseAction { @@ -68,10 +77,12 @@ export class DustAppRunAction extends BaseAction { readonly functionCallId: string | null; readonly functionCallName: string | null; readonly step: number; + readonly resultsFileId: string | null; + readonly resultsFileSnippet: string | null; readonly type = "dust_app_run_action"; constructor(blob: DustAppRunActionBlob) { - super(blob.id, "dust_app_run_action"); + super(blob.id, "dust_app_run_action", blob.generatedFiles || []); this.agentMessageId = blob.agentMessageId; this.appWorkspaceId = blob.appWorkspaceId; @@ -83,6 +94,8 @@ export class DustAppRunAction extends BaseAction { this.functionCallId = blob.functionCallId; this.functionCallName = blob.functionCallName; this.step = blob.step; + this.resultsFileId = blob.resultsFileId; + this.resultsFileSnippet = blob.resultsFileSnippet; } renderForFunctionCall(): FunctionCallType { @@ -303,6 +316,8 @@ export class DustAppRunConfigurationServerRunner extends BaseActionConfiguration functionCallName: actionConfiguration.name, agentMessageId: agentMessage.agentMessageId, step, + resultsFileId: null, + resultsFileSnippet: null, }); yield { @@ -322,6 +337,9 @@ export class DustAppRunConfigurationServerRunner extends BaseActionConfiguration functionCallName: actionConfiguration.name, agentMessageId: agentMessage.agentMessageId, step, + resultsFileId: null, + resultsFileSnippet: null, + generatedFiles: [], }), }; @@ -462,6 +480,9 @@ export class DustAppRunConfigurationServerRunner extends BaseActionConfiguration output: null, agentMessageId: agentMessage.agentMessageId, step: action.step, + resultsFileId: null, + resultsFileSnippet: null, + generatedFiles: [], }), }; } @@ -486,12 +507,79 @@ export class DustAppRunConfigurationServerRunner extends BaseActionConfiguration } } - const output = sanitizeJSONOutput(lastBlockOutput); + function isValidStructuredOutput(output: unknown): output is { + __dust_structured_output?: Array< + Record + >; + } { + return ( + typeof output === "object" && + output !== null && + "__dust_structured_output" in output && + Array.isArray(output.__dust_structured_output) && + output.__dust_structured_output.length > 0 && + output.__dust_structured_output.every( + (r) => + typeof r === "object" && + Object.values(r).every( + (v) => + typeof v === "string" || + typeof v === "number" || + typeof v === "boolean" + ) + ) + ); + } + + const sanitizedOutput = sanitizeJSONOutput(lastBlockOutput); + + const updateParams: { + resultsFileId: number | null; + resultsFileSnippet: string | null; + output: unknown | null; + } = { + resultsFileId: null, + resultsFileSnippet: null, + output: null, + }; + + let resultFile: ActionGeneratedFileType | null = null; + + // Check for structured output that should be converted to CSV + if ( + isValidStructuredOutput(sanitizedOutput) && + sanitizedOutput.__dust_structured_output + ) { + const fileTitle = getDustAppRunResultsFileTitle({ + appName: app.name, + }); + + const { file, snippet } = await getToolResultOutputCsvFileAndSnippet( + auth, + { + title: fileTitle, + conversationId: conversation.sId, + results: sanitizedOutput.__dust_structured_output, + } + ); - // Update DustAppRunAction with the output of the last block. + resultFile = { + fileId: file.sId, + title: fileTitle, + contentType: file.contentType, + snippet: file.snippet, + }; + + delete sanitizedOutput.__dust_structured_output; + updateParams.resultsFileId = file.id; + updateParams.resultsFileSnippet = snippet; + } + + // Update DustAppRunAction with the output and file references await action.update({ + ...updateParams, + output: sanitizedOutput, runId: await dustRunId, - output, }); logger.info( @@ -517,9 +605,12 @@ export class DustAppRunConfigurationServerRunner extends BaseActionConfiguration functionCallId, functionCallName: actionConfiguration.name, runningBlock: null, - output, + output: sanitizedOutput, agentMessageId: agentMessage.agentMessageId, step: action.step, + resultsFileId: resultFile?.fileId ?? null, + resultsFileSnippet: updateParams.resultsFileSnippet, + generatedFiles: resultFile ? [resultFile] : [], }), }; } @@ -580,15 +671,38 @@ async function dustAppRunActionSpecification({ // used outside of api/assistant. We allow a ModelId interface here because we don't have `sId` on // actions (the `sId` is on the `Message` object linked to the `UserMessage` parent of this action). export async function dustAppRunTypesFromAgentMessageIds( + auth: Authenticator, agentMessageIds: ModelId[] ): Promise { + const owner = auth.getNonNullableWorkspace(); + const actions = await AgentDustAppRunAction.findAll({ where: { agentMessageId: agentMessageIds, }, + include: [ + { + model: FileModel, + as: "resultsFile", + }, + ], }); return actions.map((action) => { + const resultsFile: ActionGeneratedFileType | null = action.resultsFile + ? { + fileId: FileResource.modelIdToSId({ + id: action.resultsFile.id, + workspaceId: owner.id, + }), + title: getDustAppRunResultsFileTitle({ + appName: action.appName, + }), + contentType: action.resultsFile.contentType, + snippet: action.resultsFileSnippet, + } + : null; + return new DustAppRunAction({ id: action.id, appWorkspaceId: action.appWorkspaceId, @@ -601,6 +715,17 @@ export async function dustAppRunTypesFromAgentMessageIds( functionCallName: action.functionCallName, agentMessageId: action.agentMessageId, step: action.step, + resultsFileId: resultsFile?.fileId ?? null, + resultsFileSnippet: action.resultsFileSnippet, + generatedFiles: resultsFile ? [resultsFile] : [], }); }); } + +export function getDustAppRunResultsFileTitle({ + appName, +}: { + appName: string; +}): string { + return `${appName}_output.csv`; +} diff --git a/front/lib/api/assistant/actions/result_file_helpers.ts b/front/lib/api/assistant/actions/result_file_helpers.ts new file mode 100644 index 000000000000..d7851572bfd4 --- /dev/null +++ b/front/lib/api/assistant/actions/result_file_helpers.ts @@ -0,0 +1,47 @@ +import { stringify } from "csv-stringify"; + +import { generateCSVSnippet } from "@app/lib/api/csv"; +import { internalCreateToolOutputCsvFile } from "@app/lib/api/files/tool_output"; +import type { Authenticator } from "@app/lib/auth"; +import type { FileResource } from "@app/lib/resources/file_resource"; + +export async function getToolResultOutputCsvFileAndSnippet( + auth: Authenticator, + { + title, + conversationId, + results, + }: { + title: string; + conversationId: string; + results: Array>; + } +): Promise<{ + file: FileResource; + snippet: string; +}> { + const toCsv = ( + records: Array>, + options: { header: boolean } = { header: true } + ): Promise => { + return new Promise((resolve, reject) => { + stringify(records, options, (err, data) => { + if (err) { + reject(err); + } + resolve(data); + }); + }); + }; + + const csvOutput = await toCsv(results); + + const file = await internalCreateToolOutputCsvFile(auth, { + title, + conversationId: conversationId, + content: csvOutput, + contentType: "text/csv", + }); + + return { file, snippet: generateCSVSnippet(csvOutput) }; +} diff --git a/front/lib/api/assistant/actions/tables_query.ts b/front/lib/api/assistant/actions/tables_query.ts index bd3655dedca1..27c970236794 100644 --- a/front/lib/api/assistant/actions/tables_query.ts +++ b/front/lib/api/assistant/actions/tables_query.ts @@ -19,15 +19,13 @@ import { getTablesQueryResultsFileTitle, Ok, } from "@dust-tt/types"; -import { stringify } from "csv-stringify"; import { runActionStreamed } from "@app/lib/actions/server"; import { DEFAULT_TABLES_QUERY_ACTION_NAME } from "@app/lib/api/assistant/actions/constants"; +import { getToolResultOutputCsvFileAndSnippet } from "@app/lib/api/assistant/actions/result_file_helpers"; import type { BaseActionRunParams } from "@app/lib/api/assistant/actions/types"; import { BaseActionConfigurationServerRunner } from "@app/lib/api/assistant/actions/types"; import { renderConversationForModel } from "@app/lib/api/assistant/generation"; -import { generateCSVSnippet } from "@app/lib/api/csv"; -import { internalCreateToolOutputCsvFile } from "@app/lib/api/files/tool_output"; import { getSupportedModelConfig } from "@app/lib/assistant"; import type { Authenticator } from "@app/lib/auth"; import { AgentTablesQueryAction } from "@app/lib/models/assistant/actions/tables_query"; @@ -567,7 +565,7 @@ export class TablesQueryConfigurationServerRunner extends BaseActionConfiguratio output: sanitizedOutput, }); - const { file, snippet } = await getTablesQueryOutputCsvFileAndSnippet( + const { file, snippet } = await getToolResultOutputCsvFileAndSnippet( auth, { title: queryTitle, @@ -616,43 +614,3 @@ export class TablesQueryConfigurationServerRunner extends BaseActionConfiguratio return; } } - -async function getTablesQueryOutputCsvFileAndSnippet( - auth: Authenticator, - { - title, - conversationId, - results, - }: { - title: string; - conversationId: string; - results: Array>; - } -): Promise<{ - file: FileResource; - snippet: string; -}> { - const toCsv = ( - records: Array>, - options: { header: boolean } = { header: true } - ): Promise => { - return new Promise((resolve, reject) => { - stringify(records, options, (err, data) => { - if (err) { - reject(err); - } - resolve(data); - }); - }); - }; - const csvOutput = await toCsv(results); - - const file = await internalCreateToolOutputCsvFile(auth, { - title, - conversationId: conversationId, - content: csvOutput, - contentType: "text/csv", - }); - - return { file, snippet: generateCSVSnippet(csvOutput) }; -} diff --git a/front/lib/models/assistant/actions/dust_app_run.ts b/front/lib/models/assistant/actions/dust_app_run.ts index 64679d6cb5cd..5bcb0dea2673 100644 --- a/front/lib/models/assistant/actions/dust_app_run.ts +++ b/front/lib/models/assistant/actions/dust_app_run.ts @@ -1,10 +1,11 @@ import type { DustAppParameters } from "@dust-tt/types"; -import type { CreationOptional, ForeignKey } from "sequelize"; +import type { CreationOptional, ForeignKey, NonAttribute } from "sequelize"; import { DataTypes } from "sequelize"; import { AgentConfiguration } from "@app/lib/models/assistant/agent"; import { AgentMessage } from "@app/lib/models/assistant/conversation"; import { frontSequelize } from "@app/lib/resources/storage"; +import { FileModel } from "@app/lib/resources/storage/models/files"; import { BaseModel } from "@app/lib/resources/storage/wrappers"; export class AgentDustAppRunConfiguration extends BaseModel { @@ -88,7 +89,11 @@ export class AgentDustAppRunAction extends BaseModel { declare step: number; declare agentMessageId: ForeignKey; + declare resultsFileId: ForeignKey | null; + declare resultsFileSnippet: string | null; + declare resultsFile: NonAttribute; } + AgentDustAppRunAction.init( { createdAt: { @@ -109,7 +114,6 @@ AgentDustAppRunAction.init( type: DataTypes.STRING, allowNull: false, }, - appWorkspaceId: { type: DataTypes.STRING, allowNull: false, @@ -142,6 +146,10 @@ AgentDustAppRunAction.init( type: DataTypes.INTEGER, allowNull: false, }, + resultsFileSnippet: { + type: DataTypes.TEXT, + allowNull: true, + }, }, { modelName: "agent_dust_app_run_action", @@ -151,6 +159,10 @@ AgentDustAppRunAction.init( fields: ["agentMessageId"], concurrently: true, }, + { + fields: ["resultsFileId"], + concurrently: true, + }, ], } ); @@ -162,3 +174,13 @@ AgentDustAppRunAction.belongsTo(AgentMessage, { AgentMessage.hasMany(AgentDustAppRunAction, { foreignKey: { name: "agentMessageId", allowNull: false }, }); + +FileModel.hasMany(AgentDustAppRunAction, { + foreignKey: { name: "resultsFileId", allowNull: true }, + onDelete: "SET NULL", +}); +AgentDustAppRunAction.belongsTo(FileModel, { + as: "resultsFile", + foreignKey: { name: "resultsFileId", allowNull: true }, + onDelete: "SET NULL", +}); diff --git a/front/migrations/db/migration_145.sql b/front/migrations/db/migration_145.sql new file mode 100644 index 000000000000..fdf97a5aa952 --- /dev/null +++ b/front/migrations/db/migration_145.sql @@ -0,0 +1,7 @@ +-- Migration created on Jan 21, 2025 +ALTER TABLE "public"."agent_dust_app_run_actions" + ADD COLUMN "resultsFileId" integer REFERENCES "public"."files" ("id") ON DELETE SET NULL, + ADD COLUMN "resultsFileSnippet" text; + +-- Add index for resultsFileId +CREATE INDEX CONCURRENTLY IF NOT EXISTS "agent_dust_app_run_actions_results_file_id" ON "public"."agent_dust_app_run_actions" ("resultsFileId"); diff --git a/types/src/front/assistant/actions/dust_app_run.ts b/types/src/front/assistant/actions/dust_app_run.ts index edd6e12560e7..f447c9c03853 100644 --- a/types/src/front/assistant/actions/dust_app_run.ts +++ b/types/src/front/assistant/actions/dust_app_run.ts @@ -33,6 +33,8 @@ export interface DustAppRunActionType extends BaseAction { functionCallId: string | null; functionCallName: string | null; step: number; + resultsFileId: string | null; + resultsFileSnippet: string | null; type: "dust_app_run_action"; }