Skip to content

Commit 541732f

Browse files
jguoamzjguoamz
committed
fix(chat): Update FsRead toolspec, return truncated output is file is too large
1 parent b222e0b commit 541732f

File tree

7 files changed

+99
-15
lines changed

7 files changed

+99
-15
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: 14 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,24 @@ export class FsRead {
118118
}
119119

120120
private createOutput(content: string): InvokeOutput {
121+
if (content.length > fsReadToolResponseSize) {
122+
this.logger.info(
123+
`The file is too large, truncating output to the first ${fsReadToolResponseSize} characters.`
124+
)
125+
content = this.truncateContent(content)
126+
}
121127
return {
122128
output: {
123129
kind: OutputKind.Text,
124130
content: content,
125131
},
126132
}
127133
}
134+
135+
private truncateContent(content: string): string {
136+
if (content.length > fsReadToolResponseSize) {
137+
return content.substring(0, fsReadToolResponseSize)
138+
}
139+
return content
140+
}
128141
}

‎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: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
{
22
"fsRead": {
33
"name": "fsRead",
4-
"description": "A tool for reading a file.\n * This tool returns the contents of a file, and the optional `readRange` determines what range of lines will be read from the specified file.",
4+
"description": "A tool for reading a file.\n * This tool returns the contents of a file, and the optional `readRange` determines what range of lines will be read from the specified file.\n * If the file exceeds 200K characters, this tool will only read a portion of the file.",
55
"inputSchema": {
66
"type": "object",
77
"properties": {
@@ -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. ",
1414
"items": {
1515
"type": "integer"
1616
},

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { TestFolder } from '../../testUtil'
88
import path from 'path'
99
import * as vscode from 'vscode'
1010
import sinon from 'sinon'
11+
import { fsReadToolResponseSize } from '../../../codewhispererChat/tools/toolShared'
1112

1213
describe('FsRead Tool', () => {
1314
let testFolder: TestFolder
@@ -37,6 +38,16 @@ describe('FsRead Tool', () => {
3738
assert.strictEqual(result.output.content, fileContent, 'File content should match exactly')
3839
})
3940

41+
it('truncate output if too large', async () => {
42+
const fileContent = 'A'.repeat(fsReadToolResponseSize + 1)
43+
const filePath = await testFolder.write('largeFile.txt', fileContent)
44+
const fsRead = new FsRead({ path: filePath })
45+
await fsRead.validate()
46+
const result = await fsRead.invoke(process.stdout)
47+
assert.strictEqual(result.output.kind, 'text', 'Output kind should be "text"')
48+
assert.ok(result.output.content.startsWith('The file is too large, below are the truncated output:\n'))
49+
})
50+
4051
it('reads partial lines of a file', async () => {
4152
const fileContent = 'A\nB\nC\nD\nE\nF'
4253
const filePath = await testFolder.write('partialFile.txt', fileContent)

‎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: 'listDirectory 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.FsWrite), {
220+
name: 'Error',
221+
message: 'fsWrite output exceeds maximum character limit of 100000',
180222
})
181223
})
182224
})

0 commit comments

Comments
 (0)