Set contents of editor for refinement modals instead of saving on server (#199)

* Set contents of editor for refinement modals instead of saving on server * Show New Experiment text on mobile
2023-08-27 19:52:25 -06:00
parent 28713fb3ef
commit fa87887e91
7 changed files with 105 additions and 85 deletions
--- a/app/src/components/ChangeModelModal/ChangeModelModal.tsx
+++ b/app/src/components/ChangeModelModal/ChangeModelModal.tsx
@@ -1,3 +1,4 @@
+import { useState, useMemo, useCallback } from "react";
 import {
  Button,
  HStack,
@@ -14,16 +15,18 @@ import {
  VStack,
 } from "@chakra-ui/react";
 import { type PromptVariant } from "@prisma/client";
-import { isObject, isString } from "lodash-es";
-import { useState } from "react";
+import { isString } from "lodash-es";
 import { RiExchangeFundsFill } from "react-icons/ri";
+
 import { type ProviderModel } from "~/modelProviders/types";
 import { api } from "~/utils/api";
-import { useExperiment, useHandledAsyncCallback, useVisibleScenarioIds } from "~/utils/hooks";
+import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
 import { lookupModel, modelLabel } from "~/utils/utils";
 import CompareFunctions from "../RefinePromptModal/CompareFunctions";
 import { ModelSearch } from "./ModelSearch";
 import { ModelStatsCard } from "./ModelStatsCard";
+import { maybeReportError } from "~/utils/errorHandling/maybeReportError";
+import { useAppStore } from "~/state/store";

 export const ChangeModelModal = ({
  variant,
@@ -32,48 +35,43 @@ export const ChangeModelModal = ({
  variant: PromptVariant;
  onClose: () => void;
 }) => {
+  const editorOptionsMap = useAppStore((s) => s.sharedVariantEditor.editorOptionsMap);
+  const originalPromptFn = useMemo(
+    () => editorOptionsMap[variant.uiId]?.getContent() || "",
+    [editorOptionsMap, variant.uiId],
+  );
+
  const originalModel = lookupModel(variant.modelProvider, variant.model);
  const [selectedModel, setSelectedModel] = useState({
    provider: variant.modelProvider,
    model: variant.model,
  } as ProviderModel);
  const [convertedModel, setConvertedModel] = useState<ProviderModel | undefined>();
-  const visibleScenarios = useVisibleScenarioIds();
-
-  const utils = api.useContext();
+  const [modifiedPromptFn, setModifiedPromptFn] = useState<string>();

  const experiment = useExperiment();

-  const { mutateAsync: getModifiedPromptMutateAsync, data: modifiedPromptFn } =
+  const { mutateAsync: getModifiedPromptMutateAsync } =
    api.promptVariants.getModifiedPromptFn.useMutation();

  const [getModifiedPromptFn, modificationInProgress] = useHandledAsyncCallback(async () => {
    if (!experiment) return;

-    await getModifiedPromptMutateAsync({
+    const resp = await getModifiedPromptMutateAsync({
      id: variant.id,
+      originalPromptFn,
      newModel: selectedModel,
    });
+    if (maybeReportError(resp)) return;
+    setModifiedPromptFn(resp.payload);
    setConvertedModel(selectedModel);
  }, [getModifiedPromptMutateAsync, onClose, experiment, variant, selectedModel]);

-  const replaceVariantMutation = api.promptVariants.replaceVariant.useMutation();
-
-  const [replaceVariant, replacementInProgress] = useHandledAsyncCallback(async () => {
-    if (
-      !variant.experimentId ||
-      !modifiedPromptFn ||
-      (isObject(modifiedPromptFn) && "status" in modifiedPromptFn)
-    )
-      return;
-    await replaceVariantMutation.mutateAsync({
-      id: variant.id,
-      promptConstructor: modifiedPromptFn,
-      streamScenarios: visibleScenarios,
-    });
-    await utils.promptVariants.list.invalidate();
+  const replaceVariant = useCallback(() => {
+    if (!modifiedPromptFn) return;
+    editorOptionsMap[variant.uiId]?.setContent(modifiedPromptFn);
    onClose();
-  }, [replaceVariantMutation, variant, onClose, modifiedPromptFn]);
+  }, [variant.uiId, editorOptionsMap, onClose, modifiedPromptFn]);

  const originalLabel = modelLabel(variant.modelProvider, variant.model);
  const selectedLabel = modelLabel(selectedModel.provider, selectedModel.model);
@@ -130,9 +128,9 @@ export const ChangeModelModal = ({
              colorScheme="blue"
              onClick={replaceVariant}
              minW={24}
-              isDisabled={!convertedModel || modificationInProgress || replacementInProgress}
+              isDisabled={!convertedModel || modificationInProgress}
            >
-              {replacementInProgress ? <Spinner boxSize={4} /> : <Text>Accept</Text>}
+              Accept
            </Button>
          </HStack>
        </ModalFooter>
--- a/app/src/components/OutputsTable/VariantEditor.tsx
+++ b/app/src/components/OutputsTable/VariantEditor.tsx
@@ -10,7 +10,7 @@ import {
 } from "@chakra-ui/react";
 import { useCallback, useEffect, useRef, useState } from "react";
 import { FiMaximize, FiMinimize } from "react-icons/fi";
-import { editorBackground } from "~/state/sharedVariantEditor.slice";
+import { type CreatedEditor, editorBackground } from "~/state/sharedVariantEditor.slice";
 import { useAppStore } from "~/state/store";
 import { api } from "~/utils/api";
 import {
@@ -24,8 +24,10 @@ import { type PromptVariant } from "./types";
 export default function VariantEditor(props: { variant: PromptVariant }) {
  const { canModify } = useExperimentAccess();
  const monaco = useAppStore.use.sharedVariantEditor.monaco();
-  const editorRef = useRef<ReturnType<NonNullable<typeof monaco>["editor"]["create"]> | null>(null);
+  const updateOptionsForEditor = useAppStore.use.sharedVariantEditor.updateOptionsForEditor();
+  const editorRef = useRef<CreatedEditor | null>(null);
  const containerRef = useRef<HTMLDivElement | null>(null);
+  const lastSavedFnRef = useRef(props.variant.promptConstructor);
  const [editorId] = useState(() => `editor_${Math.random().toString(36).substring(7)}`);
  const [isChanged, setIsChanged] = useState(false);

@@ -48,22 +50,18 @@ export default function VariantEditor(props: { variant: PromptVariant }) {
  }, [isFullscreen, toggleFullscreen]);

  const lastSavedFn = props.variant.promptConstructor;
+  useEffect(() => {
+    // Store in ref so that we can access it dynamically
+    lastSavedFnRef.current = lastSavedFn;
+  }, [lastSavedFn]);

  const modifierKey = useModifierKeyLabel();

  const checkForChanges = useCallback(() => {
    if (!editorRef.current) return;
    const currentFn = editorRef.current.getValue();
-    setIsChanged(currentFn.length > 0 && currentFn !== lastSavedFn);
-  }, [lastSavedFn]);
-
-  const matchUpdatedSavedFn = useCallback(() => {
-    if (!editorRef.current) return;
-    editorRef.current.setValue(lastSavedFn);
-    setIsChanged(false);
-  }, [lastSavedFn]);
-
-  useEffect(matchUpdatedSavedFn, [matchUpdatedSavedFn, lastSavedFn]);
+    setIsChanged(currentFn.length > 0 && currentFn !== lastSavedFnRef.current);
+  }, [editorRef]);

  const replaceVariant = api.promptVariants.replaceVariant.useMutation();
  const utils = api.useContext();
@@ -136,6 +134,11 @@ export default function VariantEditor(props: { variant: PromptVariant }) {
        readOnly: !canModify,
      });

+      updateOptionsForEditor(props.variant.uiId, {
+        getContent: () => editorRef.current?.getValue() || "",
+        setContent: (content) => editorRef.current?.setValue(content),
+      });
+
      // Workaround because otherwise the commands only work on whatever
      // editor was loaded on the page last.
      // https://github.com/microsoft/monaco-editor/issues/2947#issuecomment-1422265201
@@ -155,7 +158,7 @@ export default function VariantEditor(props: { variant: PromptVariant }) {
        });
      });

-      editorRef.current.onDidChangeModelContent(checkForChanges);
+      const checkForChangesListener = editorRef.current.onDidChangeModelContent(checkForChanges);

      const resizeObserver = new ResizeObserver(() => {
        editorRef.current?.layout();
@@ -164,6 +167,7 @@ export default function VariantEditor(props: { variant: PromptVariant }) {

      return () => {
        resizeObserver.disconnect();
+        checkForChangesListener.dispose();
        editorRef.current?.dispose();
      };
    }
@@ -171,7 +175,7 @@ export default function VariantEditor(props: { variant: PromptVariant }) {
    // We intentionally skip the onSave and props.savedConfig dependencies here because
    // we don't want to re-render the editor from scratch
    /* eslint-disable-next-line react-hooks/exhaustive-deps */
-  }, [monaco, editorId]);
+  }, [monaco, editorId, updateOptionsForEditor]);

  useEffect(() => {
    if (!editorRef.current) return;
--- a/app/src/components/RefinePromptModal/RefinePromptModal.tsx
+++ b/app/src/components/RefinePromptModal/RefinePromptModal.tsx
@@ -1,3 +1,4 @@
+import { useState, useMemo, useCallback } from "react";
 import {
  Button,
  Modal,
@@ -9,22 +10,23 @@ import {
  ModalOverlay,
  VStack,
  Text,
-  Spinner,
  HStack,
  Icon,
  SimpleGrid,
 } from "@chakra-ui/react";
 import { BsStars } from "react-icons/bs";
 import { api } from "~/utils/api";
-import { useHandledAsyncCallback, useVisibleScenarioIds } from "~/utils/hooks";
+import { useHandledAsyncCallback } from "~/utils/hooks";
 import { type PromptVariant } from "@prisma/client";
-import { useState } from "react";
+
 import CompareFunctions from "./CompareFunctions";
 import { CustomInstructionsInput } from "../CustomInstructionsInput";
 import { RefineAction } from "./RefineAction";
-import { isObject, isString } from "lodash-es";
+import { isString } from "lodash-es";
 import { type RefinementAction, type SupportedProvider } from "~/modelProviders/types";
 import frontendModelProviders from "~/modelProviders/frontendModelProviders";
+import { useAppStore } from "~/state/store";
+import { maybeReportError } from "~/utils/errorHandling/maybeReportError";

 export const RefinePromptModal = ({
  variant,
@@ -33,19 +35,23 @@ export const RefinePromptModal = ({
  variant: PromptVariant;
  onClose: () => void;
 }) => {
-  const utils = api.useContext();
-  const visibleScenarios = useVisibleScenarioIds();
+  const editorOptionsMap = useAppStore((s) => s.sharedVariantEditor.editorOptionsMap);
+  const originalPromptFn = useMemo(
+    () => editorOptionsMap[variant.uiId]?.getContent() || "",
+    [editorOptionsMap, variant.uiId],
+  );

  const refinementActions =
    frontendModelProviders[variant.modelProvider as SupportedProvider].refinementActions || {};

-  const { mutateAsync: getModifiedPromptMutateAsync, data: refinedPromptFn } =
+  const { mutateAsync: getModifiedPromptMutateAsync } =
    api.promptVariants.getModifiedPromptFn.useMutation();
  const [instructions, setInstructions] = useState<string>("");

  const [activeRefineActionLabel, setActiveRefineActionLabel] = useState<string | undefined>(
    undefined,
  );
+  const [refinedPromptFn, setRefinedPromptFn] = useState<string>();

  const [getModifiedPromptFn, modificationInProgress] = useHandledAsyncCallback(
    async (label?: string) => {
@@ -54,31 +60,22 @@ export const RefinePromptModal = ({
        ? (refinementActions[label] as RefinementAction).instructions
        : instructions;
      setActiveRefineActionLabel(label);
-      await getModifiedPromptMutateAsync({
+      const resp = await getModifiedPromptMutateAsync({
        id: variant.id,
+        originalPromptFn,
        instructions: updatedInstructions,
      });
+      if (maybeReportError(resp)) return;
+      setRefinedPromptFn(resp.payload);
    },
    [getModifiedPromptMutateAsync, onClose, variant, instructions, setActiveRefineActionLabel],
  );

-  const replaceVariantMutation = api.promptVariants.replaceVariant.useMutation();
-
-  const [replaceVariant, replacementInProgress] = useHandledAsyncCallback(async () => {
-    if (
-      !variant.experimentId ||
-      !refinedPromptFn ||
-      (isObject(refinedPromptFn) && "status" in refinedPromptFn)
-    )
-      return;
-    await replaceVariantMutation.mutateAsync({
-      id: variant.id,
-      promptConstructor: refinedPromptFn,
-      streamScenarios: visibleScenarios,
-    });
-    await utils.promptVariants.list.invalidate();
+  const replaceVariant = useCallback(() => {
+    if (!refinedPromptFn) return;
+    editorOptionsMap[variant.uiId]?.setContent(refinedPromptFn);
    onClose();
-  }, [replaceVariantMutation, variant, onClose, refinedPromptFn]);
+  }, [variant.uiId, editorOptionsMap, onClose, refinedPromptFn]);

  return (
    <Modal
@@ -126,7 +123,7 @@ export const RefinePromptModal = ({
              />
            </VStack>
            <CompareFunctions
-              originalFunction={variant.promptConstructor}
+              originalFunction={originalPromptFn}
              newFunction={isString(refinedPromptFn) ? refinedPromptFn : undefined}
              maxH="40vh"
            />
@@ -139,9 +136,9 @@ export const RefinePromptModal = ({
              colorScheme="blue"
              onClick={replaceVariant}
              minW={24}
-              isDisabled={replacementInProgress || !refinedPromptFn}
+              isDisabled={!refinedPromptFn}
            >
-              {replacementInProgress ? <Spinner boxSize={4} /> : <Text>Accept</Text>}
+              Accept
            </Button>
          </HStack>
        </ModalFooter>
--- a/app/src/components/experiments/ExperimentCard.tsx
+++ b/app/src/components/experiments/ExperimentCard.tsx
@@ -98,9 +98,7 @@ export const NewExperimentCard = () => {
    >
      <VStack align="center" justify="center" w="full" h="full" p={4} onClick={createExperiment}>
        <Icon as={isLoading ? Spinner : BsPlusSquare} boxSize={8} />
-        <Text display={{ base: "none", md: "block" }} ml={2}>
-          New Experiment
-        </Text>
+        <Text ml={2}>New Experiment</Text>
      </VStack>
    </Card>
  );
--- a/app/src/server/api/routers/promptVariants.router.ts
+++ b/app/src/server/api/routers/promptVariants.router.ts
@@ -298,6 +298,7 @@ export const promptVariantsRouter = createTRPCRouter({
    .input(
      z.object({
        id: z.string(),
+        originalPromptFn: z.string(),
        instructions: z.string().optional(),
        newModel: z
          .object({
@@ -315,22 +316,21 @@ export const promptVariantsRouter = createTRPCRouter({
      });
      await requireCanModifyExperiment(existing.experimentId, ctx);

-      const constructedPrompt = await parsePromptConstructor(existing.promptConstructor);
-
-      if ("error" in constructedPrompt) {
-        return error(constructedPrompt.error);
-      }
-
      const model = input.newModel
        ? modelProviders[input.newModel.provider].models[input.newModel.model]
        : undefined;

-      const promptConstructionFn = await deriveNewConstructFn(existing, model, input.instructions);
+      const promptConstructionFn = await deriveNewConstructFn(
+        existing,
+        input.originalPromptFn,
+        model,
+        input.instructions,
+      );

      // TODO: Validate promptConstructionFn
      // TODO: Record in some sort of history

-      return promptConstructionFn;
+      return success(promptConstructionFn);
    }),

  replaceVariant: protectedProcedure
--- a/app/src/server/utils/deriveNewContructFn.ts
+++ b/app/src/server/utils/deriveNewContructFn.ts
@@ -12,14 +12,20 @@ const isolate = new ivm.Isolate({ memoryLimit: 128 });

 export async function deriveNewConstructFn(
  originalVariant: PromptVariant | null,
+  originalPromptFn?: string,
  newModel?: Model,
  instructions?: string,
 ) {
-  if (originalVariant && !newModel && !instructions) {
-    return originalVariant.promptConstructor;
+  if (originalPromptFn && !newModel && !instructions) {
+    return originalPromptFn;
  }
-  if (originalVariant && (newModel || instructions)) {
-    return await requestUpdatedPromptFunction(originalVariant, newModel, instructions);
+  if (originalVariant && originalPromptFn && (newModel || instructions)) {
+    return await requestUpdatedPromptFunction(
+      originalVariant,
+      originalPromptFn,
+      newModel,
+      instructions,
+    );
  }
  return dedent`
    prompt = {
@@ -36,6 +42,7 @@ export async function deriveNewConstructFn(
 const NUM_RETRIES = 5;
 const requestUpdatedPromptFunction = async (
  originalVariant: PromptVariant,
+  originalPromptFn: string,
  newModel?: Model,
  instructions?: string,
 ) => {
@@ -55,7 +62,7 @@ const requestUpdatedPromptFunction = async (
        },
        {
          role: "user",
-          content: `This is the current prompt constructor function:\n---\n${originalVariant.promptConstructor}`,
+          content: `This is the current prompt constructor function:\n---\n${originalPromptFn}`,
        },
      ];
      if (newModel) {
--- a/app/src/state/sharedVariantEditor.slice.ts
+++ b/app/src/state/sharedVariantEditor.slice.ts
@@ -1,16 +1,26 @@
+import loader, { type Monaco } from "@monaco-editor/loader";
+
 import { type RouterOutputs } from "~/utils/api";
 import { type SliceCreator } from "./store";
-import loader from "@monaco-editor/loader";
 import formatPromptConstructor from "~/promptConstructor/format";

 export const editorBackground = "#fafafa";

+export type CreatedEditor = ReturnType<Monaco["editor"]["create"]>;
+
+type EditorOptions = {
+  getContent: () => string;
+  setContent: (content: string) => void;
+};
+
 export type SharedVariantEditorSlice = {
-  monaco: null | ReturnType<typeof loader.__getMonacoInstance>;
+  monaco: null | Monaco;
  loadMonaco: () => Promise<void>;
  scenarioVars: RouterOutputs["scenarioVars"]["list"];
  updateScenariosModel: () => void;
  setScenarioVars: (scenarioVars: RouterOutputs["scenarioVars"]["list"]) => void;
+  editorOptionsMap: Record<string, EditorOptions>;
+  updateOptionsForEditor: (uiId: string, { getContent, setContent }: EditorOptions) => void;
 };

 export const createVariantEditorSlice: SliceCreator<SharedVariantEditorSlice> = (set, get) => ({
@@ -93,4 +103,10 @@ export const createVariantEditorSlice: SliceCreator<SharedVariantEditorSlice> =
      );
    }
  },
+  editorOptionsMap: {},
+  updateOptionsForEditor: (uiId, options) => {
+    set((state) => {
+      state.sharedVariantEditor.editorOptionsMap[uiId] = options;
+    });
+  },
 });