From 0aa81e39705d8c2109c8ec30855bb5f68eae133b Mon Sep 17 00:00:00 2001
From: HugoCasa <hugo@casademont.ch>
Date: Wed, 26 Jul 2023 12:34:04 +0200
Subject: [PATCH] feat(frontend): AI edit / fix improvements (#1923)

* feat(frontend): db schema explorer + db aware AI

* fix: explorer button consistency

* fix: explorer btn really consistent

* feat: improve autocompletion regex

* feat(forntend): AI edit / fix improvements

* fix: correct typos
---
 .../lib/components/codeGen/ScriptFix.svelte   |  20 +-
 .../lib/components/codeGen/ScriptGen.svelte   |  10 +-
 frontend/src/lib/components/codeGen/lib.ts    | 142 ++++---
 .../lib/components/codeGen/prompts/edit.yaml  |  97 ++++-
 .../lib/components/codeGen/prompts/fix.yaml   | 103 ++++-
 .../lib/components/codeGen/prompts/gen.yaml   |   8 +-
 .../lib/components/codeGen/prompts/index.ts   |  13 +-
 llm/sample_answers.yaml                       | 369 ++++++++++++++----
 llm/sample_queries.yaml                       |   8 +-
 llm/src/gen_samples.py                        | 109 ++----
 10 files changed, 634 insertions(+), 245 deletions(-)

diff --git a/frontend/src/lib/components/codeGen/ScriptFix.svelte b/frontend/src/lib/components/codeGen/ScriptFix.svelte
index 3573cf10f6..411d111706 100644
--- a/frontend/src/lib/components/codeGen/ScriptFix.svelte
+++ b/frontend/src/lib/components/codeGen/ScriptFix.svelte
@@ -6,9 +6,10 @@
 	import { sendUserToast } from '$lib/toast'
 	import type Editor from '../Editor.svelte'
 	import { faCheck, faClose, faMagicWandSparkles } from '@fortawesome/free-solid-svg-icons'
-	import { existsOpenaiResourcePath } from '$lib/stores'
+	import { dbSchema, existsOpenaiResourcePath } from '$lib/stores'
 	import type DiffEditor from '../DiffEditor.svelte'
 	import { scriptLangToEditorLang } from '$lib/scripts'
+	import Popover from '../Popover.svelte'
 
 	// props
 	export let lang: SupportedLanguage
@@ -20,6 +21,7 @@
 	let genLoading: boolean = false
 	let openaiAvailable: boolean | undefined = undefined
 	let generatedCode = ''
+	let explanation = ''
 
 	async function onFix() {
 		if (!error) {
@@ -33,11 +35,14 @@
 			}
 
 			genLoading = true
-			generatedCode = await fixScript({
+			const result = await fixScript({
 				language: lang,
 				code: editor?.getCode() || '',
-				error
+				error,
+				dbSchema: lang === 'postgresql' ? $dbSchema : undefined
 			})
+			generatedCode = result.code
+			explanation = result.explanation
 		} catch (err) {
 			sendUserToast('Failed to generate code', true)
 			console.error(err)
@@ -50,11 +55,13 @@
 		editor?.setCode(diffEditor?.getModified() || '')
 		editor?.format()
 		generatedCode = ''
+		explanation = ''
 		error = ''
 	}
 
 	function rejectDiff() {
 		generatedCode = ''
+		explanation = ''
 	}
 
 	function checkIfOpenaiAvailable(
@@ -108,6 +115,13 @@
 					>
 						Accept
 					</Button>
+					{#if explanation}
+						<Popover>
+							<svelte:fragment slot="text">{explanation}</svelte:fragment>
+							<Button size="xs" color="light" variant="contained" spacingSize="xs2">Explain</Button
+							></Popover
+						>
+					{/if}
 				</div>
 			{:else}
 				<Button
diff --git a/frontend/src/lib/components/codeGen/ScriptGen.svelte b/frontend/src/lib/components/codeGen/ScriptGen.svelte
index 5c057ae80b..cce6b90316 100644
--- a/frontend/src/lib/components/codeGen/ScriptGen.svelte
+++ b/frontend/src/lib/components/codeGen/ScriptGen.svelte
@@ -48,18 +48,20 @@
 			if (isEdit && selection) {
 				const selectedCode = editor?.getSelectedLines() || ''
 				const originalCode = editor?.getCode() || ''
-				const selectionGenCode = await editScript({
+				const result = await editScript({
 					language: lang,
 					description: funcDesc,
-					selectedCode
+					selectedCode,
+					dbSchema: lang === 'postgresql' ? $dbSchema : undefined
 				})
-				generatedCode = originalCode.replace(selectedCode, selectionGenCode + '\n')
+				generatedCode = originalCode.replace(selectedCode, result.code + '\n')
 			} else {
-				generatedCode = await generateScript({
+				const result = await generateScript({
 					language: lang,
 					description: funcDesc,
 					dbSchema: lang === 'postgresql' ? $dbSchema : undefined
 				})
+				generatedCode = result.code
 			}
 			funcDesc = ''
 		} catch (err) {
diff --git a/frontend/src/lib/components/codeGen/lib.ts b/frontend/src/lib/components/codeGen/lib.ts
index da6aa0323d..8506cb1155 100644
--- a/frontend/src/lib/components/codeGen/lib.ts
+++ b/frontend/src/lib/components/codeGen/lib.ts
@@ -4,7 +4,6 @@ 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'
 
@@ -23,20 +22,6 @@ const COMMENT_TYPES = {
 	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
@@ -55,6 +40,7 @@ workspaceStore.subscribe(async (value) => {
 
 interface BaseOptions {
 	language: Script.language | 'frontend'
+	dbSchema?: object
 }
 
 interface ScriptGenerationOptions extends BaseOptions {
@@ -71,7 +57,45 @@ interface FixScriptOpions extends BaseOptions {
 	error: string
 }
 
-export async function generateScript(scriptOptions: ScriptGenerationOptions): Promise<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 (scriptOptions.language === 'postgresql' && 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')
 	}
@@ -90,15 +114,9 @@ export async function generateScript(scriptOptions: ScriptGenerationOptions): Pr
 		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)
-	}
+	prompt = await addResourceTypes(scriptOptions, workspace, prompt)
+
+	prompt = addDBSChema(scriptOptions, prompt)
 
 	if (scriptOptions.language === 'postgresql' && scriptOptions.dbSchema) {
 		const { dbSchema } = scriptOptions
@@ -137,7 +155,7 @@ export async function generateScript(scriptOptions: ScriptGenerationOptions): Pr
 		]
 	})
 
-	let result = completion.choices[0]?.message?.content
+	const result = completion.choices[0]?.message?.content
 
 	if (!result) {
 		throw new Error('No result from OpenAI')
@@ -149,27 +167,30 @@ export async function generateScript(scriptOptions: ScriptGenerationOptions): Pr
 		throw new Error('No code block found')
 	}
 
-	result = match[1]
+	const code = 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)
+		return { code: code.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()
-		)
+		return {
+			code:
+				'# shellcheck shell=bash\n' +
+				COMMENT_TYPES[scriptOptions.language] +
+				' ' +
+				WARNING_MSG +
+				'\n\n' +
+				code.trim()
+		}
 	} else {
-		return COMMENT_TYPES[scriptOptions.language] + ' ' + WARNING_MSG + '\n\n' + result.trim()
+		return {
+			code: COMMENT_TYPES[scriptOptions.language] + ' ' + WARNING_MSG + '\n\n' + code.trim()
+		}
 	}
 }
 
-export async function editScript(scriptOptions: EditScriptOptions): Promise<string> {
+export async function editScript(scriptOptions: EditScriptOptions) {
 	if (!workspace) {
 		throw new Error('No workspace selected')
 	}
@@ -183,16 +204,13 @@ export async function editScript(scriptOptions: EditScriptOptions): Promise<stri
 		}
 	})
 
-	let prompt = EDIT_CONFIG.prompt
-		.replace(
-			'{lang}',
-			scriptOptions.language === 'frontend'
-				? 'javascript'
-				: scriptLangToEditorLang(scriptOptions.language)
-		)
+	let prompt = EDIT_CONFIG.prompts[scriptOptions.language].prompt
 		.replace('{code}', scriptOptions.selectedCode)
 		.replace('{description}', scriptOptions.description)
-		.replace('{environment}', scriptLangToEnvironment(scriptOptions.language))
+
+	prompt = await addResourceTypes(scriptOptions, workspace, prompt)
+
+	prompt = addDBSChema(scriptOptions, prompt)
 
 	const completion = await openai.chat.completions.create({
 		model: 'gpt-4',
@@ -210,7 +228,7 @@ export async function editScript(scriptOptions: EditScriptOptions): Promise<stri
 		temperature: 0.5
 	})
 
-	let result = completion.choices[0]?.message?.content
+	const result = completion.choices[0]?.message?.content
 
 	if (!result) {
 		throw new Error('No result from OpenAI')
@@ -222,9 +240,9 @@ export async function editScript(scriptOptions: EditScriptOptions): Promise<stri
 		throw new Error('No code block found')
 	}
 
-	result = match[1]
+	const code = match[1]
 
-	return result
+	return { code }
 }
 
 export async function fixScript(scriptOptions: FixScriptOpions) {
@@ -241,16 +259,13 @@ export async function fixScript(scriptOptions: FixScriptOpions) {
 		}
 	})
 
-	let prompt = FIX_CONFIG.prompt
-		.replace(
-			'{lang}',
-			scriptOptions.language === 'frontend'
-				? 'javascript'
-				: scriptLangToEditorLang(scriptOptions.language)
-		)
+	let prompt = FIX_CONFIG.prompts[scriptOptions.language].prompt
 		.replace('{code}', scriptOptions.code)
 		.replace('{error}', scriptOptions.error)
-		.replace('{environment}', scriptLangToEnvironment(scriptOptions.language))
+
+	prompt = await addResourceTypes(scriptOptions, workspace, prompt)
+
+	prompt = addDBSChema(scriptOptions, prompt)
 
 	const completion = await openai.chat.completions.create({
 		model: 'gpt-4',
@@ -267,7 +282,7 @@ export async function fixScript(scriptOptions: FixScriptOpions) {
 		]
 	})
 
-	let result = completion.choices[0]?.message?.content
+	const result = completion.choices[0]?.message?.content
 
 	if (!result) {
 		throw new Error('No result from OpenAI')
@@ -279,7 +294,14 @@ export async function fixScript(scriptOptions: FixScriptOpions) {
 		throw new Error('No code block found')
 	}
 
-	result = match[1]
+	const explanationMatch = result.match(/explanation: "(.+)"/i)
 
-	return result
+	let explanation = ''
+	if (explanationMatch && explanationMatch.length > 1) {
+		explanation = explanationMatch[1]
+	}
+
+	const code = match[1]
+
+	return { code, explanation }
 }
diff --git a/frontend/src/lib/components/codeGen/prompts/edit.yaml b/frontend/src/lib/components/codeGen/prompts/edit.yaml
index 86cc8a1457..4a5656843b 100644
--- a/frontend/src/lib/components/codeGen/prompts/edit.yaml
+++ b/frontend/src/lib/components/codeGen/prompts/edit.yaml
@@ -1,13 +1,94 @@
 system: |-
-  You write code as queried by the user. Only output code. Wrap the code like that: 
+  You write code as instructed by the user. Only output code. Wrap the code like that: 
   ```language
   {code}
   ```
   Put explanations directly in the code as comments.
-prompt: |-
-  Here's my environement: {environment}
-  Here's my code: 
-  ```{lang}
-  {code}
-  ```
-  My instructions: {description}
+prompts:
+  python3:
+    prompt: |-
+      Here's my python3 code: 
+      ```python
+      {code}
+      ```
+      Additional information: We have to export a "main" function and specify the parameter types but do not call it.
+      You have access to the following resource types, if you need them, you have to define the TypedDict exactly as specified (class name has to be IN LOWERCASE) and add them as parameters: {resourceTypes}
+      Only use the ones you need. If the TypedDict name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      My instructions: {description}
+  deno:
+    prompt: |-
+      Here's my typescript code in a deno running environement: 
+      ```typescript
+      {code}
+      ```
+      Additional information: We have to export a "main" function like this: "export async function main(...)" and specify the parameter types but do not call it.
+      You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+      Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      My instructions: {description}
+  go:
+    prompt: |-
+      Here's my go code: 
+      ```go
+      {code}
+      ```
+      Additional information: We have to export a "main" function. Import the packages you need. The return type of the function has to be ({return_type}, error). The file package has to be "inner"
+      My instructions: {description}
+  bash:
+    prompt: |-
+      Here's my bash code: 
+      ```shell
+      {code}
+      ```
+      Additional information: Do not include "#!/bin/bash". Arguments are always string and can only be obtained with "var1="$1"", "var2="$2"", etc... You do not need to check if the arguments are present.
+      My instructions: {description}
+  postgresql:
+    prompt: |-
+      Here's my PostgreSQL code: 
+      ```sql
+      {code}
+      ```
+      Additional information: Arguments can be obtained directly in the statement with `$1::{type}`, `$2::{type}`, etc... Name the parameters by adding comments before the command like that: `-- $1 name1` or `-- $2 name = default` (one per row, do not include the type)
+      My instructions: {description}
+  mysql:
+    prompt: |-
+      Here's my MySQL code: 
+      ```sql
+      {code}
+      ```
+      Additional information: Arguments can be obtained directly in the statement with ?. Name the parameters by adding comments before the command like that: -- ? name1 ({type}) (one per row)
+      My instructions: {description}
+  nativets:
+    prompt: |-
+      Here's my typescript code: 
+      ```typescript
+      {code}
+      ```
+      Additional information: We have to export a "main" function like this: "export async function main(...)" and specify the parameter types but do not call it.
+      You should use fetch and are not allowed to import any libraries.
+      You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+      Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      My instructions: {description}
+  bun:
+    prompt: |-
+      Here's my typescript code in a node.js running environment: 
+      ```typescript
+      {code}
+      ```
+      Additional information: We have to export a "main" function like this: "export async function main(...)" and specify the parameter types but do not call it.
+      You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+      Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      My instructions: {description}
+  frontend:
+    prompt: |-
+      Here's my client-side javascript code: 
+      ```typescript
+      {code}
+      ```
+      Additional information: You can access the context object with the ctx global variable. 
+      The app state is a store that can be used to store data. You can access and update the state object with the state global variable like this: state.foo = 'bar'
+      You can use the goto function to navigate to a specific URL: goto(path: string, newTab?: boolean)
+      Use the setTab function to manually set the tab of a Tab component: setTab(id: string, index: string)
+      Use the recompute function to recompute a component: recompute(id: string)
+      Use the getAgGrid function to get the ag-grid instance of a table: getAgGrid(id: string)
+      The setValue function is meant to set or force the value of a component: setValue(id: string, value: any).
+      My instructions: {description}
diff --git a/frontend/src/lib/components/codeGen/prompts/fix.yaml b/frontend/src/lib/components/codeGen/prompts/fix.yaml
index 9921f0c7fe..05c9f4f270 100644
--- a/frontend/src/lib/components/codeGen/prompts/fix.yaml
+++ b/frontend/src/lib/components/codeGen/prompts/fix.yaml
@@ -3,7 +3,9 @@ system: |-
   ```language
   {code}
   ```
-  Put explanations directly in the code as comments.
+  Explain the error and the fix in the following format:
+  explanation: "Here's the explanation"
+  Also put the explanations in the code as comments.
 prompt: |-
   Here's my environement: {environment}
   Here's my code: 
@@ -11,4 +13,101 @@ prompt: |-
   {code}
   ```
   I get the following error: {error}
-  Fix it for me.
+  Fix my code.
+prompts:
+  python3:
+    prompt: |-
+      Here's my python3 code: 
+      ```python
+      {code}
+      ```
+      Additional information: We have to export a "main" function and specify the parameter types but do not call it.
+      You have access to the following resource types, if you need them, you have to define the TypedDict exactly as specified (class name has to be IN LOWERCASE) and add them as parameters: {resourceTypes}
+      Only use the ones you need. If the TypedDict name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      I get the following error: {error}
+      Fix my code.
+  deno:
+    prompt: |-
+      Here's my typescript code in a deno running environement: 
+      ```typescript
+      {code}
+      ```
+      Additional information: We have to export a "main" function like this: "export async function main(...)" and specify the parameter types but do not call it.
+      You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+      Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      I get the following error: {error}
+      Fix my code.
+  go:
+    prompt: |-
+      Here's my go code: 
+      ```go
+      {code}
+      ```
+      Additional information: We have to export a "main" function. Import the packages you need. The return type of the function has to be ({return_type}, error). The file package has to be "inner"
+      I get the following error: {error}
+      Fix my code.
+  bash:
+    prompt: |-
+      Here's my bash code: 
+      ```shell
+      {code}
+      ```
+      Additional information: Do not include "#!/bin/bash". Arguments are always string and can only be obtained with "var1="$1"", "var2="$2"", etc... You do not need to check if the arguments are present.
+      I get the following error: {error}
+      Fix my code.
+  postgresql:
+    prompt: |-
+      Here's my PostgreSQL code: 
+      ```sql
+      {code}
+      ```
+      Additional information: Arguments can be obtained directly in the statement with `$1::{type}`, `$2::{type}`, etc... Name the parameters by adding comments before the command like that: `-- $1 name1` or `-- $2 name = default` (one per row, do not include the type)
+      I get the following error: {error}
+      Fix my code.
+  mysql:
+    prompt: |-
+      Here's my MySQL code: 
+      ```sql
+      {code}
+      ```
+      Additional information: Arguments can be obtained directly in the statement with ?. Name the parameters by adding comments before the command like that: -- ? name1 ({type}) (one per row)
+      I get the following error: {error}
+      Fix my code.
+  nativets:
+    prompt: |-
+      Here's my typescript code: 
+      ```typescript
+      {code}
+      ```
+      Additional information: We have to export a "main" function like this: "export async function main(...)" and specify the parameter types but do not call it.
+      You should use fetch and are not allowed to import any libraries.
+      You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+      Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      I get the following error: {error}
+      Fix my code.
+  bun:
+    prompt: |-
+      Here's my typescript code in a node.js running environment: 
+      ```typescript
+      {code}
+      ```
+      Additional information: We have to export a "main" function like this: "export async function main(...)" and specify the parameter types but do not call it.
+      You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+      Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      I get the following error: {error}
+      Fix my code.
+  frontend:
+    prompt: |-
+      Here's my client-side javascript code: 
+      ```typescript
+      {code}
+      ```
+      Additional information: You can access the context object with the ctx global variable. 
+      The app state is a store that can be used to store data. You can access and update the state object with the state global variable like this: state.foo = 'bar'
+      You can use the goto function to navigate to a specific URL: goto(path: string, newTab?: boolean)
+      Use the setTab function to manually set the tab of a Tab component: setTab(id: string, index: string)
+      Use the recompute function to recompute a component: recompute(id: string)
+      Use the getAgGrid function to get the ag-grid instance of a table: getAgGrid(id: string)
+      The setValue function is meant to set or force the value of a component: setValue(id: string, value: any).
+      I get the following error: {error}
+      Fix my code.
diff --git a/frontend/src/lib/components/codeGen/prompts/gen.yaml b/frontend/src/lib/components/codeGen/prompts/gen.yaml
index e90444276f..9f33a62a4a 100644
--- a/frontend/src/lib/components/codeGen/prompts/gen.yaml
+++ b/frontend/src/lib/components/codeGen/prompts/gen.yaml
@@ -9,12 +9,12 @@ prompts:
     prompt: |-
       Write a function in python called "main". The function should {description}. Specify the parameter types. Do not call the main function.
       You have access to the following resource types, if you need them, you have to define the TypedDict exactly as specified (class name has to be IN LOWERCASE) and add them as parameters: {resourceTypes}
-      If the TypedDict name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      Only use the ones you need. If the TypedDict name conflicts with the imported object, rename the imported object NOT THE TYPE.
   deno:
     prompt: |-
       Write a function in typescript called "main". The function should {description}. Specify the parameter types. You are in a Deno environment. You can import deno libraries or you can also import npm libraries like that: "import ... from "npm:{package}";". Export the "main" function like this: "export async function main(...)". Do not call the main function.
       You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
-      If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
   go:
     prompt: |-
       Write a function in go called "main". The function should {description}. Import the packages you need. The return type of the function has to be ({return_type}, error). The file package has to be "inner".
@@ -31,12 +31,12 @@ prompts:
     prompt: |-
       Write a function in typescript called "main". The function should {description}. Specify the parameter types. You should use fetch and are not allowed to import any libraries. Export the "main" function like this: "export async function main(...)". Do not call the main function.
       You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
-      If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
   bun:
     prompt: |-
       Write a function in typescript called "main". The function should {description}. Specify the parameter types. You are in a Node.js environment. You can import npm libraries. Export the "main" function like this: "export async function main(...)". Do not call the main function.
       You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
-      If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+      Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
   frontend:
     prompt: |-
       Write client-side javascript code that should {description}. You have access to a few helpers:
diff --git a/frontend/src/lib/components/codeGen/prompts/index.ts b/frontend/src/lib/components/codeGen/prompts/index.ts
index 668e0877af..332a28dcf2 100644
--- a/frontend/src/lib/components/codeGen/prompts/index.ts
+++ b/frontend/src/lib/components/codeGen/prompts/index.ts
@@ -2,12 +2,7 @@ import editObj from './edit.yaml'
 import fixObj from './fix.yaml'
 import genObj from './gen.yaml'
 
-interface CommonConfig {
-	system: string
-	prompt: string
-}
-
-interface GenConfig {
+interface PromptsConfig {
 	system: string
 	prompts: {
 		[key: string]: {
@@ -16,10 +11,10 @@ interface GenConfig {
 	}
 }
 
-const EDIT_CONFIG = editObj as CommonConfig
+const EDIT_CONFIG = editObj as PromptsConfig
 
-const FIX_CONFIG = fixObj as CommonConfig
+const FIX_CONFIG = fixObj as PromptsConfig
 
-const GEN_CONFIG = genObj as GenConfig
+const GEN_CONFIG = genObj as PromptsConfig
 
 export { EDIT_CONFIG, FIX_CONFIG, GEN_CONFIG }
diff --git a/llm/sample_answers.yaml b/llm/sample_answers.yaml
index 750c456e95..252ff5b83a 100644
--- a/llm/sample_answers.yaml
+++ b/llm/sample_answers.yaml
@@ -1,6 +1,7 @@
 - answer: |-
     ```python
     import psycopg2
+    from psycopg2 import OperationalError
     from typing import TypedDict
 
     class postgresql(TypedDict):
@@ -11,36 +12,52 @@
         sslmode: str
         password: str
 
-    def main(config: postgresql):
-        # Create a connection to the database
-        conn = psycopg2.connect(
-            host=config['host'],
-            port=config['port'],
-            user=config['user'],
-            password=config['password'],
-            dbname=config['dbname'],
-            sslmode=config['sslmode']
-        )
+    def main(postgres: postgresql):
+        # Create a connection to the PostgreSQL database
+        try:
+            connection = psycopg2.connect(
+                host=postgres['host'],
+                port=postgres['port'],
+                user=postgres['user'],
+                password=postgres['password'],
+                dbname=postgres['dbname'],
+                sslmode=postgres['sslmode']
+            )
 
-        # Create a cursor object
-        cur = conn.cursor()
+            # Create a cursor object
+            cursor = connection.cursor()
 
-        # Execute the SQL query to fetch all rows from the orders table
-        cur.execute("SELECT * FROM orders")
+            # Execute the SQL query
+            cursor.execute("SELECT * FROM orders")
 
-        # Fetch all rows from the last executed SQL query
-        rows = cur.fetchall()
+            # Fetch all the rows
+            rows = cursor.fetchall()
 
-        # Print the rows
-        for row in rows:
-            print(row)
+            # Print all the rows
+            for row in rows:
+                print(row)
 
-        # Close the cursor and connection
-        cur.close()
-        conn.close()
+        except OperationalError as e:
+            print(f"The error '{e}' occurred")
+
+        finally:
+            # Close the connection
+            if connection:
+                cursor.close()
+                connection.close()
     ```
   description: Connect to postgres and list the rows in the orders table
   lang: python3
+  template_prompt: |-
+    Write a function in python called "main". The function should {description}. Specify the parameter types. Do not call the main function.
+    You have access to the following resource types, if you need them, you have to define the TypedDict exactly as specified (class name has to be IN LOWERCASE) and add them as parameters: {resourceTypes}
+    Only use the ones you need. If the TypedDict name conflicts with the imported object, rename the imported object NOT THE TYPE.
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     ```python
@@ -48,21 +65,39 @@
       # Add a condition to check if num is zero.
       # If num is zero, return a message to avoid division by zero error.
       if num == 0:
-        return "Error: Division by zero is not allowed."
-      else:
-        return 3 / num
+        raise ValueError('The input number cannot be zero')
+      return 3 / num
     ```
+    explanation: "The error occurs because you're trying to divide by zero which is not allowed in mathematics. The fix is to add a check at the beginning of the function to prevent division by zero. If the input number is zero, it raises a ValueError with an appropriate message."
   code: |-
     def main(num: int) -> int:
       return 3 / num
   error: division by zero
   lang: python3
+  template_prompt: |-
+    Here's my python3 code:
+    ```python
+    {code}
+    ```
+    Additional information: We have to export a "main" function and specify the parameter types but do not call it.
+    You have access to the following resource types, if you need them, you have to define the TypedDict exactly as specified (class name has to be IN LOWERCASE) and add them as parameters: {resourceTypes}
+    Only use the ones you need. If the TypedDict name conflicts with the imported object, rename the imported object NOT THE TYPE.
+    I get the following error: {error}
+    Fix my code.
+  template_system: |-
+    You fix the code shared by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Explain the error and the fix in the following format:
+    explanation: "Here's the explanation"
+    Also put the explanations in the code as comments.
   type: fix
 - answer: |-
     ```python
-    # This is a Python function named main. It takes an integer as input.
+    # The function main is defined here. It takes an integer 'num' as an argument.
     def main(num: int) -> int:
-      # The function returns the input number divided by 16.
+      # The function returns the result of the division of 'num' by 16.
       return num / 16
     ```
   code: |-
@@ -70,6 +105,21 @@
       return num / 16
   description: Comment my code
   lang: python3
+  template_prompt: |-
+    Here's my python3 code:
+    ```python
+    {code}
+    ```
+    Additional information: We have to export a "main" function and specify the parameter types but do not call it.
+    You have access to the following resource types, if you need them, you have to define the TypedDict exactly as specified (class name has to be IN LOWERCASE) and add them as parameters: {resourceTypes}
+    Only use the ones you need. If the TypedDict name conflicts with the imported object, rename the imported object NOT THE TYPE.
+    My instructions: {description}
+  template_system: |-
+    You write code as instructed by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: edit
 - answer: |-
     ```go
@@ -85,36 +135,51 @@
     func main(num float64) (float64, error) {
         // check if the number is 0, as division by 0 is not allowed
         if num == 0 {
-            return 0, errors.New("Division by 0 is not allowed")
+            return 0, errors.New("Cannot divide by zero")
         }
 
-        // divide the number by 16
-        result := num / 16
-
-        // return the result and nil as error
+        result := num / 16.0
         return result, nil
     }
     ```
-    In this code, we first check if the input number is 0. If it is, we return an error because division by zero is not allowed. If the number is not 0, we divide it by 16 and return the result and nil (indicating no error).
   description: Divide a number by 16
   lang: go
+  template_prompt: |-
+    Write a function in go called "main". The function should {description}. Import the packages you need. The return type of the function has to be ({return_type}, error). The file package has to be "inner".
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     ```typescript
-    // Importing the library for number to word conversion
+    // Importing the required library
     import { toWords } from "npm:number-to-words";
 
-    // Define the main function
+    // Defining the main function
     export async function main(number: number): Promise<string> {
-      // Convert the number to words
-      const words = toWords(number);
+      // Converting the number to word using the toWords function from the number-to-words library
+      const word = toWords(number);
 
-      // Return the result
-      return words;
+      // Returning the word
+      return word;
     }
     ```
+    In this function, we are using the `toWords` function from the `number-to-words` npm library to convert a number to a word. The function takes a number as input and returns a promise that resolves to a string. The string is the word representation of the number.
   description: Convert a number to a word
   lang: deno
+  template_prompt: |-
+    Write a function in typescript called "main". The function should {description}. Specify the parameter types. You are in a Deno environment. You can import deno libraries or you can also import npm libraries like that: "import ... from "npm:{package}";". Export the "main" function like this: "export async function main(...)". Do not call the main function.
+    You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+    Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     ```typescript
@@ -131,95 +196,144 @@
       password: string
     }
 
-    // Exporting the main function
+    // Defining the main function
     export async function main(postgres: Postgresql) {
-      // Creating a new PostgreSQL client with the provided configuration
+      // Creating a new client instance
       const client = new Client({
         host: postgres.host,
         port: postgres.port,
         user: postgres.user,
         password: postgres.password,
         database: postgres.dbname,
-        ssl: {
-          rejectUnauthorized: postgres.sslmode === 'require' ? true : false,
-        },
+        tls: { enforce: postgres.sslmode === "require" },
       });
 
       // Connecting to the PostgreSQL server
       await client.connect();
 
-      // Querying the orders table
-      const res = await client.query('SELECT * FROM orders');
+      // Running the query to list rows in the orders table
+      const result = await client.queryArray("SELECT * FROM orders;");
 
       // Logging the rows from the orders table
       console.log(res.rows);
 
       // Closing the connection to the PostgreSQL server
       await client.end();
+
+      // Returning the result
+      return result.rows;
     }
     ```
-    This function, when called, will connect to a PostgreSQL server using the provided configuration, query the "orders" table, log the rows from the table, and then close the connection. The function is asynchronous because it needs to wait for the connection, query, and disconnection operations to complete.
   description: Connect to postgres and list the rows in the orders table
   lang: deno
+  template_prompt: |-
+    Write a function in typescript called "main". The function should {description}. Specify the parameter types. You are in a Deno environment. You can import deno libraries or you can also import npm libraries like that: "import ... from "npm:{package}";". Export the "main" function like this: "export async function main(...)". Do not call the main function.
+    You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+    Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     Here's a simple TypeScript function called "main" that converts a number to a word using the npm library "number-to-words". I've added the library import at the top of the file. The function takes a number as an argument and returns the word representation of that number.
 
     ```typescript
-    import * as numberToWords from 'number-to-words';
+    // Importing the library to convert numbers to words
+    import * as converter from 'number-to-words';
 
+    // Defining the main function
     export async function main(num: number): Promise<string> {
-      // Convert the number to words using the number-to-words library
-      const word = numberToWords.toWords(num);
+      // Converting the number to word using the library function
+      let word = converter.toWords(num);
 
-      // Return the word
+      // Returning the word
       return word;
     }
     ```
-    The function is asynchronous and returns a Promise that resolves to a string. This is because in a real-world scenario, the conversion process might involve some asynchronous operations (like fetching data from a database or making a request to an API), and we want to be able to handle that. In this simple example, the conversion is done synchronously, but the function is still declared as async for consistency and future-proofing.
+    This function takes a number as input and returns a Promise that resolves to a string. The string is the word representation of the number. The function uses the 'number-to-words' npm library to convert the number to words.
   description: Convert a number to a word
   lang: bun
-  type: gen
-- answer: |-
-    ```sql
-    -- $1::text name
-    -- $2::text email
-    INSERT INTO users (name, email) VALUES ($1::text, $2::text);
+  template_prompt: |-
+    Write a function in typescript called "main". The function should {description}. Specify the parameter types. You are in a Node.js environment. You can import npm libraries. Export the "main" function like this: "export async function main(...)". Do not call the main function.
+    You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+    Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
     ```
-  description: Insert an email and a name in the users table
-  lang: postgresql
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     ```sql
+    -- $1 email
+    -- $2 disabled = false
+    -- $3 is_admin = false
+    -- $4 operator = false
+    -- $5 username
+    -- $6 workspace_id
+    INSERT INTO usr (email, disabled, is_admin, operator, username, workspace_id)
+    VALUES ($1::varchar, $2::bool, $3::bool, $4::bool, $5::varchar, $6::varchar);
+    ```
+    This SQL code will insert a new user into the 'usr' table in your PostgreSQL database. The values for the fields are obtained directly in the statement with `$1::type`, `$2::type`, etc. The comments before the command name the parameters. The default values for 'disabled', 'is_admin', and 'operator' are all false.
+  description: Insert an a new user
+  lang: postgresql
+  template_prompt: |-
+    Write SQL code for a PostgreSQL that should {description}. Arguments can be obtained directly in the statement with `$1::{type}`, `$2::{type}`, etc... Name the parameters by adding comments before the command like that: `-- $1 name1` or `-- $2 name = default` (one per row, do not include the type)
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
+  type: gen
+- answer: |-
+    ```SQL
     -- ? email (VARCHAR)
     -- ? name (VARCHAR)
-    INSERT INTO users (email, name)
-    VALUES (?, ?)
+    INSERT INTO users (email, name) VALUES (?, ?);
     ```
   description: Insert an email and a name in the users table
   lang: mysql
+  template_prompt: |-
+    Write SQL code for MySQL that should {description}. Arguments can be obtained directly in the statement with ?. Name the parameters by adding comments before the command like that: -- ? name1 ({type}) (one per row)
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     ```bash
     var1="$1"
 
-    # Bash only supports integer division natively.
-    # So we use bc, a command line calculator, to perform floating point division.
-    result=$(echo "scale=2; $var1 / 16" | bc)
-
-    echo $result
+    # Use bc (an arbitrary precision calculator language) to perform the division
+    # Scale is set to 2 to get output up to two decimal places
+    echo "scale=2; $var1/16" | bc
     ```
   description: Divide a number by 16
   lang: bash
+  template_prompt: |-
+    Write bash code that should {description}. Do not include "#!/bin/bash". Arguments are always string and can only be obtained with "var1="$1"", "var2="$2"", etc... You do not need to check if the arguments are present.
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     ```bash
-    # Get the arguments
+    # Assign the arguments to variables
     app_name="$1"
     python_version="$2"
 
     # Create a Dockerfile
-    cat > Dockerfile << EOF
+    cat << EOF > Dockerfile
     # Use an official Python runtime as a parent image
     FROM python:${python_version}
 
@@ -241,10 +355,24 @@
     # Run app.py when the container launches
     CMD ["python", "app.py"]
     EOF
+
+    # Build the Docker image
+    docker build -t ${app_name} .
+
+    # Run the Docker image
+    docker run -p 4000:80 ${app_name}
     ```
-    This bash script creates a Dockerfile for a Python application. The Dockerfile uses the official Python image from Docker Hub as its base. It then sets the working directory in the container to /app and copies the current directory (i.e., the directory where the Dockerfile is located) into the /app directory in the container. It installs any Python packages listed in requirements.txt, exposes port 80 for the app to communicate with the outside world, sets an environment variable with the name of the app, and finally, specifies that Docker should execute python app.py when the container is launched.
+    This bash script creates a Dockerfile for a Python application. The Dockerfile is then used to build a Docker image and finally run the Docker container. The script takes two arguments: the name of the application and the Python version. The application name is used as the Docker image name and also set as an environment variable inside the Docker container. The Python version is used to select the parent image for the Docker container. The script assumes that the current directory contains a requirements.txt file with the Python packages required by the application and an app.py file which is the entry point of the application.
   description: Write a docker image for a python app
   lang: bash
+  template_prompt: |-
+    Write bash code that should {description}. Do not include "#!/bin/bash". Arguments are always string and can only be obtained with "var1="$1"", "var2="$2"", etc... You do not need to check if the arguments are present.
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     ```typescript
@@ -253,16 +381,19 @@
       token: string
     }
 
-    // Define main function
-    export async function main(owner: string, repo: string, github: Github): Promise<number> {
-      // Set up the request headers
+    // Define the main function
+    export async function main(repo: string, github: Github): Promise<number> {
+      // Define the Github API URL
+      const url = `https://api.github.com/repos/${repo}`;
+
+      // Define the headers for the fetch request
       const headers = {
         "Authorization": `token ${github.token}`,
         "Accept": "application/vnd.github.v3+json"
       };
 
-      // Fetch data from Github API
-      const response = await fetch(`https://api.github.com/repos/${owner}/${repo}`, { headers });
+      // Make the fetch request to the Github API
+      const response = await fetch(url, { headers });
 
       // Check if response is ok
       if (!response.ok) {
@@ -272,45 +403,113 @@
       // Parse the response to JSON
       const data = await response.json();
 
-      // Return the number of stars
+      // Return the number of stars of the repo
       return data.stargazers_count;
     }
     ```
-    This function uses the fetch API to make a GET request to the GitHub API. The owner and repo parameters are the username and repository name, respectively. The github parameter is an object that contains the GitHub personal access token. The function returns a promise that resolves to the number of stars of the repository. If the request fails, it throws an error.
+    This function takes a repo string and a Github object as parameters. It constructs the Github API URL and defines the headers for the fetch request. It then makes the fetch request to the Github API and checks if the response is ok. If the response is not ok, it throws an error. If the response is ok, it parses the response as JSON and returns the number of stars of the repo.
   description: Query the github api and return the number of stars of a repo
   lang: nativets
+  template_prompt: |-
+    Write a function in typescript called "main". The function should {description}. Specify the parameter types. You should use fetch and are not allowed to import any libraries. Export the "main" function like this: "export async function main(...)". Do not call the main function.
+    You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
+    Only use the ones you need. If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     ```javascript
-    // Access the 'email' property from the global 'ctx' object
-    let email = ctx.email;
+    // Access the email from the context object
+    var email = ctx.email;
 
-    // Use the 'setValue' function to set the value of the input with id 'my_field' to the 'email' variable
+    // Use the setValue function to set the value of the input field with id 'my_field' to the email
     setValue('my_field', email);
     ```
   description:
     set the value of the input with id 'my_field' to the context variable
     email
   lang: frontend
+  template_prompt: |-
+    Write client-side javascript code that should {description}. You have access to a few helpers:
+    You can access the context object with the ctx global variable.
+    The app state is a store that can be used to store data. You can access and update the state object with the state global variable like this: state.foo = 'bar'
+    You can use the goto function to navigate to a specific URL: goto(path: string, newTab?: boolean)
+    Use the setTab function to manually set the tab of a Tab component: setTab(id: string, index: string)
+    Use the recompute function to recompute a component: recompute(id: string)
+    Use the getAgGrid function to get the ag-grid instance of a table: getAgGrid(id: string)
+    The setValue function is meant to set or force the value of a component: setValue(id: string, value: any).
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
 - answer: |-
     ```sql
     -- $1 username
     SELECT
-        DATE_TRUNC('month', completed_job.created_at) AS month,
-        COUNT(completed_job.id) AS job_count
+        DATE_TRUNC('month', timestamp) AS month,
+        COUNT(id) / COUNT(DISTINCT DATE_TRUNC('month', timestamp)) AS average_jobs
     FROM
-        completed_job
+        audit
     WHERE
-        completed_job.created_by = $1::varchar
+        username = $1::varchar AND operation = 'complete'
     GROUP BY
         month
     ORDER BY
         month;
     ```
-    In this SQL query, we are querying the `completed_job` table where the `created_by` column matches the provided username. We then group the result by month (truncated from the `created_at` timestamp), and count the number of jobs in each month. The result is ordered by month.
+    This SQL code calculates the average number of completed jobs per month for a given username.
+
+    It starts by truncating the timestamp to the month, which groups all the jobs completed in the same month together.
+
+    Then it counts the number of jobs completed in each month and divides by the total number of distinct months to get the average number of jobs per month.
+
+    The WHERE clause filters for the specified username and only considers 'complete' operations.
+
+    Finally, it orders the results by month.
   description:
     compute the average number of completed jobs per month for the given
     username
   lang: postgresql
+  template_prompt: |-
+    Write SQL code for a PostgreSQL that should {description}. Arguments can be obtained directly in the statement with `$1::{type}`, `$2::{type}`, etc... Name the parameters by adding comments before the command like that: `-- $1 name1` or `-- $2 name = default` (one per row, do not include the type)
+  template_system: |-
+    You write code as queried by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Put explanations directly in the code as comments.
   type: gen
+- answer: |-
+    explanation: "The error message indicates that there's no column named 'is_secret' in the 'account' table. From the provided schema, it can be observed that the 'is_secret' column is actually in the 'variable' table, not the 'account' table. The correct SQL query should select from the 'variable' table instead of 'account'."
+
+    ```sql
+    SELECT is_secret FROM variable
+    ```
+  code: |-
+    SELECT is_secret FROM account
+  error: 'ExecutionError: db error: ERROR: column "is_secret" does not exist'
+  lang: postgresql
+  template_prompt: |-
+    Here's my PostgreSQL code:
+    ```sql
+    {code}
+    ```
+    Additional information: Arguments can be obtained directly in the statement with `$1::{type}`, `$2::{type}`, etc... Name the parameters by adding comments before the command like that: `-- $1 name1` or `-- $2 name = default` (one per row, do not include the type)
+    I get the following error: {error}
+    Fix my code.
+  template_system: |-
+    You fix the code shared by the user. Only output code. Wrap the code like that:
+    ```language
+    {code}
+    ```
+    Explain the error and the fix in the following format:
+    explanation: "Here's the explanation"
+    Also put the explanations in the code as comments.
+  type: fix
diff --git a/llm/sample_queries.yaml b/llm/sample_queries.yaml
index 76c1739209..d063a67571 100644
--- a/llm/sample_queries.yaml
+++ b/llm/sample_queries.yaml
@@ -27,7 +27,7 @@
   description: Convert a number to a word
 - lang: postgresql
   type: gen
-  description: Insert an email and a name in the users table
+  description: Insert an a new user
 - lang: mysql
   type: gen
   description: Insert an email and a name in the users table
@@ -46,3 +46,9 @@
 - lang: postgresql
   type: gen
   description: compute the average number of completed jobs per month for the given username
+- lang: postgresql
+  type: fix
+  code: |-
+    SELECT is_secret FROM account
+  error: |-
+    ExecutionError: db error: ERROR: column "is_secret" does not exist
diff --git a/llm/src/gen_samples.py b/llm/src/gen_samples.py
index f20ab3ddf5..d25ecc5764 100644
--- a/llm/src/gen_samples.py
+++ b/llm/src/gen_samples.py
@@ -25,20 +25,15 @@ def literal_presenter(dumper, data):
 yaml.add_representer(Literal, literal_presenter)
 
 
-class GenPrompt(TypedDict):
+class Prompt(TypedDict):
     prompt: str
 
 
-class GenConfig(TypedDict):
-    prompts: dict[str, GenPrompt]
+class PromptsConfig(TypedDict):
+    prompts: dict[str, Prompt]
     system: str
 
 
-class CommonConfig(TypedDict):
-    system: str
-    prompt: str
-
-
 class Query(TypedDict):
     description: str
     type: str
@@ -47,35 +42,9 @@ class Query(TypedDict):
     error: str
 
 
-def scriptLangToCodeLang(lang: str):
-    if lang in ["deno", "bun", "nativets"]:
-        return "typescript"
-    elif lang in ["postgresql", "mysql"]:
-        return "sql"
-    elif lang == "python3":
-        return "python"
-    elif lang == "bash":
-        return "shell"
-    elif lang == "frontend":
-        return "javascript"
-    else:
-        return lang
-
-
-def scriptLangToEnvironment(lang: str):
-    if lang == "deno":
-        return "typescript in a deno running environment"
-    elif lang == "bun":
-        return "typescript in a node.js running environment"
-    elif lang == "nativets":
-        return "typescript where you should use fetch and are not allowed to import any libraries"
-    elif lang == "frontend":
-        return "client-side javascript"
-    else:
-        return lang
-
-
-def get_prompts(prompts_path: str) -> Tuple[GenConfig, CommonConfig, CommonConfig]:
+def get_prompts(
+    prompts_path: str,
+) -> Tuple[PromptsConfig, PromptsConfig, PromptsConfig]:
     GEN_CONFIG = None
     EDIT_CONFIG = None
     FIX_CONFIG = None
@@ -95,49 +64,46 @@ def get_queries(queries_path: str) -> list[Query]:
 
 def prepare_prompt(
     query: Query,
-    GEN_CONFIG: GenConfig,
-    EDIT_CONFIG: CommonConfig,
-    FIX_CONFIG: CommonConfig,
+    GEN_CONFIG: PromptsConfig,
+    EDIT_CONFIG: PromptsConfig,
+    FIX_CONFIG: PromptsConfig,
 ):
     system = ""
     prompt = ""
+    template_prompt = ""
     if query["type"] == "gen":
         system = GEN_CONFIG["system"]
-        prompt = GEN_CONFIG["prompts"][query["lang"]]["prompt"]
-        prompt = prompt.replace("{description}", query["description"])
-        if query["lang"] in ["deno", "bun", "nativets"]:
-            prompt = prompt.replace("{resourceTypes}", RESOURCE_TYPES["typescript"])
-        elif query["lang"] in ["python3"]:
-            prompt = prompt.replace("{resourceTypes}", RESOURCE_TYPES["python"])
-        if query["lang"] in ['postgresql']:
-            prompt = prompt + "\nHere's the database schema, each column is in the format [name, type, required, default?]: " + DB_SCHEMA
+        template_prompt = GEN_CONFIG["prompts"][query["lang"]]["prompt"]
+        prompt = template_prompt.replace("{description}", query["description"])
     elif query["type"] == "edit":
         system = EDIT_CONFIG["system"]
-        prompt = EDIT_CONFIG["prompt"]
-        lang = scriptLangToCodeLang(query["lang"])
-        environment = scriptLangToEnvironment(query["lang"])
-        prompt = (
-            prompt.replace("{description}", query["description"])
-            .replace("{lang}", lang)
-            .replace("{environment}", environment)
-            .replace("{code}", query["code"])
+        template_prompt = EDIT_CONFIG["prompts"][query["lang"]]["prompt"]
+        prompt = template_prompt.replace("{description}", query["description"]).replace(
+            "{code}", query["code"]
         )
     elif query["type"] == "fix":
         system = FIX_CONFIG["system"]
-        prompt = FIX_CONFIG["prompt"]
-        lang = scriptLangToCodeLang(query["lang"])
-        environment = scriptLangToEnvironment(query["lang"])
-        prompt = (
-            prompt.replace("{lang}", lang)
-            .replace("{environment}", environment)
-            .replace("{error}", query["error"])
-            .replace("{code}", query["code"])
+        template_prompt = FIX_CONFIG["prompts"][query["lang"]]["prompt"]
+        prompt = template_prompt.replace("{error}", query["error"]).replace(
+            "{code}", query["code"]
         )
 
-    return system, prompt
+    if query["lang"] in ["deno", "bun", "nativets"]:
+        prompt = prompt.replace("{resourceTypes}", RESOURCE_TYPES["typescript"])
+    elif query["lang"] in ["python3"]:
+        prompt = prompt.replace("{resourceTypes}", RESOURCE_TYPES["python"])
+
+    if query["lang"] in ["postgresql"]:
+        prompt = (
+            prompt
+            + "\nHere's the database schema, each column is in the format [name, type, required, default?]: "
+            + DB_SCHEMA
+        )
+
+    return system, prompt, template_prompt
 
 
-def format_answer(answer: str):
+def format_literal(answer: str):
     return re.sub("[^\\S\n]+\n", "\n", answer).replace("\t", "    ")
 
 
@@ -147,8 +113,11 @@ def gen_samples(queries_path: str, answers_path: str, prompts_path: str):
     queries = get_queries(queries_path)
 
     answers = []
+
     for query in tqdm(queries):
-        system, prompt = prepare_prompt(query, GEN_CONFIG, EDIT_CONFIG, FIX_CONFIG)
+        system, prompt, template_prompt = prepare_prompt(
+            query, GEN_CONFIG, EDIT_CONFIG, FIX_CONFIG
+        )
         chat_completion = openai.ChatCompletion.create(
             model="gpt-4",
             messages=[
@@ -161,11 +130,13 @@ def gen_samples(queries_path: str, answers_path: str, prompts_path: str):
 
         answer = {
             **query,
-            "answer": Literal(format_answer(chat_completion["choices"][0]["message"]["content"])),  # type: ignore
+            "answer": Literal(format_literal(chat_completion["choices"][0]["message"]["content"])),  # type: ignore
+            "template_system": Literal(format_literal(system)),
+            "template_prompt": Literal(format_literal(template_prompt)),
         }
 
         if "code" in query:
-            answer["code"] = Literal(query["code"])
+            answer["code"] = Literal(format_literal(query["code"]))
 
         answers.append(answer)