windmill-components 1.532.0 → 1.542.4

package/package/components/copilot/chat/AIChatManager.svelte.js

@@ -1,22 +1,23 @@
 import { flowTools, prepareFlowSystemMessage, prepareFlowUserMessage } from './flow/core';
 import ContextManager from './ContextManager.svelte';
 import HistoryManager from './HistoryManager.svelte';
-import { extractCodeFromMarkdown, getLatestAssistantMessage, processToolCall } from './shared';
+import {} from './shared';
 import { INLINE_CHAT_SYSTEM_PROMPT, prepareScriptSystemMessage, prepareScriptTools } from './script/core';
 import { navigatorTools, prepareNavigatorSystemMessage } from './navigator/core';
 import { loadApiTools } from './api/apiTools';
 import { prepareScriptUserMessage } from './script/core';
 import { prepareNavigatorUserMessage } from './navigator/core';
 import { sendUserToast } from '../../../toast';
-import { getCompletion, getModelContextWindow } from '../lib';
+import { getCompletion, getModelContextWindow, parseOpenAICompletion } from '../lib';
 import { dfs } from '../../flows/previousResults';
 import { getStringError } from './utils';
 import { untrack } from 'svelte';
-import { copilotSessionModel } from '../../../stores';
+import { get } from 'svelte/store';
+import { getCurrentModel, copilotInfo } from '../../../stores';
 import { askTools, prepareAskSystemMessage, prepareAskUserMessage } from './ask/core';
 import { chatState, DEFAULT_SIZE, triggerablesByAi } from './sharedChatState.svelte';
-import { get } from 'svelte/store';
 import { prepareApiSystemMessage, prepareApiUserMessage } from './api/core';
+import { getAnthropicCompletion, parseAnthropicCompletion } from './anthropic';
 // If the estimated token usage is greater than the model context window - the threshold, we delete the oldest message
 const MAX_TOKENS_THRESHOLD_PERCENTAGE = 0.05;
 const MAX_TOKENS_HARD_LIMIT = 5000;
@@ -83,7 +84,8 @@ class AIChatManager {
             }
             return acc;
         }, 0);
-        const modelContextWindow = getModelContextWindow(get(copilotSessionModel)?.model ?? '');
+        const model = getCurrentModel();
+        const modelContextWindow = getModelContextWindow(model.model);
         return (estimatedTokens >
             modelContextWindow -
                 Math.max(modelContextWindow * MAX_TOKENS_THRESHOLD_PERCENTAGE, MAX_TOKENS_HARD_LIMIT));
@@ -150,7 +152,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';
@@ -164,18 +167,8 @@ class AIChatManager {
                         args: this.scriptEditorOptions?.args ?? {}
                     };
                 },
-                getLastSuggestedCode: () => {
-                    const latestMessage = getLatestAssistantMessage(this.displayMessages);
-                    if (latestMessage) {
-                        const codeBlocks = extractCodeFromMarkdown(latestMessage);
-                        if (codeBlocks.length > 0) {
-                            return codeBlocks[codeBlocks.length - 1];
-                        }
-                    }
-                    return undefined;
-                },
-                applyCode: (code, applyAll) => {
-                    this.scriptEditorApplyCode?.(code, applyAll);
+                applyCode: (code, opts) => {
+                    this.scriptEditorApplyCode?.(code, opts);
                 }
             };
             if (options?.closeScriptSettings) {
@@ -186,23 +179,27 @@ 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 = {};
         }
