descrybe/apps/api/internal/aiaudit/aiaudit.go

// Package aiaudit records the exact prompt and response of every LLM call so
// platform admins can inspect them at /admin/ai-calls.
//
// It is deliberately a leaf package: it imports neither processing nor aiprovider,
// so processing can attach per-call context (job, product, role) without an import
// cycle while aiprovider does the actual Completer wrapping.
//
// Rows are a short-lived debugging buffer, not an audit trail — CleanupExpired
// prunes them after RetentionDays. Never build billing or reporting on them.
package aiaudit

import (
	"context"
	"log"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Roles mirror aiprovider role ids plus "probe" for admin connectivity tests.
const (
	RoleProcessing = "processing"
	RoleCategorize = "categorize"
	RoleSEOMeta    = "seo_meta"
	RoleCampaign   = "campaign"
	RoleSupport    = "support"
	RoleProbe      = "probe"
	RoleOther      = "other"
)

// RetentionDays bounds how long captured prompts live. Prompts embed the tenant's
// product copy and are large, so the window is intentionally short.
const RetentionDays = 7

// maxBodyRunes caps each stored body. A category-formula prompt is a few KB; this
// leaves room for outliers while stopping a runaway reply from bloating the table.
const maxBodyRunes = 60000

// Call is one recorded LLM exchange.
type Call struct {
	CompanyID    uuid.UUID
	UserID       uuid.UUID
	JobID        uuid.UUID
	RawProductID uuid.UUID
	Role         string
	ProviderMode string
	Model        string
	System       string
	User         string
	Response     string
	Error        string
	FinishReason string
	PromptTokens int
	OutputTokens int
	TotalTokens  int
	Duration     time.Duration
}

// CallContext is the per-call metadata processing/campaigns attach to ctx before
// invoking a Completer. The recorder merges it into every Call it writes.
type CallContext struct {
	CompanyID    uuid.UUID
	UserID       uuid.UUID
	JobID        uuid.UUID
	RawProductID uuid.UUID
	Role         string
}

type ctxKey struct{}

// WithCall returns ctx carrying call metadata for the recorder. Fields left zero
// on cc do not clear values already present, so an outer scope can set company/user
// once per job and an inner scope add the product.
func WithCall(ctx context.Context, cc CallContext) context.Context {
	cur := FromContext(ctx)
	if cc.CompanyID != uuid.Nil {
		cur.CompanyID = cc.CompanyID
	}
	if cc.UserID != uuid.Nil {
		cur.UserID = cc.UserID
	}
	if cc.JobID != uuid.Nil {
		cur.JobID = cc.JobID
	}
	if cc.RawProductID != uuid.Nil {
		cur.RawProductID = cc.RawProductID
	}
	if strings.TrimSpace(cc.Role) != "" {
		cur.Role = cc.Role
	}
	return context.WithValue(ctx, ctxKey{}, cur)
}

// FromContext reads call metadata attached by WithCall (zero value when absent).
func FromContext(ctx context.Context) CallContext {
	if ctx == nil {
		return CallContext{}
	}
	cc, _ := ctx.Value(ctxKey{}).(CallContext)
	return cc
}

// Recorder writes captured calls. A nil Recorder is a no-op, so callers never need
// to branch on whether capture is wired.
type Recorder struct {
	pool *pgxpool.Pool

	mu     sync.Mutex
	failed bool // stop log-spamming once writes are known to fail
}

// NewRecorder returns a Recorder writing to pool. nil pool yields a no-op recorder.
func NewRecorder(pool *pgxpool.Pool) *Recorder {
	if pool == nil {
		return nil
	}
	return &Recorder{pool: pool}
}

const insertCallSQL = `
	INSERT INTO ai_call_logs (
		company_id, user_id, job_id, raw_product_id, role, provider_mode, model,
		system_prompt, user_prompt, response_text, error, finish_reason,
		prompt_tokens, completion_tokens, total_tokens, duration_ms)
	VALUES (
		$1, $2, $3, $4, $5, $6, $7,
		$8, $9, $10, $11, $12,
		$13, $14, $15, $16)`

// Record persists one call, merging any CallContext on ctx. Capture must never
// break the request it is observing: every failure is swallowed after one log line.
//
// The write uses context.WithoutCancel so a cancelled/timed-out job still records
// the call that was in flight — those are exactly the ones worth inspecting.
func (r *Recorder) Record(ctx context.Context, c Call) {
	if r == nil || r.pool == nil {
		return
	}
	cc := FromContext(ctx)
	if c.CompanyID == uuid.Nil {
		c.CompanyID = cc.CompanyID
	}
	if c.UserID == uuid.Nil {
		c.UserID = cc.UserID
	}
	if c.JobID == uuid.Nil {
		c.JobID = cc.JobID
	}
	if c.RawProductID == uuid.Nil {
		c.RawProductID = cc.RawProductID
	}
	if strings.TrimSpace(c.Role) == "" {
		c.Role = cc.Role
	}
	if strings.TrimSpace(c.Role) == "" {
		c.Role = RoleOther
	}

	writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
	defer cancel()
	_, err := r.pool.Exec(writeCtx, insertCallSQL,
		nullUUID(c.CompanyID), nullUUID(c.UserID), nullUUID(c.JobID), nullUUID(c.RawProductID),
		c.Role, c.ProviderMode, c.Model,
		clamp(c.System), clamp(c.User), clamp(c.Response), clamp(c.Error), c.FinishReason,
		c.PromptTokens, c.OutputTokens, c.TotalTokens, int(c.Duration.Milliseconds()),
	)
	if err != nil {
		r.noteFailure(err)
	}
}

func (r *Recorder) noteFailure(err error) {
	r.mu.Lock()
	first := !r.failed
	r.failed = true
	r.mu.Unlock()
	if first {
		log.Printf("aiaudit: capture disabled after write error: %v", err)
	}
}

func nullUUID(id uuid.UUID) any {
	if id == uuid.Nil {
		return nil
	}
	return id
}

func clamp(s string) string {
	rs := []rune(s)
	if len(rs) <= maxBodyRunes {
		return s
	}
	return string(rs[:maxBodyRunes]) + "\n…[truncated by aiaudit]"
}

// CleanupExpired deletes captured calls older than RetentionDays and returns how
// many rows went. Called from the worker's retention tick.
func CleanupExpired(ctx context.Context, pool *pgxpool.Pool) (int64, error) {
	if pool == nil {
		return 0, nil
	}
	tag, err := pool.Exec(ctx, `
		DELETE FROM ai_call_logs
		WHERE created_at < now() - ($1 || ' days')::interval`, RetentionDays)
	if err != nil {
		return 0, err
	}
	return tag.RowsAffected(), nil
}