agentic-coding-workflow/internal/codeindex/chunker.go

package codeindex

import (
	"bytes"
	"crypto/sha256"
	"path/filepath"
	"strings"

	"github.com/yan1h/agent-coding-workflow/internal/infra/llm"
)

// Chunker config defaults. Windows are line-based; the token cap (estimated via
// llm.EstimateTokens) hard-bounds embedding cost on pathological long lines.
const (
	DefaultWindowLines    = 60
	DefaultOverlapLines   = 10
	DefaultMaxFileBytes   = 512 * 1024 // skip files larger than this
	DefaultMaxChunkTokens = 1024       // estimated; oversize windows are truncated
)

// Chunker splits file content into overlapping line windows.
type Chunker struct {
	WindowLines    int
	OverlapLines   int
	MaxFileBytes   int
	MaxChunkTokens int
}

// NewChunker returns a Chunker with platform defaults.
func NewChunker() *Chunker {
	return &Chunker{
		WindowLines:    DefaultWindowLines,
		OverlapLines:   DefaultOverlapLines,
		MaxFileBytes:   DefaultMaxFileBytes,
		MaxChunkTokens: DefaultMaxChunkTokens,
	}
}

func (c *Chunker) window() int {
	if c.WindowLines > 0 {
		return c.WindowLines
	}
	return DefaultWindowLines
}

func (c *Chunker) overlap() int {
	if c.OverlapLines >= 0 && c.OverlapLines < c.window() {
		return c.OverlapLines
	}
	return DefaultOverlapLines
}

func (c *Chunker) maxFileBytes() int {
	if c.MaxFileBytes > 0 {
		return c.MaxFileBytes
	}
	return DefaultMaxFileBytes
}

func (c *Chunker) maxChunkTokens() int {
	if c.MaxChunkTokens > 0 {
		return c.MaxChunkTokens
	}
	return DefaultMaxChunkTokens
}

// Skippable reports whether a file should not be indexed (binary, oversize, or
// vendored). Callers should check this before reading large files.
func (c *Chunker) Skippable(filePath string, size int, content []byte) bool {
	if size > c.maxFileBytes() {
		return true
	}
	if isVendored(filePath) {
		return true
	}
	if isBinary(content) {
		return true
	}
	return false
}

// Chunk splits content into overlapping line windows. Returns nil for skippable
// or empty input. ContentSHA is the sha256 of the chunk's exact content so the
// build runner can reuse already-embedded chunks across commits.
func (c *Chunker) Chunk(filePath string, content []byte) []Chunk {
	if len(content) == 0 {
		return nil
	}
	if c.Skippable(filePath, len(content), content) {
		return nil
	}
	lines := strings.Split(string(content), "\n")
	// Drop a trailing empty element from a final newline so end_line is accurate.
	if n := len(lines); n > 0 && lines[n-1] == "" {
		lines = lines[:n-1]
	}
	if len(lines) == 0 {
		return nil
	}
	lang := langForPath(filePath)
	window := c.window()
	step := window - c.overlap()
	if step <= 0 {
		step = window
	}
	maxTok := c.maxChunkTokens()

	var out []Chunk
	for start := 0; start < len(lines); start += step {
		end := start + window
		if end > len(lines) {
			end = len(lines)
		}
		body := strings.Join(lines[start:end], "\n")
		// Hard token cap: truncate runaway windows (minified/one-line files).
		if llm.EstimateTokens(body) > maxTok {
			body = truncateToTokens(body, maxTok)
		}
		sum := sha256.Sum256([]byte(body))
		out = append(out, Chunk{
			StartLine:  start + 1, // 1-based, inclusive
			EndLine:    end,       // 1-based, inclusive
			Lang:       lang,
			Content:    body,
			ContentSHA: sum[:],
		})
		if end == len(lines) {
			break
		}
	}
	return out
}

// truncateToTokens cuts s down to roughly maxTok estimated tokens (4 chars/token).
func truncateToTokens(s string, maxTok int) string {
	maxBytes := maxTok * 4
	if len(s) <= maxBytes {
		return s
	}
	// Cut on a UTF-8 boundary.
	for maxBytes > 0 && (s[maxBytes]&0xC0) == 0x80 {
		maxBytes--
	}
	return s[:maxBytes]
}

// isBinary reports whether content looks binary (NUL byte in the first 8KB).
func isBinary(content []byte) bool {
	n := len(content)
	if n > 8192 {
		n = 8192
	}
	return bytes.IndexByte(content[:n], 0) >= 0
}

// vendoredDirs are path segments whose contents are excluded from the index.
var vendoredDirs = []string{
	"vendor/", "node_modules/", ".git/", "dist/", "build/", "third_party/",
	"testdata/", ".venv/", "__pycache__/",
}

func isVendored(p string) bool {
	norm := filepath.ToSlash(p)
	for _, d := range vendoredDirs {
		if strings.HasPrefix(norm, d) || strings.Contains(norm, "/"+d) {
			return true
		}
	}
	return false
}

// langForPath maps a file extension to a coarse language label (best-effort).
func langForPath(p string) string {
	switch strings.ToLower(filepath.Ext(p)) {
	case ".go":
		return "go"
	case ".ts", ".tsx":
		return "typescript"
	case ".js", ".jsx", ".mjs", ".cjs":
		return "javascript"
	case ".vue":
		return "vue"
	case ".py":
		return "python"
	case ".rs":
		return "rust"
	case ".java":
		return "java"
	case ".c", ".h":
		return "c"
	case ".cpp", ".cc", ".hpp":
		return "cpp"
	case ".rb":
		return "ruby"
	case ".php":
		return "php"
	case ".sql":
		return "sql"
	case ".sh", ".bash":
		return "shell"
	case ".md", ".markdown":
		return "markdown"
	case ".json":
		return "json"
	case ".yaml", ".yml":
		return "yaml"
	case ".toml":
		return "toml"
	case ".html", ".htm":
		return "html"
	case ".css", ".scss":
		return "css"
	default:
		return ""
	}
}