feat: enrich task record details
This commit is contained in:
@@ -0,0 +1,221 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user