feat(ui): update node schema parsing

simplified logic thanks to backend changes

invokeai/frontend/web/src/features/nodes/types/typeGuards.ts

@@ -1,9 +0,0 @@
-import { OpenAPIV3 } from 'openapi-types';
-
-export const isReferenceObject = (
-  obj: OpenAPIV3.ReferenceObject | OpenAPIV3.SchemaObject
-): obj is OpenAPIV3.ReferenceObject => '$ref' in obj;
-
-export const isSchemaObject = (
-  obj: OpenAPIV3.ReferenceObject | OpenAPIV3.SchemaObject
-): obj is OpenAPIV3.SchemaObject => !('$ref' in obj);

invokeai/frontend/web/src/features/nodes/types/types.ts

@@ -518,6 +518,18 @@ export type InvocationBaseSchemaObject = Omit<
 > &
   InvocationSchemaExtra;
 
+export type InvocationOutputSchemaObject = Omit<
+  OpenAPIV3.SchemaObject,
+  'properties'
+> &
+  OpenAPIV3.SchemaObject['properties'] & {
+    type: Omit<OpenAPIV3.SchemaObject, 'default'> & {
+      default: AnyInvocationType;
+    };
+  } & {
+    class: 'output';
+  };
+
 export type InvocationFieldSchema = OpenAPIV3.SchemaObject & _InputField;
 
 export interface ArraySchemaObject extends InvocationBaseSchemaObject {
@@ -528,11 +540,30 @@ export interface NonArraySchemaObject extends InvocationBaseSchemaObject {
   type?: OpenAPIV3.NonArraySchemaObjectType;
 }
 
-export type InvocationSchemaObject = ArraySchemaObject | NonArraySchemaObject;
+export type InvocationSchemaObject = (
+  | ArraySchemaObject
+  | NonArraySchemaObject
+) & { class: 'invocation' };
+
+export const isSchemaObject = (
+  obj: OpenAPIV3.ReferenceObject | OpenAPIV3.SchemaObject
+): obj is OpenAPIV3.SchemaObject => !('$ref' in obj);
 
 export const isInvocationSchemaObject = (
-  obj: OpenAPIV3.ReferenceObject | InvocationSchemaObject
+  obj:
-): obj is InvocationSchemaObject => !('$ref' in obj);
+    | OpenAPIV3.ReferenceObject
+    | OpenAPIV3.SchemaObject
+    | InvocationSchemaObject
+): obj is InvocationSchemaObject =>
+  'class' in obj && obj.class === 'invocation';
+
+export const isInvocationOutputSchemaObject = (
+  obj:
+    | OpenAPIV3.ReferenceObject
+    | OpenAPIV3.SchemaObject
+    | InvocationOutputSchemaObject
+): obj is InvocationOutputSchemaObject =>
+  'class' in obj && obj.class === 'output';
 
 export const isInvocationFieldSchema = (
   obj: OpenAPIV3.ReferenceObject | OpenAPIV3.SchemaObject

invokeai/frontend/web/src/features/nodes/util/fieldTemplateBuilders.ts

@@ -1,8 +1,4 @@
-import { logger } from 'app/logging/logger';
-import { parseify } from 'common/util/serialize';
-import { reduce } from 'lodash-es';
 import { OpenAPIV3 } from 'openapi-types';
-import { isSchemaObject } from '../types/typeGuards';
 import {
   BooleanInputFieldTemplate,
   ClipInputFieldTemplate,
@@ -24,7 +20,6 @@ import {
   LatentsInputFieldTemplate,
   LoRAModelInputFieldTemplate,
   MainModelInputFieldTemplate,
-  OutputFieldTemplate,
   SDXLMainModelInputFieldTemplate,
   SDXLRefinerModelInputFieldTemplate,
   SchedulerInputFieldTemplate,
@@ -33,7 +28,6 @@ import {
   VaeInputFieldTemplate,
   VaeModelInputFieldTemplate,
   isFieldType,
-  isInvocationFieldSchema,
 } from '../types/types';
 
 export type BaseFieldProperties = 'name' | 'title' | 'description';
@@ -628,66 +622,3 @@ export const buildInputFieldTemplate = (
   }
   return;
 };
-
-/**
- * Builds invocation output fields from an invocation's output reference object.
- * @param openAPI The OpenAPI schema
- * @param refObject The output reference object
- * @returns A record of outputs
- */
-export const buildOutputFieldTemplates = (
-  refObject: OpenAPIV3.ReferenceObject,
-  openAPI: OpenAPIV3.Document
-): Record<string, OutputFieldTemplate> => {
-  // extract output schema name from ref
-  const outputSchemaName = refObject.$ref.split('/').slice(-1)[0];
-
-  if (!outputSchemaName) {
-    logger('nodes').error(
-      { refObject: parseify(refObject) },
-      'No output schema name found in ref object'
-    );
-    throw 'No output schema name found in ref object';
-  }
-
-  // get the output schema itself
-  const outputSchema = openAPI.components?.schemas?.[outputSchemaName];
-  if (!outputSchema) {
-    logger('nodes').error({ outputSchemaName }, 'Output schema not found');
-    throw 'Output schema not found';
-  }
-
-  // console.log('output', outputSchema);
-  if (isSchemaObject(outputSchema)) {
-    // console.log('isSchemaObject');
-    const outputFields = reduce(
-      outputSchema.properties as OpenAPIV3.SchemaObject,
-      (outputsAccumulator, property, propertyName) => {
-        if (
-          !['type', 'id'].includes(propertyName) &&
-          !['object'].includes(property.type) && // TODO: handle objects?
-          isInvocationFieldSchema(property)
-        ) {
-          const fieldType = getFieldType(property);
-          // console.log('output fieldType', fieldType);
-          outputsAccumulator[propertyName] = {
-            fieldKind: 'output',
-            name: propertyName,
-            title: property.title ?? '',
-            description: property.description ?? '',
-            type: fieldType,
-          };
-        } else {
-          // console.warn('Unhandled OUTPUT property', property);
-        }
-
-        return outputsAccumulator;
-      },
-      {} as Record<string, OutputFieldTemplate>
-    );
-
-    return outputFields;
-  }
-
-  return {};
-};

invokeai/frontend/web/src/features/nodes/util/parseSchema.ts

@@ -1,83 +1,123 @@
-import { filter, reduce } from 'lodash-es';
+import { logger } from 'app/logging/logger';
+import { parseify } from 'common/util/serialize';
+import { reduce } from 'lodash-es';
 import { OpenAPIV3 } from 'openapi-types';
+import { AnyInvocationType } from 'services/events/types';
 import {
   InputFieldTemplate,
   InvocationSchemaObject,
   InvocationTemplate,
+  OutputFieldTemplate,
   isInvocationFieldSchema,
+  isInvocationOutputSchemaObject,
   isInvocationSchemaObject,
 } from '../types/types';
-import {
+import { buildInputFieldTemplate, getFieldType } from './fieldTemplateBuilders';
-  buildInputFieldTemplate,
-  buildOutputFieldTemplates,
-} from './fieldTemplateBuilders';
 
 const RESERVED_FIELD_NAMES = ['id', 'type', 'metadata'];
 
-const invocationDenylist = [
+const invocationDenylist: AnyInvocationType[] = [
-  'Graph',
+  'graph',
-  'InvocationMeta',
+  'metadata_accumulator',
-  'MetadataAccumulatorInvocation',
 ];
 
+const isNotInDenylist = (schema: InvocationSchemaObject) =>
+  !invocationDenylist.includes(schema.properties.type.default);
+
 export const parseSchema = (
   openAPI: OpenAPIV3.Document
 ): Record<string, InvocationTemplate> => {
-  const filteredSchemas = filter(
+  const filteredSchemas = Object.values(openAPI.components?.schemas ?? {})
-    openAPI.components?.schemas,
+    .filter(isInvocationSchemaObject)
-    (schema, key) =>
+    .filter(isNotInDenylist);
-      key.includes('Invocation') &&
-      !key.includes('InvocationOutput') &&
-      !invocationDenylist.some((denylistItem) => key.includes(denylistItem))
-  ) as (OpenAPIV3.ReferenceObject | InvocationSchemaObject)[];
 
   const invocations = filteredSchemas.reduce<
     Record<string, InvocationTemplate>
   >((acc, schema) => {
-    if (isInvocationSchemaObject(schema)) {
+    const type = schema.properties.type.default;
-      const type = schema.properties.type.default;
+    const title = schema.title.replace('Invocation', '');
-      const title = schema.title.replace('Invocation', '');
+    const tags = schema.tags ?? [];
-      const tags = schema.tags ?? [];
+    const description = schema.description ?? '';
-      const description = schema.description ?? '';
 
-      const inputs = reduce(
+    const inputs = reduce(
-        schema.properties,
+      schema.properties,
-        (inputsAccumulator, property, propertyName) => {
+      (inputsAccumulator, property, propertyName) => {
-          if (
+        if (
-            !RESERVED_FIELD_NAMES.includes(propertyName) &&
+          !RESERVED_FIELD_NAMES.includes(propertyName) &&
-            isInvocationFieldSchema(property) &&
+          isInvocationFieldSchema(property) &&
-            !property.ui_hidden
+          !property.ui_hidden
-          ) {
+        ) {
-            const field = buildInputFieldTemplate(
+          const field = buildInputFieldTemplate(schema, property, propertyName);
-              schema,
-              property,
-              propertyName
-            );
 
-            if (field) {
+          if (field) {
-              inputsAccumulator[propertyName] = field;
+            inputsAccumulator[propertyName] = field;
-            }
           }
-          return inputsAccumulator;
+        }
-        },
+        return inputsAccumulator;
-        {} as Record<string, InputFieldTemplate>
+      },
+      {} as Record<string, InputFieldTemplate>
+    );
+
+    const outputSchemaName = schema.output.$ref.split('/').pop();
+
+    if (!outputSchemaName) {
+      logger('nodes').error(
+        { outputRefObject: parseify(schema.output) },
+        'No output schema name found in ref object'
       );
-
+      throw 'No output schema name found in ref object';
-      const rawOutput = (schema as InvocationSchemaObject).output;
-      const outputs = buildOutputFieldTemplates(rawOutput, openAPI);
-
-      const invocation: InvocationTemplate = {
-        title,
-        type,
-        tags,
-        description,
-        inputs,
-        outputs,
-      };
-
-      Object.assign(acc, { [type]: invocation });
     }
 
+    const outputSchema = openAPI.components?.schemas?.[outputSchemaName];
+    if (!outputSchema) {
+      logger('nodes').error({ outputSchemaName }, 'Output schema not found');
+      throw 'Output schema not found';
+    }
+
+    if (!isInvocationOutputSchemaObject(outputSchema)) {
+      logger('nodes').error(
+        { outputSchema: parseify(outputSchema) },
+        'Invalid output schema'
+      );
+      throw 'Invalid output schema';
+    }
+
+    const outputs = reduce(
+      outputSchema.properties as OpenAPIV3.SchemaObject,
+      (outputsAccumulator, property, propertyName) => {
+        if (
+          !['type', 'id'].includes(propertyName) &&
+          !['object'].includes(property.type) && // TODO: handle objects?
+          isInvocationFieldSchema(property)
+        ) {
+          const fieldType = getFieldType(property);
+          outputsAccumulator[propertyName] = {
+            fieldKind: 'output',
+            name: propertyName,
+            title: property.title ?? '',
+            description: property.description ?? '',
+            type: fieldType,
+          };
+        } else {
+          logger('nodes').warn({ property }, 'Unhandled output property');
+        }
+
+        return outputsAccumulator;
+      },
+      {} as Record<string, OutputFieldTemplate>
+    );
+
+    const invocation: InvocationTemplate = {
+      title,
+      type,
+      tags,
+      description,
+      inputs,
+      outputs,
+    };
+
+    Object.assign(acc, { [type]: invocation });
+
     return acc;
   }, {});