Skip to content

Commit 6baf28c

Browse files
authored
feat(tools): add support for reading PDF, DOCX, and IPYNB files in read_file tool (#4288)
- Allow specific binary formats (.pdf, .docx, .ipynb) to be processed by extractTextFromFile - Block unsupported binary files with existing "Binary file" notice - Update tests to cover both supported and unsupported binary file scenarios - Refactor test mocks for better maintainability and coverage
1 parent d87f890 commit 6baf28c

File tree

3 files changed

+150
-55
lines changed

3 files changed

+150
-55
lines changed

src/core/tools/__tests__/readFileTool.test.ts

Lines changed: 96 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -44,16 +44,14 @@ const addLineNumbersMock = jest.fn().mockImplementation((text, startLine = 1) =>
4444
return lines.map((line, i) => `${startLine + i} | ${line}`).join("\n")
4545
})
4646

47-
const extractTextFromFileMock = jest.fn().mockImplementation((_filePath) => {
48-
// Call addLineNumbersMock to register the call
49-
addLineNumbersMock(mockInputContent)
50-
return Promise.resolve(addLineNumbersMock(mockInputContent))
51-
})
47+
const extractTextFromFileMock = jest.fn()
48+
const getSupportedBinaryFormatsMock = jest.fn(() => [".pdf", ".docx", ".ipynb"])
5249

5350
// Now assign the mocks to the module
5451
const extractTextModule = jest.requireMock("../../../integrations/misc/extract-text")
5552
extractTextModule.extractTextFromFile = extractTextFromFileMock
5653
extractTextModule.addLineNumbers = addLineNumbersMock
54+
extractTextModule.getSupportedBinaryFormats = getSupportedBinaryFormatsMock
5755

