Skip to content

Commit 0877c7c

Browse files
ivicacclaude
andcommitted
4000 Expose Knowledge Base and Tags on the vector-store and search-tool
Add knowledgeBaseId (required) and tagIds (optional, scoped to the selected knowledge base via optionsLookupDependsOn) as properties on both the Knowledge Base vector-store cluster element (RAG path) and the Knowledge Base search tool cluster element. Tag options come from KnowledgeBaseTagService so only tags actually assigned to a knowledge base surface — per-KB when one is selected, across all KBs otherwise. Route tagIds through a 3-arg KnowledgeBaseVectorStoreWrapper which AND-s an OR-joined tag filter on top of the knowledge-base filter. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent ea87e3c commit 0877c7c

6 files changed

Lines changed: 364 additions & 25 deletions

File tree

server/libs/modules/components/ai/vectorstore/knowledgebase/src/main/java/com/bytechef/component/ai/vectorstore/knowledgebase/KnowledgeBaseComponentHandler.java

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseDocumentChunkService;
2626
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseDocumentService;
2727
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseService;
28+
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseTagService;
2829
import com.bytechef.component.ComponentHandler;
2930
import com.bytechef.component.ai.vectorstore.knowledgebase.action.KnowledgeBaseLoadAction;
3031
import com.bytechef.component.ai.vectorstore.knowledgebase.action.KnowledgeBaseSearchAction;
@@ -59,13 +60,13 @@ public KnowledgeBaseComponentHandler(
5960
ClusterElementDefinitionService clusterElementDefinitionService,
6061
KnowledgeBaseDocumentChunkService knowledgeBaseDocumentChunkService,
6162
KnowledgeBaseDocumentService knowledgeBaseDocumentService, KnowledgeBaseFileStorage knowledgeBaseFileStorage,
62-
KnowledgeBaseService knowledgeBaseService, TagService tagService,
63-
@Qualifier("knowledgeBasePgVectorStore") VectorStore vectorStore) {
63+
KnowledgeBaseService knowledgeBaseService, KnowledgeBaseTagService knowledgeBaseTagService,
64+
TagService tagService, @Qualifier("knowledgeBasePgVectorStore") VectorStore vectorStore) {
6465

6566
this.componentDefinition =
6667
new KnowledgeBaseVectorStoreComponentDefinitionImpl(
6768
clusterElementDefinitionService, knowledgeBaseDocumentChunkService, knowledgeBaseDocumentService,
68-
knowledgeBaseFileStorage, knowledgeBaseService, tagService, vectorStore);
69+
knowledgeBaseFileStorage, knowledgeBaseService, knowledgeBaseTagService, tagService, vectorStore);
6970
}
7071

7172
@Override
@@ -81,7 +82,7 @@ public KnowledgeBaseVectorStoreComponentDefinitionImpl(
8182
KnowledgeBaseDocumentChunkService knowledgeBaseDocumentChunkService,
8283
KnowledgeBaseDocumentService knowledgeBaseDocumentService,
8384
KnowledgeBaseFileStorage knowledgeBaseFileStorage, KnowledgeBaseService knowledgeBaseService,
84-
TagService tagService, VectorStore vectorStore) {
85+
KnowledgeBaseTagService knowledgeBaseTagService, TagService tagService, VectorStore vectorStore) {
8586

8687
super(
8788
component(KNOWLEDGE_BASE)
@@ -97,10 +98,10 @@ public KnowledgeBaseVectorStoreComponentDefinitionImpl(
9798
knowledgeBaseDocumentService, knowledgeBaseFileStorage, knowledgeBaseService),
9899
KnowledgeBaseSearchAction.of(vectorStore, knowledgeBaseService, tagService))
99100
.clusterElements(
100-
KnowledgeBaseSearchTool.of(vectorStore, clusterElementDefinitionService),
101+
KnowledgeBaseSearchTool.of(vectorStore, knowledgeBaseService, knowledgeBaseTagService),
101102
KnowledgeBaseVectorStore.of(
102-
vectorStore, clusterElementDefinitionService, knowledgeBaseDocumentChunkService,
103-
knowledgeBaseDocumentService, knowledgeBaseFileStorage, knowledgeBaseService)));
103+
vectorStore, knowledgeBaseDocumentChunkService, knowledgeBaseDocumentService,
104+
knowledgeBaseFileStorage, knowledgeBaseService, knowledgeBaseTagService)));
104105
}
105106

