feat: add --compress for tree-sitter (#57)

Author: sammcj

.github/workflows/build-release-publish.yml

@@ -11,7 +11,7 @@ on:
       - 'README.md'
 
 env:
-  GO_VERSION: '1.22'
+  GO_VERSION: '1.24.3' # Updated to match toolchain
   BINARY_NAME: 'ingest'
 
 permissions:
@@ -22,16 +22,22 @@ jobs:
   build:
     if: ${{ ! contains(github.event.head_commit.message, '[skip ci]') && ! contains(github.event.pull_request.title, '[skip ci]')}}
     name: Build
-    runs-on: ubuntu-latest
     strategy:
       matrix:
-        include:
+        target:
           - os: darwin
             arch: arm64
+            runner: macos-14
+            c_compiler_package: ""
           - os: linux
             arch: amd64
-          - os: linux
-            arch: arm64
+            runner: ubuntu-latest
+            c_compiler_package: "build-essential"
+          # - os: linux
+          #   arch: arm64
+          #   runner: ubuntu-latest-arm64 # Use native ARM64 runner
+          #   c_compiler_package: "build-essential" # Native compiler
+    runs-on: ${{ matrix.target.runner }}
 
     outputs:
       version: ${{ steps.set_version.outputs.new_tag }}
@@ -58,25 +64,32 @@ jobs:
       - name: Get dependencies
         run: go mod download
 
+      - name: Set up C compiler
+        if: startsWith(matrix.target.runner, 'ubuntu') && matrix.target.c_compiler_package != ''
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y ${{ matrix.target.c_compiler_package }}
+
       - name: Run tests
         run: go test -v ./...
 
       - name: Build
         env:
-          GOOS: ${{ matrix.os }}
-          GOARCH: ${{ matrix.arch }}
+          CGO_ENABLED: "1" # Explicitly enable CGo
+          GOOS: ${{ matrix.target.os }}
+          GOARCH: ${{ matrix.target.arch }}
           VERSION: ${{ steps.set_version.outputs.new_tag }}
 
         run: |
-          go build -v -ldflags "-X main.Version=$version" -o build/${{ env.BINARY_NAME }}-${{ matrix.os }}-${{ matrix.arch }} .
+          go build -v -ldflags "-X main.Version=$VERSION" -o build/${{ env.BINARY_NAME }}-${{ matrix.target.os }}-${{ matrix.target.arch }} .
           ls -ltarh build/
 
       - name: Upload artifact
         uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4
         if: github.event_name == 'push' && github.ref == 'refs/heads/main'
         with:
-          name: ${{ env.BINARY_NAME }}-${{ matrix.os }}-${{ matrix.arch }}
-          path: build/${{ env.BINARY_NAME }}-${{ matrix.os }}-${{ matrix.arch }}
+          name: ${{ env.BINARY_NAME }}-${{ matrix.target.os }}-${{ matrix.target.arch }}
+          path: build/${{ env.BINARY_NAME }}-${{ matrix.target.os }}-${{ matrix.target.arch }}
           retention-days: 90
 
   release:

Makefile

@@ -29,6 +29,9 @@ clean:
 	$(GOCLEAN)
 	rm -f $(BINARY_NAME)
 
+lint:
+	gofmt -s -w .
+
 test:
 	$(GOTEST) -v ./...
 

README.md

@@ -20,6 +20,7 @@ And ingest web URLs.
 
 - Traverse directory structures and generate a tree view
 - Include/exclude files based on glob patterns
+- Compress code using Tree-sitter to extract key structural information while omitting implementation details
 - Estimate vRAM requirements and check model compatibility using another package I've created called [quantest](https://github.com/sammcj/quantest)
 - Parse output directly to LLMs such as Ollama or any OpenAI compatible API for processing
 - Generate and include git diffs and logs
@@ -209,6 +210,44 @@ You can provide a prompt suffix to append to the generated prompt:
 ingest --llm -p "explain this code" /path/to/project
 ```
 
+## Code Compression with Tree-sitter
+
+**Experimental**
+
+Ingest can compress source code files by extracting key structural information while omitting implementation details. This is useful for reducing token usage while preserving the important parts of the code structure.
+
+```shell
+ingest --compress /path/to/project
+```
+
+The compression extracts:
+- Package/module declarations
+- Import statements
+- Function/method signatures (without bodies)
+- Class definitions (without method bodies)
+- Type definitions
+- Comments
+
+Currently supported languages:
+- Go
+- Python
+- JavaScript (including arrow functions and ES6 module syntax)
+- Bash
+- C
+- CSS
+
+Example of compressed JavaScript:
+
+```
+// This is a JavaScript comment
+import { something } from 'module';
+export class MyJSClass { ... } // Body removed
+constructor(name) { ... } // Body removed
+greet(message) { ... } // Body removed
+export function myJSFunction(x, y) { ... } // Body removed
+const myArrowFunc = (a, b) => { ... } // Body removed
+```
+
 ## Web Crawling & Ingestion
 
 Crawl with explicit web mode
@@ -272,6 +311,7 @@ These directories will be created automatically on first run, along with README
 
 ### Flags
 
+- `--compress`: **New** Enable code compression using Tree-sitter to extract key structural information while omitting implementation details
 - `--config`: Opens the config file in the default editor
 - `--context`: Specify the context length for VRAM estimation
 - `--exclude-from-tree`: Exclude files/folders from the source tree based on exclude patterns

config/config.go

@@ -18,9 +18,9 @@ type OllamaConfig struct {
 }
 
 type Config struct {
-	Ollama []OllamaConfig `json:"ollama"`
-	LLM    LLMConfig      `json:"llm"`
-	AutoSave bool         `json:"auto_save"`
+	Ollama   []OllamaConfig `json:"ollama"`
+	LLM      LLMConfig      `json:"llm"`
+	AutoSave bool           `json:"auto_save"`
 }
 
 type LLMConfig struct {

filesystem/filesystem.go

@@ -16,6 +16,7 @@ import (
 	"github.com/fatih/color"
 	"github.com/mitchellh/go-homedir"
 	ignore "github.com/sabhiram/go-gitignore"
+	"github.com/sammcj/ingest/internal/compressor"
 	"github.com/sammcj/ingest/pdf"
 	"github.com/sammcj/ingest/utils"
 )
@@ -28,10 +29,10 @@ type FileInfo struct {
 
 // New type to track excluded files and directories
 type ExcludedInfo struct {
-	Directories map[string]int    // Directory path -> count of excluded files
-	Extensions  map[string]int    // File extension -> count of excluded files
-	TotalFiles  int              // Total number of excluded files
-	Files       []string         // List of excluded files (if total ≤ 20)
+	Directories map[string]int // Directory path -> count of excluded files
+	Extensions  map[string]int // File extension -> count of excluded files
+	TotalFiles  int            // Total number of excluded files
+	Files       []string       // List of excluded files (if total ≤ 20)
 }
 
 type treeNode struct {
@@ -67,9 +68,9 @@ func ReadExcludePatterns(patternExclude string, noDefaultExcludes bool) ([]strin
 		// If user has a default.glob, it overrides the default patterns
 		if _, err := os.Stat(userDefaultGlob); err == nil {
 			return readGlobFile(userDefaultGlob)
-	}
+		}
 
-	// Read other user-defined patterns
+		// Read other user-defined patterns
 		userPatterns, _ := readGlobFilesFromDir(userPatternsDir)
 
 		// Combine user patterns with default patterns (if not disabled)
@@ -102,7 +103,6 @@ func trackExcludedFile(excluded *ExcludedInfo, path string, mu *sync.Mutex) {
 	}
 }
 
-
 func readGlobFile(filename string) ([]string, error) {
 	file, err := os.Open(filename)
 	if err != nil {
@@ -150,15 +150,15 @@ func trackExcludedDirectory(excluded *ExcludedInfo, path string, mu *sync.Mutex)
 	excluded.Directories[path] = 0 // Initialize directory count
 }
 
-func WalkDirectory(rootPath string, includePatterns, excludePatterns []string, patternExclude string, includePriority, lineNumber, relativePaths, excludeFromTree, noCodeblock, noDefaultExcludes bool) (string, []FileInfo, *ExcludedInfo, error) {
+func WalkDirectory(rootPath string, includePatterns, excludePatterns []string, patternExclude string, includePriority, lineNumber, relativePaths, excludeFromTree, noCodeblock, noDefaultExcludes bool, comp *compressor.GenericCompressor) (string, []FileInfo, *ExcludedInfo, error) {
 	var files []FileInfo
 	var mu sync.Mutex
 	var wg sync.WaitGroup
 
 	excluded := &ExcludedInfo{
 		Directories: make(map[string]int),
 		Extensions:  make(map[string]int),
-		Files:      make([]string, 0),
+		Files:       make([]string, 0),
 	}
 
 	// Read exclude patterns
@@ -216,7 +216,7 @@ func WalkDirectory(rootPath string, includePatterns, excludePatterns []string, p
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
-				processFile(rootPath, relPath, filepath.Dir(rootPath), lineNumber, relativePaths, noCodeblock, &mu, &files)
+				processFile(rootPath, relPath, filepath.Dir(rootPath), lineNumber, relativePaths, noCodeblock, &mu, &files, comp)
 			}()
 		} else {
 			trackExcludedFile(excluded, rootPath, &mu)
@@ -259,7 +259,7 @@ func WalkDirectory(rootPath string, includePatterns, excludePatterns []string, p
 				wg.Add(1)
 				go func(path, relPath string, info os.FileInfo) {
 					defer wg.Done()
-					processFile(path, relPath, rootPath, lineNumber, relativePaths, noCodeblock, &mu, &files)
+					processFile(path, relPath, rootPath, lineNumber, relativePaths, noCodeblock, &mu, &files, comp)
 				}(path, relPath, info)
 			}
 
@@ -281,7 +281,7 @@ func shouldExcludePath(path string, excludePatterns []string, gitignore *ignore.
 	for _, pattern := range excludePatterns {
 		if match, _ := doublestar.Match(pattern, path); match {
 			return true
-	}
+		}
 	}
 	return gitignore != nil && gitignore.MatchesPath(path)
 }
@@ -353,7 +353,7 @@ func isBinaryFile(filePath string) (bool, error) {
 	n, err := file.Read(buffer)
 	if err != nil && err != io.EOF {
 		return false, err
-}
+	}
 
 	// Use http.DetectContentType to determine the content type
 	contentType := http.DetectContentType(buffer[:n])
@@ -371,7 +371,7 @@ func PrintDefaultExcludes() {
 	fmt.Println(strings.Join(excludes, "\n"))
 }
 
-func processFile(path, relPath string, rootPath string, lineNumber, relativePaths, noCodeblock bool, mu *sync.Mutex, files *[]FileInfo) {
+func processFile(path, relPath string, rootPath string, lineNumber, relativePaths, noCodeblock bool, mu *sync.Mutex, files *[]FileInfo, comp *compressor.GenericCompressor) {
 	// Check if it's the root path being processed (explicitly provided file)
 	isExplicitFile := path == rootPath
 
@@ -428,13 +428,36 @@ func processFile(path, relPath string, rootPath string, lineNumber, relativePath
 	}
 
 	code := string(content)
+
+	// Attempt compression if compressor is provided and it's not a PDF
+	if comp != nil && !isPDF {
+		langID, err := compressor.IdentifyLanguage(path)
+		if err == nil { // Language identified
+			compressedCode, err := comp.Compress(content, langID)
+			if err == nil {
+				code = compressedCode
+				// If compressed, we might not want to add line numbers or wrap in a generic code block
+				// as the compressor might handle formatting. For now, let's assume compressed output
+				// is final for this file's content.
+				// We'll skip line numbering and code block wrapping for compressed content.
+				goto skipFormatting
+			} else {
+				utils.PrintColouredMessage("⚠️", fmt.Sprintf("Compression failed for %s: %v. Using original content.", path, err), color.FgYellow)
+			}
+		} else {
+			// Language not identified for compression, use original content
+			utils.PrintColouredMessage("ℹ️", fmt.Sprintf("Language not identified for compression for %s. Using original content.", path), color.FgBlue)
+		}
+	}
+
 	if lineNumber {
 		code = addLineNumbers(code)
 	}
 	if !noCodeblock {
 		code = wrapCodeBlock(code, filepath.Ext(path))
 	}
 
+skipFormatting:
 	filePath := path
 	if relativePaths {
 		filePath = filepath.Join(filepath.Base(rootPath), relPath)
@@ -487,9 +510,9 @@ func generateTreeString(rootPath string, excludePatterns []string) (string, erro
 					}
 					if !found {
 						newNode := &treeNode{
-							name:      part,
-							isDir:     true,
-							excluded:  true,
+							name:     part,
+							isDir:    true,
+							excluded: true,
 						}
 						current.children = append(current.children, newNode)
 						current = newNode
@@ -608,7 +631,7 @@ func isExcluded(path string, patterns []string) bool {
 	return false
 }
 
-func ProcessSingleFile(path string, lineNumber, relativePaths, noCodeblock bool) (FileInfo, error) {
+func ProcessSingleFile(path string, lineNumber, relativePaths, noCodeblock bool, comp *compressor.GenericCompressor) (FileInfo, error) {
 	// Check if it's a PDF first
 	isPDF, err := pdf.IsPDF(path)
 	if err != nil {
@@ -635,13 +658,32 @@ func ProcessSingleFile(path string, lineNumber, relativePaths, noCodeblock bool)
 	}
 
 	code := string(content)
+
+	// Attempt compression if compressor is provided and it's not a PDF
+	if comp != nil && !isPDF {
+		langID, err := compressor.IdentifyLanguage(path)
+		if err == nil { // Language identified
+			compressedCode, err := comp.Compress(content, langID)
+			if err == nil {
+				code = compressedCode
+				// Skip standard formatting for compressed content
+				goto skipSingleFileFormatting
+			} else {
+				utils.PrintColouredMessage("⚠️", fmt.Sprintf("Compression failed for %s: %v. Using original content.", path, err), color.FgYellow)
+			}
+		} else {
+			utils.PrintColouredMessage("ℹ️", fmt.Sprintf("Language not identified for compression for %s. Using original content.", path), color.FgBlue)
+		}
+	}
+
 	if lineNumber {
 		code = addLineNumbers(code)
 	}
 	if !noCodeblock {
 		code = wrapCodeBlock(code, filepath.Ext(path))
 	}
 
+skipSingleFileFormatting:
 	filePath := path
 	if relativePaths {
 		filePath = filepath.Base(path)