From 00fbacdab47e4d773b33078b366d648618a0e6c8 Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@windmill.dev>
Date: Thu, 29 Feb 2024 20:22:44 +0100
Subject: [PATCH] fix: AI copilot available in the vscode flow editor (#3314)

* all

* foo

* add ai to vscode for flows

* add ai to vscode for flows

* foo

* foo
---
 frontend/src/lib/components/Dev.svelte        | 79 +++++++++++++++++--
 .../src/lib/components/copilot/StepGen.svelte |  3 +-
 frontend/src/lib/components/copilot/lib.ts    |  4 +-
 .../content/FlowBranchesAllWrapper.svelte     | 21 ++---
 .../content/FlowBranchesOneWrapper.svelte     | 20 ++---
 .../components/flows/content/FlowLoop.svelte  | 23 +++---
 .../flows/content/FlowModuleComponent.svelte  |  3 +-
 .../flows/content/FlowModuleWrapper.svelte    |  5 ++
 .../flows/map/InsertModuleButton.svelte       |  4 +-
 9 files changed, 121 insertions(+), 41 deletions(-)

diff --git a/frontend/src/lib/components/Dev.svelte b/frontend/src/lib/components/Dev.svelte
index d77de3a91f..b90160eed5 100644
--- a/frontend/src/lib/components/Dev.svelte
+++ b/frontend/src/lib/components/Dev.svelte
@@ -1,7 +1,7 @@
 <script lang="ts">
 	import SchemaForm from '$lib/components/SchemaForm.svelte'
 	import TestJobLoader from '$lib/components/TestJobLoader.svelte'
-	import { Button } from '$lib/components/common'
+	import { Button, Drawer } from '$lib/components/common'
 	import { WindmillIcon } from '$lib/components/icons'
 	import LogPanel from '$lib/components/scriptEditor/LogPanel.svelte'
 	import {
@@ -11,10 +11,14 @@
 		OpenAPI,
 		Preview,
 		type OpenFlow,
-		type FlowModule
+		type FlowModule,
+		WorkspaceService,
+		type InputTransform,
+		RawScript,
+		type PathScript
 	} from '$lib/gen'
 	import { inferArgs } from '$lib/infer'
-	import { userStore, workspaceStore } from '$lib/stores'
+	import { copilotInfo, userStore, workspaceStore } from '$lib/stores'
 	import { emptySchema, sendUserToast } from '$lib/utils'
 	import { Pane, Splitpanes } from 'svelte-splitpanes'
 	import { onDestroy, onMount, setContext } from 'svelte'
@@ -34,6 +38,9 @@
 	import { CornerDownLeft, Play } from 'lucide-svelte'
 	import Toggle from './Toggle.svelte'
 	import { setLicense } from '$lib/enterpriseUtils'
+	import { workspacedOpenai } from './copilot/lib'
+	import type { FlowCopilotContext, FlowCopilotModule } from './copilot/flow'
+	import { pickScript } from './flows/flowStateUtils'
 
 	$: token = $page.url.searchParams.get('wm_token') ?? undefined
 	$: workspace = $page.url.searchParams.get('workspace') ?? undefined
@@ -45,8 +52,70 @@
 		OpenAPI.TOKEN = $page.url.searchParams.get('wm_token')!
 	}
 
+	let flowCopilotContext: FlowCopilotContext = {
+		drawerStore: writable<Drawer | undefined>(undefined),
+		modulesStore: writable<FlowCopilotModule[]>([]),
+		currentStepStore: writable<string | undefined>(undefined),
+		genFlow,
+		shouldUpdatePropertyType: writable<{
+			[key: string]: 'static' | 'javascript' | undefined
+		}>({}),
+		exprsToSet: writable<{
+			[key: string]: InputTransform | any | undefined
+		}>({}),
+		generatedExprs: writable<{
+			[key: string]: string | undefined
+		}>({}),
+		stepInputsLoading: writable<boolean>(false)
+	}
+	const { modulesStore } = flowCopilotContext
+
+	async function genFlow(idx: number, flowModules: FlowModule[], stepOnly = false) {
+		let module = stepOnly ? $modulesStore[0] : $modulesStore[idx]
+
+		if (module && module.selectedCompletion) {
+			const [hubScriptModule, hubScriptState] = await pickScript(
+				module.selectedCompletion.path,
+				`${module.selectedCompletion.summary} (${module.selectedCompletion.app})`,
+				module.id,
+				undefined
+			)
+			const flowModule: FlowModule & {
+				value: RawScript | PathScript
+			} = {
+				id: module.id,
+				value: hubScriptModule.value,
+				summary: hubScriptModule.summary
+			}
+
+			$flowStateStore[module.id] = hubScriptState
+
+			flowModules.splice(idx, 0, flowModule)
+			$flowStore = $flowStore
+			sendUserToast('Added module', false)
+		}
+	}
+
+	setContext('FlowCopilotContext', flowCopilotContext)
+
+	async function setCopilotInfo() {
+		if (workspace) {
+			workspacedOpenai.init(workspace, token)
+			try {
+				copilotInfo.set(await WorkspaceService.getCopilotInfo({ workspace }))
+			} catch (err) {
+				copilotInfo.set({
+					exists_openai_resource_path: false,
+					code_completion_enabled: false
+				})
+
+				console.error('Could not get copilot info')
+			}
+		}
+	}
 	$: if (workspace) {
 		$workspaceStore = workspace
+		setCopilotInfo()
 	}
 
 	$: if (workspace && token) {
@@ -242,7 +311,7 @@
 		}
 	}
 
