feat(aichat): add max tokens settings (#6613)

* add max tokens settings

* higher max

* fixes

* save max tokens in workspace settings

* cleaning

* cleaning

* cleaning

* feat(ai): add collapsible sections to ModelTokenLimits component

- Add collapsible/expandable sections for each AI provider
- Display 'Modified' badge when providers have custom settings
- Use ChevronDown/ChevronUp icons for visual feedback
- Add smooth slide transitions for better UX
- Reduce vertical space usage in workspace settings

Co-authored-by: centdix <centdix@users.noreply.github.com>

* adjust

* nit

---------

Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Co-authored-by: centdix <centdix@users.noreply.github.com>

backend/windmill-api/openapi.yaml

@@ -14514,6 +14514,12 @@ components:
           type: object
           additionalProperties:
             type: string
+        max_tokens_per_model:
+          type: object
+          additionalProperties:
+            type: integer
+            minimum: 1
+            maximum: 2000000
 
     Alert:
       type: object

backend/windmill-api/src/ai.rs

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

backend/windmill-api/src/workspaces.rs

@@ -923,6 +923,7 @@ async fn get_copilot_info(
             default_model: None,
             code_completion_model: None,
             custom_prompts: None,
+            max_tokens_per_model: None,
         }))
     }
 }

frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte

@@ -84,7 +84,9 @@
 		}
 	})
 
-	const isLastMessageTool = $derived(messages.length > 0 && messages[messages.length - 1].role === 'tool')
+	const isLastMessageTool = $derived(
+		messages.length > 0 && messages[messages.length - 1].role === 'tool'
+	)
 </script>
 
 <div class="flex flex-col h-full">

frontend/src/lib/components/copilot/chat/ProviderModelSelector.svelte

