Skip to content

Commit 4b991b3

Browse files
jguoamzjguoamz
committed
fix(chat): Update FsRead toolspec, set different output limit for tools
1 parent b222e0b commit 4b991b3

File tree

6 files changed

+85
-14
lines changed

6 files changed

+85
-14
lines changed

‎packages/core/src/codewhispererChat/controllers/chat/controller.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -746,7 +746,7 @@ export class ChatController {
746746
session.setFsWriteBackup(toolUse.toolUseId, backup)
747747
}
748748
const output = await ToolUtils.invoke(tool, chatStream)
749-
ToolUtils.validateOutput(output)
749+
ToolUtils.validateOutput(output, tool.type)
750750

751751
toolResults.push({
752752
content: [

‎packages/core/src/codewhispererChat/tools/fsRead.ts‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import * as vscode from 'vscode'
66
import { getLogger } from '../../shared/logger/logger'
77
import fs from '../../shared/fs/fs'
88
import { Writable } from 'stream'
9-
import { InvokeOutput, OutputKind, sanitizePath, CommandValidation } from './toolShared'
9+
import { InvokeOutput, OutputKind, sanitizePath, CommandValidation, fsReadToolResponseSize } from './toolShared'
1010
import { isInDirectory } from '../../shared/filesystemUtilities'
1111

1212
export interface FsReadParams {
@@ -118,11 +118,22 @@ export class FsRead {
118118
}
119119

120120
private createOutput(content: string): InvokeOutput {
121+
if (content.length > fsReadToolResponseSize) {
122+
content = 'The file is too large, below are the truncated output:\n' + content
123+
content = this.truncateContent(content)
124+
}
121125
return {
122126
output: {
123127
kind: OutputKind.Text,
124128
content: content,
125129
},
126130
}
127131
}
132+
133+
private truncateContent(content: string): string {
134+
if (content.length > fsReadToolResponseSize) {
135+
return content.substring(0, fsReadToolResponseSize)
136+
}
137+
return content
138+
}
128139
}

‎packages/core/src/codewhispererChat/tools/toolShared.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,9 @@
66
import path from 'path'
77
import fs from '../../shared/fs/fs'
88

9-
export const maxToolResponseSize = 200_000
9+
export const defaultMaxToolResponseSize = 100_000
10+
export const listDirectoryToolResponseSize = 50_000
11+
export const fsReadToolResponseSize = 200_000
1012

1113
export enum OutputKind {
1214
Text = 'text',

‎packages/core/src/codewhispererChat/tools/toolUtils.ts‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,12 @@ import { FsRead, FsReadParams } from './fsRead'
77
import { FsWrite, FsWriteParams } from './fsWrite'
88
import { CommandValidation, ExecuteBash, ExecuteBashParams } from './executeBash'
99
import { ToolResult, ToolResultContentBlock, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming'
10-
import { InvokeOutput, maxToolResponseSize } from './toolShared'
10+
import {
11+
InvokeOutput,
12+
defaultMaxToolResponseSize,
13+
listDirectoryToolResponseSize,
14+
fsReadToolResponseSize,
15+
} from './toolShared'
1116
import { ListDirectory, ListDirectoryParams } from './listDirectory'
1217

1318
export enum ToolType {
@@ -63,9 +68,20 @@ export class ToolUtils {
6368
}
6469
}
6570

66-
static validateOutput(output: InvokeOutput): void {
71+
static validateOutput(output: InvokeOutput, toolType: ToolType): void {
72+
let maxToolResponseSize = defaultMaxToolResponseSize
73+
switch (toolType) {
74+
case ToolType.FsRead:
75+
maxToolResponseSize = fsReadToolResponseSize
76+
break
77+
case ToolType.ListDirectory:
78+
maxToolResponseSize = listDirectoryToolResponseSize
79+
break
80+
default:
81+
break
82+
}
6783
if (output.output.content.length > maxToolResponseSize) {
68-
throw Error(`Tool output exceeds maximum character limit of ${maxToolResponseSize}`)
84+
throw Error(`${toolType} output exceeds maximum character limit of ${maxToolResponseSize}`)
6985
}
7086
}
7187

‎packages/core/src/codewhispererChat/tools/tool_index.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
"type": "string"
1111
},
1212
"readRange": {
13-
"description": "Optional parameter when reading files.\n * If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. Setting `[startLine, -1]` shows all lines from `startLine` to the end of the file. If the whole file is too large, try reading 4000 lines at once, for example: after reading [1, 4000], read [4000, 8000] next and repeat. You should read atleast 250 lines per invocation of the tool. In some cases, if reading a range of lines results in too many invocations instead attempt to read 4000 lines.",
13+
"description": "Optional parameter when reading files.\n * If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. If the whole file is too large, try reading 5000 lines at once, for example, use [1, 5000] first and then [5000, 10000] next and so on. You should read AT LEAST 1000 lines per invocation of the tool.",
1414
"items": {
1515
"type": "integer"
1616
},

‎packages/core/src/test/codewhispererChat/tools/toolShared.test.ts‎

Lines changed: 49 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -158,25 +158,67 @@ describe('ToolUtils', function () {
158158
})
159159

160160
describe('validateOutput', function () {
161-
it('does not throw error if output is within size limit', function () {
161+
it('does not throw error if output is within size limit for fsRead', function () {
162162
const output: InvokeOutput = {
163163
output: {
164164
kind: OutputKind.Text,
165-
content: 'a'.repeat(150_000),
165+
content: 'a'.repeat(199_000),
166166
},
167167
}
168-
assert.doesNotThrow(() => ToolUtils.validateOutput(output))
168+
assert.doesNotThrow(() => ToolUtils.validateOutput(output, ToolType.FsRead))
169169
})
170-
it('throws error if output exceeds max size', function () {
170+
it('throws error if output exceeds max size for fsRead', function () {
171171
const output: InvokeOutput = {
172172
output: {
173173
kind: OutputKind.Text,
174-
content: 'a'.repeat(200_001), // 200,001 characters
174+
content: 'a'.repeat(200_001),
175175
},
176176
}
177-
assert.throws(() => ToolUtils.validateOutput(output), {
177+
assert.throws(() => ToolUtils.validateOutput(output, ToolType.FsRead), {
178178
name: 'Error',
179-
message: 'Tool output exceeds maximum character limit of 200000',
179+
message: 'fsRead output exceeds maximum character limit of 200000',
180+
})
181+
})
182+
it('does not throw error if output is within size limit for listDirectory', function () {
183+
const output: InvokeOutput = {
184+
output: {
185+
kind: OutputKind.Text,
186+
content: 'a'.repeat(49_000),
187+
},
188+
}
189+
assert.doesNotThrow(() => ToolUtils.validateOutput(output, ToolType.ListDirectory))
190+
})
191+
it('throws error if output exceeds max size for listDirectory', function () {
192+
const output: InvokeOutput = {
193+
output: {
194+
kind: OutputKind.Text,
195+
content: 'a'.repeat(50_001),
196+
},
197+
}
198+
assert.throws(() => ToolUtils.validateOutput(output, ToolType.ListDirectory), {
199+
name: 'Error',
200+
message: 'fsRead output exceeds maximum character limit of 50000',
201+
})
202+
})
203+
it('does not throw error if output is within size limit for fsWrite', function () {
204+
const output: InvokeOutput = {
205+
output: {
206+
kind: OutputKind.Text,
207+
content: 'a'.repeat(90_000),
208+
},
209+
}
210+
assert.doesNotThrow(() => ToolUtils.validateOutput(output, ToolType.FsWrite))
211+
})
212+
it('does not throw error if output is within size limit for fsWrite', function () {
213+
const output: InvokeOutput = {
214+
output: {
215+
kind: OutputKind.Text,
216+
content: 'a'.repeat(100_001),
217+
},
218+
}
219+
assert.throws(() => ToolUtils.validateOutput(output, ToolType.ListDirectory), {
220+
name: 'Error',
221+
message: 'fsWrite output exceeds maximum character limit of 100000',
180222
})
181223
})
182224
})

0 commit comments

Comments
 (0)