Fix data mapping function detection to exclude regular and natural expression functions in data_mappings.bal

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

@@ -86,9 +86,10 @@ class FunctionSearchCommand extends SearchCommand {
     private static final String FETCH_KEY = "functions";
     private final List<String> moduleNames;
     private final Document functionsDoc;
+    private final Document dataMappingsDoc;
 
     public FunctionSearchCommand(Project project, LineRange position, Map<String, String> queryMap,
-                                 Document functionsDoc) {
+                                 Document functionsDoc, Document dataMappingsDoc) {
         super(project, position, queryMap);
 
         // Obtain the imported module names
@@ -98,6 +99,7 @@ public FunctionSearchCommand(Project project, LineRange position, Map<String, St
                 .map(moduleDependency -> moduleDependency.descriptor().name().packageName().value())
                 .toList();
         this.functionsDoc = functionsDoc;
+        this.dataMappingsDoc = dataMappingsDoc;
         // TODO: Use this method when https://github.com/ballerina-platform/ballerina-lang/issues/43695 is fixed
         // List<String> moduleNames = semanticModel.moduleSymbols().stream()
         // .filter(symbol -> symbol.kind().equals(SymbolKind.MODULE))
@@ -195,17 +197,26 @@ private void buildProjectNodes() {
         List<Item> availableTools = new ArrayList<>();
         for (Symbol symbol : functionSymbols) {
             FunctionSymbol functionSymbol = (FunctionSymbol) symbol;
+            // Skip NP functions from functions.bal
             if (functionsDoc != null
                     && CommonUtils.isNaturalExpressionBodiedFunction(functionsDoc.syntaxTree(), functionSymbol)) {
-                // Skip NP functions
+                continue;
+            }
+            // Skip NP functions from data_mappings.bal
+            if (dataMappingsDoc != null
+                    && CommonUtils.isNaturalExpressionBodiedFunction(dataMappingsDoc.syntaxTree(), functionSymbol)) {
                 continue;
             }
 
             boolean isDataMappedFunction = false;
             Optional<Location> location = symbol.getLocation();
             if (location.isPresent()) {
-                isDataMappedFunction = location.get().lineRange().fileName().equals(DATA_MAPPER_FILE_NAME);
                 LineRange fnLineRange = location.get().lineRange();
+                // Check if function is in data_mappings.bal and is a data mapping function
+                if (fnLineRange.fileName().equals(DATA_MAPPER_FILE_NAME) && dataMappingsDoc != null) {
+                    isDataMappedFunction = CommonUtils.isDataMappingFunction(
+                            dataMappingsDoc.syntaxTree(), functionSymbol);
+                }
                 if (fnLineRange.fileName().equals(position.fileName()) &&
                         PositionUtil.isWithinLineRange(fnLineRange, position)) {
                     continue;

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

@@ -62,8 +62,13 @@ public abstract class SearchCommand {
 
     public static SearchCommand from(Kind kind, Project module, LineRange position, Map<String, String> queryMap,
                                      Document functionsDoc) {
+        return from(kind, module, position, queryMap, functionsDoc, null);
+    }
+
+    public static SearchCommand from(Kind kind, Project module, LineRange position, Map<String, String> queryMap,
+                                     Document functionsDoc, Document dataMappingsDoc) {
         return switch (kind) {
-            case FUNCTION -> new FunctionSearchCommand(module, position, queryMap, functionsDoc);
+            case FUNCTION -> new FunctionSearchCommand(module, position, queryMap, functionsDoc, dataMappingsDoc);
             case CONNECTOR -> new ConnectorSearchCommand(module, position, queryMap);
             case NP_FUNCTION -> new NPFunctionSearchCommand(module, position, queryMap, functionsDoc);
             case TYPE -> new TypeSearchCommand(module, position, queryMap);

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

@@ -646,9 +646,10 @@ public CompletableFuture<FlowModelAvailableNodesResponse> search(SearchRequest r
 
                 Path projectPath = workspaceManager.projectRoot(filePath);
                 Optional<Document> functionsDoc = getDocumentFromFile(projectPath, "functions.bal");
+                Optional<Document> dataMappingsDoc = getDocumentFromFile(projectPath, "data_mappings.bal");
 
                 SearchCommand command = SearchCommand.from(searchKind, project, position, request.queryMap(),
-                        functionsDoc.orElse(null));
+                        functionsDoc.orElse(null), dataMappingsDoc.orElse(null));
                 response.setCategories(command.execute());
             } catch (Throwable e) {
                 response.setError(e);

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/search/config/functions/custom_functions_in_data_mappings.json

@@ -0,0 +1,124 @@
+{
+  "description": "Functions in data_mappings.bal - only expression-bodied functions should be marked as data mapping functions",
+  "kind": "FUNCTION",
+  "source": "custom_func_proj/main.bal",
+  "position": {
+    "startLine": {
+      "line": 2,
+      "offset": 0
+    },
+    "endLine": {
+      "line": 4,
+      "offset": 1
+    }
+  },
+  "queryMap": {},
+  "categories": [
+    {
+      "metadata": {
+        "label": "Current Integration",
+        "description": "Functions defined within the current integration",
+        "keywords": [
+          "Project",
+          "Local",
+          "Function"
+        ]
+      },
+      "items": [
+        {
+          "metadata": {
+            "label": "customHelper",
+            "description": null,
+            "data": {
+              "isDataMappedFunction": false,
+              "isAgentTool": false,
+              "isIsolatedFunction": false
+            }
+          },
+          "codedata": {
+            "node": "FUNCTION_CALL",
+            "org": "test",
+            "module": "custom_func_proj",
+            "symbol": "customHelper",
+            "version": "0.1.0"
+          },
+          "enabled": true
+        },
+        {
+          "metadata": {
+            "label": "mapPersonToEmployee",
+            "description": null,
+            "data": {
+              "isDataMappedFunction": true,
+              "isAgentTool": false,
+              "isIsolatedFunction": false
+            }
+          },
+          "codedata": {
+            "node": "FUNCTION_CALL",
+            "org": "test",
+            "module": "custom_func_proj",
+            "symbol": "mapPersonToEmployee",
+            "version": "0.1.0"
+          },
+          "enabled": true
+        },
+        {
+          "metadata": {
+            "label": "validateAge",
+            "description": null,
+            "data": {
+              "isDataMappedFunction": false,
+              "isAgentTool": false,
+              "isIsolatedFunction": true
+            }
+          },
+          "codedata": {
+            "node": "FUNCTION_CALL",
+            "org": "test",
+            "module": "custom_func_proj",
+            "symbol": "validateAge",
+            "version": "0.1.0"
+          },
+          "enabled": true
+        }
+      ]
+    },
+    {
+      "metadata": {
+        "label": "Agent Tools",
+        "description": "Functions used as agent tools",
+        "keywords": [
+          "Project",
+          "Local",
+          "Function"
+        ]
+      },
+      "items": []
+    },
+    {
+      "metadata": {
+        "label": "Imported Functions",
+        "description": "Functions imported from other integrations",
+        "keywords": [
+          "Imported",
+          "Function",
+          "Library"
+        ]
+      },
+      "items": []
+    },
+    {
+      "metadata": {
+        "label": "Available Functions",
+        "description": "Functions available in the library",
+        "keywords": [
+          "Available",
+          "Function",
+          "Library"
+        ]
+      },
+      "items": []
+    }
+  ]
+}

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/search/source/custom_func_proj/Ballerina.toml

@@ -0,0 +1,4 @@
+[package]
+org = "test"
+name = "custom_func_proj"
+version = "0.1.0"

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/search/source/custom_func_proj/data_mappings.bal

@@ -0,0 +1,36 @@
+import ballerina/ai;
+
+type Person record {
+    string name;
+    int age;
+};
+
+type Employee record {
+    string fullName;
+    int age;
+};
+
+// This is a data mapping function (expression-bodied)
+function mapPersonToEmployee(Person person) returns Employee => {
+    fullName: person.name,
+    age: person.age
+};
+
+// This is a custom function (regular function body), not a data mapping function
+function customHelper(string input) returns string {
+    return "Processed: " + input;
+}
+
+// This is another custom function (regular function body)
+isolated function validateAge(int age) returns boolean {
+    return age >= 18;
+}
+
+// This is a natural expression function - should be skipped entirely
+function generateGreeting(string name) returns string|error => natural (check ai:getDefaultModelProvider()) {
+    **What to do**
+    Generate a personalized greeting for ${name}
+
+    **Output**
+    string - A friendly greeting message
+};

flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/search/source/custom_func_proj/main.bal

@@ -0,0 +1,5 @@
+import ballerina/io;
+
+public function main() {
+    io:println("Hello World");
+}

model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/CommonUtils.java

@@ -854,6 +854,29 @@ public static boolean isNaturalExpressionBodiedFunction(SyntaxTree syntaxTree, F
         return BallerinaCompilerApi.getInstance().isNaturalExpressionBodiedFunction(functionDefNode);
     }
 
+    /**
+     * Checks if the given function is a data mapping function (expression-bodied function excluding natural expressions).
+     *
+     * @param syntaxTree the syntax tree
+     * @param functionSymbol the function symbol
+     * @return true if the function is a data mapping function, false otherwise
+     */
+    public static boolean isDataMappingFunction(SyntaxTree syntaxTree, FunctionSymbol functionSymbol) {
+        if (functionSymbol.getLocation().isEmpty()) {
+            return false;
+        }
+        NonTerminalNode node = getNode(syntaxTree, functionSymbol.getLocation().get().textRange());
+        if (node.kind() != SyntaxKind.FUNCTION_DEFINITION) {
+            return false;
+        }
+        FunctionDefinitionNode functionDefNode = (FunctionDefinitionNode) node;
+        // A function is a data mapping function if:
+        // 1. It has an expression function body
+        // 2. It's not a natural expression bodied function
+        return functionDefNode.functionBody().kind() == SyntaxKind.EXPRESSION_FUNCTION_BODY
+                && !BallerinaCompilerApi.getInstance().isNaturalExpressionBodiedFunction(functionDefNode);
+    }
+
     public static String getClassType(String packageName, String clientName) {
         String importPrefix = packageName.substring(packageName.lastIndexOf('.') + 1);
         return String.format("%s:%s", importPrefix, clientName);