ChainGraph CodeGen - Implementation Plan v2

Status: ✅ COMPLETE - All goals achieved!

Goals

1. ✅ Fix enum title formatting (HARM_CATEGORY → Harm Category) - DONE
2. ✅ Extract nested types as separate schemas - DONE
3. ✅ Generate plain TypeScript enums (no decorators) - DONE (for native enums)
4. ✅ Handle enum fields properly - DONE (literal unions use @PortEnum)
5. ✅ Proper @PortArray with schema references - DONE
6. ✅ Full type preservation (SafetySetting[] not Record<string, any>[]) - DONE

Note: Gemini SDK uses string literal unions, not native TypeScript enums, so @PortEnum with inline options is correct.

---

Desired Output (Complete Example)

Input: Gemini SDK Types

// From @google/genai
export interface GenerateContentConfig {
  httpOptions?: HttpOptions
  safetySettings?: SafetySetting[]
  temperature?: number
}

export interface HttpOptions {
  baseUrl?: string
  apiVersion?: string
}

export interface SafetySetting {
  method?: HarmBlockMethod
  category?: HarmCategory
  threshold?: HarmBlockThreshold
}

export enum HarmBlockMethod {
  HARM_BLOCK_METHOD_UNSPECIFIED = "HARM_BLOCK_METHOD_UNSPECIFIED",
  SEVERITY = "SEVERITY",
  PROBABILITY = "PROBABILITY"
}

export enum HarmCategory {
  HARM_CATEGORY_UNSPECIFIED = "HARM_CATEGORY_UNSPECIFIED",
  HARM_CATEGORY_HATE_SPEECH = "HARM_CATEGORY_HATE_SPEECH",
  // ... more
}

Generated Output (Perfect!)

import { ObjectSchema, PortObject, PortArray, PortNumber, PortString, PortEnumFromNative } from '@badaitech/chaingraph-types'

// ==========================================
// ENUMS (No decorators!)
// ==========================================

/** Harm block method */
export enum HarmBlockMethod {
  HARM_BLOCK_METHOD_UNSPECIFIED = "HARM_BLOCK_METHOD_UNSPECIFIED",
  SEVERITY = "SEVERITY",
  PROBABILITY = "PROBABILITY"
}

/** Harm category */
export enum HarmCategory {
  HARM_CATEGORY_UNSPECIFIED = "HARM_CATEGORY_UNSPECIFIED",
  HARM_CATEGORY_HATE_SPEECH = "HARM_CATEGORY_HATE_SPEECH",
  HARM_CATEGORY_DANGEROUS_CONTENT = "HARM_CATEGORY_DANGEROUS_CONTENT",
  // ... all enum values
}

/** Harm block threshold */
export enum HarmBlockThreshold {
  HARM_BLOCK_THRESHOLD_UNSPECIFIED = "HARM_BLOCK_THRESHOLD_UNSPECIFIED",
  BLOCK_LOW_AND_ABOVE = "BLOCK_LOW_AND_ABOVE",
  BLOCK_MEDIUM_AND_ABOVE = "BLOCK_MEDIUM_AND_ABOVE",
  BLOCK_ONLY_HIGH = "BLOCK_ONLY_HIGH",
  BLOCK_NONE = "BLOCK_NONE",
  OFF = "OFF"
}

// ==========================================
// OBJECT SCHEMAS (Dependencies first)
// ==========================================

/** HTTP options configuration */
@ObjectSchema({
  description: 'HTTP options to be used in requests',
  type: 'HttpOptions',
})
export class HttpOptions {
  /** The base URL for the AI platform service endpoint */
  @PortString({
    title: 'Base URL',
    description: 'The base URL for the AI platform service endpoint',
  })
  baseUrl?: string

  /** Specifies the version of the API to use */
  @PortString({
    title: 'API Version',
    description: 'Specifies the version of the API to use',
  })
  apiVersion?: string
}

/** Safety setting for content filtering */
@ObjectSchema({
  description: 'Safety settings for blocking harmful content',
  type: 'SafetySetting',
})
export class SafetySetting {
  /** Harm block method */
  @PortEnumFromNative(HarmBlockMethod, {
    title: 'Method',
    description: 'Determines if the harm block method uses probability or severity',
    defaultValue: HarmBlockMethod.HARM_BLOCK_METHOD_UNSPECIFIED,
  })
  method?: HarmBlockMethod  // ✅ Typed!

