From 565d93051db2fa9dcd26e648d2e2dc91eadb18ca Mon Sep 17 00:00:00 2001 From: bkellam Date: Tue, 29 Jul 2025 01:11:28 -0700 Subject: [PATCH] Add additional tools for repo searching and listing --- packages/web/src/features/chat/agent.ts | 17 ++--- .../components/chatThread/detailsCard.tsx | 16 +++++ .../chatThread/referencedSourcesListView.tsx | 2 +- .../tools/listAllReposToolComponent.tsx | 66 +++++++++++++++++++ .../tools/searchCodeToolComponent.tsx | 24 +++++-- .../tools/searchReposToolComponent.tsx | 63 ++++++++++++++++++ packages/web/src/features/chat/constants.ts | 4 ++ packages/web/src/features/chat/tools.ts | 65 +++++++++++++++++- packages/web/src/features/chat/types.ts | 4 +- 9 files changed, 243 insertions(+), 18 deletions(-) create mode 100644 packages/web/src/features/chat/components/chatThread/tools/listAllReposToolComponent.tsx create mode 100644 packages/web/src/features/chat/components/chatThread/tools/searchReposToolComponent.tsx diff --git a/packages/web/src/features/chat/agent.ts b/packages/web/src/features/chat/agent.ts index 73bfedf4..414f746f 100644 --- a/packages/web/src/features/chat/agent.ts +++ b/packages/web/src/features/chat/agent.ts @@ -6,7 +6,7 @@ import { ProviderOptions } from "@ai-sdk/provider-utils"; import { createLogger } from "@sourcebot/logger"; import { LanguageModel, ModelMessage, StopCondition, streamText } from "ai"; import { ANSWER_TAG, FILE_REFERENCE_PREFIX, toolNames } from "./constants"; -import { createCodeSearchTool, findSymbolDefinitionsTool, findSymbolReferencesTool, readFilesTool } from "./tools"; +import { createCodeSearchTool, findSymbolDefinitionsTool, findSymbolReferencesTool, readFilesTool, searchReposTool, listAllReposTool } from "./tools"; import { FileSource, Source } from "./types"; import { addLineNumbers, fileReferenceToString } from "./utils"; @@ -54,6 +54,13 @@ export const createAgentStream = async ({ [toolNames.readFiles]: readFilesTool, [toolNames.findSymbolReferences]: findSymbolReferencesTool, [toolNames.findSymbolDefinitions]: findSymbolDefinitionsTool, + // We only include these tools when there are no search scopes + // because the LLM will need to discover what repositories are + // available to it. + ...(searchScopeRepoNames.length === 0 ? { + [toolNames.searchRepos]: searchReposTool, + [toolNames.listAllRepos]: listAllReposTool, + } : {}), }, prepareStep: async ({ stepNumber }) => { // The first step attaches any mentioned sources to the system prompt. @@ -185,13 +192,7 @@ ${searchScopeRepoNames.map(repo => `- ${repo}`).join('\n')} -During the research phase, you have these tools available: -- \`${toolNames.searchCode}\`: Search for code patterns, functions, or text across repositories -- \`${toolNames.readFiles}\`: Read the contents of specific files -- \`${toolNames.findSymbolReferences}\`: Find where symbols are referenced -- \`${toolNames.findSymbolDefinitions}\`: Find where symbols are defined - -Use these tools to gather comprehensive context before answering. Always explain why you're using each tool. +During the research phase, use the tools available to you to gather comprehensive context before answering. Always explain why you're using each tool. Depending on the user's question, you may need to use multiple tools. If the question is vague, ask the user for more information. ${answerInstructions} diff --git a/packages/web/src/features/chat/components/chatThread/detailsCard.tsx b/packages/web/src/features/chat/components/chatThread/detailsCard.tsx index ab3b676c..6fd2fd11 100644 --- a/packages/web/src/features/chat/components/chatThread/detailsCard.tsx +++ b/packages/web/src/features/chat/components/chatThread/detailsCard.tsx @@ -12,6 +12,8 @@ import { FindSymbolDefinitionsToolComponent } from './tools/findSymbolDefinition import { FindSymbolReferencesToolComponent } from './tools/findSymbolReferencesToolComponent'; import { ReadFilesToolComponent } from './tools/readFilesToolComponent'; import { SearchCodeToolComponent } from './tools/searchCodeToolComponent'; +import { SearchReposToolComponent } from './tools/searchReposToolComponent'; +import { ListAllReposToolComponent } from './tools/listAllReposToolComponent'; import { SBChatMessageMetadata, SBChatMessagePart } from '../../types'; import { SearchScopeIcon } from '../searchScopeIcon'; @@ -181,6 +183,20 @@ export const DetailsCard = ({ part={part} /> ) + case 'tool-searchRepos': + return ( + + ) + case 'tool-listAllRepos': + return ( + + ) default: return null; } diff --git a/packages/web/src/features/chat/components/chatThread/referencedSourcesListView.tsx b/packages/web/src/features/chat/components/chatThread/referencedSourcesListView.tsx index d97709bf..61f53734 100644 --- a/packages/web/src/features/chat/components/chatThread/referencedSourcesListView.tsx +++ b/packages/web/src/features/chat/components/chatThread/referencedSourcesListView.tsx @@ -221,7 +221,7 @@ export const ReferencedSourcesListView = ({ {fileName}
- Failed to load file: {isServiceError(query.data) ? query.data.message : 'Unknown error'} + Failed to load file: {isServiceError(query.data) ? query.data.message : query.error?.message ?? 'Unknown error'}
); diff --git a/packages/web/src/features/chat/components/chatThread/tools/listAllReposToolComponent.tsx b/packages/web/src/features/chat/components/chatThread/tools/listAllReposToolComponent.tsx new file mode 100644 index 00000000..6c06146c --- /dev/null +++ b/packages/web/src/features/chat/components/chatThread/tools/listAllReposToolComponent.tsx @@ -0,0 +1,66 @@ +'use client'; + +import { ListAllReposToolUIPart } from "@/features/chat/tools"; +import { isServiceError } from "@/lib/utils"; +import { useMemo, useState } from "react"; +import { ToolHeader, TreeList } from "./shared"; +import { CodeSnippet } from "@/app/components/codeSnippet"; +import { Separator } from "@/components/ui/separator"; +import { FolderOpenIcon } from "lucide-react"; + +export const ListAllReposToolComponent = ({ part }: { part: ListAllReposToolUIPart }) => { + const [isExpanded, setIsExpanded] = useState(false); + + const label = useMemo(() => { + switch (part.state) { + case 'input-streaming': + return 'Loading all repositories...'; + case 'output-error': + return '"List all repositories" tool call failed'; + case 'input-available': + case 'output-available': + return 'Listed all repositories'; + } + }, [part]); + + return ( +
+ + {part.state === 'output-available' && isExpanded && ( + <> + {isServiceError(part.output) ? ( + + Failed with the following error: {part.output.message} + + ) : ( + <> + {part.output.length === 0 ? ( + No repositories found + ) : ( + +
+ Found {part.output.length} repositories: +
+ {part.output.map((repoName, index) => ( +
+ + {repoName} +
+ ))} + + )} + + )} + + + )} +
+ ) +} \ No newline at end of file diff --git a/packages/web/src/features/chat/components/chatThread/tools/searchCodeToolComponent.tsx b/packages/web/src/features/chat/components/chatThread/tools/searchCodeToolComponent.tsx index a7792ade..9131d8cc 100644 --- a/packages/web/src/features/chat/components/chatThread/tools/searchCodeToolComponent.tsx +++ b/packages/web/src/features/chat/components/chatThread/tools/searchCodeToolComponent.tsx @@ -11,24 +11,38 @@ import { SearchIcon } from "lucide-react"; import Link from "next/link"; import { SearchQueryParams } from "@/lib/types"; import { PlayIcon } from "@radix-ui/react-icons"; - +import { buildSearchQuery } from "@/features/chat/utils"; export const SearchCodeToolComponent = ({ part }: { part: SearchCodeToolUIPart }) => { const [isExpanded, setIsExpanded] = useState(false); const domain = useDomain(); + const displayQuery = useMemo(() => { + if (part.state !== 'input-available' && part.state !== 'output-available') { + return ''; + } + + const query = buildSearchQuery({ + query: part.input.queryRegexp, + repoNamesFilterRegexp: part.input.repoNamesFilterRegexp, + languageNamesFilter: part.input.languageNamesFilter, + fileNamesFilterRegexp: part.input.fileNamesFilterRegexp, + }); + + return query; + }, [part]); + const label = useMemo(() => { switch (part.state) { case 'input-streaming': return 'Searching...'; - case 'input-available': - return Searching for {part.input.queryRegexp}; case 'output-error': return '"Search code" tool call failed'; + case 'input-available': case 'output-available': - return Searched for {part.input.queryRegexp}; + return Searched for {displayQuery}; } - }, [part]); + }, [part, displayQuery]); return (
diff --git a/packages/web/src/features/chat/components/chatThread/tools/searchReposToolComponent.tsx b/packages/web/src/features/chat/components/chatThread/tools/searchReposToolComponent.tsx new file mode 100644 index 00000000..218cdba4 --- /dev/null +++ b/packages/web/src/features/chat/components/chatThread/tools/searchReposToolComponent.tsx @@ -0,0 +1,63 @@ +'use client'; + +import { SearchReposToolUIPart } from "@/features/chat/tools"; +import { isServiceError } from "@/lib/utils"; +import { useMemo, useState } from "react"; +import { ToolHeader, TreeList } from "./shared"; +import { CodeSnippet } from "@/app/components/codeSnippet"; +import { Separator } from "@/components/ui/separator"; +import { BookMarkedIcon } from "lucide-react"; + +export const SearchReposToolComponent = ({ part }: { part: SearchReposToolUIPart }) => { + const [isExpanded, setIsExpanded] = useState(false); + + const label = useMemo(() => { + switch (part.state) { + case 'input-streaming': + return 'Searching repositories...'; + case 'output-error': + return '"Search repositories" tool call failed'; + case 'input-available': + case 'output-available': + return Searched for repositories: {part.input.query}; + } + }, [part]); + + return ( +
+ + {part.state === 'output-available' && isExpanded && ( + <> + {isServiceError(part.output) ? ( + + Failed with the following error: {part.output.message} + + ) : ( + <> + {part.output.length === 0 ? ( + No repositories found + ) : ( + + {part.output.map((repoName, index) => ( +
+ + {repoName} +
+ ))} + + )} + + )} + + + )} +
+ ) +} diff --git a/packages/web/src/features/chat/constants.ts b/packages/web/src/features/chat/constants.ts index 82d458dd..c9da5694 100644 --- a/packages/web/src/features/chat/constants.ts +++ b/packages/web/src/features/chat/constants.ts @@ -14,6 +14,8 @@ export const toolNames = { readFiles: 'readFiles', findSymbolReferences: 'findSymbolReferences', findSymbolDefinitions: 'findSymbolDefinitions', + searchRepos: 'searchRepos', + listAllRepos: 'listAllRepos', } as const; // These part types are visible in the UI. @@ -24,4 +26,6 @@ export const uiVisiblePartTypes: SBChatMessagePart['type'][] = [ 'tool-readFiles', 'tool-findSymbolDefinitions', 'tool-findSymbolReferences', + 'tool-searchRepos', + 'tool-listAllRepos', ] as const; \ No newline at end of file diff --git a/packages/web/src/features/chat/tools.ts b/packages/web/src/features/chat/tools.ts index 3a6be519..34cf1a23 100644 --- a/packages/web/src/features/chat/tools.ts +++ b/packages/web/src/features/chat/tools.ts @@ -8,6 +8,8 @@ import { findSearchBasedSymbolDefinitions, findSearchBasedSymbolReferences } fro import { FileSourceResponse } from "../search/types"; import { addLineNumbers, buildSearchQuery } from "./utils"; import { toolNames } from "./constants"; +import { getRepos } from "@/actions"; +import Fuse from "fuse.js"; // @NOTE: When adding a new tool, follow these steps: // 1. Add the tool to the `toolNames` constant in `constants.ts`. @@ -148,7 +150,7 @@ Queries consist of space-seperated regular expressions. Wrapping expressions in \`foo\` - Match files with regex /foo/ \`foo bar\` - Match files with regex /foo/ and /bar/ \`"foo bar"\` - Match files with regex /foo bar/ -\`console\.log\` - Match files with regex /console\.log/ +\`console.log\` - Match files with regex /console.log/ Multiple expressions can be or'd together with or, negated with -, or grouped with (). Examples: \`foo or bar\` - Match files with regex /foo/ or /bar/ @@ -167,8 +169,9 @@ Multiple expressions can be or'd together with or, negated with -, or grouped wi .array(z.string()) .describe(`Filter results from filepaths that match the regex. When this option is not specified, all files are searched.`) .optional(), + limit: z.number().default(10).describe("Maximum number of matches to return (default: 100)"), }), - execute: async ({ queryRegexp: _query, repoNamesFilterRegexp, languageNamesFilter, fileNamesFilterRegexp }) => { + execute: async ({ queryRegexp: _query, repoNamesFilterRegexp, languageNamesFilter, fileNamesFilterRegexp, limit }) => { const query = buildSearchQuery({ query: _query, repoNamesFilter: selectedRepos, @@ -179,7 +182,7 @@ Multiple expressions can be or'd together with or, negated with -, or grouped wi const response = await search({ query, - matches: 100, + matches: limit ?? 100, // @todo: we can make this configurable. contextLines: 3, whole: false, @@ -210,3 +213,59 @@ export type SearchCodeTool = InferUITool export type SearchCodeToolInput = InferToolInput>; export type SearchCodeToolOutput = InferToolOutput>; export type SearchCodeToolUIPart = ToolUIPart<{ [toolNames.searchCode]: SearchCodeTool }>; + +export const searchReposTool = tool({ + description: `Search for repositories by name using fuzzy search. This helps find repositories in the codebase when you know part of their name.`, + inputSchema: z.object({ + query: z.string().describe("The search query to find repositories by name (supports fuzzy matching)"), + limit: z.number().default(10).describe("Maximum number of repositories to return (default: 10)") + }), + execute: async ({ query, limit }) => { + const reposResponse = await getRepos(SINGLE_TENANT_ORG_DOMAIN); + + if (isServiceError(reposResponse)) { + return reposResponse; + } + + // Configure Fuse.js for fuzzy searching + const fuse = new Fuse(reposResponse, { + keys: [ + { name: 'repoName', weight: 0.7 }, + { name: 'repoDisplayName', weight: 0.3 } + ], + threshold: 0.4, // Lower threshold = more strict matching + includeScore: true, + minMatchCharLength: 1, + }); + + const searchResults = fuse.search(query, { limit: limit ?? 10 }); + + searchResults.sort((a, b) => (a.score ?? 0) - (b.score ?? 0)); + + return searchResults.map(({ item }) => item.repoName); + } +}); + +export type SearchReposTool = InferUITool; +export type SearchReposToolInput = InferToolInput; +export type SearchReposToolOutput = InferToolOutput; +export type SearchReposToolUIPart = ToolUIPart<{ [toolNames.searchRepos]: SearchReposTool }>; + +export const listAllReposTool = tool({ + description: `Lists all repositories in the codebase. This provides a complete overview of all available repositories.`, + inputSchema: z.object({}), + execute: async () => { + const reposResponse = await getRepos(SINGLE_TENANT_ORG_DOMAIN); + + if (isServiceError(reposResponse)) { + return reposResponse; + } + + return reposResponse.map((repo) => repo.repoName); + } +}); + +export type ListAllReposTool = InferUITool; +export type ListAllReposToolInput = InferToolInput; +export type ListAllReposToolOutput = InferToolOutput; +export type ListAllReposToolUIPart = ToolUIPart<{ [toolNames.listAllRepos]: ListAllReposTool }>; diff --git a/packages/web/src/features/chat/types.ts b/packages/web/src/features/chat/types.ts index 18aa55a2..50a06553 100644 --- a/packages/web/src/features/chat/types.ts +++ b/packages/web/src/features/chat/types.ts @@ -3,7 +3,7 @@ import { BaseEditor, Descendant } from "slate"; import { HistoryEditor } from "slate-history"; import { ReactEditor, RenderElementProps } from "slate-react"; import { z } from "zod"; -import { FindSymbolDefinitionsTool, FindSymbolReferencesTool, ReadFilesTool, SearchCodeTool } from "./tools"; +import { FindSymbolDefinitionsTool, FindSymbolReferencesTool, ReadFilesTool, SearchCodeTool, SearchReposTool, ListAllReposTool } from "./tools"; import { toolNames } from "./constants"; import { LanguageModel } from "@sourcebot/schemas/v3/index.type"; @@ -83,6 +83,8 @@ export type SBChatMessageToolTypes = { [toolNames.readFiles]: ReadFilesTool, [toolNames.findSymbolReferences]: FindSymbolReferencesTool, [toolNames.findSymbolDefinitions]: FindSymbolDefinitionsTool, + [toolNames.searchRepos]: SearchReposTool, + [toolNames.listAllRepos]: ListAllReposTool, } export type SBChatMessageDataParts = {