Add Support for Different Types of Knowledge Bases

workspaces/ballerina/ballerina-core/src/interfaces/bi.ts

@@ -404,9 +404,9 @@ export type NodeKind =
     | "VARIABLE"
     | "VECTOR_STORE"
     | "VECTOR_STORES"
-    | "VECTOR_KNOWLEDGE_BASE"
-    | "VECTOR_KNOWLEDGE_BASE_CALL"
-    | "VECTOR_KNOWLEDGE_BASES"
+    | "KNOWLEDGE_BASE"
+    | "KNOWLEDGE_BASE_CALL"
+    | "KNOWLEDGE_BASES"
     | "EMBEDDING_PROVIDER"
     | "EMBEDDING_PROVIDERS"
     | "DATA_LOADER"

workspaces/ballerina/ballerina-core/src/interfaces/extended-lang-client.ts

@@ -905,7 +905,7 @@ export type SearchKind =
     | "MODEL_PROVIDER"
     | "VECTOR_STORE"
     | "EMBEDDING_PROVIDER"
-    | "VECTOR_KNOWLEDGE_BASE"
+    | "KNOWLEDGE_BASE"
     | "DATA_LOADER"
     | "CHUNKER"
     | "AGENT"

workspaces/ballerina/ballerina-extension/src/core/extended-language-client.ts