  /** Harm category */
  @PortEnumFromNative(HarmCategory, {
    title: 'Category',
    description: 'Harm category',
    defaultValue: HarmCategory.HARM_CATEGORY_UNSPECIFIED,
  })
  category?: HarmCategory  // ✅ Typed!

  /** Harm block threshold */
  @PortEnumFromNative(HarmBlockThreshold, {
    title: 'Threshold',
    description: 'The harm block threshold',
    defaultValue: HarmBlockThreshold.HARM_BLOCK_THRESHOLD_UNSPECIFIED,
  })
  threshold?: HarmBlockThreshold  // ✅ Typed!
}

// ==========================================
// MAIN SCHEMA
// ==========================================

/**
 * Optional model configuration parameters.
 * For more information, see Content generation parameters
 */
@ObjectSchema({
  description: 'Optional model configuration parameters',
  type: 'GenerateContentConfig',
})
export class GenerateContentConfig {
  /** HTTP request options */
  @PortObject({
    title: 'HTTP Options',
    description: 'Used to override HTTP request options',
    schema: HttpOptions,
  })
  httpOptions?: HttpOptions  // ✅ Fully typed!

  /** Safety settings for content filtering */
  @PortArray({
    title: 'Safety Settings',
    description: 'Safety settings in the request to block unsafe content',
    defaultValue: [],
    itemConfig: {
      type: 'object',
      schema: SafetySetting,  // ✅ Reference to class!
    },
    isMutable: true,
  })
  safetySettings?: SafetySetting[]  // ✅ Fully typed array!

  /**
   * Value that controls the degree of randomness in token selection.
   * Lower temperatures are good for less open-ended responses.
   */
  @PortNumber({
    title: 'Temperature',
    description: 'Value that controls the degree of randomness in token selection',
  })
  temperature?: number
}

---

Implementation Tasks

Task 1: Fix Enum Title Formatting ⚡

File: TypeMapper.ts

private formatTitle(propertyName: string): string {
  // Handle SCREAMING_SNAKE_CASE
  if (propertyName === propertyName.toUpperCase() && propertyName.includes('_')) {
    return propertyName
      .split('_')
      .map(word => word.charAt(0) + word.slice(1).toLowerCase())
      .join(' ')
  }

  // Handle camelCase/PascalCase
  const words = propertyName
    .replace(/([A-Z])/g, ' $1')
    .trim()
    .split(/[\s_-]+/)
    .filter(w => w.length > 0)

  return words
    .map(word => word.charAt(0).toUpperCase() + word.slice(1).toLowerCase())
    .join(' ')
}

Test:

it('should format SCREAMING_SNAKE_CASE', () => {
  expect(formatTitle('HARM_CATEGORY_HATE_SPEECH')).toBe('Harm Category Hate Speech')
  expect(formatTitle('BLOCK_LOW_AND_ABOVE')).toBe('Block Low And Above')
})

---

Task 2: Create TypeDependencyCollector

New File: core/TypeDependencyCollector.ts

import { Project, InterfaceDeclaration, EnumDeclaration, TypeAliasDeclaration, Type } from 'ts-morph'

export interface TypeMetadata {
  name: string
  kind: 'interface' | 'enum' | 'type-alias'
  declaration: InterfaceDeclaration | EnumDeclaration | TypeAliasDeclaration
  dependencies: string[]
}

export class TypeDependencyCollector {
  private discovered = new Map<string, TypeMetadata>()
  private queue: string[] = []
  private visited = new Set<string>()

  constructor(private project: Project) {}

  /**
   * Collect all type dependencies recursively
   */
  collect(rootTypeName: string, maxDepth: number = 5): Map<string, TypeMetadata> {
    this.queue.push({ name: rootTypeName, depth: 0 })

    while (this.queue.length > 0) {
      const { name, depth } = this.queue.shift()!

      if (this.visited.has(name) || depth > maxDepth) continue
      this.visited.add(name)

      const metadata = this.analyzeType(name)
      if (!metadata) continue

      this.discovered.set(name, metadata)

      // Add dependencies to queue
      for (const dep of metadata.dependencies) {
        this.queue.push({ name: dep, depth: depth + 1 })
      }
    }

    return this.discovered
  }

