Skip to content

Commit d0f9d43

Browse files
msukkarimsukkari
authored
feat(ask_sb): Add search context into ask sourcebot toolbar (#397)
* new context selector * ui nits * move search context fetch to server * feedback * search context for chat suggestion, nits * type nit * fix minor ui nit
1 parent e404838 commit d0f9d43

File tree

19 files changed

+522
-278
lines changed

19 files changed

+522
-278
lines changed

‎packages/web/src/actions.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1862,11 +1862,16 @@ export const getSearchContexts = async (domain: string) => sew(() =>
18621862
where: {
18631863
orgId: org.id,
18641864
},
1865+
include: {
1866+
repos: true,
1867+
},
18651868
});
18661869

18671870
return searchContexts.map((context) => ({
1871+
id: context.id,
18681872
name: context.name,
18691873
description: context.description ?? undefined,
1874+
repoNames: context.repos.map((repo) => repo.name),
18701875
}));
18711876
}, /* minRequiredRole = */ OrgRole.GUEST), /* allowAnonymousAccess = */ true
18721877
));

‎packages/web/src/app/[domain]/chat/[id]/components/chatThreadPanel.tsx

Lines changed: 54 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,17 @@
33
import { ResizablePanel } from '@/components/ui/resizable';
44
import { ChatThread } from '@/features/chat/components/chatThread';
55
import { LanguageModelInfo, SBChatMessage, SET_CHAT_STATE_QUERY_PARAM, SetChatStatePayload } from '@/features/chat/types';
6-
import { RepositoryQuery } from '@/lib/types';
6+
import { RepositoryQuery, SearchContextQuery } from '@/lib/types';
77
import { CreateUIMessage } from 'ai';
88
import { useRouter, useSearchParams } from 'next/navigation';
99
import { useEffect, useState } from 'react';
1010
import { useChatId } from '../../useChatId';
11+
import { ContextItem } from '@/features/chat/components/chatBox/contextSelector';
1112

