refactor: move backend files to backend/ directory
Reorganize project structure: - backend/cmd/openteam/ — entry point - backend/internal/ — core packages - backend/middleware/ — HTTP middleware - backend/router/ — route setup - backend/wire/ — dependency injection - backend/pkg/ — shared utilities - backend/go.mod, go.sum — Go module files Updated Makefile to work from backend/ directory. Removed old lowercase makefile.
This commit is contained in:
@@ -0,0 +1,140 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Event represents a usage event to be recorded
|
||||
type Event struct {
|
||||
UserID uint64
|
||||
ModelName string
|
||||
ChannelID uint64
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
CacheReadTokens int
|
||||
Cost float64
|
||||
IsError bool
|
||||
IsCanceled bool
|
||||
RequestID string
|
||||
}
|
||||
|
||||
// Recorder handles async usage recording
|
||||
type Recorder struct {
|
||||
usageDAO *dao.UsageDAO
|
||||
dailyDAO *dao.DailyUsageDAO
|
||||
ch chan Event
|
||||
batchSize int
|
||||
flushInterval time.Duration
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewRecorder creates a new usage recorder
|
||||
func NewRecorder(usageDAO *dao.UsageDAO, dailyDAO *dao.DailyUsageDAO) *Recorder {
|
||||
return &Recorder{
|
||||
usageDAO: usageDAO,
|
||||
dailyDAO: dailyDAO,
|
||||
ch: make(chan Event, 10000),
|
||||
batchSize: 100,
|
||||
flushInterval: 5 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// Start starts the recorder's background workers
|
||||
func (r *Recorder) Start(ctx context.Context) {
|
||||
r.wg.Add(1)
|
||||
go r.processLoop(ctx)
|
||||
}
|
||||
|
||||
// Stop gracefully stops the recorder
|
||||
func (r *Recorder) Stop() {
|
||||
close(r.ch)
|
||||
r.wg.Wait()
|
||||
}
|
||||
|
||||
// Record queues a usage event for async recording
|
||||
func (r *Recorder) Record(event Event) {
|
||||
select {
|
||||
case r.ch <- event:
|
||||
default:
|
||||
log.Printf("Usage channel full, dropping event for user %d model %s", event.UserID, event.ModelName)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) processLoop(ctx context.Context) {
|
||||
defer r.wg.Done()
|
||||
|
||||
batch := make([]Event, 0, r.batchSize)
|
||||
ticker := time.NewTicker(r.flushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
if len(batch) > 0 {
|
||||
r.flush(batch)
|
||||
}
|
||||
return
|
||||
case event, ok := <-r.ch:
|
||||
if !ok {
|
||||
if len(batch) > 0 {
|
||||
r.flush(batch)
|
||||
}
|
||||
return
|
||||
}
|
||||
batch = append(batch, event)
|
||||
if len(batch) >= r.batchSize {
|
||||
r.flush(batch)
|
||||
batch = make([]Event, 0, r.batchSize)
|
||||
}
|
||||
case <-ticker.C:
|
||||
if len(batch) > 0 {
|
||||
r.flush(batch)
|
||||
batch = make([]Event, 0, r.batchSize)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) flush(events []Event) {
|
||||
if len(events) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Batch create usage logs
|
||||
logs := make([]*store.UsageLog, 0, len(events))
|
||||
|
||||
for _, e := range events {
|
||||
status := store.UsageStatusSuccess
|
||||
if e.IsError {
|
||||
status = store.UsageStatusError
|
||||
}
|
||||
if e.IsCanceled {
|
||||
status = store.UsageStatusCanceled
|
||||
}
|
||||
|
||||
log := &store.UsageLog{
|
||||
UserID: e.UserID,
|
||||
ModelName: e.ModelName,
|
||||
ChannelID: e.ChannelID,
|
||||
InputTokens: int64(e.PromptTokens),
|
||||
OutputTokens: int64(e.CompletionTokens),
|
||||
CacheReadTokens: int64(e.CacheReadTokens),
|
||||
Cost: e.Cost,
|
||||
Status: status,
|
||||
RequestID: e.RequestID,
|
||||
}
|
||||
logs = append(logs, log)
|
||||
}
|
||||
|
||||
// Write to database
|
||||
if err := r.usageDAO.BatchCreate(context.Background(), logs); err != nil {
|
||||
log.Printf("Failed to batch create usage logs: %v", err)
|
||||
}
|
||||
|
||||
log.Printf("Flushed %d usage logs", len(logs))
|
||||
}
|
||||
Reference in New Issue
Block a user