使用 LLM 自动化生成 d.ts

2024年9月3日

使用 LLM 自动化生成 d.ts

背景

在项目里面，有很多旧的全局的方法是没有任何注释、参数类型的，但却异常稳定，有类型推断和注释在目前开发中已经是必不可少了，但因为有上千个这样的旧的全局方法，可以考虑 AI 生成，然后人工审核来降低开发成本。

方案权衡

方案优点缺点
AST + LLM 生成 d.ts没有中间结果，不需要做额外的操作生成结果不稳定，若要重复利用比 json 解析麻烦
AST + LLM 生成 JSON利用 json 存储中间结果，可以重复利用。中间 json 解析和生成 d.ts 需要额外的步骤
AST + LLM 生成 JSON + 修改源代码 JSDoc跳转到函数定义时能直观看出类型和描述需要额外的库操作，例如 magic-string

经过时间成本和技术成本的考虑，采用 AST + LLM 生成 JSON 方案。

流程设计

提取全局代码

一般来说，全局方法是 window.XXX 或者 (NameSpace).XXX 的格式，我们需要先解析 js 文件，把它全部提取到内存数组里。这里我们采用 ast-grep 这个包来提取代码。

import { js } from '@ast-grep/napi';

export async function makeDts(source: string) {
    const ast = js.parse(source);
    const root = ast.root();

    const functions = root.findAll({
        rule: {
            pattern: '$A.$B = function($$$ARGS) { $$$ }',
        },
    })
    const allFunctionsText = functions.map(node => node.text())

    for(let func of allFunctionsText) {
        // TODO call LLM
    }
    
}

/**
 * 
 * @param {string} filepath 
 */
export function makeDtsFormFile(filepath: string = '') {
    // 1. get file content from filepath
    const fileContent = fs.readFileSync(filepath, 'utf-8');
    return makeDts(fileContent);
}

调用 LLM 获取 d.ts 返回内容

这里采用文心一言来作为例子

interface Message {
    role: string
    content: string,
    name?: string
}
export interface CompletionOptions {
    messages: Message[]
    stream?: boolean
}

export interface BotAPI {
    createCompletions: (option: CompletionOptions) => Promise<any>
}

/**
 * 
 * 文心一言接口
 * 
 */
import fetch from 'node-fetch'
// @ts-ignore
import 'dotenv/config'
import { BotAPI, CompletionOptions } from './bot'
/**
 * 使用 AK，SK 生成鉴权签名（Access Token）
 * @return string 鉴权签名信息（Access Token）
 */

let accessToken: string = ''
async function getAccessToken(): Promise<string> {
    if (accessToken) {
        return accessToken
    }
    const AK = process.env.ERNIE_BOT_AK
    const SK = process.env.ERNIE_BOT_SK
    const url = 'https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=' + AK + '&client_secret=' + SK

    const response = await fetch(url, {
        method: 'post',
    })

    const data = await response.json() as { access_token: string }

    if (data.access_token) {
        accessToken = data?.access_token
    }
    return accessToken
}

export const ernie: BotAPI = {
    createCompletions: async function (option: CompletionOptions): Promise<any> {
        const accessToken = await getAccessToken()
        const url = 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=' + accessToken
        try {
            const response = await fetch(url, {
                method: 'post',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify({
                    messages: option.messages,
                    
                })
            })
            const data = await response.json() as { result: string }

            return data.result
        } catch (error) {
            console.error(error)
            return Promise.reject(error)
        }
    }
}

import { BotAPI } from "./bots/bot";

/**
 * 
 * @param {string} source
 * @returns 
 */
export async function askForDTS(source: string = '', botAPI: BotAPI) {
    const cot = `
作为一名前端工程师的 AI 辅助，我的任务是将提取 JS 源代码中的函数类型描述，生成 json 格式的文本
        示例1:
            Q: 请返回以下 JS 函数的类型描述 json 数据格式的文本, 不需要返回其他信息
               源代码如下:
                    Site.vote = function (voteId, moduleId) {
                        const $module = $(moduleId);
                        $module.find('#vote' + voteId).click();
                    }
            A: 我一步一步地分析了这个函数，它具体描述的是根据 moduleId 和 voteId 发起投票事件，接受两个参数，第一个参数是 voteId，类型是 string 或者 number，第二个参数是 moduleId，类型是 string 或者 number，返回值是 void
               所以我返回的结果为:
                {
                    "function": "Site.vote",
                    "description": "根据 moduleId 和 voteId 发起投票事件",
                    "signature": {
                        "parameters": [
                            {
                                "name": "voteId",
                                "type": "string | number",
                                "description": "投票 id"
                            },
                            {
                                "name": "moduleId",
                                "type": "string | number",
                                "description": "模块 id"
                            }
                        ],
                        "return_type": "void",
                    }
                }
    `
    const prompt = `
        请返回类型描述的 json 数据格式的文本, 不需要返回其他信息
        源代码如下:
        ${source}

        Attention!!! 结果返回应该是一个 json 格式的文本
    `
    try {
        const completion = await botAPI.createCompletions({
            messages: [
                { role: 'system', content: cot },
                { role: "user", content: prompt }
            ],
            ...BotOptions
        });
        return completion
    } catch (error) {
        throw error
    }
}

