From 7df5f5453f39afccecd859f80da4fa2d3edec21d Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Tue, 2 Sep 2025 11:27:37 +0200 Subject: [PATCH] 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 | 4 ++ backend/windmill-api/src/ai.rs | 2 + backend/windmill-api/src/workspaces.rs | 16 +++++ .../copilot/chat/AIChatManager.svelte.ts | 18 +++-- .../lib/components/copilot/chat/api/core.ts | 11 ++- .../lib/components/copilot/chat/ask/core.ts | 11 ++- .../lib/components/copilot/chat/flow/core.ts | 9 ++- .../components/copilot/chat/navigator/core.ts | 11 ++- .../components/copilot/chat/script/core.ts | 13 +++- .../workspaceSettings/AISettings.svelte | 71 ++++++++++++++++++- frontend/src/lib/stores.ts | 10 ++- .../(logged)/workspace_settings/+page.svelte | 16 +++-- 12 files changed, 167 insertions(+), 25 deletions(-) diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml index 2fa7ff8c9a..68a4781225 100644 --- a/backend/windmill-api/openapi.yaml +++ b/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 diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index a41840404d..fc60c65c3c 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -353,6 +353,8 @@ pub struct AIConfig { pub default_model: Option, #[serde(skip_serializing_if = "Option::is_none")] pub code_completion_model: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub custom_prompts: Option>, } pub fn global_service() -> Router { diff --git a/backend/windmill-api/src/workspaces.rs b/backend/windmill-api/src/workspaces.rs index fb0ea74f26..c2b9e448c2 100644 --- a/backend/windmill-api/src/workspaces.rs +++ b/backend/windmill-api/src/workspaces.rs @@ -720,6 +720,7 @@ async fn edit_deploy_to() -> Result { } 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 { let domain = email.split('@').last().unwrap(); @@ -819,6 +820,20 @@ async fn edit_copilot_config( ) -> Result { 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, })) } } diff --git a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts index dcd91e1d1f..ed7ea56157 100644 --- a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts +++ b/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 = {} } diff --git a/frontend/src/lib/components/copilot/chat/api/core.ts b/frontend/src/lib/components/copilot/chat/api/core.ts index cd7038595a..4e47baea72 100644 --- a/frontend/src/lib/components/copilot/chat/api/core.ts +++ b/frontend/src/lib/components/copilot/chat/api/core.ts @@ -57,10 +57,17 @@ export async function getApiTools(): Promise[]> { 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 } } diff --git a/frontend/src/lib/components/copilot/chat/ask/core.ts b/frontend/src/lib/components/copilot/chat/ask/core.ts index 997594cc88..f9ba219599 100644 --- a/frontend/src/lib/components/copilot/chat/ask/core.ts +++ b/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 } } diff --git a/frontend/src/lib/components/copilot/chat/flow/core.ts b/frontend/src/lib/components/copilot/chat/flow/core.ts index dae637f156..1fdd8895ff 100644 --- a/frontend/src/lib/components/copilot/chat/flow/core.ts +++ b/frontend/src/lib/components/copilot/chat/flow/core.ts @@ -885,8 +885,8 @@ export const flowTools: Tool[] = [ } ] -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 diff --git a/frontend/src/lib/components/copilot/chat/navigator/core.ts b/frontend/src/lib/components/copilot/chat/navigator/core.ts index f7488b3fa8..4c9cf38b70 100644 --- a/frontend/src/lib/components/copilot/chat/navigator/core.ts +++ b/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 } } diff --git a/frontend/src/lib/components/copilot/chat/script/core.ts b/frontend/src/lib/components/copilot/chat/script/core.ts index 9deb4f0412..ca9bde59e4 100644 --- a/frontend/src/lib/components/copilot/chat/script/core.ts +++ b/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 } } diff --git a/frontend/src/lib/components/workspaceSettings/AISettings.svelte b/frontend/src/lib/components/workspaceSettings/AISettings.svelte index 1f9127c8fd..b4cbd5eb09 100644 --- a/frontend/src/lib/components/workspaceSettings/AISettings.svelte +++ b/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 codeCompletionModel: string | undefined defaultModel: string | undefined + customPrompts: Record usingOpenaiClientCredentialsOauth: boolean } = $props() @@ -47,6 +55,9 @@ ) as Record ) + // Custom system prompt settings + let selectedAiMode = $state(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 = 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 @@ {/if} + {#if Object.keys(aiProviders).length > 0} +
+

Custom system prompts

+
+ + + +
+
+ {/if} +