  private analyzeType(typeName: string): TypeMetadata | null {
    // Find the type in project
    for (const sourceFile of this.project.getSourceFiles()) {
      // Try interface
      const interfaceDecl = sourceFile.getInterface(typeName)
      if (interfaceDecl) {
        return {
          name: typeName,
          kind: 'interface',
          declaration: interfaceDecl,
          dependencies: this.extractInterfaceDependencies(interfaceDecl),
        }
      }

      // Try enum
      const enumDecl = sourceFile.getEnum(typeName)
      if (enumDecl) {
        return {
          name: typeName,
          kind: 'enum',
          declaration: enumDecl,
          dependencies: [],  // Enums have no dependencies
        }
      }

      // Try type alias
      const typeAlias = sourceFile.getTypeAlias(typeName)
      if (typeAlias) {
        return {
          name: typeName,
          kind: 'type-alias',
          declaration: typeAlias,
          dependencies: this.extractTypeAliasDependencies(typeAlias),
        }
      }
    }

    return null
  }

  private extractInterfaceDependencies(decl: InterfaceDeclaration): string[] {
    const deps: string[] = []

    for (const prop of decl.getProperties()) {
      const propType = prop.getType()
      deps.push(...this.extractTypeNamesFromType(propType))
    }

    return [...new Set(deps)]  // Remove duplicates
  }

  private extractTypeNamesFromType(type: Type): string[] {
    const names: string[] = []

    // Get direct type name
    const typeName = this.getTypeName(type)
    if (typeName && this.isUserDefinedType(typeName)) {
      names.push(typeName)
    }

    // Handle arrays
    if (type.isArray()) {
      const elementType = type.getArrayElementType()
      if (elementType) {
        names.push(...this.extractTypeNamesFromType(elementType))
      }
    }

    // Handle unions
    if (type.isUnion()) {
      const unionTypes = type.getUnionTypes()
      for (const ut of unionTypes) {
        if (!ut.isUndefined() && !ut.isNull()) {
          names.push(...this.extractTypeNamesFromType(ut))
        }
      }
    }

    return names
  }

  private getTypeName(type: Type): string | null {
    const symbol = type.getSymbol()
    if (!symbol) return null

    return symbol.getName()
  }

  private isUserDefinedType(typeName: string): boolean {
    const builtIns = [
      'string', 'number', 'boolean', 'any', 'unknown', 'void', 'never',
      'Array', 'Record', 'Partial', 'Required', 'Pick', 'Omit',
      'AbortSignal', 'Promise', 'Date', 'RegExp',
    ]
    return !builtIns.includes(typeName)
  }
}

---

Task 3: Update CodeGenerator for Multi-Schema

File: CodeGenerator.ts

async generate(): Promise<string> {
  // 1. Find and add all relevant .d.ts files
  const dtsPath = this.resolveDtsPath(this.options.library)
  this.addDtsFilesToProject(dtsPath)

  // 2. Collect type dependencies
  const collector = new TypeDependencyCollector(this.project)
  const allTypes = collector.collect(
    this.options.typeName,
    this.options.maxDepth || 5
  )

  // 3. Generate each type
  const generatedSchemas: GeneratedSchema[] = []

  for (const [typeName, metadata] of allTypes) {
    const code = this.generateSingleType(metadata, allTypes)
    generatedSchemas.push({
      typeName,
      kind: metadata.kind,
      code,
      dependencies: metadata.dependencies,
    })
  }

  // 4. Sort schemas (dependencies first)
  const sorted = this.topologicalSort(generatedSchemas)

  // 5. Combine into single output
  return this.combineSchemas(sorted)
}

private generateSingleType(
  metadata: TypeMetadata,
  allTypes: Map<string, TypeMetadata>
): string {
  switch (metadata.kind) {
    case 'enum':
      return this.templateEngine.generateEnum(metadata)

    case 'interface':
      return this.templateEngine.generateObjectSchema(metadata, allTypes)

    default:
      // Skip type aliases for now
      return ''
  }
}

---

Task 4: Update TypeMapper for Type Name Extraction

File: TypeMapper.ts

class TypeMapper {
  // Track discovered types for dependency collection
  private discoveredTypes = new Set<string>()