@@ -339,7 +339,7 @@ enum EXTENDED_APIS {
     BI_AVAILABLE_MODEL_PROVIDERS = 'flowDesignService/getAvailableModelProviders',
     BI_AVAILABLE_VECTOR_STORES = 'flowDesignService/getAvailableVectorStores',
     BI_AVAILABLE_EMBEDDING_PROVIDERS = 'flowDesignService/getAvailableEmbeddingProviders',
-    BI_AVAILABLE_VECTOR_KNOWLEDGE_BASES = 'flowDesignService/getAvailableVectorKnowledgeBases',
+    BI_AVAILABLE_KNOWLEDGE_BASES = 'flowDesignService/getAvailableVectorKnowledgeBases',
     BI_AVAILABLE_DATA_LOADERS = 'flowDesignService/getAvailableDataLoaders',
     BI_AVAILABLE_CHUNKS = 'flowDesignService/getAvailableChunkers',
     BI_NODE_TEMPLATE = 'flowDesignService/getNodeTemplate',
@@ -1036,7 +1036,7 @@ export class ExtendedLangClient extends LanguageClient implements ExtendedLangCl
     }
 
     async getAvailableVectorKnowledgeBases(params: BIAvailableNodesRequest): Promise<BIAvailableNodesResponse> {
-        return this.sendRequest<BIAvailableNodesResponse>(EXTENDED_APIS.BI_AVAILABLE_VECTOR_KNOWLEDGE_BASES, params);
+        return this.sendRequest<BIAvailableNodesResponse>(EXTENDED_APIS.BI_AVAILABLE_KNOWLEDGE_BASES, params);
     }
 
     async getAvailableDataLoaders(params: BIAvailableNodesRequest): Promise<BIAvailableNodesResponse> {

workspaces/ballerina/ballerina-extension/src/rpc-managers/bi-diagram/rpc-manager.ts

@@ -493,16 +493,16 @@ export class BiDiagramRpcManager implements BIDiagramAPI {
     }
 
     async getAvailableVectorKnowledgeBases(params: BIAvailableNodesRequest): Promise<BIAvailableNodesResponse> {
-        console.log(">>> requesting bi available vector knowledge bases from ls", params);
+        console.log(">>> requesting bi available knowledge bases from ls", params);
         return new Promise((resolve) => {
             StateMachine.langClient()
                 .getAvailableVectorKnowledgeBases(params)
                 .then((model) => {
-                    console.log(">>> bi available vector knowledge bases from ls", model);
+                    console.log(">>> bi available knowledge bases from ls", model);
                     resolve(model);
                 })
                 .catch((error) => {
-                    console.log(">>> error fetching available vector knowledge bases from ls", error);
+                    console.log(">>> error fetching available knowledge bases from ls", error);
                     return new Promise((resolve) => {
                         resolve(undefined);
                     });

workspaces/ballerina/ballerina-side-panel/src/components/CardList/index.tsx

@@ -302,7 +302,7 @@ function CardList(props: CardListProps) {
         return (
             <S.CardsContainer>
                 {nodes.map((node, index) => (
-                    <S.Card key={node.id + index} enabled={node.enabled} onClick={() => handleCardClick(node)}>
+                    <S.Card key={node.id + index} enabled={node.enabled} onClick={() => handleCardClick(node)} title={node.description}>
                         <S.CardIcon>{node.icon ? node.icon : <LogIcon />}</S.CardIcon>
                         <S.CardContent>
                             <S.CardTitle>{node.label}</S.CardTitle>

workspaces/ballerina/ballerina-side-panel/src/components/NodeList/categoryConfig.ts

@@ -135,8 +135,8 @@ export const CATEGORY_CONFIGS: Record<string, CategoryConfig> = {
         showWhenEmpty: true,
         useConnectionContainer: true
     },
-    "Vector Knowledge Bases": {
-        title: "Vector Knowledge Bases",
+    "Knowledge Bases": {
+        title: "Knowledge Bases",
         actions: [{
             type: 'add',
             tooltip: "",

workspaces/ballerina/ballerina-side-panel/src/components/NodeList/index.tsx

@@ -544,7 +544,7 @@ export function NodeList(props: NodeListProps) {
             "Logging": { hasBackground: true },
             "Model Providers": { hasBackground: false },
             "Embedding Providers": { hasBackground: false },
-            "Vector Knowledge Bases": { hasBackground: false },
+            "Knowledge Bases": { hasBackground: false },
             "Vector Stores": { hasBackground: false },
             "Data Loaders": { hasBackground: false },
             "Chunkers": { hasBackground: false },

workspaces/ballerina/ballerina-side-panel/src/utils/localStorage.ts

@@ -53,7 +53,7 @@ export const getDefaultExpandedState = (categories: string[]): ExpandedCategorie
         "Model Providers",
         "Vector Stores",
         "Embedding Providers",
-        "Vector Knowledge Bases",
+        "Knowledge Bases",
         "Data Loaders",
         "Chunkers",
     ];

workspaces/ballerina/ballerina-visualizer/src/utils/bi.tsx

@@ -183,8 +183,13 @@ export function convertEmbeddingProviderCategoriesToSidePanelCategories(categori
     return convertModelProviderCategoriesToSidePanelCategories(categories);
 }
 
-export function convertVectorKnowledgeBaseCategoriesToSidePanelCategories(categories: Category[]): PanelCategory[] {
-    return convertModelProviderCategoriesToSidePanelCategories(categories);
+export function convertKnowledgeBaseCategoriesToSidePanelCategories(categories: Category[]): PanelCategory[] {
+    return convertCategoriesToSidePanelCategoriesWithIcon(categories, (codedata) => {
+        if ((codedata?.module as string).includes("azure")) {
+            return <AIModelIcon type="ai.azure" />;
+        }
+        return <NodeIcon type={codedata?.node} size={24} />
+    });
 }
 
 export function convertCategoriesToSidePanelCategoriesWithIcon(
@@ -403,6 +408,9 @@ export function getContainerTitle(view: SidePanelView, activeNode: FlowNode, cli
             if (!activeNode) {
                 return "";
             }
+            if (activeNode.codedata?.node === "KNOWLEDGE_BASE" && activeNode.codedata?.object === "VectorKnowledgeBase") {
+                return `ai: Vector Knowledge Base`;
+            }
             if (
                 activeNode.codedata?.node === "REMOTE_ACTION_CALL" ||
                 activeNode.codedata?.node === "RESOURCE_ACTION_CALL"

workspaces/ballerina/ballerina-visualizer/src/views/BI/FlowDiagram/PanelManager.tsx

@@ -63,7 +63,8 @@ export enum SidePanelView {
     DATA_LOADER_LIST = "DATA_LOADER_LIST",
     CHUNKERS = "CHUNKERS",
     CHUNKER_LIST = "CHUNKER_LIST",
-    VECTOR_KNOWLEDGE_BASE_LIST = "VECTOR_KNOWLEDGE_BASE_LIST",
+    KNOWLEDGE_BASES = "KNOWLEDGE_BASES",
+    KNOWLEDGE_BASE_LIST = "KNOWLEDGE_BASE_LIST",
     NEW_AGENT = "NEW_AGENT",
     ADD_TOOL = "ADD_TOOL",
     NEW_TOOL = "NEW_TOOL",
@@ -439,23 +440,35 @@ export function PanelManager(props: PanelManagerProps) {
                     />
                 );
 
-            case SidePanelView.VECTOR_KNOWLEDGE_BASE_LIST:
+            case SidePanelView.KNOWLEDGE_BASE_LIST:
                 return (
                     <NodeList
                         categories={categories}
                         onSelect={onSelectNode}
                         onAdd={onAddVectorKnowledgeBase}
-                        addButtonLabel={"Add Vector Knowledge Base"}
+                        addButtonLabel={"Add Knowledge Base"}
                         onClose={onClose}
-                        title={"Vector Knowledge Bases"}
-                        searchPlaceholder={"Search vector knowledge bases"}
+                        title={"Knowledge Bases"}
+                        searchPlaceholder={"Search knowledge bases"}
                         onSearchTextChange={(searchText) =>
                             onSearchVectorKnowledgeBase?.(searchText, FUNCTION_TYPE.REGULAR)
                         }
                         onBack={canGoBack ? onBack : undefined}
                     />
                 );
 
+            case SidePanelView.KNOWLEDGE_BASES:
+                return (
+                    <CardList
+                        categories={categories}
+                        onSelect={onSelectNode}
+                        onClose={onClose}
+                        title={"Knowledge Bases"}
+                        searchPlaceholder={"Search knowledge bases"}
+                        onBack={canGoBack ? onBack : undefined}
+                    />
+                );
+
             case SidePanelView.DATA_LOADER_LIST:
                 return (
                     <NodeList

workspaces/ballerina/ballerina-visualizer/src/views/BI/FlowDiagram/index.tsx

@@ -56,6 +56,7 @@ import {
     convertEmbeddingProviderCategoriesToSidePanelCategories,
     convertDataLoaderCategoriesToSidePanelCategories,
     convertChunkerCategoriesToSidePanelCategories,
+    convertKnowledgeBaseCategoriesToSidePanelCategories,
 } from "../../../utils/bi";
 import { useDraftNodeManager } from "./hooks/useDraftNodeManager";
 import { NodePosition, STNode } from "@wso2/syntax-tree";
@@ -375,8 +376,8 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
     };
 
     const handleVectorKnowledgeBaseAdded = async () => {
-        // Try to navigate back to VECTOR_KNOWLEDGE_BASE_LIST in the stack
-        const foundInStack = popNavigationStackUntilView(SidePanelView.VECTOR_KNOWLEDGE_BASE_LIST);
+        // Try to navigate back to KNOWLEDGE_BASE_LIST in the stack
+        const foundInStack = popNavigationStackUntilView(SidePanelView.KNOWLEDGE_BASE_LIST);
 
         if (foundInStack) {
             setShowProgressIndicator(true);
@@ -391,15 +392,15 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
                         FUNCTION_TYPE.REGULAR
                     )
                 );
-                setSidePanelView(SidePanelView.VECTOR_KNOWLEDGE_BASE_LIST);
+                setSidePanelView(SidePanelView.KNOWLEDGE_BASE_LIST);
                 setShowSidePanel(true);
             } catch (error) {
-                console.error(">>> Error refreshing vector knowledge bases", error);
+                console.error(">>> Error refreshing knowledge bases", error);
             } finally {
                 setShowProgressIndicator(false);
             }
         } else {
-            console.log(">>> VECTOR_KNOWLEDGE_BASE_LIST not found in navigation stack, closing panel");
+            console.log(">>> KNOWLEDGE_BASE_LIST not found in navigation stack, closing panel");
             closeSidePanelAndFetchUpdatedFlowModel();
         }
     };
@@ -823,8 +824,8 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
                     case "EMBEDDING_PROVIDER":
                         panelView = SidePanelView.EMBEDDING_PROVIDER_LIST;
                         break;
-                    case "VECTOR_KNOWLEDGE_BASE":
-                        panelView = SidePanelView.VECTOR_KNOWLEDGE_BASE_LIST;
+                    case "KNOWLEDGE_BASE":
+                        panelView = SidePanelView.KNOWLEDGE_BASE_LIST;
                         break;
                     case "DATA_LOADER":
                         panelView = SidePanelView.DATA_LOADER_LIST;
@@ -865,7 +866,7 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
     };
 
     const handleSearchVectorKnowledgeBase = async (searchText: string, functionType: FUNCTION_TYPE) => {
-        // await handleSearch(searchText, functionType, "VECTOR_KNOWLEDGE_BASE");
+        // await handleSearch(searchText, functionType, "KNOWLEDGE_BASE");
     };
 
     const handleSearchDataLoader = async (searchText: string, functionType: FUNCTION_TYPE) => {
@@ -1066,7 +1067,7 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
                     });
                 break;
 
-            case "VECTOR_KNOWLEDGE_BASES":
+            case "KNOWLEDGE_BASES":
                 setShowProgressIndicator(true);
                 rpcClient
                     .getBIDiagramRpcClient()
@@ -1081,7 +1082,7 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
                                 FUNCTION_TYPE.REGULAR
                             )
                         );
-                        setSidePanelView(SidePanelView.VECTOR_KNOWLEDGE_BASE_LIST);
+                        setSidePanelView(SidePanelView.KNOWLEDGE_BASE_LIST);
                         setShowSidePanel(true);
                     })
                     .finally(() => {
@@ -1423,6 +1424,14 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
                 );
                 setCategories([]);
                 setSidePanelView(SidePanelView.DATA_LOADER_LIST);
+            } else if (sidePanelView === SidePanelView.KNOWLEDGE_BASES) {
+                handleOnSelectNode(
+                    selectedNodeMetadata.current.nodeId,
+                    selectedNodeMetadata.current.metadata,
+                    selectedNodeMetadata.current.fileName
+                );
+                setCategories([]);
+                setSidePanelView(SidePanelView.KNOWLEDGE_BASE_LIST);
             } else if (sidePanelView === SidePanelView.CHUNKERS) {
                 handleOnSelectNode(
                     selectedNodeMetadata.current.nodeId,
@@ -1433,23 +1442,23 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
                 setSidePanelView(SidePanelView.CHUNKER_LIST);
             } else if (
                 sidePanelView === SidePanelView.FORM &&
-                selectedNodeMetadata.current.metadata.node.codedata.node === "VECTOR_KNOWLEDGE_BASE"
+                selectedNodeMetadata.current.metadata.node.codedata.node === "KNOWLEDGE_BASE"
             ) {
                 handleOnSelectNode(
                     selectedNodeMetadata.current.nodeId,
                     selectedNodeMetadata.current.metadata,
                     selectedNodeMetadata.current.fileName
                 );
                 setCategories([]);
-                setSidePanelView(SidePanelView.VECTOR_KNOWLEDGE_BASE_LIST);
+                setSidePanelView(SidePanelView.KNOWLEDGE_BASE_LIST);
             } else if (
                 sidePanelView === SidePanelView.FUNCTION_LIST ||
                 sidePanelView === SidePanelView.DATA_MAPPER_LIST ||
                 sidePanelView === SidePanelView.NP_FUNCTION_LIST ||
                 sidePanelView === SidePanelView.MODEL_PROVIDER_LIST ||
                 sidePanelView === SidePanelView.VECTOR_STORE_LIST ||
                 sidePanelView === SidePanelView.EMBEDDING_PROVIDER_LIST ||
-                sidePanelView === SidePanelView.VECTOR_KNOWLEDGE_BASE_LIST ||
+                sidePanelView === SidePanelView.KNOWLEDGE_BASE_LIST ||
                 sidePanelView === SidePanelView.DATA_LOADER_LIST ||
                 sidePanelView === SidePanelView.CHUNKER_LIST
             ) {
@@ -1639,22 +1648,18 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
         // Push current state to navigation stack
         pushToNavigationStack(sidePanelView, categories, selectedNodeRef.current, selectedClientName.current);
 
-        // Update the node type to VECTOR_KNOWLEDGE_BASE and get the template
-        const updatedMetadata = { ...selectedNodeMetadata.current.metadata };
-        updatedMetadata.node.codedata.node = "VECTOR_KNOWLEDGE_BASE";
-        selectedNodeMetadata.current.metadata = updatedMetadata;
-
+        // Use search to get available knowledge base types
         rpcClient
             .getBIDiagramRpcClient()
-            .getNodeTemplate({
-                position: targetRef.current.startLine,
+            .search({
+                position: { startLine: targetRef.current.startLine, endLine: targetRef.current.endLine },
                 filePath: model?.fileName,
-                id: updatedMetadata.node.codedata,
+                queryMap: undefined,
+                searchKind: "KNOWLEDGE_BASE",
             })
             .then((response) => {
-                selectedNodeRef.current = response.flowNode;
-                showEditForm.current = false;
-                setSidePanelView(SidePanelView.FORM);
+                setCategories(convertKnowledgeBaseCategoriesToSidePanelCategories(response.categories as Category[]));
+                setSidePanelView(SidePanelView.KNOWLEDGE_BASES);
                 setShowSidePanel(true);
             })
             .finally(() => {
@@ -2030,7 +2035,7 @@ export function BIFlowDiagram(props: BIFlowDiagramProps) {
     };
 
     const updateNodeWithConnection = async (selectedNode: FlowNode) => {
-        if (selectedNode.codedata.node === "VECTOR_KNOWLEDGE_BASE") {
+        if (selectedNode.codedata.node === "KNOWLEDGE_BASE") {
             setSidePanelView(SidePanelView.FORM);
             return;
         }

workspaces/ballerina/ballerina-visualizer/src/views/BI/Forms/FormGenerator/index.tsx

@@ -90,7 +90,7 @@ import MatchForm from "../MatchForm";
 import { FormSubmitOptions } from "../../FlowDiagram";
 import { getHelperPaneNew } from "../../HelperPaneNew";
 import { VariableForm } from "../DeclareVariableForm";
-import VectorKnowledgeBaseForm from "../VectorKnowledgeBaseForm";
+import KnowledgeBaseForm from "../KnowledgeBaseForm";
 import { EditorContext, StackItem, TypeHelperItem } from "@wso2/type-editor";
 import DynamicModal from "../../../../components/Modal";
 import React from "react";
@@ -1255,9 +1255,9 @@ export const FormGenerator = forwardRef<FormExpressionEditorRef, FormProps>(func
     }
 
     // handle vector knowledge base form
-    if (node?.codedata.node === "VECTOR_KNOWLEDGE_BASE") {
+    if (node?.codedata.node === "KNOWLEDGE_BASE") {
         return (
-            <VectorKnowledgeBaseForm
+            <KnowledgeBaseForm
                 fileName={fileName}
                 node={node}
                 targetLineRange={targetLineRange}