Skip to content

Commit

Permalink
feat(frontend): refactor KB selection to support multiple linked know…
Browse files Browse the repository at this point in the history
…ledge bases (#625)

part of #449
  • Loading branch information
634750802 authored Feb 18, 2025
1 parent 22abfc6 commit 4e74fb5
Show file tree
Hide file tree
Showing 9 changed files with 193 additions and 40 deletions.
3 changes: 2 additions & 1 deletion e2e/tests/bootstrap.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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' });
Expand Down
18 changes: 11 additions & 7 deletions e2e/tests/chat-engine.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand All @@ -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();
Expand Down Expand Up @@ -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/);
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 () => {
Expand Down
5 changes: 4 additions & 1 deletion e2e/utils/forms.ts
Original file line number Diff line number Diff line change
@@ -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);
});
}
Expand Down
26 changes: 23 additions & 3 deletions frontend/app/src/api/chat-engines.ts
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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 {
Expand All @@ -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({
Expand Down Expand Up @@ -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(),
Expand Down
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -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(),
Expand All @@ -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 }) {
Expand Down Expand Up @@ -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 />
Expand Down
119 changes: 119 additions & 0 deletions frontend/app/src/components/chat-engine/kb-list-select.tsx
Original file line number Diff line number Diff line change
@@ -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}
/>
);
}
23 changes: 13 additions & 10 deletions frontend/app/src/components/chat-engine/update-chat-engine-form.tsx
Original file line number Diff line number Diff line change
@@ -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();
Expand Down Expand Up @@ -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}>
Expand Down Expand Up @@ -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();
Expand Down
Loading

0 comments on commit 4e74fb5

Please sign in to comment.