feat: Give Namespace Context to AICopilot when generating a Flow

Author: Skraye

ui/src/components/ai/AiCopilot.vue

@@ -153,7 +153,8 @@
     const props = defineProps<{
         flow: string,
         conversationId: string,
-        generationType?: AiGenerationType
+        generationType?: AiGenerationType,
+        namespace?: string
     }>();
 
     const error = ref<string | undefined>(undefined);
@@ -291,13 +292,24 @@
         let aiResponse;
         try {
             const type = props.generationType ?? aiGenerationTypes.FLOW;
-            aiResponse = await aiStore.generate({
-                userPrompt: prompt.value,
-                yaml: props.flow,
-                conversationId: props.conversationId,
-                providerId: selectedProvider.value,
-                type: type
-            }) as string;
+            if (type === aiGenerationTypes.FLOW) {
+                aiResponse = await aiStore.generateFlow({
+                    userPrompt: prompt.value,
+                    yaml: props.flow,
+                    conversationId: props.conversationId,
+                    providerId: selectedProvider.value,
+                    namespace: props.namespace,
+                    type: type
+                }) as string;
+            } else {
+                aiResponse = await aiStore.generate({
+                    userPrompt: prompt.value,
+                    yaml: props.flow,
+                    conversationId: props.conversationId,
+                    providerId: selectedProvider.value,
+                    type: type
+                }) as string;
+            }
             emit("generatedYaml", aiResponse);
         } catch (e: any) {
             error.value = e.response?.data?.message ?? e.message;

ui/src/components/inputs/EditorWrapper.vue

@@ -62,6 +62,7 @@
                 @close="closeAiCopilot"
                 :flow="editorContent"
                 :conversationId="conversationId"
+                :namespace="namespace"
                 @generated-yaml="(yaml: string) => {draftSource = yaml; aiCopilotOpened = false}"
                 :generationType="aiGenerationTypes.FLOW"
             />

ui/src/stores/ai.ts

@@ -18,6 +18,19 @@ export const useAiStore = defineStore("ai", {
                 providerId
             });
 
+            return response.data;
+        },
+
+        async generateFlow({userPrompt, yaml, conversationId, providerId, namespace, tenantId}: {userPrompt: string, yaml: string, conversationId: string, providerId?: string, namespace?: string, tenantId?: string, type: AiGenerationType}) {
+            const response = await axios.post(`${apiUrl()}/ai/generate/flow`, {
+                userPrompt,
+                yaml,
+                conversationId,
+                providerId,
+                namespace,
+                tenantId
+            });
+
             return response.data;
         }
 

webserver/src/main/java/io/kestra/webserver/controllers/api/AiController.java

@@ -1,9 +1,10 @@
 package io.kestra.webserver.controllers.api;
 
-import io.kestra.webserver.models.ai.FlowGenerationPrompt;
 import io.kestra.webserver.models.ai.DashboardGenerationPrompt;
+import io.kestra.webserver.models.ai.FlowGenerationPrompt;
 import io.kestra.webserver.services.ai.AiServiceInterface;
 import io.kestra.webserver.services.ai.AiServiceManager;
+import io.kestra.core.tenant.TenantService;
 import io.micronaut.context.annotation.Requires;
 import io.micronaut.http.HttpRequest;
 import io.micronaut.http.annotation.Body;
