From 44a0c9ad1d5352e31a1323fd95b6dd5d07781ca7 Mon Sep 17 00:00:00 2001
From: Kraft-Single <cassandre.hamel.1@umontreal.ca>
Date: Tue, 17 Dec 2024 14:22:38 -0500
Subject: [PATCH] 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.
---
 chainforge/react-server/src/AiPopover.tsx     | 111 +++++++++++++++---
 .../react-server/src/TabularDataNode.tsx      |  56 +++++++--
 chainforge/react-server/src/backend/ai.ts     |  91 +++++++++++++-
 3 files changed, 229 insertions(+), 29 deletions(-)

diff --git a/chainforge/react-server/src/AiPopover.tsx b/chainforge/react-server/src/AiPopover.tsx
index 36d179c..934f3fd 100644
--- a/chainforge/react-server/src/AiPopover.tsx
+++ b/chainforge/react-server/src/AiPopover.tsx
@@ -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>
   );
diff --git a/chainforge/react-server/src/TabularDataNode.tsx b/chainforge/react-server/src/TabularDataNode.tsx
index 0083ad6..cbdef57 100644
--- a/chainforge/react-server/src/TabularDataNode.tsx
+++ b/chainforge/react-server/src/TabularDataNode.tsx
@@ -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
diff --git a/chainforge/react-server/src/backend/ai.ts b/chainforge/react-server/src/backend/ai.ts
index 1290dae..0bdf6d4 100644
--- a/chainforge/react-server/src/backend/ai.ts
+++ b/chainforge/react-server/src/backend/ai.ts
@@ -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