support configuration of formulation timeout

Chenglong-MS · Chenglong-MS · commit 96a87a62f4a2 · 2025-02-21T14:52:42.000-08:00
diff --git a/py-src/data_formulator/app.py b/py-src/data_formulator/app.py
@@ -425,6 +425,8 @@ def derive_data():
         new_fields = content["new_fields"]
         instruction = content["extra_prompt"]
 
+        max_repair_attempts = content["max_repair_attempts"] if "max_repair_attempts" in content else 1
+
         if "additional_messages" in content:
             prev_messages = content["additional_messages"]
         else:
@@ -447,7 +449,7 @@ def derive_data():
             results = agent.run(input_tables, instruction, [field['name'] for field in new_fields], prev_messages)
 
         repair_attempts = 0
-        while results[0]['status'] == 'error' and repair_attempts == 0: # only try once
+        while results[0]['status'] == 'error' and repair_attempts < max_repair_attempts: # try up to n times
             error_message = results[0]['content']
             new_instruction = f"We run into the following problem executing the code, please fix it:\n\n{error_message}\n\nPlease think step by step, reflect why the error happens and fix the code so that no more errors would occur."
 
@@ -482,7 +484,10 @@ def refine_data():
         output_fields = content["output_fields"]
         dialog = content["dialog"]
         new_instruction = content["new_instruction"]
-        
+        max_repair_attempts = content["max_repair_attempts"] if "max_repair_attempts" in content else 1
+
+        print("max_repair_attempts")
+        print(max_repair_attempts)
         print("previous dialog")
         print(dialog)
 
@@ -491,7 +496,7 @@ def refine_data():
         results = agent.followup(input_tables, dialog, [field['name'] for field in output_fields], new_instruction)
 
         repair_attempts = 0
-        while results[0]['status'] == 'error' and repair_attempts == 0: # only try once
+        while results[0]['status'] == 'error' and repair_attempts < max_repair_attempts: # only try once
             error_message = results[0]['content']
             new_instruction = f"We run into the following problem executing the code, please fix it:\n\n{error_message}\n\nPlease think step by step, reflect why the error happens and fix the code so that no more errors would occur."
             prev_dialog = results[0]['dialog']
diff --git a/src/app/App.tsx b/src/app/App.tsx
@@ -35,6 +35,7 @@ import {
     ToggleButton,
     Menu,
     MenuItem,
+    TextField,
 } from '@mui/material';
 
 
@@ -46,7 +47,7 @@ import { DataFormulatorFC } from '../views/DataFormulator';
 
 import GridViewIcon from '@mui/icons-material/GridView';
 import ViewSidebarIcon from '@mui/icons-material/ViewSidebar';
-
+import SettingsIcon from '@mui/icons-material/Settings';
 import {
     createBrowserRouter,
     RouterProvider,
@@ -295,10 +296,101 @@ const ResetDialog: React.FC = () => {
     );
 };
 
