2222import static com .bytechef .automation .knowledgebase .constant .KnowledgeBaseConstants .METADATA_KNOWLEDGE_BASE_ID ;
2323import static com .bytechef .component .ai .vectorstore .knowledgebase .constant .KnowledgeBaseVectorStoreConstants .KNOWLEDGE_BASE_ID ;
2424import 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 ;
2529import static com .bytechef .platform .component .definition .VectorStoreComponentDefinition .VECTOR_STORE ;
2630
2731import com .bytechef .automation .knowledgebase .domain .KnowledgeBase ;
3135import com .bytechef .automation .knowledgebase .service .KnowledgeBaseDocumentChunkService ;
3236import com .bytechef .automation .knowledgebase .service .KnowledgeBaseDocumentService ;
3337import com .bytechef .automation .knowledgebase .service .KnowledgeBaseService ;
38+ import com .bytechef .automation .knowledgebase .service .KnowledgeBaseTagService ;
3439import com .bytechef .component .ai .vectorstore .VectorStore ;
3540import com .bytechef .component .definition .ClusterElementDefinition ;
3641import com .bytechef .component .definition .ComponentDsl ;
42+ import com .bytechef .component .definition .Option ;
3743import com .bytechef .component .definition .Parameters ;
3844import com .bytechef .file .storage .domain .FileEntry ;
3945import com .bytechef .platform .component .definition .ParametersFactory ;
4046import 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 ;
4249import java .util .HashMap ;
50+ import java .util .LinkedHashMap ;
4351import java .util .List ;
4452import java .util .Map ;
53+ import java .util .Objects ;
4554import org .springframework .ai .document .Document ;
4655import org .springframework .ai .document .DocumentReader ;
4756import 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