mirror of
https://github.com/mudler/LocalAI.git
synced 2026-05-24 16:51:44 -04:00
feat(usage): add Source, APIKeyID, APIKeyName columns to UsageRecord
Adds three additive columns plus UsageSource* constants. The columns are auto-migrated by InitDB. APIKeyID is a nullable foreign reference to UserAPIKey.ID; APIKeyName is snapshotted on each row so revoked keys keep showing their name in history. Refs: #9862 Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
@@ -8,11 +8,27 @@ import (
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Source classification for a UsageRecord.
|
||||
const (
|
||||
UsageSourceAPIKey = "apikey" // request authenticated with a named UserAPIKey
|
||||
UsageSourceWeb = "web" // request authenticated with a session cookie (web UI)
|
||||
UsageSourceLegacy = "legacy" // request authenticated with an env-configured legacy key
|
||||
)
|
||||
|
||||
// UsageRecord represents a single API request's token usage.
|
||||
type UsageRecord struct {
|
||||
ID uint `gorm:"primaryKey;autoIncrement"`
|
||||
UserID string `gorm:"size:36;index:idx_usage_user_time"`
|
||||
UserName string `gorm:"size:255"`
|
||||
ID uint `gorm:"primaryKey;autoIncrement"`
|
||||
UserID string `gorm:"size:36;index:idx_usage_user_time"`
|
||||
UserName string `gorm:"size:255"`
|
||||
|
||||
// Source classifies how the request authenticated. One of UsageSource* constants.
|
||||
// Empty for pre-feature rows until the InitDB backfill runs.
|
||||
Source string `gorm:"size:16;index:idx_usage_source"`
|
||||
// APIKeyID is the UserAPIKey.ID when Source == UsageSourceAPIKey. Nil otherwise.
|
||||
APIKeyID *string `gorm:"size:36;index:idx_usage_apikey"`
|
||||
// APIKeyName is a snapshot of UserAPIKey.Name at write time. Survives key deletion.
|
||||
APIKeyName string `gorm:"size:255"`
|
||||
|
||||
Model string `gorm:"size:255;index"`
|
||||
Endpoint string `gorm:"size:255"`
|
||||
PromptTokens int64
|
||||
@@ -30,9 +46,12 @@ func RecordUsage(db *gorm.DB, record *UsageRecord) error {
|
||||
// UsageBucket is an aggregated time bucket for the dashboard.
|
||||
type UsageBucket struct {
|
||||
Bucket string `json:"bucket"`
|
||||
Model string `json:"model"`
|
||||
Model string `json:"model,omitempty"`
|
||||
UserID string `json:"user_id,omitempty"`
|
||||
UserName string `json:"user_name,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
APIKeyID string `json:"api_key_id,omitempty"`
|
||||
APIKeyName string `json:"api_key_name,omitempty"`
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
|
||||
@@ -158,4 +158,47 @@ var _ = Describe("Usage", func() {
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
Describe("UsageRecord with source fields", func() {
|
||||
It("persists Source, APIKeyID, APIKeyName", func() {
|
||||
db := testDB()
|
||||
keyID := "key-uuid-1"
|
||||
record := &auth.UsageRecord{
|
||||
UserID: "user-1",
|
||||
UserName: "Test User",
|
||||
Source: auth.UsageSourceAPIKey,
|
||||
APIKeyID: &keyID,
|
||||
APIKeyName: "ci-runner",
|
||||
Model: "gpt-4",
|
||||
Endpoint: "/v1/chat/completions",
|
||||
TotalTokens: 150,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
Expect(auth.RecordUsage(db, record)).To(Succeed())
|
||||
|
||||
var loaded auth.UsageRecord
|
||||
Expect(db.First(&loaded, record.ID).Error).To(Succeed())
|
||||
Expect(loaded.Source).To(Equal(auth.UsageSourceAPIKey))
|
||||
Expect(loaded.APIKeyID).ToNot(BeNil())
|
||||
Expect(*loaded.APIKeyID).To(Equal("key-uuid-1"))
|
||||
Expect(loaded.APIKeyName).To(Equal("ci-runner"))
|
||||
})
|
||||
|
||||
It("allows nil APIKeyID for web/legacy sources", func() {
|
||||
db := testDB()
|
||||
record := &auth.UsageRecord{
|
||||
UserID: "user-1",
|
||||
Source: auth.UsageSourceWeb,
|
||||
Model: "gpt-4",
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
Expect(auth.RecordUsage(db, record)).To(Succeed())
|
||||
|
||||
var loaded auth.UsageRecord
|
||||
Expect(db.First(&loaded, record.ID).Error).To(Succeed())
|
||||
Expect(loaded.Source).To(Equal(auth.UsageSourceWeb))
|
||||
Expect(loaded.APIKeyID).To(BeNil())
|
||||
Expect(loaded.APIKeyName).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user