From 8a9a2a1953c3ba5d9c3ec6fd3e22ae0165dca448 Mon Sep 17 00:00:00 2001
From: Tony Powell <apowell@arize.com>
Date: Fri, 25 Oct 2024 15:29:18 -0400
Subject: [PATCH] fix: Bubble errors up from nested invocation parameter schema
 transforms

This allows errors to be displayed in the playground UI when the invocation parameters are totally incorrect on the span
---
 .../__tests__/playgroundUtils.test.ts         | 19 ++++++++++++++++++-
 app/src/pages/playground/schemas.ts           | 19 ++++++++++++++++---
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/app/src/pages/playground/__tests__/playgroundUtils.test.ts b/app/src/pages/playground/__tests__/playgroundUtils.test.ts
index 4ed94250cc..2b108d1814 100644
--- a/app/src/pages/playground/__tests__/playgroundUtils.test.ts
+++ b/app/src/pages/playground/__tests__/playgroundUtils.test.ts
@@ -375,7 +375,7 @@ describe("transformSpanAttributesToPlaygroundInstance", () => {
       playgroundInstance: {
         ...expectedPlaygroundInstanceWithIO,
       },
-      parsingErrors: [],
+      parsingErrors: [MODEL_CONFIG_WITH_INVOCATION_PARAMETERS_PARSING_ERROR],
     });
   });
 
@@ -398,6 +398,23 @@ describe("transformSpanAttributesToPlaygroundInstance", () => {
       parsingErrors: [MODEL_CONFIG_WITH_INVOCATION_PARAMETERS_PARSING_ERROR],
     });
   });
+
+  it("should return invocation parameters parsing errors if they are malformed", () => {
+    const parsedAttributes = {
+      llm: {
+        model_name: "gpt-3.5-turbo",
+        invocation_parameters: '"invalid"',
+      },
+    };
+    expect(getModelConfigFromAttributes(parsedAttributes)).toEqual({
+      modelConfig: {
+        modelName: "gpt-3.5-turbo",
+        provider: "OPENAI",
+        invocationParameters: {},
+      },
+      parsingErrors: [MODEL_CONFIG_WITH_INVOCATION_PARAMETERS_PARSING_ERROR],
+    });
+  });
 });
 
 describe("getChatRole", () => {
diff --git a/app/src/pages/playground/schemas.ts b/app/src/pages/playground/schemas.ts
index 04b81143e9..fdf6e449e1 100644
--- a/app/src/pages/playground/schemas.ts
+++ b/app/src/pages/playground/schemas.ts
@@ -136,8 +136,8 @@ const stringToInvocationParametersSchema = z
   .string()
   .transform((s) => {
     const { json } = safelyParseJSON(s);
-    if (json == null) {
-      return {};
+    if (json == null || typeof json !== "object") {
+      return null;
     }
     // using the invocationParameterSchema as a base,
     // apply all matching keys from the input string,
@@ -156,10 +156,23 @@ const stringToInvocationParametersSchema = z
           ),
         }))
         // reparse the object to ensure the mapped keys are also validated
-        .transform(invocationParameterSchema.parse)
         .parse(json)
     );
   })
+  .transform((v, ctx) => {
+    const result = invocationParameterSchema.safeParse(v);
+    if (!result.success) {
+      // bubble errors up to the original schema
+      result.error.issues.forEach((issue) => {
+        ctx.addIssue(issue);
+      });
+      // https://zod.dev/?id=validating-during-transform
+      // ensures that this schema still infers the "success" type
+      // errors will throw instead
+      return z.NEVER;
+    }
+    return result.data;
+  })
   .default("{}");
 /**
  * The zod schema for llm model config