@beparallel/langchain-ts/src/lib/type-generator.ts

Extracts Langchain prompts and generates TypeScript types.

import { ChatPromptTemplate } from '@langchain/core/prompts'
import fs from 'fs'
import { toCamelCase, toPascalCase, toUpperCase } from './case.js'

export function generateTypes(prompts: ChatPromptTemplate[], outputPath: string) {
    const lines: string[] = [
        '// Auto-generated file by @beparallel/langchain-ts. Do not edit.\n',
        '// For more information, visit https://github.com/beparallel/langchain-ts\n\n',
        'import { z } from "zod";\n\n',
    ]

    lines.push('export enum PromptName {')

    for (const prompt of prompts) {
        const promptName = getPromptName(prompt)
        lines.push(`  ${toUpperCase(promptName)} = "${promptName}",`)
    }

    lines.push('}\n\n')

    for (const prompt of prompts) {
        const inputSchema = prompt.inputVariables
        const promptName = getPromptName(prompt)
        const promptDict = prompt.toJSON()
        // @ts-ignore
        const schema = promptDict.kwargs.schema_

        lines.push('\n')

        lines.push(`/***********************************************/`)
        lines.push(`// Prompt: ${promptName}`)
        lines.push(`/***********************************************/`)
        lines.push('\n')

        // Input variables
        lines.push(
            `export interface ${toPascalCase(promptName)}Variable {`,
            `  ${inputSchema.map((input) => `${input}: any`).join('\n  ')}`,
            '}',
        )

        lines.push('\n')

        // Output variables
        lines.push(jsonSchemaToType(schema, promptName))

        lines.push('\n')

        // Zod schema
        lines.push(jsonSchemaToZodType(schema, promptName))
    }

    fs.writeFileSync(outputPath, lines.join('\n'), 'utf-8')
}

const getPromptName = (prompt: ChatPromptTemplate) => {
    const promptName = String(prompt.metadata?.lc_hub_repo ?? '')
    return promptName
}

type PromptOutputSchema = {
    title?: string
    description?: string
    type: string
    properties?: Record<string, any>
    required?: string[]
    items?: any
    strict?: boolean
    additionalProperties?: boolean
}

function jsonSchemaToType(schema: PromptOutputSchema, promptName: string, rootName = 'Root'): string {
    const parseSchema = (schema: PromptOutputSchema, name: string, increment: number): string => {
        if (schema.type === 'object' && schema.properties) {
            const required = schema.required || []
            const props = Object.entries(schema.properties)
                .map(([key, value]: [string, any]) => {
                    const optional = required.includes(key) ? '' : '?'
                    return `${' '.repeat(increment + 1)}${key}${optional}: ${parseSchema(value, key, increment + 1)}`
                })
                .join('\n')
            return `{\n${' '.repeat(increment)}${props}${' '.repeat(increment)}\n}`
        } else if (schema.type === 'array' && schema.items) {
            return `${parseSchema(schema.items, name, increment + 1)}[]`
        } else if (schema.type === 'string') {
            return 'string'
        } else if (schema.type === 'number') {
            return 'number'
        } else if (schema.type === 'boolean') {
            return 'boolean'
        } else if (schema.type === 'null') {
            return 'null'
        } else {
            return 'any'
        }
    }

    if (!schema) {
        return ''
    }

    const typeDef = `${parseSchema(schema, rootName, 1)}`

    return [
        `export interface ${toPascalCase(promptName)}Output`,
        `  ${typeDef}`,
        '',
        '',
        `export interface ${toPascalCase(promptName)}OutputWrapper {`,
        `  ${schema.title || rootName}: ${toPascalCase(promptName)}Output`,
        '}',
    ].join('\n')
}

function jsonSchemaToZodType(schema: PromptOutputSchema, promptName: string, rootName = 'Root'): string {
    const parseSchema = (schema: PromptOutputSchema, name: string, increment: number): string => {
        if (schema.type === 'object' && schema.properties) {
            const required = schema.required || []
            const props = Object.entries(schema.properties)
                .map(([key, value]: [string, any]) => {
                    const zodType = parseSchema(value, key, increment + 1)
                    let prop = `${' '.repeat(increment + 1)}${key}: z.${zodType}`
                    if (value.description) {
                        prop += `.describe(\`${value.description}\`)`
                    }
                    if (!required.includes(key)) {
                        prop += '.optional()'
                    }
                    return prop
                })
                .join(',\n')
            return `object({\n${props}\n${' '.repeat(increment)}})`
        } else if (schema.type === 'array' && schema.items) {
            return `array(z.${parseSchema(schema.items, name, increment + 1)})`
        } else if (schema.type === 'string') {
            return 'string()'
        } else if (schema.type === 'number') {
            return 'number()'
        } else if (schema.type === 'boolean') {
            return 'boolean()'
        } else if (schema.type === 'null') {
            return 'null()'
        } else {
            return 'any()'
        }
    }

    if (!schema) {
        return ''
    }

    const zodSchema = `z.${parseSchema(schema, rootName, 1)}`
    return [
        `export const ${toCamelCase(promptName)}OutputSchema = ${zodSchema}`,
        '',
        `export const ${toPascalCase(promptName)}OutputWrapperSchema  = z.object({`,
        `  ${schema.title || rootName}: ${toCamelCase(promptName)}OutputSchema`,
        '})',
    ].join('\n')
}