根据 JSON 生成 d.ts

import fs from 'fs'
// @ts-ignore
import { outputFileSync } from 'fs-extra/esm'
import path from 'path';
import { js } from '@ast-grep/napi';
import MagicString from 'magic-string';

/**
 * get a json string and output to file 
 * @param json 
 * @param filePath 
 */
export function outputJson(json: string, filePath: string) {
    try {
        outputFileSync(filePath, json, 'utf-8');
    } catch (error) {
        console.error(error)
    }
}

interface Parameter {
    name: string
    type: string
    description: string
}

interface FunctionDTS {
    function: string
    description: string
    signature: {
        parameters: Array<Parameter>
        return_type: string
        return_description: string
    }
}

/**
 * 
 * @param {string} typesRoot
 * 
 */
export function jsonToDTS(typesRoot: string) {
    const jsonPaths = findDtsInTypesRoot(typesRoot) // 这里是文本操作找到所有的 json 文件

    for(let jsonPath of jsonPaths) {
        const json = fs.readFileSync(jsonPath, 'utf-8');
        const dts = transformJsonToDTS(json)

        if (dts !== '') { 
            // write file to ${outputPath}/${jsonPath name}.d.ts
            outputFileSync(jsonPath.replace('.dts.json', '.d.ts'), dts, 'utf-8')
        }
    }
}

function transformJsonToDTS(json: string): string {
    try {
        const functionDTS = JSON.parse(json) as FunctionDTS[]

        const namespaces: Record<string, FunctionDTS[]> = {}

        let dtsOutput= ''
        /**
         * create namespace
         */
        functionDTS.forEach((dts) => {
            const functionName = dts.function
            // get namespace
            const parts = functionName.split('.')
            const [namespaceName, name] = parts
            if (!name) return
            
            dts.function = name
            if (namespaceName in namespaces) {
                namespaces[namespaceName].push(dts)
            } else {
                namespaces[namespaceName] = [dts]
            }
        })

        for(let namespace in namespaces) {
            const functions: FunctionDTS[] = namespaces[namespace]
            dtsOutput += `
declare namespace ${namespace} {
${functions.map(transformDTS).join('\n')}
}
`            
        }
        return dtsOutput
    } catch (error) {
        console.error(error, `json: ${json.slice(0, 100)}`)
        return ''
    }
}
/**
* transform to 
*
* 注释
* declare function xxx(xxx: xxx): xxx
*/
function transformDTS(dTSObject: FunctionDTS): string {
    const { function: functionName, signature } = dTSObject
    const sign = signature ?? {  return_type: '', return_description: '' }
    const { return_type, return_description } = sign
    const functionSignature = transformSignature(sign)
    const functionParamDesc = transformParamDesc(sign)

    return `
/**
 * ${dTSObject.description}
${functionParamDesc}
 * @returns {${normalizeType(return_type)}} ${return_description}
 */
function ${functionName}(${functionSignature}): ${normalizeType(return_type)}
`
}

function transformSignature(signature: FunctionDTS['signature']): string {
    const { parameters = [] } = signature
    return parameters.map((parameter) => {
        const { name, type } = parameter
        return `${name}: ${normalizeType(type)}`
    }).join(', ')
}

function transformParamDesc(signature: FunctionDTS['signature']): string {
    const { parameters = [] } = signature
    return parameters.map((parameter) => {
        const { name, description } = parameter
        return ` * @param ${name} ${description}`
    }).join('\n')
}

function normalizeType(type: string): string {
    if (!type) return 'any'

    const transforms: Record<string, string> = {
        'Void': 'void',
        'Number': 'number',
        'String': 'string',
        'Boolean': 'boolean'
    }

    return transforms[type] ?? type 
}

生成代码如下

d.ts 示例

辅助生成 d.ts 后，可以看出体验会好很多

d.ts 示例