356 lines
11 KiB
Go
356 lines
11 KiB
Go
package workspace
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"git.gobha.me/xcaliber/chat-switchboard/knowledge"
|
|
"git.gobha.me/xcaliber/chat-switchboard/models"
|
|
"git.gobha.me/xcaliber/chat-switchboard/notifications"
|
|
"git.gobha.me/xcaliber/chat-switchboard/store"
|
|
)
|
|
|
|
// ── Configuration ────────────────────────────
|
|
|
|
const (
|
|
// indexTimeout is the max time for indexing a single file.
|
|
indexTimeout = 5 * time.Minute
|
|
|
|
// maxIndexTextLength caps text extracted from a single file (5 MB).
|
|
maxIndexTextLength = 5 * 1024 * 1024
|
|
)
|
|
|
|
// indexableExtensions defines source code and text file extensions
|
|
// eligible for chunking and embedding. MIME detection covers text/*,
|
|
// but code files often have application/* types — this map catches them.
|
|
var indexableExtensions = map[string]bool{
|
|
".go": true, ".py": true, ".js": true, ".ts": true, ".jsx": true,
|
|
".tsx": true, ".rs": true, ".c": true, ".cpp": true, ".h": true,
|
|
".hpp": true, ".java": true, ".rb": true, ".php": true, ".swift": true,
|
|
".kt": true, ".scala": true, ".sh": true, ".bash": true, ".zsh": true,
|
|
".sql": true, ".yaml": true, ".yml": true, ".toml": true, ".json": true,
|
|
".xml": true, ".css": true, ".scss": true, ".less": true,
|
|
".md": true, ".txt": true, ".csv": true, ".html": true, ".htm": true,
|
|
".dockerfile": true, ".tf": true, ".hcl": true, ".proto": true,
|
|
".graphql": true, ".vue": true, ".svelte": true,
|
|
".pl": true, ".pm": true, // Perl
|
|
".lua": true, ".zig": true, ".nim": true, ".ex": true, ".exs": true,
|
|
// Config / data
|
|
".ini": true, ".cfg": true, ".conf": true, ".env": true,
|
|
".mk": true, ".cmake": true, ".makefile": true,
|
|
".r": true, ".jl": true, // R, Julia
|
|
}
|
|
|
|
// CodeChunkConfig returns chunking settings optimized for source code.
|
|
// Larger chunks (1500 chars) because functions can be long, with
|
|
// separators tuned to split on function/class/type boundaries.
|
|
func CodeChunkConfig() knowledge.ChunkConfig {
|
|
return knowledge.ChunkConfig{
|
|
ChunkSize: 1500,
|
|
ChunkOverlap: 200,
|
|
Separators: []string{
|
|
"\n\nfunc ", "\n\nfn ", // Go, Rust
|
|
"\n\nclass ", "\n\ndef ", // Python, Ruby
|
|
"\n\ntype ", "\n\nsub ", // Go types, Perl
|
|
"\n\npub ", "\n\nimpl ", // Rust
|
|
"\n\n", "\n", " ", // fallback hierarchy
|
|
},
|
|
}
|
|
}
|
|
|
|
// isIndexable returns true if the file is eligible for text indexing.
|
|
func isIndexable(path string, contentType string) bool {
|
|
// Check extension first (catches source code with non-text MIME types)
|
|
ext := strings.ToLower(filepath.Ext(path))
|
|
if indexableExtensions[ext] {
|
|
return true
|
|
}
|
|
// Also index any text/* content type
|
|
if strings.HasPrefix(contentType, "text/") {
|
|
return true
|
|
}
|
|
// Specific filenames without extension
|
|
base := strings.ToLower(filepath.Base(path))
|
|
switch base {
|
|
case "makefile", "dockerfile", "jenkinsfile", "gemfile",
|
|
"rakefile", "procfile", "vagrantfile", ".gitignore",
|
|
".dockerignore", ".editorconfig":
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// isCodeFile returns true if the file is a source code file (uses code chunking).
|
|
func isCodeFile(path string) bool {
|
|
ext := strings.ToLower(filepath.Ext(path))
|
|
switch ext {
|
|
case ".md", ".txt", ".csv", ".html", ".htm", ".xml", ".json", ".yaml", ".yml":
|
|
return false // prose / data — use default chunking
|
|
}
|
|
return indexableExtensions[ext]
|
|
}
|
|
|
|
// ── Indexer ──────────────────────────────────
|
|
|
|
// Indexer processes workspace files through the chunking + embedding pipeline.
|
|
// It reuses the knowledge.Embedder for vector generation and shares the
|
|
// concurrency semaphore with KB ingestion to avoid overwhelming the
|
|
// embedding provider.
|
|
type Indexer struct {
|
|
stores store.Stores
|
|
embedder *knowledge.Embedder
|
|
wfs *FS
|
|
sem chan struct{} // shared semaphore with KB ingestion
|
|
wg sync.WaitGroup
|
|
enabled bool // global kill switch (WORKSPACE_INDEXING_ENABLED)
|
|
}
|
|
|
|
// NewIndexer creates a workspace indexer.
|
|
// The semaphore should be shared with the KB Ingester to cap total
|
|
// embedding concurrency.
|
|
func NewIndexer(stores store.Stores, embedder *knowledge.Embedder, wfs *FS, sem chan struct{}, enabled bool) *Indexer {
|
|
return &Indexer{
|
|
stores: stores,
|
|
embedder: embedder,
|
|
wfs: wfs,
|
|
sem: sem,
|
|
enabled: enabled,
|
|
}
|
|
}
|
|
|
|
// IsEnabled returns true if workspace indexing is globally enabled
|
|
// AND the embedding role is configured.
|
|
func (idx *Indexer) IsEnabled(ctx context.Context) bool {
|
|
return idx.enabled && idx.embedder.IsConfigured(ctx)
|
|
}
|
|
|
|
// IndexFile starts async indexing for a single workspace file.
|
|
// Called from WorkspaceFS.WriteFile after a successful write.
|
|
// No-op if indexing is disabled, the file is non-indexable, or the
|
|
// sha256 hasn't changed (content-addressed skip).
|
|
func (idx *Indexer) IndexFile(w *models.Workspace, f *models.WorkspaceFile, userID string, teamID *string) {
|
|
if !idx.enabled || !w.IndexingEnabled {
|
|
return
|
|
}
|
|
|
|
// Check if indexable
|
|
if f.IsDirectory || !isIndexable(f.Path, f.ContentType) {
|
|
// Mark as skipped so the UI knows it's intentional
|
|
ctx := context.Background()
|
|
idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusSkipped, 0)
|
|
return
|
|
}
|
|
|
|
// Mark as pending
|
|
ctx := context.Background()
|
|
idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusPending, 0)
|
|
|
|
idx.wg.Add(1)
|
|
go func() {
|
|
defer idx.wg.Done()
|
|
|
|
// Acquire semaphore
|
|
idx.sem <- struct{}{}
|
|
defer func() { <-idx.sem }()
|
|
|
|
ictx, cancel := context.WithTimeout(context.Background(), indexTimeout)
|
|
defer cancel()
|
|
|
|
if err := idx.indexFile(ictx, w, f, userID, teamID); err != nil {
|
|
log.Printf("❌ workspace index %s/%s: %v", w.ID[:8], f.Path, err)
|
|
idx.stores.Workspaces.UpdateFileIndexStatus(ictx, f.ID, models.IndexStatusError, 0)
|
|
}
|
|
}()
|
|
}
|
|
|
|
// IndexBatch starts async indexing for multiple files (e.g. after archive extract).
|
|
// Files are indexed sequentially within a single goroutine to be friendly
|
|
// to rate limits.
|
|
func (idx *Indexer) IndexBatch(w *models.Workspace, files []models.WorkspaceFile, userID string, teamID *string) {
|
|
if !idx.enabled || !w.IndexingEnabled {
|
|
return
|
|
}
|
|
|
|
// Filter to indexable files only
|
|
var indexable []models.WorkspaceFile
|
|
ctx := context.Background()
|
|
for _, f := range files {
|
|
if f.IsDirectory || !isIndexable(f.Path, f.ContentType) {
|
|
idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusSkipped, 0)
|
|
continue
|
|
}
|
|
idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusPending, 0)
|
|
indexable = append(indexable, f)
|
|
}
|
|
|
|
if len(indexable) == 0 {
|
|
return
|
|
}
|
|
|
|
idx.wg.Add(1)
|
|
go func() {
|
|
defer idx.wg.Done()
|
|
|
|
// One semaphore slot for the entire batch
|
|
idx.sem <- struct{}{}
|
|
defer func() { <-idx.sem }()
|
|
|
|
var indexed, errors int
|
|
for _, f := range indexable {
|
|
f := f
|
|
ictx, cancel := context.WithTimeout(context.Background(), indexTimeout)
|
|
if err := idx.indexFile(ictx, w, &f, userID, teamID); err != nil {
|
|
log.Printf("❌ workspace index %s/%s: %v", w.ID[:8], f.Path, err)
|
|
idx.stores.Workspaces.UpdateFileIndexStatus(ictx, f.ID, models.IndexStatusError, 0)
|
|
errors++
|
|
} else {
|
|
indexed++
|
|
}
|
|
cancel()
|
|
}
|
|
|
|
log.Printf("📑 workspace batch index %s: %d indexed, %d errors, %d skipped",
|
|
w.ID[:8], indexed, errors, len(files)-len(indexable))
|
|
|
|
// Notify on completion
|
|
if svc := notifications.Default(); svc != nil && indexed > 0 {
|
|
// Find workspace owner for notification target
|
|
if w.OwnerType == models.WorkspaceOwnerUser {
|
|
svc.Notify(context.Background(), &models.Notification{
|
|
UserID: w.OwnerID,
|
|
Type: "workspace.indexed",
|
|
Title: "Workspace indexed",
|
|
Body: fmt.Sprintf("%d files indexed in workspace %q", indexed, w.Name),
|
|
ResourceType: "workspace",
|
|
ResourceID: w.ID,
|
|
})
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
// Wait blocks until all in-flight indexing goroutines complete.
|
|
func (idx *Indexer) Wait() {
|
|
idx.wg.Wait()
|
|
}
|
|
|
|
// ── Internal Pipeline ────────────────────────
|
|
|
|
func (idx *Indexer) indexFile(ctx context.Context, w *models.Workspace, f *models.WorkspaceFile, userID string, teamID *string) error {
|
|
start := time.Now()
|
|
|
|
// Mark as indexing
|
|
idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusIndexing, 0)
|
|
|
|
// Check if embedding is configured (might have changed since enqueue)
|
|
if !idx.embedder.IsConfigured(ctx) {
|
|
return fmt.Errorf("embedding role not configured")
|
|
}
|
|
|
|
// ── Step 1: Read file content ───────────
|
|
rc, _, err := idx.wfs.ReadFile(ctx, w, f.Path)
|
|
if err != nil {
|
|
return fmt.Errorf("read: %w", err)
|
|
}
|
|
defer rc.Close()
|
|
|
|
// Cap at max text length
|
|
limited := io.LimitReader(rc, maxIndexTextLength)
|
|
data, err := io.ReadAll(limited)
|
|
if err != nil {
|
|
return fmt.Errorf("read content: %w", err)
|
|
}
|
|
text := string(data)
|
|
|
|
if strings.TrimSpace(text) == "" {
|
|
idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusSkipped, 0)
|
|
return nil
|
|
}
|
|
|
|
// ── Step 2: Chunk ───────────────────────
|
|
var cfg knowledge.ChunkConfig
|
|
if isCodeFile(f.Path) {
|
|
cfg = CodeChunkConfig()
|
|
} else {
|
|
cfg = knowledge.DefaultChunkConfig()
|
|
}
|
|
|
|
chunks := knowledge.SplitText(text, cfg)
|
|
if len(chunks) == 0 {
|
|
idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusSkipped, 0)
|
|
return nil
|
|
}
|
|
|
|
// ── Step 3: Embed ───────────────────────
|
|
texts := make([]string, len(chunks))
|
|
for i, c := range chunks {
|
|
texts[i] = c.Content
|
|
}
|
|
|
|
embedResult, err := idx.embedder.EmbedChunks(ctx, userID, teamID, texts)
|
|
if err != nil {
|
|
return fmt.Errorf("embed: %w", err)
|
|
}
|
|
|
|
// Track usage
|
|
idx.embedder.LogUsage(ctx, userID, nil, embedResult)
|
|
|
|
// ── Step 4: Delete old chunks + insert new ─
|
|
if err := idx.stores.Workspaces.DeleteChunksByFile(ctx, f.ID); err != nil {
|
|
return fmt.Errorf("delete old chunks: %w", err)
|
|
}
|
|
|
|
dbChunks := make([]models.WorkspaceChunk, len(chunks))
|
|
for i, c := range chunks {
|
|
// Compute approximate line number from character offset
|
|
lineStart := 0
|
|
if i > 0 {
|
|
charOffset := i * (cfg.ChunkSize - cfg.ChunkOverlap)
|
|
if charOffset < len(text) {
|
|
lineStart = strings.Count(text[:charOffset], "\n") + 1
|
|
}
|
|
}
|
|
|
|
dbChunks[i] = models.WorkspaceChunk{
|
|
WorkspaceID: w.ID,
|
|
FileID: f.ID,
|
|
ChunkIndex: c.Index,
|
|
Content: c.Content,
|
|
TokenCount: c.TokenCount,
|
|
Embedding: embedResult.Vectors[i],
|
|
Metadata: models.JSONMap{
|
|
"line_start": lineStart,
|
|
"language": inferLanguage(f.Path),
|
|
},
|
|
}
|
|
}
|
|
|
|
if err := idx.stores.Workspaces.InsertChunks(ctx, dbChunks); err != nil {
|
|
return fmt.Errorf("store chunks: %w", err)
|
|
}
|
|
|
|
// ── Step 5: Update file index status ────
|
|
if err := idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusReady, len(chunks)); err != nil {
|
|
log.Printf("⚠ failed to update index status for %s: %v", f.Path, err)
|
|
}
|
|
|
|
log.Printf(" 📑 indexed %s → %d chunks (%s)", f.Path, len(chunks), time.Since(start).Round(time.Millisecond))
|
|
return nil
|
|
}
|
|
|
|
// inferLanguage extracts a simple language identifier from file extension.
|
|
func inferLanguage(path string) string {
|
|
ext := strings.ToLower(filepath.Ext(path))
|
|
if ext == "" {
|
|
return ""
|
|
}
|
|
// Strip leading dot
|
|
return ext[1:]
|
|
}
|