Skip to content

Commit ca354f8

Browse files
brendan-kellambrendan-kellam
committed
add param to tool definitions
1 parent 4bab982 commit ca354f8

File tree

10 files changed

+132
-140
lines changed

10 files changed

+132
-140
lines changed

‎packages/web/src/features/chat/agent.ts‎

Lines changed: 1 addition & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -221,57 +221,7 @@ const createAgentStream = async ({
221221
return;
222222
}
223223

224-
if (toolName === readFileDefinition.name) {
225-
onWriteSource({
226-
type: 'file',
227-
repo: output.metadata.repo,
228-
path: output.metadata.path,
229-
revision: output.metadata.revision,
230-
name: output.metadata.path.split('/').pop() ?? output.metadata.path,
231-
});
232-
} else if (toolName === grepDefinition.name) {
233-
output.metadata.files.forEach((file) => {
234-
onWriteSource({
235-
type: 'file',
236-
repo: file.repo,
237-
path: file.path,
238-
revision: file.revision,
239-
name: file.path.split('/').pop() ?? file.path,
240-
});
241-
});
242-
} else if (toolName === findSymbolDefinitionsDefinition.name || toolName === findSymbolReferencesDefinition.name) {
243-
output.metadata.files.forEach((file) => {
244-
onWriteSource({
245-
type: 'file',
246-
repo: file.repo,
247-
path: file.fileName,
248-
revision: file.revision,
249-
name: file.fileName.split('/').pop() ?? file.fileName,
250-
});
251-
});
252-
} else if (toolName === listTreeDefinition.name) {
253-
output.metadata.entries
254-
.filter((entry) => entry.type === 'blob')
255-
.forEach((entry) => {
256-
onWriteSource({
257-
type: 'file',
258-
repo: output.metadata.repo,
259-
path: entry.path,
260-
revision: output.metadata.ref,
261-
name: entry.name,
262-
});
263-
});
264-
} else if (toolName === globDefinition.name) {
265-
output.metadata.files.forEach((file) => {
266-
onWriteSource({
267-
type: 'file',
268-
repo: file.repo,
269-
path: file.path,
270-
revision: file.revision,
271-
name: file.path.split('/').pop() ?? file.path,
272-
});
273-
});
274-
}
224+
output.sources?.forEach(onWriteSource);
275225
});
276226
},
277227
experimental_telemetry: {

‎packages/web/src/features/chat/types.ts‎

Lines changed: 3 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,20 +5,9 @@ import { ReactEditor, RenderElementProps } from "slate-react";
55
import { z } from "zod";
66
import { LanguageModel } from "@sourcebot/schemas/v3/index.type";
77
import { createTools } from "./tools";
8-
9-
const fileSourceSchema = z.object({
10-
type: z.literal('file'),
11-
repo: z.string(),
12-
path: z.string(),
13-
name: z.string(),
14-
revision: z.string(),
15-
});
16-
export type FileSource = z.infer<typeof fileSourceSchema>;
17-
18-
export const sourceSchema = z.discriminatedUnion('type', [
19-
fileSourceSchema,
20-
]);
21-
export type Source = z.infer<typeof sourceSchema>;
8+
export { sourceSchema } from "@/features/tools/types";
9+
export type { FileSource, Source } from "@/features/tools/types";
10+
import type { Source } from "@/features/tools/types";
2211

2312
const fileReferenceSchema = z.object({
2413
type: z.literal('file'),

‎packages/web/src/features/mcp/utils.ts‎

Lines changed: 0 additions & 61 deletions
This file was deleted.

‎packages/web/src/features/tools/findSymbolDefinitions.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,9 +99,18 @@ export const findSymbolDefinitionsDefinition: ToolDefinition<
9999
}
100100
}
101101

102+
const sources = metadata.files.map((file) => ({
103+
type: 'file' as const,
104+
repo: file.repo,
105+
path: file.fileName,
106+
name: file.fileName.split('/').pop() ?? file.fileName,
107+
revision: file.revision,
108+
}));
109+
102110
return {
103111
output: outputLines.join('\n'),
104112
metadata,
113+
sources,
105114
};
106115
},
107116
};

