easyai-ai-gateway/apps/api/internal/runner/recording.go

package runner

import (
	"strings"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type taskRecordDetails struct {
	RequestID          string
	ResolvedModel      string
	Usage              map[string]any
	Metrics            map[string]any
	BillingSummary     map[string]any
	FinalChargeAmount  float64
	ResponseStartedAt  time.Time
	ResponseFinishedAt time.Time
	ResponseDurationMS int64
}

func buildSuccessRecord(task store.GatewayTask, user *auth.User, body map[string]any, candidate store.RuntimeModelCandidate, response clients.Response, billings []any, simulated bool) taskRecordDetails {
	usage := usageToMap(response.Usage)
	metrics := taskMetrics(task, user, body, candidate, response, simulated)
	summary := summarizeBillings(billings, simulated)
	finalAmount := floatFromAny(summary["totalAmount"])
	return taskRecordDetails{
		RequestID:          response.RequestID,
		ResolvedModel:      candidate.ModelName,
		Usage:              usage,
		Metrics:            metrics,
		BillingSummary:     summary,
		FinalChargeAmount:  finalAmount,
		ResponseStartedAt:  response.ResponseStartedAt,
		ResponseFinishedAt: response.ResponseFinishedAt,
		ResponseDurationMS: response.ResponseDurationMS,
	}
}

func taskMetrics(task store.GatewayTask, user *auth.User, body map[string]any, candidate store.RuntimeModelCandidate, response clients.Response, simulated bool) map[string]any {
	metrics := map[string]any{
		"kind":            task.Kind,
		"runMode":         task.RunMode,
		"requestedModel":  task.Model,
		"resolvedModel":   candidate.ModelName,
		"modelAlias":      candidate.ModelAlias,
		"providerModel":   candidate.ProviderModelName,
		"canonicalModel":  candidate.CanonicalModelKey,
		"modelType":       candidate.ModelType,
		"provider":        candidate.Provider,
		"platformId":      candidate.PlatformID,
		"platformName":    candidate.PlatformName,
		"platformModelId": candidate.PlatformModelID,
		"clientId":        candidate.ClientID,
		"queueKey":        candidate.QueueKey,
		"requestId":       response.RequestID,
		"simulated":       simulated,
	}
	if user != nil {
		metrics["apiKeyId"] = user.APIKeyID
		metrics["apiKeyName"] = user.APIKeyName
		metrics["apiKeyPrefix"] = user.APIKeyPrefix
	}
	if response.ResponseDurationMS > 0 {
		metrics["responseDurationMs"] = response.ResponseDurationMS
	}
	switch task.Kind {
	case "chat.completions", "responses":
		metrics["stream"] = boolFromMap(body, "stream")
		metrics["messageCount"] = messageCount(body)
		copyIfPresent(metrics, body, "temperature")
		copyIfPresent(metrics, body, "max_tokens")
		copyIfPresent(metrics, body, "max_output_tokens")
	case "images.generations", "images.edits":
		metrics["imageCount"] = requestedCount(body)
		metrics["outputImageCount"] = outputDataCount(response.Result)
		metrics["inputImageCount"] = imageInputCount(body)
		metrics["hasMask"] = stringFromMap(body, "mask") != ""
		copyIfPresent(metrics, body, "size")
		copyIfPresent(metrics, body, "quality")
		copyIfPresent(metrics, body, "style")
	case "videos.generations":
		metrics["hasReferenceImage"] = imageInputCount(body) > 0
		metrics["hasReferenceVideo"] = hasAnyString(body, "video", "video_url", "videoUrl", "reference_video", "referenceVideo")
		copyIfPresent(metrics, body, "duration")
		copyIfPresent(metrics, body, "resolution")
		copyIfPresent(metrics, body, "size")
		copyIfPresent(metrics, body, "aspect_ratio")
		copyIfPresent(metrics, body, "aspectRatio")
		copyIfPresent(metrics, body, "fps")
	}
	return metrics
}

func usageToMap(usage clients.Usage) map[string]any {
	out := map[string]any{}
	if usage.InputTokens > 0 {
		out["inputTokens"] = usage.InputTokens
		out["promptTokens"] = usage.InputTokens
	}
	if usage.OutputTokens > 0 {
		out["outputTokens"] = usage.OutputTokens
		out["completionTokens"] = usage.OutputTokens
	}
	if usage.TotalTokens > 0 {
		out["totalTokens"] = usage.TotalTokens
	}
	return out
}

func summarizeBillings(billings []any, simulated bool) map[string]any {
	amountByCurrency := map[string]float64{}
	for _, raw := range billings {
		line, _ := raw.(map[string]any)
		if line == nil {
			continue
		}
		currency := strings.TrimSpace(stringFromAny(line["currency"]))
		if currency == "" {
			currency = "resource"
		}
		amountByCurrency[currency] = roundPrice(amountByCurrency[currency] + floatFromAny(line["amount"]))
	}
	currency := ""
	totalAmount := 0.0
	for key, amount := range amountByCurrency {
		if currency == "" {
			currency = key
		} else if currency != key {
			currency = "mixed"
		}
		totalAmount += amount
	}
	if currency == "" {
		currency = "resource"
	}
	totalAmount = roundPrice(totalAmount)
	return map[string]any{
		"lineCount":        len(billings),
		"totalAmount":      totalAmount,
		"amountByCurrency": amountByCurrency,
		"currency":         currency,
		"simulated":        simulated,
		"finalCharge": map[string]any{
			"amount":    totalAmount,
			"currency":  currency,
			"simulated": simulated,
		},
	}
}

func failureMetrics(err error, simulated bool) (string, map[string]any, time.Time, time.Time, int64) {
	meta := clients.ErrorResponseMetadata(err)
	metrics := map[string]any{
		"simulated": simulated,
	}
	if err != nil {
		metrics["error"] = err.Error()
		metrics["retryable"] = clients.IsRetryable(err)
	}
	if meta.StatusCode > 0 {
		metrics["statusCode"] = meta.StatusCode
	}
	if meta.RequestID != "" {
		metrics["requestId"] = meta.RequestID
	}
	if meta.ResponseDurationMS > 0 {
		metrics["responseDurationMs"] = meta.ResponseDurationMS
	}
	return meta.RequestID, metrics, meta.ResponseStartedAt, meta.ResponseFinishedAt, meta.ResponseDurationMS
}

func messageCount(body map[string]any) int {
	messages, _ := body["messages"].([]any)
	return len(messages)
}

func requestedCount(body map[string]any) int {
	count := int(floatFromAny(body["n"]))
	if count <= 0 {
		return 1
	}
	return count
}

func outputDataCount(result map[string]any) int {
	data, _ := result["data"].([]any)
	return len(data)
}

func imageInputCount(body map[string]any) int {
	count := 0
	for _, key := range []string{"image", "image_url", "imageUrl"} {
		if stringFromMap(body, key) != "" {
			count++
		}
	}
	for _, key := range []string{"image_urls", "imageUrls", "images"} {
		if values, ok := body[key].([]any); ok {
			count += len(values)
		}
	}
	return count
}

func hasAnyString(body map[string]any, keys ...string) bool {
	for _, key := range keys {
		if stringFromMap(body, key) != "" {
			return true
		}
	}
	return false
}

func copyIfPresent(target map[string]any, body map[string]any, key string) {
	if value, ok := body[key]; ok && value != nil {
		target[key] = value
	}
}