sift/internal/dcache/dcache.go

// Package dcache is the content-addressed digest + chunk-embedding SQLite cache.
// Honcho context is never stored here.
package dcache

import (
	"crypto/sha256"
	"database/sql"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"log/slog"
	"math"
	"os"
	"strings"
	"sync"
	"time"

	_ "modernc.org/sqlite"
)

const (
	PromptVersion   = "2"
	ClassFleet      = "fleet"
	ClassRestricted = "restricted"
)

func CallerClass(restricted bool) string {
	if restricted {
		return ClassRestricted
	}
	return ClassFleet
}

type DB struct {
	sql      *sql.DB
	maxBytes int64
	log      *slog.Logger

	mu        sync.Mutex
	DigestH   uint64
	DigestM   uint64
	EmbedH    uint64
	EmbedM    uint64
	Evictions uint64
}

func Open(path string, maxBytes int64, log *slog.Logger) (*DB, error) {
	if maxBytes <= 0 {
		maxBytes = 500 << 20
	}
	if err := os.MkdirAll(dirOf(path), 0o750); err != nil && path != ":memory:" && path != "" {
		// ignore: sqlite will fail if it cannot create
	}
	dsn := path
	if path == "" {
		dsn = ":memory:"
	}
	sq, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, err
	}
	sq.SetMaxOpenConns(1)
	if _, err := sq.Exec(`
		PRAGMA journal_mode=WAL;
		CREATE TABLE IF NOT EXISTS digest (
			k TEXT PRIMARY KEY,
			text TEXT NOT NULL,
			chunks_used INTEGER,
			last_hit INTEGER NOT NULL,
			nbytes INTEGER NOT NULL
		);
		CREATE TABLE IF NOT EXISTS embed (
			k TEXT PRIMARY KEY,
			vec BLOB NOT NULL,
			last_hit INTEGER NOT NULL,
			nbytes INTEGER NOT NULL
		);
	`); err != nil {
		_ = sq.Close()
		return nil, err
	}
	d := &DB{sql: sq, maxBytes: maxBytes, log: log}
	return d, nil
}

func (d *DB) Close() error {
	if d == nil || d.sql == nil {
		return nil
	}
	return d.sql.Close()
}

func NormalizeQuestion(q string) string {
	return strings.ToLower(strings.Join(strings.Fields(q), " "))
}

func DigestKey(contentHash, question, model, promptVer, callerClass string) string {
	if callerClass == "" {
		callerClass = ClassFleet
	}
	sum := sha256.Sum256([]byte(contentHash + "\x00" + NormalizeQuestion(question) + "\x00" + model + "\x00" + promptVer + "\x00" + callerClass))
	return hex.EncodeToString(sum[:])
}

func EmbedKey(contentHash string, chunkIndex int, embedModel, callerClass string) string {
	if callerClass == "" {
		callerClass = ClassFleet
	}
	sum := sha256.Sum256([]byte(fmt.Sprintf("%s\x00%d\x00%s\x00%s", contentHash, chunkIndex, embedModel, callerClass)))
	return hex.EncodeToString(sum[:])
}

func QuestionEmbedKey(question, embedModel, callerClass string) string {
	if callerClass == "" {
		callerClass = ClassFleet
	}
	sum := sha256.Sum256([]byte("q\x00" + NormalizeQuestion(question) + "\x00" + embedModel + "\x00" + callerClass))
	return hex.EncodeToString(sum[:])
}

type DigestHit struct {
	Text       string
	ChunksUsed int
}

func (d *DB) GetDigest(key string) (DigestHit, bool) {
	if d == nil {
		return DigestHit{}, false
	}
	var h DigestHit
	var n int
	err := d.sql.QueryRow(`SELECT text, chunks_used FROM digest WHERE k = ?`, key).Scan(&h.Text, &n)
	if err != nil {
		d.mu.Lock()
		d.DigestM++
		d.mu.Unlock()
		return DigestHit{}, false
	}
	h.ChunksUsed = n
	_, _ = d.sql.Exec(`UPDATE digest SET last_hit = ? WHERE k = ?`, time.Now().Unix(), key)
	d.mu.Lock()
	d.DigestH++
	d.mu.Unlock()
	return h, true
}

func (d *DB) PutDigest(key, text string, chunksUsed int) {
	if d == nil {
		return
	}
	n := len(key) + len(text)
	_, err := d.sql.Exec(`INSERT OR REPLACE INTO digest(k, text, chunks_used, last_hit, nbytes) VALUES(?,?,?,?,?)`,
		key, text, chunksUsed, time.Now().Unix(), n)
	if err != nil && d.log != nil {
		d.log.Warn("digest cache write", "err", err)
	}
	d.evict()
}

func (d *DB) GetEmbed(key string) ([]float64, bool) {
	if d == nil {
		return nil, false
	}
	var b []byte
	err := d.sql.QueryRow(`SELECT vec FROM embed WHERE k = ?`, key).Scan(&b)
	if err != nil {
		d.mu.Lock()
		d.EmbedM++
		d.mu.Unlock()
		return nil, false
	}
	_, _ = d.sql.Exec(`UPDATE embed SET last_hit = ? WHERE k = ?`, time.Now().Unix(), key)
	d.mu.Lock()
	d.EmbedH++
	d.mu.Unlock()
	return blobToVec(b), true
}

func (d *DB) PutEmbed(key string, vec []float64) {
	if d == nil {
		return
	}
	b := vecToBlob(vec)
	_, err := d.sql.Exec(`INSERT OR REPLACE INTO embed(k, vec, last_hit, nbytes) VALUES(?,?,?,?)`,
		key, b, time.Now().Unix(), len(key)+len(b))
	if err != nil && d.log != nil {
		d.log.Warn("embed cache write", "err", err)
	}
	d.evict()
}

func (d *DB) Size() int64 {
	if d == nil {
		return 0
	}
	var n sql.NullInt64
	_ = d.sql.QueryRow(`SELECT IFNULL(SUM(nbytes),0) FROM digest`).Scan(&n)
	var m sql.NullInt64
	_ = d.sql.QueryRow(`SELECT IFNULL(SUM(nbytes),0) FROM embed`).Scan(&m)
	return n.Int64 + m.Int64
}

func (d *DB) evict() {
	if d == nil {
		return
	}
	for d.Size() > d.maxBytes {
		var k string
		err := d.sql.QueryRow(`
			SELECT k FROM (
				SELECT k, last_hit, 0 AS kind FROM digest
				UNION ALL
				SELECT k, last_hit, 1 FROM embed
			) ORDER BY last_hit ASC LIMIT 1`).Scan(&k)
		if err != nil {
			return
		}
		_, _ = d.sql.Exec(`DELETE FROM digest WHERE k = ?`, k)
		_, _ = d.sql.Exec(`DELETE FROM embed WHERE k = ?`, k)
		d.mu.Lock()
		d.Evictions++
		d.mu.Unlock()
	}
}

func vecToBlob(v []float64) []byte {
	b := make([]byte, 4*len(v))
	for i, x := range v {
		binary.LittleEndian.PutUint32(b[i*4:], math.Float32bits(float32(x)))
	}
	return b
}

func blobToVec(b []byte) []float64 {
	n := len(b) / 4
	out := make([]float64, n)
	for i := 0; i < n; i++ {
		out[i] = float64(math.Float32frombits(binary.LittleEndian.Uint32(b[i*4:])))
	}
	return out
}

func dirOf(path string) string {
	i := strings.LastIndex(path, "/")
	if i <= 0 {
		return "."
	}
	return path[:i]
}