feat(playground): Extract and display variables from all message temp…

…lates as "inputs" (#4994) * feat(playground): Extract and display variables from all message templates as "inputs" * docs(playground): Add comments * fix(playground): Switch default language back to Mustache * docs(playground): Add comments and improve typing readability for template utils * feat(playground): Support text completion prompt variable substitution
Arize-ai · Oct 15, 2024 · b8cd777 · b8cd777
1 parent 0fa0c87
commit b8cd777
Show file tree

Hide file tree

Showing 10 changed files with 190 additions and 17 deletions.
diff --git a/app/src/components/templateEditor/TemplateEditor.tsx b/app/src/components/templateEditor/TemplateEditor.tsx
@@ -11,14 +11,8 @@ import { assertUnreachable } from "@phoenix/typeUtils";
 
 import { FStringTemplating } from "./language/fString";
 import { MustacheLikeTemplating } from "./language/mustacheLike";
-
-export const TemplateLanguages = {
-  FString: "f-string", // {variable}
-  Mustache: "mustache", // {{variable}}
-} as const;
-
-type TemplateLanguage =
-  (typeof TemplateLanguages)[keyof typeof TemplateLanguages];
+import { TemplateLanguages } from "./constants";
+import { TemplateLanguage } from "./types";
 
 type TemplateEditorProps = ReactCodeMirrorProps & {
   templateLanguage: TemplateLanguage;

diff --git a/app/src/components/templateEditor/constants.ts b/app/src/components/templateEditor/constants.ts
@@ -0,0 +1,15 @@
+/**
+ * Enum for the different template languages supported by the template editor
+ *
+ * - FString: `variables look like {variable}`
+ * - Mustache: `variables look like {{variable}}`
+ *
+ * @example
+ * ```tsx
+ * <TemplateEditor language={TemplateLanguages.Mustache} />
+ * ```
+ */
+export const TemplateLanguages = {
+  FString: "f-string", // {variable}
+  Mustache: "mustache", // {{variable}}
+} as const;
diff --git a/app/src/components/templateEditor/language/fString/fStringTemplating.ts b/app/src/components/templateEditor/language/fString/fStringTemplating.ts
@@ -1,7 +1,7 @@
 import { LanguageSupport, LRLanguage } from "@codemirror/language";
 import { styleTags, tags as t } from "@lezer/highlight";
 
-import { format } from "../languageUtils";
+import { extractVariables, format } from "../languageUtils";
 
 import { parser } from "./fStringTemplating.syntax.grammar";
 
@@ -65,6 +65,16 @@ export const formatFString = ({
       text.replaceAll("\\{", "{").replaceAll("{{", "{").replaceAll("}}", "}"),
   });
 
+/**
+ * Extracts the variables from an FString template
+ */
+export const extractVariablesFromFString = (text: string) => {
+  return extractVariables({
+    parser: FStringTemplatingLanguage.parser,
+    text,
+  });
+};
+
 /**
  * Creates a CodeMirror extension for the FString templating system
  */

diff --git a/app/src/components/templateEditor/language/mustacheLike/mustacheLikeTemplating.ts b/app/src/components/templateEditor/language/mustacheLike/mustacheLikeTemplating.ts
@@ -1,7 +1,7 @@
 import { LanguageSupport, LRLanguage } from "@codemirror/language";
 import { styleTags, tags as t } from "@lezer/highlight";
 
-import { format } from "../languageUtils";
+import { extractVariables, format } from "../languageUtils";
 
 import { parser } from "./mustacheLikeTemplating.syntax.grammar";
 
@@ -68,6 +68,16 @@ export const formatMustacheLike = ({
     },
   });
 
+/**
+ * Extracts the variables from a Mustache-like template
+ */
+export const extractVariablesFromMustacheLike = (text: string) => {
+  return extractVariables({
+    parser: MustacheLikeTemplatingLanguage.parser,
+    text,
+  });
+};
+
 /**
  * Creates a CodeMirror extension for the FString templating system
  */

diff --git a/app/src/components/templateEditor/templateEditorUtils.ts b/app/src/components/templateEditor/templateEditorUtils.ts
@@ -0,0 +1,53 @@
+import { assertUnreachable } from "@phoenix/typeUtils";
+
+import { extractVariablesFromFString, formatFString } from "./language/fString";
+import {
+  extractVariablesFromMustacheLike,
+  formatMustacheLike,
+} from "./language/mustacheLike";
+import { TemplateLanguages } from "./constants";
+import { TemplateLanguage } from "./types";
+
+/**
+ * A function that formats a template with the given variables
+ */
+export type FormatFn = (arg: {
+  text: string;
+  variables: Record<string, string | number | boolean | undefined>;
+}) => string;
+
+/**
+ * A function that extracts the variables from a template
+ */
+export type ExtractVariablesFn = (template: string) => string[];
+
+/**
+ * Get an object of isomorphic functions for processing templates of the given language
+ *
+ * @param templateLanguage - The language of the template to process
+ *
+ * @returns An object containing the `format` and `extractVariables` functions.
+ * These functions share the same signature despite the different underlying
+ * templating languages.
+ */
+export const getTemplateLanguageUtils = (
+  templateLanguage: TemplateLanguage
+): {
+  format: FormatFn;
+  extractVariables: ExtractVariablesFn;
+} => {
+  switch (templateLanguage) {
+    case TemplateLanguages.FString:
+      return {
+        format: formatFString,
+        extractVariables: extractVariablesFromFString,
+      };
+    case TemplateLanguages.Mustache:
+      return {
+        format: formatMustacheLike,
+        extractVariables: extractVariablesFromMustacheLike,
+      };
+    default:
+      assertUnreachable(templateLanguage);
+  }
+};
diff --git a/app/src/components/templateEditor/types.ts b/app/src/components/templateEditor/types.ts
@@ -0,0 +1,4 @@
+import { TemplateLanguages } from "./constants";
+
+export type TemplateLanguage =
+  (typeof TemplateLanguages)[keyof typeof TemplateLanguages];
diff --git a/app/src/pages/playground/Playground.tsx b/app/src/pages/playground/Playground.tsx
@@ -123,6 +123,7 @@ const playgroundInputOutputPanelContentCSS = css`
 
 function PlaygroundContent() {
   const instances = usePlaygroundContext((state) => state.instances);
+  const inputs = usePlaygroundContext((state) => state.input);
   const numInstances = instances.length;
   const isSingleInstance = numInstances === 1;
 
@@ -167,7 +168,15 @@ function PlaygroundContent() {
         <div css={playgroundInputOutputPanelContentCSS}>
           <Accordion>
             <AccordionItem title="Inputs" id="input">
-              <View padding="size-200">Inputs go here</View>
+              <View padding="size-200">
+                <pre>
+                  {JSON.stringify(
+                    "variables" in inputs ? inputs.variables : "inputs go here",
+                    null,
+                    2
+                  )}
+                </pre>
+              </View>
             </AccordionItem>
             <AccordionItem title="Output" id="output">
               <View padding="size-200" height="100%">

diff --git a/app/src/pages/playground/PlaygroundChatTemplate.tsx b/app/src/pages/playground/PlaygroundChatTemplate.tsx
@@ -19,10 +19,8 @@ import { Button, Card, Flex, Icon, Icons, View } from "@arizeai/components";
 
 import { CopyToClipboardButton } from "@phoenix/components";
 import { DragHandle } from "@phoenix/components/dnd/DragHandle";
-import {
-  TemplateEditor,
-  TemplateLanguages,
-} from "@phoenix/components/templateEditor";
+import { TemplateEditor } from "@phoenix/components/templateEditor";
+import { TemplateLanguage } from "@phoenix/components/templateEditor/types";
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
 import {
@@ -47,6 +45,9 @@ interface PlaygroundChatTemplateProps extends PlaygroundInstanceProps {}
 export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
   const id = props.playgroundInstanceId;
 
+  const templateLanguage = usePlaygroundContext(
+    (state) => state.templateLanguage
+  );
   const instances = usePlaygroundContext((state) => state.instances);
   const updateInstance = usePlaygroundContext((state) => state.updateInstance);
   const playgroundInstance = instances.find((instance) => instance.id === id);
@@ -107,6 +108,7 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
             return (
               <SortableMessageItem
                 playgroundInstanceId={id}
+                templateLanguage={templateLanguage}
                 template={template}
                 key={message.id}
                 message={message}
@@ -159,12 +161,14 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
 
 function SortableMessageItem({
   playgroundInstanceId,
+  templateLanguage,
   template,
   message,
 }: PropsWithChildren<
   PlaygroundInstanceProps & {
     template: PlaygroundChatTemplateType;
     message: ChatMessage;
+    templateLanguage: TemplateLanguage;
     index: number;
   }
 >) {
@@ -248,7 +252,7 @@ function SortableMessageItem({
             height="100%"
             value={message.content}
             aria-label="Message content"
-            templateLanguage={TemplateLanguages.Mustache}
+            templateLanguage={templateLanguage}
             onChange={(val) => {
               updateInstance({
                 instanceId: playgroundInstanceId,

diff --git a/app/src/store/playground/playgroundStore.tsx b/app/src/store/playground/playgroundStore.tsx
@@ -1,6 +1,11 @@
 import { create, StateCreator } from "zustand";
 import { devtools } from "zustand/middleware";
 
+import { TemplateLanguages } from "@phoenix/components/templateEditor/constants";
+import { getTemplateLanguageUtils } from "@phoenix/components/templateEditor/templateEditorUtils";
+import { TemplateLanguage } from "@phoenix/components/templateEditor/types";
+import { assertUnreachable } from "@phoenix/typeUtils";
+
 import {
   GenAIOperationType,
   InitialPlaygroundState,
@@ -83,7 +88,18 @@ export const createPlaygroundStore = (
   const playgroundStore: StateCreator<PlaygroundState> = (set, get) => ({
     operationType: "chat",
     inputMode: "manual",
-    input: { variables: {} },
+    input: {
+      // TODO(apowell): When implementing variable forms, we should maintain a separate
+      // map of variableName to variableValue. This will allow us to "cache" variable values
+      // as the user types and will prevent data loss if users accidentally change the variable name
+      variables: {
+        // TODO(apowell): This is hardcoded based on the default chat template
+        // Instead we should calculate this based on the template on store creation
+        // Not a huge deal since this will be overridden on the first keystroke
+        question: "",
+      },
+    },
+    templateLanguage: TemplateLanguages.Mustache,
     setInputMode: (inputMode: PlaygroundInputMode) => set({ inputMode }),
     instances: [createPlaygroundInstance()],
     setOperationType: (operationType: GenAIOperationType) => {
@@ -192,6 +208,7 @@ export const createPlaygroundStore = (
           return instance;
         }),
       });
+      get().calculateVariables();
     },
     runPlaygroundInstances: () => {
       const instances = get().instances;
@@ -232,6 +249,47 @@ export const createPlaygroundStore = (
         }),
       });
     },
+    setTemplateLanguage: (templateLanguage: TemplateLanguage) => {
+      set({ templateLanguage });
+    },
+    calculateVariables: () => {
+      const instances = get().instances;
+      const variables: Record<string, string> = {};
+      const utils = getTemplateLanguageUtils(get().templateLanguage);
+      instances.forEach((instance) => {
+        const instanceType = instance.template.__type;
+        // this double nested loop should be okay since we don't expect more than 4 instances
+        // and a handful of messages per instance
+        switch (instanceType) {
+          case "chat": {
+            // for each chat message in the instance
+            instance.template.messages.forEach((message) => {
+              // extract variables from the message content
+              const extractedVariables = utils.extractVariables(
+                message.content
+              );
+              extractedVariables.forEach((variable) => {
+                variables[variable] = "";
+              });
+            });
+            break;
+          }
+          case "text_completion": {
+            const extractedVariables = utils.extractVariables(
+              instance.template.prompt
+            );
+            extractedVariables.forEach((variable) => {
+              variables[variable] = "";
+            });
+            break;
+          }
+          default: {
+            assertUnreachable(instanceType);
+          }
+        }
+      });
+      set({ input: { variables: { ...variables } } });
+    },
     ...initialProps,
   });
   return create(devtools(playgroundStore));

diff --git a/app/src/store/playground/types.ts b/app/src/store/playground/types.ts
@@ -1,3 +1,5 @@
+import { TemplateLanguage } from "@phoenix/components/templateEditor/types";
+
 export type GenAIOperationType = "chat" | "text_completion";
 /**
  * The input mode for the playground
@@ -110,6 +112,12 @@ export interface PlaygroundProps {
    * Defaults to a single instance until a second instance is added
    */
   instances: Array<PlaygroundInstance>;
+
+  /**
+   * The current template language for all instances
+   * @default "mustache"
+   */
+  templateLanguage: TemplateLanguage;
 }
 
 export type InitialPlaygroundState = Partial<PlaygroundProps>;
@@ -163,4 +171,12 @@ export interface PlaygroundState extends PlaygroundProps {
    * Mark a given playground instance as completed
    */
   markPlaygroundInstanceComplete: (instanceId: number) => void;
+  /**
+   * Set the template language for all instances
+   */
+  setTemplateLanguage: (templateLanguage: TemplateLanguage) => void;
+  /**
+   * Calculate the variables used across all instances
+   */
+  calculateVariables: () => void;
 }