diff --git a/architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/artifactsgenerator/ModuleNodeTransformer.java b/architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/artifactsgenerator/ModuleNodeTransformer.java index 85432a885f..cbc453d46d 100644 --- a/architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/artifactsgenerator/ModuleNodeTransformer.java +++ b/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 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); } diff --git a/architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/designmodelgenerator/core/CommonUtils.java b/architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/designmodelgenerator/core/CommonUtils.java index cebc82f246..5ba04d7f35 100644 --- a/architecture-model-generator/modules/architecture-model-generator-core/src/main/java/io/ballerina/designmodelgenerator/core/CommonUtils.java +++ b/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"; diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/AiUtils.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/AiUtils.java index e3576b5518..adc8f7afe8 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/AiUtils.java +++ b/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") + )); } /** diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/AvailableNodesGenerator.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/AvailableNodesGenerator.java index 8c9ce51a82..1f7b42f65a 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/AvailableNodesGenerator.java +++ b/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 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 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 getCategory(Symbol symbol, Predicate 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 getEmbeddingProvider(Symbol symbol) { } private Optional getKnowledgeBase(Symbol symbol) { - return getCategory(symbol, CommonUtils::isAiVectorKnowledgeBase); + return getCategory(symbol, CommonUtils::isAiKnowledgeBase); } private Optional getVectorStore(Symbol symbol) { diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/CodeAnalyzer.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/CodeAnalyzer.java index 1fff25c3c6..ead62ec4eb 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/CodeAnalyzer.java +++ b/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, 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()); } diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/Constants.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/Constants.java index 2d35589b94..9344749ef1 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/Constants.java +++ b/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"; diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/LocalIndexCentral.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/LocalIndexCentral.java index 0c4a17ee10..7fbcd263d9 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/LocalIndexCentral.java +++ b/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 getMemoryStores() { return memoryStores.items(); } + public List getKnowledgeBases() { + Category knowledgeBases = readJsonResource(KNOWLEDGE_BASES_JSON, Category.class); + return knowledgeBases.items(); + } + public List getConnectorActions(Codedata codedata) { if (connectionMap == null) { initializeConnectionMap(); diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/ModelGenerator.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/ModelGenerator.java index 57c8ce32b6..f45acb6c8b 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/ModelGenerator.java +++ b/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 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 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; } } diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/Category.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/Category.java index cb7e0bff4f..97cdc33d25 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/Category.java +++ b/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", diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/NodeBuilder.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/NodeBuilder.java index 2bfadc583b..166afe6c40 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/NodeBuilder.java +++ b/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); diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/NodeKind.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/NodeKind.java index b34cc59675..98c715d486 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/NodeKind.java +++ b/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, diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/SourceBuilder.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/SourceBuilder.java index 34f146161b..3c456ec080 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/SourceBuilder.java +++ b/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) { diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/VectorKnowledgeBaseBuilder.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/KnowledgeBaseBuilder.java similarity index 72% rename from flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/VectorKnowledgeBaseBuilder.java rename to flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/KnowledgeBaseBuilder.java index 0601e801ba..86dc9b5bf3 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/VectorKnowledgeBaseBuilder.java +++ b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/KnowledgeBaseBuilder.java @@ -36,38 +36,44 @@ import java.util.Set; /** - * Represents vector knowledge base node in the flow model. + * Represents knowledge base node in the flow model. * * @since 1.1.0 */ -public class VectorKnowledgeBaseBuilder extends CallBuilder { - public static final String LABEL = "Vector Knowledge Base"; - public static final String DESCRIPTION = "Vector knowledge bases available in the integration"; +public class KnowledgeBaseBuilder extends CallBuilder { - private static final String VECTOR_KNOWLEDGE_BASE_NAME_LABEL = "Vector Knowledge Base Name"; - private static final String VECTOR_KNOWLEDGE_BASE_NAME_LABEL_DOC = "Vector knowledge-base instance name"; + public static final String LABEL = "Knowledge Base"; + public static final String DESCRIPTION = "Knowledge bases available in the integration"; + + private static final String KNOWLEDGE_BASE_NAME_LABEL = "Knowledge Base Name"; + private static final String KNOWLEDGE_BASE_NAME_LABEL_DOC = "Name of the knowledge base instance"; private static final String CHECK_ERROR_DOC = "Terminate on error"; + private static final String VECTOR_KNOWLEDGE_BASE_TYPE = "VectorKnowledgeBase"; @Override public void setConcreteConstData() { metadata().label(LABEL); - codedata().node(NodeKind.VECTOR_KNOWLEDGE_BASE); + codedata().node(NodeKind.KNOWLEDGE_BASE); } @Override protected NodeKind getFunctionNodeKind() { - return NodeKind.VECTOR_KNOWLEDGE_BASE; + return NodeKind.KNOWLEDGE_BASE; } @Override protected FunctionData.Kind getFunctionResultKind() { - return FunctionData.Kind.VECTOR_KNOWLEDGE_BASE; + return FunctionData.Kind.KNOWLEDGE_BASE; } @Override public Map> toSource(SourceBuilder sourceBuilder) { - return sourceBuilder.token().keyword(SyntaxKind.FINAL_KEYWORD).stepOut().newVariable() - .token().keyword(SyntaxKind.NEW_KEYWORD).stepOut().functionParameters(sourceBuilder.flowNode, + sourceBuilder.token().keyword(SyntaxKind.FINAL_KEYWORD).stepOut().newVariable(); + if (!sourceBuilder.flowNode.codedata().object().equals(VECTOR_KNOWLEDGE_BASE_TYPE)) { + sourceBuilder.token().keyword(SyntaxKind.CHECK_KEYWORD); + } + return sourceBuilder.token().keyword(SyntaxKind.NEW_KEYWORD).stepOut() + .functionParameters(sourceBuilder.flowNode, Set.of(Property.VARIABLE_KEY, Property.TYPE_KEY, Property.SCOPE_KEY, Property.CHECK_ERROR_KEY)) .textEdit().acceptImport().build(); } @@ -80,7 +86,7 @@ public void setConcreteTemplateData(NodeBuilder.TemplateContext context) { FunctionData functionData = new FunctionDataBuilder().parentSymbolType(codedata.object()) .name(codedata.symbol()).moduleInfo(codedataModuleInfo).userModuleInfo(moduleInfo) - .lsClientLogger(context.lsClientLogger()).functionResultKind(FunctionData.Kind.VECTOR_KNOWLEDGE_BASE) + .lsClientLogger(context.lsClientLogger()).functionResultKind(FunctionData.Kind.KNOWLEDGE_BASE) .build(); metadata().label(functionData.packageName()).description(functionData.description()) @@ -91,8 +97,8 @@ public void setConcreteTemplateData(NodeBuilder.TemplateContext context) { .version(functionData.version()); if (CommonUtils.hasReturn(functionData.returnType())) { - setReturnTypeProperties(functionData, context, VECTOR_KNOWLEDGE_BASE_NAME_LABEL, - VECTOR_KNOWLEDGE_BASE_NAME_LABEL_DOC, false); + setReturnTypeProperties(functionData, context, KNOWLEDGE_BASE_NAME_LABEL, + KNOWLEDGE_BASE_NAME_LABEL_DOC, false); } setParameterProperties(functionData); properties().scope(Property.GLOBAL_SCOPE).checkError(true, CHECK_ERROR_DOC, false); diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/VectorKnowledgeBaseCallBuilder.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/KnowledgeBaseCallBuilder.java similarity index 93% rename from flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/VectorKnowledgeBaseCallBuilder.java rename to flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/KnowledgeBaseCallBuilder.java index fe9078f34b..0ff04889e4 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/VectorKnowledgeBaseCallBuilder.java +++ b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/KnowledgeBaseCallBuilder.java @@ -38,16 +38,17 @@ import static io.ballerina.flowmodelgenerator.core.model.Property.VARIABLE_KEY; /** - * Represents a vector knowledge base call node. + * Represents a knowledge base call node. * * @since 1.1.0 */ -public class VectorKnowledgeBaseCallBuilder extends FunctionCall { +public class KnowledgeBaseCallBuilder extends FunctionCall { + private static final String RETRIEVE_METHOD_NAME = "retrieve"; @Override protected NodeKind getFunctionNodeKind() { - return NodeKind.VECTOR_KNOWLEDGE_BASE_CALL; + return NodeKind.KNOWLEDGE_BASE_CALL; } @Override @@ -59,7 +60,7 @@ public Map> toSource(SourceBuilder sourceBuilder) { } Optional connection = flowNode.getProperty(CONNECTION_KEY); if (connection.isEmpty()) { - throw new IllegalStateException("connection must be defined for a vector knowledge base"); + throw new IllegalStateException("connection must be defined for a knowledge base"); } String methodName = flowNode.metadata().label(); SourceBuilder builder = sourceBuilder.token().name(connection.get().toSourceCode()) diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/search/KnowledgeBaseSearchCommand.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/search/KnowledgeBaseSearchCommand.java new file mode 100644 index 0000000000..27c46fd653 --- /dev/null +++ b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/search/KnowledgeBaseSearchCommand.java @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2025, 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.search; + +import com.google.gson.Gson; +import com.google.gson.JsonArray; +import io.ballerina.flowmodelgenerator.core.LocalIndexCentral; +import io.ballerina.flowmodelgenerator.core.model.AvailableNode; +import io.ballerina.flowmodelgenerator.core.model.Category; +import io.ballerina.flowmodelgenerator.core.model.Item; +import io.ballerina.modelgenerator.commons.SearchResult; +import io.ballerina.projects.Project; +import io.ballerina.tools.text.LineRange; + +import java.util.Collections; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +/** + * Handles the search command for knowledge bases. + * + * @since 1.3.0 + */ +public class KnowledgeBaseSearchCommand extends SearchCommand { + + private static final Gson GSON = new Gson(); + private List cachedKnowledgeBases; + + public KnowledgeBaseSearchCommand(Project project, LineRange position, Map queryMap) { + super(project, position, queryMap); + } + + @Override + protected List defaultView() { + return getKnowledgeBases(); + } + + @Override + protected List search() { + List knowledgeBases = getKnowledgeBases(); + if (knowledgeBases.isEmpty() || !(knowledgeBases.getFirst() instanceof Category knowledgeBaseCategory)) { + return knowledgeBases; + } + + List stores = knowledgeBaseCategory.items(); + + List matchingStores = stores.stream() + .filter(item -> item instanceof AvailableNode availableNode && + (query == null || availableNode.metadata().label().toLowerCase(Locale.ROOT) + .contains(query.toLowerCase(Locale.ROOT)))) + .toList(); + + stores.clear(); + stores.addAll(matchingStores); + + return List.of(knowledgeBaseCategory); + } + + @Override + protected Map> fetchPopularItems() { + return Collections.emptyMap(); + } + + @Override + public JsonArray execute() { + List items = (query.isEmpty()) ? defaultView() : search(); + return GSON.toJsonTree(items).getAsJsonArray(); + } + + private List getKnowledgeBases() { + if (cachedKnowledgeBases == null) { + cachedKnowledgeBases = List.copyOf(LocalIndexCentral.getInstance().getKnowledgeBases()); + } + return cachedKnowledgeBases; + } +} diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/search/SearchCommand.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/search/SearchCommand.java index c10e4a42b5..e60d9cecf3 100644 --- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/search/SearchCommand.java +++ b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/search/SearchCommand.java @@ -78,6 +78,7 @@ public static SearchCommand from(Kind kind, Project module, LineRange position, case MEMORY -> new MemoryManagerSearchCommand(module, position, queryMap); case MEMORY_STORE -> new MemoryStoreSearchCommand(module, position, queryMap); case AGENT_TOOL -> new AgentToolSearchCommand(module, position, queryMap); + case KNOWLEDGE_BASE -> new KnowledgeBaseSearchCommand(module, position, queryMap); }; } @@ -206,6 +207,7 @@ public enum Kind { CLASS_INIT, MEMORY, MEMORY_STORE, - AGENT_TOOL + AGENT_TOOL, + KNOWLEDGE_BASE } } diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/main/resources/knowledge_bases.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/main/resources/knowledge_bases.json new file mode 100644 index 0000000000..9ca17bd5e6 --- /dev/null +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/main/resources/knowledge_bases.json @@ -0,0 +1,43 @@ +{ + "items": [ + { + "metadata": { + "label": "Knowledge Bases" + }, + "items": [ + { + "metadata": { + "label": "Vector Knowledge Base", + "description": "An AI knowledge base that manages document embeddings in a vector store for efficient semantic retrieval and RAG workflows." + }, + "codedata": { + "node": "KNOWLEDGE_BASE", + "org": "ballerina", + "module": "ai", + "packageName": "ai", + "object": "VectorKnowledgeBase", + "symbol": "init", + "flags": 1 + }, + "enabled": true + }, + { + "metadata": { + "label": "Azure AI Search Knowledge Base", + "description": "An AI knowledge base powered by Azure AI Search, supporting document ingestion, and semantic retrieval for RAG workflows." + }, + "codedata": { + "node": "KNOWLEDGE_BASE", + "org": "ballerinax", + "module": "ai.azure", + "packageName": "ai.azure", + "object": "AiSearchKnowledgeBase", + "symbol": "init", + "flags": 1 + }, + "enabled": true + } + ] + } + ] +} diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/agents_manager/config/available_agents.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/agents_manager/config/available_agents.json index 9d03bba8b2..64b329c09c 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/agents_manager/config/available_agents.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/agents_manager/config/available_agents.json @@ -845,11 +845,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/caller.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/caller.json index 6028a80011..7849cfdaa4 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/caller.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/caller.json @@ -178,11 +178,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector1.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector1.json index a890f6633e..dcf36f16e8 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector1.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector1.json @@ -425,11 +425,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector2.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector2.json index 281939a8f0..7b4589a719 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector2.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector2.json @@ -4543,11 +4543,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector3.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector3.json index b414251014..c2e8654ec2 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector3.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/connector3.json @@ -140,11 +140,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/force_assign.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/force_assign.json index c930f1fa33..bc7581fd26 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/force_assign.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/force_assign.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/foreach1.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/foreach1.json index fedd288afb..37ce27def5 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/foreach1.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/foreach1.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function1.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function1.json index c930f1fa33..bc7581fd26 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function1.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function1.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function2.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function2.json index c374325caa..4205d19151 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function2.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function2.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function3.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function3.json index b4bff13abc..b8617ca414 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function3.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/function3.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/inside_nested_foreach.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/inside_nested_foreach.json index 69537ca811..a8d5498a10 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/inside_nested_foreach.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/inside_nested_foreach.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/lock1.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/lock1.json index c1d9ac7153..87c5cb74de 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/lock1.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/lock1.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/match1.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/match1.json index 12fa1d829b..f77e710383 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/match1.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/match1.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/match2.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/match2.json index e91df451b7..bf80e757f0 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/match2.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/match2.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/method1.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/method1.json index 554fa2e6e7..886a403d0f 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/method1.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/method1.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/module_with_existing_ai.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/module_with_existing_ai.json index f1d390f22b..a04d817288 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/module_with_existing_ai.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/module_with_existing_ai.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/on_fail_clause1.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/on_fail_clause1.json index 41b6bd2512..fe65f89390 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/on_fail_clause1.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/on_fail_clause1.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/transaction1.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/transaction1.json index 6ac557311b..64b31e6473 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/transaction1.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/transaction1.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/transaction2.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/transaction2.json index 41b6bd2512..fe65f89390 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/transaction2.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/transaction2.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/while1.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/while1.json index aa82b7cf0d..b4ab1b040b 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/while1.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/available_nodes/config/while1.json @@ -56,11 +56,11 @@ "items": [ { "metadata": { - "label": "Vector Knowledge Base", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Base", + "description": "Knowledge bases available in the integration" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASES", + "node": "KNOWLEDGE_BASES", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_for_vector_store_knowledge_base_ingest_call.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_for_vector_store_knowledge_base_ingest_call.json index 34e5402040..dea33a4448 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_for_vector_store_knowledge_base_ingest_call.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_for_vector_store_knowledge_base_ingest_call.json @@ -8,7 +8,7 @@ "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_ai_1.1.2.png" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASE_CALL", + "node": "KNOWLEDGE_BASE_CALL", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_for_vector_store_knowledge_base_retrieve_call.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_for_vector_store_knowledge_base_retrieve_call.json index b791914237..b5ed7762ef 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_for_vector_store_knowledge_base_retrieve_call.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_for_vector_store_knowledge_base_retrieve_call.json @@ -8,7 +8,7 @@ "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_ai_1.1.2.png" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASE_CALL", + "node": "KNOWLEDGE_BASE_CALL", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_with_vector_store_knowledge_base.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_with_vector_store_knowledge_base.json index 3d3a8df1f6..00a8ad7454 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_with_vector_store_knowledge_base.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/code_generation_with_vector_store_knowledge_base.json @@ -8,7 +8,7 @@ "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_ai_1.1.0.png" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASE", + "node": "KNOWLEDGE_BASE", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/get_knowledge_bases.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/get_knowledge_bases.json index 3c0e624449..69aa0378ef 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/get_knowledge_bases.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/get_knowledge_bases.json @@ -4,8 +4,8 @@ "categories": [ { "metadata": { - "label": "Vector Knowledge Bases", - "description": "Vector knowledge bases available in the integration" + "label": "Knowledge Bases", + "description": "Knowledge bases available in the integration" }, "items": [ { @@ -20,7 +20,7 @@ "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_ai_1.5.0.png" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASE_CALL", + "node": "KNOWLEDGE_BASE_CALL", "org": "ballerina", "module": "ai", "packageName": "ai", @@ -39,7 +39,7 @@ "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_ai_1.5.0.png" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASE_CALL", + "node": "KNOWLEDGE_BASE_CALL", "org": "ballerina", "module": "ai", "packageName": "ai", @@ -58,7 +58,7 @@ "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_ai_1.5.0.png" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASE_CALL", + "node": "KNOWLEDGE_BASE_CALL", "org": "ballerina", "module": "ai", "packageName": "ai", diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/get_vector_knowledge_base_node_template.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/get_vector_knowledge_base_node_template.json index 42cf5fec6a..3cc17eff2c 100644 --- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/get_vector_knowledge_base_node_template.json +++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/knowledge_base_manager/config/get_vector_knowledge_base_node_template.json @@ -1,7 +1,7 @@ { "source": "knowledge_base/connections.bal", "codedata": { - "node": "VECTOR_KNOWLEDGE_BASE", + "node": "KNOWLEDGE_BASE", "org": "ballerina", "module": "ai", "packageName": "ai", @@ -13,12 +13,12 @@ "flowNode": { "id": "31", "metadata": { - "label": "Vector Knowledge Base", + "label": "Knowledge Base", "description": "Initializes a new `VectorKnowledgeBase` instance.\n", "icon": "https://bcentral-packageicons.azureedge.net/images/ballerina_ai_1.1.0.png" }, "codedata": { - "node": "VECTOR_KNOWLEDGE_BASE", + "node": "KNOWLEDGE_BASE", "org": "ballerina", "module": "ai", "packageName": "ai", @@ -47,8 +47,8 @@ }, "variable": { "metadata": { - "label": "Vector Knowledge Base Name", - "description": "Vector knowledge-base instance name" + "label": "Knowledge Base Name", + "description": "Name of the knowledge base instance" }, "valueType": "IDENTIFIER", "value": "aiVectorknowledgebase", diff --git a/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/CommonUtils.java b/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/CommonUtils.java index 79e70aa83c..d4b86cc998 100644 --- a/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/CommonUtils.java +++ b/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/CommonUtils.java @@ -95,7 +95,7 @@ public class CommonUtils { private static final String CENTRAL_ICON_URL = "https://bcentral-packageicons.azureedge.net/images/%s_%s_%s.png"; private static final Pattern FULLY_QUALIFIED_MODULE_ID_PATTERN = Pattern.compile("(\\w+)/([\\w.]+):([^:]+):(\\w+)[|]?"); - private static final String VECTOR_KNOWLEDGE_BASE_TYPE_NAME = "VectorKnowledgeBase"; + private static final String KNOWLEDGE_BASE_TYPE_NAME = "KnowledgeBase"; private static final String EMBEDDING_PROVIDER_TYPE_NAME = "EmbeddingProvider"; private static final String MODEL_PROVIDER_TYPE_NAME = "ModelProvider"; private static final String VECTOR_STORE_TYPE_NAME = "VectorStore"; @@ -989,10 +989,9 @@ public static boolean isAgentClass(Symbol symbol) { return symbol.getName().isPresent() && symbol.getName().get().equals(AGENT); } - public static boolean isAiVectorKnowledgeBase(Symbol symbol) { - Optional module = symbol.getModule(); - return module.isPresent() && isAiModule(module.get().id().orgName(), module.get().id().packageName()) - && symbol.getName().isPresent() && symbol.getName().get().equals(VECTOR_KNOWLEDGE_BASE_TYPE_NAME); + public static boolean isAiKnowledgeBase(Symbol symbol) { + ClassSymbol classSymbol = getClassSymbol(symbol); + return classSymbol != null && hasAiTypeInclusion(classSymbol, KNOWLEDGE_BASE_TYPE_NAME); } public static boolean isAiVectorStore(Symbol symbol) { diff --git a/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/FunctionData.java b/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/FunctionData.java index c47b0ef026..29639a230a 100644 --- a/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/FunctionData.java +++ b/model-generator-commons/src/main/java/io/ballerina/modelgenerator/commons/FunctionData.java @@ -140,7 +140,7 @@ public enum Kind { MEMORY_STORE, EMBEDDING_PROVIDER, VECTOR_STORE, - VECTOR_KNOWLEDGE_BASE, + KNOWLEDGE_BASE, DATA_LOADER, CHUNKER, MCP_TOOL_KIT, @@ -157,7 +157,7 @@ public enum Kind { FunctionData.Kind.VECTOR_STORE, FunctionData.Kind.DATA_LOADER, FunctionData.Kind.CHUNKER, - FunctionData.Kind.VECTOR_KNOWLEDGE_BASE, + FunctionData.Kind.KNOWLEDGE_BASE, FunctionData.Kind.MCP_TOOL_KIT );