@@ -298,9 +295,8 @@ class AIChatManager {
         }
     };
     chatRequest = async ({ messages, abortController, callbacks, systemMessage: systemMessageOverride }) => {
-        let addedMessages = [];
         try {
-            let completion = null;
+            let addedMessages = [];
             while (true) {
                 const systemMessage = systemMessageOverride ?? this.systemMessage;
                 const helpers = this.helpers;
@@ -324,87 +320,13 @@ class AIChatManager {
                     }
                     this.pendingPrompt = '';
                 }
-                completion = await getCompletion([systemMessage, ...messages, ...(pendingUserMessage ? [pendingUserMessage] : [])], abortController, tools.map((t) => t.def));
+                const model = getCurrentModel();
+                const completionFn = model.provider === 'anthropic' ? getAnthropicCompletion : getCompletion;
+                const parseFn = model.provider === 'anthropic' ? parseAnthropicCompletion : parseOpenAICompletion;
+                const completion = await completionFn([systemMessage, ...messages, ...(pendingUserMessage ? [pendingUserMessage] : [])], abortController, tools.map((t) => t.def));
                 if (completion) {
-                    const finalToolCalls = {};
-                    let answer = '';
-                    for await (const chunk of completion) {
-                        if (!('choices' in chunk && chunk.choices.length > 0 && 'delta' in chunk.choices[0])) {
-                            continue;
-                        }
-                        const c = chunk;
-                        const delta = c.choices[0].delta.content;
-                        if (delta) {
-                            answer += delta;
-                            callbacks.onNewToken(delta);
-                        }
-                        const toolCalls = c.choices[0].delta.tool_calls || [];
-                        if (toolCalls.length > 0 && answer) {
-                            // if tool calls are present but we have some textual content already, we need to display it to the user first
-                            callbacks.onMessageEnd();
-                            answer = '';
-                        }
-                        for (const toolCall of toolCalls) {
-                            const { index } = toolCall;
-                            let finalToolCall = finalToolCalls[index];
-                            if (!finalToolCall) {
-                                finalToolCalls[index] = toolCall;
-                            }
-                            else {
-                                if (toolCall.function?.arguments) {
-                                    if (!finalToolCall.function) {
-                                        finalToolCall.function = toolCall.function;
-                                    }
-                                    else {
-                                        finalToolCall.function.arguments =
-                                            (finalToolCall.function.arguments ?? '') + toolCall.function.arguments;
-                                    }
-                                }
-                            }
-                            finalToolCall = finalToolCalls[index];
-                            if (finalToolCall?.function) {
-                                const { function: { name: funcName }, id: toolCallId } = finalToolCall;
-                                if (funcName && toolCallId) {
-                                    const tool = tools.find((t) => t.def.function.name === funcName);
-                                    if (tool && tool.preAction) {
-                                        tool.preAction({ toolCallbacks: callbacks, toolId: toolCallId });
-                                    }
-                                }
-                            }
-                        }
-                    }
-                    if (answer) {
-                        const toAdd = { role: 'assistant', content: answer };
-                        addedMessages.push(toAdd);
-                        messages.push(toAdd);
-                    }
-                    callbacks.onMessageEnd();
-                    const toolCalls = Object.values(finalToolCalls).filter((toolCall) => toolCall.id !== undefined && toolCall.function?.arguments !== undefined);
-                    if (toolCalls.length > 0) {
-                        const toAdd = {
-                            role: 'assistant',
-                            tool_calls: toolCalls.map((t) => ({
-                                ...t,
-                                function: {
-                                    ...t.function,
-                                    arguments: t.function.arguments || '{}'
-                                }
-                            }))
-                        };
-                        messages.push(toAdd);
-                        addedMessages.push(toAdd);
-                        for (const toolCall of toolCalls) {
-                            const messageToAdd = await processToolCall({
-                                tools,
-                                toolCall,
-                                helpers,
-                                toolCallbacks: callbacks
-                            });
-                            messages.push(messageToAdd);
-                            addedMessages.push(messageToAdd);
-                        }
-                    }
-                    else {
+                    const continueCompletion = await parseFn(completion, callbacks, messages, addedMessages, tools, helpers);
+                    if (!continueCompletion) {
                         break;
                     }
                 }

package/package/components/copilot/chat/AssistantMessage.svelte

@@ -2,11 +2,7 @@
 import { gfmPlugin } from 'svelte-exmarkdown/gfm';
 import CodeDisplay from './script/CodeDisplay.svelte';
 import LinkRenderer from './LinkRenderer.svelte';
-import { setContext } from 'svelte';
 export let message;
-setContext('AssistantMessageContext', {
-    message
-});
 </script>
 
 <div

package/package/components/copilot/chat/AvailableContextList.svelte

@@ -1,16 +1,18 @@
 <script lang="ts">import FlowModuleIcon from '../../flows/FlowModuleIcon.svelte';
 import BarsStaggered from '../../icons/BarsStaggered.svelte';
 import { ContextIconMap } from './context';
-import { ArrowLeft, Diff, Database, Code, ChevronRight } from 'lucide-svelte';
-const { availableContext, selectedContext, onSelect, showAllAvailable = false, stringSearch = '', selectedIndex = 0, categorize = false } = $props();
+import { ArrowLeft, Diff, Database, ChevronRight } from 'lucide-svelte';
+const { availableContext, selectedContext, onSelect, setShowing, showAllAvailable = false, stringSearch = '', onViewChange } = $props();
 // Current view state: 'categories' or specific category type
 let currentView = $state('categories');
+// Selected index for keyboard navigation
+let itemSelectedIndex = $state(0);
+let categorySelectedIndex = $state(0);
 // Category definitions
 const categories = [
     { id: 'diffs', label: 'Diffs', icon: Diff },
     { id: 'modules', label: 'Modules', icon: BarsStaggered },
-    { id: 'databases', label: 'Databases', icon: Database },
-    { id: 'code', label: 'Code', icon: Code }
+    { id: 'databases', label: 'Databases', icon: Database }
 ];
 const filteredAvailableContext = $derived(availableContext.filter((context) => {
     const filtered = (showAllAvailable ||
@@ -23,8 +25,7 @@ const contextByCategory = $derived.by(() => {
     const grouped = {
         diffs: [],
         modules: [],
-        databases: [],
-        code: []
+        databases: []
     };
     filteredAvailableContext.forEach((context) => {
         if (context.type === 'diff')
@@ -33,82 +34,149 @@ const contextByCategory = $derived.by(() => {
             grouped.modules.push(context);
         else if (context.type === 'db')
             grouped.databases.push(context);
-        else if (context.type === 'code')
-            grouped.code.push(context);
     });
     return grouped;
 });
 const currentCategoryItems = $derived(currentView !== 'categories' ? contextByCategory[currentView] : []);
+// Filter to only show categories with items
+const availableCategories = $derived(categories.filter((cat) => contextByCategory[cat.id].length > 0));
+// Report view changes
+$effect(() => {
+    if (onViewChange) {
+        if (currentView === 'categories') {
+            onViewChange(availableCategories.length);
+        }
+        else {
+            onViewChange(currentCategoryItems.length + 1);
+        }
+    }
+});
 function handleCategoryClick(categoryId) {
     currentView = categoryId;
 }
 function handleBackClick() {
     currentView = 'categories';
+    itemSelectedIndex = 0;
+}
+function handleKeyDown(e) {
+    if (stringSearch.length > 0) {
+        // Navigation in search view (flat list)
+        if (e.key === 'ArrowDown') {
+            e.preventDefault();
+            e.stopPropagation();
+            if (filteredAvailableContext.length > 0) {
+                itemSelectedIndex = (itemSelectedIndex + 1) % filteredAvailableContext.length;
+            }
+        }
+        else if (e.key === 'ArrowUp') {
+            e.preventDefault();
+            e.stopPropagation();
+            if (filteredAvailableContext.length > 0) {
+                itemSelectedIndex =
+                    (itemSelectedIndex - 1 + filteredAvailableContext.length) %
+                        filteredAvailableContext.length;
+            }
+        }
+        else if (e.key === 'Enter' || e.key === 'Tab') {
+            if (e.key === 'Tab')
+                e.preventDefault();
+            e.stopPropagation();
+            const selectedItem = filteredAvailableContext[itemSelectedIndex];
+            if (selectedItem) {
+                onSelect(selectedItem);
+            }
+        }
+    }
+    else if (currentView === 'categories') {
+        // Navigation in categories view
+        if (e.key === 'ArrowDown') {
+            e.preventDefault();
+            e.stopPropagation();
+            categorySelectedIndex = (categorySelectedIndex + 1) % availableCategories.length;
+        }
+        else if (e.key === 'ArrowUp') {
+            e.preventDefault();
+            e.stopPropagation();
+            categorySelectedIndex =
+                (categorySelectedIndex - 1 + availableCategories.length) % availableCategories.length;
+        }
+        else if (e.key === 'Enter' || e.key === 'ArrowRight' || e.key === 'Tab') {
+            e.preventDefault();
+            e.stopPropagation();
+            const selectedCategory = availableCategories[categorySelectedIndex];
+            if (selectedCategory) {
+                handleCategoryClick(selectedCategory.id);
+            }
+        }
+        else if (e.key === 'Escape' || e.key === 'ArrowLeft') {
+            e.preventDefault();
+            e.stopPropagation();
+            setShowing?.(false);
+        }
+    }
+    else {
+        // Navigation in category items view
+        if (e.key === 'ArrowDown') {
+            e.preventDefault();
+            e.stopPropagation();
+            if (currentCategoryItems.length > 0) {
+                itemSelectedIndex = (itemSelectedIndex + 1) % currentCategoryItems.length;
+            }
+        }
+        else if (e.key === 'ArrowUp') {
+            e.preventDefault();
+            e.stopPropagation();
+            if (currentCategoryItems.length > 0) {
+                itemSelectedIndex =
+                    (itemSelectedIndex - 1 + currentCategoryItems.length) % currentCategoryItems.length;
+            }
+        }
+        else if (e.key === 'Enter' || e.key === 'Tab') {
+            if (e.key === 'Tab')
+                e.preventDefault();
+            e.stopPropagation();
+            const selectedItem = currentCategoryItems[itemSelectedIndex];
+            if (selectedItem) {
+                onSelect(selectedItem);
+                currentView = 'categories'; // Go back to categories after selection
+            }
+        }
+        else if (e.key === 'ArrowLeft' || e.key === 'Escape') {
+            e.preventDefault();
+            e.stopPropagation();
+            handleBackClick();
+        }
+    }
 }
+// Listen for keyboard events
+$effect(() => {
+    document.addEventListener('keydown', handleKeyDown);
+    return () => {
+        document.removeEventListener('keydown', handleKeyDown);
+    };
+});
+$effect(() => {
+    if (stringSearch.length > 0) {
+        itemSelectedIndex = 0;
+    }
+});
 </script>
 
-<div class="flex flex-col gap-1 text-tertiary text-xs p-1 pr-0 min-w-24 max-h-48 overflow-y-scroll">
-	{#if categorize}
-		{#if currentView === 'categories'}
-			{#each categories as category}
-				{@const itemCount = contextByCategory[category.id].length}
-				{@const Icon = category.icon}
-				{#if itemCount > 0}
-					<button
-						class="hover:bg-surface-hover rounded-md p-1 pr-0 text-left flex flex-row gap-1 items-center font-normal transition-colors"
-						onclick={() => handleCategoryClick(category.id)}
-					>
-						<Icon size={16} />
-						<span class="flex-1">{category.label}</span>
-						<ChevronRight size={16} />
-					</button>
-				{/if}
-			{/each}
-			{#if categories.every((cat) => contextByCategory[cat.id].length === 0)}
-				<div class="text-center text-tertiary text-xs py-2">No available context</div>
-			{/if}
-		{:else}
-			<!-- Category items view -->
-			<button
-				class="hover:bg-surface-hover rounded-md text-left flex flex-row gap-1 items-center font-normal transition-colors mb-1"
-				onclick={handleBackClick}
-			>
-				<ArrowLeft size={12} />
-				<span class="text-xs">Go back</span>
-			</button>
-
-			{#if currentCategoryItems.length === 0}
-				<div class="text-center text-tertiary text-xs py-2">No items in this category</div>
-			{:else}
-				{#each currentCategoryItems as element}
-					{@const Icon = ContextIconMap[element.type]}
-					<button
-						class="hover:bg-surface-hover rounded-md p-1 text-left flex flex-row gap-1 items-center font-normal transition-colors"
-						onclick={() => {
-							onSelect(element)
-							currentView = 'categories' // Go back to categories after selection
-						}}
-					>
-						{#if element.type === 'flow_module'}
-							<FlowModuleIcon module={element as FlowModule} size={16} />
-						{:else if Icon}
-							<Icon size={16} />
-						{/if}
-						<span class="truncate">
-							{element.type === 'diff' || element.type === 'flow_module'
-								? element.title.replace(/_/g, ' ')
-								: element.title}
-						</span>
-					</button>
-				{/each}
-			{/if}
-		{/if}
-	{:else}
+<div
+	class="flex flex-col gap-1 text-tertiary text-xs p-1 pr-0 min-w-24 max-h-48 overflow-y-scroll"
+	onmousedown={(e) =>
+		// avoids triggering onblur on the textinput and closing the tooltip
+		e.preventDefault()}
+	role="listbox"
+	tabindex={0}
+>
+	{#if stringSearch.length > 0}
+		<!-- Search view - show flat list -->
 		{#each filteredAvailableContext as element, i}
 			{@const Icon = ContextIconMap[element.type]}
 			<button
 				class="hover:bg-surface-hover rounded-md p-1 text-left flex flex-row gap-1 items-center font-normal transition-colors {i ===
-				selectedIndex
+				itemSelectedIndex
 					? 'bg-surface-hover'
 					: ''}"
 				onclick={() => {
@@ -127,5 +195,63 @@ function handleBackClick() {
 				</span>
 			</button>
 		{/each}
+		{#if filteredAvailableContext.length === 0}
+			<div class="text-center text-tertiary text-xs py-2">No matching context</div>
+		{/if}
+	{:else if currentView === 'categories'}
+		<!-- Categories view -->
+		{#each availableCategories as category, i}
+			{@const Icon = category.icon}
+			<button
+				class="hover:bg-surface-hover rounded-md p-1 pr-0 text-left flex flex-row gap-1 items-center font-normal transition-colors {i ===
+				categorySelectedIndex
+					? 'bg-surface-hover'
+					: ''}"
+				onclick={() => handleCategoryClick(category.id)}
+			>
+				<Icon size={16} />
+				<span class="flex-1">{category.label}</span>
+				<ChevronRight size={16} />
+			</button>
+		{/each}
+		{#if availableCategories.length === 0}
+			<div class="text-center text-tertiary text-xs py-2">No available context</div>
+		{/if}
+	{:else}
+		<!-- Category items view -->
+		<button
+			class="hover:bg-surface-hover rounded-md text-left flex flex-row gap-1 items-center font-normal transition-colors mb-1"
+			onclick={handleBackClick}
+		>
+			<ArrowLeft size={12} />
+			<span class="text-xs">Go back</span>
+		</button>
+
+		{#if currentCategoryItems.length === 0}
+			<div class="text-center text-tertiary text-xs py-2">No items in this category</div>
+		{:else}
+			{#each currentCategoryItems as element, i}
+				{@const Icon = ContextIconMap[element.type]}
+				<button
+					class="hover:bg-surface-hover rounded-md p-1 text-left flex flex-row gap-1 items-center font-normal transition-colors {i ===
+					itemSelectedIndex
+						? 'bg-surface-hover'
+						: ''}"
+					onclick={() => {
+						onSelect(element)
+						currentView = 'categories' // Go back to categories after selection
+					}}
+				>
+					{#if element.type === 'flow_module'}
+						<FlowModuleIcon module={element as FlowModule} size={16} />
+					{:else if Icon}
+						<Icon size={16} />
+					{/if}
+					<span class="truncate">
+						{element.type === 'diff' ? element.title.replace(/_/g, ' ') : element.title}
+					</span>
+				</button>
+			{/each}
+		{/if}
 	{/if}
 </div>

package/package/components/copilot/chat/AvailableContextList.svelte.d.ts

@@ -3,10 +3,10 @@ interface Props {
     availableContext: ContextElement[];
     selectedContext: ContextElement[];
     onSelect: (element: ContextElement) => void;
+    setShowing?: (showing: boolean) => void;
     showAllAvailable?: boolean;
     stringSearch?: string;
-    selectedIndex?: number;
-    categorize?: boolean;
+    onViewChange?: (newNumber: number) => void;
 }
 declare const AvailableContextList: import("svelte").Component<Props, {}, "">;
 type AvailableContextList = ReturnType<typeof AvailableContextList>;

package/package/components/copilot/chat/ContextElementBadge.svelte

@@ -29,7 +29,7 @@ const isDeletable = $derived(deletable && contextElement.deletable !== false);
 			<button onclick={isDeletable ? onDelete : undefined} class:cursor-default={!isDeletable}>
 				{#if showDelete && isDeletable}
 					<X size={16} />
-				{:else if contextElement.type === 'flow_module'}
+				{:else if contextElement.type === 'flow_module' || contextElement.type === 'flow_module_code_piece'}
 					<FlowModuleIcon module={contextElement as FlowModule} size={16} />
 				{:else}
 					{@const SvelteComponent = icon}
@@ -37,7 +37,7 @@ const isDeletable = $derived(deletable && contextElement.deletable !== false);
 				{/if}
 			</button>
 			<span class="truncate">
-				{contextElement.type === 'diff' || contextElement.type === 'flow_module'
+				{contextElement.type === 'diff'
 					? contextElement.title.replace(/_/g, ' ')
 					: contextElement.title}
 			</span>
@@ -65,7 +65,7 @@ const isDeletable = $derived(deletable && contextElement.deletable !== false);
 					<div class="text-tertiary">Not loaded yet</div>
 				{/if}
 			</div>
-		{:else if contextElement.type === 'code' || contextElement.type === 'code_piece' || contextElement.type === 'diff'}
+		{:else if contextElement.type === 'code' || contextElement.type === 'code_piece' || contextElement.type === 'diff' || contextElement.type === 'flow_module_code_piece'}
 			<div class="max-w-96 max-h-[300px] text-xs overflow-auto">
 				<HighlightCode
 					language={contextElement.lang}

package/package/components/copilot/chat/ContextManager.svelte.js

@@ -66,7 +66,7 @@ export default class ContextManager {
                 newAvailableContext.push({
                     type: 'flow_module',
                     id: module.id,
-                    title: `module_[${module.id}]`,
+                    title: `${module.id}`,
                     value: {
                         language: 'language' in module.value ? module.value.language : 'bunnative',
                         path: 'path' in module.value ? module.value.path : '',
@@ -215,22 +215,45 @@ export default class ContextManager {
         this.scriptOptions = scriptOptions;
     }
     addSelectedLinesToContext(lines, startLine, endLine, moduleId) {
-        const title = moduleId ? `[${moduleId}] L${startLine}-L${endLine}` : `L${startLine}-L${endLine}`;
+        const title = moduleId ? `${moduleId} L${startLine}-L${endLine}` : `L${startLine}-L${endLine}`;
         if (!this.scriptOptions ||
-            this.selectedContext.find((c) => c.type === 'code_piece' && c.title === title)) {
+            this.selectedContext.find((c) => (c.type === 'code_piece' && c.title === title) ||
+                (c.type === 'flow_module_code_piece' && c.id === moduleId && c.title === title))) {
             return;
         }
-        this.selectedContext = [
-            ...this.selectedContext,
-            {
-                type: 'code_piece',
-                title: title,
-                startLine,
-                endLine,
-                content: lines,
-                lang: this.scriptOptions.lang
+        if (moduleId) {
+            const module = [...this.availableContext, ...this.selectedContext].find((c) => c.type === 'flow_module' && c.id === moduleId);
+            if (!module) {
+                console.error('Module not found', moduleId);
+                return;
             }
-        ];
+            this.selectedContext = [
+                ...this.selectedContext,
+                {
+                    type: 'flow_module_code_piece',
+                    id: moduleId,
+                    title: title,
+                    startLine,
+                    endLine,
+                    content: lines,
+                    lang: this.scriptOptions.lang,
+                    value: module.value
+                }
+            ];
+        }
+        else {
+            this.selectedContext = [
+                ...this.selectedContext,
+                {
+                    type: 'code_piece',
+                    title: title,
+                    startLine,
+                    endLine,
+                    content: lines,
+                    lang: this.scriptOptions.lang
+                }
+            ];
+        }
     }
     setFixContext() {
         const codeContext = this.availableContext.find((c) => c.type === 'code');