load-generator: Add missing arg validators in scenario-runner (#43696)

GitOrigin-RevId: 4c3b04b2133a48fef9393160d13a5bb8ecd8f793

Author: Nicolapps

npm-packages/scenario-runner/convex/actions/externalNodeDeps.ts

@@ -2,6 +2,8 @@
 import { load } from "langchain/load";
 import { Tiktoken } from "tiktoken/lite";
 import { action } from "../_generated/server";
+import { v } from "convex/values";
+
 // eslint-disable-next-line @typescript-eslint/no-require-imports
 const gpt2_base = require("tiktoken/encoders/gpt2.json");
 
@@ -10,7 +12,10 @@ const gpt2_base = require("tiktoken/encoders/gpt2.json");
 load;
 
 export const encode = action({
-  handler: (_, { str }: { str: string }): number[] => {
+  args: {
+    str: v.string(),
+  },
+  handler: (_, { str }): number[] => {
     const enc = new Tiktoken(
       gpt2_base.bpe_ranks,
       gpt2_base.special_tokens,

npm-packages/scenario-runner/convex/query_index.ts

@@ -1,6 +1,7 @@
 import { DatabaseReader, query } from "./_generated/server";
 import { Id } from "./_generated/dataModel";
 import { CACHE_BREAKER_ARGS, MessagesTable } from "./common";
+import { v } from "convex/values";
 
 export const queryMessages = query({
   args: CACHE_BREAKER_ARGS,
@@ -33,20 +34,13 @@ export const queryMessagesWithSearch = query({
 });
 
 export const queryMessagesWithArgs = query({
-  handler: async (
-    { db },
-    {
-      channel,
-      rand,
-      limit,
-      table,
-    }: {
-      channel: string;
-      rand: number;
-      limit: number;
-      table: MessagesTable;
-    },
-  ) => {
+  args: {
+    channel: v.string(),
+    rand: v.number(),
+    limit: v.number(),
+    table: v.union(v.literal("messages"), v.literal("messages_with_search")),
+  },
+  handler: async ({ db }, { channel, rand, limit, table }) => {
     return await queryMessagesHelper(db, channel, rand, limit, table);
   },
 });
@@ -67,7 +61,10 @@ function queryMessagesHelper(
 }
 
 export const queryMessagesById = query({
-  handler: async ({ db }, { id }: { id: Id<MessagesTable> }) => {
-    return await db.get(id);
+  args: {
+    id: v.union(v.id("messages"), v.id("messages_with_search")),
+  },
+  handler: async ({ db }, { id }) => {
+    return await db.get(id satisfies Id<MessagesTable>);
   },
 });

npm-packages/scenario-runner/convex/setup.ts

@@ -5,10 +5,11 @@ import { v } from "convex/values";
 import { rand } from "./common";
 
 export const setupMessages = mutation({
-  handler: async (
-    { db },
-    { rows, channel }: { rows: number; channel: string },
-  ): Promise<void> => {
+  args: {
+    rows: v.number(),
+    channel: v.string(),
+  },
+  handler: async ({ db }, { rows, channel }): Promise<void> => {
     console.log(`Fill the messages table with ${rows} rows`);
     for (let i = 0; i < rows; i++) {
       // The first record always has rand=0. This is used in update_first