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 { scriptLangToEditorLang } from '$lib/scripts'

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: '//'
}

function scriptLangToEnvironment(lang: Script.language | 'frontend') {
	if (lang === Script.language.DENO) {
		return 'typescript in a deno running environment'
	} else if (lang === Script.language.BUN) {
		return 'typescript in a node.js running environment'
	} else if (lang === Script.language.NATIVETS) {
		return 'typescript where you should use fetch and are not allowed to import any libraries'
	} else if (lang === 'frontend') {
		return 'client-side javascript'
	} else {
		return lang
	}
}

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'
}

interface ScriptGenerationOptions extends BaseOptions {
	description: string
}

interface EditScriptOptions extends ScriptGenerationOptions {
	selectedCode: string
}

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

export async function generateScript(scriptOptions: ScriptGenerationOptions): Promise<string> {
	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
	)

	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)
	}

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

	let 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')
	}

	result = match[1]

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

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

export async function editScript(scriptOptions: EditScriptOptions): Promise<string> {
	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.prompt
		.replace(
			'{lang}',
			scriptOptions.language === 'frontend'
				? 'javascript'
				: scriptLangToEditorLang(scriptOptions.language)
		)
		.replace('{code}', scriptOptions.selectedCode)
		.replace('{description}', scriptOptions.description)
		.replace('{environment}', scriptLangToEnvironment(scriptOptions.language))

	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
	})

	let 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')
	}

	result = match[1]

	return result
}

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.prompt
		.replace(
			'{lang}',
			scriptOptions.language === 'frontend'
				? 'javascript'
				: scriptLangToEditorLang(scriptOptions.language)
		)
		.replace('{code}', scriptOptions.code)
		.replace('{error}', scriptOptions.error)
		.replace('{environment}', scriptLangToEnvironment(scriptOptions.language))

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

	let 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')
	}

	result = match[1]

	return result
}