Feat: Optimize database performance and fix dataset regeneration UI

.gitignore

@@ -69,6 +69,13 @@ coverage.xml
 .tox/
 .nox/
 .pytest_cache
+
+# Test files and generated data
+SeedsInstructions.json
+seeds_housing.json
+app/test_models.py
+freeform_data_*.json
+
 #old code
 app/frontend/
 app/launch_streamlit.py

app/client/src/api/Datasets/response.ts

@@ -17,6 +17,7 @@ export type DatasetResponse = {
     schema: string | null;
     custom_prompt: string;
     total_count: number;
+    completed_rows: number | null;
     num_questions: number;
     job_id: string;
     job_name: string;

app/client/src/api/api.ts

@@ -13,12 +13,12 @@ import {
 const baseUrl = import.meta.env.VITE_AMP_URL;
 
 export const usefetchTopics = (useCase: string): UseFetchApiReturn<FetchTopicsResp> => {
-    const url = `${baseUrl}/use-cases/${isEmpty(useCase) ? 'custom' : useCase}/topics`;
+    const url = isEmpty(useCase) ? '' : `${baseUrl}/use-cases/${useCase}/topics`;
     return useFetch(url);
 }
 
 export const useFetchExamples = (useCase: string): UseFetchApiReturn<FetchExamplesResp> => {
-    const url = `${baseUrl}/${isEmpty(useCase) ? 'custom' : useCase}/gen_examples`;
+    const url = isEmpty(useCase) ? '' : `${baseUrl}/${useCase}/gen_examples`;
     return useFetch(url);
 }
 
@@ -27,21 +27,25 @@ export const useFetchModels = (): UseFetchApiReturn<FetchModelsResp> => {
     return useFetch(url);
 }
 
-export const useFetchDefaultPrompt = (useCase: string, workflowType?: WorkerType): UseFetchApiReturn<FetchDefaultPromptResp> => {
-    let url = `${baseUrl}/${isEmpty(useCase) ? 'custom' : useCase}/gen_prompt`;
+export const useFetchDefaultPrompt = (useCase: string, workflowType?: string): UseFetchApiReturn<FetchDefaultPromptResp> => {
+    if (isEmpty(useCase)) {
+        return { data: null, loading: false, error: null };
+    }
+
+    let url = `${baseUrl}/${useCase}/gen_prompt`;
     if (workflowType && workflowType === 'freeform') {
-        url = `${baseUrl}/${isEmpty(useCase) ? 'custom' : useCase}/gen_freeform_prompt`;
+        url = `${baseUrl}/${useCase}/gen_freeform_prompt`;
     }
     return useFetch(url);
 }
 
-export const useFetchDefaultSchema = (): UseFetchApiReturn<FetchDefaultSchemaResp> => {
-    const url = `${baseUrl}/sql_schema`;
+export const useFetchDefaultSchema = (shouldFetch: boolean = true): UseFetchApiReturn<FetchDefaultSchemaResp> => {
+    const url = shouldFetch ? `${baseUrl}/sql_schema` : '';
     return useFetch(url);
 }
 
-export const useFetchDefaultModelParams = (): UseFetchApiReturn<FetchDefaultParamsResp> => {
-    const url = `${baseUrl}/model/parameters`;
+export const useFetchDefaultModelParams = (shouldFetch: boolean = true): UseFetchApiReturn<FetchDefaultParamsResp> => {
+    const url = shouldFetch ? `${baseUrl}/model/parameters` : '';
     return useFetch(url);
 }
 

app/client/src/api/hooks.ts

@@ -1,15 +1,23 @@
 import { useState, useMemo, useEffect } from 'react';
 import { UseDeferredFetchApiReturn, UseFetchApiReturn } from './types';
-
-
+import { useQuery } from '@tanstack/react-query';
 
 export function useFetch<T>(url: string): UseFetchApiReturn<T> {
   const [data, setData] = useState<T | null>(null);
-  const [loading, setLoading] = useState<boolean>(true);
+  const [loading, setLoading] = useState<boolean>(false);
   const [error, setError] = useState<Error | null>(null);
   const memoizedUrl = useMemo(() => url, [url]);
 
   useEffect(() => {
+    // Don't make API call if URL is empty (for regeneration scenarios)
+    if (!memoizedUrl || memoizedUrl.trim() === '') {
+      setData(null);
+      setLoading(false);
+      setError(null);
+      return;
+    }
+
+    setLoading(true);
     const fetchData = async () => {
       try {
         const response = await fetch(memoizedUrl, {
@@ -29,7 +37,12 @@ export function useFetch<T>(url: string): UseFetchApiReturn<T> {
     fetchData();
   }, [memoizedUrl]);
 
-  return { data, loading, error };
+  // Return false for loading when URL is empty
+  return { 
+    data, 
+    loading: (!memoizedUrl || memoizedUrl.trim() === '') ? false : loading, 
+    error 
+  };
 }
 
 interface UseGetApiReturn<T> {
@@ -77,6 +90,7 @@ export function useGetApi<T>(url: string): UseGetApiReturn<T> {
 
   return { data, loading, error, triggerGet };
 }
+
 export function useDeferredFetch<T>(url: string): UseDeferredFetchApiReturn<T> {
   const [data, setData] = useState<T | null>(null);
   const [loading, setLoading] = useState<boolean>(true);
@@ -192,3 +206,117 @@ export function useDeleteApi<T>(url: string): UseDeleteApiReturn<T> {
   return { data, loading, error, triggerDelete };
 }
 
+// Use case types and enums
+export enum UseCaseId {
+  CODE_GENERATION = 'code_generation',
+  TEXT2SQL = 'text2sql',
+  CUSTOM = 'custom',
+  LENDING_DATA = 'lending_data',
+  CREDIT_CARD_DATA = 'credit_card_data',
+  TICKETING_DATASET = 'ticketing_dataset',
+}
+
+export interface UseCase {
+  id: string;
+  name: string;
+}
+
+export interface UseCasesResponse {
+  usecases: UseCase[];
+}
+
+const fetchUseCases = async (): Promise<UseCasesResponse> => {
+  const BASE_API_URL = import.meta.env.VITE_AMP_URL;
+  const response = await fetch(`${BASE_API_URL}/use-cases`);
+  if (!response.ok) {
+    throw new Error('Failed to fetch use cases');
+  }
+  return response.json();
+};
+
+export const useUseCases = () => {
+  return useQuery<UseCasesResponse>({
+    queryKey: ['useCases'],
+    queryFn: fetchUseCases,
+    staleTime: 10 * 60 * 1000, // Cache for 10 minutes
+    retry: 3, // Retry 3 times on failure
+    retryDelay: (attemptIndex: number) => Math.min(1000 * 2 ** attemptIndex, 30000), // Exponential backoff
+    refetchOnWindowFocus: false, // Don't refetch when window gains focus
+    refetchOnMount: false, // Don't refetch on component mount if data exists
+  });
+};
+
+export const useUseCaseMapping = () => {
+  const { data: useCasesData, isLoading, isError, error } = useUseCases();
+
+  // Create a lookup map for fast O(1) access
+  const useCaseMap = useMemo(() => {
+    if (!useCasesData?.usecases) return {};
+
+    return (useCasesData as UseCasesResponse).usecases.reduce((acc: Record<string, string>, useCase: UseCase) => {
+      acc[useCase.id] = useCase.name;
+      return acc;
+    }, {} as Record<string, string>);
+  }, [useCasesData]);
+
+  // Helper function to get use case name with better fallback
+  const getUseCaseName = (id: string): string => {
+    if (isError) {
+      return id || 'N/A';
+    }
+
+    if (isLoading) {
+      return id || 'Loading...';
+    }
+
+    const name = useCaseMap[id];
+
+    // Log missing use cases in development
+    if (!name && id && typeof window !== 'undefined' && window.location.hostname === 'localhost') {
+      console.warn(`Missing use case mapping for: ${id}`);
+    }
+
+    return name || id || 'N/A';
+  };
+
+  // Get all use cases as array (useful for dropdowns)
+  const useCases = useMemo(() => {
+    return (useCasesData as UseCasesResponse)?.usecases || [];
+  }, [useCasesData]);
+
+  return {
+    useCaseMap,
+    useCases,
+    getUseCaseName,
+    isLoading,
+    isError,
+    error
+  };
+};
+
+// Hook to provide use case options for dropdowns and forms
+export const useUseCaseOptions = () => {
+  const { useCases, isLoading, isError } = useUseCaseMapping();
+
+  // Transform use cases to option format used in dropdowns
+  const useCaseOptions = useMemo(() => {
+    return useCases.map((useCase: UseCase) => ({
+      label: useCase.name,
+      value: useCase.id
+    }));
+  }, [useCases]);
+
+  // Helper function to get use case type/name by id (replaces getUsecaseType)
+  const getUseCaseType = (id: string): string => {
+    const useCase = useCases.find((uc: UseCase) => uc.id === id);
+    return useCase?.name || id || 'N/A';
+  };
+
+  return {
+    useCaseOptions,
+    getUseCaseType,
+    isLoading,
+    isError
+  };
+};
+

app/client/src/components/Datasets/Datasets.tsx

@@ -12,7 +12,7 @@ import { Link } from "react-router-dom";
 import { HuggingFaceIconUrl, Pages } from "../../types";
 import { blue } from '@ant-design/colors';
 import DateTime from "../DateTime/DateTime";
-import { TRANSLATIONS } from "../../constants";
+import { useUseCaseMapping } from "../../api/hooks";
 import DeleteConfirmWarningModal from './DeleteConfirmModal';
 import DatasetExportModal, { ExportResult } from '../Export/ExportModal';
 
@@ -22,6 +22,7 @@ const { Paragraph, Text } = Typography;
 export default function DatasetsComponent() {
   const datasetHistoryAPI = useGetDatasetHistory();
   const deleteDatasetHistoryAPI = useDeleteDataset();
+  const { getUseCaseName } = useUseCaseMapping();
   const [toggleDatasetDetailModal, setToggleDatasetDetailModal] = React.useState(false);
   const [toggleDatasetExportModal, setToggleDatasetExportModal] = React.useState(false);
   const [exportResult, setExportResult] = React.useState<ExportResult>();
@@ -65,19 +66,27 @@ export default function DatasetsComponent() {
       key: '3',
       title: 'Model',
       dataIndex: 'model_id',
-      render: (modelId) => <Tooltip title={modelId}><Paragraph style={{ width: 150, marginBottom: 0 }} ellipsis={{ rows: 1 }}>{modelId}</Paragraph></Tooltip>
+      render: (modelId: string) => <Tooltip title={modelId}><Paragraph style={{ width: 150, marginBottom: 0 }} ellipsis={{ rows: 1 }}>{modelId}</Paragraph></Tooltip>
     },
     {
       key: '4',
       title: 'Questions Per Topic',
       dataIndex: 'num_questions',
       width: 150
     },
+    {
+      key: '4a',
+      title: 'Completed Rows',
+      dataIndex: 'completed_rows',
+      width: 120,
+      align: 'center',
+      render: (completed_rows: number | null) => <>{completed_rows != null ? completed_rows : 'N/A'}</>
+    },
     {
       key: '5',
       title: 'Use Case',
       dataIndex: 'use_case',
-      render: (useCase) => <Paragraph style={{ width: 200, marginBottom: 0 }} ellipsis={{ rows: 1 }}>{TRANSLATIONS[useCase]}</Paragraph>
+      render: (useCase: string) => <Paragraph style={{ width: 200, marginBottom: 0 }} ellipsis={{ rows: 1 }}>{getUseCaseName(useCase)}</Paragraph>
     },
     {
       key: '6',

app/client/src/components/Evaluations/Evaluations.tsx

@@ -8,7 +8,7 @@ import DeleteIcon from '@mui/icons-material/Delete';
 import { DownOutlined, FolderViewOutlined, ThunderboltOutlined } from '@ant-design/icons';
 
 import { blue } from '@ant-design/colors';
-import { TRANSLATIONS } from "../../constants";
+import { useUseCaseMapping } from "../../api/hooks";
 import DateTime from "../DateTime/DateTime";
 import styled from "styled-components";
 import { Pages } from "../../types";
@@ -23,6 +23,7 @@ const ModalButtonGroup = styled(Flex)`
 export default function Evaluations() {
   const evaluationsHistoryAPI = useGetEvaluationsHistory();
   const deleteEvaluationHistoryAPI = useDeleteEvaluation();
+  const { getUseCaseName } = useUseCaseMapping();
   const [toggleEvaluationDetailModal, setToggleEvaluationDetailModal] = React.useState(false);
   const [evaluationDetail, setEvaluationDetail] = React.useState<EvaluationResponse>({} as EvaluationResponse);
 
@@ -46,7 +47,7 @@ export default function Evaluations() {
       key: '4',
       title: 'Use Case',
       dataIndex: 'use_case',
-      render: (useCase) => <Paragraph style={{ width: 200, marginBottom: 0 }} ellipsis={{ rows: 1 }}>{TRANSLATIONS[useCase]}</Paragraph>
+      render: (useCase: string) => <Paragraph style={{ width: 200, marginBottom: 0 }} ellipsis={{ rows: 1 }}>{getUseCaseName(useCase)}</Paragraph>
     },
     {
       key: '5',

app/client/src/pages/DataGenerator/Configure.tsx

@@ -4,6 +4,7 @@ import isFunction from 'lodash/isFunction';
 import { useEffect, useState } from 'react';
 import { Flex, Form, Input, Select, Typography } from 'antd';
 import styled from 'styled-components';
+import { useLocation } from 'react-router-dom';
 import { File, WorkflowType } from './types';
 import { useFetchModels } from '../../api/api';
 import { MODEL_PROVIDER_LABELS } from './constants';
@@ -52,6 +53,7 @@ const Configure = () => {
     const formData = Form.useWatch((values) => values, form);
     const { setIsStepValid } = useWizardCtx();
     const { data } = useFetchModels();
+    const location = useLocation();
     const [selectedFiles, setSelectedFiles] = useState(
         !isEmpty(form.getFieldValue('doc_paths')) ? form.getFieldValue('doc_paths') : []);
 
@@ -73,12 +75,19 @@ const Configure = () => {
         validateForm()
     }, [form, formData])
 
-    // keivan
+    // Only set default inference_type for completely new datasets
     useEffect(() => {
-        if (formData && formData?.inference_type === undefined) {
+        const isRegenerating = location.state?.data || location.state?.internalRedirect;
+        const existingInferenceType = form.getFieldValue('inference_type');
+
+        // Only set default if:
+        // 1. NOT regenerating an existing dataset
+        // 2. No existing inference_type value in form
+        // 3. formData watch shows undefined (initial state)
+        if (!isRegenerating && !existingInferenceType && formData && formData?.inference_type === undefined) {
             form.setFieldValue('inference_type', ModelProviders.CAII);
         }
-    }, [formData]);
+    }, [formData, location.state, form]);
 
     const labelCol = {
       span: 8