Add support for Azure AI Knowledge Base

architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/artifactsgenerator/ModuleNodeTransformer.java

@@ -50,7 +50,7 @@
 import java.util.stream.Collectors;
 
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiMemoryStore;
-import static io.ballerina.modelgenerator.commons.CommonUtils.isAiVectorKnowledgeBase;
+import static io.ballerina.modelgenerator.commons.CommonUtils.isAiKnowledgeBase;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiVectorStore;
 
 /**
@@ -240,7 +240,7 @@ private Optional<ClassSymbol> getConnection(Node node) {
             TypeReferenceTypeSymbol typeDescriptorSymbol =
                     (TypeReferenceTypeSymbol) ((VariableSymbol) symbol).typeDescriptor();
             ClassSymbol classSymbol = (ClassSymbol) typeDescriptorSymbol.typeDescriptor();
-            if (classSymbol.qualifiers().contains(Qualifier.CLIENT) || isAiVectorKnowledgeBase(classSymbol)
+            if (classSymbol.qualifiers().contains(Qualifier.CLIENT) || isAiKnowledgeBase(classSymbol)
                     || isAiVectorStore(symbol) || isAiMemoryStore(symbol)) {
                 return Optional.of(classSymbol);
             }

architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/designmodelgenerator/core/CommonUtils.java

@@ -62,6 +62,7 @@ public class CommonUtils {
     private static final String AI = "ai";
     private static final String AGENT = "Agent";
     private static final String KNOWLEDGE_BASE_TYPE_NAME = "KnowledgeBase";
+
     private static final String WSO2_MODEL_PROVIDER = "Wso2ModelProvider";
     private static final String WSO2_EMBEDDING_PROVIDER = "Wso2EmbeddingProvider";
     private static final String WSO2_ICON_SUFFIX = "?wso2_icon";

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

@@ -58,9 +58,9 @@
 import static io.ballerina.flowmodelgenerator.core.model.NodeKind.DATA_LOADERS;
 import static io.ballerina.flowmodelgenerator.core.model.NodeKind.EMBEDDING_PROVIDER;
 import static io.ballerina.flowmodelgenerator.core.model.NodeKind.EMBEDDING_PROVIDERS;
+import static io.ballerina.flowmodelgenerator.core.model.NodeKind.KNOWLEDGE_BASES;
 import static io.ballerina.flowmodelgenerator.core.model.NodeKind.MODEL_PROVIDER;
 import static io.ballerina.flowmodelgenerator.core.model.NodeKind.MODEL_PROVIDERS;
-import static io.ballerina.flowmodelgenerator.core.model.NodeKind.VECTOR_KNOWLEDGE_BASES;
 import static io.ballerina.flowmodelgenerator.core.model.NodeKind.VECTOR_STORE;
 import static io.ballerina.flowmodelgenerator.core.model.NodeKind.VECTOR_STORES;
 
@@ -94,7 +94,7 @@ public class AiUtils {
 
     static {
         versionToFeatures.put("1.0.0",
-                Set.of(MODEL_PROVIDERS, EMBEDDING_PROVIDERS, VECTOR_STORES, VECTOR_KNOWLEDGE_BASES));
+                Set.of(MODEL_PROVIDERS, EMBEDDING_PROVIDERS, VECTOR_STORES, KNOWLEDGE_BASES));
         versionToFeatures.put("1.3.0", Set.of(CHUNKERS, DATA_LOADERS));
 
         dependentModules.put("1.0.0", List.of(
@@ -140,6 +140,10 @@ public class AiUtils {
                 new Module("ballerinax", "ai.pgvector", "1.0.0"),
                 new Module("ballerinax", "ai.weaviate", "1.0.0")
         ));
+
+        dependentModules.put("1.6.0", List.of(
+                new Module("ballerinax", "ai.azure", "1.2.0")
+        ));
     }
 
     /**

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

@@ -44,9 +44,9 @@
 import io.ballerina.flowmodelgenerator.core.model.node.ChunkerBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.DataLoaderBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.EmbeddingProviderBuilder;
+import io.ballerina.flowmodelgenerator.core.model.node.KnowledgeBaseBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.ModelProviderBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.NPFunctionCall;
-import io.ballerina.flowmodelgenerator.core.model.node.VectorKnowledgeBaseBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.VectorStoreBuilder;
 import io.ballerina.modelgenerator.commons.CommonUtils;
 import io.ballerina.modelgenerator.commons.FunctionData;
@@ -71,7 +71,7 @@
 import static io.ballerina.flowmodelgenerator.core.Constants.NaturalFunctions;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiEmbeddingProvider;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiModelProvider;
-import static io.ballerina.modelgenerator.commons.CommonUtils.isAiVectorKnowledgeBase;
+import static io.ballerina.modelgenerator.commons.CommonUtils.isAiKnowledgeBase;
 
 /**
  * Generates available nodes for a given position in the diagram.
@@ -134,7 +134,7 @@ public JsonArray getAvailableVectorStores(LinePosition position) {
     }
 
     public JsonArray getAvailableVectorKnowledgeBases(LinePosition position) {
-        return this.getAvailableItemsByCategory(position, Category.Name.VECTOR_KNOWLEDGE_BASE, this::getKnowledgeBase);
+        return this.getAvailableItemsByCategory(position, Category.Name.KNOWLEDGE_BASE, this::getKnowledgeBase);
     }
 
     public JsonArray getAvailableDataLoaders(LinePosition position) {
@@ -268,13 +268,13 @@ private List<Item> getAiNodes(boolean disableBallerinaAiNodes) {
         Category directLlmCategory = new Category.Builder(null).name(Category.Name.DIRECT_LLM)
                 .items(List.of(modelProvider, npFunction)).build();
 
-        AvailableNode vectorKnowledgeBase = new AvailableNode(
-                new Metadata.Builder<>(null).label(VectorKnowledgeBaseBuilder.LABEL)
-                        .description(VectorKnowledgeBaseBuilder.DESCRIPTION).build(),
-                new Codedata.Builder<>(null).node(NodeKind.VECTOR_KNOWLEDGE_BASES).org(Ai.BALLERINA_ORG)
+        AvailableNode knowledgeBase = new AvailableNode(
+                new Metadata.Builder<>(null).label(KnowledgeBaseBuilder.LABEL)
+                        .description(KnowledgeBaseBuilder.DESCRIPTION).build(),
+                new Codedata.Builder<>(null).node(NodeKind.KNOWLEDGE_BASES).org(Ai.BALLERINA_ORG)
                         .module(Ai.AI_PACKAGE).packageName(Ai.AI_PACKAGE)
-                        .object(Ai.VECTOR_KNOWLEDGE_BASE_TYPE_NAME).symbol("init").version(aiPackageVersion).build(),
-                !disableBallerinaAiNodes && supportedFeatures.contains(NodeKind.VECTOR_KNOWLEDGE_BASES));
+                        .object(Ai.KNOWLEDGE_BASE_TYPE_NAME).symbol("init").version(aiPackageVersion).build(),
+                !disableBallerinaAiNodes && supportedFeatures.contains(NodeKind.KNOWLEDGE_BASES));
 
         AvailableNode recursiveDocumentChunker = new AvailableNode(new Metadata.Builder<>(null)
                 .label(Ai.RECURSIVE_DOCUMENT_CHUNKER_LABEL).build(), new Codedata.Builder<>(null)
@@ -312,7 +312,7 @@ private List<Item> getAiNodes(boolean disableBallerinaAiNodes) {
                 !disableBallerinaAiNodes && supportedFeatures.contains(NodeKind.DATA_LOADERS));
 
         Category ragCategory = new Category.Builder(null).name(Category.Name.RAG)
-                .items(List.of(vectorKnowledgeBase, dataLoaders, recursiveDocumentChunker, chunkers, augmentUserQuery,
+                .items(List.of(knowledgeBase, dataLoaders, recursiveDocumentChunker, chunkers, augmentUserQuery,
                         vectorStore, embeddingProvider)).build();
 
         AvailableNode agentCall = new AvailableNode(
@@ -410,8 +410,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 (kind == FunctionData.Kind.FUNCTION && isAiVectorKnowledgeBase(classSymbol)) {
-                        nodeBuilder = NodeBuilder.getNodeFromKind(NodeKind.VECTOR_KNOWLEDGE_BASE_CALL);
+                    } else if (kind == FunctionData.Kind.FUNCTION && isAiKnowledgeBase(classSymbol)) {
+                        nodeBuilder = NodeBuilder.getNodeFromKind(NodeKind.KNOWLEDGE_BASE_CALL);
                     } else if (kind == FunctionData.Kind.FUNCTION) {
                         nodeBuilder = NodeBuilder.getNodeFromKind(NodeKind.METHOD_CALL);
                     } else {
@@ -461,7 +461,7 @@ private Optional<Category> getEmbeddingProvider(Symbol symbol) {
     }
 
     private Optional<Category> getKnowledgeBase(Symbol symbol) {
-        return getCategory(symbol, CommonUtils::isAiVectorKnowledgeBase);
+        return getCategory(symbol, CommonUtils::isAiKnowledgeBase);
     }
 
     private Optional<Category> getVectorStore(Symbol symbol) {

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

@@ -186,12 +186,12 @@
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiChunker;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiDataLoader;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiEmbeddingProvider;
+import static io.ballerina.modelgenerator.commons.CommonUtils.isAiKnowledgeBase;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiMemory;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiMemoryStore;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiMcpBaseToolKit;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiModelModule;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiModelProvider;
-import static io.ballerina.modelgenerator.commons.CommonUtils.isAiVectorKnowledgeBase;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiVectorStore;
 
 /**
@@ -1522,8 +1522,8 @@ private NodeKind resolveNodeKind(ClassSymbol classSymbol) {
         if (isAiEmbeddingProvider(classSymbol)) {
             return NodeKind.EMBEDDING_PROVIDER;
         }
-        if (isAiVectorKnowledgeBase(classSymbol)) {
-            return NodeKind.VECTOR_KNOWLEDGE_BASE;
+        if (isAiKnowledgeBase(classSymbol)) {
+            return NodeKind.KNOWLEDGE_BASE;
         }
         if (isAiVectorStore(classSymbol)) {
             return NodeKind.VECTOR_STORE;
@@ -1556,7 +1556,7 @@ private FunctionData.Kind getFunctionResultKind(ClassSymbol classSymbol) {
         Map<Predicate<ClassSymbol>, FunctionData.Kind> kindMappings = Map.of(
                 CommonUtils::isAiModelProvider, FunctionData.Kind.MODEL_PROVIDER,
                 CommonUtils::isAiEmbeddingProvider, FunctionData.Kind.EMBEDDING_PROVIDER,
-                CommonUtils::isAiVectorKnowledgeBase, FunctionData.Kind.VECTOR_KNOWLEDGE_BASE,
+                CommonUtils::isAiKnowledgeBase, FunctionData.Kind.KNOWLEDGE_BASE,
                 CommonUtils::isAiVectorStore, FunctionData.Kind.VECTOR_STORE,
                 CommonUtils::isAiDataLoader, FunctionData.Kind.DATA_LOADER,
                 CommonUtils::isAiChunker, FunctionData.Kind.CHUNKER,
@@ -1809,8 +1809,8 @@ public void visit(MethodCallExpressionNode methodCallExpressionNode) {
         if (isAgentClass(classSymbol)) {
             startNode(NodeKind.AGENT_CALL, expressionNode.parent());
             populateAgentMetaData(expressionNode, classSymbol);
-        } else if (isAiVectorKnowledgeBase(classSymbol)) {
-            startNode(NodeKind.VECTOR_KNOWLEDGE_BASE_CALL, expressionNode.parent());
+        } else if (isAiKnowledgeBase(classSymbol)) {
+            startNode(NodeKind.KNOWLEDGE_BASE_CALL, expressionNode.parent());
         } else {
             startNode(NodeKind.METHOD_CALL, methodCallExpressionNode.parent());
         }

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

@@ -105,7 +105,7 @@ private Ai() {
         public static final String RECURSIVE_DOCUMENT_CHUNKER_LABEL = "Recursive Document Chunker";
         public static final String AUGMENT_QUERY_LABEL = "Augment Query";
 
-        public static final String VECTOR_KNOWLEDGE_BASE_TYPE_NAME = "VectorKnowledgeBase";
+        public static final String KNOWLEDGE_BASE_TYPE_NAME = "VectorKnowledgeBase";
         public static final String AGENT_TYPE_NAME = "Agent";
 
         public static final String AGENT_RUN_METHOD_NAME = "run";

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

@@ -62,6 +62,7 @@ public class LocalIndexCentral {
     private static final String CLASS_INITS_JSON = "class_inits.json";
     private static final String MEMORY_MANAGERS_JSON = "memory_managers.json";
     private static final String MEMORY_STORES_JSON = "memory_stores.json";
+    private static final String KNOWLEDGE_BASES_JSON = "knowledge_bases.json";
 
     private static final class CentralProxyHolder {
 
@@ -116,6 +117,11 @@ public List<Item> getMemoryStores() {
         return memoryStores.items();
     }
 
+    public List<Item> getKnowledgeBases() {
+        Category knowledgeBases = readJsonResource(KNOWLEDGE_BASES_JSON, Category.class);
+        return knowledgeBases.items();
+    }
+
     public List<Item> getConnectorActions(Codedata codedata) {
         if (connectionMap == null) {
             initializeConnectionMap();

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

@@ -78,8 +78,8 @@
 import java.util.stream.Stream;
 
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAgentClass;
+import static io.ballerina.modelgenerator.commons.CommonUtils.isAiKnowledgeBase;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiMemoryStore;
-import static io.ballerina.modelgenerator.commons.CommonUtils.isAiVectorKnowledgeBase;
 import static io.ballerina.modelgenerator.commons.CommonUtils.isAiVectorStore;
 
 /**
@@ -544,7 +544,6 @@ private Optional<FlowNode> buildFlowNode(Symbol symbol) {
         return codeAnalyzer.getFlowNodes().stream().findFirst();
     }
 
-
     /**
      * Gets the assignment statement node if a ClassFieldSymbol is initialized in an init() function.
      *
@@ -607,9 +606,8 @@ private Optional<NonTerminalNode> getInitAssignmentNode(ClassFieldSymbol classFi
 
     private boolean isClassOrObject(TypeSymbol typeSymbol) {
         if (typeSymbol.kind() == SymbolKind.CLASS) {
-            if (((ClassSymbol) typeSymbol).qualifiers().contains(Qualifier.CLIENT) || isAgentClass(typeSymbol)
-                    || isAiVectorStore(typeSymbol) || isAiVectorKnowledgeBase(typeSymbol) ||
-                    isAiMemoryStore(typeSymbol)) {
+            if (((ClassSymbol) typeSymbol).qualifiers().contains(Qualifier.CLIENT) || isAgentClass(typeSymbol) ||
+                    isAiVectorStore(typeSymbol) || isAiKnowledgeBase(typeSymbol) || isAiMemoryStore(typeSymbol)) {
                 return true;
             }
         }

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

@@ -62,8 +62,8 @@ public enum Name {
                 List.of("Imported", "Function", "Library")),
         IMPORTED_TYPES("Imported Types", "Types imported from other integrations",
                 List.of("Imported", "Type", "Library")),
-        VECTOR_KNOWLEDGE_BASE("Vector Knowledge Bases",
-                "Vector knowledge bases available in the integration", null),
+        KNOWLEDGE_BASE("Knowledge Bases",
+                "Knowledge bases available in the integration", null),
         MODEL_PROVIDER("Model Providers",
                 "Model providers used in the integration to connect to LLMs", null),
         EMBEDDING_PROVIDER("Embedding Providers",

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

@@ -50,6 +50,8 @@
 import io.ballerina.flowmodelgenerator.core.model.node.FunctionDefinitionBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.IfBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.JsonPayloadBuilder;
+import io.ballerina.flowmodelgenerator.core.model.node.KnowledgeBaseBuilder;
+import io.ballerina.flowmodelgenerator.core.model.node.KnowledgeBaseCallBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.LockBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.MatchBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.MemoryBuilder;
@@ -70,8 +72,6 @@
 import io.ballerina.flowmodelgenerator.core.model.node.StopBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.TransactionBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.VariableBuilder;
-import io.ballerina.flowmodelgenerator.core.model.node.VectorKnowledgeBaseBuilder;
-import io.ballerina.flowmodelgenerator.core.model.node.VectorKnowledgeBaseCallBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.VectorStoreBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.WaitBuilder;
 import io.ballerina.flowmodelgenerator.core.model.node.WhileBuilder;
@@ -163,8 +163,8 @@ public abstract class NodeBuilder implements DiagnosticHandler.DiagnosticCapable
         put(NodeKind.MODEL_PROVIDER, ModelProviderBuilder::new);
         put(NodeKind.EMBEDDING_PROVIDER, EmbeddingProviderBuilder::new);
         put(NodeKind.VECTOR_STORE, VectorStoreBuilder::new);
-        put(NodeKind.VECTOR_KNOWLEDGE_BASE, VectorKnowledgeBaseBuilder::new);
-        put(NodeKind.VECTOR_KNOWLEDGE_BASE_CALL, VectorKnowledgeBaseCallBuilder::new);
+        put(NodeKind.KNOWLEDGE_BASE, KnowledgeBaseBuilder::new);
+        put(NodeKind.KNOWLEDGE_BASE_CALL, KnowledgeBaseCallBuilder::new);
         put(NodeKind.DATA_LOADER, DataLoaderBuilder::new);
         put(NodeKind.CHUNKER, ChunkerBuilder::new);
         put(NodeKind.MCP_TOOL_KIT, McpToolKitBuilder::new);

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

@@ -86,9 +86,9 @@ public enum NodeKind {
     EMBEDDING_PROVIDERS,
     VECTOR_STORE,
     VECTOR_STORES,
-    VECTOR_KNOWLEDGE_BASE,
-    VECTOR_KNOWLEDGE_BASE_CALL,
-    VECTOR_KNOWLEDGE_BASES,
+    KNOWLEDGE_BASE,
+    KNOWLEDGE_BASE_CALL,
+    KNOWLEDGE_BASES,
     DATA_LOADERS,
     DATA_LOADER,
     CHUNKERS,

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

@@ -116,12 +116,12 @@ public SourceBuilder(FlowNode flowNode, WorkspaceManager workspaceManager, Path
     private Path resolvePath(Path inputPath, NodeKind node, LineRange lineRange, Boolean isNew) {
         if (Boolean.TRUE.equals(isNew) || lineRange == null) {
             String defaultFile = switch (node) {
-                case NEW_CONNECTION, MODEL_PROVIDER, EMBEDDING_PROVIDER, VECTOR_STORE, VECTOR_KNOWLEDGE_BASE,
-                     DATA_LOADER, CHUNKER, CLASS_INIT, MEMORY_STORE -> CONNECTIONS_BAL;
+                case NEW_CONNECTION, MODEL_PROVIDER, EMBEDDING_PROVIDER, VECTOR_STORE, KNOWLEDGE_BASE,
+                     DATA_LOADER, CHUNKER, CLASS_INIT -> CONNECTIONS_BAL;
                 case DATA_MAPPER_DEFINITION -> DATA_MAPPINGS_BAL;
                 case FUNCTION_DEFINITION, NP_FUNCTION, NP_FUNCTION_DEFINITION -> FUNCTIONS_BAL;
                 case AUTOMATION -> AUTOMATION_BAL;
-                case AGENT, MEMORY, MCP_TOOL_KIT -> AGENTS_BAL;
+                case AGENT, MEMORY, MEMORY_STORE, MCP_TOOL_KIT -> AGENTS_BAL;
                 default -> null;
             };
             if (defaultFile == null) {