From 13c92f134c9b6863e4cb52f35182a227b647de5b Mon Sep 17 00:00:00 2001
From: Lucas Massemin <lucasmassemin@gmail.com>
Date: Thu, 5 Dec 2024 12:41:45 +0100
Subject: [PATCH] 1725 prompt versioning (#9105)

* Basic versioning (not sticking to figma yet)

* Remembering edits on latest version

* better button positioning

* Re-positioning advanced settings

* Now remembering all version edits

* New date formatter

* Fixed scroll area height -> now dynamic

---------

Co-authored-by: Lucas <lucas@dust.tt>
---
 .../assistant_builder/AssistantBuilder.tsx    |   1 +
 .../assistant_builder/InstructionScreen.tsx   | 147 +++++++++++++++++-
 front/lib/swr/assistants.ts                   |  31 ++++
 .../[aId]/history/index.ts                    |  87 +++++++++++
 .../internal/agent_configuration.ts           |   4 +
 5 files changed, 269 insertions(+), 1 deletion(-)
 create mode 100644 front/pages/api/w/[wId]/assistant/agent_configurations/[aId]/history/index.ts

diff --git a/front/components/assistant_builder/AssistantBuilder.tsx b/front/components/assistant_builder/AssistantBuilder.tsx
index 458ca0a6c9d2..bda114636a67 100644
--- a/front/components/assistant_builder/AssistantBuilder.tsx
+++ b/front/components/assistant_builder/AssistantBuilder.tsx
@@ -450,6 +450,7 @@ export default function AssistantBuilder({
                         instructionsError={instructionsError}
                         doTypewriterEffect={doTypewriterEffect}
                         setDoTypewriterEffect={setDoTypewriterEffect}
+                        agentConfigurationId={agentConfigurationId}
                       />
                     );
                   case "actions":
diff --git a/front/components/assistant_builder/InstructionScreen.tsx b/front/components/assistant_builder/InstructionScreen.tsx
index d74da80d643b..d59c1f1baefb 100644
--- a/front/components/assistant_builder/InstructionScreen.tsx
+++ b/front/components/assistant_builder/InstructionScreen.tsx
@@ -15,6 +15,7 @@ import type {
   APIError,
   AssistantCreativityLevel,
   BuilderSuggestionsType,
+  LightAgentConfigurationType,
   ModelConfigurationType,
   ModelIdType,
   PlanType,
@@ -42,7 +43,13 @@ import { History } from "@tiptap/extension-history";
 import Text from "@tiptap/extension-text";
 import type { Editor, JSONContent } from "@tiptap/react";
 import { EditorContent, useEditor } from "@tiptap/react";
-import React, { useEffect, useMemo, useRef, useState } from "react";
+import React, {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
 
 import type { AssistantBuilderState } from "@app/components/assistant_builder/types";
 import {
@@ -56,6 +63,7 @@ import {
   tipTapContentFromPlainText,
 } from "@app/lib/client/assistant_builder/instructions";
 import { isUpgraded } from "@app/lib/plans/plan_codes";
+import { useAgentConfigurationHistory } from "@app/lib/swr/assistants";
 import { classNames } from "@app/lib/utils";
 import { debounce } from "@app/lib/utils/debounce";
 
@@ -111,6 +119,7 @@ export function InstructionScreen({
   instructionsError,
   doTypewriterEffect,
   setDoTypewriterEffect,
+  agentConfigurationId,
 }: {
   owner: WorkspaceType;
   plan: PlanType;
@@ -124,6 +133,7 @@ export function InstructionScreen({
   instructionsError: string | null;
   doTypewriterEffect: boolean;
   setDoTypewriterEffect: (doTypewriterEffect: boolean) => void;
+  agentConfigurationId: string | null;
 }) {
   const editor = useEditor({
     extensions: [
@@ -153,6 +163,44 @@ export function InstructionScreen({
   });
   const editorService = useInstructionEditorService(editor);
 
+  const { agentConfigurationHistory } = useAgentConfigurationHistory({
+    workspaceId: owner.sId,
+    agentConfigurationId,
+    disabled: !agentConfigurationId,
+    limit: 30,
+  });
+  // Keep a memory of overriden versions, to not lose them when switching back and forth
+  const [currentConfig, setCurrentConfig] =
+    useState<LightAgentConfigurationType | null>(null);
+  // versionNumber -> instructions
+  const [overridenConfigInstructions, setOverridenConfigInstructions] =
+    useState<{
+      [key: string]: string;
+    }>({});
+
+  // Deduplicate configs based on instructions
+  const configsWithUniqueInstructions: LightAgentConfigurationType[] =
+    useMemo(() => {
+      const uniqueInstructions = new Set<string>();
+      const configs: LightAgentConfigurationType[] = [];
+      agentConfigurationHistory?.forEach((config) => {
+        if (
+          !config.instructions ||
+          uniqueInstructions.has(config.instructions)
+        ) {
+          return;
+        } else {
+          uniqueInstructions.add(config.instructions);
+          configs.push(config);
+        }
+      });
+      return configs;
+    }, [agentConfigurationHistory]);
+
+  useEffect(() => {
+    setCurrentConfig(agentConfigurationHistory?.[0] || null);
+  }, [agentConfigurationHistory]);
+
   const [letterIndex, setLetterIndex] = useState(0);
 
   // Beware that using this useEffect will cause a lot of re-rendering until we finished the visual effect
@@ -251,6 +299,45 @@ export function InstructionScreen({
       </div>
       <div className="flex h-full flex-col gap-1">
         <div className="relative h-full min-h-[240px] grow gap-1 p-px">
+          {configsWithUniqueInstructions &&
+            configsWithUniqueInstructions.length > 1 &&
+            currentConfig && (
+              <div className="absolute right-2 top-2 z-10">
+                <PromptHistory
+                  history={configsWithUniqueInstructions}
+                  onConfigChange={(config) => {
+                    // Remember the instructions of the version we're leaving, if overriden
+                    if (
+                      currentConfig &&
+                      currentConfig.instructions !== builderState.instructions
+                    ) {
+                      setOverridenConfigInstructions((prev) => ({
+                        ...prev,
+                        [currentConfig.version]: builderState.instructions,
+                      }));
+                    }
+
+                    // Bring new version's instructions to the editor, fetch overriden instructions if any
+                    setCurrentConfig(config);
+                    editorService.resetContent(
+                      tipTapContentFromPlainText(
+                        overridenConfigInstructions[config.version] ||
+                          config.instructions ||
+                          ""
+                      )
+                    );
+                    setBuilderState((state) => ({
+                      ...state,
+                      instructions:
+                        overridenConfigInstructions[config.version] ||
+                        config.instructions ||
+                        "",
+                    }));
+                  }}
+                  currentConfig={currentConfig}
+                />
+              </div>
+            )}
           <EditorContent
             editor={editor}
             className="absolute bottom-0 left-0 right-0 top-0"
@@ -461,6 +548,64 @@ function AdvancedSettings({
   );
 }
 
+function PromptHistory({
+  history,
+  onConfigChange,
+  currentConfig,
+}: {
+  history: LightAgentConfigurationType[];
+  onConfigChange: (config: LightAgentConfigurationType) => void;
+  currentConfig: LightAgentConfigurationType;
+}) {
+  const [isOpen, setIsOpen] = useState(false);
+  const latestConfig = history[0];
+
+  const getStringRepresentation = useCallback(
+    (config: LightAgentConfigurationType) => {
+      const dateFormatter = new Intl.DateTimeFormat(navigator.language, {
+        year: "numeric",
+        month: "short",
+        day: "numeric",
+        hour: "numeric",
+        minute: "numeric",
+      });
+      return config.version === latestConfig?.version
+        ? "Latest Version"
+        : config.versionCreatedAt
+          ? dateFormatter.format(new Date(config.versionCreatedAt))
+          : `v${config.version}`;
+    },
+    [latestConfig]
+  );
+
+  return (
+    <DropdownMenu>
+      <DropdownMenuTrigger asChild>
+        <Button
+          label={getStringRepresentation(currentConfig)}
+          variant="outline"
+          size="sm"
+          isSelect
+          onClick={() => setIsOpen(!isOpen)}
+        />
+      </DropdownMenuTrigger>
+      <DropdownMenuContent>
+        <ScrollArea className="flex max-h-[300px] flex-col">
+          {history.map((config) => (
+            <DropdownMenuItem
+              key={config.version}
+              label={getStringRepresentation(config)}
+              onClick={() => {
+                onConfigChange(config);
+              }}
+            />
+          ))}
+        </ScrollArea>
+      </DropdownMenuContent>
+    </DropdownMenu>
+  );
+}
+
 const STATIC_SUGGESTIONS = [
   "Break down your instructions into steps to leverage the model's reasoning capabilities.",
   "Give context on how you'd like the assistant to act, e.g. 'Act like a senior analyst'.",
diff --git a/front/lib/swr/assistants.ts b/front/lib/swr/assistants.ts
index 54be2c383dbf..dc80fcded342 100644
--- a/front/lib/swr/assistants.ts
+++ b/front/lib/swr/assistants.ts
@@ -242,6 +242,37 @@ export function useAgentConfiguration({
   };
 }
 
+export function useAgentConfigurationHistory({
+  workspaceId,
+  agentConfigurationId,
+  limit,
+  disabled,
+}: {
+  workspaceId: string;
+  agentConfigurationId: string | null;
+  limit?: number;
+  disabled?: boolean;
+}) {
+  const agentConfigurationHistoryFetcher: Fetcher<{
+    history: AgentConfigurationType[];
+  }> = fetcher;
+
+  const queryParams = limit ? `?limit=${limit}` : "";
+  const { data, error, mutate } = useSWRWithDefaults(
+    agentConfigurationId
+      ? `/api/w/${workspaceId}/assistant/agent_configurations/${agentConfigurationId}/history${queryParams}`
+      : null,
+    agentConfigurationHistoryFetcher,
+    { disabled }
+  );
+
+  return {
+    agentConfigurationHistory: data?.history,
+    isAgentConfigurationHistoryLoading: !error && !data,
+    isAgentConfigurationHistoryError: error,
+    mutateAgentConfigurationHistory: mutate,
+  };
+}
 export function useAgentConfigurationLastAuthor({
   workspaceId,
   agentConfigurationId,
diff --git a/front/pages/api/w/[wId]/assistant/agent_configurations/[aId]/history/index.ts b/front/pages/api/w/[wId]/assistant/agent_configurations/[aId]/history/index.ts
new file mode 100644
index 000000000000..ad961c9e1a68
--- /dev/null
+++ b/front/pages/api/w/[wId]/assistant/agent_configurations/[aId]/history/index.ts
@@ -0,0 +1,87 @@
+import type {
+  LightAgentConfigurationType,
+  WithAPIErrorResponse,
+} from "@dust-tt/types";
+import { GetAgentConfigurationsHistoryQuerySchema } from "@dust-tt/types";
+import { isLeft } from "fp-ts/lib/Either";
+import * as reporter from "io-ts-reporters";
+import type { NextApiRequest, NextApiResponse } from "next";
+
+import { getAgentConfigurations } from "@app/lib/api/assistant/configuration";
+import { withSessionAuthenticationForWorkspace } from "@app/lib/api/auth_wrappers";
+import type { Authenticator } from "@app/lib/auth";
+import { apiError } from "@app/logger/withlogging";
+
+export type GetAgentConfigurationsResponseBody = {
+  history: LightAgentConfigurationType[];
+};
+
+async function handler(
+  req: NextApiRequest,
+  res: NextApiResponse<
+    WithAPIErrorResponse<GetAgentConfigurationsResponseBody | void>
+  >,
+  auth: Authenticator
+): Promise<void> {
+  switch (req.method) {
+    case "GET":
+      // extract the limit from the query parameters
+      const queryValidation = GetAgentConfigurationsHistoryQuerySchema.decode({
+        ...req.query,
+        limit:
+          typeof req.query.limit === "string"
+            ? parseInt(req.query.limit, 10)
+            : undefined,
+      });
+      if (isLeft(queryValidation)) {
+        const pathError = reporter.formatValidationErrors(queryValidation.left);
+        return apiError(req, res, {
+          status_code: 400,
+          api_error: {
+            type: "invalid_request_error",
+            message: `Invalid query parameters: ${pathError}`,
+          },
+        });
+      }
+
+      const { limit } = queryValidation.right;
+
+      const agentConfigurations = await getAgentConfigurations({
+        auth,
+        agentsGetView: {
+          agentIds: [req.query.aId as string],
+          allVersions: true,
+        },
+        variant: "light",
+        // Return the latest versions first
+        sort: "updatedAt",
+        limit,
+      });
+
+      if (
+        !agentConfigurations ||
+        (agentConfigurations[0].scope === "private" &&
+          agentConfigurations[0].versionAuthorId !== auth.user()?.id)
+      ) {
+        return apiError(req, res, {
+          status_code: 404,
+          api_error: {
+            type: "agent_configuration_not_found",
+            message: "The Assistant you're trying to access was not found.",
+          },
+        });
+      }
+
+      return res.status(200).json({ history: agentConfigurations });
+    default:
+      return apiError(req, res, {
+        status_code: 405,
+        api_error: {
+          type: "method_not_supported_error",
+          message: "The method passed is not supported, GET is expected.",
+        },
+      });
+  }
+}
+
+export default withSessionAuthenticationForWorkspace(handler);
diff --git a/types/src/front/api_handlers/internal/agent_configuration.ts b/types/src/front/api_handlers/internal/agent_configuration.ts
index b4a867b566b8..239affa64373 100644
--- a/types/src/front/api_handlers/internal/agent_configuration.ts
+++ b/types/src/front/api_handlers/internal/agent_configuration.ts
@@ -33,6 +33,10 @@ export const GetAgentConfigurationsQuerySchema = t.type({
   ]),
 });
 
+export const GetAgentConfigurationsHistoryQuerySchema = t.type({
+  limit: t.union([LimitCodec, t.undefined]),
+});
+
 export const GetAgentConfigurationsLeaderboardQuerySchema = t.type({
   view: t.union([
     t.literal("list"),