-	let mode: 'script' | 'flow' = 'script'
+	let mode: 'script' | 'flow' = 'flow' //'script'
 
 	const flowStore = writable({
 		summary: '',
@@ -478,7 +547,7 @@
 					</Pane>
 					<Pane size={33}>
 						{#key reload}
-							<FlowEditorPanel noEditor />
+							<FlowEditorPanel enableAi noEditor />
 						{/key}
 					</Pane>
 				</Splitpanes>
diff --git a/frontend/src/lib/components/copilot/StepGen.svelte b/frontend/src/lib/components/copilot/StepGen.svelte
index 3cdecd599e..72fc3ea187 100644
--- a/frontend/src/lib/components/copilot/StepGen.svelte
+++ b/frontend/src/lib/components/copilot/StepGen.svelte
@@ -17,6 +17,7 @@
 	export let funcDesc: string
 	export let modules: FlowModule[]
 	export let trigger = false
+	export let disableAi = false
 
 	// state
 	let input: HTMLInputElement | undefined
@@ -131,7 +132,7 @@
 				/>
 			{/if}
 		</div>
-		{#if funcDesc.length > 0}
+		{#if !disableAi && funcDesc.length > 0}
 			<ul class="transition-all divide-y">
 				<li>
 					<button
diff --git a/frontend/src/lib/components/copilot/lib.ts b/frontend/src/lib/components/copilot/lib.ts
index 9e6a1381ca..cbd9ae5249 100644
--- a/frontend/src/lib/components/copilot/lib.ts
+++ b/frontend/src/lib/components/copilot/lib.ts
@@ -27,13 +27,13 @@ const openaiConfig: ChatCompletionCreateParamsStreaming = {
 class WorkspacedOpenai {
 	private client: OpenAI | undefined
 
-	init(workspace: string) {
+	init(workspace: string, token: string | undefined = undefined) {
 		const baseURL = `${location.origin}${OpenAPI.BASE}/w/${workspace}/openai/proxy`
 		this.client = new OpenAI({
 			baseURL,
 			apiKey: 'fakekey',
 			defaultHeaders: {
-				Authorization: ''
+				Authorization: token ? `Bearer ${token}` : ''
 			},
 			dangerouslyAllowBrowser: true
 		})
diff --git a/frontend/src/lib/components/flows/content/FlowBranchesAllWrapper.svelte b/frontend/src/lib/components/flows/content/FlowBranchesAllWrapper.svelte
index c093b17faa..14d5673af7 100644
--- a/frontend/src/lib/components/flows/content/FlowBranchesAllWrapper.svelte
+++ b/frontend/src/lib/components/flows/content/FlowBranchesAllWrapper.svelte
@@ -30,16 +30,17 @@
 		<SplitPanesWrapper>
 			<Splitpanes horizontal>
 				<Pane size={flowModule ? 60 : 100}>
-					<Alert
-						type="info"
-						title="All branches will be run"
-						tooltip="Branch all"
-						documentationLink="https://www.windmill.dev/docs/flows/flow_branches#branch-all"
-						class="m-4"
-					>
-						The result of this step is the list of the result of each branch.
-					</Alert>
-
+					{#if !noEditor}
+						<Alert
+							type="info"
+							title="All branches will be run"
+							tooltip="Branch all"
+							documentationLink="https://www.windmill.dev/docs/flows/flow_branches#branch-all"
+							class="m-4"
+						>
+							The result of this step is the list of the result of each branch.
+						</Alert>
+					{/if}
 					<div class="p-4 mt-4 w-full">
 						<h3 class="mb-4"
 							>{value.branches.length} branch{value.branches.length > 1 ? 'es' : ''}</h3
diff --git a/frontend/src/lib/components/flows/content/FlowBranchesOneWrapper.svelte b/frontend/src/lib/components/flows/content/FlowBranchesOneWrapper.svelte
index 0420f30d2c..6ccdcbac95 100644
--- a/frontend/src/lib/components/flows/content/FlowBranchesOneWrapper.svelte
+++ b/frontend/src/lib/components/flows/content/FlowBranchesOneWrapper.svelte
@@ -32,15 +32,17 @@
 		<SplitPanesWrapper>
 			<Splitpanes horizontal>
 				<Pane size={flowModule ? 60 : 100}>
-					<Alert
-						type="info"
-						title="Only first branch whose condition is true will be run"
-						tooltip="Branch one"
-						documentationLink="https://www.windmill.dev/docs/flows/flow_branches#branch-one"
-						class="m-4"
-					>
-						The result of this step is the result of the branch.
-					</Alert>
+					{#if !noEditor}
+						<Alert
+							type="info"
+							title="Only first branch whose condition is true will be run"
+							tooltip="Branch one"
+							documentationLink="https://www.windmill.dev/docs/flows/flow_branches#branch-one"
+							class="m-4"
+						>
+							The result of this step is the result of the branch.
+						</Alert>
+					{/if}
 					<div class="p-4">
 						<h3 class="my-4">
 							{value.branches.length + 1} branch{value.branches.length + 1 > 1 ? 'es' : ''}
diff --git a/frontend/src/lib/components/flows/content/FlowLoop.svelte b/frontend/src/lib/components/flows/content/FlowLoop.svelte
index ddda211ab9..b66c83618e 100644
--- a/frontend/src/lib/components/flows/content/FlowLoop.svelte
+++ b/frontend/src/lib/components/flows/content/FlowLoop.svelte
@@ -76,16 +76,19 @@
 
 		<Splitpanes horizontal class="!max-h-[calc(100%-48px)]">
 			<Pane size={60} minSize={20} class="p-4">
-				<Alert
-					type="info"
-					title="For loops"
-					tooltip="For loops"
-					documentationLink="https://www.windmill.dev/docs/flows/flow_loops"
-					class="mb-4"
-				>
-					Add steps inside the loop and specify an iterator expression that defines the sequence
-					over which your subsequent steps will iterate.
-				</Alert>
+				{#if !noEditor}
+					<Alert
+						type="info"
+						title="For loops"
+						tooltip="For loops"
+						documentationLink="https://www.windmill.dev/docs/flows/flow_loops"
+						class="mb-4"
+					>
+						Add steps inside the loop and specify an iterator expression that defines the sequence
+						over which your subsequent steps will iterate.
+					</Alert>
+				{/if}
+
 				{#if mod.value.type === 'forloopflow'}
 					<div class="flex flex-row gap-8 mt-2 mb-6">
 						<div>
diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index cecfdeec60..f550dff29a 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -56,6 +56,7 @@
 	export let scriptKind: 'script' | 'trigger' | 'approval' = 'script'
 	export let scriptTemplate: 'pgsql' | 'mysql' | 'script' | 'docker' | 'powershell' = 'script'
 	export let noEditor: boolean
+	export let enableAi: boolean
 
 	let editor: Editor
 	let diffEditor: DiffEditor
@@ -334,7 +335,7 @@
 											previousModuleId={previousModule?.id}
 											bind:args={flowModule.value.input_transforms}
 											extraLib={stepPropPicker.extraLib}
-											enableAi
+											{enableAi}
 										/>
 									</PropPickerWrapper>
 								</div>
diff --git a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
index d66caa4f4c..340dc7dce0 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
@@ -144,11 +144,13 @@
 			failureModule={$selectedId === 'failure'}
 			{scriptKind}
 			{scriptTemplate}
+			{enableAi}
 		/>
 	{/if}
 {:else if flowModule.value.type === 'forloopflow'}
 	{#each flowModule.value.modules as submodule, index (index)}
 		<svelte:self
+			{noEditor}
 			bind:flowModule={submodule}
 			bind:parentModule={flowModule}
 			previousModule={flowModule.value.modules[index - 1]}
@@ -164,6 +166,7 @@
 	{:else}
 		{#each flowModule.value.default as submodule, index}
 			<svelte:self
+				{noEditor}
 				bind:flowModule={submodule}
 				bind:parentModule={flowModule}
 				previousModule={flowModule.value.default[index - 1]}
@@ -183,6 +186,7 @@
 		{:else}
 			{#each branch.modules as submodule, index}
 				<svelte:self
+					{noEditor}
 					bind:flowModule={submodule}
 					bind:parentModule={flowModule}
 					previousModule={flowModule.value.branches[branchIndex].modules[index - 1]}
@@ -198,6 +202,7 @@
 		{:else}
 			{#each branch.modules as submodule, index}
 				<svelte:self
+					{noEditor}
 					bind:flowModule={submodule}
 					bind:parentModule={flowModule}
 					previousModule={flowModule.value.branches[branchIndex].modules[index - 1]}
diff --git a/frontend/src/lib/components/flows/map/InsertModuleButton.svelte b/frontend/src/lib/components/flows/map/InsertModuleButton.svelte
index 56cd8678b6..04e0d93efb 100644
--- a/frontend/src/lib/components/flows/map/InsertModuleButton.svelte
+++ b/frontend/src/lib/components/flows/map/InsertModuleButton.svelte
@@ -36,9 +36,7 @@
 		<Cross class="mx-[5px]" size={15} />
 	</button>
 	<div id="flow-editor-insert-module">
-		{#if !disableAi}
-			<StepGen on:insert {index} bind:funcDesc bind:open {close} {modules} />
-		{/if}
+		<StepGen on:insert {index} bind:funcDesc bind:open {close} {modules} {disableAi} />
 
 		{#if funcDesc.length === 0}
 			<div class="font-mono divide-y text-xs w-full text-secondary">