1213
interface ChatThreadPanelProps {
1314
languageModels: LanguageModelInfo[];
1415
repos: RepositoryQuery[];
16+
searchContexts: SearchContextQuery[];
1517
order: number;
1618
messages: SBChatMessage[];
1719
isChatReadonly: boolean;
@@ -20,6 +22,7 @@ interface ChatThreadPanelProps {
2022
export const ChatThreadPanel = ({
2123
languageModels,
2224
repos,
25+
searchContexts,
2326
order,
2427
messages,
2528
isChatReadonly,
@@ -31,8 +34,31 @@ export const ChatThreadPanel = ({
3134
const searchParams = useSearchParams();
3235
const [inputMessage, setInputMessage] = useState<CreateUIMessage<SBChatMessage> | undefined>(undefined);
3336

34-
// Use the last user's last message to determine what repos we should select by default.
35-
const [selectedRepos, setSelectedRepos] = useState<string[]>(messages.findLast((message) => message.role === "user")?.metadata?.selectedRepos ?? []);
37+
// Use the last user's last message to determine what repos and contexts we should select by default.
38+
const lastUserMessage = messages.findLast((message) => message.role === "user");
39+
const defaultSelectedRepos = lastUserMessage?.metadata?.selectedRepos ?? [];
40+
const defaultSelectedContexts = lastUserMessage?.metadata?.selectedContexts ?? [];
41+
42+
const [selectedItems, setSelectedItems] = useState<ContextItem[]>([
43+
...defaultSelectedRepos.map(repoName => {
44+
const repoInfo = repos.find(r => r.repoName === repoName);
45+
return {
46+
type: 'repo' as const,
47+
value: repoName,
48+
name: repoInfo?.repoDisplayName || repoName.split('/').pop() || repoName,
49+
codeHostType: repoInfo?.codeHostType || ''
50+
};
51+
}),
52+
...defaultSelectedContexts.map(contextName => {
53+
const context = searchContexts.find(c => c.name === contextName);
54+
return {
55+
type: 'context' as const,
56+
value: contextName,
57+
name: contextName,
58+
repoCount: context?.repoNames.length || 0
59+
};
60+
})
61+
]);
3662

3763
useEffect(() => {
3864
const setChatState = searchParams.get(SET_CHAT_STATE_QUERY_PARAM);
@@ -41,9 +67,28 @@ export const ChatThreadPanel = ({
4167
}
4268

4369
try {
44-
const { inputMessage, selectedRepos } = JSON.parse(setChatState) as SetChatStatePayload;
70+
const { inputMessage, selectedRepos, selectedContexts } = JSON.parse(setChatState) as SetChatStatePayload;
4571
setInputMessage(inputMessage);
46-
setSelectedRepos(selectedRepos);
72+
setSelectedItems([
73+
...selectedRepos.map(repoName => {
74+
const repoInfo = repos.find(r => r.repoName === repoName);
75+
return {
76+
type: 'repo' as const,
77+
value: repoName,
78+
name: repoInfo?.repoDisplayName || repoName.split('/').pop() || repoName,
79+
codeHostType: repoInfo?.codeHostType || ''
80+
};
81+
}),
82+
...selectedContexts.map(contextName => {
83+
const context = searchContexts.find(c => c.name === contextName);
84+
return {
85+
type: 'context' as const,
86+
value: contextName,
87+
name: contextName,
88+
repoCount: context?.repoNames.length || 0
89+
};
90+
})
91+
]);
4792
} catch {
4893
console.error('Invalid message in URL');
4994
}
@@ -52,7 +97,7 @@ export const ChatThreadPanel = ({
5297
const newSearchParams = new URLSearchParams(searchParams.toString());
5398
newSearchParams.delete(SET_CHAT_STATE_QUERY_PARAM);
5499
router.replace(`?${newSearchParams.toString()}`, { scroll: false });
55-
}, [searchParams, router]);
100+
}, [searchParams, router, repos, searchContexts]);
56101

57102
return (
58103
<ResizablePanel
@@ -67,8 +112,9 @@ export const ChatThreadPanel = ({
67112
inputMessage={inputMessage}
68113
languageModels={languageModels}
69114
repos={repos}
70-
selectedRepos={selectedRepos}
71-
onSelectedReposChange={setSelectedRepos}
115+
searchContexts={searchContexts}
116+
selectedItems={selectedItems}
117+
onSelectedItemsChange={setSelectedItems}
72118
isChatReadonly={isChatReadonly}
73119
/>
74120
</div>

‎packages/web/src/app/[domain]/chat/[id]/page.tsx

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { getRepos } from '@/actions';
1+
import { getRepos, getSearchContexts } from '@/actions';
22
import { getUserChatHistory, getConfiguredLanguageModelsInfo, getChatInfo } from '@/features/chat/actions';
33
import { ServiceErrorException } from '@/lib/serviceError';
44
import { isServiceError } from '@/lib/utils';
@@ -22,6 +22,7 @@ interface PageProps {
2222
export default async function Page({ params }: PageProps) {
2323
const languageModels = await getConfiguredLanguageModelsInfo();
2424
const repos = await getRepos(params.domain);
25+
const searchContexts = await getSearchContexts(params.domain);
2526
const chatInfo = await getChatInfo({ chatId: params.id }, params.domain);
2627
const session = await auth();
2728
const chatHistory = session ? await getUserChatHistory(params.domain) : [];
@@ -34,6 +35,10 @@ export default async function Page({ params }: PageProps) {
3435
throw new ServiceErrorException(repos);
3536
}
3637

38+
if (isServiceError(searchContexts)) {
39+
throw new ServiceErrorException(searchContexts);
40+
}
41+
3742
if (isServiceError(chatInfo)) {
3843
if (chatInfo.statusCode === StatusCodes.NOT_FOUND) {
3944
return notFound();
@@ -74,6 +79,7 @@ export default async function Page({ params }: PageProps) {
7479
<ChatThreadPanel
7580
languageModels={languageModels}
7681
repos={indexedRepos}
82+
searchContexts={searchContexts}
7783
messages={messages}
7884
order={2}
7985
isChatReadonly={isReadonly}

‎packages/web/src/app/[domain]/chat/components/newChatPanel.tsx

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,29 +6,32 @@ import { ChatBoxToolbar } from "@/features/chat/components/chatBox/chatBoxToolba
66
import { CustomSlateEditor } from "@/features/chat/customSlateEditor";
77
import { useCreateNewChatThread } from "@/features/chat/useCreateNewChatThread";
88
import { LanguageModelInfo } from "@/features/chat/types";
9-
import { RepositoryQuery } from "@/lib/types";
9+
import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
1010
import { useCallback, useState } from "react";
1111
import { Descendant } from "slate";
1212
import { useLocalStorage } from "usehooks-ts";
13+
import { ContextItem } from "@/features/chat/components/chatBox/contextSelector";
1314

1415
interface NewChatPanelProps {
1516
languageModels: LanguageModelInfo[];
1617
repos: RepositoryQuery[];
18+
searchContexts: SearchContextQuery[];
1719
order: number;
1820
}
1921

2022
export const NewChatPanel = ({
2123
languageModels,
2224
repos,
25+
searchContexts,
2326
order,
2427
}: NewChatPanelProps) => {
25-
const [selectedRepos, setSelectedRepos] = useLocalStorage<string[]>("selectedRepos", [], { initializeWithValue: false });
28+
const [selectedItems, setSelectedItems] = useLocalStorage<ContextItem[]>("selectedContextItems", [], { initializeWithValue: false });
2629
const { createNewChatThread, isLoading } = useCreateNewChatThread();
27-
const [isRepoSelectorOpen, setIsRepoSelectorOpen] = useState(false);
30+
const [isContextSelectorOpen, setIsContextSelectorOpen] = useState(false);
2831

2932
const onSubmit = useCallback((children: Descendant[]) => {
30-
createNewChatThread(children, selectedRepos);
31-
}, [createNewChatThread, selectedRepos]);
33+
createNewChatThread(children, selectedItems);
34+
}, [createNewChatThread, selectedItems]);
3235

3336

3437
return (
@@ -47,17 +50,19 @@ export const NewChatPanel = ({
4750
preferredSuggestionsBoxPlacement="bottom-start"
4851
isRedirecting={isLoading}
4952
languageModels={languageModels}
50-
selectedRepos={selectedRepos}
51-
onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
53+
selectedItems={selectedItems}
54+
searchContexts={searchContexts}
55+
onContextSelectorOpenChanged={setIsContextSelectorOpen}
5256
/>
5357
<div className="w-full flex flex-row items-center bg-accent rounded-b-md px-2">
5458
<ChatBoxToolbar
5559
languageModels={languageModels}
5660
repos={repos}
57-
selectedRepos={selectedRepos}
58-
onSelectedReposChange={setSelectedRepos}
59-
isRepoSelectorOpen={isRepoSelectorOpen}
60-
onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
61+
searchContexts={searchContexts}
62+
selectedItems={selectedItems}
63+
onSelectedItemsChange={setSelectedItems}
64+
isContextSelectorOpen={isContextSelectorOpen}
65+
onContextSelectorOpenChanged={setIsContextSelectorOpen}
6166
/>
6267
</div>
6368
</CustomSlateEditor>

‎packages/web/src/app/[domain]/chat/page.tsx

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { getRepos } from "@/actions";
1+
import { getRepos, getSearchContexts } from "@/actions";
22
import { getUserChatHistory, getConfiguredLanguageModelsInfo } from "@/features/chat/actions";
33
import { ServiceErrorException } from "@/lib/serviceError";
44
import { isServiceError } from "@/lib/utils";
@@ -18,6 +18,7 @@ interface PageProps {
1818
export default async function Page({ params }: PageProps) {
1919
const languageModels = await getConfiguredLanguageModelsInfo();
2020
const repos = await getRepos(params.domain);
21+
const searchContexts = await getSearchContexts(params.domain);
2122
const session = await auth();
2223
const chatHistory = session ? await getUserChatHistory(params.domain) : [];
2324

@@ -29,6 +30,10 @@ export default async function Page({ params }: PageProps) {
2930
throw new ServiceErrorException(repos);
3031
}
3132

33+
if (isServiceError(searchContexts)) {
34+
throw new ServiceErrorException(searchContexts);
35+
}
36+
3237
const indexedRepos = repos.filter((repo) => repo.indexedAt !== undefined);
3338

3439
return (
@@ -48,6 +53,7 @@ export default async function Page({ params }: PageProps) {
4853
<AnimatedResizableHandle />
4954
<NewChatPanel
5055
languageModels={languageModels}
56+
searchContexts={searchContexts}
5157
repos={indexedRepos}
5258
order={2}
5359
/>

‎packages/web/src/app/[domain]/components/homepage/agenticSearch.tsx

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,15 @@ import { LanguageModelInfo } from "@/features/chat/types";
88
import { useCreateNewChatThread } from "@/features/chat/useCreateNewChatThread";
99
import { resetEditor } from "@/features/chat/utils";
1010
import { useDomain } from "@/hooks/useDomain";
11-
import { RepositoryQuery } from "@/lib/types";
11+
import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
1212
import { getDisplayTime } from "@/lib/utils";
1313
import { BrainIcon, FileIcon, LucideIcon, SearchIcon } from "lucide-react";
1414
import Link from "next/link";
1515
import { ReactNode, useCallback, useEffect, useRef, useState } from "react";
1616
import { ReactEditor, useSlate } from "slate-react";
1717
import { SearchModeSelector, SearchModeSelectorProps } from "./toolbar";
1818
import { useLocalStorage } from "usehooks-ts";
19+
import { ContextItem } from "@/features/chat/components/chatBox/contextSelector";
1920

2021
// @todo: we should probably rename this to a different type since it sort-of clashes
2122
// with the Suggestion system we have built into the chat box.
@@ -109,6 +110,7 @@ interface AgenticSearchProps {
109110
searchModeSelectorProps: SearchModeSelectorProps;
110111
languageModels: LanguageModelInfo[];
111112
repos: RepositoryQuery[];
113+
searchContexts: SearchContextQuery[];
112114
chatHistory: {
113115
id: string;
114116
createdAt: Date;
@@ -120,15 +122,16 @@ export const AgenticSearch = ({
120122
searchModeSelectorProps,
121123
languageModels,
122124
repos,
125+
searchContexts,
123126
chatHistory,
124127
}: AgenticSearchProps) => {
125128
const [selectedSuggestionType, _setSelectedSuggestionType] = useState<SuggestionType | undefined>(undefined);
126129
const { createNewChatThread, isLoading } = useCreateNewChatThread();
127130
const dropdownRef = useRef<HTMLDivElement>(null);
128131
const editor = useSlate();
129-
const [selectedRepos, setSelectedRepos] = useLocalStorage<string[]>("selectedRepos", [], { initializeWithValue: false });
132+
const [selectedItems, setSelectedItems] = useLocalStorage<ContextItem[]>("selectedContextItems", [], { initializeWithValue: false });
130133
const domain = useDomain();
131-
const [isRepoSelectorOpen, setIsRepoSelectorOpen] = useState(false);
134+
const [isContextSelectorOpen, setIsContextSelectorOpen] = useState(false);
132135

133136
const setSelectedSuggestionType = useCallback((type: SuggestionType | undefined) => {
134137
_setSelectedSuggestionType(type);
@@ -158,24 +161,26 @@ export const AgenticSearch = ({
158161
>
159162
<ChatBox
160163
onSubmit={(children) => {
161-
createNewChatThread(children, selectedRepos);
164+
createNewChatThread(children, selectedItems);
162165
}}
163166
className="min-h-[50px]"
164167
isRedirecting={isLoading}
165168
languageModels={languageModels}
166-
selectedRepos={selectedRepos}
167-
onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
169+
selectedItems={selectedItems}
170+
searchContexts={searchContexts}
171+
onContextSelectorOpenChanged={setIsContextSelectorOpen}
168172
/>
169173
<Separator />
170174
<div className="relative">
171175
<div className="w-full flex flex-row items-center bg-accent rounded-b-md px-2">
172176
<ChatBoxToolbar
173177
languageModels={languageModels}
174178
repos={repos}
175-
selectedRepos={selectedRepos}
176-
onSelectedReposChange={setSelectedRepos}
177-
isRepoSelectorOpen={isRepoSelectorOpen}
178-
onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
179+
searchContexts={searchContexts}
180+
selectedItems={selectedItems}
181+
onSelectedItemsChange={setSelectedItems}
182+
isContextSelectorOpen={isContextSelectorOpen}
183+
onContextSelectorOpenChanged={setIsContextSelectorOpen}
179184
/>
180185
<SearchModeSelector
181186
{...searchModeSelectorProps}
@@ -201,7 +206,7 @@ export const AgenticSearch = ({
201206
setSelectedSuggestionType(undefined);
202207

203208
if (openRepoSelector) {
204-
setIsRepoSelectorOpen(true);
209+
setIsContextSelectorOpen(true);
205210
} else {
206211
ReactEditor.focus(editor);
207212
}

‎packages/web/src/app/[domain]/components/homepage/index.tsx

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import { SourcebotLogo } from "@/app/components/sourcebotLogo";
44
import { LanguageModelInfo } from "@/features/chat/types";
5-
import { RepositoryQuery } from "@/lib/types";
5+
import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
66
import { useHotkeys } from "react-hotkeys-hook";
77
import { AgenticSearch } from "./agenticSearch";
88
import { PreciseSearch } from "./preciseSearch";
@@ -13,6 +13,7 @@ import { useCallback, useState } from "react";
1313

1414
interface HomepageProps {
1515
initialRepos: RepositoryQuery[];
16+
searchContexts: SearchContextQuery[];
1617
languageModels: LanguageModelInfo[];
1718
chatHistory: {
1819
id: string;
@@ -25,6 +26,7 @@ interface HomepageProps {
2526

2627
export const Homepage = ({
2728
initialRepos,
29+
searchContexts,
2830
languageModels,
2931
chatHistory,
3032
initialSearchMode,
@@ -82,6 +84,7 @@ export const Homepage = ({
8284
}}
8385
languageModels={languageModels}
8486
repos={initialRepos}
87+
searchContexts={searchContexts}
8588
chatHistory={chatHistory}
8689
/>
8790
</CustomSlateEditor>

0 commit comments

Comments
 (0)