‎packages/web/src/features/tools/findSymbolReferences.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,9 +109,18 @@ export const findSymbolReferencesDefinition: ToolDefinition<
109109
}
110110
}
111111

112+
const sources = metadata.files.map((file) => ({
113+
type: 'file' as const,
114+
repo: file.repo,
115+
path: file.fileName,
116+
name: file.fileName.split('/').pop() ?? file.fileName,
117+
revision: file.revision,
118+
}));
119+
112120
return {
113121
output: outputLines.join('\n'),
114122
metadata,
123+
sources,
115124
};
116125
},
117126
};

‎packages/web/src/features/tools/glob.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,9 +184,18 @@ export const globDefinition: ToolDefinition<'glob', typeof globShape, GlobMetada
184184
outputLines.push(TRUNCATION_MESSAGE);
185185
}
186186

187+
const sources = files.map((file) => ({
188+
type: 'file' as const,
189+
repo: file.repo,
190+
path: file.path,
191+
name: file.name,
192+
revision: file.revision,
193+
}));
194+
187195
return {
188196
output: outputLines.join('\n'),
189197
metadata,
198+
sources,
190199
};
191200
},
192201
};

‎packages/web/src/features/tools/grep.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,14 @@ export const grepDefinition: ToolDefinition<'grep', typeof grepShape, GrepMetada
179179
};
180180
}
181181

182+
const sources = files.map((file) => ({
183+
type: 'file' as const,
184+
repo: file.repo,
185+
path: file.path,
186+
name: file.path.split('/').pop() ?? file.path,
187+
revision: file.revision,
188+
}));
189+
182190
if (groupByRepo) {
183191
const repoCounts = new Map<string, { matches: number; files: number }>();
184192
for (const file of response.files) {
@@ -203,6 +211,7 @@ export const grepDefinition: ToolDefinition<'grep', typeof grepShape, GrepMetada
203211
return {
204212
output: outputLines.join('\n'),
205213
metadata,
214+
sources,
206215
};
207216
} else {
208217
const outputLines: string[] = [
@@ -234,6 +243,7 @@ export const grepDefinition: ToolDefinition<'grep', typeof grepShape, GrepMetada
234243
return {
235244
output: outputLines.join('\n'),
236245
metadata,
246+
sources,
237247
};
238248
}
239249
},

‎packages/web/src/features/tools/listTree.ts‎

Lines changed: 65 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
1-
import { z } from "zod";
1+
import { getRepoInfoByName } from "@/actions";
2+
import { FileTreeNode, getTree } from "@/features/git";
23
import { isServiceError } from "@/lib/utils";
3-
import { getTree } from "@/features/git";
4-
import { buildTreeNodeIndex, joinTreePath, normalizeTreePath, sortTreeEntries } from "@/features/mcp/utils";
5-
import { ToolDefinition } from "./types";
6-
import { logger } from "./logger";
7-
import description from "./listTree.txt";
84
import { CodeHostType } from "@sourcebot/db";
9-
import { getRepoInfoByName } from "@/actions";
5+
import { z } from "zod";
6+
import description from "./listTree.txt";
7+
import { logger } from "./logger";
8+
import { ToolDefinition } from "./types";
109

1110
const DEFAULT_TREE_DEPTH = 1;
1211
const MAX_TREE_DEPTH = 10;
@@ -42,7 +41,6 @@ export type ListTreeMetadata = {
4241
repoInfo: ListTreeRepoInfo;
4342
ref: string;
4443
path: string;
45-
entries: ListTreeEntry[];
4644
totalReturned: number;
4745
truncated: boolean;
4846
};
@@ -152,8 +150,10 @@ export const listTreeDefinition: ToolDefinition<'list_tree', typeof listTreeShap
152150

153151
const sortedEntries = sortTreeEntries(entries);
154152
const metadata: ListTreeMetadata = {
155-
repo, repoInfo, ref, path: normalizedPath,
156-
entries: sortedEntries,
153+
repo,
154+
repoInfo,
155+
ref,
156+
path: normalizedPath,
157157
totalReturned: sortedEntries.length,
158158
truncated,
159159
};
@@ -190,6 +190,60 @@ export const listTreeDefinition: ToolDefinition<'list_tree', typeof listTreeShap
190190
outputLines.push(`(truncated — showing first ${normalizedMaxEntries} entries)`);
191191
}
192192

