feat: include data action

dust-tt · Nov 8, 2024 · 71ee64f · 71ee64f
1 parent 8add322
commit 71ee64f
Show file tree

Hide file tree

Showing 9 changed files with 107 additions and 70 deletions.
diff --git a/front/components/assistant_builder/actions/ProcessAction.tsx b/front/components/assistant_builder/actions/ProcessAction.tsx
@@ -567,7 +567,7 @@ export function ActionProcess({
           }}
         />
         <TimeUnitDropdown
-          actionConfiguration={actionConfiguration}
+          timeFrame={actionConfiguration.timeFrame}
           updateAction={updateAction}
           onEdit={() => setEdited(true)}
         />

diff --git a/front/components/assistant_builder/actions/RetrievalAction.tsx b/front/components/assistant_builder/actions/RetrievalAction.tsx
@@ -1,3 +1,4 @@
+import { Checkbox } from "@dust-tt/sparkle";
 import type { SpaceType, WorkspaceType } from "@dust-tt/types";
 import { useEffect, useState } from "react";
 
@@ -7,6 +8,8 @@ import DataSourceSelectionSection from "@app/components/assistant_builder/DataSo
 import type {
   AssistantBuilderActionConfiguration,
   AssistantBuilderRetrievalConfiguration,
+  AssistantBuilderRetrievalExhaustiveConfiguration,
+  AssistantBuilderTimeFrame,
 } from "@app/components/assistant_builder/types";
 import { classNames } from "@app/lib/utils";
 
