feat(aichat): allow custom system prompt for each mode (#6500)

* add ui to add custom system prompt

* implement backend changes

* add custom prompt in system prompt

* add in openapi

* add max length

* add backend validation

backend/windmill-api/openapi.yaml

@@ -14137,6 +14137,10 @@ components:
           $ref: "#/components/schemas/AIProviderModel"
         code_completion_model:
           $ref: "#/components/schemas/AIProviderModel"
+        custom_prompts:
+          type: object
+          additionalProperties:
+            type: string
 
     Alert:
       type: object

backend/windmill-api/src/ai.rs

@@ -353,6 +353,8 @@ pub struct AIConfig {
     pub default_model: Option<ProviderModel>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub code_completion_model: Option<ProviderModel>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub custom_prompts: Option<HashMap<String, String>>,
 }
 
 pub fn global_service() -> Router {

backend/windmill-api/src/workspaces.rs

@@ -720,6 +720,7 @@ async fn edit_deploy_to() -> Result<String> {
 }
 
 pub const BANNED_DOMAINS: &str = include_str!("../banned_domains.txt");
+pub const MAX_CUSTOM_PROMPT_LENGTH: usize = 5000;
 
 async fn is_allowed_auto_domain(ApiAuthed { email, .. }: ApiAuthed) -> JsonResult<bool> {
     let domain = email.split('@').last().unwrap();
@@ -819,6 +820,20 @@ async fn edit_copilot_config(
 ) -> Result<String> {
     require_admin(is_admin, &username)?;
 
+    // Validate custom prompts length
+    if let Some(ref custom_prompts) = ai_config.custom_prompts {
+        for (mode, prompt) in custom_prompts.iter() {
+            if prompt.len() > MAX_CUSTOM_PROMPT_LENGTH {
+                return Err(Error::BadRequest(format!(
+                    "Custom prompt for mode '{}' exceeds maximum length of {} characters (current: {})",
+                    mode,
+                    MAX_CUSTOM_PROMPT_LENGTH,
+                    prompt.len()
+                )));
+            }
+        }
+    }
+
     let mut tx = db.begin().await?;
 
     sqlx::query!(
@@ -887,6 +902,7 @@ async fn get_copilot_info(
             providers: None,
             default_model: None,
             code_completion_model: None,
+            custom_prompts: None,
         }))
     }
 }

frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts

@@ -37,7 +37,8 @@ import { getStringError } from './utils'
 import type { FlowModuleState, FlowState } from '$lib/components/flows/flowState'
 import type { CurrentEditor, ExtendedOpenFlow } from '$lib/components/flows/types'
 import { untrack } from 'svelte'
-import { getCurrentModel, type DBSchemas } from '$lib/stores'
+import { get } from 'svelte/store'
+import { getCurrentModel, type DBSchemas, copilotInfo } from '$lib/stores'
 import { askTools, prepareAskSystemMessage, prepareAskUserMessage } from './ask/core'
 import { chatState, DEFAULT_SIZE, triggerablesByAi } from './sharedChatState.svelte'
 import type { ContextElement } from './context'
@@ -208,7 +209,8 @@ class AIChatManager {
 		this.mode = mode
 		this.pendingPrompt = pendingPrompt ?? ''
 		if (mode === AIMode.SCRIPT) {
-			this.systemMessage = prepareScriptSystemMessage()
+			const customPrompt = get(copilotInfo).customPrompts?.[mode]
+			this.systemMessage = prepareScriptSystemMessage(customPrompt)
 			this.systemMessage.content = this.NAVIGATION_SYSTEM_PROMPT + this.systemMessage.content
 			const context = this.contextManager.getSelectedContext()
 			const lang = this.scriptEditorOptions?.lang ?? 'bun'
@@ -243,20 +245,24 @@ class AIChatManager {
 				}
 			}
 		} else if (mode === AIMode.FLOW) {
-			this.systemMessage = prepareFlowSystemMessage()
+			const customPrompt = get(copilotInfo).customPrompts?.[mode]
+			this.systemMessage = prepareFlowSystemMessage(customPrompt)
 			this.systemMessage.content = this.NAVIGATION_SYSTEM_PROMPT + this.systemMessage.content
 			this.tools = [this.changeModeTool, ...flowTools]
 			this.helpers = this.flowAiChatHelpers
 		} else if (mode === AIMode.NAVIGATOR) {
-			this.systemMessage = prepareNavigatorSystemMessage()
+			const customPrompt = get(copilotInfo).customPrompts?.[mode]
+			this.systemMessage = prepareNavigatorSystemMessage(customPrompt)
 			this.tools = [this.changeModeTool, ...navigatorTools]
 			this.helpers = {}
 		} else if (mode === AIMode.ASK) {
-			this.systemMessage = prepareAskSystemMessage()
+			const customPrompt = get(copilotInfo).customPrompts?.[mode]
+			this.systemMessage = prepareAskSystemMessage(customPrompt)
 			this.tools = [...askTools]
 			this.helpers = {}
 		} else if (mode === AIMode.API) {
-			this.systemMessage = prepareApiSystemMessage()
+			const customPrompt = get(copilotInfo).customPrompts?.[mode]
+			this.systemMessage = prepareApiSystemMessage(customPrompt)
 			this.tools = [...this.apiTools]
 			this.helpers = {}
 		}

frontend/src/lib/components/copilot/chat/api/core.ts

@@ -57,10 +57,17 @@ export async function getApiTools(): Promise<Tool<{}>[]> {
 
 export const apiTools: Tool<{}>[] = [getDocumentationTool]
 
-export function prepareApiSystemMessage(): ChatCompletionSystemMessageParam {
+export function prepareApiSystemMessage(customPrompt?: string): ChatCompletionSystemMessageParam {
+	let content = CHAT_SYSTEM_PROMPT(get(userStore)?.username ?? '')
+
+	// If there's a custom prompt, append it to the system prompt
+	if (customPrompt?.trim()) {
+		content = `${content}\n\nUSER GIVEN INSTRUCTIONS:\n${customPrompt.trim()}`
+	}
+
 	return {
 		role: 'system',
-		content: CHAT_SYSTEM_PROMPT(get(userStore)?.username ?? '')
+		content
 	}
 }
 

frontend/src/lib/components/copilot/chat/ask/core.ts

@@ -25,10 +25,17 @@ GENERAL PRINCIPLES:
 
 export const askTools: Tool<{}>[] = [getDocumentationTool]
 
-export function prepareAskSystemMessage(): ChatCompletionSystemMessageParam {
+export function prepareAskSystemMessage(customPrompt?: string): ChatCompletionSystemMessageParam {
+	let content = CHAT_SYSTEM_PROMPT
+
+	// If there's a custom prompt, append it to the system prompt
+	if (customPrompt?.trim()) {
+		content = `${content}\n\nUSER GIVEN INSTRUCTIONS:\n${customPrompt.trim()}`
+	}
+
 	return {
 		role: 'system',
-		content: CHAT_SYSTEM_PROMPT
+		content
 	}
 }
 

frontend/src/lib/components/copilot/chat/flow/core.ts

@@ -885,8 +885,8 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 	}
 ]
 
-export function prepareFlowSystemMessage(): ChatCompletionSystemMessageParam {
-	const content = `You are a helpful assistant that creates and edits workflows on the Windmill platform. You're provided with a bunch of tools to help you edit the flow.
+export function prepareFlowSystemMessage(customPrompt?: string): ChatCompletionSystemMessageParam {
+	let content = `You are a helpful assistant that creates and edits workflows on the Windmill platform. You're provided with a bunch of tools to help you edit the flow.
 Follow the user instructions carefully.
 Go step by step, and explain what you're doing as you're doing it.
 DO NOT wait for user confirmation before performing an action. Only do it if the user explicitly asks you to wait in their initial instructions.
@@ -1012,6 +1012,11 @@ If the user needs a resource as flow input, you should set the property type in
 If the user wants a specific resource as step input, you should set the step value to a static string in the following format: "$res:path/to/resource".
 `
 
+	// If there's a custom prompt, append it to the system prompt
+	if (customPrompt?.trim()) {
+		content = `${content}\n\nUSER GIVEN INSTRUCTIONS:\n${customPrompt.trim()}`
+	}
+
 	return {
 		role: 'system',
 		content

frontend/src/lib/components/copilot/chat/navigator/core.ts

@@ -350,10 +350,17 @@ export const navigatorTools: Tool<{}>[] = [
 	getAvailableResourcesTool
 ]
 
-export function prepareNavigatorSystemMessage(): ChatCompletionSystemMessageParam {
+export function prepareNavigatorSystemMessage(customPrompt?: string): ChatCompletionSystemMessageParam {
+	let content = CHAT_SYSTEM_PROMPT
+
+	// If there's a custom prompt, append it to the system prompt
+	if (customPrompt?.trim()) {
+		content = `${content}\n\nUSER GIVEN INSTRUCTIONS:\n${customPrompt.trim()}`
+	}
+
 	return {
 		role: 'system',
-		content: CHAT_SYSTEM_PROMPT
+		content
 	}
 }
 

frontend/src/lib/components/copilot/chat/script/core.ts

@@ -452,10 +452,19 @@ WINDMILL LANGUAGE CONTEXT:
 
 `
 
-export function prepareScriptSystemMessage(): ChatCompletionSystemMessageParam {
+export function prepareScriptSystemMessage(
+	customPrompt?: string
+): ChatCompletionSystemMessageParam {
+	let content = CHAT_SYSTEM_PROMPT
+
+	// If there's a custom prompt, prepend it to the system prompt
+	if (customPrompt?.trim()) {
+		content = `${content}\n\nUSER GIVEN INSTRUCTIONS:\n${customPrompt.trim()}`
+	}
+
 	return {
 		role: 'system',
-		content: CHAT_SYSTEM_PROMPT
+		content
 	}
 }
 

frontend/src/lib/components/workspaceSettings/AISettings.svelte

@@ -14,6 +14,10 @@
 	import { safeSelectItems } from '../select/utils.svelte'
 	import Badge from '../common/badge/Badge.svelte'
 	import Tooltip from '../Tooltip.svelte'
+	import { AIMode } from '../copilot/chat/AIChatManager.svelte'
+	import ToggleButtonGroup from '../common/toggleButton-v2/ToggleButtonGroup.svelte'
+	import ToggleButton from '../common/toggleButton-v2/ToggleButton.svelte'
+	import autosize from '$lib/autosize'
 
 	const aiProviderLabels: [AIProvider, string][] = [
 		['openai', 'OpenAI'],
@@ -28,15 +32,19 @@
 		['customai', 'Custom AI']
 	]
 
+	const MAX_CUSTOM_PROMPT_LENGTH = 5000
+
 	let {
 		aiProviders = $bindable(),
 		codeCompletionModel = $bindable(),
 		defaultModel = $bindable(),
+		customPrompts = $bindable(),
 		usingOpenaiClientCredentialsOauth = $bindable()
 	}: {
 		aiProviders: Exclude<AIConfig['providers'], undefined>
 		codeCompletionModel: string | undefined
 		defaultModel: string | undefined
+		customPrompts: Record<string, string>
 		usingOpenaiClientCredentialsOauth: boolean
 	} = $props()
 
@@ -47,6 +55,9 @@
 		) as Record<AIProvider, string[]>
 	)
 
+	// Custom system prompt settings
+	let selectedAiMode = $state<AIMode>(AIMode.ASK)
+
 	let selectedAiModels = $derived(Object.values(aiProviders).flatMap((p) => p.models))
 	let modelProviderMap = $derived(
 		Object.fromEntries(
@@ -94,10 +105,16 @@
 				defaultModel && modelProviderMap[defaultModel]
 					? { model: defaultModel, provider: modelProviderMap[defaultModel] }
 					: undefined
+			// Convert customPrompts to include only non-empty prompts
+			const custom_prompts: Record<string, string> = Object.entries(customPrompts)
+				.filter(([_, prompt]) => prompt.trim().length > 0)
+				.reduce((acc, [mode, prompt]) => ({ ...acc, [mode]: prompt }), {})
+
 			const config: AIConfig = {
 				providers: aiProviders,
 				code_completion_model,
-				default_model
+				default_model,
+				custom_prompts: Object.keys(custom_prompts).length > 0 ? custom_prompts : undefined
 			}
 			await WorkspaceService.editCopilotConfig({
 				workspace: $workspaceStore!,
@@ -313,6 +330,58 @@
 		</div>
 	{/if}
 
+	{#if Object.keys(aiProviders).length > 0}
+		<div class="flex flex-col gap-2">
+			<p class="font-semibold">Custom system prompts</p>
+			<div class="flex flex-col gap-4">
+				<Label label="AI Mode">
+					<ToggleButtonGroup
+						bind:selected={selectedAiMode}
+						on:selected={({ detail }) => {
+							selectedAiMode = detail
+						}}
+					>
+						{#snippet children({ item })}
+							{#each Object.values(AIMode) as mode}
+								<div class="relative">
+									<ToggleButton
+										value={mode}
+										label={mode.charAt(0).toUpperCase() + mode.slice(1)}
+										{item}
+									/>
+									{#if customPrompts[mode]?.length > 0}
+										<div
+											class="absolute -top-1 -right-1 w-2 h-2 bg-blue-500 rounded-full border border-surface"
+										></div>
+									{/if}
+								</div>
+							{/each}
+						{/snippet}
+					</ToggleButtonGroup>
+				</Label>
+
+				<Label
+					label="Custom system prompt for {selectedAiMode.charAt(0).toUpperCase() +
+						selectedAiMode.slice(1)} Mode"
+				>
+					<textarea
+						bind:value={customPrompts[selectedAiMode]}
+						placeholder="Enter a custom system prompt for {selectedAiMode} mode."
+						class="w-full min-h-24 p-2 border border-gray-200 dark:border-gray-700 rounded-md bg-surface text-primary resize-y"
+						rows="4"
+						maxlength={MAX_CUSTOM_PROMPT_LENGTH}
+						use:autosize
+					></textarea>
+					<div class="flex justify-end mt-1">
+						<span class="text-xs text-secondary">
+							{(customPrompts[selectedAiMode] ?? '').length}/{MAX_CUSTOM_PROMPT_LENGTH} characters
+						</span>
+					</div>
+				</Label>
+			</div>
+		</div>
+	{/if}
+
 	<Button
 		wrapperClasses="self-start"
 		disabled={!Object.values(aiProviders).every((p) => p.resource_path) ||

frontend/src/lib/stores.ts

@@ -100,11 +100,13 @@ export const copilotInfo = writable<{
 	codeCompletionModel?: AIProviderModel
 	defaultModel?: AIProviderModel
 	aiModels: AIProviderModel[]
+	customPrompts?: Record<string, string>
 }>({
 	enabled: false,
 	codeCompletionModel: undefined,
 	defaultModel: undefined,
-	aiModels: []
+	aiModels: [],
+	customPrompts: {}
 })
 
 export async function loadCopilot(workspace: string) {
@@ -139,7 +141,8 @@ export function setCopilotInfo(aiConfig: AIConfig) {
 			enabled: true,
 			codeCompletionModel: aiConfig.code_completion_model,
 			defaultModel: aiConfig.default_model,
-			aiModels: aiModels
+			aiModels: aiModels,
+			customPrompts: aiConfig.custom_prompts ?? {}
 		})
 	} else {
 		copilotSessionModel.set(undefined)
@@ -148,7 +151,8 @@ export function setCopilotInfo(aiConfig: AIConfig) {
 			enabled: false,
 			codeCompletionModel: undefined,
 			defaultModel: undefined,
-			aiModels: []
+			aiModels: [],
+			customPrompts: {}
 		})
 	}
 }

frontend/src/routes/(root)/(logged)/workspace_settings/+page.svelte

@@ -19,9 +19,7 @@
 		ResourceService,
 		SettingService,
 		type AIConfig,
-
 		type ErrorHandler
-
 	} from '$lib/gen'
 	import {
 		enterpriseLicense,
@@ -57,6 +55,7 @@
 		convertDucklakeSettingsFromBackend,
 		type DucklakeSettingsType
 	} from '$lib/components/workspaceSettings/DucklakeSettings.svelte'
+	import { AIMode } from '$lib/components/copilot/chat/AIChatManager.svelte'
 
 	let slackInitialPath: string = $state('')
 	let slackScriptPath: string = $state('')
@@ -81,6 +80,7 @@
 	let aiProviders: Exclude<AIConfig['providers'], undefined> = $state({})
 	let codeCompletionModel: string | undefined = $state(undefined)
 	let defaultModel: string | undefined = $state(undefined)
+	let customPrompts: Record<string, string> = $state({})
 
 	let s3ResourceSettings: S3ResourceSettings = $state({
 		resourceType: 's3',
@@ -245,7 +245,12 @@
 		aiProviders = settings.ai_config?.providers ?? {}
 		defaultModel = settings.ai_config?.default_model?.model
 		codeCompletionModel = settings.ai_config?.code_completion_model?.model
-
+		customPrompts = settings.ai_config?.custom_prompts ?? {}
+		for (const mode of Object.values(AIMode)) {
+			if (!(mode in customPrompts)) {
+				customPrompts[mode] = ''
+			}
+		}
 		errorHandlerItemKind = settings.error_handler
 			? (settings.error_handler.split('/')[0] as 'flow' | 'script')
 			: 'script'
@@ -318,7 +323,7 @@
 				requestBody: {
 					error_handler: `${errorHandlerItemKind}/${errorHandlerScriptPath}`,
 					error_handler_extra_args: errorHandlerExtraArgs,
-					error_handler_muted_on_cancel: errorHandlerMutedOnCancel,
+					error_handler_muted_on_cancel: errorHandlerMutedOnCancel
 				}
 			})
 			sendUserToast(`workspace error handler set to ${errorHandlerScriptPath}`)
@@ -328,7 +333,7 @@
 				requestBody: {
 					error_handler: undefined,
 					error_handler_extra_args: undefined,
-					error_handler_muted_on_cancel: undefined,
+					error_handler_muted_on_cancel: undefined
 				}
 			})
 			sendUserToast(`workspace error handler removed`)
@@ -804,6 +809,7 @@
 				bind:aiProviders
 				bind:codeCompletionModel
 				bind:defaultModel
+				bind:customPrompts
 				bind:usingOpenaiClientCredentialsOauth
 			/>
 		{:else if tab == 'windmill_lfs'}