@@ -26,7 +26,7 @@
 			<div class="text-tertiary text-xs flex flex-row items-center font-normal gap-0.5">
 				<span class={`truncate ${multipleModels ? '' : 'pr-2'}`}>{providerModel.model}</span>
 				{#if multipleModels}
-					<div class="shrink-0 pr-1">
+					<div class="shrink-0">
 						<ChevronDown size={16} />
 					</div>
 				{/if}

frontend/src/lib/components/copilot/lib.ts

@@ -1,5 +1,6 @@
 import type { AIProvider, AIProviderModel } from '$lib/gen'
 import {
+	copilotInfo,
 	getCurrentModel,
 	workspaceStore,
 	type DBSchema,
@@ -157,7 +158,7 @@ export async function fetchAvailableModels(
 	return data?.data.map((m) => m.id) ?? []
 }
 
-function getModelMaxTokens(provider: AIProvider, model: string) {
+export function getModelMaxTokens(provider: AIProvider, model: string) {
 	if (model.startsWith('gpt-5')) {
 		return 128000
 	} else if ((provider === 'azure_openai' || provider === 'openai') && model.startsWith('o')) {
@@ -196,7 +197,10 @@ function getModelSpecificConfig(
 	modelProvider: AIProviderModel,
 	tools?: OpenAI.Chat.Completions.ChatCompletionTool[]
 ) {
-	const maxTokens = getModelMaxTokens(modelProvider.provider, modelProvider.model)
+	const defaultMaxTokens = getModelMaxTokens(modelProvider.provider, modelProvider.model)
+	const modelKey = `${modelProvider.provider}:${modelProvider.model}`
+	const customMaxTokensStore = get(copilotInfo)?.maxTokensPerModel
+	const maxTokens = customMaxTokensStore?.[modelKey] ?? defaultMaxTokens
 	if (
 		(modelProvider.provider === 'openai' || modelProvider.provider === 'azure_openai') &&
 		(modelProvider.model.startsWith('o') || modelProvider.model.startsWith('gpt-5'))

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

@@ -18,6 +18,7 @@
 	import ToggleButtonGroup from '../common/toggleButton-v2/ToggleButtonGroup.svelte'
 	import ToggleButton from '../common/toggleButton-v2/ToggleButton.svelte'
 	import autosize from '$lib/autosize'
+	import ModelTokenLimits from './ModelTokenLimits.svelte'
 
 	const MAX_CUSTOM_PROMPT_LENGTH = 5000
 
@@ -26,12 +27,14 @@
 		codeCompletionModel = $bindable(),
 		defaultModel = $bindable(),
 		customPrompts = $bindable(),
+		maxTokensPerModel = $bindable(),
 		usingOpenaiClientCredentialsOauth = $bindable()
 	}: {
 		aiProviders: Exclude<AIConfig['providers'], undefined>
 		codeCompletionModel: string | undefined
 		defaultModel: string | undefined
 		customPrompts: Record<string, string>
+		maxTokensPerModel: Record<string, number>
 		usingOpenaiClientCredentialsOauth: boolean
 	} = $props()
 
@@ -101,7 +104,8 @@
 				providers: aiProviders,
 				code_completion_model,
 				default_model,
-				custom_prompts: Object.keys(custom_prompts).length > 0 ? custom_prompts : undefined
+				custom_prompts: Object.keys(custom_prompts).length > 0 ? custom_prompts : undefined,
+				max_tokens_per_model: Object.keys(maxTokensPerModel).length > 0 ? maxTokensPerModel : undefined
 			}
 			await WorkspaceService.editCopilotConfig({
 				workspace: $workspaceStore!,
@@ -317,6 +321,10 @@
 		</div>
 	{/if}
 
+	{#if Object.keys(aiProviders).length > 0}
+		<ModelTokenLimits {aiProviders} bind:maxTokensPerModel />
+	{/if}
+
 	{#if Object.keys(aiProviders).length > 0}
 		<div class="flex flex-col gap-2">
 			<p class="font-semibold">Custom system prompts</p>

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

@@ -0,0 +1,192 @@
+<script lang="ts">
+	import type { AIConfig, AIProvider } from '$lib/gen'
+	import { getModelMaxTokens } from '../copilot/lib'
+	import { ChevronDown, ChevronUp } from 'lucide-svelte'
+	import { slide } from 'svelte/transition'
+
+	const MAX_TOKENS_LIMIT = 2000000
+
+	let {
+		aiProviders,
+		maxTokensPerModel = $bindable()
+	}: {
+		aiProviders: Exclude<AIConfig['providers'], undefined>
+		maxTokensPerModel: Record<string, number>
+	} = $props()
+
+	let errors = $state<Record<string, string>>({})
+	let collapsedProviders = $state<Record<string, boolean>>({})
+
+	// Group available models by provider
+	const modelsByProvider = $derived(
+		Object.entries(aiProviders).reduce(
+			(acc, [provider, config]) => {
+				acc[provider] = config.models.map((model) => ({
+					provider: provider as AIProvider,
+					model
+				}))
+				return acc
+			},
+			{} as Record<string, Array<{ provider: AIProvider; model: string }>>
+		)
+	)
+
+	function getModelKey(provider: AIProvider, model: string): string {
+		return `${provider}:${model}`
+	}
+
+	function getDefaultTokensForModel(provider: AIProvider, model: string): number {
+		return getModelMaxTokens(provider, model)
+	}
+
+	function getCurrentTokensForModel(provider: AIProvider, model: string): number {
+		const modelKey = getModelKey(provider, model)
+		return maxTokensPerModel[modelKey] ?? getDefaultTokensForModel(provider, model)
+	}
+
+	function updateTokensForModel(provider: AIProvider, model: string, tokens: number) {
+		const modelKey = getModelKey(provider, model)
+		if (tokens < 1 || tokens > MAX_TOKENS_LIMIT) {
+			errors[modelKey] = 'Token limit must be between 1 and ' + MAX_TOKENS_LIMIT
+			return
+		}
+
+		const defaultTokens = getDefaultTokensForModel(provider, model)
+
+		if (tokens === defaultTokens) {
+			// Remove from object if it's the default value
+			const newSettings = { ...maxTokensPerModel }
+			delete newSettings[modelKey]
+			maxTokensPerModel = newSettings
+		} else {
+			maxTokensPerModel = {
+				...maxTokensPerModel,
+				[modelKey]: tokens
+			}
+		}
+		errors[modelKey] = ''
+	}
+
+	function resetModelToDefault(provider: AIProvider, model: string) {
+		const modelKey = getModelKey(provider, model)
+		const newSettings = { ...maxTokensPerModel }
+		delete newSettings[modelKey]
+		maxTokensPerModel = newSettings
+	}
+
+	function isModelAtDefault(provider: AIProvider, model: string): boolean {
+		const currentTokens = getCurrentTokensForModel(provider, model)
+		const defaultTokens = getDefaultTokensForModel(provider, model)
+		return currentTokens === defaultTokens
+	}
+
+	function toggleProvider(provider: string) {
+		collapsedProviders[provider] = !collapsedProviders[provider]
+	}
+
+	function hasCustomSettings(provider: string, models: Array<{ model: string }>): boolean {
+		return models.some((m) => !isModelAtDefault(provider as AIProvider, m.model))
+	}
+
+	$effect(() => {
+		// Initialize collapsedProviders to true for all providers
+		collapsedProviders = {
+			...Object.fromEntries(Object.keys(aiProviders).map((provider) => [provider, true]))
+		}
+	})
+</script>
+
+{#if Object.keys(aiProviders).length > 0}
+	<div class="flex flex-col gap-4">
+		<div class="flex flex-col gap-1">
+			<p class="font-semibold">Model Output Limits</p>
+			<p class="text-xs text-secondary">
+				Configure maximum token limits for each model. These limits apply to all AI chat
+				interactions in the workspace.
+			</p>
+		</div>
+
+		<div class="flex flex-col gap-3">
+			{#each Object.entries(modelsByProvider).filter(([provider, models]) => models.length > 0) as [provider, models]}
+				{@const isExpanded = !collapsedProviders[provider]}
+				{@const hasCustom = hasCustomSettings(provider, models)}
+				<div class="border border-gray-200 dark:border-gray-700 rounded-md">
+					<button
+						type="button"
+						onclick={() => toggleProvider(provider)}
+						class="w-full px-4 py-3 flex items-center justify-between bg-surface-secondary hover:bg-surface-hover transition-colors rounded-md"
+					>
+						<div class="flex items-center gap-2">
+							<h4 class="font-medium text-sm capitalize">{provider}</h4>
+							{#if hasCustom}
+								<span
+									class="text-xs px-2 py-0.5 bg-blue-100 dark:bg-blue-900 text-blue-700 dark:text-blue-300 rounded"
+								>
+									Modified
+								</span>
+							{/if}
+						</div>
+						{#if isExpanded}
+							<ChevronUp size={16} class="text-gray-500" />
+						{:else}
+							<ChevronDown size={16} class="text-gray-500" />
+						{/if}
+					</button>
+
+					{#if isExpanded}
+						<div transition:slide|local={{ duration: 200 }} class="p-3">
+							<div class="space-y-3">
+								{#each models as { model }}
+									{@const currentTokens = getCurrentTokensForModel(provider as AIProvider, model)}
+									{@const defaultTokens = getDefaultTokensForModel(provider as AIProvider, model)}
+									{@const isAtDefault = isModelAtDefault(provider as AIProvider, model)}
+									<div class="flex flex-col gap-1">
+										<div class="flex items-center gap-3">
+											<div class="flex-1 min-w-0">
+												<span class="text-sm text-primary truncate block">{model}</span>
+											</div>
+											<div class="flex items-center gap-2">
+												<input
+													type="number"
+													min="1"
+													max={MAX_TOKENS_LIMIT}
+													value={currentTokens}
+													oninput={(e) => {
+														const value = parseInt(e.currentTarget.value)
+														if (!isNaN(value)) {
+															updateTokensForModel(provider as AIProvider, model, value)
+														}
+													}}
+													class="w-20 px-2 py-1 text-xs text-center border border-gray-200 dark:border-gray-700 rounded bg-surface focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-transparent"
+												/>
+												<span class="text-xs text-secondary whitespace-nowrap">tokens</span>
+											</div>
+										</div>
+										{#if !isAtDefault}
+											<div class="text-xs text-tertiary flex flex-row items-center gap-1">
+												<span>Default: {defaultTokens} tokens</span>
+												<button
+													type="button"
+													onclick={() => resetModelToDefault(provider as AIProvider, model)}
+													class="text-xs text-blue-600 hover:text-blue-800 dark:text-blue-400 dark:hover:text-blue-300 whitespace-nowrap"
+													title="Reset to default ({defaultTokens})"
+												>
+													Reset
+												</button>
+											</div>
+											{#if errors[getModelKey(provider as AIProvider, model)]}
+												<div class="text-xs text-red-500"
+													>{errors[getModelKey(provider as AIProvider, model)]}</div
+												>
+											{/if}
+										{/if}
+									</div>
+								{/each}
+							</div>
+						</div>
+					{/if}
+				</div>
+			{/each}
+		</div>
+	</div>
+{/if}

frontend/src/lib/stores.ts

@@ -102,12 +102,14 @@ export const copilotInfo = writable<{
 	defaultModel?: AIProviderModel
 	aiModels: AIProviderModel[]
 	customPrompts?: Record<string, string>
+	maxTokensPerModel?: Record<string, number>
 }>({
 	enabled: false,
 	codeCompletionModel: undefined,
 	defaultModel: undefined,
 	aiModels: [],
-	customPrompts: {}
+	customPrompts: {},
+	maxTokensPerModel: {}
 })
 
 export async function loadCopilot(workspace: string) {
@@ -143,7 +145,8 @@ export function setCopilotInfo(aiConfig: AIConfig) {
 			codeCompletionModel: aiConfig.code_completion_model,
 			defaultModel: aiConfig.default_model,
 			aiModels: aiModels,
-			customPrompts: aiConfig.custom_prompts ?? {}
+			customPrompts: aiConfig.custom_prompts ?? {},
+			maxTokensPerModel: aiConfig.max_tokens_per_model ?? {}
 		})
 	} else {
 		copilotSessionModel.set(undefined)
@@ -153,7 +156,8 @@ export function setCopilotInfo(aiConfig: AIConfig) {
 			codeCompletionModel: undefined,
 			defaultModel: undefined,
 			aiModels: [],
-			customPrompts: {}
+			customPrompts: {},
+			maxTokensPerModel: {}
 		})
 	}
 }
@@ -197,6 +201,7 @@ export const copilotSessionModel = writable<AIProviderModel | undefined>(
 			}
 		: undefined
 )
+
 export const usedTriggerKinds = writable<string[]>([])
 
 type SQLBaseSchema = {

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

@@ -81,6 +81,7 @@
 	let codeCompletionModel: string | undefined = $state(undefined)
 	let defaultModel: string | undefined = $state(undefined)
 	let customPrompts: Record<string, string> = $state({})
+	let maxTokensPerModel: Record<string, number> = $state({})
 
 	let s3ResourceSettings: S3ResourceSettings = $state({
 		resourceType: 's3',
@@ -246,6 +247,7 @@
 		defaultModel = settings.ai_config?.default_model?.model
 		codeCompletionModel = settings.ai_config?.code_completion_model?.model
 		customPrompts = settings.ai_config?.custom_prompts ?? {}
+		maxTokensPerModel = settings.ai_config?.max_tokens_per_model ?? {}
 		for (const mode of Object.values(AIMode)) {
 			if (!(mode in customPrompts)) {
 				customPrompts[mode] = ''
@@ -813,6 +815,7 @@
 				bind:codeCompletionModel
 				bind:defaultModel
 				bind:customPrompts
+				bind:maxTokensPerModel
 				bind:usingOpenaiClientCredentialsOauth
 			/>
 		{:else if tab == 'windmill_lfs'}