rony-chat-bot/internal/portfolio/indexer.go

// Package portfolio loads and indexes the user's project markdowns for RAG.
//
// Schema decisions are documented in docs/architecture.md §4.0 (driver) and
// §4 (tokenizer, chunking). The tokenize/driver choices are validated by
// ./bench/; chunking was changed from size-based to heading-based after
// inspecting real data/templates in data/projects/.
package portfolio

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"time"

	_ "modernc.org/sqlite"
)

// SearchResult is one hit from the FTS5 index, with a relevance score.
type SearchResult struct {
	ID         string
	ProjectID  string
	SourceFile string
	Section    string
	Index      int
	Content    string
	Score      float64
}

const schema = `
CREATE VIRTUAL TABLE IF NOT EXISTS portfolio_chunks USING fts5(
    id UNINDEXED,
    project_id UNINDEXED,
    source_file UNINDEXED,
    section UNINDEXED,
    chunk_index UNINDEXED,
    content,
    tokenize = 'unicode61 remove_diacritics 2'
);
`

// Store wraps a SQLite FTS5 database with the portfolio schema.
type Store struct {
	db        *sql.DB
	chunkSize int // legacy field kept for compat; not used by heading chunker
}

func OpenStore(dbPath string) (*Store, error) {
	dir := filepath.Dir(dbPath)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, fmt.Errorf("create db dir: %w", err)
	}
	dsn := dbPath + "?_pragma=journal_mode(WAL)&_pragma=synchronous(NORMAL)&_pragma=foreign_keys(1)"
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, fmt.Errorf("open sqlite: %w", err)
	}
	db.SetMaxOpenConns(1) // SQLite + concurrent writers doesn't help
	if _, err := db.ExecContext(context.Background(), schema); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("create schema: %w", err)
	}
	if _, err := db.ExecContext(context.Background(), conversationSchema); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("create conversation schema: %w", err)
	}
	return &Store{db: db, chunkSize: 500}, nil
}

func (s *Store) Close() error { return s.db.Close() }

// DB exposes the underlying *sql.DB for callers that need to run their own
// queries (e.g. the health check). Don't use for hot-path code: go through
// the Search / Reindex methods.
func (s *Store) DB() *sql.DB { return s.db }
func (s *Store) Reindex(ctx context.Context, dataPath string, cfg ChunkerConfig) (files, chunks int, err error) {
	matches, err := filepath.Glob(filepath.Join(dataPath, "*.md"))
	if err != nil {
		return 0, 0, fmt.Errorf("glob: %w", err)
	}

	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return 0, 0, err
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx, `DELETE FROM portfolio_chunks`); err != nil {
		return 0, 0, fmt.Errorf("clear: %w", err)
	}

	stmt, err := tx.PrepareContext(ctx,
		`INSERT INTO portfolio_chunks (id, project_id, source_file, section, chunk_index, content) VALUES (?,?,?,?,?,?)`)
	if err != nil {
		return 0, 0, err
	}
	defer stmt.Close()

	for _, file := range matches {
		body, err := os.ReadFile(file)
		if err != nil {
			slog.Warn("read file failed", "file", file, "err", err)
			continue
		}
		projectID := strings.TrimSuffix(filepath.Base(file), ".md")
		sections := SplitMarkdownSections(string(body), cfg)
		for idx, sec := range sections {
			id := fmt.Sprintf("%s-%s-%d", projectID, slugify(sec.Heading), idx)
			if _, err := stmt.ExecContext(ctx, id, projectID, file, sec.Heading, idx, sec.Body); err != nil {
				return len(matches), chunks, fmt.Errorf("insert %s: %w", id, err)
			}
			chunks++
		}
		files++
	}

	if err := tx.Commit(); err != nil {
		return 0, 0, err
	}
	return files, chunks, nil
}

func slugify(s string) string {
	out := make([]byte, 0, len(s))
	for i := 0; i < len(s); i++ {
		c := s[i]
		switch {
		case c >= 'a' && c <= 'z', c >= '0' && c <= '9':
			out = append(out, c)
		case c >= 'A' && c <= 'Z':
			out = append(out, c+32)
		case c == ' ' || c == '-' || c == '_':
			out = append(out, '_')
		}
	}
	return string(out)
}

// Search returns up to topK chunks ordered by BM25 score.
func (s *Store) Search(ctx context.Context, query string, topK int) ([]SearchResult, error) {
	if topK <= 0 {
		topK = 5
	}
	rows, err := s.db.QueryContext(ctx, fmt.Sprintf(`
		SELECT id, project_id, source_file, section, chunk_index, content, bm25(portfolio_chunks) AS score
		FROM portfolio_chunks
		WHERE portfolio_chunks MATCH '%s'
		ORDER BY score
		LIMIT %d
	`, sanitizeFTS5(query), topK))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var hits []SearchResult
	for rows.Next() {
		var r SearchResult
		if err := rows.Scan(&r.ID, &r.ProjectID, &r.SourceFile, &r.Section, &r.Index, &r.Content, &r.Score); err != nil {
			return nil, err
		}
		hits = append(hits, r)
	}
	return hits, rows.Err()
}

// sanitizeFTS5 escapes special chars, adds prefix-match wildcards, and joins
// tokens with OR (Q&A behavior: "what database" should match a doc that
// contains "database" even when it doesn't contain "what"). FTS5 doesn't
// accept `?` placeholders for MATCH in driver-prepared statements; this
// inlines the escaped query.
func sanitizeFTS5(q string) string {
	tokens := strings.FieldsFunc(strings.ToLower(q), func(r rune) bool {
		return !(r == '-' || r == '_' || r == '.' || r == '+' ||
			(r >= '0' && r <= '9') ||
			(r >= 'a' && r <= 'z') ||
			r > 0x7F)
	})
	if len(tokens) == 0 {
		return `""`
	}
	// Drop a tiny stopword list so "what is" doesn't dominate the OR
	// (these match every doc and dilute BM25 ranking).
	keep := tokens[:0]
	for _, t := range tokens {
		switch t {
		case "what", "which", "who", "how", "when", "where", "is", "are",
			"do", "does", "can", "tell", "about", "the", "a", "an":
			continue
		}
		keep = append(keep, t)
	}
	if len(keep) == 0 {
		// All tokens were stopwords — fall back to the original set.
		keep = tokens
	}
	for i, t := range keep {
		keep[i] = `"` + t + `"*`
	}
	return strings.Join(keep, " OR ")
}

// ReindexOnDisk is a small convenience that opens the store, reindexes, and
// closes — used by the CLI subcommand.
func ReindexOnDisk(dbPath, dataPath string, cfg ChunkerConfig) (time.Duration, int, int, error) {
	start := time.Now()
	store, err := OpenStore(dbPath)
	if err != nil {
		return 0, 0, 0, err
	}
	defer store.Close()
	files, chunks, err := store.Reindex(context.Background(), dataPath, cfg)
	return time.Since(start), files, chunks, err
}