From b8cd7772e59455b8b3c691609f4baee7a6157a24 Mon Sep 17 00:00:00 2001 From: Anthony Powell Date: Mon, 14 Oct 2024 22:15:53 -0400 Subject: [PATCH] feat(playground): Extract and display variables from all message templates 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 --- .../templateEditor/TemplateEditor.tsx | 10 +--- .../components/templateEditor/constants.ts | 15 +++++ .../language/fString/fStringTemplating.ts | 12 +++- .../mustacheLike/mustacheLikeTemplating.ts | 12 +++- .../templateEditor/templateEditorUtils.ts | 53 ++++++++++++++++ app/src/components/templateEditor/types.ts | 4 ++ app/src/pages/playground/Playground.tsx | 11 +++- .../playground/PlaygroundChatTemplate.tsx | 14 +++-- app/src/store/playground/playgroundStore.tsx | 60 ++++++++++++++++++- app/src/store/playground/types.ts | 16 +++++ 10 files changed, 190 insertions(+), 17 deletions(-) create mode 100644 app/src/components/templateEditor/constants.ts create mode 100644 app/src/components/templateEditor/templateEditorUtils.ts create mode 100644 app/src/components/templateEditor/types.ts diff --git a/app/src/components/templateEditor/TemplateEditor.tsx b/app/src/components/templateEditor/TemplateEditor.tsx index dfc36beb97..2628940809 100644 --- 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 new file mode 100644 index 0000000000..803266784d --- /dev/null +++ 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 + * + * ``` + */ +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 index d51510e43d..ec11fe8236 100644 --- 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 index 6d472c6c2d..b1f9b3015a 100644 --- 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 new file mode 100644 index 0000000000..1dd1510ce7 --- /dev/null +++ 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; + +/** + * 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 new file mode 100644 index 0000000000..ff45d4dc5e --- /dev/null +++ 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 index e15fe4f220..6c5155e183 100644 --- 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() {
- Inputs go here + + 
+                  {JSON.stringify(
+                    "variables" in inputs ? inputs.variables : "inputs go here",
+                    null,
+                    2
+                  )}
+
+ diff --git a/app/src/pages/playground/PlaygroundChatTemplate.tsx b/app/src/pages/playground/PlaygroundChatTemplate.tsx index c99b7c0c60..d0361d0a1f 100644 --- 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 ( ) { @@ -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 index 7011dc41c0..11f6e21502 100644 --- 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 = (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 = {}; + 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 index 1b36d2f5be..892da8cc45 100644 --- 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; + + /** + * The current template language for all instances + * @default "mustache" + */ + templateLanguage: TemplateLanguage; } export type InitialPlaygroundState = Partial; @@ -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; }