Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
msukkari merged 8 commits into from
Jul 26, 2025
Merged

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

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
d409738
new context selector
msukkari Jul 26, 2025
9bb7865
ui nits
msukkari Jul 26, 2025
b6e8961
Merge branch 'main' into msukkarieh/ask_context
msukkari Jul 26, 2025
0b80e02
move search context fetch to server
msukkari Jul 26, 2025
f47ffc3
feedback
msukkari Jul 26, 2025
4324bde
search context for chat suggestion, nits
msukkari Jul 26, 2025
1c769c9
type nit
msukkari Jul 26, 2025
d7151f2
fix minor ui nit
msukkari Jul 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions packages/web/src/actions.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
));
Expand Down
41 changes: 34 additions & 7 deletions packages/web/src/app/[domain]/chat/[id]/components/chatThreadPanel.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ 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[];
Expand All @@ -31,8 +32,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);
Expand All @@ -41,9 +57,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');
}
Expand All @@ -52,7 +79,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
Expand All @@ -67,8 +94,8 @@ export const ChatThreadPanel = ({
inputMessage={inputMessage}
languageModels={languageModels}
repos={repos}
selectedRepos={selectedRepos}
onSelectedReposChange={setSelectedRepos}
selectedItems={selectedItems}
onSelectedItemsChange={setSelectedItems}
isChatReadonly={isChatReadonly}
/>
</div>
Expand Down
21 changes: 11 additions & 10 deletions packages/web/src/app/[domain]/chat/components/newChatPanel.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import { RepositoryQuery } 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[];
Expand All @@ -22,13 +23,13 @@ export const NewChatPanel = ({
repos,
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 (
Expand All @@ -47,17 +48,17 @@ 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}
selectedItems={selectedItems}
onSelectedItemsChange={setSelectedItems}
isContextSelectorOpen={isContextSelectorOpen}
onContextSelectorOpenChanged={setIsContextSelectorOpen}
/>
</div>
</CustomSlateEditor>
Expand Down
21 changes: 11 additions & 10 deletions packages/web/src/app/[domain]/components/homepage/agenticSearch.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -126,9 +127,9 @@ export const AgenticSearch = ({
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);
Expand Down Expand Up @@ -158,24 +159,24 @@ 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}
selectedItems={selectedItems}
onSelectedItemsChange={setSelectedItems}
isContextSelectorOpen={isContextSelectorOpen}
onContextSelectorOpenChanged={setIsContextSelectorOpen}
/>
<SearchModeSelector
{...searchModeSelectorProps}
Expand All @@ -201,7 +202,7 @@ export const AgenticSearch = ({
setSelectedSuggestionType(undefined);

if (openRepoSelector) {
setIsRepoSelectorOpen(true);
setIsContextSelectorOpen(true);
} else {
ReactEditor.focus(editor);
}
Expand Down
29 changes: 26 additions & 3 deletions packages/web/src/app/api/(server)/chat/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,11 +64,12 @@ export async function POST(req: Request) {
return serviceErrorResponse(schemaValidationError(parsed.error));
}

const { messages, id, selectedRepos, languageModelId } = parsed.data;
const { messages, id, selectedRepos, selectedContexts, languageModelId } = parsed.data;
const response = await chatHandler({
messages,
id,
selectedRepos,
selectedContexts,
languageModelId,
}, domain);

Expand All @@ -93,10 +94,11 @@ interface ChatHandlerProps {
messages: SBChatMessage[];
id: string;
selectedRepos: string[];
selectedContexts?: string[];
languageModelId: string;
}

const chatHandler = ({ messages, id, selectedRepos, languageModelId }: ChatHandlerProps, domain: string) => 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({
Expand Down Expand Up @@ -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',
Expand Down
23 changes: 12 additions & 11 deletions packages/web/src/features/chat/components/chatBox/chatBox.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import { Suggestion } from "./types";
import { useSuggestionModeAndQuery } from "./useSuggestionModeAndQuery";
import { useSuggestionsData } from "./useSuggestionsData";
import { useToast } from "@/components/hooks/use-toast";
import { ContextItem } from "./contextSelector";

interface ChatBoxProps {
onSubmit: (children: Descendant[], editor: CustomEditor) => void;
Expand All @@ -27,8 +28,8 @@ interface ChatBoxProps {
isRedirecting?: boolean;
isGenerating?: boolean;
languageModels: LanguageModelInfo[];
selectedRepos: string[];
onRepoSelectorOpenChanged: (isOpen: boolean) => void;
selectedItems: ContextItem[];
onContextSelectorOpenChanged: (isOpen: boolean) => void;
}

export const ChatBox = ({
Expand All @@ -39,8 +40,8 @@ export const ChatBox = ({
isRedirecting,
isGenerating,
languageModels,
selectedRepos,
onRepoSelectorOpenChanged,
selectedItems,
onContextSelectorOpenChanged,
}: ChatBoxProps) => {
const suggestionsBoxRef = useRef<HTMLDivElement>(null);
const [index, setIndex] = useState(0);
Expand All @@ -49,7 +50,7 @@ export const ChatBox = ({
const { suggestions, isLoading } = useSuggestionsData({
suggestionMode,
suggestionQuery,
selectedRepos,
selectedRepos: selectedItems.filter(item => item.type === 'repo').map(item => item.value),
});
const { selectedLanguageModel } = useSelectedLanguageModel({
initialLanguageModel: languageModels.length > 0 ? languageModels[0] : undefined,
Expand Down Expand Up @@ -113,7 +114,7 @@ export const ChatBox = ({
}
}

if (selectedRepos.length === 0) {
if (selectedItems.length === 0) {
return {
isSubmitDisabled: true,
isSubmitDisabledReason: "no-repos-selected",
Expand All @@ -137,25 +138,25 @@ export const ChatBox = ({
editor.children,
isRedirecting,
isGenerating,
selectedRepos.length,
selectedItems.length,
selectedLanguageModel,
])

const onSubmit = useCallback(() => {
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) {
Expand Down Expand Up @@ -322,7 +323,7 @@ export const ChatBox = ({
<TooltipContent>
<div className="flex flex-row items-center">
<TriangleAlertIcon className="h-4 w-4 text-warning mr-1" />
<span className="text-destructive">One or more repositories must be selected.</span>
<span className="text-destructive">One or more repositories or search contexts must be selected.</span>
</div>
</TooltipContent>
)}
Expand Down
Loading