Skip to content

Commit 7b3fec6

Browse files
MohamedSabtharMohamedSabthar
authored
Merge pull request #128 from MohamedSabthar/ai-nodes
Introduces new NodeKinds related to AI components and expands relevant APIs
2 parents a66466c + 0a7393a commit 7b3fec6

File tree

95 files changed

+6585
-1711
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

95 files changed

+6585
-1711
lines changed

‎build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@ subprojects {
9393
ballerinaStdLibs "io.ballerina:observe-ballerina:${observeInternalVersion}"
9494
ballerinaStdLibs "io.ballerina.stdlib:observe-ballerina:${observeVersion}"
9595
ballerinaStdLibs "io.ballerina.stdlib:tcp-ballerina:${stdlibTcpVersion}"
96+
ballerinaStdLibs "io.ballerina.stdlib:ai-ballerina:${stdlibAiVersion}"
9697

9798
/* Ballerina Persist Tool */
9899
ballerinaStdLibs "io.ballerina:persist-tools:${persistToolVersion}"

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

Lines changed: 167 additions & 46 deletions
Large diffs are not rendered by default.

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,4 +71,27 @@ private NaturalFunctions() {}
7171
public static final String ICON =
7272
"https://gist.github.com/user-attachments/assets/903c5c16-7d67-4af8-8113-ce7c59ccdaab";
7373
}
74+
75+
// Constants used for AI
76+
public static final class Ai {
77+
private Ai() {}
78+
79+
public static final String BALLERINA_ORG = "ballerina";
80+
public static final String AI_PACKAGE = "ai";
81+
public static final String VERSION = "1.1.0";
82+
83+
public static final String RECURSIVE_DOCUMENT_CHUNKER_LABEL = "Recursive Document Chunker";
84+
public static final String AUGMENT_QUERY_LABEL = "Augment Query";
85+
86+
public static final String MODEL_PROVIDER_TYPE_NAME = "ModelProvider";
87+
public static final String EMBEDDING_PROVIDER_TYPE_NAME = "EmbeddingProvider";
88+
public static final String KNOWLEDGE_BASE_TYPE_NAME = "KnowledgeBase";
89+
public static final String VECTOR_KNOWLEDGE_BASE_TYPE_NAME = "VectorKnowledgeBase";
90+
public static final String VECTOR_STORE_TYPE_NAME = "VectorStore";
91+
public static final String AGENT_TYPE_NAME = "Agent";
92+
93+
public static final String AGENT_RUN_METHOD_NAME = "run";
94+
public static final String CHUNK_DOCUMENT_RECURSIVELY_METHOD_NAME = "chunkDocumentRecursively";
95+
public static final String AUGMENT_USER_QUERY_METHOD_NAME = "augmentUserQuery";
96+
}
7497
}

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

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ public class LocalIndexCentral {
5656
private Map<String, List<Item>> connectionMap;
5757
private static final String NODE_TEMPLATES_JSON = "node_templates.json";
5858
private static final String CONNECTORS_JSON = "connectors.json";
59+
private static final String MODEL_PROVIDERS_JSON = "model_providers.json";
60+
private static final String VECTOR_STORES_JSON = "vector_stores.json";
61+
private static final String EMBEDDING_PROVIDERS_JSON = "embedding_providers.json";
5962
private static final String CONNECTIONS_JSON = "connections.json";
6063
private static final String FUNCTIONS_JSON = "functions.json";
6164

@@ -87,6 +90,21 @@ public List<Item> getConnectors() {
8790
return connectors.items();
8891
}
8992

93+
public List<Item> getModelProviders() {
94+
Category modelProviders = readJsonResource(MODEL_PROVIDERS_JSON, Category.class);
95+
return modelProviders.items();
96+
}
97+
98+
public List<Item> getEmbeddingProviders() {
99+
Category embeddingProviders = readJsonResource(EMBEDDING_PROVIDERS_JSON, Category.class);
100+
return embeddingProviders.items();
101+
}
102+
103+
public List<Item> getVectorStores() {
104+
Category vectorStores = readJsonResource(VECTOR_STORES_JSON, Category.class);
105+
return vectorStores.items();
106+
}
107+
90108
public List<Item> getFunctions() {
91109
Category functions = readJsonResource(FUNCTIONS_JSON, Category.class);
92110
return functions.items();
@@ -119,7 +137,7 @@ public List<AvailableNode> getConnectors(Map<String, String> queryMap) {
119137
.limit(limit)
120138
.toList();
121139
}
122-
140+
123141
private List<AvailableNode> getAvailableNodesFromCategory(Category category) {
124142
List<AvailableNode> availableNodes = new ArrayList<>();
125143
for (Item item : category.items()) {
@@ -140,7 +158,7 @@ private void initializeConnectionMap() {
140158
connectionMap = readJsonResource(CONNECTIONS_JSON, new ConnectionTypeToken().getType());
141159
}
142160

143-
public <T> T readJsonResource(String resourcePath, Type type) {
161+
public <T> T readJsonResource(String resourcePath, Type type) {
144162
InputStream resourceStream = getClass().getClassLoader().getResourceAsStream(resourcePath);
145163
if (resourceStream == null) {
146164
throw new IllegalArgumentException("Resource not found: " + resourcePath);

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,11 @@
2424
* @param metadata the metadata of the node
2525
* @param codedata the codedata of the node
2626
* @param enabled whether the node is enabled
27+
* @param more indicates whether the node should be placed under the "more" section
2728
* @since 1.0.0
2829
*/
29-
public record AvailableNode(Metadata metadata, Codedata codedata, boolean enabled)
30-
implements Item {
31-
30+
public record AvailableNode(Metadata metadata, Codedata codedata, boolean enabled, boolean more) implements Item {
31+
public AvailableNode(Metadata metadata, Codedata codedata, boolean enabled) {
32+
this(metadata, codedata, enabled, false);
33+
}
3234
}

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,15 @@ public enum Name {
6262
IMPORTED_TYPES("Imported Types", "Types imported from other integrations",
6363
List.of("Imported", "Type", "Library")),
6464
AVAILABLE_TYPES("Available Types", "Types available in the library",
65-
List.of("Available", "Type", "Library"));
65+
List.of("Available", "Type", "Library")),
66+
VECTOR_KNOWLEDGE_BASE("Vector Knowledge Bases",
67+
"Vector knowledge bases available in the integration", null),
68+
MODEL_PROVIDER("Model Providers",
69+
"Model providers used in the integration to connect to LLMs", null),
70+
EMBEDDING_PROVIDER("Embedding Providers",
71+
"Embedding providers used in the integration to connect to embedding models", null),
72+
VECTOR_STORE("Vector Stores", "Vector stores used in the integration", null),
73+
AI("AI Components", "AI components available in the flow", null);
6674

6775
final String name;
6876
final String description;

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import io.ballerina.flowmodelgenerator.core.model.node.DataMapperBuilder;
3838
import io.ballerina.flowmodelgenerator.core.model.node.DataMapperCallBuilder;
3939
import io.ballerina.flowmodelgenerator.core.model.node.DataMapperDefinitionBuilder;
40+
import io.ballerina.flowmodelgenerator.core.model.node.EmbeddingProviderBuilder;
4041
import io.ballerina.flowmodelgenerator.core.model.node.ErrorHandlerBuilder;
4142
import io.ballerina.flowmodelgenerator.core.model.node.EventStartBuilder;
4243
import io.ballerina.flowmodelgenerator.core.model.node.ExpressionBuilder;
@@ -50,6 +51,7 @@
5051
import io.ballerina.flowmodelgenerator.core.model.node.LockBuilder;
5152
import io.ballerina.flowmodelgenerator.core.model.node.MatchBuilder;
5253
import io.ballerina.flowmodelgenerator.core.model.node.MethodCall;
54+
import io.ballerina.flowmodelgenerator.core.model.node.ModelProviderBuilder;
5355
import io.ballerina.flowmodelgenerator.core.model.node.NPFunctionCall;
5456
import io.ballerina.flowmodelgenerator.core.model.node.NPFunctionDefinitionBuilder;
5557
import io.ballerina.flowmodelgenerator.core.model.node.NewConnectionBuilder;
@@ -64,6 +66,8 @@
6466
import io.ballerina.flowmodelgenerator.core.model.node.StopBuilder;
6567
import io.ballerina.flowmodelgenerator.core.model.node.TransactionBuilder;
6668
import io.ballerina.flowmodelgenerator.core.model.node.VariableBuilder;
69+
import io.ballerina.flowmodelgenerator.core.model.node.VectorKnowledgeBaseBuilder;
70+
import io.ballerina.flowmodelgenerator.core.model.node.VectorStoreBuilder;
6771
import io.ballerina.flowmodelgenerator.core.model.node.WaitBuilder;
6872
import io.ballerina.flowmodelgenerator.core.model.node.WhileBuilder;
6973
import io.ballerina.flowmodelgenerator.core.model.node.XmlPayloadBuilder;
@@ -149,6 +153,10 @@ public abstract class NodeBuilder implements DiagnosticHandler.DiagnosticCapable
149153
put(NodeKind.AGENT, AgentBuilder::new);
150154
put(NodeKind.AGENT_CALL, AgentCallBuilder::new);
151155
put(NodeKind.CLASS_INIT, ClassInitBuilder::new);
156+
put(NodeKind.MODEL_PROVIDER, ModelProviderBuilder::new);
157+
put(NodeKind.EMBEDDING_PROVIDER, EmbeddingProviderBuilder::new);
158+
put(NodeKind.VECTOR_STORE, VectorStoreBuilder::new);
159+
put(NodeKind.VECTOR_KNOWLEDGE_BASE, VectorKnowledgeBaseBuilder::new);
152160
}};
153161

154162
public static NodeBuilder getNodeFromKind(NodeKind kind) {

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,4 +79,13 @@ public enum NodeKind {
7979
AGENT,
8080
AGENT_CALL,
8181
CLASS_INIT,
82+
83+
MODEL_PROVIDER,
84+
MODEL_PROVIDERS,
85+
EMBEDDING_PROVIDER,
86+
EMBEDDING_PROVIDERS,
87+
VECTOR_STORE,
88+
VECTOR_STORES,
89+
VECTOR_KNOWLEDGE_BASE,
90+
VECTOR_KNOWLEDGE_BASES
8291
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,9 @@ public SourceBuilder(FlowNode flowNode, WorkspaceManager workspaceManager, Path
116116
private Path resolvePath(Path inputPath, NodeKind node, LineRange lineRange, Boolean isNew) {
117117
if (Boolean.TRUE.equals(isNew) || lineRange == null) {
118118
String defaultFile = switch (node) {
119-
case NEW_CONNECTION -> CONNECTIONS_BAL;
119+
case NEW_CONNECTION,
120+
MODEL_PROVIDER, EMBEDDING_PROVIDER,
121+
VECTOR_STORE, VECTOR_KNOWLEDGE_BASE -> CONNECTIONS_BAL;
120122
case DATA_MAPPER_DEFINITION -> DATA_MAPPINGS_BAL;
121123
case FUNCTION_DEFINITION, NP_FUNCTION, NP_FUNCTION_DEFINITION -> FUNCTIONS_BAL;
122124
case AUTOMATION -> AUTOMATION_BAL;

‎flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/model/node/EmbeddingProviderBuilder.java‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
/*
2+
* Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com)
3+
*
4+
* WSO2 LLC. licenses this file to you under the Apache License,
5+
* Version 2.0 (the "License"); you may not use this file except
6+
* in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing,
12+
* software distributed under the License is distributed on an
13+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
* KIND, either express or implied. See the License for the
15+
* specific language governing permissions and limitations
16+
* under the License.
17+
*/
18+
19+
package io.ballerina.flowmodelgenerator.core.model.node;
20+
21+
import io.ballerina.compiler.syntax.tree.SyntaxKind;
22+
import io.ballerina.flowmodelgenerator.core.model.Codedata;
23+
import io.ballerina.flowmodelgenerator.core.model.NodeBuilder;
24+
import io.ballerina.flowmodelgenerator.core.model.NodeKind;
25+
import io.ballerina.flowmodelgenerator.core.model.Property;
26+
import io.ballerina.flowmodelgenerator.core.model.SourceBuilder;
27+
import io.ballerina.modelgenerator.commons.CommonUtils;
28+
import io.ballerina.modelgenerator.commons.FunctionData;
29+
import io.ballerina.modelgenerator.commons.FunctionDataBuilder;
30+
import io.ballerina.modelgenerator.commons.ModuleInfo;
31+
import org.eclipse.lsp4j.TextEdit;
32+
33+
import java.nio.file.Path;
34+
import java.util.List;
35+
import java.util.Map;
36+
import java.util.Set;
37+
38+
import static io.ballerina.modelgenerator.commons.FunctionDataBuilder.GET_DEFAULT_EMBEDDING_PROVIDER_FUNCTION_NAME;
39+
40+
/**
41+
* Represents embedding provider node in the flow model.
42+
*
43+
* @since 1.1.0
44+
*/
45+
public class EmbeddingProviderBuilder extends CallBuilder {
46+
public static final String LABEL = "Embedding Provider";
47+
public static final String DESCRIPTION = "Embedding providers available within the integration for connecting" +
48+
" to an embedding model";
49+
50+
private static final String EMBEDDING_PROVIDER_NAME_LABEL = "Embedding Provider Name";
51+
private static final String EMBEDDING_PROVIDER_NAME_LABEL_DOC = "Name of the embedding-provider connection";
52+
private static final String CHECK_ERROR_DOC = "Terminate on error";
53+
54+
@Override
55+
public void setConcreteConstData() {
56+
metadata().label(LABEL);
57+
codedata().node(NodeKind.EMBEDDING_PROVIDER);
58+
}
59+
60+
@Override
61+
protected NodeKind getFunctionNodeKind() {
62+
return NodeKind.EMBEDDING_PROVIDER;
63+
}
64+
65+
@Override
66+
protected FunctionData.Kind getFunctionResultKind() {
67+
return FunctionData.Kind.EMBEDDING_PROVIDER;
68+
}
69+
70+
@Override
71+
public Map<Path, List<TextEdit>> toSource(SourceBuilder sourceBuilder) {
72+
sourceBuilder.token().keyword(SyntaxKind.FINAL_KEYWORD).stepOut().newVariable();
73+
sourceBuilder.token().keyword(SyntaxKind.CHECK_KEYWORD);
74+
if (sourceBuilder.flowNode.codedata().symbol().equals(GET_DEFAULT_EMBEDDING_PROVIDER_FUNCTION_NAME)) {
75+
sourceBuilder.token().name(methodCallWithModulePrefix(sourceBuilder));
76+
} else {
77+
sourceBuilder.token().keyword(SyntaxKind.NEW_KEYWORD);
78+
}
79+
sourceBuilder.functionParameters(sourceBuilder.flowNode,
80+
Set.of(Property.VARIABLE_KEY, Property.TYPE_KEY, Property.SCOPE_KEY, Property.CHECK_ERROR_KEY));
81+
sourceBuilder.textEdit();
82+
sourceBuilder.acceptImport();
83+
return sourceBuilder.build();
84+
}
85+
86+
private static String methodCallWithModulePrefix(SourceBuilder sourceBuilder) {
87+
String module = sourceBuilder.flowNode.codedata().module();
88+
String methodCallPrefix = (module != null) ? module.substring(module.lastIndexOf('.') + 1) + ":" : "";
89+
return methodCallPrefix + GET_DEFAULT_EMBEDDING_PROVIDER_FUNCTION_NAME;
90+
}
91+
92+
@Override
93+
public void setConcreteTemplateData(NodeBuilder.TemplateContext context) {
94+
Codedata codedata = context.codedata();
95+
ModuleInfo codedataModuleInfo = new ModuleInfo(codedata.org(), codedata.packageName(),
96+
codedata.module(), codedata.version());
97+
98+
FunctionData functionData = new FunctionDataBuilder().moduleInfo(codedataModuleInfo).userModuleInfo(moduleInfo)
99+
.parentSymbolType(codedata.object()).name(codedata.symbol())
100+
.lsClientLogger(context.lsClientLogger()).functionResultKind(FunctionData.Kind.EMBEDDING_PROVIDER)
101+
.build();
102+
103+
metadata().label(functionData.packageName()).description(functionData.description())
104+
.icon(CommonUtils.generateIcon(functionData.org(), functionData.packageName(), functionData.version()));
105+
106+
codedata().node(NodeKind.EMBEDDING_PROVIDER)
107+
.org(functionData.org()).module(functionData.moduleName())
108+
.packageName(functionData.packageName()).version(functionData.version())
109+
.symbol(codedata.symbol());
110+
111+
if (!functionData.name().equals(GET_DEFAULT_EMBEDDING_PROVIDER_FUNCTION_NAME)) {
112+
codedata().object(functionData.name());
113+
}
114+
115+
if (CommonUtils.hasReturn(functionData.returnType())) {
116+
setReturnTypeProperties(functionData, context, EMBEDDING_PROVIDER_NAME_LABEL,
117+
EMBEDDING_PROVIDER_NAME_LABEL_DOC, false);
118+
}
119+
setParameterProperties(functionData);
120+
properties().scope(Property.GLOBAL_SCOPE).checkError(true, CHECK_ERROR_DOC, false);
121+
}
122+
}

0 commit comments

Comments
 (0)