Skip to content

Commit f85a4da

Browse files
committed
perf: add read_file tool to drastically reduce token bloat for general chat context
1 parent 260ce64 commit f85a4da

2 files changed

Lines changed: 27 additions & 17 deletions

File tree

src/main/java/com/secai/ai/ToolExecutor.java

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,4 +112,18 @@ public static String webSearch(String query) {
112112
return "Web search failed: " + e.getMessage();
113113
}
114114
}
115+
116+
public static String readFile(String targetPath) {
117+
Path path = Paths.get(targetPath);
118+
if (!Files.exists(path)) {
119+
return "Error: File " + targetPath + " does not exist.";
120+
}
121+
System.out.println("\033[36m[AI reading file: " + targetPath + " ...]\033[0m");
122+
try {
123+
String content = Files.readString(path);
124+
return "File Contents of " + targetPath + ":\n```\n" + content + "\n```";
125+
} catch (IOException e) {
126+
return "Error reading file: " + e.getMessage();
127+
}
128+
}
115129
}

src/main/java/com/secai/cli/ChatCommand.java

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,9 @@ public Integer call() {
8484
"3. web_search\n" +
8585
" args: query\n" +
8686
" description: Searches the web.\n" +
87+
"4. read_file\n" +
88+
" args: path\n" +
89+
" description: Reads the contents of a file to gain context.\n" +
8790
"Only call ONE tool per response.";
8891

8992
String context = String.format("The user is asking questions about the following security finding:\n" +
@@ -107,24 +110,8 @@ public Integer call() {
107110
StringBuilder contextBuilder = new StringBuilder();
108111
contextBuilder.append("The user is asking questions about a recent security scan. The scan found the following issues:\n\n");
109112
for (Finding f : findings) {
110-
String fileContent = "";
111-
if (f.getFile() != null && !f.getFile().isEmpty()) {
112-
try {
113-
java.nio.file.Path targetPath = java.nio.file.Paths.get(projectPath, f.getFile());
114-
if (java.nio.file.Files.exists(targetPath)) {
115-
fileContent = java.nio.file.Files.readString(targetPath);
116-
}
117-
} catch (Exception e) {
118-
// ignore
119-
}
120-
}
121-
contextBuilder.append(String.format("- ID %s: %s\n Severity: %s\n File: %s\n Description: %s\n",
113+
contextBuilder.append(String.format("- ID %s: %s\n Severity: %s\n File: %s\n Description: %s\n\n",
122114
f.getId(), f.getTitle(), f.getSeverity(), f.getFile(), f.getDescription()));
123-
if (!fileContent.isEmpty()) {
124-
contextBuilder.append(" File Contents:\n```\n").append(fileContent).append("\n```\n\n");
125-
} else {
126-
contextBuilder.append("\n");
127-
}
128115
}
129116

130117
String toolInstructions = "\n\nTo call a tool, you MUST use this exact XML format in your response:\n" +
@@ -148,6 +135,9 @@ public Integer call() {
148135
"3. web_search\n" +
149136
" args: query\n" +
150137
" description: Searches the web.\n" +
138+
"4. read_file\n" +
139+
" args: path\n" +
140+
" description: Reads the contents of a file to gain context.\n" +
151141
"Only call ONE tool per response.";
152142

153143
contextBuilder.append("IMPORTANT INSTRUCTION: You are acting as an expert Penetration Tester and Security Educator. ");
@@ -213,6 +203,12 @@ public Integer call() {
213203
String query = extractXmlTag(argsXml, "query");
214204

215205
toolResult = com.secai.ai.ToolExecutor.webSearch(query);
206+
} else if ("read_file".equals(toolName)) {
207+
String argsXml = extractXmlTag(toolCallXml, "args");
208+
String path = extractXmlTag(argsXml, "path");
209+
path = java.nio.file.Paths.get(projectPath, path).toString();
210+
211+
toolResult = com.secai.ai.ToolExecutor.readFile(path);
216212
} else {
217213
toolResult = "Error: Unknown tool " + toolName;
218214
}

0 commit comments

Comments
 (0)