Adding a meta column to datasets to link prompts to them (#1900)

Helicone · May 17, 2024 · 9ac4914 · 9ac4914
1 parent 089bf48
commit 9ac4914
Show file tree

Hide file tree

Showing 16 changed files with 135 additions and 35 deletions.
diff --git a/docs/swagger.json b/docs/swagger.json
@@ -1485,6 +1485,21 @@
 					}
 				]
 			},
+			"DatasetMetadata": {
+				"properties": {
+					"promptId": {
+						"type": "string"
+					},
+					"inputRecordsIds": {
+						"items": {
+							"type": "string"
+						},
+						"type": "array"
+					}
+				},
+				"type": "object",
+				"additionalProperties": false
+			},
 			"NewDatasetParams": {
 				"properties": {
 					"datasetName": {
@@ -1495,6 +1510,9 @@
 							"type": "string"
 						},
 						"type": "array"
+					},
+					"meta": {
+						"$ref": "#/components/schemas/DatasetMetadata"
 					}
 				},
 				"required": [
@@ -1608,20 +1626,16 @@
 					"name": {
 						"type": "string"
 					},
-					"request_ids": {
-						"items": {
-							"type": "string"
-						},
-						"type": "array"
-					},
 					"created_at": {
 						"type": "string"
+					},
+					"meta": {
+						"$ref": "#/components/schemas/DatasetMetadata"
 					}
 				},
 				"required": [
 					"id",
 					"name",
-					"request_ids",
 					"created_at"
 				],
 				"type": "object",
