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
	// CachedPromptTokens is the cached subset of PromptTokens — billed far cheaper,
	// so cost accounting must know about it.
	CachedPromptTokens 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

	// prices caches ai_model_prices for cost rollup (see cost.go).
	prices priceCache
}

// 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)
	}
	// Roll the same call into the durable daily cost aggregate. ai_call_logs is
	// pruned after a week; spend has to survive for years (see cost.go).
	r.recordUsage(ctx, c)
}

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
}