windmill/frontend/src/lib/components/codeGen/lib.ts

import { OpenAI } from 'openai'
import { OpenAPI } from '../../gen/core/OpenAPI'
import { ResourceService, Script, WorkspaceService } from '../../gen'

import { existsOpenaiResourcePath, workspaceStore } from '$lib/stores'
import { formatResourceTypes } from './utils'

import { EDIT_CONFIG, FIX_CONFIG, GEN_CONFIG } from './prompts'

const WARNING_MSG =
	'WARNING: this code was generated by OpenAI, it might not be accurate. Use at your own risk.'

const COMMENT_TYPES = {
	[Script.language.PYTHON3]: '#',
	[Script.language.DENO]: '//',
	[Script.language.GO]: '//',
	[Script.language.BASH]: '#',
	[Script.language.POSTGRESQL]: '--',
	[Script.language.MYSQL]: '--',
	[Script.language.NATIVETS]: '//',
	[Script.language.BUN]: '//',
	frontend: '//'
}

export const SUPPORTED_LANGUAGES = new Set(Object.keys(GEN_CONFIG.prompts))

let workspace: string | undefined = undefined

workspaceStore.subscribe(async (value) => {
	workspace = value
	if (workspace) {
		try {
			existsOpenaiResourcePath.set(await WorkspaceService.existsOpenaiResourcePath({ workspace }))
		} catch (err) {
			existsOpenaiResourcePath.set(false)
			console.error('Could not get if OpenAI resource exists')
		}
	}
})

interface BaseOptions {
	language: Script.language | 'frontend'
	dbSchema?: object
}

interface ScriptGenerationOptions extends BaseOptions {
	description: string
	dbSchema?: object
}

interface EditScriptOptions extends ScriptGenerationOptions {
	selectedCode: string
}

interface FixScriptOpions extends BaseOptions {
	code: string
	error: string
}

async function addResourceTypes(scriptOptions: BaseOptions, workspace: string, prompt: string) {
	if (['deno', 'bun', 'nativets'].includes(scriptOptions.language)) {
		const resourceTypes = await ResourceService.listResourceType({ workspace })
		const resourceTypesText = formatResourceTypes(resourceTypes, 'typescript')
		prompt = prompt.replace('{resourceTypes}', resourceTypesText)
	} else if (scriptOptions.language === 'python3') {
		const resourceTypes = await ResourceService.listResourceType({ workspace })
		const resourceTypesText = formatResourceTypes(resourceTypes, 'python3')
		prompt = prompt.replace('{resourceTypes}', resourceTypesText)
	}
	return prompt
}

function addDBSChema(scriptOptions: BaseOptions, prompt: string) {
	if (['mysql', 'postgresql'].includes(scriptOptions.language) && scriptOptions.dbSchema) {
		const { dbSchema } = scriptOptions
		const smallerSchema = {}
		for (const schemaKey in dbSchema) {
			for (const tableKey in dbSchema[schemaKey]) {
				smallerSchema[tableKey] = []
				for (const colKey in dbSchema[schemaKey][tableKey]) {
					const col = dbSchema[schemaKey][tableKey][colKey]
					const p = [colKey, col.type, col.required]
					if (col.default) {
						p.push(col.default)
					}
					smallerSchema[tableKey].push(p)
				}
			}
		}
		prompt =
			prompt +
			"\nHere's the database schema, each column is in the format [name, type, required, default?]: " +
			JSON.stringify(smallerSchema)
	}
	return prompt
}

