Skip to content

Commit 5c0b44b

Browse files
appflowyappflowy
authored
fix(chat): use full workspace context for root AI chats (#365)
1 parent 694b5e4 commit 5c0b44b

9 files changed

Lines changed: 280 additions & 71 deletions

File tree

‎src/components/ai-chat/__tests__/chat-settings.test.ts‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import { View, ViewLayout } from '@/application/types';
2+
import { buildInitialAIChatSettings, isWorkspaceRootView } from '@/components/ai-chat/chat-settings';
3+
4+
function view(overrides: Partial<View>): View {
5+
return {
6+
view_id: 'view-id',
7+
name: 'View',
8+
icon: null,
9+
layout: ViewLayout.Document,
10+
extra: { is_space: false },
11+
children: [],
12+
is_published: false,
13+
is_private: false,
14+
...overrides,
15+
};
16+
}
17+
18+
describe('AI chat initial settings', () => {
19+
it('uses full workspace context for chats created under a workspace root', () => {
20+
const parent = view({
21+
view_id: 'space-id',
22+
extra: { is_space: true },
23+
});
24+
25+
expect(isWorkspaceRootView(parent)).toBe(true);
26+
expect(
27+
buildInitialAIChatSettings({
28+
parent,
29+
query: 'appflowy',
30+
sourceIds: ['doc-1', 'doc-2'],
31+
})
32+
).toEqual({
33+
full_workspace: true,
34+
rag_ids: [],
35+
metadata: { initial_prompt: 'appflowy' },
36+
});
37+
});
38+
39+
it('keeps explicit source context for chats created under a page', () => {
40+
expect(
41+
buildInitialAIChatSettings({
42+
parent: view({ view_id: 'page-id' }),
43+
sourceIds: ['doc-1', 'doc-1', '', 'doc-2'],
44+
})
45+
).toEqual({
46+
full_workspace: false,
47+
rag_ids: ['doc-1', 'doc-2'],
48+
});
49+
});
50+
51+
it('preserves the initial prompt without adding source context for non-root chats', () => {
52+
expect(
53+
buildInitialAIChatSettings({
54+
parent: view({ view_id: 'page-id' }),
55+
query: 'summarize this',
56+
})
57+
).toEqual({
58+
metadata: { initial_prompt: 'summarize this' },
59+
});
60+
});
61+
});

‎src/components/ai-chat/chat-settings.ts‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
import { View } from '@/application/types';
2+
import type { ChatSettings } from '@/components/chat/types';
3+
4+
export function isWorkspaceRootView(view: View | undefined): boolean {
5+
return Boolean(view?.extra?.is_space);
6+
}
7+
8+
export function buildInitialAIChatSettings({
9+
parent,
10+
query,
11+
sourceIds,
12+
}: {
13+
parent?: View;
14+
query?: string;
15+
sourceIds?: string[];
16+
}): Partial<Pick<ChatSettings, 'rag_ids' | 'metadata' | 'full_workspace'>> {
17+
const metadata = query ? { initial_prompt: query } : undefined;
18+
19+
if (isWorkspaceRootView(parent)) {
20+
return {
21+
full_workspace: true,
22+
rag_ids: [],
23+
...(metadata ? { metadata } : {}),
24+
};
25+
}
26+
27+
const uniqueSourceIds = Array.from(new Set(sourceIds || [])).filter(Boolean);
28+
29+
return {
30+
...(uniqueSourceIds.length > 0 ? { full_workspace: false, rag_ids: uniqueSourceIds } : {}),
31+
...(metadata ? { metadata } : {}),
32+
};
33+
}

‎src/components/app/search/BestMatch.tsx‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -298,16 +298,17 @@ function BestMatch({
298298
onClose={onClose}
299299
onLoadMore={handleLoadMore}
300300
header={
301-
<SearchAIOverview
302-
aiEnabled={aiEnabled}
303-
askingAI={askingAI}
304-
loading={loading || summaryLoading}
305-
query={searchValue}
306-
sources={overviewSources}
307-
summary={summary}
308-
onClose={onClose}
309-
onAskAI={(sourceIds) => onAskAI(searchValue, sourceIds)}
310-
/>
301+
aiEnabled ? (
302+
<SearchAIOverview
303+
askingAI={askingAI}
304+
loading={loading || summaryLoading}
305+
query={searchValue}
306+
sources={overviewSources}
307+
summary={summary}
308+
onClose={onClose}
309+
onAskAI={(sourceIds) => onAskAI(searchValue, sourceIds)}
310+
/>
311+
) : undefined
311312
}
312313
/>
313314
);

‎src/components/app/search/Search.tsx‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { ReactComponent as CloseIcon } from '@/assets/icons/close.svg';
77
import { ReactComponent as SearchIcon } from '@/assets/icons/search.svg';
88
import { notify } from '@/components/_shared/notify';
99
import { findAncestors } from '@/components/_shared/outline/utils';
10+
import { buildInitialAIChatSettings } from '@/components/ai-chat/chat-settings';
1011
import {
1112
useAIEnabled,
1213
useAppOperations,
@@ -71,10 +72,10 @@ export function Search() {
7172
name: query || t('chat.newChat', { defaultValue: 'New chat' }),
7273
prev_view_id: parent.children?.[parent.children.length - 1]?.view_id,
7374
});
74-
const uniqueSourceIds = Array.from(new Set(sourceIds || [])).filter(Boolean);
75+
const initialSettings = buildInitialAIChatSettings({ parent, query, sourceIds });
7576
let settingsError: unknown;
7677

77-
if (uniqueSourceIds.length > 0 || query) {
78+
if (Object.keys(initialSettings).length > 0) {
7879
try {
7980
const [{ ChatRequest }, { getAxiosInstance }] = await Promise.all([
8081
import('@/components/chat/request'),
@@ -88,10 +89,7 @@ export function Search() {
8889

8990
const request = new ChatRequest(currentWorkspaceId, created.view_id, axiosInstance);
9091

91-
await request.updateChatSettings({
92-
...(uniqueSourceIds.length > 0 ? { rag_ids: uniqueSourceIds } : {}),
93-
...(query ? { metadata: { initial_prompt: query } } : {}),
94-
});
92+
await request.updateChatSettings(initialSettings);
9593
} catch (error) {
9694
settingsError = error;
9795
}

‎src/components/app/search/SearchAIOverview.tsx‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ export interface SearchOverviewSource {
2121
}
2222

2323
interface SearchAIOverviewProps {
24-
aiEnabled: boolean;
2524
askingAI: boolean;
2625
loading: boolean;
2726
query: string;
@@ -202,7 +201,6 @@ function AskAIButton({ askingAI, query, onAskAI }: { askingAI: boolean; query: s
202201
}
203202

204203
export function SearchAIOverview({
205-
aiEnabled,
206204
askingAI,
207205
loading,
208206
query,
@@ -213,8 +211,6 @@ export function SearchAIOverview({
213211
}: SearchAIOverviewProps) {
214212
const { t } = useTranslation();
215213

216-
if (!aiEnabled) return null;
217-
218214
if (loading) {
219215
return (
220216
<div className='px-2 py-1'>

‎src/components/app/search/__tests__/BestMatch.test.tsx‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
import { describe, expect, it, beforeEach, jest } from '@jest/globals';
2+
import { render, screen } from '@testing-library/react';
3+
4+
import BestMatch from '../BestMatch';
5+
6+
import type { ReactNode } from 'react';
7+
8+
const mockUseAIEnabled = jest.fn();
9+
const mockSearchWorkspaceDocumentPage = jest.fn();
10+
const mockGenerateSearchSummary = jest.fn();
11+
const mockAppOutline: never[] = [];
12+
const mockT = (key: string, options?: { defaultValue?: string }) => options?.defaultValue ?? key;
13+
14+
jest.mock('react-i18next', () => ({
15+
useTranslation: () => ({
16+
t: mockT,
17+
}),
18+
}));
19+
20+
jest.mock('@/components/app/app.hooks', () => ({
21+
useAIEnabled: () => mockUseAIEnabled(),
22+
useAppOutline: () => mockAppOutline,
23+
useCurrentWorkspaceId: () => 'workspace-id',
24+
}));
25+
26+
jest.mock('@/application/services/domains', () => ({
27+
SearchService: {
28+
searchWorkspaceDocumentPage: (...args: unknown[]) => mockSearchWorkspaceDocumentPage(...args),
29+
generateSearchSummary: (...args: unknown[]) => mockGenerateSearchSummary(...args),
30+
},
31+
}));
32+
33+
jest.mock('@/components/app/search/SearchAIOverview', () => ({
34+
SearchAIOverview: () => <div data-testid='ai-overview' />,
35+
}));
36+
37+
jest.mock('@/components/app/search/ViewList', () => ({
38+
__esModule: true,
39+
default: ({ header }: { header?: ReactNode }) => (
40+
<div data-testid='view-list'>{header ? <div data-testid='search-header'>{header}</div> : null}</div>
41+
),
42+
}));
43+
44+
function renderBestMatch() {
45+
return render(<BestMatch askingAI={false} searchValue='' onAskAI={jest.fn()} onClose={jest.fn()} />);
46+
}
47+
48+
describe('BestMatch', () => {
49+
beforeEach(() => {
50+
jest.clearAllMocks();
51+
mockUseAIEnabled.mockReturnValue(true);
52+
mockSearchWorkspaceDocumentPage.mockResolvedValue({
53+
has_more: false,
54+
items: [],
55+
next_offset: null,
56+
});
57+
mockGenerateSearchSummary.mockResolvedValue({ summaries: [] });
58+
});
59+
60+
it('does not mount the AI overview header when server info disables AI', () => {
61+
mockUseAIEnabled.mockReturnValue(false);
62+
63+
renderBestMatch();
64+
65+
expect(screen.queryByTestId('search-header')).toBeNull();
66+
expect(screen.queryByTestId('ai-overview')).toBeNull();
67+
});
68+
69+
it('mounts the AI overview header when AI is enabled', () => {
70+
renderBestMatch();
71+
72+
expect(screen.getByTestId('search-header')).toBeTruthy();
73+
expect(screen.getByTestId('ai-overview')).toBeTruthy();
74+
});
75+
});

‎src/components/app/view-actions/AddPageActions.tsx‎

Lines changed: 55 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,14 @@ import { toast } from 'sonner';
55
import { View, ViewLayout } from '@/application/types';
66
import { ReactComponent as UploadIcon } from '@/assets/icons/upload.svg';
77
import { ViewIcon } from '@/components/_shared/view-icon';
8-
import { useAIEnabled, useAppOperations, useOpenPageModal, useToView } from '@/components/app/app.hooks';
8+
import { buildInitialAIChatSettings } from '@/components/ai-chat/chat-settings';
9+
import {
10+
useAIEnabled,
11+
useAppOperations,
12+
useCurrentWorkspaceId,
13+
useOpenPageModal,
14+
useToView,
15+
} from '@/components/app/app.hooks';
916
import { DropdownMenuGroup, DropdownMenuItem } from '@/components/ui/dropdown-menu';
1017
import { Tooltip, TooltipContent, TooltipTrigger } from '@/components/ui/tooltip';
1118

@@ -15,31 +22,63 @@ function AddPageActions({ view, onImportClick }: { view: View; onImportClick?: (
1522
const openPageModal = useOpenPageModal();
1623
const toView = useToView();
1724
const aiEnabled = useAIEnabled();
25+
const currentWorkspaceId = useCurrentWorkspaceId();
1826
const lastChildViewId = view.children?.[view.children.length - 1]?.view_id;
1927

2028
const handleAddPage = useCallback(
2129
async (layout: ViewLayout, name?: string) => {
2230
if (!addPage) return;
2331
if (layout === ViewLayout.AIChat && !aiEnabled) return;
24-
toast.loading(t('document.creating'));
32+
33+
const loadingToastId = toast.loading(t('document.creating'));
34+
2535
try {
2636
// Append after the last child so the new page appears at the bottom.
2737
// When prev_view_id is omitted the backend prepends (inserts at index 0).
2838
const response = await addPage(view.view_id, { layout, name, prev_view_id: lastChildViewId });
2939

40+
if (layout === ViewLayout.AIChat && currentWorkspaceId) {
41+
const initialSettings = buildInitialAIChatSettings({ parent: view });
42+
43+
if (Object.keys(initialSettings).length > 0) {
44+
try {
45+
const [{ ChatRequest }, { getAxiosInstance }] = await Promise.all([
46+
import('@/components/chat/request'),
47+
import('@/application/services/js-services/http'),
48+
]);
49+
const axiosInstance = getAxiosInstance();
50+
51+
if (!axiosInstance) {
52+
throw new Error('Missing axios instance');
53+
}
54+
55+
const request = new ChatRequest(currentWorkspaceId, response.view_id, axiosInstance);
56+
57+
await request.updateChatSettings(initialSettings);
58+
} catch {
59+
toast.error(
60+
t('search.updateAIChatSettingsFailed', {
61+
defaultValue: 'AI chat was created, but the context could not be attached',
62+
})
63+
);
64+
}
65+
}
66+
}
67+
3068
if (layout === ViewLayout.Document) {
3169
void openPageModal?.(response.view_id);
3270
} else {
3371
void toView(response.view_id);
3472
}
3573

36-
toast.dismiss();
74+
toast.dismiss(loadingToastId);
3775
// eslint-disable-next-line
3876
} catch (e: any) {
77+
toast.dismiss(loadingToastId);
3978
toast.error(e.message);
4079
}
4180
},
42-
[addPage, aiEnabled, openPageModal, t, toView, view.view_id, lastChildViewId]
81+
[addPage, aiEnabled, currentWorkspaceId, openPageModal, t, toView, view, lastChildViewId]
4382
);
4483

4584
const actions: {
@@ -80,16 +119,18 @@ function AddPageActions({ view, onImportClick }: { view: View; onImportClick?: (
80119
void handleAddPage(ViewLayout.Calendar, t('document.plugins.database.newDatabase'));
81120
},
82121
},
83-
...(aiEnabled ? [
84-
{
85-
label: t('chat.newChat'),
86-
icon: <ViewIcon layout={ViewLayout.AIChat} size={'small'} />,
87-
testId: 'add-ai-chat-button',
88-
onSelect: () => {
89-
void handleAddPage(ViewLayout.AIChat, t('menuAppHeader.defaultNewPageName'));
90-
},
91-
},
92-
] : []),
122+
...(aiEnabled
123+
? [
124+
{
125+
label: t('chat.newChat'),
126+
icon: <ViewIcon layout={ViewLayout.AIChat} size={'small'} />,
127+
testId: 'add-ai-chat-button',
128+
onSelect: () => {
129+
void handleAddPage(ViewLayout.AIChat, t('menuAppHeader.defaultNewPageName'));
130+
},
131+
},
132+
]
133+
: []),
93134
{
94135
label: t('chart.menuName'),
95136
icon: <ViewIcon layout={ViewLayout.Chart} size={'small'} />,

0 commit comments

Comments
 (0)