@@ -85,18 +88,19 @@ export function hasErrorActionRetrievalExhaustive(
 ): string | null {
   return action.type === "RETRIEVAL_EXHAUSTIVE" &&
     Object.keys(action.configuration.dataSourceConfigurations).length > 0 &&
-    !!action.configuration.timeFrame.value
+    // The time frame is optional for exhaustive retrieval, but if it is set, it must be valid.
+    (!action.configuration.timeFrame || !!action.configuration.timeFrame.value)
     ? null
-    : "Please select at least one data source and set a timeframe";
+    : "Please select at least one data source and set a valid timeframe";
 }
 
 type ActionRetrievalExhaustiveProps = {
   owner: WorkspaceType;
-  actionConfiguration: AssistantBuilderRetrievalConfiguration | null;
+  actionConfiguration: AssistantBuilderRetrievalExhaustiveConfiguration | null;
   allowedSpaces: SpaceType[];
   updateAction: (
     setNewAction: (
-      previousAction: AssistantBuilderRetrievalConfiguration
+      previousAction: AssistantBuilderRetrievalExhaustiveConfiguration
     ) => AssistantBuilderRetrievalConfiguration
   ) => void;
   setEdited: (edited: boolean) => void;
@@ -112,11 +116,25 @@ export function ActionRetrievalExhaustive({
   const [showDataSourcesModal, setShowDataSourcesModal] = useState(false);
   const [timeFrameError, setTimeFrameError] = useState<string | null>(null);
 
+  const [defaultTimeFrame, setDefaultTimeFrame] =
+    useState<AssistantBuilderTimeFrame>({
+      value: 1,
+      unit: "month",
+    });
+
   useEffect(() => {
     if (actionConfiguration) {
-      if (!actionConfiguration.timeFrame.value) {
+      if (
+        actionConfiguration.timeFrame &&
+        !actionConfiguration.timeFrame.value
+      ) {
         setTimeFrameError("Timeframe must be a number");
       } else {
+        // Set the default time frame to the current time frame if it exists,
+        // so if the user unchecks the checkbox, it won't reset to initial value
+        if (actionConfiguration.timeFrame) {
+          setDefaultTimeFrame(actionConfiguration.timeFrame);
+        }
         setTimeFrameError(null);
       }
     }
@@ -125,6 +143,8 @@ export function ActionRetrievalExhaustive({
   if (!actionConfiguration) {
     return null;
   }
+  const timeFrame = actionConfiguration.timeFrame || defaultTimeFrame;
+  const timeFrameDisabled = !actionConfiguration.timeFrame;
 
   return (
     <>
@@ -157,7 +177,22 @@ export function ActionRetrievalExhaustive({
         viewType={"documents"}
       />
       <div className={"flex flex-row items-center gap-4 pb-4"}>
-        <div className="text-sm font-semibold text-element-900">
+        <Checkbox
+          checked={!!actionConfiguration.timeFrame}
+          onCheckedChange={(checked) => {
+            setEdited(true);
+            updateAction((previousAction) => ({
+              ...previousAction,
+              timeFrame: checked ? defaultTimeFrame : undefined,
+            }));
+          }}
+        />
+        <div
+          className={classNames(
+            "text-sm font-semibold",
+            timeFrameDisabled ? "text-slate-400" : "text-element-900"
+          )}
+        >
           Collect data from the last
         </div>
         <input
@@ -167,9 +202,10 @@ export function ActionRetrievalExhaustive({
             !timeFrameError
               ? "focus:border-action-500 focus:ring-action-500"
               : "border-red-500 focus:border-red-500 focus:ring-red-500",
-            "bg-structure-50 stroke-structure-50"
+            "bg-structure-50 stroke-structure-50",
+            timeFrameDisabled ? "text-slate-400" : ""
           )}
-          value={actionConfiguration.timeFrame.value || ""}
+          value={timeFrame.value || ""}
           onChange={(e) => {
             const value = parseInt(e.target.value, 10);
             if (!isNaN(value) || !e.target.value) {
@@ -178,16 +214,18 @@ export function ActionRetrievalExhaustive({
                 ...previousAction,
                 timeFrame: {
                   value,
-                  unit: previousAction.timeFrame.unit,
+                  unit: timeFrame.unit,
                 },
               }));
             }
           }}
+          disabled={timeFrameDisabled}
         />
         <TimeUnitDropdown
-          actionConfiguration={actionConfiguration}
+          timeFrame={timeFrame}
           updateAction={updateAction}
           onEdit={() => setEdited(true)}
+          disabled={timeFrameDisabled}
         />
       </div>
     </>

diff --git a/front/components/assistant_builder/actions/TimeDropdown.tsx b/front/components/assistant_builder/actions/TimeDropdown.tsx
@@ -8,27 +8,29 @@ import {
 import type { TimeframeUnit } from "@dust-tt/types";
 
 import { TIME_FRAME_UNIT_TO_LABEL } from "@app/components/assistant_builder/shared";
-import type { AssistantBuilderBaseConfiguration } from "@app/components/assistant_builder/types";
+import type { AssistantBuilderTimeFrame } from "@app/components/assistant_builder/types";
 
-interface TimeUnitDropdownProps<T extends AssistantBuilderBaseConfiguration> {
-  actionConfiguration: T;
+interface TimeUnitDropdownProps<
+  T extends { timeFrame?: AssistantBuilderTimeFrame },
+> {
+  timeFrame: AssistantBuilderTimeFrame;
+  disabled?: boolean;
   onEdit: () => void;
   updateAction: (setNewAction: (previousAction: T) => T) => void;
 }
 
-export function TimeUnitDropdown<T extends AssistantBuilderBaseConfiguration>({
-  actionConfiguration,
-  updateAction,
-  onEdit,
-}: TimeUnitDropdownProps<T>) {
+export function TimeUnitDropdown<
+  T extends { timeFrame?: AssistantBuilderTimeFrame },
+>({ timeFrame, updateAction, onEdit, disabled }: TimeUnitDropdownProps<T>) {
   return (
     <NewDropdownMenu>
       <NewDropdownMenuTrigger asChild>
         <Button
           isSelect
-          label={TIME_FRAME_UNIT_TO_LABEL[actionConfiguration.timeFrame.unit]}
+          label={TIME_FRAME_UNIT_TO_LABEL[timeFrame.unit]}
           variant="outline"
           size="sm"
+          disabled={disabled}
         />
       </NewDropdownMenuTrigger>
       <NewDropdownMenuContent>
@@ -41,7 +43,7 @@ export function TimeUnitDropdown<T extends AssistantBuilderBaseConfiguration>({
               updateAction((previousAction) => ({
                 ...previousAction,
                 timeFrame: {
-                  value: previousAction.timeFrame.value,
+                  value: timeFrame.value,
                   unit: key as TimeframeUnit,
                 },
               }));

diff --git a/front/components/assistant_builder/server_side_props_helpers.ts b/front/components/assistant_builder/server_side_props_helpers.ts
@@ -120,9 +120,10 @@ async function getRetrievalActionConfiguration(
       : getDefaultRetrievalExhaustiveActionConfiguration();
   if (
     action.relativeTimeFrame !== "auto" &&
-    action.relativeTimeFrame !== "none"
+    action.relativeTimeFrame !== "none" &&
+    "timeFrame" in retrievalConfiguration
   ) {
-    retrievalConfiguration.configuration.timeFrame = {
+    retrievalConfiguration.timeFrame = {
       value: action.relativeTimeFrame.duration,
       unit: action.relativeTimeFrame.unit,
     };

diff --git a/front/components/assistant_builder/submitAssistantBuilderForm.ts b/front/components/assistant_builder/submitAssistantBuilderForm.ts
@@ -3,6 +3,7 @@ import type {
   LightAgentConfigurationType,
   PostOrPatchAgentConfigurationRequestBody,
   Result,
+  RetrievalTimeframe,
   WorkspaceType,
 } from "@dust-tt/types";
 import { assertNever, Err, Ok } from "@dust-tt/types";
@@ -61,6 +62,24 @@ export async function submitAssistantBuilderForm({
   >;
 
   const map: (a: AssistantBuilderActionConfiguration) => ActionsType = (a) => {
+    let timeFrame: RetrievalTimeframe = "auto";
+
+    if (a.type === "RETRIEVAL_EXHAUSTIVE") {
+      if (a.configuration.timeFrame) {
+        timeFrame = {
+          duration: a.configuration.timeFrame.value,
+          unit: a.configuration.timeFrame.unit,
+        };
+      } else {
+        timeFrame = "none";
+      }
+    } else if (a.type === "PROCESS") {
+      timeFrame = {
+        duration: a.configuration.timeFrame.value,
+        unit: a.configuration.timeFrame.unit,
+      };
+    }
+
     switch (a.type) {
       case "RETRIEVAL_SEARCH":
       case "RETRIEVAL_EXHAUSTIVE":
@@ -70,13 +89,7 @@ export async function submitAssistantBuilderForm({
             name: a.name,
             description: a.description,
             query: a.type === "RETRIEVAL_SEARCH" ? "auto" : "none",
-            relativeTimeFrame:
-              a.type === "RETRIEVAL_EXHAUSTIVE"
-                ? {
-                    duration: a.configuration.timeFrame.value,
-                    unit: a.configuration.timeFrame.unit,
-                  }
-                : "auto",
+            relativeTimeFrame: timeFrame,
             topK: "auto",
             dataSources: Object.values(
               a.configuration.dataSourceConfigurations
@@ -173,10 +186,7 @@ export async function submitAssistantBuilderForm({
               },
             })),
             tagsFilter: a.configuration.tagsFilter,
-            relativeTimeFrame: {
-              duration: a.configuration.timeFrame.value,
-              unit: a.configuration.timeFrame.unit,
-            },
+            relativeTimeFrame: timeFrame,
             schema: a.configuration.schema,
           },
         ];

diff --git a/front/components/assistant_builder/types.ts b/front/components/assistant_builder/types.ts
@@ -38,10 +38,6 @@ export const ACTION_MODES = [
 
 // Retrieval configuration
 
-export interface AssistantBuilderBaseConfiguration {
-  timeFrame: AssistantBuilderTimeFrame;
-}
-
 export type AssistantBuilderTimeFrame = {
   value: number;
   unit: TimeframeUnit;
@@ -51,10 +47,13 @@ export type AssistantBuilderTagsFilter = {
   in: string[];
 };
 
-export type AssistantBuilderRetrievalConfiguration =
-  AssistantBuilderBaseConfiguration & {
-    dataSourceConfigurations: DataSourceViewSelectionConfigurations;
-  };
+export type AssistantBuilderRetrievalConfiguration = {
+  dataSourceConfigurations: DataSourceViewSelectionConfigurations;
+};
+
+export type AssistantBuilderRetrievalExhaustiveConfiguration = {
+  timeFrame?: AssistantBuilderTimeFrame;
+} & AssistantBuilderRetrievalConfiguration;
 
 // DustAppRun configuration
 
@@ -69,12 +68,13 @@ export type AssistantBuilderTableConfiguration =
 
 // Process configuration
 
-export type AssistantBuilderProcessConfiguration =
-  AssistantBuilderBaseConfiguration & {
-    dataSourceConfigurations: DataSourceViewSelectionConfigurations;
-    tagsFilter: AssistantBuilderTagsFilter | null;
-    schema: ProcessSchemaPropertyType[];
-  };
+export type AssistantBuilderProcessConfiguration = {
+  timeFrame: AssistantBuilderTimeFrame;
+} & {
+  dataSourceConfigurations: DataSourceViewSelectionConfigurations;
+  tagsFilter: AssistantBuilderTagsFilter | null;
+  schema: ProcessSchemaPropertyType[];
+};
 
 // Websearch configuration
 export type AssistantBuilderWebNavigationConfiguration = Record<string, never>; // no relevant params identified yet
@@ -85,9 +85,13 @@ export type AssistantBuilderVisualizationConfiguration = Record<string, never>;
 
 export type AssistantBuilderActionConfiguration = (
   | {
-      type: "RETRIEVAL_SEARCH" | "RETRIEVAL_EXHAUSTIVE";
+      type: "RETRIEVAL_SEARCH";
       configuration: AssistantBuilderRetrievalConfiguration;
     }
+  | {
+      type: "RETRIEVAL_EXHAUSTIVE";
+      configuration: AssistantBuilderRetrievalExhaustiveConfiguration;
+    }
   | {
       type: "DUST_APP_RUN";
       configuration: AssistantBuilderDustAppConfiguration;
@@ -221,11 +225,7 @@ export function getDefaultRetrievalExhaustiveActionConfiguration() {
     type: "RETRIEVAL_EXHAUSTIVE",
     configuration: {
       dataSourceConfigurations: {},
-      timeFrame: {
-        value: 1,
-        unit: "month",
-      },
-    } as AssistantBuilderRetrievalConfiguration,
+    } as AssistantBuilderRetrievalExhaustiveConfiguration,
     name: DEFAULT_RETRIEVAL_NO_QUERY_ACTION_NAME,
     description: "",
   } satisfies AssistantBuilderActionConfiguration;

diff --git a/front/lib/api/assistant/actions/names.ts b/front/lib/api/assistant/actions/names.ts
@@ -3,7 +3,7 @@ export const DEFAULT_BROWSE_ACTION_NAME = "browse";
 export const DEFAULT_PROCESS_ACTION_NAME =
   "extract_structured_data_from_data_sources";
 export const DEFAULT_RETRIEVAL_ACTION_NAME = "search_data_sources";
-// The name below is used by the assistant builder when the user selects "Most recent data".
+// The name below is used by the assistant builder when the user selects "Include data".
 export const DEFAULT_RETRIEVAL_NO_QUERY_ACTION_NAME =
   "most_recent_in_data_sources";
 export const DEFAULT_WEBSEARCH_ACTION_NAME = "web_search";

diff --git a/front/lib/api/assistant/actions/utils.ts b/front/lib/api/assistant/actions/utils.ts
@@ -33,7 +33,7 @@ export const ACTION_SPECIFICATIONS: Record<
   }
 > = {
   RETRIEVAL_EXHAUSTIVE: {
-    label: "Most recent data",
+    label: "Include data",
     description: "Include as much data as possible",
     cardIcon: TimeIcon,
     dropDownIcon: TimeIcon,

diff --git a/types/src/front/assistant/actions/guards.ts b/types/src/front/assistant/actions/guards.ts
@@ -135,20 +135,6 @@ export function throwIfInvalidAgentConfiguration(
   configation: AgentConfigurationType | TemplateAgentConfigurationType
 ) {
   configation.actions.forEach((action) => {
-    if (isRetrievalConfiguration(action)) {
-      if (action.query === "none") {
-        if (
-          action.relativeTimeFrame === "auto" ||
-          action.relativeTimeFrame === "none"
-        ) {
-          /** Should never happen. Throw loudly if it does */
-          throw new Error(
-            "Invalid configuration: exhaustive retrieval must have a definite time frame"
-          );
-        }
-      }
-    }
-
     if (isProcessConfiguration(action)) {
       if (
         action.relativeTimeFrame === "auto" ||