  mapType(type: Type, propertyName: string, context?: {
    collectDependencies?: boolean
  }): IPortConfig {
    // ... existing logic

    // When mapping objects, extract the type name
    if (type.isObject() && !type.isArray()) {
      const typeName = this.extractTypeName(type)

      if (typeName && context?.collectDependencies) {
        this.discoveredTypes.add(typeName)

        // Return a reference config
        return {
          type: 'object',
          title: this.formatTitle(propertyName),
          schema: typeName,  // ✅ Type name, not inline schema
        }
      }

      // Fallback to inline schema
      return this.mapObjectInline(type, propertyName)
    }
  }

  private extractTypeName(type: Type): string | null {
    const symbol = type.getSymbol()
    if (!symbol) return null

    const declarations = symbol.getDeclarations()
    if (declarations.length === 0) return null

    const decl = declarations[0]

    // Check if it's a named interface or type alias
    if (Node.isInterfaceDeclaration(decl) || Node.isTypeAliasDeclaration(decl)) {
      return decl.getName()
    }

    return null
  }

  // For enum types
  mapEnumField(enumType: Type, propertyName: string): IPortConfig {
    const enumName = this.extractTypeName(enumType)

    if (enumName) {
      this.discoveredTypes.add(enumName)

      // Get first enum value as default
      const enumDecl = this.findEnumDeclaration(enumName)
      const firstMember = enumDecl?.getMembers()[0]
      const defaultValue = firstMember ? `${enumName}.${firstMember.getName()}` : undefined

      return {
        type: 'enum-native',  // New type!
        title: this.formatTitle(propertyName),
        enumType: enumName,
        defaultValue,
      }
    }

    // Fallback to inline enum options
    return this.mapInlineEnum(enumType, propertyName)
  }
}

---

Task 5: Update TemplateEngine

File: TemplateEngine.ts

New Methods:

/**
 * Generate plain TypeScript enum (NO @ObjectSchema!)
 */
generateEnum(metadata: TypeMetadata): string {
  const enumDecl = metadata.declaration as EnumDeclaration
  const members = enumDecl.getMembers()
  const jsDocs = enumDecl.getJsDocs()
  const description = this.extractDescription(jsDocs)

  const comment = description ? `/** ${description} */\n` : ''

  const memberCode = members.map(member => {
    const name = member.getName()
    const value = member.getValue()
    const memberDoc = member.getJsDocs()[0]?.getComment()
    const memberComment = memberDoc ? `  /** ${memberDoc} */\n  ` : '  '

    return `${memberComment}${name} = "${value}"`
  }).join(',\n')

  return `${comment}export enum ${metadata.name} {\n${memberCode}\n}`
}

/**
 * Generate object schema with type references
 */
generateObjectSchema(
  metadata: TypeMetadata,
  allTypes: Map<string, TypeMetadata>
): string {
  const interfaceDecl = metadata.declaration as InterfaceDeclaration
  const properties = interfaceDecl.getProperties()

  const propertyCode = properties.map(prop => {
    const name = prop.getName()
    const type = prop.getType()
    const optional = prop.hasQuestionToken()

    // Extract property metadata
    const portConfig = this.typeMapper.mapType(type, name, {
      collectDependencies: true,
      allTypes
    })

    // Special handling for enum types
    if (this.isEnumType(type, allTypes)) {
      return this.generateEnumProperty(prop, type, allTypes)
    }

    // Special handling for object types
    if (this.isObjectType(type, allTypes)) {
      return this.generateObjectProperty(prop, type, allTypes)
    }

    // Special handling for array types
    if (type.isArray()) {
      return this.generateArrayProperty(prop, type, allTypes)
    }

    // Default property generation
    return this.generatePropertyDeclaration({
      name,
      portConfig,
      description: this.extractDescription(prop.getJsDocs()),
      optional,
    })
  }).join('\n\n')

  return `
@ObjectSchema({
  description: '${this.extractDescription(interfaceDecl.getJsDocs())}',
  type: '${metadata.name}',
})
export class ${metadata.name} {
${this.indent(propertyCode, 2)}
}
  `.trim()
}

/**
 * Generate enum property using @PortEnumFromNative
 */
