From 6807aa8ea2130954d0f8961aa0895d336551e2cb Mon Sep 17 00:00:00 2001
From: JSv4 <scrudato@umich.edu>
Date: Wed, 6 Nov 2024 03:41:34 -0800
Subject: [PATCH] Existing output type definitions now loaded correctly.

---
 .../components/widgets/ModelFieldBuilder.tsx  |  2 +-
 .../widgets/modals/CreateColumnModal.tsx      | 33 +++++++++----------
 .../modals/sections/OutputTypeSection.tsx     |  6 ++--
 frontend/src/utils/parseOutputType.ts         | 23 +++++++++++++
 4 files changed, 44 insertions(+), 20 deletions(-)

diff --git a/frontend/src/components/widgets/ModelFieldBuilder.tsx b/frontend/src/components/widgets/ModelFieldBuilder.tsx
index c762c465..57015f67 100644
--- a/frontend/src/components/widgets/ModelFieldBuilder.tsx
+++ b/frontend/src/components/widgets/ModelFieldBuilder.tsx
@@ -3,7 +3,7 @@ import { Button, Form, Grid } from "semantic-ui-react";
 import { motion, AnimatePresence } from "framer-motion";
 import styled from "styled-components";
 
-interface FieldType {
+export interface FieldType {
   fieldName: string;
   fieldType: string;
   id: string; // Added for stable animations
diff --git a/frontend/src/components/widgets/modals/CreateColumnModal.tsx b/frontend/src/components/widgets/modals/CreateColumnModal.tsx
index 1588e476..b0ddddef 100644
--- a/frontend/src/components/widgets/modals/CreateColumnModal.tsx
+++ b/frontend/src/components/widgets/modals/CreateColumnModal.tsx
@@ -7,6 +7,8 @@ import { AdvancedOptionsSection } from "./sections/AdvancedOptionsSection";
 import { LooseObject } from "../../types";
 import styled from "styled-components";
 import { ColumnType } from "../../../types/graphql-api";
+import { parsePydanticModel } from "../../../utils/parseOutputType";
+import { FieldType } from "../ModelFieldBuilder";
 
 interface CreateColumnModalProps {
   open: boolean;
@@ -15,11 +17,6 @@ interface CreateColumnModalProps {
   onSubmit: (data: any) => void;
 }
 
-interface FieldType {
-  fieldName: string;
-  fieldType: string;
-}
-
 interface RequiredFields {
   query: string;
   primitiveType?: string;
@@ -55,14 +52,10 @@ export const CreateColumnModal: React.FC<CreateColumnModalProps> = ({
   );
   const [isSubmitting, setIsSubmitting] = useState(false);
 
-  // Determine if the existing_column.outputType is a primitive type
-  const isPrimitiveType = ["str", "int", "float", "bool"].includes(
-    existing_column?.outputType || ""
-  );
-
-  const [outputTypeOption, setOutputTypeOption] = useState<string>(
-    isPrimitiveType ? "primitive" : "custom"
-  );
+  const [outputTypeOption, setOutputTypeOption] = useState<string>("primitive");
+  const [primitiveType, setPrimitiveType] = useState<string>("str");
+  const [extractIsList, setExtractIsList] = useState<boolean>(false);
+  const [initialFields, setInitialFields] = useState<FieldType[]>([]);
 
   useEffect(() => {
     if (existing_column) {
@@ -71,9 +64,15 @@ export const CreateColumnModal: React.FC<CreateColumnModalProps> = ({
         existing_column.outputType || ""
       );
       setOutputTypeOption(isPrimitiveType ? "primitive" : "custom");
+      setPrimitiveType(existing_column.outputType);
+      setExtractIsList(Boolean(existing_column.extractIsList));
+      setInitialFields(parsePydanticModel(existing_column.outputType));
     } else {
       setFormData({});
       setOutputTypeOption("primitive");
+      setPrimitiveType("str");
+      setExtractIsList(false);
+      setInitialFields([]);
     }
   }, [existing_column]);
 
@@ -94,10 +93,9 @@ export const CreateColumnModal: React.FC<CreateColumnModalProps> = ({
     data: any
   ) => {
     setOutputTypeOption(data.value);
-    // Reset outputType in formData when outputTypeOption changes
     setFormData((prev) => ({
       ...prev,
-      outputType: data.value === "primitive" ? "" : prev.outputType,
+      outputType: data.value === "primitive" ? primitiveType : "",
     }));
   };
 
@@ -146,11 +144,12 @@ export const CreateColumnModal: React.FC<CreateColumnModalProps> = ({
             />
             <OutputTypeSection
               outputTypeOption={outputTypeOption}
-              extractIsList={formData.extractIsList || false}
-              primitiveType={formData.primitiveType || ""}
+              extractIsList={extractIsList}
+              primitiveType={primitiveType}
               handleOutputTypeChange={handleOutputTypeChange}
               handleChange={handleChange}
               setFormData={setFormData}
+              initialFields={initialFields}
             />
             <ExtractionConfigSection
               query={formData.query || ""}
diff --git a/frontend/src/components/widgets/modals/sections/OutputTypeSection.tsx b/frontend/src/components/widgets/modals/sections/OutputTypeSection.tsx
index 005c2925..0a0843e5 100644
--- a/frontend/src/components/widgets/modals/sections/OutputTypeSection.tsx
+++ b/frontend/src/components/widgets/modals/sections/OutputTypeSection.tsx
@@ -1,7 +1,7 @@
 import React from "react";
 import { Grid, Form } from "semantic-ui-react";
 import { FormSection, SectionTitle, StyledFormField } from "../styled";
-import { ModelFieldBuilder } from "../../ModelFieldBuilder";
+import { FieldType, ModelFieldBuilder } from "../../ModelFieldBuilder";
 
 interface OutputTypeSectionProps {
   outputTypeOption: string;
@@ -19,6 +19,7 @@ interface OutputTypeSectionProps {
   setFormData: (
     updater: (prev: Record<string, any>) => Record<string, any>
   ) => void;
+  initialFields?: FieldType[];
 }
 
 /**
@@ -47,6 +48,7 @@ export const OutputTypeSection: React.FC<OutputTypeSectionProps> = ({
   handleOutputTypeChange,
   handleChange,
   setFormData,
+  initialFields = [],
 }) => {
   const handleFieldsChange = (fields: any[]) => {
     setFormData((prev) => ({
@@ -126,7 +128,7 @@ export const OutputTypeSection: React.FC<OutputTypeSectionProps> = ({
             <Grid.Column width={16}>
               <ModelFieldBuilder
                 onFieldsChange={handleFieldsChange}
-                initialFields={[]}
+                initialFields={initialFields}
               />
             </Grid.Column>
           </Grid.Row>
diff --git a/frontend/src/utils/parseOutputType.ts b/frontend/src/utils/parseOutputType.ts
index 878e4ae4..87cc6419 100644
--- a/frontend/src/utils/parseOutputType.ts
+++ b/frontend/src/utils/parseOutputType.ts
@@ -1,4 +1,5 @@
 import { JSONSchema7, JSONSchema7TypeName } from "json-schema";
+import { FieldType } from "../components/widgets/ModelFieldBuilder";
 
 /**
  * Parse a string representation of a Python Pydantic model or a primitive type
@@ -72,3 +73,25 @@ export function parseOutputType(outputType: string): JSONSchema7 {
 
   throw new Error(`Invalid model or primitive type: "${outputType}"`);
 }
+
+/**
+ * Parses a Pydantic model string and extracts field definitions.
+ *
+ * @param modelStr - The Pydantic model as a string.
+ * @returns An array of field definitions.
+ */
+export function parsePydanticModel(modelStr: string): FieldType[] {
+  const lines = modelStr.split("\n").map((line) => line.trim());
+  const fieldLines = lines.filter(
+    (line) => line && line.includes(":") && !line.startsWith("class")
+  );
+  const fields: FieldType[] = fieldLines.map((line) => {
+    const [fieldName, fieldType] = line.split(":").map((part) => part.trim());
+    return {
+      id: Math.random().toString(),
+      fieldName,
+      fieldType,
+    };
+  });
+  return fields;
+}