@@ -2836,7 +2850,11 @@
 					"content": {
 						"application/json": {
 							"schema": {
-								"properties": {},
+								"properties": {
+									"promptId": {
+										"type": "string"
+									}
+								},
 								"type": "object"
 							}
 						}

diff --git a/helicone-cron/src/db/database.types.ts b/helicone-cron/src/db/database.types.ts
@@ -250,18 +250,21 @@ export type Database = {
         Row: {
           created_at: string | null
           id: string
+          meta: Json | null
           name: string | null
           organization: string
         }
         Insert: {
           created_at?: string | null
           id?: string
+          meta?: Json | null
           name?: string | null
           organization: string
         }
         Update: {
           created_at?: string | null
           id?: string
+          meta?: Json | null
           name?: string | null
           organization?: string
         }

diff --git a/supabase/migrations/20240516222421_experiments-meta.sql b/supabase/migrations/20240516222421_experiments-meta.sql
@@ -0,0 +1,3 @@
+alter table "public"."experiment_dataset_v2" add column "meta" jsonb;
+
+
diff --git a/valhalla/jawn/src/controllers/public/experimentDatasetController.ts b/valhalla/jawn/src/controllers/public/experimentDatasetController.ts
@@ -18,16 +18,22 @@ type DatasetFilterNode =
   | DatasetFilterBranch
   | "all";
 
+export interface DatasetMetadata {
+  promptId?: string;
+  inputRecordsIds?: string[];
+}
+
 export interface NewDatasetParams {
   datasetName: string;
   requestIds: string[];
+  meta?: DatasetMetadata;
 }
 
 export interface DatasetResult {
   id: string;
   name: string;
-  request_ids: string[];
   created_at: string;
+  meta?: DatasetMetadata;
 }
 
 export interface RandomDatasetParams {
@@ -99,11 +105,13 @@ export class ExperimentDatasetController extends Controller {
   @Post("/query")
   public async getDatasets(
     @Body()
-    requestBody: {},
+    requestBody: {
+      promptId?: string;
+    },
     @Request() request: JawnAuthenticatedRequest
   ): Promise<Result<DatasetResult[], string>> {
     const datasetManager = new DatasetManager(request.authParams);
-    const result = await datasetManager.getDatasets();
+    const result = await datasetManager.getDatasets(requestBody.promptId);
     if (result.error || !result.data) {
       this.setStatus(500);
     } else {

diff --git a/valhalla/jawn/src/lib/db/database.types.ts b/valhalla/jawn/src/lib/db/database.types.ts
@@ -250,18 +250,21 @@ export type Database = {
         Row: {
           created_at: string | null
           id: string
+          meta: Json | null
           name: string | null
           organization: string
         }
         Insert: {
           created_at?: string | null
           id?: string
+          meta?: Json | null
           name?: string | null
           organization: string
         }
         Update: {
           created_at?: string | null
           id?: string
+          meta?: Json | null
           name?: string | null
           organization?: string
         }

diff --git a/valhalla/jawn/src/managers/dataset/DatasetManager.ts b/valhalla/jawn/src/managers/dataset/DatasetManager.ts
@@ -1,5 +1,6 @@
 // src/users/usersService.ts
 import {
+  DatasetMetadata,
   DatasetResult,
   NewDatasetParams,
   RandomDatasetParams,
@@ -19,28 +20,32 @@ import { buildFilterPostgres } from "../../lib/shared/filters/filters";
 import { resultMap } from "../../lib/shared/result";
 import { User } from "../../models/user";
 import { BaseManager } from "../BaseManager";
+import { Json } from "../../lib/db/database.types";
 
 // A post request should not contain an id.
 export type UserCreationParams = Pick<User, "email" | "name" | "phoneNumbers">;
 
 export class DatasetManager extends BaseManager {
-  async getDatasets(): Promise<Result<DatasetResult[], string>> {
+  async getDatasets(
+    promptId?: string
+  ): Promise<Result<DatasetResult[], string>> {
     const result = dbExecute<{
       id: string;
       name: string;
-      request_ids: string[];
       created_at: string;
+      meta: DatasetMetadata;
     }>(
       `
     SELECT 
       id,
       name,
-      created_at
+      created_at,
+      meta
     FROM experiment_dataset_v2
-    WHERE organization = $1
+    WHERE organization = $1 ${promptId ? "AND meta->>'promptId' = $2" : ""}
     LIMIT 100
     `,
-      [this.authParams.organizationId]
+      [this.authParams.organizationId].concat(promptId ? [promptId] : [])
     );
     return result;
   }
@@ -51,6 +56,7 @@ export class DatasetManager extends BaseManager {
       .insert({
         name: params.datasetName,
         organization: this.authParams.organizationId,
+        meta: (params.meta ?? null) as Json,
       })
       .select("*")
       .single();

diff --git a/valhalla/jawn/src/managers/inputs/InputsManager.ts b/valhalla/jawn/src/managers/inputs/InputsManager.ts
@@ -16,6 +16,7 @@ import { resultMap } from "../../lib/shared/result";
 import { User } from "../../models/user";
 import { BaseManager } from "../BaseManager";
 import { S3Client } from "../../lib/shared/db/s3Client";
+import { RequestResponseBodyStore } from "../../lib/stores/request/RequestResponseBodyStore";
 
 async function fetchImageAsBase64(url: string): Promise<string> {
   try {
@@ -113,12 +114,18 @@ export class InputsManager extends BaseManager {
       `,
       [this.authParams.organizationId, promptVersion, limit]
     );
+    const bodyStore = new RequestResponseBodyStore(
+      this.authParams.organizationId
+    );
 
     return promiseResultMap(result, async (data) => {
       return Promise.all(
         data.map(async (record) => {
           return {
             ...record,
+            response_body:
+              (await bodyStore.getRequestResponseBody(record.source_request))
+                .data?.response ?? {},
             inputs: await getAllSignedURLsFromInputs(
               record.inputs,
               this.authParams.organizationId,

diff --git a/valhalla/jawn/src/tsoa-build/public/routes.ts b/valhalla/jawn/src/tsoa-build/public/routes.ts
@@ -545,11 +545,21 @@ const models: TsoaRoute.Models = {
         "type": {"dataType":"union","subSchemas":[{"ref":"ResultSuccess__datasetId-string__"},{"ref":"ResultError_string_"}],"validators":{}},
     },
     // WARNING: This file was auto-generated with tsoa. Please do not modify it. Re-run tsoa to re-generate this file: https://github.com/lukeautry/tsoa
+    "DatasetMetadata": {
+        "dataType": "refObject",
+        "properties": {
+            "promptId": {"dataType":"string"},
+            "inputRecordsIds": {"dataType":"array","array":{"dataType":"string"}},
+        },
+        "additionalProperties": false,
+    },
+    // WARNING: This file was auto-generated with tsoa. Please do not modify it. Re-run tsoa to re-generate this file: https://github.com/lukeautry/tsoa
     "NewDatasetParams": {
         "dataType": "refObject",
         "properties": {
             "datasetName": {"dataType":"string","required":true},
             "requestIds": {"dataType":"array","array":{"dataType":"string"},"required":true},
+            "meta": {"ref":"DatasetMetadata"},
         },
         "additionalProperties": false,
     },
@@ -595,8 +605,8 @@ const models: TsoaRoute.Models = {
         "properties": {
             "id": {"dataType":"string","required":true},
             "name": {"dataType":"string","required":true},
-            "request_ids": {"dataType":"array","array":{"dataType":"string"},"required":true},
             "created_at": {"dataType":"string","required":true},
+            "meta": {"ref":"DatasetMetadata"},
         },
         "additionalProperties": false,
     },
@@ -1233,7 +1243,7 @@ export function RegisterRoutes(app: Router) {
 
             function ExperimentDatasetController_getDatasets(request: ExRequest, response: ExResponse, next: any) {
             const args: Record<string, TsoaRoute.ParameterSchema> = {
-                    requestBody: {"in":"body","name":"requestBody","required":true,"dataType":"nestedObjectLiteral","nestedProperties":{}},
+                    requestBody: {"in":"body","name":"requestBody","required":true,"dataType":"nestedObjectLiteral","nestedProperties":{"promptId":{"dataType":"string"}}},
                     request: {"in":"request","name":"request","required":true,"dataType":"object"},
             };
 

diff --git a/valhalla/jawn/src/tsoa-build/public/swagger.json b/valhalla/jawn/src/tsoa-build/public/swagger.json
@@ -1485,6 +1485,21 @@
 					}
 				]
 			},
+			"DatasetMetadata": {
+				"properties": {
+					"promptId": {
+						"type": "string"
+					},
+					"inputRecordsIds": {
+						"items": {
+							"type": "string"
+						},
+						"type": "array"
+					}
+				},
+				"type": "object",
+				"additionalProperties": false
+			},
 			"NewDatasetParams": {
 				"properties": {
 					"datasetName": {
@@ -1495,6 +1510,9 @@
 							"type": "string"
 						},
 						"type": "array"
+					},
+					"meta": {
+						"$ref": "#/components/schemas/DatasetMetadata"
 					}
 				},
 				"required": [
@@ -1608,20 +1626,16 @@
 					"name": {
 						"type": "string"
 					},
-					"request_ids": {
-						"items": {
-							"type": "string"
-						},
-						"type": "array"
-					},
 					"created_at": {
 						"type": "string"
+					},
+					"meta": {
+						"$ref": "#/components/schemas/DatasetMetadata"
 					}
 				},
 				"required": [
 					"id",
 					"name",
-					"request_ids",
 					"created_at"
 				],
 				"type": "object",
@@ -2836,7 +2850,11 @@
 					"content": {
 						"application/json": {
 							"schema": {
-								"properties": {},
+								"properties": {
+									"promptId": {
+										"type": "string"
+									}
+								},
 								"type": "object"
 							}
 						}

diff --git a/web/components/templates/prompts/experiments/id/experimentIdPage.tsx b/web/components/templates/prompts/experiments/id/experimentIdPage.tsx
@@ -274,8 +274,8 @@ const ExperimentIdPage = (props: PromptIdPageProps) => {
                                     `w-max items-center rounded-lg px-2 py-1 -my-1 text-xs font-medium ring-1 ring-inset`
                                   )}
                                 >
-                                  {run.testResult.response?.completionTokens}{" "}
-                                  input tokens
+                                  {run.testResult.response?.promptTokens} input
+                                  tokens
                                 </span>
 
                                 <span
@@ -284,8 +284,8 @@ const ExperimentIdPage = (props: PromptIdPageProps) => {
                                     `w-max items-center rounded-lg px-2 py-1 -my-1 text-xs font-medium ring-1 ring-inset`
                                   )}
                                 >
-                                  {run.testResult.response?.promptTokens} prompt
-                                  tokens
+                                  {run.testResult.response?.completionTokens}{" "}
+                                  output tokens
                                 </span>
                                 <ModelPill
                                   model={run.testResult.response?.model ?? ""}

diff --git a/web/components/templates/prompts/id/promptNewExperiment.tsx b/web/components/templates/prompts/id/promptNewExperiment.tsx
@@ -98,7 +98,7 @@ const PromptNewExperimentPage = (props: PromptIdPageProps) => {
     datasets: datasets,
     isLoading: isDataSetsLoading,
     refetch: refetchDataSets,
-  } = useGetDataSets();
+  } = useGetDataSets(id);
 
   const selectedDataset = datasets.find(
     (dataset) => dataset.id === selectedDatasetId
@@ -448,6 +448,9 @@ const PromptNewExperimentPage = (props: PromptIdPageProps) => {
         open={openConfirmModal}
         setOpen={setOpenConfirmModal}
         requestIds={requestIds}
+        meta={{
+          promptVersionId: prompt?.id,
+        }}
         onSuccess={(datasetId) => {
           setSelectedDatasetId(datasetId);