5856
jest.mock("../../ignore/RooIgnoreController", () => ({
5957
RooIgnoreController: class {
@@ -128,6 +126,9 @@ describe("read_file tool with maxReadFileLine setting", () => {
128126
mockCline.say = jest.fn().mockResolvedValue(undefined)
129127
mockCline.ask = jest.fn().mockResolvedValue({ response: "yesButtonClicked" })
130128
mockCline.presentAssistantMessage = jest.fn()
129+
mockCline.handleError = jest.fn().mockResolvedValue(undefined)
130+
mockCline.pushToolResult = jest.fn()
131+
mockCline.removeClosingTag = jest.fn((tag, content) => content)
131132

132133
mockCline.fileContextTracker = {
133134
trackFileContext: jest.fn().mockResolvedValue(undefined),
@@ -410,6 +411,13 @@ describe("read_file tool XML output structure", () => {
410411
mockedPathResolve.mockReturnValue(absoluteFilePath)
411412
mockedIsBinaryFile.mockResolvedValue(false)
412413

414+
// Set default implementation for extractTextFromFile
415+
mockedExtractTextFromFile.mockImplementation((filePath) => {
416+
// Call addLineNumbersMock to register the call
417+
addLineNumbersMock(mockInputContent)
418+
return Promise.resolve(addLineNumbersMock(mockInputContent))
419+
})
420+
413421
mockInputContent = fileContent
414422

415423
// Setup mock provider with default maxReadFileLine
@@ -1126,34 +1134,101 @@ describe("read_file tool XML output structure", () => {
11261134
const textPath = "test/text.txt"
11271135
const binaryPath = "test/binary.pdf"
11281136
const numberedContent = "1 | Text file content"
1137+
const pdfContent = "1 | PDF content extracted"
1138+
1139+
// Mock path.resolve to return the expected paths
1140+
mockedPathResolve.mockImplementation((cwd, relPath) => `/${relPath}`)
11291141

11301142
// Mock binary file detection
1131-
mockedIsBinaryFile.mockImplementationOnce(() => Promise.resolve(false))
1132-
mockedIsBinaryFile.mockImplementationOnce(() => Promise.resolve(true))
1143+
mockedIsBinaryFile.mockImplementation((path) => {
1144+
if (path.includes("text.txt")) return Promise.resolve(false)
1145+
if (path.includes("binary.pdf")) return Promise.resolve(true)
1146+
return Promise.resolve(false)
1147+
})
1148+
1149+
mockedCountFileLines.mockImplementation((path) => {
1150+
return Promise.resolve(1)
1151+
})
11331152

1134-
// Mock content based on file type
11351153
mockedExtractTextFromFile.mockImplementation((path) => {
1136-
if (path.includes("binary")) {
1137-
return Promise.resolve("")
1154+
if (path.includes("binary.pdf")) {
1155+
return Promise.resolve(pdfContent)
11381156
}
11391157
return Promise.resolve(numberedContent)
11401158
})
1141-
mockedCountFileLines.mockImplementation((path) => {
1142-
return Promise.resolve(path.includes("binary") ? 0 : 1)
1143-
})
1159+
1160+
// Configure mocks for the test
11441161
mockProvider.getState.mockResolvedValue({ maxReadFileLine: -1 })
11451162

1146-
// Execute
1147-
const result = await executeReadFileTool(
1148-
{
1163+
// Create standalone mock functions
1164+
const mockAskApproval = jest.fn().mockResolvedValue({ response: "yesButtonClicked" })
1165+
const mockHandleError = jest.fn().mockResolvedValue(undefined)
1166+
const mockPushToolResult = jest.fn()
1167+
const mockRemoveClosingTag = jest.fn((tag, content) => content)
1168+
1169+
// Create a tool use object directly
1170+
const toolUse: ReadFileToolUse = {
1171+
type: "tool_use",
1172+
name: "read_file",
1173+
params: {
11491174
args: `<file><path>${textPath}</path></file><file><path>${binaryPath}</path></file>`,
11501175
},
1151-
{ totalLines: 1 },
1176+
partial: false,
1177+
}
1178+
1179+
// Call readFileTool directly
1180+
await readFileTool(
1181+
mockCline,
1182+
toolUse,
1183+
mockAskApproval,
1184+
mockHandleError,
1185+
mockPushToolResult,
1186+
mockRemoveClosingTag,
11521187
)
11531188

1154-
// Verify
1155-
expect(result).toBe(
1156-
`<files>\n<file><path>${textPath}</path>\n<content lines="1-1">\n${numberedContent}</content>\n</file>\n<file><path>${binaryPath}</path>\n<notice>Binary file</notice>\n</file>\n</files>`,
1189+
// Check the result
1190+
expect(mockPushToolResult).toHaveBeenCalledWith(
1191+
`<files>\n<file><path>${textPath}</path>\n<content lines="1-1">\n${numberedContent}</content>\n</file>\n<file><path>${binaryPath}</path>\n<content lines="1-1">\n${pdfContent}</content>\n</file>\n</files>`,
1192+
)
1193+
})
1194+
1195+
it("should block unsupported binary files", async () => {
1196+
// Setup
1197+
const unsupportedBinaryPath = "test/binary.exe"
1198+
1199+
mockedIsBinaryFile.mockImplementation(() => Promise.resolve(true))
1200+
mockedCountFileLines.mockImplementation(() => Promise.resolve(1))
1201+
mockProvider.getState.mockResolvedValue({ maxReadFileLine: -1 })
1202+
1203+
// Create standalone mock functions
1204+
const mockAskApproval = jest.fn().mockResolvedValue({ response: "yesButtonClicked" })
1205+
const mockHandleError = jest.fn().mockResolvedValue(undefined)
1206+
const mockPushToolResult = jest.fn()
1207+
const mockRemoveClosingTag = jest.fn((tag, content) => content)
1208+
1209+
// Create a tool use object directly
1210+
const toolUse: ReadFileToolUse = {
1211+
type: "tool_use",
1212+
name: "read_file",
1213+
params: {
1214+
args: `<file><path>${unsupportedBinaryPath}</path></file>`,
1215+
},
1216+
partial: false,
1217+
}
1218+
1219+
// Call readFileTool directly
1220+
await readFileTool(
1221+
mockCline,
1222+
toolUse,
1223+
mockAskApproval,
1224+
mockHandleError,
1225+
mockPushToolResult,
1226+
mockRemoveClosingTag,
1227+
)
1228+
1229+
// Check the result
1230+
expect(mockPushToolResult).toHaveBeenCalledWith(
1231+
`<files>\n<file><path>${unsupportedBinaryPath}</path>\n<notice>Binary file</notice>\n</file>\n</files>`,
11571232
)
11581233
})
11591234
})
@@ -1165,6 +1240,7 @@ describe("read_file tool XML output structure", () => {
11651240
const maxReadFileLine = -1
11661241
const totalLines = 0
11671242
mockedCountFileLines.mockResolvedValue(totalLines)
1243+
mockedIsBinaryFile.mockResolvedValue(false) // Ensure empty file is not detected as binary
11681244

11691245
// Execute
11701246
const result = await executeReadFileTool({}, { maxReadFileLine, totalLines })

src/core/tools/readFileTool.ts

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import { isPathOutsideWorkspace } from "../../utils/pathUtils"
1111
import { getReadablePath } from "../../utils/path"
1212
import { countFileLines } from "../../integrations/misc/line-counter"
1313
import { readLines } from "../../integrations/misc/read-lines"
14-
import { extractTextFromFile, addLineNumbers } from "../../integrations/misc/extract-text"
14+
import { extractTextFromFile, addLineNumbers, getSupportedBinaryFormats } from "../../integrations/misc/extract-text"
1515
import { parseSourceCodeDefinitionsForFile } from "../../services/tree-sitter"
1616
import { parseXml } from "../../utils/xml"
1717

@@ -435,13 +435,19 @@ export async function readFileTool(
435435
try {
436436
const [totalLines, isBinary] = await Promise.all([countFileLines(fullPath), isBinaryFile(fullPath)])
437437

438-
// Handle binary files
438+
// Handle binary files (but allow specific file types that extractTextFromFile can handle)
439439
if (isBinary) {
440-
updateFileResult(relPath, {
441-
notice: "Binary file",
442-
xmlContent: `<file><path>${relPath}</path>\n<notice>Binary file</notice>\n</file>`,
443-
})
444-
continue
440+
const fileExtension = path.extname(relPath).toLowerCase()
441+
const supportedBinaryFormats = getSupportedBinaryFormats()
442+
443+
if (!supportedBinaryFormats.includes(fileExtension)) {
444+
updateFileResult(relPath, {
445+
notice: "Binary file",
446+
xmlContent: `<file><path>${relPath}</path>\n<notice>Binary file</notice>\n</file>`,
447+
})
448+
continue
449+
}
450+
// For supported binary formats (.pdf, .docx, .ipynb), continue to extractTextFromFile
445451
}
446452

447453
// Handle range reads (bypass maxReadFileLine)

src/integrations/misc/extract-text.ts

Lines changed: 41 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -5,34 +5,6 @@ import mammoth from "mammoth"
55
import fs from "fs/promises"
66
import { isBinaryFile } from "isbinaryfile"
77

8-
export async function extractTextFromFile(filePath: string): Promise<string> {
9-
try {
10-
await fs.access(filePath)
11-
} catch (error) {
12-
throw new Error(`File not found: ${filePath}`)
13-
}
14-
15-
const fileExtension = path.extname(filePath).toLowerCase()
16-
17-
switch (fileExtension) {
18-
case ".pdf":
19-
return extractTextFromPDF(filePath)
20-
case ".docx":
21-
return extractTextFromDOCX(filePath)
22-
case ".ipynb":
23-
return extractTextFromIPYNB(filePath)
24-
default: {
25-
const isBinary = await isBinaryFile(filePath).catch(() => false)
26-
27-
if (!isBinary) {
28-
return addLineNumbers(await fs.readFile(filePath, "utf8"))
29-
} else {
30-
throw new Error(`Cannot read text for file type: ${fileExtension}`)
31-
}
32-
}
33-
}
34-
}
35-
368
async function extractTextFromPDF(filePath: string): Promise<string> {
379
const dataBuffer = await fs.readFile(filePath)
3810
const data = await pdf(dataBuffer)
@@ -58,6 +30,47 @@ async function extractTextFromIPYNB(filePath: string): Promise<string> {
5830
return addLineNumbers(extractedText)
5931
}
6032

33+
/**
34+
* Map of supported binary file formats to their extraction functions
35+
*/
36+
const SUPPORTED_BINARY_FORMATS = {
37+
".pdf": extractTextFromPDF,
38+
".docx": extractTextFromDOCX,
39+
".ipynb": extractTextFromIPYNB,
40+
} as const
41+
42+
/**
43+
* Returns the list of supported binary file formats that can be processed by extractTextFromFile
44+
*/
45+
export function getSupportedBinaryFormats(): string[] {
46+
return Object.keys(SUPPORTED_BINARY_FORMATS)
47+
}
48+
49+
export async function extractTextFromFile(filePath: string): Promise<string> {
50+
try {
51+
await fs.access(filePath)
52+
} catch (error) {
53+
throw new Error(`File not found: ${filePath}`)
54+
}
55+
56+
const fileExtension = path.extname(filePath).toLowerCase()
57+
58+
// Check if we have a specific extractor for this format
59+
const extractor = SUPPORTED_BINARY_FORMATS[fileExtension as keyof typeof SUPPORTED_BINARY_FORMATS]
60+
if (extractor) {
61+
return extractor(filePath)
62+
}
63+
64+
// Handle other files
65+
const isBinary = await isBinaryFile(filePath).catch(() => false)
66+
67+
if (!isBinary) {
68+
return addLineNumbers(await fs.readFile(filePath, "utf8"))
69+
} else {
70+
throw new Error(`Cannot read text for file type: ${fileExtension}`)
71+
}
72+
}
73+
6174
export function addLineNumbers(content: string, startLine: number = 1): string {
6275
// If content is empty, return empty string - empty files should not have line numbers
6376
// If content is empty but startLine > 1, return "startLine | " because we know the file is not empty

0 commit comments

Comments
 (0)