193-
return { output: outputLines.join('\n'), metadata };
193+
const sources = sortedEntries
194+
.filter((entry) => entry.type === 'blob')
195+
.map((entry) => ({
196+
type: 'file' as const,
197+
repo,
198+
path: entry.path,
199+
name: entry.name,
200+
revision: ref,
201+
}));
202+
203+
return { output: outputLines.join('\n'), metadata, sources };
194204
},
195205
};
206+
207+
const normalizeTreePath = (path: string): string => {
208+
const withoutLeading = path.replace(/^\/+/, '');
209+
return withoutLeading.replace(/\/+$/, '');
210+
}
211+
212+
const joinTreePath = (parentPath: string, name: string): string => {
213+
if (!parentPath) {
214+
return name;
215+
}
216+
return `${parentPath}/${name}`;
217+
}
218+
219+
const buildTreeNodeIndex = (root: FileTreeNode): Map<string, FileTreeNode> => {
220+
const nodeIndex = new Map<string, FileTreeNode>();
221+
222+
const visit = (node: FileTreeNode, currentPath: string) => {
223+
nodeIndex.set(currentPath, node);
224+
for (const child of node.children) {
225+
visit(child, joinTreePath(currentPath, child.name));
226+
}
227+
};
228+
229+
visit(root, '');
230+
return nodeIndex;
231+
}
232+
233+
const sortTreeEntries = (entries: ListTreeEntry[]): ListTreeEntry[] => {
234+
const collator = new Intl.Collator(undefined, { sensitivity: 'base' });
235+
236+
return [...entries].sort((a, b) => {
237+
const parentCompare = collator.compare(a.parentPath, b.parentPath);
238+
if (parentCompare !== 0) return parentCompare;
239+
240+
if (a.type !== b.type) {
241+
return a.type === 'tree' ? -1 : 1;
242+
}
243+
244+
const nameCompare = collator.compare(a.name, b.name);
245+
if (nameCompare !== 0) return nameCompare;
246+
247+
return collator.compare(a.path, b.path);
248+
});
249+
}

‎packages/web/src/features/tools/readFile.ts‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@ export type ReadFileMetadata = {
3535
path: string;
3636
repo: string;
3737
repoInfo: ReadFileRepoInfo;
38-
language: string;
3938
startLine: number;
4039
endLine: number;
4140
isTruncated: boolean;
@@ -120,13 +119,22 @@ export const readFileDefinition: ToolDefinition<"read_file", typeof readFileShap
120119
path: fileSource.path,
121120
repo: fileSource.repo,
122121
repoInfo,
123-
language: fileSource.language,
124122
startLine,
125123
endLine: lastReadLine,
126124
isTruncated: truncatedByBytes || truncatedByLines,
127125
revision,
128126
};
129127

130-
return { output, metadata };
128+
return {
129+
output,
130+
metadata,
131+
sources: [{
132+
type: 'file',
133+
repo: fileSource.repo,
134+
path: fileSource.path,
135+
name: fileSource.path.split('/').pop() ?? fileSource.path,
136+
revision,
137+
}],
138+
};
131139
},
132140
};

‎packages/web/src/features/tools/types.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,19 @@
11
import { z } from "zod";
22

3+
const fileSourceSchema = z.object({
4+
type: z.literal('file'),
5+
repo: z.string(),
6+
path: z.string(),
7+
name: z.string(),
8+
revision: z.string(),
9+
});
10+
export type FileSource = z.infer<typeof fileSourceSchema>;
11+
12+
export const sourceSchema = z.discriminatedUnion('type', [
13+
fileSourceSchema,
14+
]);
15+
export type Source = z.infer<typeof sourceSchema>;
16+
317
export interface ToolContext {
418
source?: string;
519
selectedRepos?: string[];
@@ -22,4 +36,5 @@ export interface ToolDefinition<
2236
export interface ToolResult<TMetadata = Record<string, unknown>> {
2337
output: string;
2438
metadata: TMetadata;
39+
sources?: Source[];
2540
}

0 commit comments

Comments
 (0)