llmgw/internal/store/batches.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"time"
)

// File is an uploaded or written file of the Files API, without its bytes.
type File struct {
	ID        string
	Owner     string
	Purpose   string
	Filename  string
	Bytes     int64
	CreatedAt time.Time
	// ExpiresAt is when it is deleted by itself; zero is never.
	ExpiresAt time.Time
}

// PutFile keeps a file and its bytes.
func (s *DB) PutFile(ctx context.Context, f File, content []byte) error {
	_, err := s.db.ExecContext(ctx, `INSERT INTO files (id, owner, purpose, filename, bytes, content, created_at, expires_at)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, f.ID, f.Owner, f.Purpose, f.Filename, int64(len(content)), content,
		millis(f.CreatedAt), optMillis(f.ExpiresAt))
	return err
}

const fileCols = `id, owner, purpose, filename, bytes, created_at, expires_at`

func scanFile(sc interface{ Scan(...any) error }) (File, error) {
	var f File
	var created, expires int64
	if err := sc.Scan(&f.ID, &f.Owner, &f.Purpose, &f.Filename, &f.Bytes, &created, &expires); err != nil {
		return File{}, err
	}
	f.CreatedAt = time.UnixMilli(created)
	if expires > 0 {
		f.ExpiresAt = time.UnixMilli(expires)
	}
	return f, nil
}

// FileByID finds a file, ErrNotFound when there is none.
func (s *DB) FileByID(ctx context.Context, id string) (File, error) {
	f, err := scanFile(s.db.QueryRowContext(ctx, `SELECT `+fileCols+` FROM files WHERE id = ?`, id))
	if errors.Is(err, sql.ErrNoRows) {
		return File{}, ErrNotFound
	}
	return f, err
}

// FileContent is a file's bytes.
func (s *DB) FileContent(ctx context.Context, id string) ([]byte, error) {
	var b []byte
	err := s.db.QueryRowContext(ctx, `SELECT content FROM files WHERE id = ?`, id).Scan(&b)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, ErrNotFound
	}
	return b, err
}

// Files lists an owner's files, the newest first; purpose narrows it.
func (s *DB) Files(ctx context.Context, owner, purpose string, limit int) ([]File, error) {
	q := `SELECT ` + fileCols + ` FROM files WHERE owner = ?`
	args := []any{owner}
	if purpose != "" {
		q += ` AND purpose = ?`
		args = append(args, purpose)
	}
	if limit <= 0 || limit > 10000 {
		limit = 10000
	}
	q += ` ORDER BY created_at DESC LIMIT ?`
	args = append(args, limit)
	rows, err := s.db.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []File
	for rows.Next() {
		f, err := scanFile(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, f)
	}
	return out, rows.Err()
}

// DeleteFile removes a file; false when there was none.
func (s *DB) DeleteFile(ctx context.Context, id string) (bool, error) {
	res, err := s.db.ExecContext(ctx, `DELETE FROM files WHERE id = ?`, id)
	if err != nil {
		return false, err
	}
	n, _ := res.RowsAffected()
	return n > 0, nil
}

// DeleteExpiredFiles removes the files past their expiry.
func (s *DB) DeleteExpiredFiles(ctx context.Context, now time.Time) (int64, error) {
	res, err := s.db.ExecContext(ctx, `DELETE FROM files WHERE expires_at > 0 AND expires_at < ?`, millis(now))
	if err != nil {
		return 0, err
	}
	return res.RowsAffected()
}

// Batch is one batch of the Batch API.
type Batch struct {
	ID               string
	Owner            string
	KeyID            string
	Endpoint         string
	InputFileID      string
	CompletionWindow string
	Status           string
	OutputFileID     string
	ErrorFileID      string
	Total            int
	Completed        int
	Failed           int
	CreatedAt        time.Time
	InProgressAt     time.Time
	FinalizingAt     time.Time
	CompletedAt      time.Time
	FailedAt         time.Time
	ExpiredAt        time.Time
	CancellingAt     time.Time
	CancelledAt      time.Time
	ExpiresAt        time.Time
	// Metadata is the client's, as JSON; Errors why validation failed, as
	// JSON.
	Metadata string
	Errors   string
}

// Batch states, as OpenAI names them.
const (
	BatchValidating = "validating"
	BatchFailed     = "failed"
	BatchInProgress = "in_progress"
	BatchFinalizing = "finalizing"
	BatchCompleted  = "completed"
	BatchExpired    = "expired"
	BatchCancelling = "cancelling"
	BatchCancelled  = "cancelled"
)

// BatchDone is whether a state is an end.
func BatchDone(status string) bool {
	switch status {
	case BatchFailed, BatchCompleted, BatchExpired, BatchCancelled:
		return true
	}
	return false
}

const batchCols = `id, owner, key_id, endpoint, input_file_id, completion_window, status, output_file_id,
	error_file_id, total, completed, failed, created_at, in_progress_at, finalizing_at, completed_at,
	failed_at, expired_at, cancelling_at, cancelled_at, expires_at, metadata, errors`

func scanBatch(sc interface{ Scan(...any) error }) (Batch, error) {
	var b Batch
	var t [10]int64
	err := sc.Scan(&b.ID, &b.Owner, &b.KeyID, &b.Endpoint, &b.InputFileID, &b.CompletionWindow, &b.Status,
		&b.OutputFileID, &b.ErrorFileID, &b.Total, &b.Completed, &b.Failed,
		&t[0], &t[1], &t[2], &t[3], &t[4], &t[5], &t[6], &t[7], &t[8], &b.Metadata, &b.Errors)
	if err != nil {
		return Batch{}, err
	}
	at := func(v int64) time.Time {
		if v == 0 {
			return time.Time{}
		}
		return time.UnixMilli(v)
	}
	b.CreatedAt, b.InProgressAt, b.FinalizingAt, b.CompletedAt = at(t[0]), at(t[1]), at(t[2]), at(t[3])
	b.FailedAt, b.ExpiredAt, b.CancellingAt, b.CancelledAt, b.ExpiresAt = at(t[4]), at(t[5]), at(t[6]), at(t[7]), at(t[8])
	return b, nil
}

// PutBatch records a new batch.
func (s *DB) PutBatch(ctx context.Context, b Batch) error {
	_, err := s.db.ExecContext(ctx, `INSERT INTO batches (id, owner, key_id, endpoint, input_file_id,
		completion_window, status, created_at, expires_at, metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		b.ID, b.Owner, b.KeyID, b.Endpoint, b.InputFileID, b.CompletionWindow, b.Status,
		millis(b.CreatedAt), millis(b.ExpiresAt), b.Metadata)
	return err
}

// UpdateBatch writes a batch's progress: its state, counts, files, times
// and errors.
func (s *DB) UpdateBatch(ctx context.Context, b Batch) error {
	_, err := s.db.ExecContext(ctx, `UPDATE batches SET status = ?, output_file_id = ?, error_file_id = ?,
		total = ?, completed = ?, failed = ?, in_progress_at = ?, finalizing_at = ?, completed_at = ?,
		failed_at = ?, expired_at = ?, cancelling_at = ?, cancelled_at = ?, errors = ? WHERE id = ?`,
		b.Status, b.OutputFileID, b.ErrorFileID, b.Total, b.Completed, b.Failed,
		optMillis(b.InProgressAt), optMillis(b.FinalizingAt), optMillis(b.CompletedAt), optMillis(b.FailedAt),
		optMillis(b.ExpiredAt), optMillis(b.CancellingAt), optMillis(b.CancelledAt), b.Errors, b.ID)
	return err
}

// CancelBatch asks for a batch to stop: cancelling, unless it has ended
// already. It says the batch as it stands.
func (s *DB) CancelBatch(ctx context.Context, id string, now time.Time) (Batch, error) {
	if _, err := s.db.ExecContext(ctx, `UPDATE batches SET status = ?, cancelling_at = ?
		WHERE id = ? AND status IN (?, ?, ?)`, BatchCancelling, millis(now), id,
		BatchValidating, BatchInProgress, BatchFinalizing); err != nil {
		return Batch{}, err
	}
	return s.BatchByID(ctx, id)
}

// BatchByID finds a batch, ErrNotFound when there is none.
func (s *DB) BatchByID(ctx context.Context, id string) (Batch, error) {
	b, err := scanBatch(s.db.QueryRowContext(ctx, `SELECT `+batchCols+` FROM batches WHERE id = ?`, id))
	if errors.Is(err, sql.ErrNoRows) {
		return Batch{}, ErrNotFound
	}
	return b, err
}

// Batches lists an owner's batches, the newest first, those made before
// the batch after when it is given.
func (s *DB) Batches(ctx context.Context, owner, after string, limit int) ([]Batch, error) {
	q := `SELECT ` + batchCols + ` FROM batches WHERE owner = ?`
	args := []any{owner}
	if after != "" {
		q += ` AND created_at < (SELECT created_at FROM batches WHERE id = ?)`
		args = append(args, after)
	}
	if limit <= 0 || limit > 100 {
		limit = 20
	}
	q += ` ORDER BY created_at DESC LIMIT ?`
	args = append(args, limit)
	rows, err := s.db.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Batch
	for rows.Next() {
		b, err := scanBatch(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, b)
	}
	return out, rows.Err()
}

// ClaimBatch takes the lease on one batch that has work left and that no
// gateway holds, for owner until the time given. ok is false when there is
// none.
func (s *DB) ClaimBatch(ctx context.Context, owner string, now, until time.Time) (Batch, bool, error) {
	var id string
	err := s.db.QueryRowContext(ctx, `UPDATE batches SET lease_owner = ?, lease_until = ?
		WHERE id = (SELECT id FROM batches WHERE status IN (?, ?, ?, ?) AND lease_until < ?
		            ORDER BY created_at LIMIT 1)
		  AND lease_until < ?
		RETURNING id`, owner, millis(until), BatchValidating, BatchInProgress, BatchFinalizing, BatchCancelling,
		millis(now), millis(now)).Scan(&id)
	if errors.Is(err, sql.ErrNoRows) {
		return Batch{}, false, nil
	}
	if err != nil {
		return Batch{}, false, err
	}
	b, err := s.BatchByID(ctx, id)
	return b, err == nil, err
}

// RenewBatchLease keeps a lease held; false when it is held no longer.
func (s *DB) RenewBatchLease(ctx context.Context, id, owner string, until time.Time) (bool, error) {
	res, err := s.db.ExecContext(ctx, `UPDATE batches SET lease_until = ? WHERE id = ? AND lease_owner = ?`,
		millis(until), id, owner)
	if err != nil {
		return false, err
	}
	n, _ := res.RowsAffected()
	return n > 0, nil
}

// ReleaseBatch gives a lease up.
func (s *DB) ReleaseBatch(ctx context.Context, id, owner string) error {
	_, err := s.db.ExecContext(ctx, `UPDATE batches SET lease_until = 0, lease_owner = '' WHERE id = ? AND lease_owner = ?`, id, owner)
	return err
}

// BatchResult is one line of a batch, served.
type BatchResult struct {
	Line int
	OK   bool
	Data string
}

// PutBatchResult keeps one served line; a line kept already is left as it
// was.
func (s *DB) PutBatchResult(ctx context.Context, batchID string, r BatchResult) error {
	ok := 0
	if r.OK {
		ok = 1
	}
	_, err := s.db.ExecContext(ctx, `INSERT INTO batch_results (batch_id, line, ok, data) VALUES (?, ?, ?, ?)
		ON CONFLICT (batch_id, line) DO NOTHING`, batchID, r.Line, ok, r.Data)
	return err
}

// BatchResults are a batch's served lines, in order.
func (s *DB) BatchResults(ctx context.Context, batchID string) ([]BatchResult, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT line, ok, data FROM batch_results WHERE batch_id = ? ORDER BY line`, batchID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []BatchResult
	for rows.Next() {
		var r BatchResult
		var ok int
		if err := rows.Scan(&r.Line, &ok, &r.Data); err != nil {
			return nil, err
		}
		r.OK = ok == 1
		out = append(out, r)
	}
	return out, rows.Err()
}

// DeleteBatchResults forgets a batch's lines, once its files are written.
func (s *DB) DeleteBatchResults(ctx context.Context, batchID string) error {
	_, err := s.db.ExecContext(ctx, `DELETE FROM batch_results WHERE batch_id = ?`, batchID)
	return err
}

func optMillis(t time.Time) int64 {
	if t.IsZero() {
		return 0
	}
	return t.UnixMilli()
}