Update LS to Support Agent Evaluations in BI

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

@@ -242,6 +242,50 @@ public static Property copyAsOptionalAdvanced(Property original) {
         );
     }
 
+    /**
+     * Creates a copy of a property with updated metadata label while preserving all other fields.
+     *
+     * @param original the property to copy from
+     * @param newLabel the new label to set in metadata
+     * @return the new property with updated metadata label
+     */
+    public static Property createPropertyWithUpdatedLabel(Property original, String newLabel) {
+        if (original == null) {
+            throw new IllegalArgumentException("Original property cannot be null");
+        }
+        if (original.metadata() == null) {
+            throw new IllegalArgumentException("Original property metadata cannot be null");
+        }
+
+        Metadata updatedMetadata = new Metadata(
+                newLabel,
+                original.metadata().description(),
+                original.metadata().keywords(),
+                original.metadata().icon(),
+                original.metadata().functionKind(),
+                original.metadata().data()
+        );
+
+        return new Property(
+                updatedMetadata,
+                original.types(),
+                original.value(),
+                original.oldValue(),
+                original.placeholder(),
+                original.optional(),
+                original.editable(),
+                original.advanced(),
+                original.hidden(),
+                original.modified(),
+                original.diagnostics(),
+                original.codedata(),
+                original.advancedValue(),
+                original.imports(),
+                original.defaultValue(),
+                original.comment()
+        );
+    }
+
     /**
      * Adds a property to a NodeBuilder by copying all attributes from an existing property with an optional custom
      * value.
@@ -265,26 +309,25 @@ public static void addPropertyFromTemplate(NodeBuilder nodeBuilder, String key,
         Object valueToUse = customValue != null ? customValue : property.value();
         boolean hidden = isHidden || property.hidden();
 
-        nodeBuilder.properties().custom()
-                .metadata()
-                    .label(property.metadata().label())
-                    .description(property.metadata().description())
-                    .stepOut()
-                .types(property.types())
-                .placeholder(property.placeholder())
-                .value(valueToUse)
-                .defaultValue(property.defaultValue())
-                .imports(property.imports() != null ? property.imports().toString() : null)
-                .optional(property.optional())
-                .editable(property.editable())
-                .advanced(property.advanced())
-                .hidden(hidden)
-                .modified(property.modified())
-                .codedata()
-                    .kind(property.codedata() != null ? property.codedata().kind() : "")
-                    .stepOut()
-                .stepOut()
-                .addProperty(key);
+        Property copied = new Property(
+                property.metadata(),
+                property.types(),
+                valueToUse,
+                property.oldValue(),
+                property.placeholder(),
+                property.optional(),
+                property.editable(),
+                property.advanced(),
+                hidden,
+                property.modified(),
+                property.diagnostics(),
+                property.codedata(),
+                property.advancedValue(),
+                property.imports(),
+                property.defaultValue(),
+                property.comment()
+        );
+        nodeBuilder.properties().build().put(key, copied);
     }
 
     /**

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/AgentCallBuilder.java

@@ -29,6 +29,7 @@
 import io.ballerina.flowmodelgenerator.core.Constants;
 import io.ballerina.flowmodelgenerator.core.model.Codedata;
 import io.ballerina.flowmodelgenerator.core.model.FlowNode;
+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.Property;
@@ -66,6 +67,7 @@
 public class AgentCallBuilder extends CallBuilder {
 
     private static final String BALLERINA = "ballerina";
+    private Set<String> cachedVisibleSymbolNames;
 
     // Agent Properties
     public static final String AGENT = "AGENT";
@@ -98,6 +100,9 @@ public class AgentCallBuilder extends CallBuilder {
     // Cache for agent templates to avoid expensive repeated creation
     private static final Map<String, FlowNode> agentTemplateCache = new ConcurrentHashMap<>();
 
+    // Cache for agent call function templates to avoid repeated FunctionDataBuilder.build() calls
+    private static final Map<String, FlowNode> agentCallFnCache = new ConcurrentHashMap<>();
+
     @Override
     protected NodeKind getFunctionNodeKind() {
         return NodeKind.AGENT_CALL;
@@ -118,12 +123,61 @@ public void setConcreteConstData() {
     public void setConcreteTemplateData(TemplateContext context) {
         setAgentProperties(this, context, null);
         setAdditionalAgentProperties(this, null);
-        super.setConcreteTemplateData(context);
+
+        FlowNode callTemplate = getOrCreateCallFunctionTemplate(context);
+        restoreFromTemplate(callTemplate);
+
+        Codedata contextCd = context.codedata();
+        codedata().lineRange(contextCd.lineRange()).sourceCode(contextCd.sourceCode());
+
         // TODO: This is a temporary solution until we have a proper plan for handling all generic types.
         makeInferredTypePropertyOptional();
         overrideVariableName(context);
     }
 
+    private FlowNode getOrCreateCallFunctionTemplate(TemplateContext context) {
+        Codedata cd = context.codedata();
+        String cacheKey = String.format("%s|%s|%s|%s|%s",
+                cd.org(), cd.packageName(), cd.version(), cd.symbol(), cd.object());
+        return agentCallFnCache.computeIfAbsent(cacheKey, k -> {
+            AgentCallBuilder temp = new AgentCallBuilder();
+            temp.defaultModuleName(moduleInfo);
+            temp.callSuperSetConcreteTemplateData(context);
+            return temp.build();
+        });
+    }
+
+    void callSuperSetConcreteTemplateData(TemplateContext context) {
+        super.setConcreteTemplateData(context);
+    }
+
+    private void restoreFromTemplate(FlowNode template) {
+        Metadata md = template.metadata();
+        if (md != null) {
+            metadata().label(md.label()).description(md.description());
+            if (md.icon() != null) {
+                metadata().icon(md.icon());
+            }
+        }
+
+        Codedata cd = template.codedata();
+        if (cd != null) {
+            codedata().node(cd.node()).org(cd.org()).module(cd.module())
+                    .packageName(cd.packageName()).object(cd.object())
+                    .version(cd.version()).symbol(cd.symbol())
+                    .inferredReturnType(cd.inferredReturnType());
+        }
+
+        if (template.properties() != null) {
+            Map<String, Property> currentProps = properties().build();
+            template.properties().forEach(currentProps::put);
+        }
+
+        if (template.flags() != 0) {
+            flag(template.flags());
+        }
+    }
+
     private void makeInferredTypePropertyOptional() {
         if (formBuilder == null) {
             return;
@@ -133,11 +187,32 @@ private void makeInferredTypePropertyOptional() {
             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));
+                Property updatedProp = AiUtils.copyAsOptionalAdvanced(prop);
+                // Update the label to "Type Descriptor" for better clarity
+                if ("td".equals(entry.getKey()) && updatedProp.metadata() != null) {
+                    updatedProp = AiUtils.createPropertyWithUpdatedLabel(updatedProp, "Type Descriptor");
+                }
+                props.put(entry.getKey(), updatedProp);
             }
         }
     }
 
+    private Set<String> getVisibleSymbolNames(TemplateContext context) {
+        if (cachedVisibleSymbolNames == null) {
+            cachedVisibleSymbolNames = context.getAllVisibleSymbolNames();
+        }
+        return cachedVisibleSymbolNames;
+    }
+
+    @Override
+    protected void setReturnTypeProperties(FunctionData functionData, TemplateContext context,
+                                           String label, String doc, boolean hidden) {
+        properties()
+                .type(functionData.returnType(), false, functionData.importStatements(), hidden,
+                        Property.RESULT_TYPE_LABEL)
+                .data(functionData.returnType(), getVisibleSymbolNames(context), label, doc);
+    }
+
     private void overrideVariableName(TemplateContext context) {
         if (formBuilder == null) {
             return;
@@ -147,7 +222,7 @@ private void overrideVariableName(TemplateContext context) {
         if (variableProp == null) {
             return;
         }
-        String uniqueVarName = NameUtil.generateVariableName("string", context.getAllVisibleSymbolNames());
+        String uniqueVarName = NameUtil.generateVariableName("string", getVisibleSymbolNames(context));
         props.put(Property.VARIABLE_KEY, AiUtils.createUpdatedProperty(variableProp, uniqueVarName));
     }