feat(frontend): refactor KB selection to support multiple linked knowledge bases

e2e/tests/bootstrap.ts

@@ -240,8 +240,9 @@ test('Bootstrap', async ({ browser, page }) => {
     await page.getByRole('link', { name: 'default' }).click();
 
     await page.getByRole('tab', { name: 'Retrieval' }).click();
-    await page.getByRole('button', { name: 'Knowledge Base', exact: true }).click();
+    await page.getByRole('button', { name: 'Linked Knowledge Bases', exact: true }).click();
     await page.getByRole('option').filter({ has: page.getByText('My Knowledge Base') }).click();
+    await page.click('body');
 
     await page.getByRole('button', { name: 'Save', exact: true }).click();
     await page.getByRole('button', { name: 'Save', exact: true }).waitFor({ state: 'detached' });

e2e/tests/chat-engine.spec.ts

@@ -22,7 +22,7 @@ test.describe('Chat Engine', () => {
         await page.getByRole('tab', { name: 'Retrieval' }).click();
 
         // Select default knowledge base
-        await selectOption(page, 'Select Knowledge Base', /My Knowledge Base/);
+        await selectOption(page, 'Linked Knowledge Bases', /My Knowledge Base/, true);
       });
 
       const chatEngineId = await test.step('Create', async () => {
@@ -39,10 +39,13 @@ test.describe('Chat Engine', () => {
         expect(chatEngine.name).toBe(name);
         expect(chatEngine.engine_options).toStrictEqual({
           knowledge_base: {
-            linked_knowledge_base: {
+            linked_knowledge_bases: [{
               id: 1,
-            },
+            }],
           },
+          knowledge_graph: {
+            enabled: false,
+          }
         });
         expect(chatEngine.llm_id).toBeNull();
         expect(chatEngine.fast_llm_id).toBeNull();
@@ -77,7 +80,7 @@ test.describe('Chat Engine', () => {
         await page.getByRole('tab', { name: 'Retrieval' }).click();
 
         // Select default knowledge base
-        await selectOption(page, 'Select Knowledge Base', /My Knowledge Base/);
+        await selectOption(page, 'Linked Knowledge Bases', /My Knowledge Base/, true);
 
         // Select Reranker
         await selectOption(page, 'Reranker', /My Reranker/);
@@ -108,11 +111,12 @@ test.describe('Chat Engine', () => {
         expect(chatEngine.name).toBe(name);
         expect(chatEngine.engine_options).toStrictEqual({
           knowledge_base: {
-            linked_knowledge_base: {
+            linked_knowledge_bases: [{
               id: 1,
-            },
+            }],
           },
           knowledge_graph: {
+            enabled: false,
             depth: 1,
             include_meta: true,
             using_intent_search: true,
@@ -147,7 +151,7 @@ test.describe('Chat Engine', () => {
           await page.getByRole('tab', { name: 'Retrieval' }).click();
 
           // Select default knowledge base
-          await selectOption(page, 'Select Knowledge Base', /My Knowledge Base/);
+          await selectOption(page, 'Linked Knowledge Bases', /My Knowledge Base/, true);
         });
 
         const chatEngineId = await test.step('Create', async () => {

e2e/utils/forms.ts

@@ -1,9 +1,12 @@
 import { expect, type Page, test } from '@playwright/test';
 
-export async function selectOption (page: Page, name: string, value: string | RegExp) {
+export async function selectOption (page: Page, name: string, value: string | RegExp, clickWindow = false) {
   await test.step(`Select field ${name}`, async () => {
     await page.getByRole('button', { name: name, exact: true }).click();
     await page.getByRole('option', { name: value }).click();
+    if (clickWindow) {
+      await page.click('body');
+    }
     await expect(page.getByRole('button', { name: name, exact: true })).toHaveText(value);
   });
 }

frontend/app/src/api/chat-engines.ts

@@ -1,6 +1,6 @@
 import { authenticationHeaders, handleErrors, handleResponse, type Page, type PageParams, requestUrl, zodPage } from '@/lib/request';
 import { zodJsonDate } from '@/lib/zod';
-import { number, z, type ZodType } from 'zod';
+import { z, type ZodType } from 'zod';
 
 export interface ChatEngine {
   id: number;
@@ -37,7 +37,11 @@ export interface ChatEngineOptions {
 }
 
 export interface ChatEngineKnowledgeBaseOptions {
+  /**
+   * @deprecated
+   */
   linked_knowledge_base?: LinkedKnowledgeBaseOptions | null;
+  linked_knowledge_bases?: { id: number }[] | null;
 }
 
 export interface ChatEngineKnowledgeGraphOptions {
@@ -60,12 +64,16 @@ export type ChatEngineLLMOptions = {
   further_questions_prompt?: string | null
 }
 
+/**
+ * @deprecated
+ */
 export interface LinkedKnowledgeBaseOptions {
   id?: number | null;
 }
 
 const kbOptionsSchema = z.object({
-  linked_knowledge_base: z.object({ id: number().nullable().optional() }).nullable().optional(),
+  linked_knowledge_base: z.object({ id: z.number().nullable().optional() }).nullable().optional(),
+  linked_knowledge_bases: z.object({ id: z.number() }).array().nullable().optional(),
 }).passthrough();
 
 const kgOptionsSchema = z.object({
@@ -103,7 +111,19 @@ const chatEngineOptionsSchema = z.object({
   post_verification_url: z.string().nullable().optional(),
   post_verification_token: z.string().nullable().optional(),
   hide_sources: z.boolean().nullable().optional(),
-}).passthrough() satisfies ZodType<ChatEngineOptions, any, any>;
+}).passthrough()
+  .refine(option => {
+    if (!option.knowledge_base?.linked_knowledge_bases?.length) {
+      if (option.knowledge_base?.linked_knowledge_base?.id != null) {
+        // Frontend temporary migration. Should be removed after backend removed linked_knowledge_base field.
+        option.knowledge_base.linked_knowledge_bases = [{
+          id: option.knowledge_base.linked_knowledge_base.id,
+        }];
+        delete option.knowledge_base.linked_knowledge_base;
+      }
+    }
+    return option;
+  }) satisfies ZodType<ChatEngineOptions, any, any>;
 
 const chatEngineSchema = z.object({
   id: z.number(),

frontend/app/src/components/chat-engine/create-chat-engine-form.tsx

@@ -1,8 +1,9 @@
 'use client';
 
 import { type ChatEngineOptions, createChatEngine } from '@/api/chat-engines';
+import { KBListSelectForObjectValue } from '@/components/chat-engine/kb-list-select';
 import { FormSection, FormSectionsProvider, useFormSectionFields } from '@/components/form-sections';
-import { KBSelect, LLMSelect, RerankerSelect } from '@/components/form/biz';
+import { LLMSelect, RerankerSelect } from '@/components/form/biz';
 import { FormCheckbox, FormInput, FormSwitch } from '@/components/form/control-widget';
 import { formFieldLayout } from '@/components/form/field-layout';
 import { FormRootError } from '@/components/form/root-error';
@@ -27,9 +28,9 @@ const schema = z.object({
   reranker_id: z.number().optional(),
   engine_options: z.object({
     knowledge_base: z.object({
-      linked_knowledge_base: z.object({
+      linked_knowledge_bases: z.object({
         id: z.number(),
-      }),
+      }).array().min(1),
     }),
     knowledge_graph: z.object({
       depth: z.number().min(1).nullable().optional(),
@@ -41,7 +42,7 @@ const schema = z.object({
 const field = formFieldLayout<typeof schema>();
 
 const nameSchema = z.string().min(1);
-const kbSchema = z.number();
+const kbSchema = z.object({ id: z.number() }).array().min(1);
 const kgGraphDepthSchema = z.number().min(1).optional();
 
 export function CreateChatEngineForm ({ defaultChatEngineOptions }: { defaultChatEngineOptions: ChatEngineOptions }) {
@@ -102,8 +103,8 @@ export function CreateChatEngineForm ({ defaultChatEngineOptions }: { defaultCha
               </SubSection>
             </Section>
             <Section title="Retrieval">
-              <field.Basic required name="engine_options.knowledge_base.linked_knowledge_base.id" label="Select Knowledge Base" validators={{ onChange: kbSchema, onSubmit: kbSchema }}>
-                <KBSelect />
+              <field.Basic required name="engine_options.knowledge_base.linked_knowledge_bases" label="Linked Knowledge Bases" validators={{ onChange: kbSchema, onSubmit: kbSchema }}>
+                <KBListSelectForObjectValue />
               </field.Basic>
               <field.Basic name="reranker_id" label="Reranker">
                 <RerankerSelect />

frontend/app/src/components/chat-engine/kb-list-select.tsx

@@ -0,0 +1,119 @@
+import { type FormControlWidgetProps } from '@/components/form/control-widget';
+import { useAllKnowledgeBases } from '@/components/knowledge-base/hooks';
+import { Command, CommandEmpty, CommandGroup, CommandInput, CommandItem, CommandList } from '@/components/ui/command';
+import { Popover, PopoverContent } from '@/components/ui/popover';
+import { getErrorMessage } from '@/lib/errors';
+import { cn } from '@/lib/utils';
+import * as PopoverPrimitive from '@radix-ui/react-popover';
+import { AlertTriangleIcon, CheckIcon, DotIcon } from 'lucide-react';
+import * as React from 'react';
+import { useState } from 'react';
+
+export function KBListSelect ({ ref, disabled, value, onChange, ...props }: FormControlWidgetProps<number[]>) {
+  const [open, setOpen] = useState(false);
+  const { data: knowledgeBases, isLoading, error } = useAllKnowledgeBases();
+  const isConfigReady = !isLoading && !error;
+
+  const current = value?.map(id => knowledgeBases?.find(kb => kb.id === id));
+
+  return (
+    <Popover open={open} onOpenChange={setOpen}>
+      <div className={cn('flex items-center gap-2')}>
+        <PopoverPrimitive.Trigger
+          ref={ref}
+          disabled={disabled || !isConfigReady}
+          className={cn(
+            'flex flex-col min-h-10 w-full text-left items-stretch justify-start rounded-md border border-input bg-background px-3 py-1 text-sm ring-offset-background placeholder:text-muted-foreground focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2 disabled:cursor-not-allowed disabled:opacity-50',
+          )}
+          {...props}
+        >
+          {isLoading
+            ? <span>Loading options...</span>
+            : !!error
+              ? <span className="text-destructive">{getErrorMessage(error)}</span>
+              : !!current?.length
+                ? current.map((option, index) => (
+                  option ? (
+                    <div key={option.id} className="w-full block border-t first-of-type:border-t-0 py-2">
+                      <span>{option.name}</span>
+                      <div className="text-xs text-muted-foreground ml-2 inline-flex gap-1 items-center">
+                      <span>
+                        {(option.documents_total ?? 0) || <><AlertTriangleIcon className="text-warning inline-flex size-3 mr-0.5" /> no</>} documents
+                      </span>
+                        <DotIcon className="size-4" />
+                        <span className="text-xs text-muted-foreground">
+                        {(option.data_sources_total ?? 0) || <><AlertTriangleIcon className="inline-flex size-3 mr-0.5" /> no</>} data sources
+                      </span>
+                      </div>
+                    </div>
+                  ) : <span key={value?.[index]}>UNKNOWN KNOWLEDGE BASE {value?.[index]}</span>
+                )) : <span className="pt-1 text-muted-foreground">Select Knowledge Bases</span>
+          }
+        </PopoverPrimitive.Trigger>
+      </div>
+      <PopoverContent className={cn('p-0 focus:outline-none w-[--radix-popover-trigger-width]')} align="start" collisionPadding={8}>
+        <Command>
+          <CommandInput />
+          <CommandList>
+            <CommandGroup>
+              {knowledgeBases?.map(option => (
+                <CommandItem
+                  key={option.id}
+                  value={String(option.id)}
+                  keywords={[option.name, option.description ?? '']}
+                  className={cn('group')}
+                  onSelect={idValue => {
+                    const id = knowledgeBases?.find(option => String(option.id) === idValue)?.id;
+                    if (id) {
+                      if (value?.includes(id)) {
+                        onChange?.(value.filter(v => v !== id));
+                      } else {
+                        onChange?.([...(value ?? []), id]);
+                      }
+                    }
+                  }}
+                >
+                  <div className="space-y-1">
+                    <div>
+                      <strong>
+                        {option.name}
+                      </strong>
+                    </div>
+                    <div className="text-xs text-muted-foreground flex gap-1 items-center">
+                      <span>
+                        {(option.documents_total ?? 0) || <><AlertTriangleIcon className="text-warning inline-flex size-3 mr-0.5" /> no</>} documents
+                      </span>
+                      <DotIcon className="size-4" />
+                      <span>
+                        {(option.data_sources_total ?? 0) || <><AlertTriangleIcon className="inline-flex size-3 mr-0.5" /> no</>} data sources
+                      </span>
+                    </div>
+                    <div className="text-xs text-muted-foreground">
+                      {option.description}
+                    </div>
+                  </div>
+                  <CheckIcon className={cn('ml-auto size-4 opacity-0 flex-shrink-0', value?.includes(option.id) && 'opacity-100')} />
+                </CommandItem>
+              ))}
+            </CommandGroup>
+            <CommandEmpty className="text-muted-foreground/50 text-xs p-4 text-center">
+              Empty List
+            </CommandEmpty>
+          </CommandList>
+        </Command>
+      </PopoverContent>
+    </Popover>
+  );
+}
+
+export function KBListSelectForObjectValue ({ value, onChange, ...props }: FormControlWidgetProps<{ id: number }[], true>) {
+  return (
+    <KBListSelect
+      value={value?.map(v => v.id) ?? []}
+      onChange={value => {
+        onChange?.((value as number[]).map(id => ({ id })));
+      }}
+      {...props}
+    />
+  );
+}

frontend/app/src/components/chat-engine/update-chat-engine-form.tsx

@@ -1,20 +1,21 @@
 'use client';
 
 import { type ChatEngine, type ChatEngineKnowledgeGraphOptions, type ChatEngineLLMOptions, type ChatEngineOptions, updateChatEngine } from '@/api/chat-engines';
-import { KBSelect, LLMSelect, RerankerSelect } from '@/components/form/biz';
+import { KBListSelect } from '@/components/chat-engine/kb-list-select';
+import { LLMSelect, RerankerSelect } from '@/components/form/biz';
 import { FormCheckbox, FormInput, FormSwitch } from '@/components/form/control-widget';
 import { formFieldLayout } from '@/components/form/field-layout';
 import { PromptInput } from '@/components/form/widgets/PromptInput';
 import { SecondaryNavigatorItem, SecondaryNavigatorLayout, SecondaryNavigatorList, SecondaryNavigatorMain } from '@/components/secondary-navigator-list';
-import { fieldAccessor, type GeneralSettingsFieldAccessor, GeneralSettingsField as GeneralSettingsField, GeneralSettingsForm, shallowPick } from '@/components/settings-form';
+import { fieldAccessor, GeneralSettingsField as GeneralSettingsField, type GeneralSettingsFieldAccessor, GeneralSettingsForm, shallowPick } from '@/components/settings-form';
 import type { KeyOfType } from '@/lib/typing-utils';
 import { capitalCase } from 'change-case-all';
 import { format } from 'date-fns';
 import { useRouter } from 'next/navigation';
 import { type ReactNode, useTransition } from 'react';
 import { z } from 'zod';
 
-const field = formFieldLayout<{ value: any }>();
+const field = formFieldLayout<{ value: any | any[] }>();
 
 export function UpdateChatEngineForm ({ chatEngine, defaultChatEngineOptions }: { chatEngine: ChatEngine, defaultChatEngineOptions: ChatEngineOptions }) {
   const [transitioning, startTransition] = useTransition();
@@ -106,8 +107,8 @@ export function UpdateChatEngineForm ({ chatEngine, defaultChatEngineOptions }:
 
         <Section title="Retrieval">
           <GeneralSettingsField accessor={kbAccessor} schema={kbSchema}>
-            <field.Basic required name="value" label="Knowledge Base">
-              <KBSelect />
+            <field.Basic required name="value" label="Linked Knowledge Bases">
+              <KBListSelect />
             </field.Basic>
           </GeneralSettingsField>
           <GeneralSettingsField accessor={rerankerIdAccessor} schema={idSchema}>
@@ -299,24 +300,26 @@ const llmIdAccessor = getIdAccessor('llm_id');
 const fastLlmIdAccessor = getIdAccessor('fast_llm_id');
 const rerankerIdAccessor = getIdAccessor('reranker_id');
 
-const kbAccessor: GeneralSettingsFieldAccessor<ChatEngine, number | null> = {
+const kbAccessor: GeneralSettingsFieldAccessor<ChatEngine, number[] | null> = {
   path: ['engine_options'],
   get (data) {
-    return data.engine_options.knowledge_base?.linked_knowledge_base?.id ?? null;
+    console.log(data.engine_options.knowledge_base?.linked_knowledge_bases?.map(kb => kb.id) ?? null);
+    return data.engine_options.knowledge_base?.linked_knowledge_bases?.map(kb => kb.id) ?? null;
   },
-  set (data, id) {
+  set (data, value) {
     return {
       ...data,
       engine_options: {
         ...data.engine_options,
         knowledge_base: {
-          linked_knowledge_base: { id },
+          linked_knowledge_base: undefined,
+          linked_knowledge_bases: value?.map(id => ({ id })) ?? null,
         },
       },
     };
   },
 };
-const kbSchema = z.number();
+const kbSchema = z.number().array().min(1);
 
 const kgEnabledAccessor = kgOptionAccessor('enabled');
 const kgEnabledSchema = z.boolean().nullable();