+const ConfigDialog: React.FC = () => {
+    const [open, setOpen] = useState(false);
+    const dispatch = useDispatch();
+    const config = useSelector((state: DataFormulatorState) => state.config);
+
+    const [formulateTimeoutSeconds, setFormulateTimeoutSeconds] = useState(config.formulateTimeoutSeconds);
+    const [maxRepairAttempts, setMaxRepairAttempts] = useState(config.maxRepairAttempts);
+
+    // Add check for changes
+    const hasChanges = formulateTimeoutSeconds !== config.formulateTimeoutSeconds || 
+                      maxRepairAttempts !== config.maxRepairAttempts;
+
+    return (
+        <>
+            <Button variant="text" sx={{textTransform: 'none'}} onClick={() => setOpen(true)} startIcon={<SettingsIcon />}>
+                 <Box component="span" sx={{lineHeight: 1.2, display: 'flex', flexDirection: 'column', alignItems: 'left'}}>
+                    <Box component="span" sx={{py: 0, my: 0, fontSize: '10px', mr: 'auto'}}>timeout={config.formulateTimeoutSeconds}s</Box>
+                    <Box component="span" sx={{py: 0, my: 0, fontSize: '10px', mr: 'auto'}}>max_repair={config.maxRepairAttempts}</Box>
+                </Box>
+            </Button>
+            <Dialog onClose={() => setOpen(false)} open={open}>
+                <DialogTitle>Data Formulator Configuration</DialogTitle>
+                <DialogContent>
+                    <DialogContentText>
+                        <Box sx={{ 
+                            display: 'flex', 
+                            flexDirection: 'column', 
+                            gap: 3,
+                            my: 2,
+                            maxWidth: 400
+                        }}>
+                            <Box sx={{ display: 'flex', alignItems: 'center', gap: 2 }}>
+                                <Box sx={{ flex: 1 }}>
+                                    <TextField
+                                        label="formulate timeout (seconds)"
+                                        type="number"
+                                        variant="outlined"
+                                        value={formulateTimeoutSeconds}
+                                        onChange={(e) => setFormulateTimeoutSeconds(parseInt(e.target.value))}
+                                        fullWidth
+                                    />
+                                    <Typography variant="caption" color="text.secondary" sx={{ mt: 1, display: 'block' }}>
+                                        Maximum time allowed for the formulation process before timing out. 
+                                    </Typography>
+                                    <Typography variant="caption" color="text.secondary" sx={{ mt: 1, display: 'block' }}>
+                                        Smaller values (&lt;30s) make the model fails fast thus providing a smoother UI experience. Increase this value for slow models.
+                                    </Typography>
+                                </Box>
+                            </Box>
+                            <Box sx={{ display: 'flex', alignItems: 'center', gap: 2 }}>
+                                <Box sx={{ flex: 1 }}>
+                                    <TextField
+                                        label="max repair attempts"
+                                        type="number"
+                                        variant="outlined"
+                                        value={maxRepairAttempts}
+                                        onChange={(e) => setMaxRepairAttempts(parseInt(e.target.value))}
+                                        fullWidth
+                                    />
+                                    <Typography variant="caption" color="text.secondary" sx={{ mt: 1, display: 'block' }}>
+                                        Maximum number of times the LLM will attempt to repair code if generated code fails to execute (recommended = 1).
+                                    </Typography>
+                                    <Typography variant="caption" color="text.secondary" sx={{ mt: 1, display: 'block' }}>
+                                        Higher values <span style={{fontStyle: 'italic'}}>might</span> increase chances of success but may take longer. Repair time is considered as part of the formulate timeout.
+                                    </Typography>
+                                </Box>
+                            </Box>
+                        </Box>
+                    </DialogContentText>
+                </DialogContent>
+                <DialogActions sx={{'.MuiButton-root': {textTransform: 'none'}}}>
+                    <Button sx={{marginRight: 'auto'}} onClick={() => {
+                        setFormulateTimeoutSeconds(30);
+                        setMaxRepairAttempts(1);
+                    }}>Reset to default</Button>
+                    <Button onClick={() => setOpen(false)}>Cancel</Button>
+                    <Button 
+                        variant={hasChanges ? "contained" : "text"}
+                        onClick={() => {
+                            dispatch(dfActions.setConfig({formulateTimeoutSeconds, maxRepairAttempts}));
+                            setOpen(false);
+                        }}
+                    >
+                        Apply
+                    </Button>
+                </DialogActions>
+            </Dialog>
+        </>
+    );  
+}
+
 export const AppFC: FC<AppFCProps> = function AppFC(appProps) {
 
     const visViewMode = useSelector((state: DataFormulatorState) => state.visViewMode);
-    const betaMode = useSelector((state: DataFormulatorState) => state.betaMode);
+    const config = useSelector((state: DataFormulatorState) => state.config);
     const tables = useSelector((state: DataFormulatorState) => state.tables);
 
     // if the user has logged in
@@ -410,7 +502,7 @@ export const AppFC: FC<AppFCProps> = function AppFC(appProps) {
 
     let appBar = [
         <AppBar className="app-bar" position="static" key="app-bar-main">
-            <Toolbar variant="dense" sx={{ backgroundColor: betaMode ? 'lavender' : '' }}>
+            <Toolbar variant="dense">
                 <Button href={"/"} sx={{
                     display: "flex", flexDirection: "row", textTransform: "none",
                     backgroundColor: 'transparent',
@@ -420,7 +512,7 @@ export const AppFC: FC<AppFCProps> = function AppFC(appProps) {
                 }} color="inherit">
                     <Box component="img" sx={{ height: 32, marginRight: "12px" }} alt="" src={dfLogo} />
                     <Typography variant="h6" noWrap component="h1" sx={{ fontWeight: 300, display: { xs: 'none', sm: 'block' } }}>
-                        {toolName} {betaMode ? "β" : ""} {process.env.NODE_ENV == "development" ? "" : ""}
+                        {toolName} {process.env.NODE_ENV == "development" ? "" : ""}
                     </Typography>
                 </Button>
                 <Box sx={{ flexGrow: 1, textAlign: 'center', display: 'flex', justifyContent: 'center' }} >
@@ -432,6 +524,8 @@ export const AppFC: FC<AppFCProps> = function AppFC(appProps) {
                         about
                     </Button>
                     <Divider orientation="vertical" variant="middle" flexItem /> */}
+                    <ConfigDialog />
+                    <Divider orientation="vertical" variant="middle" flexItem />
                     <ModelSelectionButton />
                     <Divider orientation="vertical" variant="middle" flexItem />
                     <Typography sx={{ display: 'flex', fontSize: 14, alignItems: 'center', gap: 1 }}>
diff --git a/src/app/dfSlice.tsx b/src/app/dfSlice.tsx
@@ -65,7 +65,10 @@ export interface DataFormulatorState {
 
     chartSynthesisInProgress: string[];
 
-    betaMode: boolean;
+    config: {
+        formulateTimeoutSeconds: number;
+        maxRepairAttempts: number;
+    }
 }
 
 // Define the initial state using that type
@@ -98,7 +101,10 @@ const initialState: DataFormulatorState = {
 
     chartSynthesisInProgress: [],
 
-    betaMode: false,
+    config: {
+        formulateTimeoutSeconds: 30,
+        maxRepairAttempts: 1,
+    }
 }
 
 let getUnrefedDerivedTableIds = (state: DataFormulatorState) => {
@@ -247,7 +253,11 @@ export const dataFormulatorSlice = createSlice({
 
             state.chartSynthesisInProgress = [];
 
-            state.betaMode = false;
+            // avoid resetting config
+            // state.config = {
+            //     formulateTimeoutSeconds: 30,
+            //     repairAttempts: 1,
+            // }
         },
         loadState: (state, action: PayloadAction<any>) => {
 
@@ -274,10 +284,10 @@ export const dataFormulatorSlice = createSlice({
 
             state.chartSynthesisInProgress = [];
 
-            state.betaMode = savedState.betaMode;
+            state.config = savedState.config;
         },
-        toggleBetaMode: (state, action: PayloadAction<boolean>) => {
-            state.betaMode = action.payload;
+        setConfig: (state, action: PayloadAction<{formulateTimeoutSeconds: number, maxRepairAttempts: number}>) => {
+            state.config = action.payload;
         },
         selectModel: (state, action: PayloadAction<string | undefined>) => {
             state.selectedModelId = action.payload;
diff --git a/src/views/EncodingShelfCard.tsx b/src/views/EncodingShelfCard.tsx
@@ -308,6 +308,7 @@ export const EncodingShelfCard: FC<EncodingShelfCardProps> = function ({ chartId
     // reference to states
     const tables = useSelector((state: DataFormulatorState) => state.tables);
     const charts = useSelector((state: DataFormulatorState) => state.charts);
+    const config = useSelector((state: DataFormulatorState) => state.config);
 
     let existMultiplePossibleBaseTables = tables.filter(t => t.derive == undefined).length > 1;
 
@@ -421,7 +422,8 @@ export const EncodingShelfCard: FC<EncodingShelfCardProps> = function ({ chartId
             input_tables: actionTables.map(t => {return { name: t.id.replace(/\.[^/.]+$/ , ""), rows: t.rows }}),
             new_fields: activeBaseFields.map(f => { return {name: f.name} }),
             extra_prompt: instruction,
-            model: activeModel
+            model: activeModel,
+            max_repair_attempts: config.maxRepairAttempts
         }) 
         let engine = getUrls().SERVER_DERIVE_DATA_URL;
 
@@ -449,7 +451,8 @@ export const EncodingShelfCard: FC<EncodingShelfCardProps> = function ({ chartId
                         new_fields: activeBaseFields.map(f => { return {name: f.name} }),
                         extra_prompt: instruction,
                         model: activeModel,
-                        additional_messages: additionalMessages
+                        additional_messages: additionalMessages,
+                        max_repair_attempts: config.maxRepairAttempts
                     });
                     engine = getUrls().SERVER_DERIVE_DATA_URL;
                 } else {
@@ -460,7 +463,8 @@ export const EncodingShelfCard: FC<EncodingShelfCardProps> = function ({ chartId
                         output_fields: activeBaseFields.map(f => { return {name: f.name} }),
                         dialog: currentTable.derive?.dialog,
                         new_instruction: instruction,
-                        model: activeModel
+                        model: activeModel,
+                        max_repair_attempts: config.maxRepairAttempts
                     })
                     engine = getUrls().SERVER_REFINE_DATA_URL;
                 }
@@ -479,7 +483,7 @@ export const EncodingShelfCard: FC<EncodingShelfCardProps> = function ({ chartId
 
         // timeout the request after 30 seconds
         const controller = new AbortController();
-        const timeoutId = setTimeout(() => controller.abort(), 30000);
+        const timeoutId = setTimeout(() => controller.abort(), config.formulateTimeoutSeconds * 1000);
     
         fetch(engine, {...message, signal: controller.signal })
             .then((response) => response.json())
@@ -650,7 +654,7 @@ export const EncodingShelfCard: FC<EncodingShelfCardProps> = function ({ chartId
                     dispatch(dfActions.addMessages({
                         "timestamp": Date.now(),
                         "type": "error",
-                        "value": "Data formulation timed out after 30 seconds. Please try again.",
+                        "value": `Data formulation timed out after ${config.formulateTimeoutSeconds} seconds. Consider breaking down the task, using a different model or prompt, or increasing the timeout limit.`,
                         "detail": "Request exceeded timeout limit"
                     }));
                 } else {
diff --git a/src/views/EncodingShelfThread.tsx b/src/views/EncodingShelfThread.tsx
@@ -155,7 +155,8 @@ export const EncodingShelfThread: FC<EncodingShelfThreadProps> = function ({ cha
     const charts = useSelector((state: DataFormulatorState) => state.charts);
     let activeThreadChartId = useSelector((state: DataFormulatorState) => state.activeThreadChartId);
     let activeModel = useSelector(dfSelectors.getActiveModel);
-    
+    const config = useSelector((state: DataFormulatorState) => state.config);
+
     let [reformulateRunning, setReformulteRunning] = useState<boolean>(false);
 
     let activeThreadChart = charts.find(c => c.id == activeThreadChartId);
@@ -215,7 +216,8 @@ export const EncodingShelfThread: FC<EncodingShelfThreadProps> = function ({ cha
             input_tables: baseTables.map(t => {return { name: t.id.replace(/\.[^/.]+$/ , ""), rows: t.rows }}),
             new_fields: activeBaseFields.map(f => { return {name: f.name} }),
             extra_prompt: prompt,
-            model: activeModel
+            model: activeModel,
+            max_repair_attempts: config.maxRepairAttempts
         }) 
         let engine = getUrls().SERVER_DERIVE_DATA_URL;
 
@@ -231,7 +233,8 @@ export const EncodingShelfThread: FC<EncodingShelfThreadProps> = function ({ cha
                 new_fields: activeBaseFields.map(f => { return {name: f.name} }),
                 extra_prompt: prompt,
                 additional_messages: triggerTable.derive?.dialog,
-                model: activeModel
+                model: activeModel,
+                max_repair_attempts: config.maxRepairAttempts
             }) 
             engine = getUrls().SERVER_DERIVE_DATA_URL;
             // messageBody = JSON.stringify({
@@ -241,7 +244,8 @@ export const EncodingShelfThread: FC<EncodingShelfThreadProps> = function ({ cha
             //     output_fields: activeBaseFields.map(f => { return {name: f.name } }),
             //     dialog: triggerTable.derive?.dialog,
             //     new_instruction: prompt,
-            //     model: activeModel
+            //     model: activeModel,
+            //     max_repair_attempts: config.maxRepairAttempts
             // })
             //engine = getUrls().SERVER_REFINE_DATA_URL;
         }
@@ -258,7 +262,7 @@ export const EncodingShelfThread: FC<EncodingShelfThreadProps> = function ({ cha
 
         // timeout the request after 30 seconds
         const controller = new AbortController();
-        const timeoutId = setTimeout(() => controller.abort(), 30000);
+        const timeoutId = setTimeout(() => controller.abort(), config.formulateTimeoutSeconds * 1000);
 
         setReformulteRunning(true);
     
@@ -386,7 +390,7 @@ export const EncodingShelfThread: FC<EncodingShelfThreadProps> = function ({ cha
                     dispatch(dfActions.addMessages({
                         "timestamp": Date.now(),
                         "type": "error",
-                        "value": "Data formulation timed out after 30 seconds. Please try again.",
+                        "value": `Data formulation timed out after ${config.formulateTimeoutSeconds} seconds. Consider breaking down the task, using a different model or prompt, or increasing the timeout limit.`,
                         "detail": "Request exceeded timeout limit"
                     }));
                 } else {
diff --git a/src/views/ModelSelectionDialog.tsx b/src/views/ModelSelectionDialog.tsx
@@ -436,7 +436,7 @@ export const ModelSelectionButton: React.FC<{}> = ({ }) => {
 
     return <>
         <Tooltip title="select model">
-            <Button sx={{fontSize: "inherit", textTransform: "none"}} variant="text" color="primary" onClick={()=>{setModelDialogOpen(true)}} endIcon={selectedModelId ? <SettingsIcon /> : ''}>
+            <Button sx={{fontSize: "inherit", textTransform: "none"}} variant="text" color="primary" onClick={()=>{setModelDialogOpen(true)}}>
                 {selectedModelId ? `Model: ${(models.find(m => m.id == selectedModelId) as any)?.model}` : 'Select A Model'}
             </Button>
         </Tooltip>
