Skip to content

Commit 08b6513

Browse files
sazidsazid
committed
Add LineCount and FileSize tools for measuring files before reading
These tools help the agent assess file size before reading to avoid wasting context window space on large files: - LineCount: Fast SIMD-optimized line counting using bytes.Count (~5-10 GB/s throughput with 1MB buffer) - FileSize: Returns file size in bytes with human-readable format Both tools are registered and the system prompt updated to instruct the agent to measure files before reading large content. Co-Authored-By: BitCode <https://github.com/sazid/bitcode>
1 parent 2fe2f6e commit 08b6513

6 files changed

Lines changed: 598 additions & 0 deletions

File tree

‎app/main.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,8 @@ func main() {
6363
toolManager.Register(&tools.GlobTool{})
6464
toolManager.Register(&tools.BashTool{})
6565
toolManager.Register(&tools.WebSearchTool{})
66+
toolManager.Register(&tools.LineCountTool{})
67+
toolManager.Register(&tools.FileSizeTool{})
6668

6769
todoStore := tools.NewTodoStore()
6870
toolManager.Register(&tools.TodoReadTool{Store: todoStore})

‎app/system_prompt.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,10 @@ IMPORTANT: You must NEVER generate or guess URLs for the user unless you are con
8181
- To create files use Write instead of cat with heredoc or echo redirection
8282
- To search for files use Glob instead of find or ls
8383
- Reserve using the Bash exclusively for system commands and terminal operations that require shell execution.
84+
- ALWAYS measure files before reading to avoid wasting context:
85+
- Before reading a file, use FileSize and/or LineCount to check the file size
86+
- If a file is large (>500 lines or >50KB), consider using offset/limit parameters or searching for specific text patterns instead of reading the entire file
87+
- Large files consume significant context window space - be judicious about when to read whole files
8488
8589
# Communication style
8690
- When starting work on a user request, ALWAYS begin by briefly restating what you understand the user wants in your own words (1-2 sentences). This lets the user confirm you're on the right track before you dive in.

‎internal/tools/filesize.go‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
package tools
2+
3+
import (
4+
"encoding/json"
5+
"fmt"
6+
"os"
7+
"path"
8+
"path/filepath"
9+
"strings"
10+
11+
"github.com/sazid/bitcode/internal"
12+
)
13+
14+
type FileSizeInput struct {
15+
FilePath string `json:"file_path"`
16+
}
17+
18+
type FileSizeTool struct{}
19+
20+
var _ Tool = (*FileSizeTool)(nil)
21+
22+
func (f *FileSizeTool) Name() string {
23+
return "FileSize"
24+
}
25+
26+
func (f *FileSizeTool) Description() string {
27+
return `Gets the size of a file in bytes.
28+
29+
IMPORTANT:
30+
- Use this tool BEFORE reading files to assess their size and avoid wasting context.
31+
- Supports both absolute and relative paths.
32+
- Returns size in bytes and human-readable format (KB, MB, GB).
33+
- This tool can only get file sizes, not directory sizes.
34+
35+
Parameters:
36+
- file_path (required): The path to the file (absolute or relative to current working directory)`
37+
}
38+
39+
func (f *FileSizeTool) ParametersSchema() map[string]any {
40+
return map[string]any{
41+
"type": "object",
42+
"properties": map[string]any{
43+
"file_path": map[string]any{
44+
"type": "string",
45+
"description": "The path to the file (absolute or relative to current working directory)",
46+
},
47+
},
48+
"required": []string{"file_path"},
49+
}
50+
}
51+
52+
func (f *FileSizeTool) Execute(input json.RawMessage, eventsCh chan<- internal.Event) (ToolResult, error) {
53+
var params FileSizeInput
54+
if err := json.Unmarshal(input, &params); err != nil {
55+
return ToolResult{}, fmt.Errorf("invalid input: %w", err)
56+
}
57+
58+
wd, err := os.Getwd()
59+
if err != nil {
60+
return ToolResult{}, fmt.Errorf("failed to get working directory: %w", err)
61+
}
62+
63+
fullPath := params.FilePath
64+
if !filepath.IsAbs(fullPath) {
65+
fullPath = path.Join(wd, fullPath)
66+
}
67+
68+
cleanPath := filepath.Clean(fullPath)
69+
if strings.Contains(cleanPath, "..") {
70+
return ToolResult{}, fmt.Errorf("file_path cannot contain '..' for security reasons")
71+
}
72+
73+
stat, err := os.Stat(cleanPath)
74+
if err != nil {
75+
if os.IsNotExist(err) {
76+
return ToolResult{}, fmt.Errorf("file does not exist: %w", err)
77+
}
78+
return ToolResult{}, fmt.Errorf("failed to stat file: %w", err)
79+
}
80+
81+
if stat.IsDir() {
82+
return ToolResult{}, fmt.Errorf("path is a directory, not a file: %s", cleanPath)
83+
}
84+
85+
size := stat.Size()
86+
humanReadable := formatSize(size)
87+
88+
info := fmt.Sprintf("%s (%d bytes)", humanReadable, size)
89+
eventsCh <- internal.Event{
90+
Name: f.Name(),
91+
Args: []string{cleanPath},
92+
Message: info,
93+
}
94+
95+
return ToolResult{
96+
Content: fmt.Sprintf("%d bytes (%s) %s", size, humanReadable, cleanPath),
97+
}, nil
98+
}
99+
100+
// formatSize converts a size in bytes to a human-readable string.
101+
func formatSize(bytes int64) string {
102+
const (
103+
KB = 1024
104+
MB = KB * 1024
105+
GB = MB * 1024
106+
)
107+
108+
switch {
109+
case bytes >= GB:
110+
return fmt.Sprintf("%.2f GB", float64(bytes)/float64(GB))
111+
case bytes >= MB:
112+
return fmt.Sprintf("%.2f MB", float64(bytes)/float64(MB))
113+
case bytes >= KB:
114+
return fmt.Sprintf("%.2f KB", float64(bytes)/float64(KB))
115+
default:
116+
return fmt.Sprintf("%d B", bytes)
117+
}
118+
}

‎internal/tools/filesize_test.go‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
package tools
2+
3+
import (
4+
"encoding/json"
5+
"os"
6+
"strings"
7+
"testing"
8+
9+
"github.com/sazid/bitcode/internal"
10+
)
11+
12+
func executeFileSize(t *testing.T, input FileSizeInput) (ToolResult, error) {
13+
t.Helper()
14+
raw, err := json.Marshal(input)
15+
if err != nil {
16+
t.Fatalf("failed to marshal input: %v", err)
17+
}
18+
tool := &FileSizeTool{}
19+
ch := makeEventsCh()
20+
result, err := tool.Execute(raw, ch)
21+
close(ch)
22+
return result, err
23+
}
24+
25+
func TestFileSizeTool_GetsSize(t *testing.T) {
26+
content := "Hello, World!"
27+
path := writeTempFile(t, content)
28+
result, err := executeFileSize(t, FileSizeInput{FilePath: path})
29+
if err != nil {
30+
t.Fatalf("unexpected error: %v", err)
31+
}
32+
// 13 bytes
33+
if !strings.Contains(result.Content, "13") {
34+
t.Errorf("expected 13 bytes, got: %q", result.Content)
35+
}
36+
}
37+
38+
func TestFileSizeTool_EmptyFile(t *testing.T) {
39+
path := writeTempFile(t, "")
40+
result, err := executeFileSize(t, FileSizeInput{FilePath: path})
41+
if err != nil {
42+
t.Fatalf("unexpected error: %v", err)
43+
}
44+
if !strings.Contains(result.Content, "0") {
45+
t.Errorf("expected 0 bytes, got: %q", result.Content)
46+
}
47+
}
48+
49+
func TestFileSizeTool_LargeFile(t *testing.T) {
50+
// Create a 1MB file
51+
content := strings.Repeat("x", 1024*1024)
52+
path := writeTempFile(t, content)
53+
result, err := executeFileSize(t, FileSizeInput{FilePath: path})
54+
if err != nil {
55+
t.Fatalf("unexpected error: %v", err)
56+
}
57+
// Check for 1.00 MB in human-readable format
58+
if !strings.Contains(result.Content, "1.00 MB") && !strings.Contains(result.Content, "1048576") {
59+
t.Errorf("expected ~1MB size, got: %q", result.Content)
60+
}
61+
}
62+
63+
func TestFileSizeTool_KBSize(t *testing.T) {
64+
// Create a 5KB file
65+
content := strings.Repeat("a", 5*1024)
66+
path := writeTempFile(t, content)
67+
result, err := executeFileSize(t, FileSizeInput{FilePath: path})
68+
if err != nil {
69+
t.Fatalf("unexpected error: %v", err)
70+
}
71+
if !strings.Contains(result.Content, "5") {
72+
t.Errorf("expected 5KB size, got: %q", result.Content)
73+
}
74+
}
75+
76+
func TestFileSizeTool_FileNotFound(t *testing.T) {
77+
_, err := executeFileSize(t, FileSizeInput{FilePath: "/nonexistent/path/to/file.txt"})
78+
if err == nil {
79+
t.Fatal("expected error for non-existent file, got nil")
80+
}
81+
}
82+
83+
func TestFileSizeTool_DirectoryError(t *testing.T) {
84+
// Try to get size of a directory
85+
dir := t.TempDir()
86+
_, err := executeFileSize(t, FileSizeInput{FilePath: dir})
87+
if err == nil {
88+
t.Fatal("expected error for directory, got nil")
89+
}
90+
if !strings.Contains(err.Error(), "directory") {
91+
t.Errorf("expected directory error, got: %v", err)
92+
}
93+
}
94+
95+
func TestFileSizeTool_PathTraversal(t *testing.T) {
96+
_, err := executeFileSize(t, FileSizeInput{FilePath: "../../etc/passwd"})
97+
// Should fail due to path traversal check
98+
if err == nil {
99+
t.Fatal("expected error for path traversal, got nil")
100+
}
101+
}
102+
103+
func TestFileSizeTool_EmitsEvent(t *testing.T) {
104+
content := "test content"
105+
path := writeTempFile(t, content)
106+
raw, _ := json.Marshal(FileSizeInput{FilePath: path})
107+
tool := &FileSizeTool{}
108+
ch := makeEventsCh()
109+
_, err := tool.Execute(raw, ch)
110+
close(ch)
111+
if err != nil {
112+
t.Fatalf("unexpected error: %v", err)
113+
}
114+
115+
var events []internal.Event
116+
for e := range ch {
117+
events = append(events, e)
118+
}
119+
if len(events) != 1 {
120+
t.Fatalf("expected 1 event, got %d", len(events))
121+
}
122+
if events[0].Name != "FileSize" {
123+
t.Errorf("expected event name 'FileSize', got %q", events[0].Name)
124+
}
125+
if events[0].Message == "" {
126+
t.Error("expected event message to not be empty")
127+
}
128+
}
129+
130+
func TestFileSizeTool_RelativePath(t *testing.T) {
131+
dir := t.TempDir()
132+
filePath := dir + "/relative_test.txt"
133+
if err := os.WriteFile(filePath, []byte("content"), 0o600); err != nil {
134+
t.Fatalf("failed to write file: %v", err)
135+
}
136+
137+
// Change working directory so relative path resolves.
138+
orig, err := os.Getwd()
139+
if err != nil {
140+
t.Fatalf("getwd: %v", err)
141+
}
142+
t.Cleanup(func() { os.Chdir(orig) }) //nolint:errcheck
143+
if err := os.Chdir(dir); err != nil {
144+
t.Fatalf("chdir: %v", err)
145+
}
146+
147+
result, err := executeFileSize(t, FileSizeInput{FilePath: "relative_test.txt"})
148+
if err != nil {
149+
t.Fatalf("unexpected error: %v", err)
150+
}
151+
// "content" = 7 bytes
152+
if !strings.Contains(result.Content, "7") {
153+
t.Errorf("expected 7 bytes, got: %q", result.Content)
154+
}
155+
}
156+
157+
func TestFormatSize(t *testing.T) {
158+
tests := []struct {
159+
bytes int64
160+
expected string
161+
}{
162+
{0, "0 B"},
163+
{100, "100 B"},
164+
{1024, "1.00 KB"},
165+
{1024 * 1024, "1.00 MB"},
166+
{1024 * 1024 * 1024, "1.00 GB"},
167+
{1536, "1.50 KB"}, // 1.5 KB
168+
}
169+
170+
for _, tt := range tests {
171+
result := formatSize(tt.bytes)
172+
if result != tt.expected {
173+
t.Errorf("formatSize(%d) = %q, want %q", tt.bytes, result, tt.expected)
174+
}
175+
}
176+
}

0 commit comments

Comments
 (0)