kyle
4372673184
Default GET /metrics is Prometheus text; JSON stays at /metrics.json and ?format=json. :8443 needs the metrics caller bearer; loopback stays open. Restricted callers have no per-caller series and no tavily:restricted gauge.
358 lines
8.5 KiB
Go
358 lines
8.5 KiB
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
|
|
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)
|
|
);
|
|
`); 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 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)
|
|
return n.Int64 + m.Int64 + s.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 FROM digest
|
|
UNION ALL
|
|
SELECT k, last_hit FROM embed
|
|
UNION ALL
|
|
SELECT k, last_hit FROM search
|
|
) 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.sql.Exec(`DELETE FROM search 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]
|
|
}
|