Skip to content

Commit

Permalink
Add script to update assistants models. (#4045)
Browse files Browse the repository at this point in the history
* Add script to update assistants models.

* ✨
  • Loading branch information
flvndvd authored Feb 29, 2024
1 parent aef5d87 commit dc263e7
Showing 1 changed file with 90 additions and 0 deletions.
90 changes: 90 additions & 0 deletions front/scripts/update_assistants_models.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
import type { ModelId, SupportedModel } from "@dust-tt/types";
import { SUPPORTED_MODEL_CONFIGS } from "@dust-tt/types";

import {
AgentConfiguration,
AgentGenerationConfiguration,
Workspace,
} from "@app/lib/models";
import { makeScript } from "@app/scripts/helpers";

type SupportedModelIds = SupportedModel["modelId"];

async function updateWorkspaceAssistants(
workspaceId: ModelId,
fromModel: string,
toModel: SupportedModelIds,
execute: boolean
) {
const agentConfigurations = await AgentConfiguration.findAll({
where: { workspaceId },
});

for (const ac of agentConfigurations) {
if (!ac.generationConfigurationId) {
console.log(
`Skipping ${ac.name}(${ac.sId}): (no generation configuration).`
);
continue;
}

const generationConfiguration = await AgentGenerationConfiguration.findOne({
where: { id: ac.generationConfigurationId },
});

if (!generationConfiguration) {
throw new Error(
`Generation configuration ${ac.generationConfigurationId} not found.`
);
}

if (generationConfiguration.modelId === fromModel) {
if (execute) {
await generationConfiguration.update({ modelId: toModel });
}

console.log(
`${execute ? "" : "[DRYRUN]"} Updated ${ac.name}(${ac.sId}) from ${
generationConfiguration.modelId
} to ${toModel}.`
);
}
}
}

makeScript(
{
fromModel: {
type: "string",
demandOption: true,
},
toModel: {
type: "string",
choices: SUPPORTED_MODEL_CONFIGS.map((m) => m.modelId),
demandOption: true,
},
workspaceIds: {
type: "array",
demandOption: false,
default: [],
description:
"List of workspace identifiers, separated by a space, for which the feature flag should be toggled.",
},
},
async ({ fromModel, toModel, workspaceIds, execute }) => {
const whereClause = workspaceIds.length > 0 ? { sId: workspaceIds } : {};
const workspaces = await Workspace.findAll({
attributes: ["id", "name", "sId"],
where: whereClause,
});

for (const workspace of workspaces) {
await updateWorkspaceAssistants(
workspace.id,
fromModel,
toModel as SupportedModelIds,
execute
);
}
}
);

0 comments on commit dc263e7

Please sign in to comment.