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
	Kind       string // KindProject | KindDoc
	SourceFile string
	Section    string
	Index      int
	Content    string
	Score      float64
}

// Chunk kinds. A "project" is a piece of Victor's portfolio and shows up in
// the catalogue the bot injects into the prompt; a "doc" (his CV, an about
// page, a FAQ) is retrievable evidence that is not itself a project and must
// never be listed as one.
const (
	KindProject = "project"
	KindDoc     = "doc"
)

// Source is one directory of markdown to index, and what the documents in it
// mean. See KindProject / KindDoc.
type Source struct {
	Path string
	Kind string
}

// SourcesFor is the standard mapping from the two configured directories to
// index sources. docsPath may be empty.
func SourcesFor(dataPath, docsPath string) []Source {
	return []Source{
		{Path: dataPath, Kind: KindProject},
		{Path: docsPath, Kind: KindDoc},
	}
}

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

// Vectors live in an ordinary table keyed by chunk id. There is no ANN index:
// a portfolio is hundreds of chunks, not millions, so a full scan with a dot
// product is microseconds and needs no extension.
// content_hash is what makes a vector verifiable. A chunk's id is derived
// from file + heading + position, so editing the *body* of a section leaves
// the id untouched — and a vector keyed only by id would keep describing the
// text that used to be there. Nothing errors; semantic ranking just silently
// scores the chunk by a meaning it no longer has. Storing the hash of the
// embedded text lets the search ignore rows whose source has moved on.
const vectorSchema = `
CREATE TABLE IF NOT EXISTS portfolio_vectors (
    chunk_id     TEXT PRIMARY KEY,
    content_hash TEXT NOT NULL,
    dim          INTEGER NOT NULL,
    vec          BLOB NOT NULL
);
`

// ensureChunkSchema creates portfolio_chunks, and rebuilds it when an older
// database is missing a column (FTS5 has no ALTER TABLE ADD COLUMN).
//
// Dropping is safe precisely because this table is a derived index: every row
// is regenerated from the markdown on the next Reindex. It deliberately
// touches only portfolio_chunks — conversations live in the same file and are
// real user data.
func ensureChunkSchema(ctx context.Context, db *sql.DB) error {
	var existing string
	err := db.QueryRowContext(ctx,
		`SELECT sql FROM sqlite_master WHERE type='table' AND name='portfolio_chunks'`).Scan(&existing)
	switch {
	case err == sql.ErrNoRows:
		// Fresh database; fall through to CREATE.
	case err != nil:
		return fmt.Errorf("inspect chunk schema: %w", err)
	case !strings.Contains(existing, "kind"):
		slog.Info("portfolio index predates the kind column, rebuilding it (run reindex to repopulate)")
		if _, err := db.ExecContext(ctx, `DROP TABLE portfolio_chunks`); err != nil {
			return fmt.Errorf("drop stale chunk table: %w", err)
		}
	}
	if _, err := db.ExecContext(ctx, schema); err != nil {
		return fmt.Errorf("create schema: %w", err)
	}
	// Same rebuild-on-mismatch rule as the chunk table: vectors are derived
	// data, regenerated by the next reindex.
	var vecSQL string
	switch err := db.QueryRowContext(ctx,
		`SELECT sql FROM sqlite_master WHERE type='table' AND name='portfolio_vectors'`).Scan(&vecSQL); {
	case err == sql.ErrNoRows:
	case err != nil:
		return fmt.Errorf("inspect vector schema: %w", err)
	case !strings.Contains(vecSQL, "content_hash"):
		slog.Info("vector index predates content_hash, rebuilding it (run reindex to repopulate)")
		if _, err := db.ExecContext(ctx, `DROP TABLE portfolio_vectors`); err != nil {
			return fmt.Errorf("drop stale vector table: %w", err)
		}
	}
	if _, err := db.ExecContext(ctx, vectorSchema); err != nil {
		return fmt.Errorf("create vector schema: %w", err)
	}
	return nil
}

// Store wraps a SQLite FTS5 database with the portfolio schema.
type Store struct {
	db *sql.DB
}

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 := ensureChunkSchema(context.Background(), db); err != nil {
		_ = db.Close()
		return nil, 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}, 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 }

// isReadme reports whether a file is a directory's README rather than content.
// These directories are checked into the repo with instructions for whoever
// fills them, and those instructions are not one of Victor's projects: without
// this, `data/projects/README.md` was indexed as a project and the catalogue
// injected into every prompt announced "README" and "README.es" to visitors.
// Localised variants (README.es.md) are covered by matching the first segment.
func isReadme(path string) bool {
	base := filepath.Base(path)
	name, _, _ := strings.Cut(base, ".")
	return strings.EqualFold(name, "README")
}

