[editor] Run Prompt (non-streaming) (#667)

[editor] Run Prompt (non-streaming) # [editor] Run Prompt (non-streaming) With non-streaming /run endpoint properly returning the response now, we can update the client config with the resulting output. We can also leverage the client-side `_ui` state for maintaining 'isRunning' state to show a spinner / disable the execute button for now (TODO: we may need to handle canceling execution as well). https://github.com/lastmile-ai/aiconfig/assets/5060851/67fa793a-9a60-4285-a7d1-69fff633d6a6 --- Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/lastmile-ai/aiconfig/pull/667). * __->__ #667 * #666 * #665 * #662
lastmile-ai · Dec 29, 2023 · 164f0b0 · 164f0b0
2 parents 8c35024 + 4a1b8ab
commit 164f0b0
Show file tree

Hide file tree

Showing 6 changed files with 150 additions and 74 deletions.
diff --git a/python/src/aiconfig/editor/client/src/components/EditorContainer.tsx b/python/src/aiconfig/editor/client/src/components/EditorContainer.tsx
@@ -38,7 +38,7 @@ export type AIConfigCallbacks = {
   ) => Promise<{ aiconfig: AIConfig }>;
   deletePrompt: (promptName: string) => Promise<void>;
   getModels: (search: string) => Promise<string[]>;
-  runPrompt: (promptName: string) => Promise<void>;
+  runPrompt: (promptName: string) => Promise<{ aiconfig: AIConfig }>;
   save: (aiconfig: AIConfig) => Promise<void>;
   updateModel: (
     promptName?: string,
@@ -96,7 +96,7 @@ export default function EditorContainer({
   const onSave = useCallback(async () => {
     setIsSaving(true);
     try {
-      await callbacks.save(clientConfigToAIConfig(aiconfigState));
+      await callbacks.save(clientConfigToAIConfig(stateRef.current));
     } catch (err: any) {
       showNotification({
         title: "Error saving",
@@ -106,7 +106,7 @@ export default function EditorContainer({
     } finally {
       setIsSaving(false);
     }
-  }, [aiconfigState, callbacks.save]);
+  }, [callbacks.save]);
 
   const debouncedUpdatePrompt = useMemo(
     () =>
@@ -119,18 +119,18 @@ export default function EditorContainer({
   );
 
   const onChangePromptInput = useCallback(
-    async (promptIndex: number, newPromptInput: PromptInput) => {
+    async (promptId: string, newPromptInput: PromptInput) => {
       const action: AIConfigReducerAction = {
         type: "UPDATE_PROMPT_INPUT",
-        index: promptIndex,
+        id: promptId,
         input: newPromptInput,
       };
 
       dispatch(action);
 
       try {
         const prompt = clientPromptToAIConfigPrompt(
-          aiconfigState.prompts[promptIndex]
+          getPrompt(stateRef.current, promptId)!
         );
         const serverConfigRes = await debouncedUpdatePrompt(prompt.name, {
           ...prompt,
@@ -154,10 +154,10 @@ export default function EditorContainer({
   );
 
   const onChangePromptName = useCallback(
-    async (promptIndex: number, newName: string) => {
+    async (promptId: string, newName: string) => {
       const action: AIConfigReducerAction = {
         type: "UPDATE_PROMPT_NAME",
-        index: promptIndex,
+        id: promptId,
         name: newName,
       };
 
@@ -177,10 +177,10 @@ export default function EditorContainer({
   );
 
   const onUpdatePromptModelSettings = useCallback(
-    async (promptIndex: number, newModelSettings: any) => {
+    async (promptId: string, newModelSettings: any) => {
       dispatch({
         type: "UPDATE_PROMPT_MODEL_SETTINGS",
-        index: promptIndex,
+        id: promptId,
         modelSettings: newModelSettings,
       });
       // TODO: Call server-side endpoint to update model
@@ -189,16 +189,16 @@ export default function EditorContainer({
   );
 
   const onUpdatePromptModel = useCallback(
-    async (promptIndex: number, newModel?: string) => {
+    async (promptId: string, newModel?: string) => {
       dispatch({
         type: "UPDATE_PROMPT_MODEL",
-        index: promptIndex,
+        id: promptId,
         modelName: newModel,
       });
 
       try {
         const prompt = clientPromptToAIConfigPrompt(
-          aiconfigState.prompts[promptIndex]
+          getPrompt(stateRef.current, promptId)!
         );
         const currentModel = prompt.metadata?.model;
         let modelData: string | ModelMetadata | undefined = newModel;
@@ -224,10 +224,10 @@ export default function EditorContainer({
   );
 
   const onUpdatePromptParameters = useCallback(
-    async (promptIndex: number, newParameters: any) => {
+    async (promptId: string, newParameters: any) => {
       dispatch({
         type: "UPDATE_PROMPT_PARAMETERS",
-        index: promptIndex,
+        id: promptId,
         parameters: newParameters,
       });
       // TODO: Call server-side endpoint to update prompt parameters
@@ -306,10 +306,23 @@ export default function EditorContainer({
   );
 
   const onRunPrompt = useCallback(
-    async (promptIndex: number) => {
-      const promptName = aiconfigState.prompts[promptIndex].name;
+    async (promptId: string) => {
+      const action: AIConfigReducerAction = {
+        type: "RUN_PROMPT",
+        id: promptId,
+      };
+
+      dispatch(action);
+
       try {
-        await callbacks.runPrompt(promptName);
+        const promptName = getPrompt(stateRef.current, promptId)!.name;
+        const serverConfigRes = await callbacks.runPrompt(promptName);
+
+        dispatch({
+          type: "CONSOLIDATE_AICONFIG",
+          action,
+          config: serverConfigRes.aiconfig,
+        });
       } catch (err: any) {
         showNotification({
           title: "Error running prompt",
@@ -345,7 +358,6 @@ export default function EditorContainer({
                   onDeletePrompt={() => onDeletePrompt(prompt._ui.id)}
                 />
                 <PromptContainer
-                  index={i}
                   prompt={prompt}
                   getModels={callbacks.getModels}
                   onChangePromptInput={onChangePromptInput}

diff --git a/python/src/aiconfig/editor/client/src/components/aiconfigReducer.ts b/python/src/aiconfig/editor/client/src/components/aiconfigReducer.ts
@@ -9,6 +9,7 @@ export type AIConfigReducerAction =
 export type MutateAIConfigAction =
   | AddPromptAction
   | DeletePromptAction
+  | RunPromptAction
   | UpdatePromptInputAction
   | UpdatePromptNameAction
   | UpdatePromptModelAction
@@ -32,57 +33,61 @@ export type DeletePromptAction = {
   id: string;
 };
 
-// TODO: Update index to prompt id for all existing-prompt actions
+export type RunPromptAction = {
+  type: "RUN_PROMPT";
+  id: string;
+};
+
 export type UpdatePromptInputAction = {
   type: "UPDATE_PROMPT_INPUT";
-  index: number;
+  id: string;
   input: PromptInput;
 };
 
 export type UpdatePromptNameAction = {
   type: "UPDATE_PROMPT_NAME";
-  index: number;
+  id: string;
   name: string;
 };
 
 export type UpdatePromptModelAction = {
   type: "UPDATE_PROMPT_MODEL";
-  index: number;
+  id: string;
   modelName?: string;
 };
 
 export type UpdatePromptModelSettingsAction = {
   type: "UPDATE_PROMPT_MODEL_SETTINGS";
-  index: number;
+  id: string;
   modelSettings: JSONObject;
 };
 
 // TODO: saqadri - can likely use this same action for global parameters update
 export type UpdatePromptParametersAction = {
   type: "UPDATE_PROMPT_PARAMETERS";
-  index: number;
+  id: string;
   parameters: JSONObject;
 };
 
 function reduceReplacePrompt(
   state: ClientAIConfig,
-  index: number,
+  id: string,
   replacerFn: (prompt: ClientPrompt) => ClientPrompt
 ): ClientAIConfig {
   return {
     ...state,
-    prompts: state.prompts.map((prompt, i) =>
-      i === index ? replacerFn(prompt) : prompt
+    prompts: state.prompts.map((prompt) =>
+      prompt._ui.id === id ? replacerFn(prompt) : prompt
     ),
   };
 }
 
 function reduceReplaceInput(
   state: ClientAIConfig,
-  index: number,
+  id: string,
   replacerFn: (input: PromptInput) => PromptInput
 ): ClientAIConfig {
-  return reduceReplacePrompt(state, index, (prompt) => ({
+  return reduceReplacePrompt(state, id, (prompt) => ({
     ...prompt,
     input: replacerFn(prompt.input),
   }));
@@ -108,23 +113,49 @@ function reduceConsolidateAIConfig(
   action: MutateAIConfigAction,
   responseConfig: AIConfig
 ): ClientAIConfig {
+  // Make sure prompt structure is properly updated. Client input and metadata takes precedence
+  // since it may have been updated by the user while the request was in flight
+  const consolidatePrompt = (statePrompt: ClientPrompt) => {
+    const responsePrompt = responseConfig.prompts.find(
+      (resPrompt) => resPrompt.name === statePrompt.name
+    );
+    return {
+      ...responsePrompt,
+      ...statePrompt,
+      metadata: {
+        ...responsePrompt!.metadata,
+        ...statePrompt.metadata,
+      },
+    } as ClientPrompt;
+  };
+
   switch (action.type) {
-    case "ADD_PROMPT_AT_INDEX":
-    case "UPDATE_PROMPT_INPUT": {
-      // Make sure prompt structure is properly updated. Client input and metadata takes precedence
-      // since it may have been updated by the user while the request was in flight
-      return reduceReplacePrompt(state, action.index, (prompt) => {
-        const responsePrompt = responseConfig.prompts[action.index];
+    case "ADD_PROMPT_AT_INDEX": {
+      return reduceReplacePrompt(
+        state,
+        action.prompt._ui.id,
+        consolidatePrompt
+      );
+    }
+    case "RUN_PROMPT": {
+      return reduceReplacePrompt(state, action.id, (prompt) => {
+        const responsePrompt = responseConfig.prompts.find(
+          (resPrompt) => resPrompt.name === prompt.name
+        );
+
         return {
-          ...responsePrompt,
           ...prompt,
-          metadata: {
-            ...responsePrompt.metadata,
-            ...prompt.metadata,
+          _ui: {
+            ...prompt._ui,
+            isRunning: false,
           },
-        } as ClientPrompt;
+          outputs: responsePrompt!.outputs,
+        };
       });
     }
+    case "UPDATE_PROMPT_INPUT": {
+      return reduceReplacePrompt(state, action.id, consolidatePrompt);
+    }
     default: {
       return state;
     }
@@ -145,17 +176,26 @@ export default function aiconfigReducer(
         prompts: state.prompts.filter((prompt) => prompt._ui.id !== action.id),
       };
     }
+    case "RUN_PROMPT": {
+      return reduceReplacePrompt(state, action.id, (prompt) => ({
+        ...prompt,
+        _ui: {
+          ...prompt._ui,
+          isRunning: true,
+        },
+      }));
+    }
     case "UPDATE_PROMPT_INPUT": {
-      return reduceReplaceInput(state, action.index, () => action.input);
+      return reduceReplaceInput(state, action.id, () => action.input);
     }
     case "UPDATE_PROMPT_NAME": {
-      return reduceReplacePrompt(state, action.index, (prompt) => ({
+      return reduceReplacePrompt(state, action.id, (prompt) => ({
         ...prompt,
         name: action.name,
       }));
     }
     case "UPDATE_PROMPT_MODEL": {
-      return reduceReplacePrompt(state, action.index, (prompt) => ({
+      return reduceReplacePrompt(state, action.id, (prompt) => ({
         ...prompt,
         metadata: {
           ...prompt.metadata,
@@ -169,7 +209,7 @@ export default function aiconfigReducer(
       }));
     }
     case "UPDATE_PROMPT_MODEL_SETTINGS": {
-      return reduceReplacePrompt(state, action.index, (prompt) => ({
+      return reduceReplacePrompt(state, action.id, (prompt) => ({
         ...prompt,
         metadata: {
           ...prompt.metadata,
@@ -186,7 +226,7 @@ export default function aiconfigReducer(
       }));
     }
     case "UPDATE_PROMPT_PARAMETERS": {
-      return reduceReplacePrompt(state, action.index, (prompt) => ({
+      return reduceReplacePrompt(state, action.id, (prompt) => ({
         ...prompt,
         metadata: {
           ...prompt.metadata,

diff --git a/python/src/aiconfig/editor/client/src/components/prompt/PromptActionBar.tsx b/python/src/aiconfig/editor/client/src/components/prompt/PromptActionBar.tsx
@@ -99,7 +99,11 @@ export default memo(function PromptActionBar({
               <IconClearAll />
             </ActionIcon>
           </Flex>
-          <RunPromptButton runPrompt={onRunPrompt} size="compact" />
+          <RunPromptButton
+            runPrompt={onRunPrompt}
+            isRunning={prompt._ui.isRunning}
+            size="compact"
+          />
         </Flex>
       )}
     </Flex>