Update LS to Support Agent Evaluations in BI

flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/AvailableNodesGenerator.java

@@ -40,7 +40,7 @@
 import io.ballerina.flowmodelgenerator.core.model.Metadata;
 import io.ballerina.flowmodelgenerator.core.model.NodeBuilder;
 import io.ballerina.flowmodelgenerator.core.model.NodeKind;
-import io.ballerina.flowmodelgenerator.core.model.node.AgentBuilder;
+import io.ballerina.flowmodelgenerator.core.model.node.AgentRunBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.ChunkerBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.DataLoaderBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.EmbeddingProviderBuilder;
@@ -73,9 +73,10 @@
 import static io.ballerina.modelgenerator.commons.CommonUtils.PERSIST;
 import static io.ballerina.modelgenerator.commons.CommonUtils.PERSIST_MODEL_FILE;
 import static io.ballerina.modelgenerator.commons.CommonUtils.getPersistModelFilePath;
+import static io.ballerina.modelgenerator.commons.CommonUtils.isAgentClass;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiEmbeddingProvider;
-import static io.ballerina.modelgenerator.commons.CommonUtils.isAiModelProvider;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiKnowledgeBase;
+import static io.ballerina.modelgenerator.commons.CommonUtils.isAiModelProvider;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isPersistClient;
 
 /**
@@ -126,6 +127,10 @@ public JsonArray getAvailableNodes(LinePosition position) {
         return getAvailableNodes(true, position);
     }
 
+    public JsonArray getAvailableAgents(LinePosition position) {
+        return this.getAvailableItemsByCategory(position, Category.Name.AGENT, this::getAgent);
+    }
+
     public JsonArray getAvailableModelProviders(LinePosition position) {
         return this.getAvailableItemsByCategory(position, Category.Name.MODEL_PROVIDER, this::getModelProvider);
     }
@@ -320,16 +325,13 @@ private List<Item> getAiNodes(boolean disableBallerinaAiNodes) {
                 .items(List.of(knowledgeBase, dataLoaders, recursiveDocumentChunker, chunkers, augmentUserQuery,
                         vectorStore, embeddingProvider)).build();
 
-        AvailableNode agentCall = new AvailableNode(
-                new Metadata.Builder<>(null).label(AgentBuilder.LABEL)
-                        .description(AgentBuilder.DESCRIPTION).build(),
-                new Codedata.Builder<>(null).node(NodeKind.AGENT_CALL).
-                        org(disableBallerinaAiNodes ? BALLERINAX : BALLERINA).module(Ai.AI_PACKAGE)
-                        .packageName(Ai.AI_PACKAGE).symbol(Ai.AGENT_RUN_METHOD_NAME)
-                        .object(Ai.AGENT_TYPE_NAME).build(), true);
+        AvailableNode agent = new AvailableNode(
+                new Metadata.Builder<>(null).label(AgentRunBuilder.LABEL)
+                        .description(AgentRunBuilder.DESCRIPTION).build(),
+                new Codedata.Builder<>(null).node(NodeKind.AGENTS).build(), true);
 
         Category agentCategory = new Category.Builder(null).name(Category.Name.AGENT)
-                .items(List.of(agentCall)).build();
+                .items(List.of(agent)).build();
 
         return List.of(directLlmCategory, ragCategory, agentCategory);
     }
@@ -416,6 +418,8 @@ private Optional<Category> getCategory(Symbol symbol, Predicate<ClassSymbol> con
                     FunctionData.Kind kind = methodFunction.kind();
                     if (kind == FunctionData.Kind.REMOTE) {
                         nodeBuilder = NodeBuilder.getNodeFromKind(NodeKind.REMOTE_ACTION_CALL);
+                    } else if (isAgentClass(classSymbol) && label.equals(Ai.AGENT_RUN_METHOD_NAME)) {
+                        nodeBuilder = NodeBuilder.getNodeFromKind(NodeKind.AGENT_RUN);
                     } else if (kind == FunctionData.Kind.FUNCTION && isAiKnowledgeBase(classSymbol)) {
                         nodeBuilder = NodeBuilder.getNodeFromKind(NodeKind.KNOWLEDGE_BASE_CALL);
                     } else if (kind == FunctionData.Kind.FUNCTION) {
@@ -460,6 +464,16 @@ private Optional<Category> getCategory(Symbol symbol, Predicate<ClassSymbol> con
         }
     }
 
+    private Optional<Category> getAgent(Symbol symbol) {
+        return getCategory(symbol, classSymbol -> {
+            try {
+                return isAgentClass(classSymbol);
+            } catch (Exception e) {
+                return false;
+            }
+        });
+    }
+
     private Optional<Category> getModelProvider(Symbol symbol) {
         return getCategory(symbol, classSymbol -> classSymbol.qualifiers().contains(Qualifier.CLIENT)
                 && isAiModelProvider(classSymbol)

flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/NodeBuilder.java

@@ -25,6 +25,7 @@
 import io.ballerina.flowmodelgenerator.core.DiagnosticHandler;
 import io.ballerina.flowmodelgenerator.core.model.node.AgentBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.AgentCallBuilder;
+import io.ballerina.flowmodelgenerator.core.model.node.AgentRunBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.AssignBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.AutomationBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.BinaryBuilder;
@@ -161,6 +162,7 @@ public abstract class NodeBuilder implements DiagnosticHandler.DiagnosticCapable
         put(NodeKind.WAIT, WaitBuilder::new);
         put(NodeKind.AGENT, AgentBuilder::new);
         put(NodeKind.AGENT_CALL, AgentCallBuilder::new);
+        put(NodeKind.AGENT_RUN, AgentRunBuilder::new);
         put(NodeKind.CLASS_INIT, ClassInitBuilder::new);
         put(NodeKind.MEMORY, MemoryBuilder::new);
         put(NodeKind.MEMORY_STORE, MemoryStoreBuilder::new);

flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/NodeKind.java

@@ -80,7 +80,9 @@ public enum NodeKind {
 
     AGENT,
     AGENT_CALL,
+    AGENT_RUN,
     CLASS_INIT,
+    AGENTS,
 
     MODEL_PROVIDER,
     MODEL_PROVIDERS,

flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/AgentRunBuilder.java

@@ -0,0 +1,195 @@
+/*
+ *  Copyright (c) 2026, WSO2 LLC. (http://www.wso2.com)
+ *
+ *  WSO2 LLC. licenses this file to you under the Apache License,
+ *  Version 2.0 (the "License"); you may not use this file except
+ *  in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing,
+ *  software distributed under the License is distributed on an
+ *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ *  KIND, either express or implied.  See the License for the
+ *  specific language governing permissions and limitations
+ *  under the License.
+ */
+
+package io.ballerina.flowmodelgenerator.core.model.node;
+
+import io.ballerina.compiler.syntax.tree.SyntaxKind;
+import io.ballerina.flowmodelgenerator.core.AiUtils;
+import io.ballerina.flowmodelgenerator.core.model.FlowNode;
+import io.ballerina.flowmodelgenerator.core.model.NodeKind;
+import io.ballerina.flowmodelgenerator.core.model.Property;
+import io.ballerina.flowmodelgenerator.core.model.SourceBuilder;
+import io.ballerina.flowmodelgenerator.core.utils.FlowNodeUtil;
+import io.ballerina.modelgenerator.commons.FunctionData;
+import io.ballerina.modelgenerator.commons.ParameterData;
+import org.ballerinalang.langserver.common.utils.NameUtil;
+import org.eclipse.lsp4j.TextEdit;
+
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Represents a function call node.
+ *
+ * @since 1.5.1
+ */
+public class AgentRunBuilder extends CallBuilder {
+
+    private static final String BALLERINA = "ballerina";
+
+    // Agent Call Properties
+    public static final String QUERY = "query";
+    public static final String SESSION_ID = "sessionId";
+    public static final String CONTEXT = "context";
+
+    public static final String LABEL = "Agent";
+    public static final String DESCRIPTION = "Create or reuse an Agent.";
+    static final Set<String> AGENT_CALL_PARAMS_TO_SHOW = Set.of(QUERY, SESSION_ID, CONTEXT);
+
+    @Override
+    protected NodeKind getFunctionNodeKind() {
+        return NodeKind.AGENT_RUN;
+    }
+
+    @Override
+    protected FunctionData.Kind getFunctionResultKind() {
+        return FunctionData.Kind.FUNCTION;
+    }
+
+    @Override
+    public void setConcreteConstData() {
+        codedata().node(NodeKind.AGENT_RUN);
+        metadata().description(DESCRIPTION);
+    }
+
+    @Override
+    public void setConcreteTemplateData(TemplateContext context) {
+        super.setConcreteTemplateData(context);
+        // TODO: This is a temporary solution until we have a proper plan for handling all generic types.
+        makeInferredTypePropertyOptional();
+        overrideVariableName(context);
+    }
+
+    private void makeInferredTypePropertyOptional() {
+        if (formBuilder == null) {
+            return;
+        }
+        Map<String, Property> props = formBuilder.build();
+        for (Map.Entry<String, Property> entry : props.entrySet()) {
+            Property prop = entry.getValue();
+            if (prop.codedata() != null &&
+                    ParameterData.Kind.PARAM_FOR_TYPE_INFER.name().equals(prop.codedata().kind())) {
+                props.put(entry.getKey(), AiUtils.copyAsOptionalAdvanced(prop));
+            }
+        }
+    }
+
+    private void overrideVariableName(TemplateContext context) {
+        if (formBuilder == null) {
+            return;
+        }
+        Map<String, Property> props = formBuilder.build();
+        Property variableProp = props.get(Property.VARIABLE_KEY);
+        if (variableProp == null) {
+            return;
+        }
+        String uniqueVarName = NameUtil.generateVariableName("string", context.getAllVisibleSymbolNames());
+        props.put(Property.VARIABLE_KEY, AiUtils.createUpdatedProperty(variableProp, uniqueVarName));
+    }
+
+    private void newVariableWithInferredTypeAndDefault(SourceBuilder sourceBuilder) {
+        FlowNode flowNode = sourceBuilder.flowNode;
+        Optional<Property> optionalType = sourceBuilder.getProperty(Property.TYPE_KEY);
+        Optional<Property> variable = sourceBuilder.getProperty(Property.VARIABLE_KEY);
+
+        if (optionalType.isEmpty() || variable.isEmpty()) {
+            return;
+        }
+
+        Property type = optionalType.get();
+        String typeName = type.value().toString();
+
+        if (flowNode.codedata().inferredReturnType() != null) {
+            Optional<Property> inferredParam = flowNode.properties().values().stream()
+                    .filter(property -> property.codedata() != null && property.codedata().kind() != null &&
+                            property.codedata().kind().equals(ParameterData.Kind.PARAM_FOR_TYPE_INFER.name()))
+                    .findFirst();
+            if (inferredParam.isPresent()) {
+                String returnType = flowNode.codedata().inferredReturnType();
+                Object inferredValue = inferredParam.get().value();
+                // Default to "string" when the inferred type value is null or empty
+                String inferredType = (inferredValue != null && !inferredValue.toString().isEmpty())
+                        ? inferredValue.toString()
+                        : "string";
+                String inferredTypeDef = inferredParam.get()
+                        .codedata().originalName();
+                typeName = returnType.replace(inferredTypeDef, inferredType);
+            }
+        }
+
+        sourceBuilder.token().expressionWithType(typeName, variable.get()).keyword(SyntaxKind.EQUAL_TOKEN);
+    }
+
+    @Override
+    public Map<Path, List<TextEdit>> toSource(SourceBuilder sourceBuilder) {
+        // Use custom variable declaration with inferred type handling and default to "string"
+        newVariableWithInferredTypeAndDefault(sourceBuilder);
+
+        FlowNode agentRunNode = sourceBuilder.flowNode;
+        Map<Path, List<TextEdit>> allTextEdits = new HashMap<>();
+
+        generateAgentCallSource(sourceBuilder, agentRunNode, allTextEdits);
+
+        return allTextEdits;
+    }
+
+    private void generateAgentCallSource(SourceBuilder sourceBuilder, FlowNode agentRunNode,
+                                         Map<Path, List<TextEdit>> allTextEdits) {
+        Optional<Property> connection = agentRunNode.getProperty(Property.CONNECTION_KEY);
+        if (connection.isEmpty()) {
+            throw new IllegalStateException("Agent variable must be defined for an agent call node");
+        }
+
+        if (FlowNodeUtil.hasCheckKeyFlagSet(agentRunNode)) {
+            sourceBuilder.token().keyword(SyntaxKind.CHECK_KEYWORD);
+        }
+
+        Set<String> excludeKeys = getExcludeKeys(agentRunNode);
+        Map<Path, List<TextEdit>> callTextEdits = sourceBuilder.token()
+                .name(connection.get().toSourceCode())
+                .keyword(BALLERINA.equals(agentRunNode.codedata().org()) ?
+                        SyntaxKind.DOT_TOKEN : SyntaxKind.RIGHT_ARROW_TOKEN)
+                .name(agentRunNode.metadata().label())
+                .stepOut()
+                .functionParameters(agentRunNode, excludeKeys)
+                .textEdit()
+                .build();
+
+        callTextEdits.forEach((path, textEdits) ->
+                allTextEdits.merge(path, textEdits, (existing, incoming) -> {
+                    List<TextEdit> merged = new ArrayList<>(existing);
+                    merged.addAll(incoming);
+                    return merged;
+                }));
+    }
+
+    private Set<String> getExcludeKeys(FlowNode agentRunNode) {
+        return agentRunNode.properties() != null
+                ? agentRunNode.properties().keySet().stream()
+                .filter(key -> !AGENT_CALL_PARAMS_TO_SHOW.contains(key))
+                .collect(Collectors.toSet())
+                : new HashSet<>();
+    }
+}

