improve batch insert (#3287)

Helicone · Feb 19, 2025 · 59b4c71 · 59b4c71
1 parent a1f348f
commit 59b4c71
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 10 deletions.
diff --git a/valhalla/jawn/src/managers/experiment/ExperimentV2Manager.ts b/valhalla/jawn/src/managers/experiment/ExperimentV2Manager.ts
@@ -431,16 +431,11 @@ export class ExperimentV2Manager extends BaseManager {
       }
 
       const inputManager = new InputsManager(this.authParams);
-
-      await Promise.all(
-        inputs.map(async (input) => {
-          await inputManager.createInputRecord(
-            experiment.copied_original_prompt_version ?? "",
-            input,
-            undefined,
-            experimentId
-          );
-        })
+      await inputManager.createInputRecords(
+        experiment.copied_original_prompt_version ?? "",
+        inputs,
+        undefined,
+        experimentId
       );
 
       return ok(null);

diff --git a/valhalla/jawn/src/managers/inputs/InputsManager.ts b/valhalla/jawn/src/managers/inputs/InputsManager.ts
@@ -126,6 +126,61 @@ export class InputsManager extends BaseManager {
     return ok(inputRecordId);
   }
 
+  async createInputRecords(
+    promptVersionId: string,
+    inputs: Record<string, string>[],
+    sourceRequest?: string,
+    experimentId?: string
+  ): Promise<Result<string[], string>> {
+    const existingPrompt = await supabaseServer.client
+      .from("prompts_versions")
+      .select("*")
+      .eq("id", promptVersionId)
+      .eq("organization", this.authParams.organizationId)
+      .single();
+
+    if (existingPrompt.error || !existingPrompt.data) {
+      return err(existingPrompt.error?.message ?? "Prompt version not found");
+    }
+
+    // Generate UUIDs for each input record
+    const inputRecordIds = inputs.map(() => randomUUID());
+
+    // Create the VALUES part of the query dynamically
+    const values = inputs
+      .map(
+        (_, index) =>
+          `($${index * 5 + 1}, $${index * 5 + 2}, $${index * 5 + 3}, $${
+            index * 5 + 4
+          }, $${index * 5 + 5})`
+      )
+      .join(",");
+
+    // Flatten parameters array
+    const params = inputs.flatMap((input, index) => [
+      inputRecordIds[index],
+      JSON.stringify(input),
+      sourceRequest,
+      promptVersionId,
+      experimentId,
+    ]);
+
+    const result = await dbExecute<PromptInputRecord>(
+      `
+      INSERT INTO prompt_input_record (id, inputs, source_request, prompt_version, experiment_id)
+      VALUES ${values}
+      RETURNING id
+      `,
+      params
+    );
+
+    if (result.error) {
+      return err(result.error);
+    }
+
+    return ok(inputRecordIds);
+  }
+
   async updateInputRecord(
     inputRecordId: string,
     inputs: Record<string, string>