diff --git a/build.gradle b/build.gradle
index e6d3b839b..9965f8ad3 100644
--- a/build.gradle
+++ b/build.gradle
@@ -126,7 +126,7 @@ def targetSequenceDiagramGeneratorLSExt = file("$project.rootDir/sequence-model-
 def targetTestManagerServiceLSExt = file("$project.rootDir/test-manager-service/modules/test-manager-service-ls-extension/build/libs/test-manager-service-ls-extension-${project.version}.jar")
 
 // TODO: Remove this once the workspace manager is refactored to import modules where necessary.
-def pullBallerinaModule(String packageName) {
+def pullBallerinaModule(String org, String packageName) {
     tasks.register("pullBallerinaModule_${packageName.replace('/', '_')}") {
         doLast {
             def errOutput = new ByteArrayOutputStream()
@@ -134,7 +134,7 @@ def pullBallerinaModule(String packageName) {
                 ignoreExitValue = true
 
                 // Check if the package exists in the ballerina user 
-                def centralRepoDir = new File(System.getProperty("user.home"), ".ballerina/repositories/central.ballerina.io/bala/ballerinax")
+                def centralRepoDir = new File(System.getProperty("user.home"), ".ballerina/repositories/central.ballerina.io/bala/${org}")
                 if (centralRepoDir.exists()) {
                     def pkgDir = new File(centralRepoDir, packageName)
                     if (pkgDir.exists()) {
@@ -143,7 +143,7 @@ def pullBallerinaModule(String packageName) {
                     }
                 }
 
-                def ballPullCommand = "bal pull ballerinax/${packageName}"
+                def ballPullCommand = "bal pull ${org}/${packageName}"
                 if (org.gradle.internal.os.OperatingSystem.current().isWindows()) {
                     commandLine 'cmd', '/c', ballPullCommand
                 } else {
diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/build.gradle b/flow-model-generator/modules/flow-model-generator-ls-extension/build.gradle
index 98059b29e..2ea5e16db 100644
--- a/flow-model-generator/modules/flow-model-generator-ls-extension/build.gradle
+++ b/flow-model-generator/modules/flow-model-generator-ls-extension/build.gradle
@@ -102,13 +102,13 @@ task copyStdlibs() {
     }
 }
 
-test.dependsOn rootProject.pullBallerinaModule('redis')
-test.dependsOn rootProject.pullBallerinaModule('trigger.salesforce')
-test.dependsOn rootProject.pullBallerinaModule('github')
-test.dependsOn rootProject.pullBallerinaModule('snowflake')
-test.dependsOn rootProject.pullBallerinaModule('docusign.dsadmin')
-test.dependsOn rootProject.pullBallerinaModule('sql')
-test.dependsOn rootProject.pullBallerinaModule('wso2.controlplane')
+test.dependsOn rootProject.pullBallerinaModule('ballerinax','redis')
+test.dependsOn rootProject.pullBallerinaModule('ballerinax','trigger.salesforce')
+test.dependsOn rootProject.pullBallerinaModule('ballerinax','github')
+test.dependsOn rootProject.pullBallerinaModule('ballerinax','snowflake')
+test.dependsOn rootProject.pullBallerinaModule('ballerinax','docusign.dsadmin')
+test.dependsOn rootProject.pullBallerinaModule('ballerina','sql')
+test.dependsOn rootProject.pullBallerinaModule('ballerinax','wso2.controlplane')
 
 test {
     dependsOn {
diff --git a/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/ServiceDatabaseManager.java b/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/ServiceDatabaseManager.java
index 17d7fb590..c5b3167b2 100644
--- a/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/ServiceDatabaseManager.java
+++ b/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/ServiceDatabaseManager.java
@@ -45,6 +45,7 @@ public class ServiceDatabaseManager {
     private final String dbPath;
 
     private static class Holder {
+
         private static final ServiceDatabaseManager INSTANCE = new ServiceDatabaseManager();
     }
 
@@ -111,7 +112,8 @@ public Optional<FunctionData> getListener(String module) {
                         null,
                         null,
                         rs.getBoolean("return_error"),
-                        false);
+                        false,
+                        null);
                 functionData.setPackageId(rs.getString("package_id"));
                 return Optional.of(functionData);
             }
@@ -200,6 +202,7 @@ public LinkedHashMap<String, ParameterData> getFunctionParametersAsMap(int liste
 
     // Helper builder class
     private static class ParameterDataBuilder {
+
         int parameterId;
         String name;
         String type;
diff --git a/service-model-generator/modules/service-model-generator-ls-extension/build.gradle b/service-model-generator/modules/service-model-generator-ls-extension/build.gradle
index 897215c81..95a381f29 100644
--- a/service-model-generator/modules/service-model-generator-ls-extension/build.gradle
+++ b/service-model-generator/modules/service-model-generator-ls-extension/build.gradle
@@ -106,8 +106,8 @@ task copyStdlibs() {
     }
 }
 
-test.dependsOn rootProject.pullBallerinaModule('kafka')
-test.dependsOn rootProject.pullBallerinaModule('rabbitmq')
+test.dependsOn rootProject.pullBallerinaModule('ballerinax','kafka')
+test.dependsOn rootProject.pullBallerinaModule('ballerinax','rabbitmq')
 
 test {
     dependsOn {