descrybe/apps/api/internal/aiprovider/audit.go

package aiprovider

import (
	"context"
	"strings"
	"time"

	"github.com/descrybe/descrybe-v2/apps/api/internal/aiaudit"
	"github.com/descrybe/descrybe-v2/apps/api/internal/processing"
	"github.com/google/uuid"
)

// Every Completer this service hands out is wrapped here, so the admin AI
// inspector sees processing, categorize, SEO meta, campaign and support calls from
// one capture point instead of each caller remembering to log.
//
// Callers attach per-call detail (job, product, role) with aiaudit.WithCall; the
// wrapper fills in what the resolver already knows (company, provider mode, model).

// Audit enables prompt/response capture. Nil disables it and the wrapper is skipped
// entirely, so a service built without a recorder behaves exactly as before.
func (s *Service) WithAudit(rec *aiaudit.Recorder) *Service {
	if s == nil {
		return s
	}
	s.Audit = rec
	return s
}

type auditingCompleter struct {
	inner        processing.Completer
	rec          *aiaudit.Recorder
	companyID    uuid.UUID
	role         string
	providerMode string
}

// Enabled forwards the wrapped client's gate so CompleterEnabled keeps working
// (processing checks for the EnableChecker interface, not a concrete type).
func (a *auditingCompleter) Enabled() bool {
	if c, ok := a.inner.(processing.EnableChecker); ok {
		return c.Enabled()
	}
	return a.inner != nil
}

// ProviderModeLabel forwards the wrapped client's analytics label.
func (a *auditingCompleter) ProviderModeLabel() string {
	type labeler interface{ ProviderModeLabel() string }
	if c, ok := a.inner.(labeler); ok {
		return c.ProviderModeLabel()
	}
	return a.providerMode
}

func (a *auditingCompleter) Complete(ctx context.Context, system, user string) (processing.Completion, error) {
	return a.record(ctx, system, user, func() (processing.Completion, error) {
		return a.inner.Complete(ctx, system, user)
	})
}

// auditingCompleterWithOptions is used when the wrapped client honours
// CompleteOptions. processing.CompleteOnce prefers that interface and only falls
// back to Complete, so a wrapper that did not forward it would silently drop
// MaxTokens / Temperature / ReasoningEffort from every enhance call.
type auditingCompleterWithOptions struct {
	*auditingCompleter
	inner processing.CompleterWithOptions
}

func (a *auditingCompleterWithOptions) CompleteWithOptions(ctx context.Context, system, user string, opts processing.CompleteOptions) (processing.Completion, error) {
	return a.record(ctx, system, user, func() (processing.Completion, error) {
		return a.inner.CompleteWithOptions(ctx, system, user, opts)
	})
}

func (a *auditingCompleter) record(ctx context.Context, system, user string, call func() (processing.Completion, error)) (processing.Completion, error) {
	started := time.Now()
	comp, err := call()
	// The call site knows the role better than the resolver does: one processing
	// Completer serves both the enhance and categorize steps.
	role := a.role
	if cc := aiaudit.FromContext(ctx); strings.TrimSpace(cc.Role) != "" {
		role = cc.Role
	}
	rec := aiaudit.Call{
		CompanyID:          a.companyID,
		Role:               role,
		ProviderMode:       a.providerMode,
		Model:              comp.Model,
		System:             system,
		User:               user,
		Response:           comp.Text,
		FinishReason:       finishReason(comp),
		PromptTokens:       comp.PromptTokens,
		CachedPromptTokens: comp.CachedPromptTokens,
		OutputTokens:       comp.OutputTokens,
		TotalTokens:        comp.TotalTokens,
		Duration:           time.Since(started),
	}
	if err != nil {
		// The provider error text is already redacted by processing's OpenAI client.
		rec.Error = err.Error()
	}
	a.rec.Record(ctx, rec)
	return comp, err
}

func finishReason(comp processing.Completion) string {
	m, ok := comp.Raw.(map[string]any)
	if !ok || m == nil {
		return ""
	}
	fr, _ := m["finish_reason"].(string)
	return strings.TrimSpace(fr)
}

// WrapAudit decorates an externally supplied Completer with the same capture the
// resolvers apply. Local tools that pin a provider endpoint (cmd/formula-e2e) use
// it so their runs show up in the admin inspector too.
func (s *Service) WrapAudit(c processing.Completer, companyID uuid.UUID, role, providerMode string) processing.Completer {
	return s.wrapAudit(c, companyID, role, providerMode)
}

// wrapAudit decorates a resolved Completer with capture. It is idempotent (a
// re-wrap keeps the original) and a no-op when capture is off or c is nil, so the
// public resolvers can all call it on their way out.
func (s *Service) wrapAudit(c processing.Completer, companyID uuid.UUID, role, providerMode string) processing.Completer {
	if s == nil || s.Audit == nil || c == nil {
		return c
	}
	switch c.(type) {
	case *auditingCompleter, *auditingCompleterWithOptions:
		return c
	}
	base := &auditingCompleter{
		inner:        c,
		rec:          s.Audit,
		companyID:    companyID,
		role:         auditRole(role),
		providerMode: providerMode,
	}
	// Only advertise the options interface when the wrapped client actually has it,
	// so a plain Completer is not handed options it cannot honour.
	if withOpts, ok := c.(processing.CompleterWithOptions); ok {
		return &auditingCompleterWithOptions{auditingCompleter: base, inner: withOpts}
	}
	return base
}

// auditRole maps an aiprovider role onto the aiaudit vocabulary shown in the admin
// filter. Unknown roles are recorded rather than dropped.
func auditRole(role string) string {
	switch strings.TrimSpace(role) {
	case "", RoleProcessing:
		return aiaudit.RoleProcessing
	case RoleSupport:
		return aiaudit.RoleSupport
	case RoleDocsAPI:
		return aiaudit.RoleOther
	default:
		return strings.TrimSpace(role)
	}
}