@@ -34,6 +35,9 @@ public class AiController {
     @Inject
     protected HttpClientAddressResolver httpClientAddressResolver;
 
+    @Inject
+    protected TenantService tenantService;
+
     @ExecuteOn(TaskExecutors.IO)
     @Post(uri = "/generate/flow", produces = "application/yaml")
     @Operation(tags = {"AI"}, summary = "Generate or regenerate a flow based on a prompt")
@@ -43,7 +47,7 @@ public String generateFlow(
     ) {
         AiServiceInterface service = aiServiceManager.getAiService(flowGenerationPrompt.providerId());
 
-        return service.generateFlow(httpClientAddressResolver.resolve(httpRequest), flowGenerationPrompt);
+        return service.generateFlow(httpClientAddressResolver.resolve(httpRequest), flowGenerationPrompt, tenantService.resolveTenant());
     }
 
     @ExecuteOn(TaskExecutors.IO)

webserver/src/main/java/io/kestra/webserver/models/ai/FlowGenerationPrompt.java

@@ -2,5 +2,10 @@
 
 import jakarta.validation.constraints.NotNull;
 
-public record FlowGenerationPrompt(@NotNull String conversationId, @NotNull String userPrompt, String yaml, String providerId) {}
-
+public record FlowGenerationPrompt(
+    @NotNull String conversationId,
+    @NotNull String userPrompt,
+    String yaml,
+    String providerId,
+    String namespace
+) {}

webserver/src/main/java/io/kestra/webserver/services/ai/AiService.java

@@ -24,6 +24,7 @@ public abstract class AiService<T extends AiConfiguration> implements AiServiceI
     private final T aiConfiguration;
     private final FlowAiCopilot flowAiCopilot;
     private final DashboardAiCopilot dashboardAiCopilot;
+    private final NamespaceContextTool namespaceContextTool;
     private final String instanceUid;
     private final String aiProvider;
     private final String displayName;
@@ -51,14 +52,17 @@ protected FlowYamlBuilder flowYamlBuilder(String conversationId) {
         return AiServices.builder(FlowYamlBuilder.class)
             .chatModel(this.chatModel(
                 this.listeners("FlowYamlBuilder", conversationId)
-            )).build();
+            ))
+            .tools(namespaceContextTool)
+            .build();
     }
 
     protected DashboardYamlBuilder dashboardYamlBuilder(String conversationId) {
         return AiServices.builder(DashboardYamlBuilder.class)
             .chatModel(this.chatModel(
                 this.listeners("DashboardYamlBuilder", conversationId)
-            )).build();
+            ))
+            .build();
     }
 
     public AiService(
@@ -67,6 +71,7 @@ public AiService(
         final VersionProvider versionProvider,
         final InstanceService instanceService,
         final PosthogService postHogService,
+        final NamespaceContextTool namespaceContextTool,
         final String aiProvider,
         final String displayName,
         final List<ChatModelListener> listeners,
@@ -78,13 +83,14 @@ public AiService(
         this.displayName = displayName;
         this.listeners = listeners;
         this.aiConfiguration = aiConfiguration;
+        this.namespaceContextTool = namespaceContextTool;
 
         this.flowAiCopilot = new FlowAiCopilot(jsonSchemaGenerator, pluginRegistry, versionProvider.getVersion());
         this.dashboardAiCopilot = new DashboardAiCopilot(jsonSchemaGenerator, pluginRegistry, versionProvider.getVersion());
     }
 
     @Override
-    public String generateFlow(String ip, FlowGenerationPrompt flowGenerationPrompt) {
+    public String generateFlow(String ip, FlowGenerationPrompt flowGenerationPrompt, String tenantId) {
         AiService.GenerationContext ctx = this.beforeGeneration(ip, flowGenerationPrompt.conversationId(), "FlowGeneration", Map.of(
             "flowYaml", flowGenerationPrompt.yaml(),
             "userPrompt", flowGenerationPrompt.userPrompt()
@@ -93,7 +99,8 @@ public String generateFlow(String ip, FlowGenerationPrompt flowGenerationPrompt)
         String generatedFlow = flowAiCopilot.generateFlow(
             this.pluginFinder(flowGenerationPrompt.conversationId()),
             this.flowYamlBuilder(flowGenerationPrompt.conversationId()),
-            flowGenerationPrompt
+            flowGenerationPrompt,
+            tenantId
         );
 
         return this.afterGeneration(ctx, "FlowGenerationResult", Map.of("generatedFlow", generatedFlow), generatedFlow, "generatedFlow");

webserver/src/main/java/io/kestra/webserver/services/ai/AiServiceInterface.java

@@ -11,7 +11,7 @@
  */
 @WebServerEnabled
 public interface AiServiceInterface {
-    String generateFlow(String ip, FlowGenerationPrompt flowGenerationPrompt);
+    String generateFlow(String ip, FlowGenerationPrompt flowGenerationPrompt, String tenantId);
 
     String generateDashboard(String ip, DashboardGenerationPrompt dashboardGenerationPrompt);
 

webserver/src/main/java/io/kestra/webserver/services/ai/AiServiceManager.java

@@ -22,6 +22,7 @@ public class AiServiceManager {
     private final Map<String, AiServiceInterface> aiServices = new HashMap<>();
     private final AiProvidersConfiguration providersConfiguration;
     private String defaultProviderId;
+    protected final NamespaceContextTool namespaceContextTool;
 
     public AiServiceManager(
         AiProvidersConfiguration providersConfiguration,
@@ -32,9 +33,11 @@ public AiServiceManager(
         VersionProvider versionProvider,
         InstanceService instanceService,
         PosthogService posthogService,
-        List<dev.langchain4j.model.chat.listener.ChatModelListener> listeners
+        List<dev.langchain4j.model.chat.listener.ChatModelListener> listeners,
+        NamespaceContextTool namespaceContextTool
     ) {
         this.providersConfiguration = providersConfiguration;
+        this.namespaceContextTool = namespaceContextTool;
 
         List<AiProviderConfiguration> configs = new java.util.ArrayList<>(
             providersConfiguration.providers() != null ? providersConfiguration.providers() : List.of()
@@ -95,7 +98,7 @@ protected AiServiceInterface createAiService(
         try {
             if (type.equals("gemini")) {
                 GeminiConfiguration geminiConfig = JacksonMapper.ofJson().convertValue(configMap, GeminiConfiguration.class);
-                return new GeminiAiService(pluginRegistry, jsonSchemaGenerator, versionProvider, instanceService, posthogService, provider.displayName(), listeners, geminiConfig);
+                return new GeminiAiService(pluginRegistry, jsonSchemaGenerator, versionProvider, instanceService, posthogService, namespaceContextTool, provider.displayName(), listeners, geminiConfig);
             }
             log.warn("Unknown AI type: {}", type);
             return null;
@@ -131,6 +134,3 @@ public String getDefaultProviderId() {
         return defaultProviderId;
     }
 }
-
-
-

webserver/src/main/java/io/kestra/webserver/services/ai/FlowAiCopilot.java

@@ -56,13 +56,13 @@ protected List<String> excludedPluginTypes() {
         return EXCLUDED_PLUGIN_TYPES;
     }
 
-    public String generateFlow(PluginFinder pluginFinder, FlowYamlBuilder flowYamlBuilder, FlowGenerationPrompt flowGenerationPrompt) {
+    public String generateFlow(PluginFinder pluginFinder, FlowYamlBuilder flowYamlBuilder, FlowGenerationPrompt flowGenerationPrompt, String tenantId) {
         String enhancedPrompt = String.format("Current Flow YAML:\n```yaml\n%s\n```\n\nUser's prompt:\n``\n%s\n```", java.util.Optional.ofNullable(flowGenerationPrompt.yaml()).orElse(""), flowGenerationPrompt.userPrompt());
 
         List<String> mostRelevantPlugins = this.mostRelevantPlugins(pluginFinder, enhancedPrompt, this.excludedPluginTypes());
 
         return this.generateYaml(
-            flowYamlBuilder::buildFlow,
+            (schemaJson, genErr, userPr) -> flowYamlBuilder.buildFlow(schemaJson, genErr, java.util.Optional.ofNullable(flowGenerationPrompt.yaml()).orElse(""), flowGenerationPrompt.namespace(), tenantId, userPr),
             Flow.class,
             mostRelevantPlugins,
             NON_REQUEST_ERROR,

webserver/src/main/java/io/kestra/webserver/services/ai/FlowYamlBuilder.java

@@ -7,11 +7,15 @@
 public interface FlowYamlBuilder {
     // Note, there is a newline within a pebble example because {{...}} are parsed by langchain4j that tries to find a variable. This is a hack to workaround the regex and avoid errors.
     @SystemMessage("""
-        You are an expert in generating Kestra Flow YAML. Your task is to generate a valid Kestra Flow YAML that follows user's requirements strictly following the following json schema:
+        You are an expert in generating Kestra Flow YAML. Your task is to generate a valid Kestra Flow YAML that follows user's requirements strictly following the following json schema (when provided):
         ```
         {_{flowSchema}_}
         ```
-        
+
+        Additional runtime inputs available to you as variables (preferred over embedding data inside the schema):
+        - {_{namespace}_}: An explicit namespace string provided separately. If present, prefer this over any namespace found in the YAML snippet.
+        - {_{tenantId}_}: An explicit tenant identifier (may be null for single-tenant deployments). If present, prefer this over any tenantId found in the YAML snippet.
+
         Here are the rules:
         - Use examples, properties, and outputs only as specified in the schema.
         - If the user asks for troubleshooting, try to fix any related expression or task.
@@ -37,21 +41,62 @@ Avoid duplicating existing intent (e.g., if the Flow logs "hi" and the user want
         - Triggers expose some variables that can be accessed through `trigger.outputName` in expressions. The only variables available are those defined in the trigger's outputs.
         - Unless specified by the user, never assume a local port to serve any content, always use a remote URL (like a public HTTP server) to fetch content.
         - Unless specified by the user, do not use any authenticated API, always use public APIs or those that don't require authentication.
-        - To avoid escaping quotes, use double quotes first and if you need quotes inside, use single ones. Only escape them if you have 3+ level quotes, for example: `message: "Hello {{inputs.userJson | jq('.name')}}"` is preferred but `message: "Hello \\"Bob\\""` may still be used.
+        - To avoid escaping quotes, use double quotes first and if you need quotes inside, use single ones. Only escape them if you have 3+ level quotes, for example: `message: "Hello {{inputs.userJson | jq('.name')}}"` is preferred but `message: "Hello \"Bob\""` may still be used.
         - A property key is unique within each type.
         - When fetching data from the JDBC plugin, always use fetchType: STORE.
         - Manipulating date in pebble expressions can be done through `dateAdd` (`{{now()|dateAdd(-1,'DAYS')}}`) and `date` filters (`{{"July 24, 2001"|date("yyyy-MM-dd",existingFormat="MMMM dd, yyyy")}}`). Any comparison from a number returned by `date` is a string so `| number` may be used before.
         - Current date is `{{current_date_time}}`.
         - Always preserve root-level `id` and `namespace` if provided.
         - Don't add any Schedule trigger unless a regular occurrence is asked.
-        - If the user uses vague references ("it", "that"), infer context from the current Flow YAML.
+        - If the user uses vague references ("it", "that"), infer context from the current Flow YAML or the explicit `namespace`/`tenantId` variables.
         - Except for error scenarios, output only the raw YAML, with no explanation or additional text.
-        
+        - If you have any other information to share to the user, add them as comments in the YAML using `#` at the beginning of the raw YAML.
+        - Never add raw text in the response
+
+        Available Tools for Context Retrieval:
+
+        You have access to tools that retrieve namespace configuration on-demand. Prefer explicit variables `{_{namespace}_}` and `{_{tenantId}_}` when provided. Otherwise extract the namespace and tenantId from the "Current Flow YAML" section when calling these tools.
+
+        Namespace Context Tools :
+
+        1. KV Store Keys (getKvStoreKeys):
+           - Call when user wants to interact with KV Store (read/write/list keys)
+           - Returns: JSON with list of existing KV keys, descriptions, and update dates
+           - Usage in flows: {{kv('keyName')}} to read, KV tasks (Get, Put, Delete) to manage
+
+        2. Plugin Defaults (getPluginDefaults) [EE Only]:
+           - Call when user asks to integrate with specific technology (e.g., MongoDB, PostgreSQL)
+           - Returns: JSON mapping plugin types to their default values
+           - Apply defaults in generated tasks
+
+        3. Namespace Variables (getNamespaceVariables) [EE Only]:
+           - Call when user references variables
+           - Returns: JSON mapping variable names to values
+           - Prefer namespace variables using {{vars.variableName}} over hardcoded values
+
+        4. Secret Names (getSecretNames) [EE Only]:
+           - Call when user mentions secrets or credentials, or the technology being integrated typically requires credentials
+           - Returns: JSON array of secret names (NOT values)
+           - Reference using {{secret('secretName')}} for sensitive data
+
+        5. Complete Context (getAllNamespaceContext) [EE Only]:
+           - Call for comprehensive namespace information
+           - Returns: Combined context including KV keys, plugin defaults, variables, and secrets
+
+        Tool Calling Guidelines:
+        - Use explicit `{_{namespace}_}` and `{_{tenantId}_}` placeholders, if not provided, do not call the tool.
+        - Call tools BEFORE generating YAML to provide accurate suggestions
+        - Only call tools relevant to the user's request
+        - If namespace is not in YAML and flow is new, inform user namespace is required
+
         IMPORTANT: If the user prompt cannot be fulfilled with the schema, instead of generating a Flow, reply: `{_{flowGenerationError}_}`.
         Do not invent properties or types. Strictly follow the provided schema.""")
     String buildFlow(
         @V("flowSchema") String flowSchema,
         @V("flowGenerationError") String flowGenerationError,
+        @V("currentFlowYaml") String currentFlowYaml,
+        @V("namespace") String namespace,
+        @V("tenantId") String tenantId,
         @UserMessage String userPrompt
     );
 }