106107
@Override

server/libs/modules/components/ai/vectorstore/knowledgebase/src/main/java/com/bytechef/component/ai/vectorstore/knowledgebase/cluster/KnowledgeBaseSearchTool.java

Lines changed: 97 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,27 +20,39 @@
2020
import static com.bytechef.component.ai.vectorstore.constant.VectorStoreConstants.SEARCH_PROPERTIES;
2121
import static com.bytechef.component.ai.vectorstore.knowledgebase.cluster.KnowledgeBaseVectorStore.createVectorStore;
2222
import static com.bytechef.component.ai.vectorstore.knowledgebase.constant.KnowledgeBaseVectorStoreConstants.KNOWLEDGE_BASE;
23+
import static com.bytechef.component.ai.vectorstore.knowledgebase.constant.KnowledgeBaseVectorStoreConstants.KNOWLEDGE_BASE_ID;
24+
import static com.bytechef.component.ai.vectorstore.knowledgebase.constant.KnowledgeBaseVectorStoreConstants.TAG_IDS;
25+
import static com.bytechef.component.definition.ComponentDsl.array;
26+
import static com.bytechef.component.definition.ComponentDsl.integer;
27+
import static com.bytechef.component.definition.ComponentDsl.option;
2328
import static com.bytechef.component.definition.ai.agent.BaseToolFunction.TOOLS;
2429

30+
import com.bytechef.automation.knowledgebase.domain.KnowledgeBase;
31+
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseService;
32+
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseTagService;
2533
import com.bytechef.component.ai.vectorstore.VectorStore;
2634
import com.bytechef.component.definition.ClusterElementDefinition;
2735
import com.bytechef.component.definition.ComponentDsl;
36+
import com.bytechef.component.definition.Option;
2837
import com.bytechef.platform.component.ComponentConnection;
2938
import com.bytechef.platform.component.definition.ParametersFactory;
3039
import com.bytechef.platform.component.definition.ai.agent.MultipleConnectionsToolFunction;
31-
import com.bytechef.platform.component.service.ClusterElementDefinitionService;
40+
import com.bytechef.platform.tag.domain.Tag;
41+
import java.util.ArrayList;
42+
import java.util.LinkedHashMap;
43+
import java.util.List;
3244
import java.util.Map;
45+
import java.util.Objects;
3346
import java.util.stream.Stream;
3447

