llmgw/internal/store/guard.go

package store

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
)

// GuardAttachments lists which policy is attached where.
func (s *DB) GuardAttachments(ctx context.Context) ([]guard.Attachment, error) {
	return queryAll(ctx, s.db, `SELECT scope, scope_id, policy FROM guard_attachments ORDER BY scope, scope_id, policy`, nil,
		func(rows *sql.Rows) (guard.Attachment, error) {
			var a guard.Attachment
			err := rows.Scan(&a.Scope, &a.ScopeID, &a.Policy)
			return a, err
		})
}

// AttachGuard attaches a policy to a scope; attaching it again changes
// nothing.
func (s *DB) AttachGuard(ctx context.Context, a guard.Attachment, by string) error {
	_, err := s.db.ExecContext(ctx, `INSERT INTO guard_attachments (scope, scope_id, policy, added_at, added_by)
		VALUES (?, ?, ?, ?, ?) ON CONFLICT (scope, scope_id, policy) DO NOTHING`,
		a.Scope, a.ScopeID, a.Policy, millis(time.Now()), by)
	return err
}

// DetachGuard removes one attachment, reporting whether it was there.
func (s *DB) DetachGuard(ctx context.Context, a guard.Attachment) (bool, error) {
	res, err := s.db.ExecContext(ctx, `DELETE FROM guard_attachments WHERE scope = ? AND scope_id = ? AND policy = ?`,
		a.Scope, a.ScopeID, a.Policy)
	if err != nil {
		return false, err
	}
	n, _ := res.RowsAffected()
	return n > 0, nil
}

// GuardEvent is one request the guardrails acted on, by name only.
type GuardEvent struct {
	At        time.Time
	RequestID string
	KeyID     string
	UserID    string
	ProjectID string
	Model     string
	Status    int
	Action    string
	Policy    string
	Category  string
	GuardMs   int64
}

// GuardEvents are the newest requests the guardrails did something to —
// flagged, redacted, blocked, or would have in shadow — since a moment.
func (s *DB) GuardEvents(ctx context.Context, since time.Time, limit int) ([]GuardEvent, error) {
	return queryAll(ctx, s.db, `SELECT ts, request_id, key_id, COALESCE(user_id, ''), project_id, model, status,
		guard_action, guard_policy, guard_category, guard_ms
		FROM usage WHERE ts >= ? AND guard_action <> '' ORDER BY ts DESC LIMIT ?`,
		[]any{millis(since), limit},
		func(rows *sql.Rows) (GuardEvent, error) {
			var e GuardEvent
			var ts int64
			err := rows.Scan(&ts, &e.RequestID, &e.KeyID, &e.UserID, &e.ProjectID, &e.Model, &e.Status,
				&e.Action, &e.Policy, &e.Category, &e.GuardMs)
			e.At = fromMillis(ts)
			return e, err
		})
}

// GuardCount is how many requests one combination covers.
type GuardCount struct {
	ProjectID, Action, Policy, Category string
	Count                               int64
}

// GuardCounts totals what the guardrails did since a moment, by project,
// action, policy and category.
func (s *DB) GuardCounts(ctx context.Context, since time.Time) ([]GuardCount, error) {
	return queryAll(ctx, s.db, `SELECT project_id, guard_action, guard_policy, guard_category, COUNT(*)
		FROM usage WHERE ts >= ? AND guard_action <> ''
		GROUP BY project_id, guard_action, guard_policy, guard_category
		ORDER BY COUNT(*) DESC`, []any{millis(since)},
		func(rows *sql.Rows) (GuardCount, error) {
			var c GuardCount
			err := rows.Scan(&c.ProjectID, &c.Action, &c.Policy, &c.Category, &c.Count)
			return c, err
		})
}