1363 lines
34 KiB
Go
1363 lines
34 KiB
Go
package repository
|
||||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"database/sql"
|
|||
|
|
"encoding/json"
|
|||
|
|
"errors"
|
|||
|
|
"fmt"
|
|||
|
|
"strconv"
|
|||
|
|
"strings"
|
|||
|
|
"sync/atomic"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
dbent "github.com/Wei-Shaw/sub2api/ent"
|
|||
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
|||
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// usageLogInsertArgTypes must stay in the same order as:
|
|||
|
|
// 1. prepareUsageLogInsert().args
|
|||
|
|
// 2. every INSERT/CTE VALUES column list in this file
|
|||
|
|
// 3. execUsageLogInsertNoResult placeholder positions
|
|||
|
|
// 4. scanUsageLog selected column order (via usageLogSelectColumns)
|
|||
|
|
//
|
|||
|
|
// When adding a usage_logs column, update all of those call sites together.
|
|||
|
|
var usageLogInsertArgTypes = [...]string{
|
|||
|
|
"bigint", // user_id
|
|||
|
|
"bigint", // api_key_id
|
|||
|
|
"bigint", // account_id
|
|||
|
|
"text", // request_id
|
|||
|
|
"text", // model
|
|||
|
|
"text", // requested_model
|
|||
|
|
"text", // upstream_model
|
|||
|
|
"text", // upstream_response_model
|
|||
|
|
"boolean", // upstream_model_mismatch
|
|||
|
|
"bigint", // group_id
|
|||
|
|
"bigint", // subscription_id
|
|||
|
|
"integer", // input_tokens
|
|||
|
|
"integer", // output_tokens
|
|||
|
|
"integer", // cache_creation_tokens
|
|||
|
|
"integer", // cache_read_tokens
|
|||
|
|
"integer", // cache_creation_5m_tokens
|
|||
|
|
"integer", // cache_creation_1h_tokens
|
|||
|
|
"integer", // image_output_tokens
|
|||
|
|
"numeric", // image_output_cost
|
|||
|
|
"integer", // image_input_tokens
|
|||
|
|
"numeric", // image_input_cost
|
|||
|
|
"numeric", // input_cost
|
|||
|
|
"numeric", // output_cost
|
|||
|
|
"numeric", // cache_creation_cost
|
|||
|
|
"numeric", // cache_read_cost
|
|||
|
|
"numeric", // total_cost
|
|||
|
|
"numeric", // actual_cost
|
|||
|
|
"numeric", // rate_multiplier
|
|||
|
|
"numeric", // account_rate_multiplier
|
|||
|
|
"smallint", // billing_type
|
|||
|
|
"smallint", // request_type
|
|||
|
|
"boolean", // stream
|
|||
|
|
"boolean", // openai_ws_mode
|
|||
|
|
"integer", // duration_ms
|
|||
|
|
"integer", // first_token_ms
|
|||
|
|
"text", // user_agent
|
|||
|
|
"text", // ip_address
|
|||
|
|
"integer", // image_count
|
|||
|
|
"text", // image_size
|
|||
|
|
"text", // image_input_size
|
|||
|
|
"text", // image_output_size
|
|||
|
|
"text", // image_size_source
|
|||
|
|
"jsonb", // image_size_breakdown
|
|||
|
|
"integer", // video_count
|
|||
|
|
"text", // video_resolution
|
|||
|
|
"integer", // video_duration_seconds
|
|||
|
|
"text", // service_tier
|
|||
|
|
"text", // reasoning_effort
|
|||
|
|
"text", // inbound_endpoint
|
|||
|
|
"text", // upstream_endpoint
|
|||
|
|
"boolean", // cache_ttl_overridden
|
|||
|
|
"boolean", // long_context_billing_applied
|
|||
|
|
"bigint", // channel_id
|
|||
|
|
"text", // model_mapping_chain
|
|||
|
|
"text", // billing_tier
|
|||
|
|
"text", // billing_mode
|
|||
|
|
"numeric", // account_stats_cost
|
|||
|
|
"text", // session_id
|
|||
|
|
"timestamptz", // created_at
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
usageLogCreateBatchMaxSize = 64
|
|||
|
|
usageLogCreateBatchWindow = 3 * time.Millisecond
|
|||
|
|
usageLogCreateBatchQueueCap = 4096
|
|||
|
|
usageLogCreateCancelWait = 2 * time.Second
|
|||
|
|
|
|||
|
|
usageLogBestEffortBatchMaxSize = 256
|
|||
|
|
usageLogBestEffortBatchWindow = 20 * time.Millisecond
|
|||
|
|
usageLogBestEffortBatchQueueCap = 32768
|
|||
|
|
usageLogBestEffortRecentTTL = 30 * time.Second
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
type usageLogCreateRequest struct {
|
|||
|
|
log *service.UsageLog
|
|||
|
|
prepared usageLogInsertPrepared
|
|||
|
|
shared *usageLogCreateShared
|
|||
|
|
resultCh chan usageLogCreateResult
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type usageLogCreateResult struct {
|
|||
|
|
inserted bool
|
|||
|
|
err error
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type usageLogBestEffortRequest struct {
|
|||
|
|
prepared usageLogInsertPrepared
|
|||
|
|
apiKeyID int64
|
|||
|
|
resultCh chan error
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type usageLogInsertPrepared struct {
|
|||
|
|
createdAt time.Time
|
|||
|
|
requestID string
|
|||
|
|
rateMultiplier float64
|
|||
|
|
requestType int16
|
|||
|
|
args []any
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type usageLogBatchState struct {
|
|||
|
|
ID int64
|
|||
|
|
CreatedAt time.Time
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type usageLogBatchRow struct {
|
|||
|
|
RequestID string `json:"request_id"`
|
|||
|
|
APIKeyID int64 `json:"api_key_id"`
|
|||
|
|
ID int64 `json:"id"`
|
|||
|
|
CreatedAt time.Time `json:"created_at"`
|
|||
|
|
Inserted bool `json:"inserted"`
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type usageLogCreateShared struct {
|
|||
|
|
state atomic.Int32
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
usageLogCreateStateQueued int32 = iota
|
|||
|
|
usageLogCreateStateProcessing
|
|||
|
|
usageLogCreateStateCompleted
|
|||
|
|
usageLogCreateStateCanceled
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) Create(ctx context.Context, log *service.UsageLog) (bool, error) {
|
|||
|
|
if log == nil {
|
|||
|
|
return false, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if tx := dbent.TxFromContext(ctx); tx != nil {
|
|||
|
|
return r.createSingle(ctx, tx.Client(), log)
|
|||
|
|
}
|
|||
|
|
requestID := strings.TrimSpace(log.RequestID)
|
|||
|
|
if requestID == "" {
|
|||
|
|
return r.createSingle(ctx, r.sql, log)
|
|||
|
|
}
|
|||
|
|
log.RequestID = requestID
|
|||
|
|
return r.createBatched(ctx, log)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) CreateBestEffort(ctx context.Context, log *service.UsageLog) error {
|
|||
|
|
if log == nil {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if tx := dbent.TxFromContext(ctx); tx != nil {
|
|||
|
|
_, err := r.createSingle(ctx, tx.Client(), log)
|
|||
|
|
return err
|
|||
|
|
}
|
|||
|
|
if r.db == nil {
|
|||
|
|
_, err := r.createSingle(ctx, r.sql, log)
|
|||
|
|
return err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
r.ensureBestEffortBatcher()
|
|||
|
|
if r.bestEffortBatchCh == nil {
|
|||
|
|
_, err := r.createSingle(ctx, r.sql, log)
|
|||
|
|
return err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
req := usageLogBestEffortRequest{
|
|||
|
|
prepared: prepareUsageLogInsert(log),
|
|||
|
|
apiKeyID: log.APIKeyID,
|
|||
|
|
resultCh: make(chan error, 1),
|
|||
|
|
}
|
|||
|
|
if key, ok := r.bestEffortRecentKey(req.prepared.requestID, req.apiKeyID); ok {
|
|||
|
|
if _, exists := r.bestEffortRecent.Get(key); exists {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 队列满时阻塞等待而非立即丢弃:批处理器持续排空队列,短暂等待即可入队。
|
|||
|
|
// 立即丢弃会造成“已扣费但无 usage_log”的永久数据缺口(issue #3656);
|
|||
|
|
// 阻塞上限由调用方 ctx 期限约束,超时后由上层同步兜底。
|
|||
|
|
select {
|
|||
|
|
case r.bestEffortBatchCh <- req:
|
|||
|
|
case <-ctx.Done():
|
|||
|
|
return service.MarkUsageLogCreateDropped(ctx.Err())
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
select {
|
|||
|
|
case err := <-req.resultCh:
|
|||
|
|
return err
|
|||
|
|
case <-ctx.Done():
|
|||
|
|
return service.MarkUsageLogCreateDropped(ctx.Err())
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, log *service.UsageLog) (bool, error) {
|
|||
|
|
prepared := prepareUsageLogInsert(log)
|
|||
|
|
if sqlq == nil {
|
|||
|
|
sqlq = r.sql
|
|||
|
|
}
|
|||
|
|
if ctx != nil && ctx.Err() != nil {
|
|||
|
|
return false, service.MarkUsageLogCreateNotPersisted(ctx.Err())
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
query := `
|
|||
|
|
INSERT INTO usage_logs (
|
|||
|
|
user_id,
|
|||
|
|
api_key_id,
|
|||
|
|
account_id,
|
|||
|
|
request_id,
|
|||
|
|
model,
|
|||
|
|
requested_model,
|
|||
|
|
upstream_model,
|
|||
|
|
upstream_response_model,
|
|||
|
|
upstream_model_mismatch,
|
|||
|
|
group_id,
|
|||
|
|
subscription_id,
|
|||
|
|
input_tokens,
|
|||
|
|
output_tokens,
|
|||
|
|
cache_creation_tokens,
|
|||
|
|
cache_read_tokens,
|
|||
|
|
cache_creation_5m_tokens,
|
|||
|
|
cache_creation_1h_tokens,
|
|||
|
|
image_output_tokens,
|
|||
|
|
image_output_cost,
|
|||
|
|
image_input_tokens,
|
|||
|
|
image_input_cost,
|
|||
|
|
input_cost,
|
|||
|
|
output_cost,
|
|||
|
|
cache_creation_cost,
|
|||
|
|
cache_read_cost,
|
|||
|
|
total_cost,
|
|||
|
|
actual_cost,
|
|||
|
|
rate_multiplier,
|
|||
|
|
account_rate_multiplier,
|
|||
|
|
billing_type,
|
|||
|
|
request_type,
|
|||
|
|
stream,
|
|||
|
|
openai_ws_mode,
|
|||
|
|
duration_ms,
|
|||
|
|
first_token_ms,
|
|||
|
|
user_agent,
|
|||
|
|
ip_address,
|
|||
|
|
image_count,
|
|||
|
|
image_size,
|
|||
|
|
image_input_size,
|
|||
|
|
image_output_size,
|
|||
|
|
image_size_source,
|
|||
|
|
image_size_breakdown,
|
|||
|
|
video_count,
|
|||
|
|
video_resolution,
|
|||
|
|
video_duration_seconds,
|
|||
|
|
service_tier,
|
|||
|
|
reasoning_effort,
|
|||
|
|
inbound_endpoint,
|
|||
|
|
upstream_endpoint,
|
|||
|
|
cache_ttl_overridden,
|
|||
|
|
long_context_billing_applied,
|
|||
|
|
channel_id,
|
|||
|
|
model_mapping_chain,
|
|||
|
|
billing_tier,
|
|||
|
|
billing_mode,
|
|||
|
|
account_stats_cost,
|
|||
|
|
session_id,
|
|||
|
|
created_at
|
|||
|
|
) VALUES (
|
|||
|
|
$1, $2, $3, $4, $5, $6, $7, $8, $9,
|
|||
|
|
$10, $11,
|
|||
|
|
$12, $13, $14, $15,
|
|||
|
|
$16, $17, $18, $19,
|
|||
|
|
$20, $21, $22, $23, $24, $25,
|
|||
|
|
$26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57, $58, $59
|
|||
|
|
)
|
|||
|
|
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
|||
|
|
RETURNING id, created_at
|
|||
|
|
`
|
|||
|
|
|
|||
|
|
if err := scanSingleRow(ctx, sqlq, query, prepared.args, &log.ID, &log.CreatedAt); err != nil {
|
|||
|
|
if errors.Is(err, sql.ErrNoRows) && prepared.requestID != "" {
|
|||
|
|
selectQuery := "SELECT id, created_at FROM usage_logs WHERE request_id = $1 AND api_key_id = $2"
|
|||
|
|
if err := scanSingleRow(ctx, sqlq, selectQuery, []any{prepared.requestID, log.APIKeyID}, &log.ID, &log.CreatedAt); err != nil {
|
|||
|
|
return false, err
|
|||
|
|
}
|
|||
|
|
log.RateMultiplier = prepared.rateMultiplier
|
|||
|
|
return false, nil
|
|||
|
|
} else {
|
|||
|
|
return false, err
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
log.RateMultiplier = prepared.rateMultiplier
|
|||
|
|
return true, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) createBatched(ctx context.Context, log *service.UsageLog) (bool, error) {
|
|||
|
|
if r.db == nil {
|
|||
|
|
return r.createSingle(ctx, r.sql, log)
|
|||
|
|
}
|
|||
|
|
r.ensureCreateBatcher()
|
|||
|
|
if r.createBatchCh == nil {
|
|||
|
|
return r.createSingle(ctx, r.sql, log)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
req := usageLogCreateRequest{
|
|||
|
|
log: log,
|
|||
|
|
prepared: prepareUsageLogInsert(log),
|
|||
|
|
shared: &usageLogCreateShared{},
|
|||
|
|
resultCh: make(chan usageLogCreateResult, 1),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 队列满时阻塞等待而非立即报错:本路径是 best-effort 丢弃后的最后兜底,
|
|||
|
|
// 立即失败会让日志永久丢失;阻塞上限由调用方 ctx 期限约束。
|
|||
|
|
select {
|
|||
|
|
case r.createBatchCh <- req:
|
|||
|
|
case <-ctx.Done():
|
|||
|
|
return false, service.MarkUsageLogCreateNotPersisted(ctx.Err())
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
select {
|
|||
|
|
case res := <-req.resultCh:
|
|||
|
|
return res.inserted, res.err
|
|||
|
|
case <-ctx.Done():
|
|||
|
|
if req.shared != nil && req.shared.state.CompareAndSwap(usageLogCreateStateQueued, usageLogCreateStateCanceled) {
|
|||
|
|
return false, service.MarkUsageLogCreateNotPersisted(ctx.Err())
|
|||
|
|
}
|
|||
|
|
timer := time.NewTimer(usageLogCreateCancelWait)
|
|||
|
|
defer timer.Stop()
|
|||
|
|
select {
|
|||
|
|
case res := <-req.resultCh:
|
|||
|
|
return res.inserted, res.err
|
|||
|
|
case <-timer.C:
|
|||
|
|
return false, ctx.Err()
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) ensureCreateBatcher() {
|
|||
|
|
if r == nil || r.db == nil {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
// nil 检查必须在 Once 内部:在外层做无同步快路径读会与 Once 内的写构成数据竞争。
|
|||
|
|
r.createBatchOnce.Do(func() {
|
|||
|
|
if r.createBatchCh == nil {
|
|||
|
|
r.createBatchCh = make(chan usageLogCreateRequest, usageLogCreateBatchQueueCap)
|
|||
|
|
go r.runCreateBatcher(r.db)
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) ensureBestEffortBatcher() {
|
|||
|
|
if r == nil || r.db == nil {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
// 同 ensureCreateBatcher:nil 检查放在 Once 内部以避免数据竞争。
|
|||
|
|
r.bestEffortBatchOnce.Do(func() {
|
|||
|
|
if r.bestEffortBatchCh == nil {
|
|||
|
|
r.bestEffortBatchCh = make(chan usageLogBestEffortRequest, usageLogBestEffortBatchQueueCap)
|
|||
|
|
go r.runBestEffortBatcher(r.db)
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) runCreateBatcher(db *sql.DB) {
|
|||
|
|
for {
|
|||
|
|
first, ok := <-r.createBatchCh
|
|||
|
|
if !ok {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
batch := make([]usageLogCreateRequest, 0, usageLogCreateBatchMaxSize)
|
|||
|
|
batch = append(batch, first)
|
|||
|
|
|
|||
|
|
timer := time.NewTimer(usageLogCreateBatchWindow)
|
|||
|
|
batchLoop:
|
|||
|
|
for len(batch) < usageLogCreateBatchMaxSize {
|
|||
|
|
select {
|
|||
|
|
case req, ok := <-r.createBatchCh:
|
|||
|
|
if !ok {
|
|||
|
|
break batchLoop
|
|||
|
|
}
|
|||
|
|
batch = append(batch, req)
|
|||
|
|
case <-timer.C:
|
|||
|
|
break batchLoop
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
if !timer.Stop() {
|
|||
|
|
select {
|
|||
|
|
case <-timer.C:
|
|||
|
|
default:
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
r.flushCreateBatch(db, batch)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) runBestEffortBatcher(db *sql.DB) {
|
|||
|
|
for {
|
|||
|
|
first, ok := <-r.bestEffortBatchCh
|
|||
|
|
if !ok {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
batch := make([]usageLogBestEffortRequest, 0, usageLogBestEffortBatchMaxSize)
|
|||
|
|
batch = append(batch, first)
|
|||
|
|
|
|||
|
|
timer := time.NewTimer(usageLogBestEffortBatchWindow)
|
|||
|
|
bestEffortLoop:
|
|||
|
|
for len(batch) < usageLogBestEffortBatchMaxSize {
|
|||
|
|
select {
|
|||
|
|
case req, ok := <-r.bestEffortBatchCh:
|
|||
|
|
if !ok {
|
|||
|
|
break bestEffortLoop
|
|||
|
|
}
|
|||
|
|
batch = append(batch, req)
|
|||
|
|
case <-timer.C:
|
|||
|
|
break bestEffortLoop
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
if !timer.Stop() {
|
|||
|
|
select {
|
|||
|
|
case <-timer.C:
|
|||
|
|
default:
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
r.flushBestEffortBatch(db, batch)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) flushCreateBatch(db *sql.DB, batch []usageLogCreateRequest) {
|
|||
|
|
if len(batch) == 0 {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
uniqueOrder := make([]string, 0, len(batch))
|
|||
|
|
preparedByKey := make(map[string]usageLogInsertPrepared, len(batch))
|
|||
|
|
requestsByKey := make(map[string][]usageLogCreateRequest, len(batch))
|
|||
|
|
fallback := make([]usageLogCreateRequest, 0)
|
|||
|
|
|
|||
|
|
for _, req := range batch {
|
|||
|
|
if req.log == nil {
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil})
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
if req.shared != nil && !req.shared.state.CompareAndSwap(usageLogCreateStateQueued, usageLogCreateStateProcessing) {
|
|||
|
|
if req.shared.state.Load() == usageLogCreateStateCanceled {
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{
|
|||
|
|
inserted: false,
|
|||
|
|
err: service.MarkUsageLogCreateNotPersisted(context.Canceled),
|
|||
|
|
})
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
prepared := req.prepared
|
|||
|
|
if prepared.requestID == "" {
|
|||
|
|
fallback = append(fallback, req)
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
key := usageLogBatchKey(prepared.requestID, req.log.APIKeyID)
|
|||
|
|
if _, exists := requestsByKey[key]; !exists {
|
|||
|
|
uniqueOrder = append(uniqueOrder, key)
|
|||
|
|
preparedByKey[key] = prepared
|
|||
|
|
}
|
|||
|
|
requestsByKey[key] = append(requestsByKey[key], req)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if len(uniqueOrder) > 0 {
|
|||
|
|
insertedMap, stateMap, safeFallback, err := r.batchInsertUsageLogs(db, uniqueOrder, preparedByKey)
|
|||
|
|
if err != nil {
|
|||
|
|
if safeFallback {
|
|||
|
|
for _, key := range uniqueOrder {
|
|||
|
|
fallback = append(fallback, requestsByKey[key]...)
|
|||
|
|
}
|
|||
|
|
} else {
|
|||
|
|
for _, key := range uniqueOrder {
|
|||
|
|
reqs := requestsByKey[key]
|
|||
|
|
state, hasState := stateMap[key]
|
|||
|
|
inserted := insertedMap[key]
|
|||
|
|
for idx, req := range reqs {
|
|||
|
|
req.log.RateMultiplier = preparedByKey[key].rateMultiplier
|
|||
|
|
if hasState {
|
|||
|
|
req.log.ID = state.ID
|
|||
|
|
req.log.CreatedAt = state.CreatedAt
|
|||
|
|
}
|
|||
|
|
switch {
|
|||
|
|
case inserted && idx == 0:
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: true, err: nil})
|
|||
|
|
case inserted:
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil})
|
|||
|
|
case hasState:
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil})
|
|||
|
|
case idx == 0:
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: err})
|
|||
|
|
default:
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil})
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
} else {
|
|||
|
|
for _, key := range uniqueOrder {
|
|||
|
|
reqs := requestsByKey[key]
|
|||
|
|
state, ok := stateMap[key]
|
|||
|
|
if !ok {
|
|||
|
|
for _, req := range reqs {
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{
|
|||
|
|
inserted: false,
|
|||
|
|
err: fmt.Errorf("usage log batch state missing for key=%s", key),
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
for idx, req := range reqs {
|
|||
|
|
req.log.ID = state.ID
|
|||
|
|
req.log.CreatedAt = state.CreatedAt
|
|||
|
|
req.log.RateMultiplier = preparedByKey[key].rateMultiplier
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{
|
|||
|
|
inserted: idx == 0 && insertedMap[key],
|
|||
|
|
err: nil,
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if len(fallback) == 0 {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
for _, req := range fallback {
|
|||
|
|
fallbackCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|||
|
|
inserted, err := r.createSingle(fallbackCtx, db, req.log)
|
|||
|
|
cancel()
|
|||
|
|
completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: inserted, err: err})
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) flushBestEffortBatch(db *sql.DB, batch []usageLogBestEffortRequest) {
|
|||
|
|
if len(batch) == 0 {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type bestEffortGroup struct {
|
|||
|
|
prepared usageLogInsertPrepared
|
|||
|
|
apiKeyID int64
|
|||
|
|
key string
|
|||
|
|
reqs []usageLogBestEffortRequest
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
groupsByKey := make(map[string]*bestEffortGroup, len(batch))
|
|||
|
|
groupOrder := make([]*bestEffortGroup, 0, len(batch))
|
|||
|
|
preparedList := make([]usageLogInsertPrepared, 0, len(batch))
|
|||
|
|
|
|||
|
|
for idx, req := range batch {
|
|||
|
|
prepared := req.prepared
|
|||
|
|
key := fmt.Sprintf("__best_effort_%d", idx)
|
|||
|
|
if prepared.requestID != "" {
|
|||
|
|
key = usageLogBatchKey(prepared.requestID, req.apiKeyID)
|
|||
|
|
}
|
|||
|
|
group, exists := groupsByKey[key]
|
|||
|
|
if !exists {
|
|||
|
|
group = &bestEffortGroup{
|
|||
|
|
prepared: prepared,
|
|||
|
|
apiKeyID: req.apiKeyID,
|
|||
|
|
key: key,
|
|||
|
|
}
|
|||
|
|
groupsByKey[key] = group
|
|||
|
|
groupOrder = append(groupOrder, group)
|
|||
|
|
preparedList = append(preparedList, prepared)
|
|||
|
|
}
|
|||
|
|
group.reqs = append(group.reqs, req)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if len(preparedList) == 0 {
|
|||
|
|
for _, req := range batch {
|
|||
|
|
sendUsageLogBestEffortResult(req.resultCh, nil)
|
|||
|
|
}
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|||
|
|
defer cancel()
|
|||
|
|
|
|||
|
|
query, args := buildUsageLogBestEffortInsertQuery(preparedList)
|
|||
|
|
if _, err := db.ExecContext(ctx, query, args...); err != nil {
|
|||
|
|
logger.LegacyPrintf("repository.usage_log", "best-effort batch insert failed: %v", err)
|
|||
|
|
for _, group := range groupOrder {
|
|||
|
|
singleErr := execUsageLogInsertNoResult(ctx, db, group.prepared)
|
|||
|
|
if singleErr != nil {
|
|||
|
|
logger.LegacyPrintf("repository.usage_log", "best-effort single fallback insert failed: %v", singleErr)
|
|||
|
|
} else if group.prepared.requestID != "" && r != nil && r.bestEffortRecent != nil {
|
|||
|
|
r.bestEffortRecent.SetDefault(group.key, struct{}{})
|
|||
|
|
}
|
|||
|
|
for _, req := range group.reqs {
|
|||
|
|
sendUsageLogBestEffortResult(req.resultCh, singleErr)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
for _, group := range groupOrder {
|
|||
|
|
if group.prepared.requestID != "" && r != nil && r.bestEffortRecent != nil {
|
|||
|
|
r.bestEffortRecent.SetDefault(group.key, struct{}{})
|
|||
|
|
}
|
|||
|
|
for _, req := range group.reqs {
|
|||
|
|
sendUsageLogBestEffortResult(req.resultCh, nil)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func sendUsageLogBestEffortResult(ch chan error, err error) {
|
|||
|
|
if ch == nil {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
select {
|
|||
|
|
case ch <- err:
|
|||
|
|
default:
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func completeUsageLogCreateRequest(req usageLogCreateRequest, res usageLogCreateResult) {
|
|||
|
|
if req.shared != nil {
|
|||
|
|
req.shared.state.Store(usageLogCreateStateCompleted)
|
|||
|
|
}
|
|||
|
|
sendUsageLogCreateResult(req.resultCh, res)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) batchInsertUsageLogs(db *sql.DB, keys []string, preparedByKey map[string]usageLogInsertPrepared) (map[string]bool, map[string]usageLogBatchState, bool, error) {
|
|||
|
|
if len(keys) == 0 {
|
|||
|
|
return map[string]bool{}, map[string]usageLogBatchState{}, false, nil
|
|||
|
|
}
|
|||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|||
|
|
defer cancel()
|
|||
|
|
|
|||
|
|
query, args := buildUsageLogBatchInsertQuery(keys, preparedByKey)
|
|||
|
|
var payload []byte
|
|||
|
|
if err := db.QueryRowContext(ctx, query, args...).Scan(&payload); err != nil {
|
|||
|
|
return nil, nil, true, err
|
|||
|
|
}
|
|||
|
|
var rows []usageLogBatchRow
|
|||
|
|
if err := json.Unmarshal(payload, &rows); err != nil {
|
|||
|
|
return nil, nil, false, err
|
|||
|
|
}
|
|||
|
|
insertedMap := make(map[string]bool, len(keys))
|
|||
|
|
stateMap := make(map[string]usageLogBatchState, len(keys))
|
|||
|
|
for _, row := range rows {
|
|||
|
|
key := usageLogBatchKey(row.RequestID, row.APIKeyID)
|
|||
|
|
insertedMap[key] = row.Inserted
|
|||
|
|
stateMap[key] = usageLogBatchState{
|
|||
|
|
ID: row.ID,
|
|||
|
|
CreatedAt: row.CreatedAt,
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
if len(stateMap) != len(keys) {
|
|||
|
|
return insertedMap, stateMap, false, fmt.Errorf("usage log batch state count mismatch: got=%d want=%d", len(stateMap), len(keys))
|
|||
|
|
}
|
|||
|
|
return insertedMap, stateMap, false, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usageLogInsertPrepared) (string, []any) {
|
|||
|
|
var query strings.Builder
|
|||
|
|
_, _ = query.WriteString(`
|
|||
|
|
WITH input (
|
|||
|
|
input_idx,
|
|||
|
|
user_id,
|
|||
|
|
api_key_id,
|
|||
|
|
account_id,
|
|||
|
|
request_id,
|
|||
|
|
model,
|
|||
|
|
requested_model,
|
|||
|
|
upstream_model,
|
|||
|
|
upstream_response_model,
|
|||
|
|
upstream_model_mismatch,
|
|||
|
|
group_id,
|
|||
|
|
subscription_id,
|
|||
|
|
input_tokens,
|
|||
|
|
output_tokens,
|
|||
|
|
cache_creation_tokens,
|
|||
|
|
cache_read_tokens,
|
|||
|
|
cache_creation_5m_tokens,
|
|||
|
|
cache_creation_1h_tokens,
|
|||
|
|
image_output_tokens,
|
|||
|
|
image_output_cost,
|
|||
|
|
image_input_tokens,
|
|||
|
|
image_input_cost,
|
|||
|
|
input_cost,
|
|||
|
|
output_cost,
|
|||
|
|
cache_creation_cost,
|
|||
|
|
cache_read_cost,
|
|||
|
|
total_cost,
|
|||
|
|
actual_cost,
|
|||
|
|
rate_multiplier,
|
|||
|
|
account_rate_multiplier,
|
|||
|
|
billing_type,
|
|||
|
|
request_type,
|
|||
|
|
stream,
|
|||
|
|
openai_ws_mode,
|
|||
|
|
duration_ms,
|
|||
|
|
first_token_ms,
|
|||
|
|
user_agent,
|
|||
|
|
ip_address,
|
|||
|
|
image_count,
|
|||
|
|
image_size,
|
|||
|
|
image_input_size,
|
|||
|
|
image_output_size,
|
|||
|
|
image_size_source,
|
|||
|
|
image_size_breakdown,
|
|||
|
|
video_count,
|
|||
|
|
video_resolution,
|
|||
|
|
video_duration_seconds,
|
|||
|
|
service_tier,
|
|||
|
|
reasoning_effort,
|
|||
|
|
inbound_endpoint,
|
|||
|
|
upstream_endpoint,
|
|||
|
|
cache_ttl_overridden,
|
|||
|
|
long_context_billing_applied,
|
|||
|
|
channel_id,
|
|||
|
|
model_mapping_chain,
|
|||
|
|
billing_tier,
|
|||
|
|
billing_mode,
|
|||
|
|
account_stats_cost,
|
|||
|
|
session_id,
|
|||
|
|
created_at
|
|||
|
|
) AS (VALUES `)
|
|||
|
|
|
|||
|
|
// Each batch row prepends the synthetic input_index before the 59
|
|||
|
|
// usage-log column values.
|
|||
|
|
args := make([]any, 0, len(keys)*60)
|
|||
|
|
argPos := 1
|
|||
|
|
for idx, key := range keys {
|
|||
|
|
if idx > 0 {
|
|||
|
|
_, _ = query.WriteString(",")
|
|||
|
|
}
|
|||
|
|
_, _ = query.WriteString("(")
|
|||
|
|
_, _ = query.WriteString("$")
|
|||
|
|
_, _ = query.WriteString(strconv.Itoa(argPos))
|
|||
|
|
args = append(args, idx)
|
|||
|
|
argPos++
|
|||
|
|
prepared := preparedByKey[key]
|
|||
|
|
for i := 0; i < len(prepared.args); i++ {
|
|||
|
|
_, _ = query.WriteString(",")
|
|||
|
|
_, _ = query.WriteString("$")
|
|||
|
|
_, _ = query.WriteString(strconv.Itoa(argPos))
|
|||
|
|
if i < len(usageLogInsertArgTypes) {
|
|||
|
|
_, _ = query.WriteString("::")
|
|||
|
|
_, _ = query.WriteString(usageLogInsertArgTypes[i])
|
|||
|
|
}
|
|||
|
|
argPos++
|
|||
|
|
}
|
|||
|
|
_, _ = query.WriteString(")")
|
|||
|
|
args = append(args, prepared.args...)
|
|||
|
|
}
|
|||
|
|
_, _ = query.WriteString(`
|
|||
|
|
),
|
|||
|
|
inserted AS (
|
|||
|
|
INSERT INTO usage_logs (
|
|||
|
|
user_id,
|
|||
|
|
api_key_id,
|
|||
|
|
account_id,
|
|||
|
|
request_id,
|
|||
|
|
model,
|
|||
|
|
requested_model,
|
|||
|
|
upstream_model,
|
|||
|
|
upstream_response_model,
|
|||
|
|
upstream_model_mismatch,
|
|||
|
|
group_id,
|
|||
|
|
subscription_id,
|
|||
|
|
input_tokens,
|
|||
|
|
output_tokens,
|
|||
|
|
cache_creation_tokens,
|
|||
|
|
cache_read_tokens,
|
|||
|
|
cache_creation_5m_tokens,
|
|||
|
|
cache_creation_1h_tokens,
|
|||
|
|
image_output_tokens,
|
|||
|
|
image_output_cost,
|
|||
|
|
image_input_tokens,
|
|||
|
|
image_input_cost,
|
|||
|
|
input_cost,
|
|||
|
|
output_cost,
|
|||
|
|
cache_creation_cost,
|
|||
|
|
cache_read_cost,
|
|||
|
|
total_cost,
|
|||
|
|
actual_cost,
|
|||
|
|
rate_multiplier,
|
|||
|
|
account_rate_multiplier,
|
|||
|
|
billing_type,
|
|||
|
|
request_type,
|
|||
|
|
stream,
|
|||
|
|
openai_ws_mode,
|
|||
|
|
duration_ms,
|
|||
|
|
first_token_ms,
|
|||
|
|
user_agent,
|
|||
|
|
ip_address,
|
|||
|
|
image_count,
|
|||
|
|
image_size,
|
|||
|
|
image_input_size,
|
|||
|
|
image_output_size,
|
|||
|
|
image_size_source,
|
|||
|
|
image_size_breakdown,
|
|||
|
|
video_count,
|
|||
|
|
video_resolution,
|
|||
|
|
video_duration_seconds,
|
|||
|
|
service_tier,
|
|||
|
|
reasoning_effort,
|
|||
|
|
inbound_endpoint,
|
|||
|
|
upstream_endpoint,
|
|||
|
|
cache_ttl_overridden,
|
|||
|
|
long_context_billing_applied,
|
|||
|
|
channel_id,
|
|||
|
|
model_mapping_chain,
|
|||
|
|
billing_tier,
|
|||
|
|
billing_mode,
|
|||
|
|
account_stats_cost,
|
|||
|
|
session_id,
|
|||
|
|
created_at
|
|||
|
|
)
|
|||
|
|
SELECT
|
|||
|
|
user_id,
|
|||
|
|
api_key_id,
|
|||
|
|
account_id,
|
|||
|
|
request_id,
|
|||
|
|
model,
|
|||
|
|
requested_model,
|
|||
|
|
upstream_model,
|
|||
|
|
upstream_response_model,
|
|||
|
|
upstream_model_mismatch,
|
|||
|
|
group_id,
|
|||
|
|
subscription_id,
|
|||
|
|
input_tokens,
|
|||
|
|
output_tokens,
|
|||
|
|
cache_creation_tokens,
|
|||
|
|
cache_read_tokens,
|
|||
|
|
cache_creation_5m_tokens,
|
|||
|
|
cache_creation_1h_tokens,
|
|||
|
|
image_output_tokens,
|
|||
|
|
image_output_cost,
|
|||
|
|
image_input_tokens,
|
|||
|
|
image_input_cost,
|
|||
|
|
input_cost,
|
|||
|
|
output_cost,
|
|||
|
|
cache_creation_cost,
|
|||
|
|
cache_read_cost,
|
|||
|
|
total_cost,
|
|||
|
|
actual_cost,
|
|||
|
|
rate_multiplier,
|
|||
|
|
account_rate_multiplier,
|
|||
|
|
billing_type,
|
|||
|
|
request_type,
|
|||
|
|
stream,
|
|||
|
|
openai_ws_mode,
|
|||
|
|
duration_ms,
|
|||
|
|
first_token_ms,
|
|||
|
|
user_agent,
|
|||
|
|
ip_address,
|
|||
|
|
image_count,
|
|||
|
|
image_size,
|
|||
|
|
image_input_size,
|
|||
|
|
image_output_size,
|
|||
|
|
image_size_source,
|
|||
|
|
image_size_breakdown,
|
|||
|
|
video_count,
|
|||
|
|
video_resolution,
|
|||
|
|
video_duration_seconds,
|
|||
|
|
service_tier,
|
|||
|
|
reasoning_effort,
|
|||
|
|
inbound_endpoint,
|
|||
|
|
upstream_endpoint,
|
|||
|
|
cache_ttl_overridden,
|
|||
|
|
long_context_billing_applied,
|
|||
|
|
channel_id,
|
|||
|
|
model_mapping_chain,
|
|||
|
|
billing_tier,
|
|||
|
|
billing_mode,
|
|||
|
|
account_stats_cost,
|
|||
|
|
session_id,
|
|||
|
|
created_at
|
|||
|
|
FROM input
|
|||
|
|
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
|||
|
|
RETURNING request_id, api_key_id, id, created_at
|
|||
|
|
),
|
|||
|
|
resolved AS (
|
|||
|
|
SELECT
|
|||
|
|
input.input_idx,
|
|||
|
|
input.request_id,
|
|||
|
|
input.api_key_id,
|
|||
|
|
COALESCE(inserted.id, existing.id) AS id,
|
|||
|
|
COALESCE(inserted.created_at, existing.created_at) AS created_at,
|
|||
|
|
(inserted.id IS NOT NULL) AS inserted
|
|||
|
|
FROM input
|
|||
|
|
LEFT JOIN inserted
|
|||
|
|
ON inserted.request_id = input.request_id
|
|||
|
|
AND inserted.api_key_id = input.api_key_id
|
|||
|
|
LEFT JOIN usage_logs existing
|
|||
|
|
ON existing.request_id = input.request_id
|
|||
|
|
AND existing.api_key_id = input.api_key_id
|
|||
|
|
)
|
|||
|
|
SELECT COALESCE(
|
|||
|
|
json_agg(
|
|||
|
|
json_build_object(
|
|||
|
|
'request_id', resolved.request_id,
|
|||
|
|
'api_key_id', resolved.api_key_id,
|
|||
|
|
'id', resolved.id,
|
|||
|
|
'created_at', resolved.created_at,
|
|||
|
|
'inserted', resolved.inserted
|
|||
|
|
)
|
|||
|
|
ORDER BY resolved.input_idx
|
|||
|
|
),
|
|||
|
|
'[]'::json
|
|||
|
|
)
|
|||
|
|
FROM resolved
|
|||
|
|
`)
|
|||
|
|
return query.String(), args
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (string, []any) {
|
|||
|
|
var query strings.Builder
|
|||
|
|
_, _ = query.WriteString(`
|
|||
|
|
WITH input (
|
|||
|
|
user_id,
|
|||
|
|
api_key_id,
|
|||
|
|
account_id,
|
|||
|
|
request_id,
|
|||
|
|
model,
|
|||
|
|
requested_model,
|
|||
|
|
upstream_model,
|
|||
|
|
upstream_response_model,
|
|||
|
|
upstream_model_mismatch,
|
|||
|
|
group_id,
|
|||
|
|
subscription_id,
|
|||
|
|
input_tokens,
|
|||
|
|
output_tokens,
|
|||
|
|
cache_creation_tokens,
|
|||
|
|
cache_read_tokens,
|
|||
|
|
cache_creation_5m_tokens,
|
|||
|
|
cache_creation_1h_tokens,
|
|||
|
|
image_output_tokens,
|
|||
|
|
image_output_cost,
|
|||
|
|
image_input_tokens,
|
|||
|
|
image_input_cost,
|
|||
|
|
input_cost,
|
|||
|
|
output_cost,
|
|||
|
|
cache_creation_cost,
|
|||
|
|
cache_read_cost,
|
|||
|
|
total_cost,
|
|||
|
|
actual_cost,
|
|||
|
|
rate_multiplier,
|
|||
|
|
account_rate_multiplier,
|
|||
|
|
billing_type,
|
|||
|
|
request_type,
|
|||
|
|
stream,
|
|||
|
|
openai_ws_mode,
|
|||
|
|
duration_ms,
|
|||
|
|
first_token_ms,
|
|||
|
|
user_agent,
|
|||
|
|
ip_address,
|
|||
|
|
image_count,
|
|||
|
|
image_size,
|
|||
|
|
image_input_size,
|
|||
|
|
image_output_size,
|
|||
|
|
image_size_source,
|
|||
|
|
image_size_breakdown,
|
|||
|
|
video_count,
|
|||
|
|
video_resolution,
|
|||
|
|
video_duration_seconds,
|
|||
|
|
service_tier,
|
|||
|
|
reasoning_effort,
|
|||
|
|
inbound_endpoint,
|
|||
|
|
upstream_endpoint,
|
|||
|
|
cache_ttl_overridden,
|
|||
|
|
long_context_billing_applied,
|
|||
|
|
channel_id,
|
|||
|
|
model_mapping_chain,
|
|||
|
|
billing_tier,
|
|||
|
|
billing_mode,
|
|||
|
|
account_stats_cost,
|
|||
|
|
session_id,
|
|||
|
|
created_at
|
|||
|
|
) AS (VALUES `)
|
|||
|
|
|
|||
|
|
args := make([]any, 0, len(preparedList)*59)
|
|||
|
|
argPos := 1
|
|||
|
|
for idx, prepared := range preparedList {
|
|||
|
|
if idx > 0 {
|
|||
|
|
_, _ = query.WriteString(",")
|
|||
|
|
}
|
|||
|
|
_, _ = query.WriteString("(")
|
|||
|
|
for i := 0; i < len(prepared.args); i++ {
|
|||
|
|
if i > 0 {
|
|||
|
|
_, _ = query.WriteString(",")
|
|||
|
|
}
|
|||
|
|
_, _ = query.WriteString("$")
|
|||
|
|
_, _ = query.WriteString(strconv.Itoa(argPos))
|
|||
|
|
if i < len(usageLogInsertArgTypes) {
|
|||
|
|
_, _ = query.WriteString("::")
|
|||
|
|
_, _ = query.WriteString(usageLogInsertArgTypes[i])
|
|||
|
|
}
|
|||
|
|
argPos++
|
|||
|
|
}
|
|||
|
|
_, _ = query.WriteString(")")
|
|||
|
|
args = append(args, prepared.args...)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
_, _ = query.WriteString(`
|
|||
|
|
)
|
|||
|
|
INSERT INTO usage_logs (
|
|||
|
|
user_id,
|
|||
|
|
api_key_id,
|
|||
|
|
account_id,
|
|||
|
|
request_id,
|
|||
|
|
model,
|
|||
|
|
requested_model,
|
|||
|
|
upstream_model,
|
|||
|
|
upstream_response_model,
|
|||
|
|
upstream_model_mismatch,
|
|||
|
|
group_id,
|
|||
|
|
subscription_id,
|
|||
|
|
input_tokens,
|
|||
|
|
output_tokens,
|
|||
|
|
cache_creation_tokens,
|
|||
|
|
cache_read_tokens,
|
|||
|
|
cache_creation_5m_tokens,
|
|||
|
|
cache_creation_1h_tokens,
|
|||
|
|
image_output_tokens,
|
|||
|
|
image_output_cost,
|
|||
|
|
image_input_tokens,
|
|||
|
|
image_input_cost,
|
|||
|
|
input_cost,
|
|||
|
|
output_cost,
|
|||
|
|
cache_creation_cost,
|
|||
|
|
cache_read_cost,
|
|||
|
|
total_cost,
|
|||
|
|
actual_cost,
|
|||
|
|
rate_multiplier,
|
|||
|
|
account_rate_multiplier,
|
|||
|
|
billing_type,
|
|||
|
|
request_type,
|
|||
|
|
stream,
|
|||
|
|
openai_ws_mode,
|
|||
|
|
duration_ms,
|
|||
|
|
first_token_ms,
|
|||
|
|
user_agent,
|
|||
|
|
ip_address,
|
|||
|
|
image_count,
|
|||
|
|
image_size,
|
|||
|
|
image_input_size,
|
|||
|
|
image_output_size,
|
|||
|
|
image_size_source,
|
|||
|
|
image_size_breakdown,
|
|||
|
|
video_count,
|
|||
|
|
video_resolution,
|
|||
|
|
video_duration_seconds,
|
|||
|
|
service_tier,
|
|||
|
|
reasoning_effort,
|
|||
|
|
inbound_endpoint,
|
|||
|
|
upstream_endpoint,
|
|||
|
|
cache_ttl_overridden,
|
|||
|
|
long_context_billing_applied,
|
|||
|
|
channel_id,
|
|||
|
|
model_mapping_chain,
|
|||
|
|
billing_tier,
|
|||
|
|
billing_mode,
|
|||
|
|
account_stats_cost,
|
|||
|
|
session_id,
|
|||
|
|
created_at
|
|||
|
|
)
|
|||
|
|
SELECT
|
|||
|
|
user_id,
|
|||
|
|
api_key_id,
|
|||
|
|
account_id,
|
|||
|
|
request_id,
|
|||
|
|
model,
|
|||
|
|
requested_model,
|
|||
|
|
upstream_model,
|
|||
|
|
upstream_response_model,
|
|||
|
|
upstream_model_mismatch,
|
|||
|
|
group_id,
|
|||
|
|
subscription_id,
|
|||
|
|
input_tokens,
|
|||
|
|
output_tokens,
|
|||
|
|
cache_creation_tokens,
|
|||
|
|
cache_read_tokens,
|
|||
|
|
cache_creation_5m_tokens,
|
|||
|
|
cache_creation_1h_tokens,
|
|||
|
|
image_output_tokens,
|
|||
|
|
image_output_cost,
|
|||
|
|
image_input_tokens,
|
|||
|
|
image_input_cost,
|
|||
|
|
input_cost,
|
|||
|
|
output_cost,
|
|||
|
|
cache_creation_cost,
|
|||
|
|
cache_read_cost,
|
|||
|
|
total_cost,
|
|||
|
|
actual_cost,
|
|||
|
|
rate_multiplier,
|
|||
|
|
account_rate_multiplier,
|
|||
|
|
billing_type,
|
|||
|
|
request_type,
|
|||
|
|
stream,
|
|||
|
|
openai_ws_mode,
|
|||
|
|
duration_ms,
|
|||
|
|
first_token_ms,
|
|||
|
|
user_agent,
|
|||
|
|
ip_address,
|
|||
|
|
image_count,
|
|||
|
|
image_size,
|
|||
|
|
image_input_size,
|
|||
|
|
image_output_size,
|
|||
|
|
image_size_source,
|
|||
|
|
image_size_breakdown,
|
|||
|
|
video_count,
|
|||
|
|
video_resolution,
|
|||
|
|
video_duration_seconds,
|
|||
|
|
service_tier,
|
|||
|
|
reasoning_effort,
|
|||
|
|
inbound_endpoint,
|
|||
|
|
upstream_endpoint,
|
|||
|
|
cache_ttl_overridden,
|
|||
|
|
long_context_billing_applied,
|
|||
|
|
channel_id,
|
|||
|
|
model_mapping_chain,
|
|||
|
|
billing_tier,
|
|||
|
|
billing_mode,
|
|||
|
|
account_stats_cost,
|
|||
|
|
session_id,
|
|||
|
|
created_at
|
|||
|
|
FROM input
|
|||
|
|
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
|||
|
|
`)
|
|||
|
|
|
|||
|
|
return query.String(), args
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared usageLogInsertPrepared) error {
|
|||
|
|
_, err := sqlq.ExecContext(ctx, `
|
|||
|
|
INSERT INTO usage_logs (
|
|||
|
|
user_id,
|
|||
|
|
api_key_id,
|
|||
|
|
account_id,
|
|||
|
|
request_id,
|
|||
|
|
model,
|
|||
|
|
requested_model,
|
|||
|
|
upstream_model,
|
|||
|
|
upstream_response_model,
|
|||
|
|
upstream_model_mismatch,
|
|||
|
|
group_id,
|
|||
|
|
subscription_id,
|
|||
|
|
input_tokens,
|
|||
|
|
output_tokens,
|
|||
|
|
cache_creation_tokens,
|
|||
|
|
cache_read_tokens,
|
|||
|
|
cache_creation_5m_tokens,
|
|||
|
|
cache_creation_1h_tokens,
|
|||
|
|
image_output_tokens,
|
|||
|
|
image_output_cost,
|
|||
|
|
image_input_tokens,
|
|||
|
|
image_input_cost,
|
|||
|
|
input_cost,
|
|||
|
|
output_cost,
|
|||
|
|
cache_creation_cost,
|
|||
|
|
cache_read_cost,
|
|||
|
|
total_cost,
|
|||
|
|
actual_cost,
|
|||
|
|
rate_multiplier,
|
|||
|
|
account_rate_multiplier,
|
|||
|
|
billing_type,
|
|||
|
|
request_type,
|
|||
|
|
stream,
|
|||
|
|
openai_ws_mode,
|
|||
|
|
duration_ms,
|
|||
|
|
first_token_ms,
|
|||
|
|
user_agent,
|
|||
|
|
ip_address,
|
|||
|
|
image_count,
|
|||
|
|
image_size,
|
|||
|
|
image_input_size,
|
|||
|
|
image_output_size,
|
|||
|
|
image_size_source,
|
|||
|
|
image_size_breakdown,
|
|||
|
|
video_count,
|
|||
|
|
video_resolution,
|
|||
|
|
video_duration_seconds,
|
|||
|
|
service_tier,
|
|||
|
|
reasoning_effort,
|
|||
|
|
inbound_endpoint,
|
|||
|
|
upstream_endpoint,
|
|||
|
|
cache_ttl_overridden,
|
|||
|
|
long_context_billing_applied,
|
|||
|
|
channel_id,
|
|||
|
|
model_mapping_chain,
|
|||
|
|
billing_tier,
|
|||
|
|
billing_mode,
|
|||
|
|
account_stats_cost,
|
|||
|
|
session_id,
|
|||
|
|
created_at
|
|||
|
|
) VALUES (
|
|||
|
|
$1, $2, $3, $4, $5, $6, $7, $8, $9,
|
|||
|
|
$10, $11,
|
|||
|
|
$12, $13, $14, $15,
|
|||
|
|
$16, $17, $18, $19,
|
|||
|
|
$20, $21, $22, $23, $24, $25,
|
|||
|
|
$26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57, $58, $59
|
|||
|
|
)
|
|||
|
|
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
|||
|
|
`, prepared.args...)
|
|||
|
|
return err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
|
|||
|
|
createdAt := log.CreatedAt
|
|||
|
|
if createdAt.IsZero() {
|
|||
|
|
createdAt = time.Now()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
requestID := strings.TrimSpace(log.RequestID)
|
|||
|
|
log.RequestID = requestID
|
|||
|
|
|
|||
|
|
rateMultiplier := log.RateMultiplier
|
|||
|
|
log.SyncRequestTypeAndLegacyFields()
|
|||
|
|
requestType := int16(log.RequestType)
|
|||
|
|
|
|||
|
|
groupID := nullInt64(log.GroupID)
|
|||
|
|
subscriptionID := nullInt64(log.SubscriptionID)
|
|||
|
|
duration := nullInt(log.DurationMs)
|
|||
|
|
firstToken := nullInt(log.FirstTokenMs)
|
|||
|
|
userAgent := nullString(log.UserAgent)
|
|||
|
|
ipAddress := nullString(log.IPAddress)
|
|||
|
|
imageSize := nullString(log.ImageSize)
|
|||
|
|
imageInputSize := nullString(log.ImageInputSize)
|
|||
|
|
imageOutputSize := nullString(log.ImageOutputSize)
|
|||
|
|
imageSizeSource := nullString(log.ImageSizeSource)
|
|||
|
|
imageSizeBreakdown := nullStringIntMapJSON(log.ImageSizeBreakdown)
|
|||
|
|
videoResolution := nullString(log.VideoResolution)
|
|||
|
|
videoDurationSeconds := nullInt(log.VideoDurationSeconds)
|
|||
|
|
serviceTier := nullString(log.ServiceTier)
|
|||
|
|
reasoningEffort := nullString(log.ReasoningEffort)
|
|||
|
|
inboundEndpoint := nullString(log.InboundEndpoint)
|
|||
|
|
upstreamEndpoint := nullString(log.UpstreamEndpoint)
|
|||
|
|
channelID := nullInt64(log.ChannelID)
|
|||
|
|
modelMappingChain := nullString(log.ModelMappingChain)
|
|||
|
|
billingTier := nullString(log.BillingTier)
|
|||
|
|
billingMode := nullString(log.BillingMode)
|
|||
|
|
sessionID := nullString(log.SessionID)
|
|||
|
|
requestedModel := strings.TrimSpace(log.RequestedModel)
|
|||
|
|
if requestedModel == "" {
|
|||
|
|
requestedModel = strings.TrimSpace(log.Model)
|
|||
|
|
}
|
|||
|
|
upstreamModel := nullString(log.UpstreamModel)
|
|||
|
|
upstreamResponseModel := nullString(log.UpstreamResponseModel)
|
|||
|
|
upstreamModelMismatch := nullBool(log.UpstreamModelMismatch)
|
|||
|
|
|
|||
|
|
var requestIDArg any
|
|||
|
|
if requestID != "" {
|
|||
|
|
requestIDArg = requestID
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return usageLogInsertPrepared{
|
|||
|
|
createdAt: createdAt,
|
|||
|
|
requestID: requestID,
|
|||
|
|
rateMultiplier: rateMultiplier,
|
|||
|
|
requestType: requestType,
|
|||
|
|
args: []any{
|
|||
|
|
log.UserID,
|
|||
|
|
log.APIKeyID,
|
|||
|
|
log.AccountID,
|
|||
|
|
requestIDArg,
|
|||
|
|
log.Model,
|
|||
|
|
nullString(&requestedModel),
|
|||
|
|
upstreamModel,
|
|||
|
|
upstreamResponseModel,
|
|||
|
|
upstreamModelMismatch,
|
|||
|
|
groupID,
|
|||
|
|
subscriptionID,
|
|||
|
|
log.InputTokens,
|
|||
|
|
log.OutputTokens,
|
|||
|
|
log.CacheCreationTokens,
|
|||
|
|
log.CacheReadTokens,
|
|||
|
|
log.CacheCreation5mTokens,
|
|||
|
|
log.CacheCreation1hTokens,
|
|||
|
|
log.ImageOutputTokens,
|
|||
|
|
log.ImageOutputCost,
|
|||
|
|
log.ImageInputTokens,
|
|||
|
|
log.ImageInputCost,
|
|||
|
|
log.InputCost,
|
|||
|
|
log.OutputCost,
|
|||
|
|
log.CacheCreationCost,
|
|||
|
|
log.CacheReadCost,
|
|||
|
|
log.TotalCost,
|
|||
|
|
log.ActualCost,
|
|||
|
|
rateMultiplier,
|
|||
|
|
log.AccountRateMultiplier,
|
|||
|
|
log.BillingType,
|
|||
|
|
requestType,
|
|||
|
|
log.Stream,
|
|||
|
|
log.OpenAIWSMode,
|
|||
|
|
duration,
|
|||
|
|
firstToken,
|
|||
|
|
userAgent,
|
|||
|
|
ipAddress,
|
|||
|
|
log.ImageCount,
|
|||
|
|
imageSize,
|
|||
|
|
imageInputSize,
|
|||
|
|
imageOutputSize,
|
|||
|
|
imageSizeSource,
|
|||
|
|
imageSizeBreakdown,
|
|||
|
|
log.VideoCount,
|
|||
|
|
videoResolution,
|
|||
|
|
videoDurationSeconds,
|
|||
|
|
serviceTier,
|
|||
|
|
reasoningEffort,
|
|||
|
|
inboundEndpoint,
|
|||
|
|
upstreamEndpoint,
|
|||
|
|
log.CacheTTLOverridden,
|
|||
|
|
log.LongContextBillingApplied,
|
|||
|
|
channelID,
|
|||
|
|
modelMappingChain,
|
|||
|
|
billingTier,
|
|||
|
|
billingMode,
|
|||
|
|
log.AccountStatsCost, // account_stats_cost
|
|||
|
|
sessionID, // session_id
|
|||
|
|
createdAt,
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func usageLogBatchKey(requestID string, apiKeyID int64) string {
|
|||
|
|
return requestID + "\x1f" + strconv.FormatInt(apiKeyID, 10)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func sendUsageLogCreateResult(ch chan usageLogCreateResult, res usageLogCreateResult) {
|
|||
|
|
if ch == nil {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
select {
|
|||
|
|
case ch <- res:
|
|||
|
|
default:
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *usageLogRepository) bestEffortRecentKey(requestID string, apiKeyID int64) (string, bool) {
|
|||
|
|
requestID = strings.TrimSpace(requestID)
|
|||
|
|
if requestID == "" || r == nil || r.bestEffortRecent == nil {
|
|||
|
|
return "", false
|
|||
|
|
}
|
|||
|
|
return usageLogBatchKey(requestID, apiKeyID), true
|
|||
|
|
}
|