export async function generateScript(scriptOptions: ScriptGenerationOptions) {
	if (!workspace) {
		throw new Error('No workspace selected')
	}

	const baseURL = `${location.origin}${OpenAPI.BASE}/w/${workspace}/openai/proxy`
	const openai = new OpenAI({
		baseURL,
		apiKey: 'fakekey',
		defaultHeaders: {
			Authorization: ''
		}
	})

	let prompt = GEN_CONFIG.prompts[scriptOptions.language].prompt.replace(
		'{description}',
		scriptOptions.description
	)

	prompt = await addResourceTypes(scriptOptions, workspace, prompt)

	prompt = addDBSChema(scriptOptions, prompt)

	const completion = await openai.chat.completions.create({
		model: 'gpt-4',
		max_tokens: 2048,
		messages: [
			{
				role: 'system',
				content: GEN_CONFIG.system
			},
			{
				role: 'user',
				content: prompt
			}
		]
	})

	const result = completion.choices[0]?.message?.content

	if (!result) {
		throw new Error('No result from OpenAI')
	}

	const match = result.match(/```[a-zA-z]+\n([\s\S]*?)\n```/)

	if (!match || match.length < 2) {
		throw new Error('No code block found')
	}

	const code = match[1]

	if (scriptOptions.language == Script.language.GO) {
		const warning = COMMENT_TYPES[scriptOptions.language] + ' ' + WARNING_MSG + '\n'

		return { code: code.trim().replace('package inner\n', 'package inner\n' + warning) }
	} else if (scriptOptions.language == Script.language.BASH) {
		return {
			code:
				'# shellcheck shell=bash\n' +
				COMMENT_TYPES[scriptOptions.language] +
				' ' +
				WARNING_MSG +
				'\n\n' +
				code.trim()
		}
	} else {
		return {
			code: COMMENT_TYPES[scriptOptions.language] + ' ' + WARNING_MSG + '\n\n' + code.trim()
		}
	}
}

export async function editScript(scriptOptions: EditScriptOptions) {
	if (!workspace) {
		throw new Error('No workspace selected')
	}

	const baseURL = `${location.origin}${OpenAPI.BASE}/w/${workspace}/openai/proxy`
	const openai = new OpenAI({
		baseURL,
		apiKey: 'fakekey',
		defaultHeaders: {
			Authorization: ''
		}
	})

	let prompt = EDIT_CONFIG.prompts[scriptOptions.language].prompt
		.replace('{code}', scriptOptions.selectedCode)
		.replace('{description}', scriptOptions.description)

	prompt = await addResourceTypes(scriptOptions, workspace, prompt)

	prompt = addDBSChema(scriptOptions, prompt)

	const completion = await openai.chat.completions.create({
		model: 'gpt-4',
		max_tokens: 2048,
		messages: [
			{
				role: 'system',
				content: EDIT_CONFIG.system
			},
			{
				role: 'user',
				content: prompt
			}
		],
		temperature: 0.5
	})

	const result = completion.choices[0]?.message?.content

	if (!result) {
		throw new Error('No result from OpenAI')
	}

	const match = result.match(/```[a-zA-z]+\n([\s\S]*?)\n```/)

	if (!match || match.length < 2) {
		throw new Error('No code block found')
	}

	const code = match[1]

	return { code }
}

export async function fixScript(scriptOptions: FixScriptOpions) {
	if (!workspace) {
		throw new Error('No workspace selected')
	}

	const baseURL = `${location.origin}${OpenAPI.BASE}/w/${workspace}/openai/proxy`
	const openai = new OpenAI({
		baseURL,
		apiKey: 'fakekey',
		defaultHeaders: {
			Authorization: ''
		}
	})

	let prompt = FIX_CONFIG.prompts[scriptOptions.language].prompt
		.replace('{code}', scriptOptions.code)
		.replace('{error}', scriptOptions.error)

	prompt = await addResourceTypes(scriptOptions, workspace, prompt)

	prompt = addDBSChema(scriptOptions, prompt)

	const completion = await openai.chat.completions.create({
		model: 'gpt-4',
		max_tokens: 2048,
		messages: [
			{
				role: 'system',
				content: FIX_CONFIG.system
			},
			{
				role: 'user',
				content: prompt
			}
		]
	})

	const result = completion.choices[0]?.message?.content

	if (!result) {
		throw new Error('No result from OpenAI')
	}

	const match = result.match(/```[a-zA-z]+\n([\s\S]*?)\n```/)

	if (!match || match.length < 2) {
		throw new Error('No code block found')
	}

	const explanationMatch = result.match(/explanation: "(.+)"/i)

	let explanation = ''
	if (explanationMatch && explanationMatch.length > 1) {
		explanation = explanationMatch[1]
	}

	const code = match[1]

	return { code, explanation }
}