Spaces:
Paused
Paused
| import readline from "readline"; | |
| import minimist from "minimist"; | |
| // @ts-expect-error: vite-node makes the var available but the typescript compiler doesn't see them | |
| import { env } from "$env/dynamic/private"; | |
| import { faker } from "@faker-js/faker"; | |
| import { ObjectId } from "mongodb"; | |
| // @ts-expect-error: vite-node makes the var available but the typescript compiler doesn't see them | |
| import { collections } from "$lib/server/database"; | |
| import { models } from "../src/lib/server/models.ts"; | |
| import type { User } from "../src/lib/types/User"; | |
| import type { Assistant } from "../src/lib/types/Assistant"; | |
| import type { Conversation } from "../src/lib/types/Conversation"; | |
| import type { Settings } from "../src/lib/types/Settings"; | |
| import { defaultEmbeddingModel } from "../src/lib/server/embeddingModels.ts"; | |
| import { Message } from "../src/lib/types/Message.ts"; | |
| import { addChildren } from "../src/lib/utils/tree/addChildren.ts"; | |
| import { generateSearchTokens } from "../src/lib/utils/searchTokens.ts"; | |
| const rl = readline.createInterface({ | |
| input: process.stdin, | |
| output: process.stdout, | |
| }); | |
| rl.on("close", function () { | |
| process.exit(0); | |
| }); | |
| const possibleFlags = ["reset", "all", "users", "settings", "assistants", "conversations"]; | |
| const argv = minimist(process.argv.slice(2)); | |
| const flags = argv["_"].filter((flag) => possibleFlags.includes(flag)); | |
| async function generateMessages(preprompt?: string): Promise<Message[]> { | |
| const isLinear = faker.datatype.boolean(0.5); | |
| const isInterrupted = faker.datatype.boolean(0.05); | |
| const messages: Message[] = []; | |
| messages.push({ | |
| id: crypto.randomUUID(), | |
| from: "system", | |
| content: preprompt ?? "", | |
| createdAt: faker.date.recent({ days: 30 }), | |
| updatedAt: faker.date.recent({ days: 30 }), | |
| }); | |
| let isUser = true; | |
| let lastId = messages[0].id; | |
| if (isLinear) { | |
| const convLength = faker.number.int({ min: 1, max: 25 }) * 2; // must always be even | |
| for (let i = 0; i < convLength; i++) { | |
| lastId = addChildren( | |
| { | |
| messages, | |
| rootMessageId: messages[0].id, | |
| }, | |
| { | |
| from: isUser ? "user" : "assistant", | |
| content: faker.lorem.sentence({ | |
| min: 10, | |
| max: isUser ? 50 : 200, | |
| }), | |
| createdAt: faker.date.recent({ days: 30 }), | |
| updatedAt: faker.date.recent({ days: 30 }), | |
| interrupted: i === convLength - 1 && isInterrupted, | |
| }, | |
| lastId | |
| ); | |
| isUser = !isUser; | |
| } | |
| } else { | |
| const convLength = faker.number.int({ min: 2, max: 200 }); | |
| for (let i = 0; i < convLength; i++) { | |
| addChildren( | |
| { | |
| messages, | |
| rootMessageId: messages[0].id, | |
| }, | |
| { | |
| from: isUser ? "user" : "assistant", | |
| content: faker.lorem.sentence({ | |
| min: 10, | |
| max: isUser ? 50 : 200, | |
| }), | |
| createdAt: faker.date.recent({ days: 30 }), | |
| updatedAt: faker.date.recent({ days: 30 }), | |
| interrupted: i === convLength - 1 && isInterrupted, | |
| }, | |
| faker.helpers.arrayElement([ | |
| messages[0].id, | |
| ...messages.filter((m) => m.from === (isUser ? "assistant" : "user")).map((m) => m.id), | |
| ]) | |
| ); | |
| isUser = !isUser; | |
| } | |
| } | |
| return messages; | |
| } | |
| async function seed() { | |
| console.log("Seeding..."); | |
| const modelIds = models.map((model) => model.id); | |
| if (flags.includes("reset")) { | |
| console.log("Starting reset of DB"); | |
| await collections.users.deleteMany({}); | |
| await collections.settings.deleteMany({}); | |
| await collections.assistants.deleteMany({}); | |
| await collections.conversations.deleteMany({}); | |
| console.log("Reset done"); | |
| } | |
| if (flags.includes("users") || flags.includes("all")) { | |
| console.log("Creating 100 new users"); | |
| const newUsers: User[] = Array.from({ length: 100 }, () => ({ | |
| _id: new ObjectId(), | |
| createdAt: faker.date.recent({ days: 30 }), | |
| updatedAt: faker.date.recent({ days: 30 }), | |
| username: faker.internet.userName(), | |
| name: faker.person.fullName(), | |
| hfUserId: faker.string.alphanumeric(24), | |
| avatarUrl: faker.image.avatar(), | |
| })); | |
| await collections.users.insertMany(newUsers); | |
| console.log("Done creating users."); | |
| } | |
| const users = await collections.users.find().toArray(); | |
| if (flags.includes("settings") || flags.includes("all")) { | |
| console.log("Updating settings for all users"); | |
| users.forEach(async (user) => { | |
| const settings: Settings = { | |
| userId: user._id, | |
| shareConversationsWithModelAuthors: faker.datatype.boolean(0.25), | |
| hideEmojiOnSidebar: faker.datatype.boolean(0.25), | |
| ethicsModalAcceptedAt: faker.date.recent({ days: 30 }), | |
| activeModel: faker.helpers.arrayElement(modelIds), | |
| createdAt: faker.date.recent({ days: 30 }), | |
| updatedAt: faker.date.recent({ days: 30 }), | |
| customPrompts: {}, | |
| assistants: [], | |
| }; | |
| await collections.settings.updateOne( | |
| { userId: user._id }, | |
| { $set: { ...settings } }, | |
| { upsert: true } | |
| ); | |
| }); | |
| console.log("Done updating settings."); | |
| } | |
| if (flags.includes("assistants") || flags.includes("all")) { | |
| console.log("Creating assistants for all users"); | |
| await Promise.all( | |
| users.map(async (user) => { | |
| const name = faker.animal.insect(); | |
| const assistants = faker.helpers.multiple<Assistant>( | |
| () => ({ | |
| _id: new ObjectId(), | |
| name, | |
| createdById: user._id, | |
| createdByName: user.username, | |
| createdAt: faker.date.recent({ days: 30 }), | |
| updatedAt: faker.date.recent({ days: 30 }), | |
| userCount: faker.number.int({ min: 1, max: 100000 }), | |
| featured: faker.datatype.boolean(0.25), | |
| modelId: faker.helpers.arrayElement(modelIds), | |
| description: faker.lorem.sentence(), | |
| preprompt: faker.hacker.phrase(), | |
| exampleInputs: faker.helpers.multiple(() => faker.lorem.sentence(), { | |
| count: faker.number.int({ min: 0, max: 4 }), | |
| }), | |
| searchTokens: generateSearchTokens(name), | |
| last24HoursCount: faker.number.int({ min: 0, max: 1000 }), | |
| }), | |
| { count: faker.number.int({ min: 3, max: 10 }) } | |
| ); | |
| await collections.assistants.insertMany(assistants); | |
| await collections.settings.updateOne( | |
| { userId: user._id }, | |
| { $set: { assistants: assistants.map((a) => a._id.toString()) } }, | |
| { upsert: true } | |
| ); | |
| }) | |
| ); | |
| console.log("Done creating assistants."); | |
| } | |
| if (flags.includes("conversations") || flags.includes("all")) { | |
| console.log("Creating conversations for all users"); | |
| await Promise.all( | |
| users.map(async (user) => { | |
| const conversations = faker.helpers.multiple( | |
| async () => { | |
| const settings = await collections.settings.findOne<Settings>({ userId: user._id }); | |
| const assistantId = | |
| settings?.assistants && settings.assistants.length > 0 && faker.datatype.boolean(0.1) | |
| ? faker.helpers.arrayElement<ObjectId>(settings.assistants) | |
| : undefined; | |
| const preprompt = | |
| (assistantId | |
| ? await collections.assistants | |
| .findOne({ _id: assistantId }) | |
| .then((assistant: Assistant) => assistant?.preprompt ?? "") | |
| : faker.helpers.maybe(() => faker.hacker.phrase(), { probability: 0.5 })) ?? ""; | |
| const messages = await generateMessages(preprompt); | |
| const conv = { | |
| _id: new ObjectId(), | |
| userId: user._id, | |
| assistantId, | |
| preprompt, | |
| createdAt: faker.date.recent({ days: 145 }), | |
| updatedAt: faker.date.recent({ days: 145 }), | |
| model: faker.helpers.arrayElement(modelIds), | |
| title: faker.internet.emoji() + " " + faker.hacker.phrase(), | |
| embeddingModel: defaultEmbeddingModel.id, | |
| messages, | |
| rootMessageId: messages[0].id, | |
| } satisfies Conversation; | |
| return conv; | |
| }, | |
| { count: faker.number.int({ min: 10, max: 200 }) } | |
| ); | |
| await collections.conversations.insertMany(await Promise.all(conversations)); | |
| }) | |
| ); | |
| console.log("Done creating conversations."); | |
| } | |
| } | |
| // run seed | |
| (async () => { | |
| try { | |
| rl.question( | |
| "You're about to run a seeding script on the following MONGODB_URL: \x1b[31m" + | |
| env.MONGODB_URL + | |
| "\x1b[0m\n\n With the following flags: \x1b[31m" + | |
| flags.join("\x1b[0m , \x1b[31m") + | |
| "\x1b[0m\n \n\n Are you sure you want to continue? (yes/no): ", | |
| async (confirm) => { | |
| if (confirm !== "yes") { | |
| console.log("Not 'yes', exiting."); | |
| rl.close(); | |
| process.exit(0); | |
| } | |
| console.log("Starting seeding..."); | |
| await seed(); | |
| console.log("Seeding done."); | |
| rl.close(); | |
| } | |
| ); | |
| } catch (e) { | |
| console.error(e); | |
| process.exit(1); | |
| } | |
| })(); | |