diff --git a/e2e/tests/bootstrap.ts b/e2e/tests/bootstrap.ts index 74b944c8..223c0087 100644 --- a/e2e/tests/bootstrap.ts +++ b/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' }); diff --git a/e2e/tests/chat-engine.spec.ts b/e2e/tests/chat-engine.spec.ts index 63dac789..5ad0ead7 100644 --- a/e2e/tests/chat-engine.spec.ts +++ b/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 () => { diff --git a/e2e/utils/forms.ts b/e2e/utils/forms.ts index 0a3f9695..ee8d269d 100644 --- a/e2e/utils/forms.ts +++ b/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); }); } diff --git a/frontend/app/src/api/chat-engines.ts b/frontend/app/src/api/chat-engines.ts index 70b03fe8..b17439b3 100644 --- a/frontend/app/src/api/chat-engines.ts +++ b/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; +}).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; const chatEngineSchema = z.object({ id: z.number(), diff --git a/frontend/app/src/components/chat-engine/create-chat-engine-form.tsx b/frontend/app/src/components/chat-engine/create-chat-engine-form.tsx index eda9c715..27f57f3e 100644 --- a/frontend/app/src/components/chat-engine/create-chat-engine-form.tsx +++ b/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(); 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
- - + + diff --git a/frontend/app/src/components/chat-engine/kb-list-select.tsx b/frontend/app/src/components/chat-engine/kb-list-select.tsx new file mode 100644 index 00000000..be13c6b8 --- /dev/null +++ b/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) { + 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 ( + +
+ + {isLoading + ? Loading options... + : !!error + ? {getErrorMessage(error)} + : !!current?.length + ? current.map((option, index) => ( + option ? ( +
+ {option.name} +
+ + {(option.documents_total ?? 0) || <> no} documents + + + + {(option.data_sources_total ?? 0) || <> no} data sources + +
+
+ ) : UNKNOWN KNOWLEDGE BASE {value?.[index]} + )) : Select Knowledge Bases + } + +
+ + + + + + {knowledgeBases?.map(option => ( + { + 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]); + } + } + }} + > +
+
+ + {option.name} + +
+
+ + {(option.documents_total ?? 0) || <> no} documents + + + + {(option.data_sources_total ?? 0) || <> no} data sources + +
+
+ {option.description} +
+
+ + + ))} + + + Empty List + + + + + + ); +} + +export function KBListSelectForObjectValue ({ value, onChange, ...props }: FormControlWidgetProps<{ id: number }[], true>) { + return ( + v.id) ?? []} + onChange={value => { + onChange?.((value as number[]).map(id => ({ id }))); + }} + {...props} + /> + ); +} diff --git a/frontend/app/src/components/chat-engine/update-chat-engine-form.tsx b/frontend/app/src/components/chat-engine/update-chat-engine-form.tsx index b53135a9..b36bd09e 100644 --- a/frontend/app/src/components/chat-engine/update-chat-engine-form.tsx +++ b/frontend/app/src/components/chat-engine/update-chat-engine-form.tsx @@ -1,12 +1,13 @@ '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'; @@ -14,7 +15,7 @@ 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 }:
- - + + @@ -299,24 +300,26 @@ const llmIdAccessor = getIdAccessor('llm_id'); const fastLlmIdAccessor = getIdAccessor('fast_llm_id'); const rerankerIdAccessor = getIdAccessor('reranker_id'); -const kbAccessor: GeneralSettingsFieldAccessor = { +const kbAccessor: GeneralSettingsFieldAccessor = { 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(); diff --git a/frontend/app/src/components/chat/knowledge-graph-debug-info.tsx b/frontend/app/src/components/chat/knowledge-graph-debug-info.tsx index d6fb82e6..3582f781 100644 --- a/frontend/app/src/components/chat/knowledge-graph-debug-info.tsx +++ b/frontend/app/src/components/chat/knowledge-graph-debug-info.tsx @@ -5,8 +5,6 @@ import type { OngoingState } from '@/components/chat/chat-message-controller'; import { AppChatStreamState, type StackVMState } from '@/components/chat/chat-stream-state'; import { NetworkViewer } from '@/components/graph/components/NetworkViewer'; import { useNetwork } from '@/components/graph/useNetwork'; -import { PencilIcon } from 'lucide-react'; -import Link from 'next/link'; import { useEffect } from 'react'; import useSWR from 'swr'; @@ -14,8 +12,8 @@ export function KnowledgeGraphDebugInfo ({ group }: { group: ChatMessageGroup }) const { engine_options } = useChatInfo(useCurrentChatController()) ?? {}; const auth = useAuth(); const ongoing = useChatMessageStreamState(group.assistant); - const kbId = engine_options?.knowledge_base?.linked_knowledge_base?.id; - const canEdit = !!auth.me?.is_superuser && kbId != null; + const kbLinked = !!engine_options?.knowledge_base?.linked_knowledge_bases?.length; + const canEdit = !!auth.me?.is_superuser && kbLinked; const shouldFetch = (!ongoing || ongoing.finished || couldFetchKnowledgeGraphDebugInfo(ongoing)); const { data: span, isLoading, mutate, error } = useSWR( @@ -42,13 +40,17 @@ export function KnowledgeGraphDebugInfo ({ group }: { group: ChatMessageGroup }) loading={!shouldFetch || isLoading} loadingTitle={shouldFetch ? 'Loading knowledge graph...' : 'Waiting knowledge graph request...'} network={network} - Details={() => canEdit - ? ( - - - Edit graph - - ) : null} + Details={ + () => null + //// TODO: Don't know which KB subgraph to edit for now. + // () => canEdit + // ? ( + // + // + // Edit graph + // + // ) : null + } /> ); } diff --git a/frontend/app/src/components/form/control-widget.tsx b/frontend/app/src/components/form/control-widget.tsx index 6964ea45..b1192ae4 100644 --- a/frontend/app/src/components/form/control-widget.tsx +++ b/frontend/app/src/components/form/control-widget.tsx @@ -165,7 +165,7 @@ export function FormCombobox> ({ r FormCombobox.displayName = 'FormCombobox'; -function FormComboboxClearButton ({ onClick }: { onClick?: () => void }) { +export function FormComboboxClearButton ({ onClick }: { onClick?: () => void }) { return (