AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,282 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type Event struct {
|
||||
TenantID *string
|
||||
RequestID string
|
||||
APIKeyID *string
|
||||
ProviderCode string
|
||||
Model string
|
||||
Protocol string
|
||||
StatusCode int
|
||||
PromptTokens int64
|
||||
CompletionTokens int64
|
||||
CostMicrounits int64
|
||||
PriceID string
|
||||
Currency string
|
||||
LatencyMS int
|
||||
Labels map[string]any
|
||||
RecordedAt time.Time
|
||||
}
|
||||
|
||||
type Stats struct {
|
||||
Accepted uint64
|
||||
Dropped uint64
|
||||
Written uint64
|
||||
Failures uint64
|
||||
}
|
||||
|
||||
type Recorder struct {
|
||||
pool *pgxpool.Pool
|
||||
logger *slog.Logger
|
||||
queue chan Event
|
||||
batchSize int
|
||||
flushInterval time.Duration
|
||||
accepted atomic.Uint64
|
||||
dropped atomic.Uint64
|
||||
written atomic.Uint64
|
||||
failures atomic.Uint64
|
||||
}
|
||||
|
||||
func NewRecorder(pool *pgxpool.Pool, logger *slog.Logger, queueSize, batchSize int, flushInterval time.Duration) *Recorder {
|
||||
if queueSize < 1 {
|
||||
queueSize = 4096
|
||||
}
|
||||
if batchSize < 1 {
|
||||
batchSize = 200
|
||||
}
|
||||
if flushInterval <= 0 {
|
||||
flushInterval = time.Second
|
||||
}
|
||||
return &Recorder{pool: pool, logger: logger, queue: make(chan Event, queueSize), batchSize: batchSize, flushInterval: flushInterval}
|
||||
}
|
||||
|
||||
func (r *Recorder) Record(event Event) bool {
|
||||
if r == nil || r.pool == nil {
|
||||
return false
|
||||
}
|
||||
if event.RecordedAt.IsZero() {
|
||||
event.RecordedAt = time.Now().UTC()
|
||||
}
|
||||
select {
|
||||
case r.queue <- event:
|
||||
r.accepted.Add(1)
|
||||
return true
|
||||
default:
|
||||
r.dropped.Add(1)
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) Stats() Stats {
|
||||
if r == nil {
|
||||
return Stats{}
|
||||
}
|
||||
return Stats{Accepted: r.accepted.Load(), Dropped: r.dropped.Load(), Written: r.written.Load(), Failures: r.failures.Load()}
|
||||
}
|
||||
|
||||
func (r *Recorder) Prometheus() string {
|
||||
stats := r.Stats()
|
||||
return fmt.Sprintf("gateway_audit_events_accepted_total %d\ngateway_audit_events_dropped_total %d\ngateway_audit_events_written_total %d\ngateway_audit_flush_failures_total %d\n", stats.Accepted, stats.Dropped, stats.Written, stats.Failures)
|
||||
}
|
||||
|
||||
func (r *Recorder) Run(ctx context.Context) {
|
||||
ticker := time.NewTicker(r.flushInterval)
|
||||
defer ticker.Stop()
|
||||
pending := make([]Event, 0, r.batchSize)
|
||||
flushAllowed := true
|
||||
for {
|
||||
var events <-chan Event = r.queue
|
||||
if len(pending) >= cap(r.queue)+r.batchSize {
|
||||
events = nil
|
||||
}
|
||||
select {
|
||||
case event := <-events:
|
||||
pending = append(pending, event)
|
||||
if len(pending) >= r.batchSize && flushAllowed {
|
||||
if r.flush(ctx, pending) == nil {
|
||||
pending = pending[:0]
|
||||
} else {
|
||||
flushAllowed = false
|
||||
}
|
||||
}
|
||||
case <-ticker.C:
|
||||
flushAllowed = true
|
||||
if len(pending) > 0 {
|
||||
if r.flush(ctx, pending) == nil {
|
||||
pending = pending[:0]
|
||||
} else {
|
||||
flushAllowed = false
|
||||
}
|
||||
}
|
||||
case <-ctx.Done():
|
||||
for {
|
||||
select {
|
||||
case event := <-r.queue:
|
||||
pending = append(pending, event)
|
||||
default:
|
||||
flushCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
if len(pending) > 0 {
|
||||
_ = r.flush(flushCtx, pending)
|
||||
}
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type dailyKey struct {
|
||||
date, apiKeyID, provider, model string
|
||||
}
|
||||
|
||||
type dailyValue struct {
|
||||
requests, failed, prompt, completion, cost int64
|
||||
}
|
||||
|
||||
type policyAlert struct {
|
||||
eventID, requestID string
|
||||
tenantID *string
|
||||
payload []byte
|
||||
}
|
||||
|
||||
func (r *Recorder) flush(ctx context.Context, events []Event) error {
|
||||
if len(events) == 0 {
|
||||
return nil
|
||||
}
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
rows := make([][]any, 0, len(events))
|
||||
daily := make(map[dailyKey]dailyValue)
|
||||
alerts := make([]policyAlert, 0)
|
||||
for _, event := range events {
|
||||
id, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
labels, _ := json.Marshal(event.Labels)
|
||||
if policies, matched := event.Labels["content_policies"]; matched {
|
||||
alertID, idErr := platformid.NewUUID()
|
||||
if idErr != nil {
|
||||
return r.flushError(idErr)
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"request_id": event.RequestID, "protocol": event.Protocol, "status_code": event.StatusCode,
|
||||
"provider_code": event.ProviderCode, "model": event.Model, "content_policies": policies,
|
||||
"content_redacted": event.Labels["content_redacted"], "recorded_at": event.RecordedAt,
|
||||
})
|
||||
alerts = append(alerts, policyAlert{eventID: alertID, requestID: event.RequestID, tenantID: event.TenantID, payload: payload})
|
||||
}
|
||||
rows = append(rows, []any{
|
||||
uuidValue(id), uuidPointer(event.TenantID), event.RequestID, nil, uuidPointer(event.APIKeyID),
|
||||
nullString(event.ProviderCode), nullString(event.Model), event.Protocol, event.StatusCode,
|
||||
nullToken(event.PromptTokens), nullToken(event.CompletionTokens), nullCost(event.CostMicrounits), event.LatencyMS,
|
||||
nil, nil, labels, event.RecordedAt,
|
||||
})
|
||||
if event.APIKeyID != nil && *event.APIKeyID != "" {
|
||||
key := dailyKey{date: event.RecordedAt.UTC().Format("2006-01-02"), apiKeyID: *event.APIKeyID, provider: event.ProviderCode, model: event.Model}
|
||||
value := daily[key]
|
||||
value.requests++
|
||||
if event.StatusCode >= 400 {
|
||||
value.failed++
|
||||
}
|
||||
value.prompt += max(event.PromptTokens, 0)
|
||||
value.completion += max(event.CompletionTokens, 0)
|
||||
value.cost += max(event.CostMicrounits, 0)
|
||||
daily[key] = value
|
||||
}
|
||||
}
|
||||
_, err = tx.CopyFrom(ctx, pgx.Identifier{"gateway", "audit_events"}, []string{
|
||||
"id", "tenant_id", "request_id", "actor_id", "api_key_id", "provider_code", "model", "protocol",
|
||||
"status_code", "prompt_tokens", "completion_tokens", "cost_microunits", "latency_ms",
|
||||
"request_preview", "response_preview", "labels", "recorded_at",
|
||||
}, pgx.CopyFromRows(rows))
|
||||
if err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
batch := &pgx.Batch{}
|
||||
for key, value := range daily {
|
||||
batch.Queue(`INSERT INTO gateway.usage_daily(usage_date,api_key_id,provider_code,model,requests,failed_requests,prompt_tokens,completion_tokens,cost_microunits)
|
||||
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9)
|
||||
ON CONFLICT(usage_date,api_key_id,provider_code,model) DO UPDATE SET
|
||||
requests=gateway.usage_daily.requests+EXCLUDED.requests,
|
||||
failed_requests=gateway.usage_daily.failed_requests+EXCLUDED.failed_requests,
|
||||
prompt_tokens=gateway.usage_daily.prompt_tokens+EXCLUDED.prompt_tokens,
|
||||
completion_tokens=gateway.usage_daily.completion_tokens+EXCLUDED.completion_tokens,
|
||||
cost_microunits=gateway.usage_daily.cost_microunits+EXCLUDED.cost_microunits,
|
||||
updated_at=clock_timestamp()`, key.date, key.apiKeyID, key.provider, key.model, value.requests, value.failed, value.prompt, value.completion, value.cost)
|
||||
}
|
||||
for _, alert := range alerts {
|
||||
batch.Queue(`INSERT INTO gateway.outbox_events(event_id,event_type,event_version,tenant_id,aggregate_type,aggregate_id,payload)
|
||||
VALUES($1,'content_policy.matched',1,$2,'gateway_request',$3,$4)`, alert.eventID, uuidPointer(alert.tenantID), alert.requestID, alert.payload)
|
||||
}
|
||||
results := tx.SendBatch(ctx, batch)
|
||||
if err := results.Close(); err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
r.written.Add(uint64(len(events)))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Recorder) flushError(err error) error {
|
||||
r.failures.Add(1)
|
||||
if r.logger != nil {
|
||||
r.logger.Error("audit batch flush failed", "error", err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func uuidValue(value string) pgtype.UUID {
|
||||
var result pgtype.UUID
|
||||
_ = result.Scan(value)
|
||||
return result
|
||||
}
|
||||
|
||||
func uuidPointer(value *string) any {
|
||||
if value == nil || strings.TrimSpace(*value) == "" {
|
||||
return nil
|
||||
}
|
||||
return uuidValue(*value)
|
||||
}
|
||||
|
||||
func nullString(value string) any {
|
||||
if value == "" {
|
||||
return nil
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func nullToken(value int64) any {
|
||||
if value <= 0 {
|
||||
return nil
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func nullCost(value int64) any {
|
||||
if value <= 0 {
|
||||
return nil
|
||||
}
|
||||
return value
|
||||
}
|
||||
Reference in New Issue
Block a user