private generateEnumProperty(
  prop: PropertySignature,
  type: Type,
  allTypes: Map<string, TypeMetadata>
): string {
  const name = prop.getName()
  const optional = prop.hasQuestionToken()
  const enumTypeName = this.extractTypeName(type)
  const description = this.extractDescription(prop.getJsDocs())

  // Get default value (first enum member)
  const enumMeta = allTypes.get(enumTypeName)
  const firstMember = (enumMeta?.declaration as EnumDeclaration)?.getMembers()[0]
  const defaultValue = firstMember ? `${enumTypeName}.${firstMember.getName()}` : undefined

  const comment = description ? this.formatJSDocComment(description, 2) : ''

  return `${comment}  @PortEnumFromNative(${enumTypeName}, {
    title: '${this.formatTitle(name)}',
    description: '${this.escapeString(description || '')}',
    ${defaultValue ? `defaultValue: ${defaultValue},` : ''}
  })
  ${name}${optional ? '?' : ''}: ${enumTypeName}`
}

/**
 * Generate array property with schema reference
 */
private generateArrayProperty(
  prop: PropertySignature,
  type: Type,
  allTypes: Map<string, TypeMetadata>
): string {
  const name = prop.getName()
  const optional = prop.hasQuestionToken()
  const elementType = type.getArrayElementType()!
  const description = this.extractDescription(prop.getJsDocs())

  const comment = description ? this.formatJSDocComment(description, 2) : ''

  // Check if element is an object type
  const elementTypeName = this.extractTypeName(elementType)

  if (elementTypeName && allTypes.has(elementTypeName)) {
    // Reference to generated schema
    return `${comment}  @PortArray({
    title: '${this.formatTitle(name)}',
    description: '${this.escapeString(description || '')}',
    defaultValue: [],
    itemConfig: {
      type: 'object',
      schema: ${elementTypeName},  // ✅ Class reference!
    },
    isMutable: true,
  })
  ${name}${optional ? '?' : ''}: ${elementTypeName}[]`  // ✅ Typed array!
  }

  // Fallback to primitive array
  const itemConfig = this.typeMapper.mapType(elementType, `${name}_item`)

  return `${comment}  @PortArray({
    title: '${this.formatTitle(name)}',
    description: '${this.escapeString(description || '')}',
    defaultValue: [],
    itemConfig: ${JSON.stringify(itemConfig)},
    isMutable: true,
  })
  ${name}${optional ? '?' : ''}: ${this.getTypeScriptType(type)}`
}

---

Task 6: Combine Schemas

File: TemplateEngine.ts

combineSchemas(schemas: GeneratedSchema[]): string {
  // Collect used decorators
  const usedDecorators = new Set<string>()
  usedDecorators.add('ObjectSchema')

  for (const schema of schemas) {
    if (schema.kind === 'object-schema') {
      // Scan code for used decorators
      if (schema.code.includes('@PortEnumFromNative')) usedDecorators.add('PortEnumFromNative')
      if (schema.code.includes('@PortNumber')) usedDecorators.add('PortNumber')
      if (schema.code.includes('@PortString')) usedDecorators.add('PortString')
      if (schema.code.includes('@PortArray')) usedDecorators.add('PortArray')
      if (schema.code.includes('@PortObject')) usedDecorators.add('PortObject')
      if (schema.code.includes('@PortBoolean')) usedDecorators.add('PortBoolean')
    }
  }

  const imports = `import { ${Array.from(usedDecorators).sort().join(', ')} } from '@badaitech/chaingraph-types'`

  // Group schemas by kind
  const enums = schemas.filter(s => s.kind === 'enum')
  const objects = schemas.filter(s => s.kind === 'object-schema')

  const enumsSection = enums.length > 0 ? `
// ==========================================
// ENUMS
// ==========================================

${enums.map(s => s.code).join('\n\n')}
` : ''

  const objectsSection = objects.length > 0 ? `
// ==========================================
// OBJECT SCHEMAS
// ==========================================

${objects.map(s => s.code).join('\n\n')}
` : ''

  return `${imports}\n${enumsSection}${objectsSection}`.trim()
}

---

Summary

This plan will transform the generator to produce clean, typed, reusable schemas that follow ChainGraph patterns exactly:

✅ Enums as plain TypeScript (no decorators) ✅ @PortEnumFromNative for enum fields ✅ @PortArray with schema references for object arrays ✅ Separate schemas for each type ✅ Full type preservation ✅ Proper dependency ordering

Ready to implement?