llmgw/internal/store/requestlog.go

package store

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

// RequestLog is how a project keeps its requests and answers: for how long,
// and what share of them. Zero Retain keeps none.
type RequestLog struct {
	Retain time.Duration
	// Sample is the share kept, above 0 and at most 1.
	Sample float64
}

// SetRequestLog turns keeping requests and answers on for a project, or off
// with hours 0. A sample outside (0, 1] keeps every request.
func (s *DB) SetRequestLog(ctx context.Context, projectID string, hours int, sample float64) error {
	var h, sm any
	if hours > 0 {
		h = hours
		if sample <= 0 || sample > 1 {
			sample = 1
		}
		sm = sample
	}
	res, err := s.db.ExecContext(ctx,
		`UPDATE projects SET request_log_hours = ?, request_log_sample = ? WHERE id = ?`, h, sm, projectID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return fmt.Errorf("store: no project %q", projectID)
	}
	return nil
}

// RequestLogOf is a project's setting; zero when it keeps none.
func (s *DB) RequestLogOf(ctx context.Context, projectID string) (RequestLog, error) {
	var hours sql.NullInt64
	var sample sql.NullFloat64
	err := s.db.QueryRowContext(ctx,
		`SELECT request_log_hours, request_log_sample FROM projects WHERE id = ?`, projectID).Scan(&hours, &sample)
	if errors.Is(err, sql.ErrNoRows) || err == nil && (!hours.Valid || hours.Int64 <= 0) {
		return RequestLog{}, nil
	}
	if err != nil {
		return RequestLog{}, err
	}
	l := RequestLog{Retain: time.Duration(hours.Int64) * time.Hour, Sample: 1}
	if sample.Valid && sample.Float64 > 0 && sample.Float64 <= 1 {
		l.Sample = sample.Float64
	}
	return l, nil
}

// RequestBody is one kept request and its answer.
type RequestBody struct {
	RequestID   string    `json:"request_id"`
	ProjectID   string    `json:"project_id"`
	At          time.Time `json:"at"`
	ExpiresAt   time.Time `json:"expires_at"`
	Path        string    `json:"path"`
	Model       string    `json:"model"`
	ServedModel string    `json:"served_model"`
	Status      int       `json:"status"`
	Request     string    `json:"request"`
	// Response is the answer's text as the client was sent it; Truncated
	// says it was longer than what is kept.
	Response  string `json:"response"`
	Truncated bool   `json:"truncated,omitempty"`
}

// PutRequestBody keeps one.
func (s *DB) PutRequestBody(ctx context.Context, b RequestBody) error {
	tr := 0
	if b.Truncated {
		tr = 1
	}
	_, err := s.db.ExecContext(ctx, `INSERT INTO request_bodies (request_id, project_id, ts, expires_at, path,
		model, served_model, status, request, response, truncated) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT (request_id) DO NOTHING`,
		b.RequestID, b.ProjectID, millis(b.At), millis(b.ExpiresAt), b.Path, b.Model, b.ServedModel, b.Status,
		b.Request, b.Response, tr)
	return err
}

// RequestBodyByID finds one, ErrNotFound when none is kept.
func (s *DB) RequestBodyByID(ctx context.Context, id string) (RequestBody, error) {
	var b RequestBody
	var ts, exp int64
	var tr int
	err := s.db.QueryRowContext(ctx, `SELECT request_id, project_id, ts, expires_at, path, model, served_model,
		status, request, response, truncated FROM request_bodies WHERE request_id = ?`, id).
		Scan(&b.RequestID, &b.ProjectID, &ts, &exp, &b.Path, &b.Model, &b.ServedModel, &b.Status,
			&b.Request, &b.Response, &tr)
	if errors.Is(err, sql.ErrNoRows) {
		return RequestBody{}, ErrNotFound
	}
	if err != nil {
		return RequestBody{}, err
	}
	b.At, b.ExpiresAt, b.Truncated = time.UnixMilli(ts), time.UnixMilli(exp), tr == 1
	return b, nil
}

// PruneRequestBodies removes those past their retention, and those of
// projects that keep none any more.
func (s *DB) PruneRequestBodies(ctx context.Context, now time.Time) (int64, error) {
	res, err := s.db.ExecContext(ctx, `DELETE FROM request_bodies WHERE expires_at < ?
		OR project_id IN (SELECT id FROM projects WHERE request_log_hours IS NULL OR request_log_hours <= 0)`,
		millis(now))
	if err != nil {
		return 0, err
	}
	return res.RowsAffected()
}

// SetInCluster says whether a project's requests stay in the cluster: no
// model served by an outside provider, asked for or as a fallback.
func (s *DB) SetInCluster(ctx context.Context, projectID string, on bool) error {
	v := 0
	if on {
		v = 1
	}
	res, err := s.db.ExecContext(ctx, `UPDATE projects SET in_cluster = ? WHERE id = ?`, v, projectID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return fmt.Errorf("store: no project %q", projectID)
	}
	return nil
}

// InCluster is whether a project's requests stay in the cluster.
func (s *DB) InCluster(ctx context.Context, projectID string) (bool, error) {
	var v int
	err := s.db.QueryRowContext(ctx, `SELECT in_cluster FROM projects WHERE id = ?`, projectID).Scan(&v)
	if errors.Is(err, sql.ErrNoRows) {
		return false, nil
	}
	return v == 1, err
}