From d0f9d4362452087ce80bca31be5fc0bda7727c75 Mon Sep 17 00:00:00 2001 From: Michael Sukkarieh Date: Sat, 26 Jul 2025 16:16:07 -0700 Subject: [PATCH] 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 --- packages/web/src/actions.ts | 5 + .../chat/[id]/components/chatThreadPanel.tsx | 62 +++- .../web/src/app/[domain]/chat/[id]/page.tsx | 8 +- .../[domain]/chat/components/newChatPanel.tsx | 27 +- packages/web/src/app/[domain]/chat/page.tsx | 8 +- .../components/homepage/agenticSearch.tsx | 27 +- .../[domain]/components/homepage/index.tsx | 5 +- packages/web/src/app/[domain]/page.tsx | 8 +- .../web/src/app/api/(server)/chat/route.ts | 29 +- .../chat/components/chatBox/chatBox.tsx | 39 ++- .../components/chatBox/chatBoxToolbar.tsx | 39 +-- .../components/chatBox/contextSelector.tsx | 279 ++++++++++++++++++ .../chat/components/chatBox/repoSelector.tsx | 191 ------------ .../chat/components/chatThread/chatThread.tsx | 44 ++- packages/web/src/features/chat/types.ts | 3 + .../features/chat/useCreateNewChatThread.ts | 11 +- packages/web/src/features/chat/utils.ts | 3 +- packages/web/src/lib/schemas.ts | 7 + packages/web/src/lib/types.ts | 5 +- 19 files changed, 522 insertions(+), 278 deletions(-) create mode 100644 packages/web/src/features/chat/components/chatBox/contextSelector.tsx delete mode 100644 packages/web/src/features/chat/components/chatBox/repoSelector.tsx diff --git a/packages/web/src/actions.ts b/packages/web/src/actions.ts index a78c9e14..430c003b 100644 --- a/packages/web/src/actions.ts +++ b/packages/web/src/actions.ts @@ -1862,11 +1862,16 @@ export const getSearchContexts = async (domain: string) => sew(() => where: { orgId: org.id, }, + include: { + repos: true, + }, }); return searchContexts.map((context) => ({ + id: context.id, name: context.name, description: context.description ?? undefined, + repoNames: context.repos.map((repo) => repo.name), })); }, /* minRequiredRole = */ OrgRole.GUEST), /* allowAnonymousAccess = */ true )); diff --git a/packages/web/src/app/[domain]/chat/[id]/components/chatThreadPanel.tsx b/packages/web/src/app/[domain]/chat/[id]/components/chatThreadPanel.tsx index 38d12781..48824d76 100644 --- a/packages/web/src/app/[domain]/chat/[id]/components/chatThreadPanel.tsx +++ b/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,31 @@ export const ChatThreadPanel = ({ const searchParams = useSearchParams(); const [inputMessage, setInputMessage] = useState | undefined>(undefined); - // Use the last user's last message to determine what repos we should select by default. - const [selectedRepos, setSelectedRepos] = useState(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([ + ...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(contextName => { + const context = searchContexts.find(c => c.name === contextName); + return { + type: 'context' as const, + value: contextName, + name: contextName, + repoCount: context?.repoNames.length || 0 + }; + }) + ]); useEffect(() => { const setChatState = searchParams.get(SET_CHAT_STATE_QUERY_PARAM); @@ -41,9 +67,28 @@ 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(contextName => { + const context = searchContexts.find(c => c.name === contextName); + return { + type: 'context' as const, + value: contextName, + name: contextName, + repoCount: context?.repoNames.length || 0 + }; + }) + ]); } catch { console.error('Invalid message in URL'); } @@ -52,7 +97,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, searchContexts]); return ( diff --git a/packages/web/src/app/[domain]/chat/[id]/page.tsx b/packages/web/src/app/[domain]/chat/[id]/page.tsx index 67331eb4..ab7592ba 100644 --- a/packages/web/src/app/[domain]/chat/[id]/page.tsx +++ b/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) { { - const [selectedRepos, setSelectedRepos] = useLocalStorage("selectedRepos", [], { initializeWithValue: false }); + const [selectedItems, setSelectedItems] = useLocalStorage("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,19 @@ export const NewChatPanel = ({ preferredSuggestionsBoxPlacement="bottom-start" isRedirecting={isLoading} languageModels={languageModels} - selectedRepos={selectedRepos} - onRepoSelectorOpenChanged={setIsRepoSelectorOpen} + selectedItems={selectedItems} + searchContexts={searchContexts} + onContextSelectorOpenChanged={setIsContextSelectorOpen} />
diff --git a/packages/web/src/app/[domain]/chat/page.tsx b/packages/web/src/app/[domain]/chat/page.tsx index f15a11cb..4b78cff8 100644 --- a/packages/web/src/app/[domain]/chat/page.tsx +++ b/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) { diff --git a/packages/web/src/app/[domain]/components/homepage/agenticSearch.tsx b/packages/web/src/app/[domain]/components/homepage/agenticSearch.tsx index 957b6c5f..cbd6999b 100644 --- a/packages/web/src/app/[domain]/components/homepage/agenticSearch.tsx +++ b/packages/web/src/app/[domain]/components/homepage/agenticSearch.tsx @@ -8,7 +8,7 @@ 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"; @@ -16,6 +16,7 @@ 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(undefined); const { createNewChatThread, isLoading } = useCreateNewChatThread(); const dropdownRef = useRef(null); const editor = useSlate(); - const [selectedRepos, setSelectedRepos] = useLocalStorage("selectedRepos", [], { initializeWithValue: false }); + const [selectedItems, setSelectedItems] = useLocalStorage("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,13 +161,14 @@ export const AgenticSearch = ({ > { - createNewChatThread(children, selectedRepos); + createNewChatThread(children, selectedItems); }} className="min-h-[50px]" isRedirecting={isLoading} languageModels={languageModels} - selectedRepos={selectedRepos} - onRepoSelectorOpenChanged={setIsRepoSelectorOpen} + selectedItems={selectedItems} + searchContexts={searchContexts} + onContextSelectorOpenChanged={setIsContextSelectorOpen} />
@@ -172,10 +176,11 @@ export const AgenticSearch = ({ diff --git a/packages/web/src/app/[domain]/page.tsx b/packages/web/src/app/[domain]/page.tsx index 607bc143..a4455f34 100644 --- a/packages/web/src/app/[domain]/page.tsx +++ b/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 sew(async () => +const chatHandler = ({ messages, id, selectedRepos, selectedContexts, languageModelId }: ChatHandlerProps, domain: string) => sew(async () => withAuth((userId) => withOrgMembership(userId, domain, async ({ org }) => { const chat = await prisma.chat.findUnique({ @@ -186,13 +188,34 @@ const chatHandler = ({ messages, id, selectedRepos, languageModelId }: ChatHandl const startTime = new Date(); + // Expand search contexts to repos + let expandedRepos = [...selectedRepos]; + if (selectedContexts && selectedContexts.length > 0) { + const searchContexts = await prisma.searchContext.findMany({ + where: { + orgId: org.id, + name: { in: selectedContexts } + }, + include: { + repos: true + } + }); + + const contextRepos = searchContexts.flatMap(context => + context.repos.map(repo => repo.name) + ); + + // Combine and deduplicate repos + expandedRepos = Array.from(new Set([...selectedRepos, ...contextRepos])); + } + const researchStream = await createAgentStream({ model, providerOptions, headers, inputMessages: messageHistory, inputSources: sources, - selectedRepos, + selectedRepos: expandedRepos, onWriteSource: (source) => { writer.write({ type: 'data-source', diff --git a/packages/web/src/features/chat/components/chatBox/chatBox.tsx b/packages/web/src/features/chat/components/chatBox/chatBox.tsx index 47c94f6a..35bb8301 100644 --- a/packages/web/src/features/chat/components/chatBox/chatBox.tsx +++ b/packages/web/src/features/chat/components/chatBox/chatBox.tsx @@ -18,6 +18,8 @@ import { Suggestion } from "./types"; import { useSuggestionModeAndQuery } from "./useSuggestionModeAndQuery"; import { useSuggestionsData } from "./useSuggestionsData"; import { useToast } from "@/components/hooks/use-toast"; +import { ContextItem } from "./contextSelector"; +import { SearchContextQuery } from "@/lib/types"; interface ChatBoxProps { onSubmit: (children: Descendant[], editor: CustomEditor) => void; @@ -27,8 +29,9 @@ interface ChatBoxProps { isRedirecting?: boolean; isGenerating?: boolean; languageModels: LanguageModelInfo[]; - selectedRepos: string[]; - onRepoSelectorOpenChanged: (isOpen: boolean) => void; + selectedItems: ContextItem[]; + searchContexts: SearchContextQuery[]; + onContextSelectorOpenChanged: (isOpen: boolean) => void; } export const ChatBox = ({ @@ -39,8 +42,9 @@ export const ChatBox = ({ isRedirecting, isGenerating, languageModels, - selectedRepos, - onRepoSelectorOpenChanged, + selectedItems, + searchContexts, + onContextSelectorOpenChanged, }: ChatBoxProps) => { const suggestionsBoxRef = useRef(null); const [index, setIndex] = useState(0); @@ -49,7 +53,20 @@ export const ChatBox = ({ const { suggestions, isLoading } = useSuggestionsData({ suggestionMode, suggestionQuery, - selectedRepos, + selectedRepos: selectedItems.map((item) => { + if (item.type === 'repo') { + return [item.value]; + } + + if (item.type === 'context') { + const context = searchContexts.find((context) => context.name === item.value); + if (context) { + return context.repoNames; + } + } + + return []; + }).flat(), }); const { selectedLanguageModel } = useSelectedLanguageModel({ initialLanguageModel: languageModels.length > 0 ? languageModels[0] : undefined, @@ -113,7 +130,7 @@ export const ChatBox = ({ } } - if (selectedRepos.length === 0) { + if (selectedItems.length === 0) { return { isSubmitDisabled: true, isSubmitDisabledReason: "no-repos-selected", @@ -137,7 +154,7 @@ export const ChatBox = ({ editor.children, isRedirecting, isGenerating, - selectedRepos.length, + selectedItems.length, selectedLanguageModel, ]) @@ -145,17 +162,17 @@ export const ChatBox = ({ if (isSubmitDisabled) { if (isSubmitDisabledReason === "no-repos-selected") { toast({ - description: "⚠️ One or more repositories must be selected.", + description: "⚠️ One or more repositories or search contexts must be selected.", variant: "destructive", }); - onRepoSelectorOpenChanged(true); + onContextSelectorOpenChanged(true); } return; } _onSubmit(editor.children, editor); - }, [_onSubmit, editor, isSubmitDisabled, isSubmitDisabledReason, toast, onRepoSelectorOpenChanged]); + }, [_onSubmit, editor, isSubmitDisabled, isSubmitDisabledReason, toast, onContextSelectorOpenChanged]); const onInsertSuggestion = useCallback((suggestion: Suggestion) => { switch (suggestion.type) { @@ -322,7 +339,7 @@ export const ChatBox = ({
- One or more repositories must be selected. + One or more repositories or search contexts must be selected.
)} diff --git a/packages/web/src/features/chat/components/chatBox/chatBoxToolbar.tsx b/packages/web/src/features/chat/components/chatBox/chatBoxToolbar.tsx index 0e83a19c..8744e9ac 100644 --- a/packages/web/src/features/chat/components/chatBox/chatBoxToolbar.tsx +++ b/packages/web/src/features/chat/components/chatBox/chatBoxToolbar.tsx @@ -5,34 +5,36 @@ import { Button } from "@/components/ui/button"; import { Separator } from "@/components/ui/separator"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { LanguageModelInfo } from "@/features/chat/types"; -import { RepositoryQuery } from "@/lib/types"; +import { RepositoryQuery, SearchContextQuery } from "@/lib/types"; import { AtSignIcon } from "lucide-react"; import { useCallback } from "react"; import { useHotkeys } from "react-hotkeys-hook"; import { ReactEditor, useSlate } from "slate-react"; import { useSelectedLanguageModel } from "../../useSelectedLanguageModel"; import { LanguageModelSelector } from "./languageModelSelector"; -import { RepoSelector } from "./repoSelector"; +import { ContextSelector, type ContextItem } from "./contextSelector"; export interface ChatBoxToolbarProps { languageModels: LanguageModelInfo[]; repos: RepositoryQuery[]; - selectedRepos: string[]; - onSelectedReposChange: (repos: string[]) => void; - isRepoSelectorOpen: boolean; - onRepoSelectorOpenChanged: (isOpen: boolean) => void; + searchContexts: SearchContextQuery[]; + selectedItems: ContextItem[]; + onSelectedItemsChange: (items: ContextItem[]) => void; + isContextSelectorOpen: boolean; + onContextSelectorOpenChanged: (isOpen: boolean) => void; } export const ChatBoxToolbar = ({ languageModels, repos, - selectedRepos, - onSelectedReposChange, - isRepoSelectorOpen, - onRepoSelectorOpenChanged, + searchContexts, + selectedItems, + onSelectedItemsChange, + isContextSelectorOpen, + onContextSelectorOpenChanged, }: ChatBoxToolbarProps) => { const editor = useSlate(); - + const onAddContext = useCallback(() => { editor.insertText("@"); ReactEditor.focus(editor); @@ -76,17 +78,18 @@ export const ChatBoxToolbar = ({ - repo.repoName)} - selectedRepos={selectedRepos} - onSelectedReposChange={onSelectedReposChange} - isOpen={isRepoSelectorOpen} - onOpenChanged={onRepoSelectorOpenChanged} + repos={repos} + searchContexts={searchContexts} + selectedItems={selectedItems} + onSelectedItemsChange={onSelectedItemsChange} + isOpen={isContextSelectorOpen} + onOpenChanged={onContextSelectorOpenChanged} /> - Repositories to scope conversation to. + Search contexts and repositories to scope conversation to. {languageModels.length > 0 && ( diff --git a/packages/web/src/features/chat/components/chatBox/contextSelector.tsx b/packages/web/src/features/chat/components/chatBox/contextSelector.tsx new file mode 100644 index 00000000..e2ebe18c --- /dev/null +++ b/packages/web/src/features/chat/components/chatBox/contextSelector.tsx @@ -0,0 +1,279 @@ +// Adapted from: web/src/components/ui/multi-select.tsx + +import * as React from "react"; +import { + CheckIcon, + ChevronDown, + FolderIcon, + LayersIcon, + LibraryBigIcon, +} from "lucide-react"; +import Image from "next/image"; + +import { cn, getCodeHostIcon } from "@/lib/utils"; +import { RepositoryQuery, SearchContextQuery } from "@/lib/types"; +import { Button } from "@/components/ui/button"; +import { Badge } from "@/components/ui/badge"; +import { + Popover, + PopoverContent, + PopoverTrigger, +} from "@/components/ui/popover"; +import { + Command, + CommandEmpty, + CommandGroup, + CommandInput, + CommandItem, + CommandList, + CommandSeparator, +} from "@/components/ui/command"; + +export type RepoContextItem = { + type: 'repo'; + value: string; + name: string; + codeHostType: string; +} + +export type SearchContextItem = { + type: 'context'; + value: string; + name: string; + repoCount: number; +} + +export type ContextItem = RepoContextItem | SearchContextItem; + +interface ContextSelectorProps extends React.ButtonHTMLAttributes { + repos: RepositoryQuery[]; + searchContexts: SearchContextQuery[]; + selectedItems: ContextItem[]; + onSelectedItemsChange: (items: ContextItem[]) => void; + className?: string; + isOpen: boolean; + onOpenChanged: (isOpen: boolean) => void; +} + +export const ContextSelector = React.forwardRef< + HTMLButtonElement, + ContextSelectorProps +>( + ( + { + repos, + searchContexts, + onSelectedItemsChange, + className, + selectedItems, + isOpen, + onOpenChanged, + ...props + }, + ref + ) => { + const scrollContainerRef = React.useRef(null); + const scrollPosition = React.useRef(0); + + const handleInputKeyDown = ( + event: React.KeyboardEvent + ) => { + if (event.key === "Enter") { + onOpenChanged(true); + } else if (event.key === "Backspace" && !event.currentTarget.value) { + const newSelectedItems = [...selectedItems]; + newSelectedItems.pop(); + onSelectedItemsChange(newSelectedItems); + } + }; + + const toggleItem = (item: ContextItem) => { + // Store current scroll position before state update + if (scrollContainerRef.current) { + scrollPosition.current = scrollContainerRef.current.scrollTop; + } + + const isSelected = selectedItems.some( + (selected) => selected.type === item.type && selected.value === item.value + ); + + const newSelectedItems = isSelected + ? selectedItems.filter( + (selected) => !(selected.type === item.type && selected.value === item.value) + ) + : [...selectedItems, item]; + onSelectedItemsChange(newSelectedItems); + }; + + const handleClear = () => { + onSelectedItemsChange([]); + }; + + const handleTogglePopover = () => { + onOpenChanged(!isOpen); + }; + + const allItems = React.useMemo(() => { + const contextItems: ContextItem[] = searchContexts.map(context => ({ + type: 'context' as const, + value: context.name, + name: context.name, + repoCount: context.repoNames.length + })); + + const repoItems: ContextItem[] = repos.map(repo => ({ + type: 'repo' as const, + value: repo.repoName, + name: repo.repoDisplayName || repo.repoName.split('/').pop() || repo.repoName, + codeHostType: repo.codeHostType, + })); + + return [...contextItems, ...repoItems]; + }, [repos, searchContexts]); + + const sortedItems = React.useMemo(() => { + return allItems + .map((item) => ({ + item, + isSelected: selectedItems.some( + (selected) => selected.type === item.type && selected.value === item.value + ) + })) + .sort((a, b) => { + // Selected items first + if (a.isSelected && !b.isSelected) return -1; + if (!a.isSelected && b.isSelected) return 1; + // Then contexts before repos + if (a.item.type === 'context' && b.item.type === 'repo') return -1; + if (a.item.type === 'repo' && b.item.type === 'context') return 1; + return 0; + }) + }, [allItems, selectedItems]); + + // Restore scroll position after re-render + React.useEffect(() => { + if (scrollContainerRef.current && scrollPosition.current > 0) { + scrollContainerRef.current.scrollTop = scrollPosition.current; + } + }, [sortedItems]); + + return ( + + + + + onOpenChanged(false)} + > + + + + No results found. + + {sortedItems.map(({ item, isSelected }) => { + return ( + toggleItem(item)} + className="cursor-pointer" + > +
+ +
+
+ {item.type === 'context' ? ( + + ) : ( + // Render code host icon for repos + (() => { + const codeHostIcon = item.codeHostType ? getCodeHostIcon(item.codeHostType) : null; + return codeHostIcon ? ( + {`${item.codeHostType} + ) : ( + + ); + })() + )} +
+
+ + {item.name} + + {item.type === 'context' && ( + + {item.repoCount} repo{item.repoCount === 1 ? '' : 's'} + + )} +
+
+
+ + ); + })} + + + {selectedItems.length > 0 && ( + <> + + + Clear + + + )} + + + + ); + } +); + +ContextSelector.displayName = "ContextSelector"; \ No newline at end of file diff --git a/packages/web/src/features/chat/components/chatBox/repoSelector.tsx b/packages/web/src/features/chat/components/chatBox/repoSelector.tsx deleted file mode 100644 index 8b2d7b13..00000000 --- a/packages/web/src/features/chat/components/chatBox/repoSelector.tsx +++ /dev/null @@ -1,191 +0,0 @@ -// Adapted from: web/src/components/ui/multi-select.tsx - -import * as React from "react"; -import { - CheckIcon, - ChevronDown, - BookMarkedIcon, -} from "lucide-react"; - -import { cn } from "@/lib/utils"; -import { Button } from "@/components/ui/button"; -import { - Popover, - PopoverContent, - PopoverTrigger, -} from "@/components/ui/popover"; -import { - Command, - CommandEmpty, - CommandGroup, - CommandInput, - CommandItem, - CommandList, - CommandSeparator, -} from "@/components/ui/command"; - -interface RepoSelectorProps extends React.ButtonHTMLAttributes { - repos: string[]; - selectedRepos: string[]; - onSelectedReposChange: (repos: string[]) => void; - className?: string; - isOpen: boolean; - onOpenChanged: (isOpen: boolean) => void; -} - -export const RepoSelector = React.forwardRef< - HTMLButtonElement, - RepoSelectorProps ->( - ( - { - repos, - onSelectedReposChange, - className, - selectedRepos, - isOpen, - onOpenChanged, - ...props - }, - ref - ) => { - const scrollContainerRef = React.useRef(null); - const scrollPosition = React.useRef(0); - - const handleInputKeyDown = ( - event: React.KeyboardEvent - ) => { - if (event.key === "Enter") { - onOpenChanged(true); - } else if (event.key === "Backspace" && !event.currentTarget.value) { - const newSelectedRepos = [...selectedRepos]; - newSelectedRepos.pop(); - onSelectedReposChange(newSelectedRepos); - } - }; - - const toggleRepo = (repo: string) => { - // Store current scroll position before state update - if (scrollContainerRef.current) { - scrollPosition.current = scrollContainerRef.current.scrollTop; - } - - const newSelectedValues = selectedRepos.includes(repo) - ? selectedRepos.filter((value) => value !== repo) - : [...selectedRepos, repo]; - onSelectedReposChange(newSelectedValues); - }; - - const handleClear = () => { - onSelectedReposChange([]); - }; - - const handleTogglePopover = () => { - onOpenChanged(!isOpen); - }; - - const sortedRepos = React.useMemo(() => { - return repos - .map((repo) => ({ - repo, - isSelected: selectedRepos.includes(repo) - })) - .sort((a, b) => { - if (a.isSelected && !b.isSelected) return -1; - if (!a.isSelected && b.isSelected) return 1; - return 0; - }) - }, [repos, selectedRepos]); - - // Restore scroll position after re-render - React.useEffect(() => { - if (scrollContainerRef.current && scrollPosition.current > 0) { - scrollContainerRef.current.scrollTop = scrollPosition.current; - } - }, [sortedRepos]); - - return ( - - - - - onOpenChanged(false)} - > - - - - No results found. - - {sortedRepos.map(({ repo, isSelected }) => { - return ( - toggleRepo(repo)} - className="cursor-pointer" - > -
- -
- {repo} - - ); - })} - - - {selectedRepos.length > 0 && ( - <> - - - Clear - - - )} - - - - ); - } -); - -RepoSelector.displayName = "RepoSelector"; \ No newline at end of file diff --git a/packages/web/src/features/chat/components/chatThread/chatThread.tsx b/packages/web/src/features/chat/components/chatThread/chatThread.tsx index 1e4cd2d5..acbb2023 100644 --- a/packages/web/src/features/chat/components/chatThread/chatThread.tsx +++ b/packages/web/src/features/chat/components/chatThread/chatThread.tsx @@ -12,7 +12,7 @@ import { useChat } from '@ai-sdk/react'; import { CreateUIMessage, DefaultChatTransport } from 'ai'; import { ArrowDownIcon } from 'lucide-react'; import { useNavigationGuard } from 'next-navigation-guard'; -import { Fragment, useCallback, useEffect, useRef, useState } from 'react'; +import { Fragment, useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { Descendant } from 'slate'; import { useMessagePairs } from '../../useMessagePairs'; import { useSelectedLanguageModel } from '../../useSelectedLanguageModel'; @@ -22,7 +22,8 @@ import { ChatThreadListItem } from './chatThreadListItem'; import { ErrorBanner } from './errorBanner'; import { useRouter } from 'next/navigation'; import { usePrevious } from '@uidotdev/usehooks'; -import { RepositoryQuery } from '@/lib/types'; +import { RepositoryQuery, SearchContextQuery } from '@/lib/types'; +import { ContextItem } from '../chatBox/contextSelector'; type ChatHistoryState = { scrollOffset?: number; @@ -34,8 +35,9 @@ interface ChatThreadProps { inputMessage?: CreateUIMessage; languageModels: LanguageModelInfo[]; repos: RepositoryQuery[]; - selectedRepos: string[]; - onSelectedReposChange: (repos: string[]) => void; + searchContexts: SearchContextQuery[]; + selectedItems: ContextItem[]; + onSelectedItemsChange: (items: ContextItem[]) => void; isChatReadonly: boolean; } @@ -45,8 +47,9 @@ export const ChatThread = ({ inputMessage, languageModels, repos, - selectedRepos, - onSelectedReposChange, + searchContexts, + selectedItems, + onSelectedItemsChange, isChatReadonly, }: ChatThreadProps) => { const domain = useDomain(); @@ -57,7 +60,13 @@ export const ChatThread = ({ const [isAutoScrollEnabled, setIsAutoScrollEnabled] = useState(false); const { toast } = useToast(); const router = useRouter(); - const [isRepoSelectorOpen, setIsRepoSelectorOpen] = useState(false); + const [isContextSelectorOpen, setIsContextSelectorOpen] = useState(false); + + const { selectedRepos, selectedContexts } = useMemo(() => { + const repos = selectedItems.filter(item => item.type === 'repo').map(item => item.value); + const contexts = selectedItems.filter(item => item.type === 'context').map(item => item.value); + return { selectedRepos: repos, selectedContexts: contexts }; + }, [selectedItems]); // Initial state is from attachments that exist in in the chat history. const [sources, setSources] = useState( @@ -114,10 +123,11 @@ export const ChatThread = ({ _sendMessage(message, { body: { selectedRepos, + selectedContexts, languageModelId: selectedLanguageModel.model, } satisfies AdditionalChatRequestParams, }); - }, [_sendMessage, selectedLanguageModel, selectedRepos, toast]); + }, [_sendMessage, selectedLanguageModel, toast, selectedRepos, selectedContexts]); const messagePairs = useMessagePairs(messages); @@ -233,13 +243,13 @@ export const ChatThread = ({ const text = slateContentToString(children); const mentions = getAllMentionElements(children); - const message = createUIMessage(text, mentions.map(({ data }) => data), selectedRepos); + const message = createUIMessage(text, mentions.map(({ data }) => data), selectedRepos, selectedContexts); sendMessage(message); setIsAutoScrollEnabled(true); resetEditor(editor); - }, [sendMessage, selectedRepos]); + }, [sendMessage, selectedRepos, selectedContexts]); return ( <> @@ -317,17 +327,19 @@ export const ChatThread = ({ isGenerating={status === "streaming" || status === "submitted"} onStop={stop} languageModels={languageModels} - selectedRepos={selectedRepos} - onRepoSelectorOpenChanged={setIsRepoSelectorOpen} + selectedItems={selectedItems} + searchContexts={searchContexts} + onContextSelectorOpenChanged={setIsContextSelectorOpen} />
diff --git a/packages/web/src/features/chat/types.ts b/packages/web/src/features/chat/types.ts index 2bbea798..b4c6bc8d 100644 --- a/packages/web/src/features/chat/types.ts +++ b/packages/web/src/features/chat/types.ts @@ -51,6 +51,7 @@ export const sbChatMessageMetadataSchema = z.object({ userId: z.string(), })).optional(), selectedRepos: z.array(z.string()).optional(), + selectedContexts: z.array(z.string()).optional(), traceId: z.string().optional(), }); @@ -139,6 +140,7 @@ export const SET_CHAT_STATE_QUERY_PARAM = 'setChatState'; export type SetChatStatePayload = { inputMessage: CreateUIMessage; selectedRepos: string[]; + selectedContexts: string[]; } @@ -156,5 +158,6 @@ export type LanguageModelInfo = { export const additionalChatRequestParamsSchema = z.object({ languageModelId: z.string(), selectedRepos: z.array(z.string()), + selectedContexts: z.array(z.string()), }); export type AdditionalChatRequestParams = z.infer; \ No newline at end of file diff --git a/packages/web/src/features/chat/useCreateNewChatThread.ts b/packages/web/src/features/chat/useCreateNewChatThread.ts index 155be337..54aaf14d 100644 --- a/packages/web/src/features/chat/useCreateNewChatThread.ts +++ b/packages/web/src/features/chat/useCreateNewChatThread.ts @@ -11,6 +11,7 @@ import { createChat } from "./actions"; import { isServiceError } from "@/lib/utils"; import { createPathWithQueryParams } from "@/lib/utils"; import { SET_CHAT_STATE_QUERY_PARAM, SetChatStatePayload } from "./types"; +import { ContextItem } from "./components/chatBox/contextSelector"; export const useCreateNewChatThread = () => { const domain = useDomain(); @@ -18,10 +19,15 @@ export const useCreateNewChatThread = () => { const { toast } = useToast(); const router = useRouter(); - const createNewChatThread = useCallback(async (children: Descendant[], selectedRepos: string[]) => { + const createNewChatThread = useCallback(async (children: Descendant[], selectedItems: ContextItem[]) => { const text = slateContentToString(children); const mentions = getAllMentionElements(children); - const inputMessage = createUIMessage(text, mentions.map((mention) => mention.data), selectedRepos); + + // Extract repos and contexts from selectedItems + const selectedRepos = selectedItems.filter(item => item.type === 'repo').map(item => item.value); + const selectedContexts = selectedItems.filter(item => item.type === 'context').map(item => item.value); + + const inputMessage = createUIMessage(text, mentions.map((mention) => mention.data), selectedRepos, selectedContexts); setIsLoading(true); const response = await createChat(domain); @@ -37,6 +43,7 @@ export const useCreateNewChatThread = () => { [SET_CHAT_STATE_QUERY_PARAM, JSON.stringify({ inputMessage, selectedRepos, + selectedContexts, } satisfies SetChatStatePayload)], ); diff --git a/packages/web/src/features/chat/utils.ts b/packages/web/src/features/chat/utils.ts index f9651115..e2495664 100644 --- a/packages/web/src/features/chat/utils.ts +++ b/packages/web/src/features/chat/utils.ts @@ -172,7 +172,7 @@ export const addLineNumbers = (source: string, lineOffset = 1) => { return source.split('\n').map((line, index) => `${index + lineOffset}:${line}`).join('\n'); } -export const createUIMessage = (text: string, mentions: MentionData[], selectedRepos: string[]): CreateUIMessage => { +export const createUIMessage = (text: string, mentions: MentionData[], selectedRepos: string[], selectedContexts: string[]): CreateUIMessage => { // Converts applicable mentions into sources. const sources: Source[] = mentions .map((mention) => { @@ -206,6 +206,7 @@ export const createUIMessage = (text: string, mentions: MentionData[], selectedR ], metadata: { selectedRepos, + selectedContexts, }, } } diff --git a/packages/web/src/lib/schemas.ts b/packages/web/src/lib/schemas.ts index 09bdae70..6c66e0fc 100644 --- a/packages/web/src/lib/schemas.ts +++ b/packages/web/src/lib/schemas.ts @@ -28,6 +28,13 @@ export const repositoryQuerySchema = z.object({ repoIndexingStatus: z.nativeEnum(RepoIndexingStatus), }); +export const searchContextQuerySchema = z.object({ + id: z.number(), + name: z.string(), + description: z.string().optional(), + repoNames: z.array(z.string()), +}); + export const verifyCredentialsRequestSchema = z.object({ email: z.string().email(), password: z.string().min(8), diff --git a/packages/web/src/lib/types.ts b/packages/web/src/lib/types.ts index 57946242..043b27df 100644 --- a/packages/web/src/lib/types.ts +++ b/packages/web/src/lib/types.ts @@ -1,5 +1,5 @@ import { z } from "zod"; -import { getVersionResponseSchema, repositoryQuerySchema } from "./schemas"; +import { getVersionResponseSchema, repositoryQuerySchema, searchContextQuerySchema } from "./schemas"; import { tenancyModeSchema } from "@/env.mjs"; export type KeymapType = "default" | "vim"; @@ -25,4 +25,5 @@ export type NewsItem = { } export type TenancyMode = z.infer; -export type RepositoryQuery = z.infer; \ No newline at end of file +export type RepositoryQuery = z.infer; +export type SearchContextQuery = z.infer; \ No newline at end of file