flow-model-generator/modules/flow-model-generator-ls-extension/src/main/java/io/ballerina/flowmodelgenerator/extension/FlowModelGeneratorService.java

@@ -348,6 +348,13 @@ public CompletableFuture<FlowModelAvailableNodesResponse> getAvailableVectorKnow
                 generator -> generator.getAvailableVectorKnowledgeBases(request.position()));
     }
 
+    @JsonRequest
+    public CompletableFuture<FlowModelAvailableNodesResponse> getAvailableAgents(
+            FlowModelAvailableNodesRequest request) {
+        return handleAvailableNodesRequest(request,
+                generator -> generator.getAvailableAgents(request.position()));
+    }
+
     @JsonRequest
     public CompletableFuture<FlowModelAvailableNodesResponse> getAvailableModelProviders(
             FlowModelAvailableNodesRequest request) {

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/agents_manager/config/available_agents.json

@@ -935,15 +935,10 @@
             {
               "metadata": {
                 "label": "Agent",
-                "description": "Create new agent"
+                "description": "Create or reuse an Agent."
               },
               "codedata": {
-                "node": "AGENT_CALL",
-                "org": "ballerinax",
-                "module": "ai",
-                "packageName": "ai",
-                "object": "Agent",
-                "symbol": "run"
+                "node": "AGENTS"
               },
               "enabled": true
             }

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/caller.json

@@ -268,15 +268,10 @@
             {
               "metadata": {
                 "label": "Agent",
-                "description": "Create new agent"
+                "description": "Create or reuse an Agent."
               },
               "codedata": {
-                "node": "AGENT_CALL",
-                "org": "ballerina",
-                "module": "ai",
-                "packageName": "ai",
-                "object": "Agent",
-                "symbol": "run"
+                "node": "AGENTS"
               },
               "enabled": true
             }

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector1.json

@@ -515,15 +515,10 @@
             {
               "metadata": {
                 "label": "Agent",
-                "description": "Create new agent"
+                "description": "Create or reuse an Agent."
               },
               "codedata": {
-                "node": "AGENT_CALL",
-                "org": "ballerina",
-                "module": "ai",
-                "packageName": "ai",
-                "object": "Agent",
-                "symbol": "run"
+                "node": "AGENTS"
               },
               "enabled": true
             }

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector2.json

@@ -4633,15 +4633,10 @@
             {
               "metadata": {
                 "label": "Agent",
-                "description": "Create new agent"
+                "description": "Create or reuse an Agent."
               },
               "codedata": {
-                "node": "AGENT_CALL",
-                "org": "ballerina",
-                "module": "ai",
-                "packageName": "ai",
-                "object": "Agent",
-                "symbol": "run"
+                "node": "AGENTS"
               },
               "enabled": true
             }