Existing output type definitions now loaded correctly.

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

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 || ""}

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>

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;
+}