// Reindex rebuilds the whole index from the given sources. Each source
// contributes its `.md` and `.mdx` files; `.mdx` is included because content
// written for an Astro or Next site (a CV, an about page) is usually authored
// there, and its JSX is harmless to full-text search.
//
// A source with an empty Path is skipped, so callers can pass an optional
// docs directory without branching.
func (s *Store) Reindex(ctx context.Context, sources []Source, cfg ChunkerConfig) (files, chunks int, err error) {
	type doc struct {
		path string
		kind string
	}
	var docs []doc
	for _, src := range sources {
		if strings.TrimSpace(src.Path) == "" {
			continue
		}
		kind := src.Kind
		if kind == "" {
			kind = KindProject
		}
		for _, ext := range []string{"*.md", "*.mdx"} {
			matches, err := filepath.Glob(filepath.Join(src.Path, ext))
			if err != nil {
				return 0, 0, fmt.Errorf("glob %s: %w", ext, err)
			}
			for _, m := range matches {
				if isReadme(m) {
					continue
				}
				docs = append(docs, doc{path: m, kind: kind})
			}
		}
	}

	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, kind, source_file, section, chunk_index, content) VALUES (?,?,?,?,?,?,?)`)
	if err != nil {
		return 0, 0, err
	}
	defer stmt.Close()

	for _, d := range docs {
		body, err := os.ReadFile(d.path)
		if err != nil {
			slog.Warn("read file failed", "file", d.path, "err", err)
			continue
		}
		docID := strings.TrimSuffix(strings.TrimSuffix(filepath.Base(d.path), ".mdx"), ".md")
		sections := SplitMarkdownSections(string(body), cfg)
		for idx, sec := range sections {
			id := fmt.Sprintf("%s-%s-%d", docID, slugify(sec.Heading), idx)
			if _, err := stmt.ExecContext(ctx, id, docID, d.kind, d.path, sec.Heading, idx, sec.Body); err != nil {
				return len(docs), 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)
}

// CatalogEntry is one project in the portfolio, identified by its filename
// stem and its human-readable H1 title.
type CatalogEntry struct {
	ProjectID string
	Title     string
}

// Catalog lists every indexed project with its title, cheaply and
// deterministically (no BM25, no query).
//
// It exists because retrieval alone can't answer "what projects does Victor
// have?": top-K search returns the K best-matching *chunks*, which for a
// broad question is a handful of sections from two or three documents, and a
// small model asked to enumerate from that will confidently fill the gaps
// with invented project names. Injecting the full catalogue into the system
// prompt turns enumeration into a copy job instead of a recall job. The
// portfolio is a few dozen documents at most, so the whole list costs a
// trivial number of tokens.
//
// Title falls back to the project ID when a document has no H1.
func (s *Store) Catalog(ctx context.Context) ([]CatalogEntry, error) {
	// The chunker emits the frontmatter block first (when present) and the
	// H1 title section next, so the lowest-indexed non-frontmatter section
	// carries the document's title.
	rows, err := s.db.QueryContext(ctx, `
		SELECT project_id, section, MIN(chunk_index)
		FROM portfolio_chunks
		WHERE section <> 'frontmatter' AND kind = ?
		GROUP BY project_id
		ORDER BY project_id`, KindProject)
	if err != nil {
		return nil, fmt.Errorf("catalog query: %w", err)
	}
	defer rows.Close()

	var out []CatalogEntry
	for rows.Next() {
		var e CatalogEntry
		var idx int
		if err := rows.Scan(&e.ProjectID, &e.Title, &idx); err != nil {
			return nil, fmt.Errorf("catalog scan: %w", err)
		}
		if strings.TrimSpace(e.Title) == "" {
			e.Title = e.ProjectID
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

// 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, kind, source_file, section, chunk_index, content, bm25(portfolio_chunks) AS score
		FROM portfolio_chunks
		WHERE portfolio_chunks MATCH '%s' AND section <> 'frontmatter'
		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.Kind, &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.
// embedder may be nil, in which case no vectors are written and retrieval
// stays keyword-only. An embedding failure is logged and swallowed: the
// keyword index is already committed by then, and a bot that answers from
// keywords alone beats a reindex that reports failure and leaves nothing.
func ReindexOnDisk(dbPath string, sources []Source, cfg ChunkerConfig, embedder Embedder) (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(), sources, cfg)
	if err != nil {
		return time.Since(start), files, chunks, err
	}
	if embedder != nil {
		n, err := store.EmbedChunks(context.Background(), embedder)
		if err != nil {
			slog.Error("embedding failed; keyword search still works", "err", err)
		} else {
			slog.Info("embedded chunks", "chunks", n)
		}
	}
	return time.Since(start), files, chunks, nil
}