feat(ask_sb): Add search context into ask sourcebot toolbar

packages/web/src/actions.ts

@@ -1862,11 +1862,15 @@ export const getSearchContexts = async (domain: string) => sew(() =>
                 where: {
                     orgId: org.id,
                 },
+                include: {
+                    repos: true,
+                },
             });
 
             return searchContexts.map((context) => ({
                 name: context.name,
                 description: context.description ?? undefined,
+                repoCount: context.repos.length,
             }));
         }, /* minRequiredRole = */ OrgRole.GUEST), /* allowAnonymousAccess = */ true
     ));

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

@@ -3,15 +3,17 @@
 import { ResizablePanel } from '@/components/ui/resizable';
 import { ChatThread } from '@/features/chat/components/chatThread';
 import { LanguageModelInfo, SBChatMessage, SET_CHAT_STATE_QUERY_PARAM, SetChatStatePayload } from '@/features/chat/types';
-import { RepositoryQuery } from '@/lib/types';
+import { RepositoryQuery, SearchContextQuery } from '@/lib/types';
 import { CreateUIMessage } from 'ai';
 import { useRouter, useSearchParams } from 'next/navigation';
 import { useEffect, useState } from 'react';
 import { useChatId } from '../../useChatId';
+import { ContextItem } from '@/features/chat/components/chatBox/contextSelector';
 
 interface ChatThreadPanelProps {
     languageModels: LanguageModelInfo[];
     repos: RepositoryQuery[];
+    searchContexts: SearchContextQuery[];
     order: number;
     messages: SBChatMessage[];
     isChatReadonly: boolean;
@@ -20,6 +22,7 @@ interface ChatThreadPanelProps {
 export const ChatThreadPanel = ({
     languageModels,
     repos,
+    searchContexts,
     order,
     messages,
     isChatReadonly,
@@ -31,8 +34,23 @@ export const ChatThreadPanel = ({
     const searchParams = useSearchParams();
     const [inputMessage, setInputMessage] = useState<CreateUIMessage<SBChatMessage> | undefined>(undefined);
 
-    // Use the last user's last message to determine what repos we should select by default.
-    const [selectedRepos, setSelectedRepos] = useState<string[]>(messages.findLast((message) => message.role === "user")?.metadata?.selectedRepos ?? []);
+    // Use the last user's last message to determine what repos and contexts we should select by default.
+    const lastUserMessage = messages.findLast((message) => message.role === "user");
+    const defaultSelectedRepos = lastUserMessage?.metadata?.selectedRepos ?? [];
+    const defaultSelectedContexts = lastUserMessage?.metadata?.selectedContexts ?? [];
+
+    const [selectedItems, setSelectedItems] = useState<ContextItem[]>([
+        ...defaultSelectedRepos.map(repoName => {
+            const repoInfo = repos.find(r => r.repoName === repoName);
+            return {
+                type: 'repo' as const,
+                value: repoName,
+                name: repoInfo?.repoDisplayName || repoName.split('/').pop() || repoName,
+                codeHostType: repoInfo?.codeHostType
+            };
+        }),
+        ...defaultSelectedContexts.map(context => ({ type: 'context' as const, value: context, name: context }))
+    ]);
 
     useEffect(() => {
         const setChatState = searchParams.get(SET_CHAT_STATE_QUERY_PARAM);
@@ -41,9 +59,20 @@ export const ChatThreadPanel = ({
         }
 
         try {
-            const { inputMessage, selectedRepos } = JSON.parse(setChatState) as SetChatStatePayload;
+            const { inputMessage, selectedRepos, selectedContexts } = JSON.parse(setChatState) as SetChatStatePayload;
             setInputMessage(inputMessage);
-            setSelectedRepos(selectedRepos);
+            setSelectedItems([
+                ...selectedRepos.map(repoName => {
+                    const repoInfo = repos.find(r => r.repoName === repoName);
+                    return {
+                        type: 'repo' as const,
+                        value: repoName,
+                        name: repoInfo?.repoDisplayName || repoName.split('/').pop() || repoName,
+                        codeHostType: repoInfo?.codeHostType
+                    };
+                }),
+                ...(selectedContexts || []).map(context => ({ type: 'context' as const, value: context, name: context }))
+            ]);
         } catch {
             console.error('Invalid message in URL');
         }
@@ -52,7 +81,7 @@ export const ChatThreadPanel = ({
         const newSearchParams = new URLSearchParams(searchParams.toString());
         newSearchParams.delete(SET_CHAT_STATE_QUERY_PARAM);
         router.replace(`?${newSearchParams.toString()}`, { scroll: false });
-    }, [searchParams, router]);
+    }, [searchParams, router, repos]);
 
     return (
         <ResizablePanel
@@ -67,8 +96,9 @@ export const ChatThreadPanel = ({
                     inputMessage={inputMessage}
                     languageModels={languageModels}
                     repos={repos}
-                    selectedRepos={selectedRepos}
-                    onSelectedReposChange={setSelectedRepos}
+                    searchContexts={searchContexts}
+                    selectedItems={selectedItems}
+                    onSelectedItemsChange={setSelectedItems}
                     isChatReadonly={isChatReadonly}
                 />
             </div>

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

@@ -1,4 +1,4 @@
-import { getRepos } from '@/actions';
+import { getRepos, getSearchContexts } from '@/actions';
 import { getUserChatHistory, getConfiguredLanguageModelsInfo, getChatInfo } from '@/features/chat/actions';
 import { ServiceErrorException } from '@/lib/serviceError';
 import { isServiceError } from '@/lib/utils';
@@ -22,6 +22,7 @@ interface PageProps {
 export default async function Page({ params }: PageProps) {
     const languageModels = await getConfiguredLanguageModelsInfo();
     const repos = await getRepos(params.domain);
+    const searchContexts = await getSearchContexts(params.domain);
     const chatInfo = await getChatInfo({ chatId: params.id }, params.domain);
     const session = await auth();
     const chatHistory = session ? await getUserChatHistory(params.domain) : [];
@@ -34,6 +35,10 @@ export default async function Page({ params }: PageProps) {
         throw new ServiceErrorException(repos);
     }
 
+    if (isServiceError(searchContexts)) {
+        throw new ServiceErrorException(searchContexts);
+    }
+
     if (isServiceError(chatInfo)) {
         if (chatInfo.statusCode === StatusCodes.NOT_FOUND) {
             return notFound();
@@ -74,6 +79,7 @@ export default async function Page({ params }: PageProps) {
                 <ChatThreadPanel
                     languageModels={languageModels}
                     repos={indexedRepos}
+                    searchContexts={searchContexts}
                     messages={messages}
                     order={2}
                     isChatReadonly={isReadonly}

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

@@ -6,29 +6,32 @@ import { ChatBoxToolbar } from "@/features/chat/components/chatBox/chatBoxToolba
 import { CustomSlateEditor } from "@/features/chat/customSlateEditor";
 import { useCreateNewChatThread } from "@/features/chat/useCreateNewChatThread";
 import { LanguageModelInfo } from "@/features/chat/types";
-import { RepositoryQuery } from "@/lib/types";
+import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
 import { useCallback, useState } from "react";
 import { Descendant } from "slate";
 import { useLocalStorage } from "usehooks-ts";
+import { ContextItem } from "@/features/chat/components/chatBox/contextSelector";
 
 interface NewChatPanelProps {
     languageModels: LanguageModelInfo[];
     repos: RepositoryQuery[];
+    searchContexts: SearchContextQuery[];
     order: number;
 }
 
 export const NewChatPanel = ({
     languageModels,
     repos,
+    searchContexts,
     order,
 }: NewChatPanelProps) => {
-    const [selectedRepos, setSelectedRepos] = useLocalStorage<string[]>("selectedRepos", [], { initializeWithValue: false });
+    const [selectedItems, setSelectedItems] = useLocalStorage<ContextItem[]>("selectedContextItems", [], { initializeWithValue: false });
     const { createNewChatThread, isLoading } = useCreateNewChatThread();
-    const [isRepoSelectorOpen, setIsRepoSelectorOpen] = useState(false);
+    const [isContextSelectorOpen, setIsContextSelectorOpen] = useState(false);
 
     const onSubmit = useCallback((children: Descendant[]) => {
-        createNewChatThread(children, selectedRepos);
-    }, [createNewChatThread, selectedRepos]);
+        createNewChatThread(children, selectedItems);
+    }, [createNewChatThread, selectedItems]);
 
 
     return (
@@ -47,17 +50,18 @@ export const NewChatPanel = ({
                             preferredSuggestionsBoxPlacement="bottom-start"
                             isRedirecting={isLoading}
                             languageModels={languageModels}
-                            selectedRepos={selectedRepos}
-                            onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                            selectedItems={selectedItems}
+                            onContextSelectorOpenChanged={setIsContextSelectorOpen}
                         />
                         <div className="w-full flex flex-row items-center bg-accent rounded-b-md px-2">
                             <ChatBoxToolbar
                                 languageModels={languageModels}
                                 repos={repos}
-                                selectedRepos={selectedRepos}
-                                onSelectedReposChange={setSelectedRepos}
-                                isRepoSelectorOpen={isRepoSelectorOpen}
-                                onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                                searchContexts={searchContexts}
+                                selectedItems={selectedItems}
+                                onSelectedItemsChange={setSelectedItems}
+                                isContextSelectorOpen={isContextSelectorOpen}
+                                onContextSelectorOpenChanged={setIsContextSelectorOpen}
                             />
                         </div>
                     </CustomSlateEditor>

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

@@ -1,4 +1,4 @@
-import { getRepos } from "@/actions";
+import { getRepos, getSearchContexts } from "@/actions";
 import { getUserChatHistory, getConfiguredLanguageModelsInfo } from "@/features/chat/actions";
 import { ServiceErrorException } from "@/lib/serviceError";
 import { isServiceError } from "@/lib/utils";
@@ -18,6 +18,7 @@ interface PageProps {
 export default async function Page({ params }: PageProps) {
     const languageModels = await getConfiguredLanguageModelsInfo();
     const repos = await getRepos(params.domain);
+    const searchContexts = await getSearchContexts(params.domain);
     const session = await auth();
     const chatHistory = session ? await getUserChatHistory(params.domain) : [];
 
@@ -29,6 +30,10 @@ export default async function Page({ params }: PageProps) {
         throw new ServiceErrorException(repos);
     }
 
+    if (isServiceError(searchContexts)) {
+        throw new ServiceErrorException(searchContexts);
+    }
+
     const indexedRepos = repos.filter((repo) => repo.indexedAt !== undefined);
 
     return (
@@ -48,6 +53,7 @@ export default async function Page({ params }: PageProps) {
                 <AnimatedResizableHandle />
                 <NewChatPanel
                     languageModels={languageModels}
+                    searchContexts={searchContexts}
                     repos={indexedRepos}
                     order={2}
                 />

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

@@ -8,14 +8,15 @@ import { LanguageModelInfo } from "@/features/chat/types";
 import { useCreateNewChatThread } from "@/features/chat/useCreateNewChatThread";
 import { resetEditor } from "@/features/chat/utils";
 import { useDomain } from "@/hooks/useDomain";
-import { RepositoryQuery } from "@/lib/types";
+import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
 import { getDisplayTime } from "@/lib/utils";
 import { BrainIcon, FileIcon, LucideIcon, SearchIcon } from "lucide-react";
 import Link from "next/link";
 import { ReactNode, useCallback, useEffect, useRef, useState } from "react";
 import { ReactEditor, useSlate } from "slate-react";
 import { SearchModeSelector, SearchModeSelectorProps } from "./toolbar";
 import { useLocalStorage } from "usehooks-ts";
+import { ContextItem } from "@/features/chat/components/chatBox/contextSelector";
 
 // @todo: we should probably rename this to a different type since it sort-of clashes
 // with the Suggestion system we have built into the chat box.
@@ -109,6 +110,7 @@ interface AgenticSearchProps {
     searchModeSelectorProps: SearchModeSelectorProps;
     languageModels: LanguageModelInfo[];
     repos: RepositoryQuery[];
+    searchContexts: SearchContextQuery[];
     chatHistory: {
         id: string;
         createdAt: Date;
@@ -120,15 +122,16 @@ export const AgenticSearch = ({
     searchModeSelectorProps,
     languageModels,
     repos,
+    searchContexts,
     chatHistory,
 }: AgenticSearchProps) => {
     const [selectedSuggestionType, _setSelectedSuggestionType] = useState<SuggestionType | undefined>(undefined);
     const { createNewChatThread, isLoading } = useCreateNewChatThread();
     const dropdownRef = useRef<HTMLDivElement>(null);
     const editor = useSlate();
-    const [selectedRepos, setSelectedRepos] = useLocalStorage<string[]>("selectedRepos", [], { initializeWithValue: false });
+    const [selectedItems, setSelectedItems] = useLocalStorage<ContextItem[]>("selectedContextItems", [], { initializeWithValue: false });
     const domain = useDomain();
-    const [isRepoSelectorOpen, setIsRepoSelectorOpen] = useState(false);
+    const [isContextSelectorOpen, setIsContextSelectorOpen] = useState(false);
 
     const setSelectedSuggestionType = useCallback((type: SuggestionType | undefined) => {
         _setSelectedSuggestionType(type);
@@ -158,24 +161,25 @@ export const AgenticSearch = ({
             >
                 <ChatBox
                     onSubmit={(children) => {
-                        createNewChatThread(children, selectedRepos);
+                        createNewChatThread(children, selectedItems);
                     }}
                     className="min-h-[50px]"
                     isRedirecting={isLoading}
                     languageModels={languageModels}
-                    selectedRepos={selectedRepos}
-                    onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                    selectedItems={selectedItems}
+                    onContextSelectorOpenChanged={setIsContextSelectorOpen}
                 />
                 <Separator />
                 <div className="relative">
                     <div className="w-full flex flex-row items-center bg-accent rounded-b-md px-2">
                         <ChatBoxToolbar
                             languageModels={languageModels}
                             repos={repos}
-                            selectedRepos={selectedRepos}
-                            onSelectedReposChange={setSelectedRepos}
-                            isRepoSelectorOpen={isRepoSelectorOpen}
-                            onRepoSelectorOpenChanged={setIsRepoSelectorOpen}
+                            searchContexts={searchContexts}
+                            selectedItems={selectedItems}
+                            onSelectedItemsChange={setSelectedItems}
+                            isContextSelectorOpen={isContextSelectorOpen}
+                            onContextSelectorOpenChanged={setIsContextSelectorOpen}
                         />
                         <SearchModeSelector
                             {...searchModeSelectorProps}
@@ -201,7 +205,7 @@ export const AgenticSearch = ({
                                         setSelectedSuggestionType(undefined);
 
                                         if (openRepoSelector) {
-                                            setIsRepoSelectorOpen(true);
+                                            setIsContextSelectorOpen(true);
                                         } else {
                                             ReactEditor.focus(editor);
                                         }

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

@@ -2,7 +2,7 @@
 
 import { SourcebotLogo } from "@/app/components/sourcebotLogo";
 import { LanguageModelInfo } from "@/features/chat/types";
-import { RepositoryQuery } from "@/lib/types";
+import { RepositoryQuery, SearchContextQuery } from "@/lib/types";
 import { useHotkeys } from "react-hotkeys-hook";
 import { AgenticSearch } from "./agenticSearch";
 import { PreciseSearch } from "./preciseSearch";
@@ -13,6 +13,7 @@ import { useCallback, useState } from "react";
 
 interface HomepageProps {
     initialRepos: RepositoryQuery[];
+    searchContexts: SearchContextQuery[];
     languageModels: LanguageModelInfo[];
     chatHistory: {
         id: string;
@@ -25,6 +26,7 @@ interface HomepageProps {
 
 export const Homepage = ({
     initialRepos,
+    searchContexts,
     languageModels,
     chatHistory,
     initialSearchMode,
@@ -82,6 +84,7 @@ export const Homepage = ({
                         }}
                         languageModels={languageModels}
                         repos={initialRepos}
+                        searchContexts={searchContexts}
                         chatHistory={chatHistory}
                     />
                 </CustomSlateEditor>

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

@@ -1,4 +1,4 @@
-import { getRepos } from "@/actions";
+import { getRepos, getSearchContexts } from "@/actions";
 import { Footer } from "@/app/components/footer";
 import { getOrgFromDomain } from "@/data/org";
 import { getConfiguredLanguageModelsInfo, getUserChatHistory } from "@/features/chat/actions";
@@ -22,12 +22,17 @@ export default async function Home({ params: { domain } }: { params: { domain: s
 
     const models = await getConfiguredLanguageModelsInfo();
     const repos = await getRepos(domain);
+    const searchContexts = await getSearchContexts(domain);
     const chatHistory = session ? await getUserChatHistory(domain) : [];
 
     if (isServiceError(repos)) {
         throw new ServiceErrorException(repos);
     }
 
+    if (isServiceError(searchContexts)) {
+        throw new ServiceErrorException(searchContexts);
+    }
+
     if (isServiceError(chatHistory)) {
         throw new ServiceErrorException(chatHistory);
     }
@@ -52,6 +57,7 @@ export default async function Home({ params: { domain } }: { params: { domain: s
 
             <Homepage
                 initialRepos={indexedRepos}
+                searchContexts={searchContexts}
                 languageModels={models}
                 chatHistory={chatHistory}
                 initialSearchMode={initialSearchMode}