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"
	"sort"
	"strconv"
	"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
	SearchH   uint64
	SearchM   uint64
	PDFH      uint64
	PDFM      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
		);
		CREATE TABLE IF NOT EXISTS search (
			k TEXT PRIMARY KEY,
			payload TEXT NOT NULL,
			expires INTEGER NOT NULL,
			last_hit INTEGER NOT NULL,
			nbytes INTEGER NOT NULL
		);
		CREATE TABLE IF NOT EXISTS search_cap (
			source TEXT NOT NULL,
			day TEXT NOT NULL,
			credits INTEGER NOT NULL,
			PRIMARY KEY (source, day)
		);
		CREATE TABLE IF NOT EXISTS pdf_doc (
			k TEXT PRIMARY KEY,
			content_hash,
			etag,
			last_modified,
			fetched_at INTEGER,
			meta BLOB,
			last_hit INTEGER,
			nbytes INTEGER
		);
		CREATE TABLE IF NOT EXISTS pdf_page (
			k TEXT,
			page INTEGER,
			text TEXT,
			PRIMARY KEY (k, page)
		);
	`); 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[:])
}

func SearchKey(query string, sources []string, k int, hour time.Time, callerClass string) string {
	if callerClass == "" {
		callerClass = ClassFleet
	}
	src := append([]string(nil), sources...)
	sort.Strings(src)
	bucket := hour.UTC().Format("2006-01-02T15")
	sum := sha256.Sum256([]byte(NormalizeQuestion(query) + "\x00" + strings.Join(src, ",") + "\x00" + strconv.Itoa(k) + "\x00" + bucket + "\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) Stats() (digestH, digestM, embedH, embedM, searchH, searchM, evictions uint64) {
	if d == nil {
		return 0, 0, 0, 0, 0, 0, 0
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.DigestH, d.DigestM, d.EmbedH, d.EmbedM, d.SearchH, d.SearchM, d.Evictions
}

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

func (d *DB) GetSearch(key string) (payload []byte, ok bool) {
	if d == nil {
		return nil, false
	}
	var b string
	var exp int64
	err := d.sql.QueryRow(`SELECT payload, expires FROM search WHERE k = ?`, key).Scan(&b, &exp)
	if err != nil {
		d.mu.Lock()
		d.SearchM++
		d.mu.Unlock()
		return nil, false
	}
	if exp > 0 && time.Now().Unix() > exp {
		_, _ = d.sql.Exec(`DELETE FROM search WHERE k = ?`, key)
		d.mu.Lock()
		d.SearchM++
		d.mu.Unlock()
		return nil, false
	}
	_, _ = d.sql.Exec(`UPDATE search SET last_hit = ? WHERE k = ?`, time.Now().Unix(), key)
	d.mu.Lock()
	d.SearchH++
	d.mu.Unlock()
	return []byte(b), true
}

func (d *DB) PutSearch(key string, payload []byte, ttl time.Duration) {
	if d == nil {
		return
	}
	exp := time.Now().Add(ttl).Unix()
	n := len(key) + len(payload)
	_, err := d.sql.Exec(`INSERT OR REPLACE INTO search(k, payload, expires, last_hit, nbytes) VALUES(?,?,?,?,?)`,
		key, string(payload), exp, time.Now().Unix(), n)
	if err != nil && d.log != nil {
		d.log.Warn("search cache write", "err", err)
	}
	d.evict()
}

func (d *DB) GetSearchCredits(source, day string) int {
	if d == nil {
		return 0
	}
	var n int
	err := d.sql.QueryRow(`SELECT credits FROM search_cap WHERE source = ? AND day = ?`, source, day).Scan(&n)
	if err != nil {
		return 0
	}
	return n
}

func (d *DB) AddSearchCredits(source, day string, delta int) int {
	if d == nil || delta == 0 {
		if d == nil {
			return 0
		}
		return d.GetSearchCredits(source, day)
	}
	_, err := d.sql.Exec(`INSERT INTO search_cap(source, day, credits) VALUES(?,?,?)
		ON CONFLICT(source, day) DO UPDATE SET credits = credits + excluded.credits`, source, day, delta)
	if err != nil && d.log != nil {
		d.log.Warn("search cap write", "err", err)
	}
	return d.GetSearchCredits(source, day)
}

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, rowid AS rid FROM digest
				UNION ALL
				SELECT k, last_hit, rowid AS rid FROM embed
				UNION ALL
				SELECT k, last_hit, rowid AS rid FROM search
				UNION ALL
				SELECT k, last_hit, rowid AS rid FROM pdf_doc
			) ORDER BY last_hit ASC, rid 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.sql.Exec(`DELETE FROM search WHERE k = ?`, k)
		_, _ = d.sql.Exec(`DELETE FROM pdf_doc WHERE k = ?`, k)
		_, _ = d.sql.Exec(`DELETE FROM pdf_page 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]
}