From d007a2083948f0ecdabca1da4d1089b0d42822e6 Mon Sep 17 00:00:00 2001
From: Tony Powell <apowell@arize.com>
Date: Tue, 15 Oct 2024 17:41:12 -0400
Subject: [PATCH] refactor(playground): Add and use type guards for playground
 input type

---
 app/src/store/playground/playgroundStore.tsx |  7 ++++---
 app/src/store/playground/types.ts            | 18 +++++++++++++++++-
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/app/src/store/playground/playgroundStore.tsx b/app/src/store/playground/playgroundStore.tsx
index 1b5d864c9a..22cc9965c3 100644
--- a/app/src/store/playground/playgroundStore.tsx
+++ b/app/src/store/playground/playgroundStore.tsx
@@ -9,6 +9,7 @@ import { assertUnreachable } from "@phoenix/typeUtils";
 import {
   GenAIOperationType,
   InitialPlaygroundState,
+  isManualInput,
   PlaygroundChatTemplate,
   PlaygroundInputMode,
   PlaygroundInstance,
@@ -334,7 +335,7 @@ export const createPlaygroundStore = (
     },
     setVariableValue: (key: string, value: string) => {
       const input = get().input;
-      if ("variablesValueCache" in input) {
+      if (isManualInput(input)) {
         set({
           input: {
             ...input,
@@ -356,7 +357,7 @@ export type PlaygroundStore = ReturnType<typeof createPlaygroundStore>;
  * @returns the variable keys
  */
 export const selectInputVariableKeys = (state: PlaygroundState) => {
-  if ("variableKeys" in state.input) {
+  if (isManualInput(state.input)) {
     return state.input.variableKeys;
   }
   return [];
@@ -368,7 +369,7 @@ export const selectInputVariableKeys = (state: PlaygroundState) => {
  * @returns the derived input variables
  */
 export const selectDerivedInputVariables = (state: PlaygroundState) => {
-  if ("variableKeys" in state.input) {
+  if (isManualInput(state.input)) {
     const input = state.input;
     const variableKeys = input.variableKeys;
     const variablesValueCache = input.variablesValueCache;
diff --git a/app/src/store/playground/types.ts b/app/src/store/playground/types.ts
index 55d2fe2781..12cc4cfb3d 100644
--- a/app/src/store/playground/types.ts
+++ b/app/src/store/playground/types.ts
@@ -49,7 +49,7 @@ type DatasetInput = {
 };
 
 type ManualInput = {
-  variablesValueCache: Record<string, string>;
+  variablesValueCache: Record<string, string | undefined>;
   variableKeys: string[];
 };
 
@@ -194,3 +194,19 @@ export interface PlaygroundState extends PlaygroundProps {
 
   setVariableValue: (key: string, value: string) => void;
 }
+
+/**
+ * Check if the input is manual
+ */
+export const isManualInput = (input: PlaygroundInput): input is ManualInput => {
+  return "variablesValueCache" in input && "variableKeys" in input;
+};
+
+/**
+ * Check if the input is a dataset
+ */
+export const isDatasetInput = (
+  input: PlaygroundInput
+): input is DatasetInput => {
+  return "datasetId" in input;
+};