API cleanup

Author: arminmeh

docs/data/data-grid/prompt/PromptWithDataSampling.js

@@ -22,10 +22,7 @@ function ToolbarWithPromptInput() {
   const [prompt, setPrompt] = React.useState('');
 
   const context = React.useMemo(
-    () =>
-      apiRef.current.unstable_getPromptContext(
-        apiRef.current.unstable_collectSampleData(),
-      ),
+    () => apiRef.current.unstable_getPromptContext(true),
     [apiRef],
   );
 
@@ -35,7 +32,7 @@ function ToolbarWithPromptInput() {
       setPrompt(selectedPrompt);
       if (selectedPrompt) {
         apiRef.current.setLoading(true);
-        const result = await mockPromptResolver(context, selectedPrompt);
+        const result = await mockPromptResolver(selectedPrompt, context);
         apiRef.current.unstable_applyPromptResult(result);
         apiRef.current.setLoading(false);
       }

docs/data/data-grid/prompt/PromptWithDataSampling.tsx

@@ -22,10 +22,7 @@ function ToolbarWithPromptInput() {
   const [prompt, setPrompt] = React.useState('');
 
   const context = React.useMemo(
-    () =>
-      apiRef.current.unstable_getPromptContext(
-        apiRef.current.unstable_collectSampleData(),
-      ),
+    () => apiRef.current.unstable_getPromptContext(true),
     [apiRef],
   );
 
@@ -35,7 +32,7 @@ function ToolbarWithPromptInput() {
       setPrompt(selectedPrompt);
       if (selectedPrompt) {
         apiRef.current.setLoading(true);
-        const result = await mockPromptResolver(context, selectedPrompt);
+        const result = await mockPromptResolver(selectedPrompt, context);
         apiRef.current.unstable_applyPromptResult(result);
         apiRef.current.setLoading(false);
       }

docs/data/data-grid/prompt/PromptWithDataSource.js

@@ -44,10 +44,7 @@ function ToolbarWithPromptInput() {
   const [prompt, setPrompt] = React.useState('');
 
   const context = React.useMemo(
-    () =>
-      apiRef.current.unstable_getPromptContext(
-        apiRef.current.unstable_collectSampleData(),
-      ),
+    () => apiRef.current.unstable_getPromptContext(true),
     [apiRef],
   );
 
@@ -57,7 +54,7 @@ function ToolbarWithPromptInput() {
       setPrompt(selectedPrompt);
       if (selectedPrompt) {
         apiRef.current.setLoading(true);
-        const result = await mockPromptResolver(context, selectedPrompt);
+        const result = await mockPromptResolver(selectedPrompt, context);
         apiRef.current.unstable_applyPromptResult(result);
       }
     },

docs/data/data-grid/prompt/PromptWithDataSource.tsx

@@ -46,10 +46,7 @@ function ToolbarWithPromptInput() {
   const [prompt, setPrompt] = React.useState('');
 
   const context = React.useMemo(
-    () =>
-      apiRef.current.unstable_getPromptContext(
-        apiRef.current.unstable_collectSampleData(),
-      ),
+    () => apiRef.current.unstable_getPromptContext(true),
     [apiRef],
   );
 
@@ -59,7 +56,7 @@ function ToolbarWithPromptInput() {
       setPrompt(selectedPrompt);
       if (selectedPrompt) {
         apiRef.current.setLoading(true);
-        const result = await mockPromptResolver(context, selectedPrompt);
+        const result = await mockPromptResolver(selectedPrompt, context);
         apiRef.current.unstable_applyPromptResult(result);
       }
     },

docs/data/data-grid/prompt/PromptWithExamples.js

@@ -44,7 +44,7 @@ function ToolbarWithPromptInput() {
       setPrompt(selectedPrompt);
       if (selectedPrompt) {
         apiRef.current.setLoading(true);
-        const result = await mockPromptResolver(context, selectedPrompt);
+        const result = await mockPromptResolver(selectedPrompt, context);
         apiRef.current.unstable_applyPromptResult(result);
         apiRef.current.setLoading(false);
       }

docs/data/data-grid/prompt/PromptWithExamples.tsx

@@ -44,7 +44,7 @@ function ToolbarWithPromptInput() {
       setPrompt(selectedPrompt);
       if (selectedPrompt) {
         apiRef.current.setLoading(true);
-        const result = await mockPromptResolver(context, selectedPrompt);
+        const result = await mockPromptResolver(selectedPrompt, context);
         apiRef.current.unstable_applyPromptResult(result);
         apiRef.current.setLoading(false);
       }

docs/data/data-grid/prompt/prompt.md

@@ -86,13 +86,14 @@ const PROMPT_RESOLVER_PROXY_BASE_URL =
     ? 'http://localhost:3000'
     : 'https://api.my-proxy.com';
 
-function processPrompt(context: string, query: string) {
-  const extendedContext = `The rows represent: List of employees with their company, position and start date\n\n${context}`;
+function processPrompt(query: string, context: string) {
+  const additionalContext = `The rows represent: List of employees with their company, position and start date`;
 
   return unstable_gridDefaultPromptResolver(
     `${PROMPT_RESOLVER_PROXY_BASE_URL}/api/datagrid/prompt`,
-    extendedContext,
     query,
+    context,
+    additionalContext,
   );
 }
 ```

docs/pages/x/api/data-grid/grid-api.json

@@ -471,13 +471,8 @@
       "required": true,
       "isPremiumPlan": true
     },
-    "unstable_collectSampleData": {
-      "type": { "description": "(sampleCount?: number) => Record<string, any[]>" },
-      "required": true,
-      "isPremiumPlan": true
-    },
     "unstable_getPromptContext": {
-      "type": { "description": "(examples?: Record<string, any[]>) => string" },
+      "type": { "description": "(allowDataSampling?: boolean) => string" },
       "required": true,
       "isPremiumPlan": true
     },

docs/translations/api-docs/data-grid/grid-api.json

@@ -242,9 +242,6 @@
     "unstable_applyPromptResult": {
       "description": "Use the prompt processing result to update all relevant parts of the grid state."
     },
-    "unstable_collectSampleData": {
-      "description": "Collects sample data from the grid that can be used to improve the prompt context.<br />It takes random values for all columns but from different rows to make the samples more random."
-    },
     "unstable_getPromptContext": { "description": "Get the context for the prompt." },
     "unstable_replaceRows": { "description": "Replace a set of rows with new rows." },
     "unstable_setColumnVirtualization": { "description": "Enable/disable column virtualization." },

packages/x-data-grid-generator/src/services/prompt-resolver.ts

@@ -1,7 +1,7 @@
 import { unstable_PromptResponse as PromptResponse } from '@mui/x-data-grid-premium';
 import { mockPrompts } from '../constants/prompts';
 
-export const mockPromptResolver = (_: string, query: string) => {
+export const mockPromptResolver = (query: string, _: string) => {
   const resolved = mockPrompts.get(query.toLowerCase().trim());
 
   return new Promise<PromptResponse>((resolve, reject) => {

packages/x-data-grid-premium/src/components/promptControl/GridToolbarPromptControl.tsx

@@ -42,7 +42,7 @@ const GridToolbarPromptControlRoot = styled('div', {
   flexDirection: 'row',
 });
 
-type GridToolbarPromptControlProps = {
+export type GridToolbarPromptControlProps = {
   /**
    * Allow taking couple of random cell values from each column to improve the prompt context.
    * If allowed, samples are taken from different rows.
@@ -58,11 +58,11 @@ type GridToolbarPromptControlProps = {
   /**
    * Called when the new prompt is ready to be processed.
    * Provides the prompt and the data context and expects the grid state updates to be returned.
-   * @param {string} context The context of the prompt
    * @param {string} query The query to process
+   * @param {string} context The context of the prompt
    * @returns {Promise<PromptResponse>} The grid state updates
    */
-  onPrompt: (context: string, query: string) => Promise<PromptResponse>;
+  onPrompt: (query: string, context: string) => Promise<PromptResponse>;
   /**
    * Called when an error occurs.
    * @param {string} error The error message
@@ -82,17 +82,17 @@ function GridToolbarPromptControl(props: GridToolbarPromptControlProps) {
   const ownerState = { classes: rootProps.classes, recording: isRecording };
   const classes = useUtilityClasses(ownerState);
 
-  const context = React.useMemo(() => {
-    const examples = allowDataSampling ? apiRef.current.unstable_collectSampleData() : undefined;
-    return apiRef.current.unstable_getPromptContext(examples);
-  }, [apiRef, allowDataSampling]);
+  const context = React.useMemo(
+    () => apiRef.current.unstable_getPromptContext(allowDataSampling),
+    [apiRef, allowDataSampling],
+  );
 
   const processPrompt = React.useCallback(() => {
     setLoading(true);
     setError(null);
     apiRef.current.setLoading(true);
 
-    onPrompt(context, query)
+    onPrompt(query, context)
       .then(apiRef.current.unstable_applyPromptResult)
       .catch((promptError) => {
         setError(apiRef.current.getLocaleText('toolbarPromptControlErrorMessage'));

packages/x-data-grid-premium/src/hooks/features/prompt/api.ts

@@ -2,7 +2,21 @@ import { PromptResponse } from './gridPromptInterfaces';
 
 type Result<T> = { ok: false; message: string } | { ok: true; data: T };
 
-export function gridDefaultPromptResolver(url: string, context: string, query: string) {
+/**
+ * Prompt resolver for the <DataGridPremium /> component.
+ * @param {string} url The URL to send the request to.
+ * @param {string} query The query to be processed.
+ * @param {string} context The prompt context containing necessary information about the current columns definition.
+ * Either use the `context` parameter of the `onPrompt` callback or the return value of the `unstable_getPromptContext()` API method.
+ * @param {string} additionalContext Optional, additional context to make the processing results more accurate.
+ * @returns {Promise<PromptResponse>} The grid state updates to be applied.
+ */
+export function gridDefaultPromptResolver(
+  url: string,
+  query: string,
+  context: string,
+  additionalContext = '',
+) {
   return fetch(url, {
     mode: 'cors',
     method: 'POST',
@@ -13,16 +27,14 @@ export function gridDefaultPromptResolver(url: string, context: string, query: s
     body: JSON.stringify({
       context,
       query,
+      additionalContext,
     }),
   })
     .then((result) => result.json())
     .then((result: Result<PromptResponse>) => {
       if (result.ok === false) {
         return Promise.reject(new Error(result.message));
       }
-      if (result.data.error) {
-        return Promise.reject(new Error(result.data.error));
-      }
       return result.data;
     });
 }

packages/x-data-grid-premium/src/hooks/features/prompt/gridPromptInterfaces.ts

@@ -32,21 +32,15 @@ export type PromptResponse = {
  * The prompt API interface that is available in the grid [[apiRef]].
  */
 export interface GridPromptApi {
-  /**
-   * Collects sample data from the grid that can be used to improve the prompt context.
-   * It takes random values for all columns but from different rows to make the samples more random.
-   * @param {number} sampleCount The number of samples (randomized rows) to collect.
-   * Default is 5. Maximum is 20.
-   * If the number of rows in the grid is less than the sample count, all rows are used.
-   * @returns {Record<string, any[]>} The sample data.
-   */
-  unstable_collectSampleData: (sampleCount?: number) => Record<string, any[]>;
   /**
    * Get the context for the prompt.
-   * @param {Record<string, any[]>} examples Column examples to be used to make the context more detailed and the response more accurate.
+   * @param {boolean} allowDataSampling Whether to use grid data as examples in the context.
+   * If true, random cell values from each column are used to make the context more detailed and the response more accurate.
+   * If false, the samples are only generated from the unstable_examples property of the column(s).
+   * Default is false.
    * @returns {string} The context for the prompt.
    */
-  unstable_getPromptContext: (examples?: Record<string, any[]>) => string;
+  unstable_getPromptContext: (allowDataSampling?: boolean) => string;
   /**
    * Use the prompt processing result to update all relevant parts of the grid state.
    * @param {PromptResponse} result The result of the prompt.

packages/x-data-grid-premium/src/hooks/features/prompt/useGridPrompt.tsx

@@ -13,47 +13,46 @@ import { getValueOptions, getVisibleRows } from '@mui/x-data-grid-pro/internals'
 import { GridPrivateApiPremium } from '../../../models/gridApiPremium';
 import { PromptResponse } from './gridPromptInterfaces';
 
-const MAX_SAMPLE_COUNT = 20;
+const DEFAULT_SAMPLE_COUNT = 5;
 
 export const useGridPrompt = (apiRef: RefObject<GridPrivateApiPremium>) => {
   const columnsLookup = gridColumnLookupSelector(apiRef);
   const columns = Object.values(columnsLookup);
   const rows = Object.values(gridRowsLookupSelector(apiRef));
 
-  const collectSampleData = React.useCallback(
-    (sampleCount = 5) => {
-      const columnExamples: Record<string, any[]> = {};
-
-      columns.forEach((column) => {
-        columnExamples[column.field] = Array.from({
-          length: Math.min(sampleCount, rows.length, MAX_SAMPLE_COUNT),
-        }).map(() => {
-          const row = rows[Math.floor(Math.random() * rows.length)];
-          if (column.valueGetter) {
-            return column.valueGetter(row[column.field] as never, row, column, apiRef);
-          }
-          return row[column.field];
-        });
+  const collectSampleData = React.useCallback(() => {
+    const columnExamples: Record<string, any[]> = {};
+
+    columns.forEach((column) => {
+      columnExamples[column.field] = Array.from({
+        length: Math.min(DEFAULT_SAMPLE_COUNT, rows.length),
+      }).map(() => {
+        const row = rows[Math.floor(Math.random() * rows.length)];
+        if (column.valueGetter) {
+          return column.valueGetter(row[column.field] as never, row, column, apiRef);
+        }
+        return row[column.field];
       });
+    });
 
-      return columnExamples;
-    },
-    [apiRef, columns, rows],
-  );
+    return columnExamples;
+  }, [apiRef, columns, rows]);
 
   const getPromptContext = React.useCallback(
-    (examples?: Record<string, any[]>) => {
+    (allowDataSampling = false) => {
+      const examples = allowDataSampling ? collectSampleData() : {};
+
       const columnsContext = columns.map((column) => ({
         field: column.field,
         description: column.description ?? null,
-        examples: examples?.[column.field] ?? column.unstable_examples ?? [],
+        examples: examples[column.field] ?? column.unstable_examples ?? [],
         type: column.type ?? 'string',
         allowedOperators: column.filterOperators?.map((operator) => operator.value) ?? [],
       }));
 
-      return `The columns are described by the following JSON:\n${JSON.stringify(columnsContext)}`;
+      return JSON.stringify(columnsContext);
     },
-    [columns],
+    [columns, collectSampleData],
   );
 
   const applyPromptResult = React.useCallback(
@@ -119,7 +118,6 @@ export const useGridPrompt = (apiRef: RefObject<GridPrivateApiPremium>) => {
   useGridApiMethod(
     apiRef,
     {
-      unstable_collectSampleData: collectSampleData,
       unstable_getPromptContext: getPromptContext,
       unstable_applyPromptResult: applyPromptResult,
     },