Added "add column" prompt & button in TablePopover

modified:   src/AiPopover.tsx
            Added handleGenerateColumn so that
            a column can be generated given
            a prompt.
            Added changes to the TablePopover UI
            Now extend is diveded into AddRow
            and AddColumn sections.
modified:   src/TabularDataNode.tsx
            Modified addColumns so that its safer.
            Added optional pass of rowValue to
            support generateColumn.
modified:   src/backend/ai.ts
            Added generateColumn and it's
            corresponding system message.
Cleaned up some comments and added missing commas.
This commit is contained in:
Kraft-Single 2024-12-17 14:22:38 -05:00
parent 44ad3e1c9f
commit 44a0c9ad1d
3 changed files with 229 additions and 29 deletions

View File

@ -10,10 +10,12 @@ import {
Badge,
Textarea,
Alert,
Divider,
} from "@mantine/core";
import {
autofill,
autofillTable,
generateColumn,
generateAndReplace,
AIError,
getAIFeaturesModels,
@ -246,7 +248,10 @@ export interface AIGenReplaceTablePopoverProps {
rows: TabularDataRowType[],
) => void;
// Function to add new columns
onAddColumns: (newColumns: TabularDataColType[]) => void;
onAddColumns: (
newColumns: TabularDataColType[],
rowValues?: string[] // Optional row values
) => void;
// Indicates if values are loading
areValuesLoading: boolean;
// Callback to set loading state
@ -283,6 +288,11 @@ export function AIGenReplaceTablePopover({
const [didGenerateAndReplaceTableError, setDidGenerateAndReplaceTableError] =
useState(false);
// Generate Column state
const [isGenerateColumnLoading, setIsGenerateColumnLoading] = useState(false);
const [generateColumnPrompt, setGenerateColumnPrompt] = useState("");
const [didGenerateColumnError, setDidGenerateColumnError] = useState(false);
// Check if there are any non-empty rows
const nonEmptyRows = useMemo(
() =>
@ -350,22 +360,21 @@ export function AIGenReplaceTablePopover({
setDidCommandFillError(false);
try {
// Extract columns from the values, excluding the __uid column
const tableColumns = Object.keys(values[0] || {}).filter(
(col) => col !== "__uid"
(col) => col !== "__uid",
);
// Extract rows as strings, excluding the __uid column and handling empty rows
const tableRows = values
.slice(0, -1) // Remove the last empty row
.map((row) =>
tableColumns.map((col) => row[col]?.trim() || "").join(" | ")
tableColumns.map((col) => row[col]?.trim() || "").join(" | "),
);
const tableInput = {
cols: tableColumns,
rows: tableRows
rows: tableRows,
};
// Fetch new rows from the autofillTable function
@ -396,6 +405,51 @@ export function AIGenReplaceTablePopover({
}
};
const handleGenerateColumn = async () => {
setDidGenerateColumnError(false);
setIsGenerateColumnLoading(true);
try {
// Extract columns from the values, excluding the __uid column
const tableColumns = Object.keys(values[0] || {}).filter(
(col) => col !== "__uid",
);
// Extract rows as strings, excluding the __uid column and handling empty rows
const tableRows = values
.slice(0, -1) // Remove the last empty row
.map((row) =>
tableColumns.map((col) => row[col]?.trim() || "").join(" | "),
);
const tableInput = {
cols: tableColumns,
rows: tableRows,
};
// Fetch the generated column
const generatedColumn = await generateColumn(
tableInput,
generateColumnPrompt,
aiFeaturesProvider,
apiKeys,
);
const rowValues = generatedColumn.rows;
// Append the new column to the existing columns
onAddColumns(
[{ key: `col-${values.length}`, header: generatedColumn.col }],
rowValues
);
} catch (error) {
console.error("Error generating column:", error);
setDidGenerateColumnError(true);
showAlert && showAlert("Failed to generate a new column. Please try again.");
} finally {
setIsGenerateColumnLoading(false);
}
};
const extendUI = (
<Stack>
{didCommandFillError && (
@ -403,29 +457,54 @@ export function AIGenReplaceTablePopover({
Failed to generate rows. Please try again.
</Text>
)}
<NumberInput
label="Rows to add"
mt={5}
min={1}
max={10}
value={commandFillNumber}
onChange={(num) => setCommandFillNumber(num || 1)}
/>
<div style={{ display: "flex", alignItems: "center", gap: "0.5rem" }}>
<NumberInput
label="Rows to add"
mt={5}
min={1}
max={10}
value={commandFillNumber}
onChange={(num) => setCommandFillNumber(num || 1)}
style={{ flex: 1 }}
/>
<Button
size="sm"
variant="light"
color="grape"
onClick={handleCommandFill}
disabled={!enoughRowsForSuggestions}
loading={isCommandFillLoading}
style={{ marginTop: "1.5rem", flex: 1 }}
>
Extend
</Button>
</div>
{showWarning && (
<Text size="xs" color="grape">
You may want to add more fields for better suggestions.
</Text>
)}
<Divider label="OR" labelPosition="center" />
{didGenerateColumnError && (
<Text size="xs" color="red">
Failed to generate column. Please try again.
</Text>
)}
<Textarea
label="Generate a column for..."
value={generateColumnPrompt}
onChange={(e) => setGenerateColumnPrompt(e.currentTarget.value)}
/>
<Button
size="sm"
variant="light"
color="grape"
fullWidth
onClick={handleCommandFill}
onClick={handleGenerateColumn}
disabled={!enoughRowsForSuggestions}
loading={isCommandFillLoading}
loading={isGenerateColumnLoading}
>
Extend
Add Column
</Button>
</Stack>
);

View File

@ -480,21 +480,53 @@ const TabularDataNode: React.FC<TabularDataNodeProps> = ({ data, id }) => {
tableData.map((row) => row.value || ""),
);
// This function will add the new columns to the right of the existing columns
const addColumns = (newColumns: TabularDataColType[]) => {
const updatedColumns = [...tableColumns, ...newColumns];
// Function to add new columns to the right of the existing columns (with optional row values)
const addColumns = (
newColumns: TabularDataColType[],
rowValues?: string[] // If values are passed, they will be used to populate the new columns
) => {
setTableColumns((prevColumns) => {
const updatedColumns = [...prevColumns, ...newColumns];
// Add blank values for the new column in each row
const updatedRows = tableData.map((row) => {
const updatedRow = { ...row };
newColumns.forEach((col) => {
updatedRow[col.key] = ""; // Default value for new column
setTableData((prevData) => {
let updatedRows: TabularDataRowType[] = [];
if (prevData.length > 0) {
// Update the existing rows with the new column values
updatedRows = prevData.map((row, rowIndex) => {
const updatedRow = { ...row };
newColumns.forEach((col) => {
// If rowValues, use them for the new column
updatedRow[col.key] =
rowValues && rowValues[rowIndex] !== undefined
? rowValues[rowIndex]
: ""; // Default to empty value
});
return updatedRow;
});
} else if (rowValues) {
// If no rows exist but rowValues are passed, create new rows
updatedRows = rowValues.map((value) => {
const newRow: TabularDataRowType = { __uid: uuidv4() };
newColumns.forEach((col) => {
newRow[col.key] = value || "";
});
return newRow;
});
} else {
// Create a single blank row
const blankRow: TabularDataRowType = { __uid: uuidv4() };
newColumns.forEach((col) => {
blankRow[col.key] = "";
});
updatedRows.push(blankRow);
}
return updatedRows; // Update table rows
});
return updatedRow;
});
setTableColumns(updatedColumns);
setTableData(updatedRows);
return updatedColumns; // Update table columns
});
};
// Function to add multiple rows to the table

View File

@ -114,6 +114,7 @@ function autofillSystemMessage(
/**
* Generate the system message used for autofillingTables.
* @param n number of rows to generate
* @param templateVariables list of template variables to use
*/
function autofillTableSystemMessage(
n: number,
@ -122,6 +123,18 @@ function autofillTableSystemMessage(
return `Here is a table. Generate ${n} more commands or items following the pattern. You must format your response as a markdown table with labeled columns and a divider with only the next ${n} generated commands or items of the table. ${templateVariables && templateVariables.length > 0 ? templateVariableMessage(templateVariables) : ""}`;
}
/**
* Generate the system message used for generate column.
* @param templateVariables list of template variables to use
* @param prompt description or pattern for the column content
*/
function generateColumnSystemMessage(
templateVariables?: string[],
prompt?: string,
): string {
return `Here is a table. Generate exactly one column with an appropriate header given the prompt: ${prompt} and one divider with the appropriate commands or items and the same amount of rows. Output the formatted labled markdown column, with each command or item as a new row. ${templateVariables && templateVariables.length > 0 ? templateVariableMessage(templateVariables) : ""}`;
}
/**
* Generate the system message used for generate and replace (GAR).
*/
@ -386,7 +399,7 @@ export async function autofillTable(
// Return the updated table with "n" number of rows
return {
cols: input.cols,
rows: newRows // Return the new rows generated by the LLM
rows: newRows, // Return the new rows generated by the LLM
};
} catch (error) {
console.error("Error in autofillTable:", error);
@ -396,6 +409,82 @@ export async function autofillTable(
}
}
/**
* Uses an LLM to generate one new column with data based on the pattern explained in `prompt`.
* @param prompt Description or pattern for the column content.
* @param provider The LLM provider to use (e.g., OpenAI, Bedrock).
* @param apiKeys API keys required for the LLM query.
* @returns A promise resolving to an array of strings (column values).
*/
export async function generateColumn(
tableData: { cols: string[]; rows: Row[] },
prompt: string,
provider: string,
apiKeys: Dict,
): Promise<{col: string; rows: string[]}> {
// Build a unique ID based on the input prompt
const id = JSON.stringify([prompt]);
// Encode the input table to a markdown table
const encoded = encodeTable(tableData.cols, tableData.rows);
// Extract template variables from the columns and rows
const templateVariables = [
...new Set([
...new StringTemplate(tableData.rows.join("\n")).get_vars(),
...new StringTemplate(tableData.cols.join("\n")).get_vars(),
]),
];
// History for the system message to set up LLM behavior
const history: ChatHistoryInfo[] = [
{
messages: [
{
role: "system",
content: generateColumnSystemMessage(templateVariables, prompt),
},
],
fill_history: {},
},
];
try {
// Query the LLM
const result = await queryLLM(
id,
getAIFeaturesModels(provider).small, // Use the small model
1,
encoded,
{},
history,
apiKeys,
true,
);
// Handle any errors in the response
if (result.errors && Object.keys(result.errors).length > 0) {
throw new Error(Object.values(result.errors)[0].toString());
}
console.log("LLM said: ", result.responses[0].responses[0]);
// Decode the LLM response as a column in a markdown table
const columnValues = decodeTable(result.responses[0].responses[0] as string);
return {
col:columnValues.cols[0],
rows: columnValues.rows
};
} catch (error) {
console.error("Error in generateColumn:", error);
throw new AIError(
`Failed to generate column. Details: ${(error as Error).message || error}`,
);
}
}
/**
* Uses an LLM to generate `n` new rows based on the pattern explained in `prompt`.
* @param prompt