3548
/**
3649
* @author Ivica Cardic
3750
*/
3851
public class KnowledgeBaseSearchTool {
3952

40-
@SuppressWarnings("PMD.UnusedFormalParameter")
4153
public static ClusterElementDefinition<MultipleConnectionsToolFunction> of(
42-
org.springframework.ai.vectorstore.VectorStore vectorStore,
43-
ClusterElementDefinitionService clusterElementDefinitionService) {
54+
org.springframework.ai.vectorstore.VectorStore vectorStore, KnowledgeBaseService knowledgeBaseService,
55+
KnowledgeBaseTagService knowledgeBaseTagService) {
4456

4557
VectorStore kbVectorStore = createVectorStore(vectorStore);
4658

@@ -50,7 +62,24 @@ public static ClusterElementDefinition<MultipleConnectionsToolFunction> of(
5062
.type(TOOLS)
5163
.properties(
5264
Stream
53-
.of(Stream.of(QUERY_PROPERTY), SEARCH_PROPERTIES.stream())
65+
.of(
66+
Stream.of(
67+
integer(KNOWLEDGE_BASE_ID)
68+
.label("Knowledge Base")
69+
.description("The knowledge base to search.")
70+
.options(getKnowledgeBaseOptions(knowledgeBaseService))
71+
.required(true),
72+
array(TAG_IDS)
73+
.label("Tags")
74+
.description(
75+
"Filter results by tags. Documents with ANY of the selected tags will be " +
76+
"returned (OR logic).")
77+
.items(integer())
78+
.options(getTagOptions(knowledgeBaseTagService))
79+
.optionsLookupDependsOn(KNOWLEDGE_BASE_ID)
80+
.required(false),
81+
QUERY_PROPERTY),
82+
SEARCH_PROPERTIES.stream())
5483
.flatMap(stream -> stream)
5584
.toList())
5685
.object(() -> (MultipleConnectionsToolFunction) (
@@ -65,4 +94,67 @@ public static ClusterElementDefinition<MultipleConnectionsToolFunction> of(
6594
null);
6695
});
6796
}
97+
98+
private static ClusterElementDefinition.OptionsFunction<Long> getKnowledgeBaseOptions(
99+
KnowledgeBaseService knowledgeBaseService) {
100+
101+
return (inputParameters, connectionParameters, lookupDependsOnPaths, searchText, context) -> {
102+
List<Option<Long>> options = new ArrayList<>();
103+
104+
List<KnowledgeBase> knowledgeBases = knowledgeBaseService.getKnowledgeBases();
105+
106+
for (KnowledgeBase knowledgeBase : knowledgeBases) {
107+
String name = knowledgeBase.getName();
108+
109+
if (searchText == null || name.toLowerCase()
110+
.contains(searchText.toLowerCase())) {
111+
112+
options.add(option(name, knowledgeBase.getId()
113+
.longValue()));
114+
}
115+
}
116+
117+
return options;
118+
};
119+
}
120+
121+
private static ClusterElementDefinition.OptionsFunction<Long> getTagOptions(
122+
KnowledgeBaseTagService knowledgeBaseTagService) {
123+
124+
return (inputParameters, connectionParameters, lookupDependsOnPaths, searchText, context) -> {
125+
Long knowledgeBaseId = inputParameters.getLong(KNOWLEDGE_BASE_ID);
126+
127+
List<Tag> tags;
128+
129+
if (knowledgeBaseId == null) {
130+
tags = knowledgeBaseTagService.getAllTags();
131+
} else {
132+
tags = knowledgeBaseTagService.getTagsByKnowledgeBaseId()
133+
.getOrDefault(knowledgeBaseId, List.of());
134+
}
135+
136+
Map<Long, Tag> unique = new LinkedHashMap<>();
137+
138+
for (Tag tag : tags) {
139+
Long tagId = Objects.requireNonNull(tag.getId());
140+
141+
unique.putIfAbsent(tagId, tag);
142+
}
143+
144+
List<Option<Long>> options = new ArrayList<>();
145+
146+
for (Tag tag : unique.values()) {
147+
String tagName = tag.getName();
148+
149+
if (searchText == null || tagName.toLowerCase()
150+
.contains(searchText.toLowerCase())) {
151+
152+
options.add(option(tagName, Objects.requireNonNull(tag.getId())
153+
.longValue()));
154+
}
155+
}
156+
157+
return options;
158+
};
159+
}
68160
}

server/libs/modules/components/ai/vectorstore/knowledgebase/src/main/java/com/bytechef/component/ai/vectorstore/knowledgebase/cluster/KnowledgeBaseVectorStore.java

Lines changed: 93 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,10 @@
2222
import static com.bytechef.automation.knowledgebase.constant.KnowledgeBaseConstants.METADATA_KNOWLEDGE_BASE_ID;
2323
import static com.bytechef.component.ai.vectorstore.knowledgebase.constant.KnowledgeBaseVectorStoreConstants.KNOWLEDGE_BASE_ID;
2424
import static com.bytechef.component.ai.vectorstore.knowledgebase.constant.KnowledgeBaseVectorStoreConstants.QUERY;
25+
import static com.bytechef.component.ai.vectorstore.knowledgebase.constant.KnowledgeBaseVectorStoreConstants.TAG_IDS;
26+
import static com.bytechef.component.definition.ComponentDsl.array;
27+
import static com.bytechef.component.definition.ComponentDsl.integer;
28+
import static com.bytechef.component.definition.ComponentDsl.option;
2529
import static com.bytechef.platform.component.definition.VectorStoreComponentDefinition.VECTOR_STORE;
2630

2731
import com.bytechef.automation.knowledgebase.domain.KnowledgeBase;
@@ -31,17 +35,22 @@
3135
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseDocumentChunkService;
3236
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseDocumentService;
3337
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseService;
38+
import com.bytechef.automation.knowledgebase.service.KnowledgeBaseTagService;
3439
import com.bytechef.component.ai.vectorstore.VectorStore;
3540
import com.bytechef.component.definition.ClusterElementDefinition;
3641
import com.bytechef.component.definition.ComponentDsl;
42+
import com.bytechef.component.definition.Option;
3743
import com.bytechef.component.definition.Parameters;
3844
import com.bytechef.file.storage.domain.FileEntry;
3945
import com.bytechef.platform.component.definition.ParametersFactory;
4046
import com.bytechef.platform.component.definition.ai.agent.VectorStoreFunction;
41-
import com.bytechef.platform.component.service.ClusterElementDefinitionService;
47+
import com.bytechef.platform.tag.domain.Tag;
48+
import java.util.ArrayList;
4249
import java.util.HashMap;
50+
import java.util.LinkedHashMap;
4351
import java.util.List;
4452
import java.util.Map;
53+
import java.util.Objects;
4554
import org.springframework.ai.document.Document;
4655
import org.springframework.ai.document.DocumentReader;
4756
import org.springframework.ai.document.DocumentTransformer;
@@ -57,12 +66,12 @@ public final class KnowledgeBaseVectorStore {
5766
private KnowledgeBaseVectorStore() {
5867
}
5968

69+
@SuppressWarnings("PMD.UnusedFormalParameter")
6070
public static ClusterElementDefinition<VectorStoreFunction> of(
6171
org.springframework.ai.vectorstore.VectorStore vectorStore,
62-
ClusterElementDefinitionService clusterElementDefinitionService,
6372
KnowledgeBaseDocumentChunkService knowledgeBaseDocumentChunkService,
6473
KnowledgeBaseDocumentService knowledgeBaseDocumentService, KnowledgeBaseFileStorage knowledgeBaseFileStorage,
65-
KnowledgeBaseService knowledgeBaseService) {
74+
KnowledgeBaseService knowledgeBaseService, KnowledgeBaseTagService knowledgeBaseTagService) {
6675

6776
VectorStore kbVectorStore = createVectorStore(
6877
knowledgeBaseDocumentChunkService, knowledgeBaseDocumentService, knowledgeBaseFileStorage,
@@ -72,12 +81,89 @@ public static ClusterElementDefinition<VectorStoreFunction> of(
7281
.title("Knowledge Base VectorStore")
7382
.description("Knowledge Base VectorStore.")
7483
.type(VectorStoreFunction.VECTOR_STORE)
84+
.properties(
85+
integer(KNOWLEDGE_BASE_ID)
86+
.label("Knowledge Base")
87+
.description("The knowledge base to retrieve documents from.")
88+
.options(getKnowledgeBaseOptions(knowledgeBaseService))
89+
.required(true),
90+
array(TAG_IDS)
91+
.label("Tags")
92+
.description(
93+
"Filter results by tags. Documents with ANY of the selected tags will be returned (OR logic).")
94+
.items(integer())
95+
.options(getTagOptions(knowledgeBaseTagService))
96+
.optionsLookupDependsOn(KNOWLEDGE_BASE_ID)
97+
.required(false))
7598
.object(() -> (
7699
inputParameters, connectionParameters, extensions,
77100
componentConnections) -> kbVectorStore.createVectorStore(
78101
inputParameters, ParametersFactory.create(connectionParameters), null));
79102
}
80103

104+
private static ClusterElementDefinition.OptionsFunction<Long> getKnowledgeBaseOptions(
105+
KnowledgeBaseService knowledgeBaseService) {
106+
107+
return (inputParameters, connectionParameters, lookupDependsOnPaths, searchText, context) -> {
108+
List<Option<Long>> options = new ArrayList<>();
109+
110+
List<KnowledgeBase> knowledgeBases = knowledgeBaseService.getKnowledgeBases();
111+
112+
for (KnowledgeBase knowledgeBase : knowledgeBases) {
113+
String name = knowledgeBase.getName();
114+
115+
if (searchText == null || name.toLowerCase()
116+
.contains(searchText.toLowerCase())) {
117+
118+
options.add(option(name, knowledgeBase.getId()
119+
.longValue()));
120+
}
121+
}
122+
123+
return options;
124+
};
125+
}
126+
127+
private static ClusterElementDefinition.OptionsFunction<Long> getTagOptions(
128+
KnowledgeBaseTagService knowledgeBaseTagService) {
129+
130+
return (inputParameters, connectionParameters, lookupDependsOnPaths, searchText, context) -> {
131+
Long knowledgeBaseId = inputParameters.getLong(KNOWLEDGE_BASE_ID);
132+
133+
List<Tag> tags;
134+
135+
if (knowledgeBaseId == null) {
136+
tags = knowledgeBaseTagService.getAllTags();
137+
} else {
138+
tags = knowledgeBaseTagService.getTagsByKnowledgeBaseId()
139+
.getOrDefault(knowledgeBaseId, List.of());
140+
}
141+
142+
Map<Long, Tag> unique = new LinkedHashMap<>();
143+
144+
for (Tag tag : tags) {
145+
Long tagId = Objects.requireNonNull(tag.getId());
146+
147+
unique.putIfAbsent(tagId, tag);
148+
}
149+
150+
List<Option<Long>> options = new ArrayList<>();
151+
152+
for (Tag tag : unique.values()) {
153+
String tagName = tag.getName();
154+
155+
if (searchText == null || tagName.toLowerCase()
156+
.contains(searchText.toLowerCase())) {
157+
158+
options.add(option(tagName, Objects.requireNonNull(tag.getId())
159+
.longValue()));
160+
}
161+
}
162+
163+
return options;
164+
};
165+
}
166+
81167
public static VectorStore createVectorStore(org.springframework.ai.vectorstore.VectorStore vectorStore) {
82168
return new VectorStoreImpl(null, null, null, null, vectorStore);
83169
}
@@ -134,8 +220,9 @@ public org.springframework.ai.vectorstore.VectorStore createVectorStore(
134220
Parameters inputParameters, Parameters connectionParameters, EmbeddingModel embeddingModel) {
135221

136222
Long knowledgeBaseId = inputParameters.getRequiredLong(KNOWLEDGE_BASE_ID);
223+
List<Long> tagIds = inputParameters.getList(TAG_IDS, Long.class);
137224

138-
return new KnowledgeBaseVectorStoreWrapper(vectorStore, knowledgeBaseId);
225+
return new KnowledgeBaseVectorStoreWrapper(vectorStore, knowledgeBaseId, tagIds);
139226
}
140227

141228
@Override
@@ -221,9 +308,10 @@ public List<Document> search(
221308
Parameters inputParameters, Parameters connectionParameters, EmbeddingModel embeddingModel) {
222309

223310
Long knowledgeBaseId = inputParameters.getRequiredLong(KNOWLEDGE_BASE_ID);
311+
List<Long> tagIds = inputParameters.getList(TAG_IDS, Long.class);
224312

225313
org.springframework.ai.vectorstore.VectorStore wrappedVectorStore =
226-
new KnowledgeBaseVectorStoreWrapper(vectorStore, knowledgeBaseId);
314+
new KnowledgeBaseVectorStoreWrapper(vectorStore, knowledgeBaseId, tagIds);
227315

228316
return wrappedVectorStore.similaritySearch(inputParameters.getRequiredString(QUERY));
229317
}

0 commit comments

Comments
 (0)