Skip to content

Commit

Permalink
overwrite Devin's mess
Browse files Browse the repository at this point in the history
  • Loading branch information
Henry Fontanier committed Jan 21, 2025
1 parent 5043aff commit 3f81c7c
Show file tree
Hide file tree
Showing 6 changed files with 218 additions and 57 deletions.
147 changes: 136 additions & 11 deletions front/lib/api/assistant/actions/dust_app_run.ts
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
import { DustAPI } from "@dust-tt/client";
import type {
ActionGeneratedFileType,
AgentActionSpecification,
DatasetSchema,
DustAppParameters,
DustAppRunActionType,
DustAppRunBlockEvent,
DustAppRunConfigurationType,
DustAppRunErrorEvent,
Expand All @@ -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";
Expand All @@ -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";

Expand All @@ -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 {
Expand All @@ -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;
Expand All @@ -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 {
Expand Down Expand Up @@ -303,6 +316,8 @@ export class DustAppRunConfigurationServerRunner extends BaseActionConfiguration
functionCallName: actionConfiguration.name,
agentMessageId: agentMessage.agentMessageId,
step,
resultsFileId: null,
resultsFileSnippet: null,
});

yield {
Expand All @@ -322,6 +337,9 @@ export class DustAppRunConfigurationServerRunner extends BaseActionConfiguration
functionCallName: actionConfiguration.name,
agentMessageId: agentMessage.agentMessageId,
step,
resultsFileId: null,
resultsFileSnippet: null,
generatedFiles: [],
}),
};

Expand Down Expand Up @@ -462,6 +480,9 @@ export class DustAppRunConfigurationServerRunner extends BaseActionConfiguration
output: null,
agentMessageId: agentMessage.agentMessageId,
step: action.step,
resultsFileId: null,
resultsFileSnippet: null,
generatedFiles: [],
}),
};
}
Expand All @@ -486,12 +507,79 @@ export class DustAppRunConfigurationServerRunner extends BaseActionConfiguration
}
}

const output = sanitizeJSONOutput(lastBlockOutput);
function isValidStructuredOutput(output: unknown): output is {
__dust_structured_output?: Array<
Record<string, string | number | boolean>
>;
} {
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(
Expand All @@ -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] : [],
}),
};
}
Expand Down Expand Up @@ -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<DustAppRunActionType[]> {
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,
Expand All @@ -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`;
}
47 changes: 47 additions & 0 deletions front/lib/api/assistant/actions/result_file_helpers.ts
View file
Original file line number Diff line number Diff line change
@@ -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<Record<string, string | number | boolean>>;
}
): Promise<{
file: FileResource;
snippet: string;
}> {
const toCsv = (
records: Array<Record<string, string | number | boolean>>,
options: { header: boolean } = { header: true }
): Promise<string> => {
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) };
}
46 changes: 2 additions & 44 deletions front/lib/api/assistant/actions/tables_query.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -567,7 +565,7 @@ export class TablesQueryConfigurationServerRunner extends BaseActionConfiguratio
output: sanitizedOutput,
});

const { file, snippet } = await getTablesQueryOutputCsvFileAndSnippet(
const { file, snippet } = await getToolResultOutputCsvFileAndSnippet(
auth,
{
title: queryTitle,
Expand Down Expand Up @@ -616,43 +614,3 @@ export class TablesQueryConfigurationServerRunner extends BaseActionConfiguratio
return;
}
}

async function getTablesQueryOutputCsvFileAndSnippet(
auth: Authenticator,
{
title,
conversationId,
results,
}: {
title: string;
conversationId: string;
results: Array<Record<string, string | number | boolean>>;
}
): Promise<{
file: FileResource;
snippet: string;
}> {
const toCsv = (
records: Array<Record<string, string | number | boolean>>,
options: { header: boolean } = { header: true }
): Promise<string> => {
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) };
}
Loading

0 comments on commit 3f81c7c

Please sign in to comment.