llama-link/internal/db/models.go

package db

import (
	"database/sql/driver"
	"encoding/json"
	"time"

	"github.com/google/uuid"
	"gorm.io/gorm"
)

type StringArray []string

func (a StringArray) Value() (driver.Value, error) {
	return json.Marshal(a)
}

func (a *StringArray) Scan(value interface{}) error {
	if value == nil {
		*a = []string{}
		return nil
	}
	b, ok := value.([]byte)
	if !ok {
		return nil
	}
	return json.Unmarshal(b, a)
}

type AdminUser struct {
	ID           uuid.UUID  `gorm:"type:uuid;primaryKey" json:"id"`
	Username     string     `gorm:"size:64;uniqueIndex;not null" json:"username"`
	PasswordHash string     `gorm:"size:255;not null" json:"-"`
	IsActive     bool       `gorm:"default:true" json:"is_active"`
	CreatedAt   time.Time  `json:"created_at"`
	UpdatedAt   time.Time  `json:"updated_at"`
	LastLoginAt *time.Time `json:"last_login_at,omitempty"`
}

func (AdminUser) TableName() string { return "admin_users" }

func (u *AdminUser) BeforeCreate(tx *gorm.DB) error {
	if u.ID == uuid.Nil {
		u.ID = uuid.New()
	}
	return nil
}

type ApiKey struct {
	ID          uuid.UUID  `gorm:"type:uuid;primaryKey" json:"id"`
	Name        string     `gorm:"size:255;not null" json:"name"`
	KeyHash     string     `gorm:"size:255;not null" json:"-"`
	KeyPrefix   string     `gorm:"size:8;not null;index" json:"key_prefix"`
	Scopes      StringArray `gorm:"type:text;serializer:json" json:"scopes"`
	IsActive    bool       `gorm:"default:true" json:"is_active"`
	IsAdmin     bool       `gorm:"default:false" json:"is_admin"`
	OwnerLabel  *string    `gorm:"size:255" json:"owner_label,omitempty"`
	CreatedAt   time.Time  `json:"created_at"`
	UpdatedAt   time.Time  `json:"updated_at"`
	LastUsedAt  *time.Time `json:"last_used_at,omitempty"`
	ExpiresAt   *time.Time `json:"expires_at,omitempty"`
	DeletedAt   gorm.DeletedAt `gorm:"index" json:"-"`

	UsageLogs []UsageLog `gorm:"foreignKey:ApiKeyID" json:"-"`
	Quotas    []Quota    `gorm:"foreignKey:ApiKeyID" json:"-"`
	Webhooks  []Webhook  `gorm:"foreignKey:ApiKeyID" json:"-"`
}

func (ApiKey) TableName() string { return "api_keys" }

func (k *ApiKey) BeforeCreate(tx *gorm.DB) error {
	if k.ID == uuid.Nil {
		k.ID = uuid.New()
	}
	return nil
}

type Model struct {
	ID          uuid.UUID  `gorm:"type:uuid;primaryKey" json:"id"`
	Name        string     `gorm:"size:255;uniqueIndex;not null" json:"name"`
	ModelPath   string     `gorm:"size:1024;not null" json:"model_path"`
	Alias       string     `gorm:"size:255;not null" json:"alias"`
	CtxSize     int        `gorm:"default:8192" json:"ctx_size"`
	NGPULayers  int        `gorm:"default:-1" json:"n_gpu_layers"`
	ExtraArgs   StringArray `gorm:"type:text;serializer:json" json:"extra_args,omitempty"`
	IsDefault   bool       `gorm:"default:false" json:"is_default"`
	IsEnabled   bool       `gorm:"default:true" json:"is_enabled"`
	IsActive    bool       `gorm:"default:false" json:"is_active"`
	LoadedAt    *time.Time `json:"loaded_at,omitempty"`
	CreatedAt   time.Time  `json:"created_at"`
	UpdatedAt   time.Time  `json:"updated_at"`
	DeletedAt   gorm.DeletedAt `gorm:"index" json:"-"`
}

func (Model) TableName() string { return "models" }

func (m *Model) BeforeCreate(tx *gorm.DB) error {
	if m.ID == uuid.Nil {
		m.ID = uuid.New()
	}
	return nil
}

type UsageLog struct {
	ID              uuid.UUID  `gorm:"type:uuid;primaryKey" json:"id"`
	ApiKeyID        uuid.UUID  `gorm:"type:uuid;not null;index" json:"api_key_id"`
	ModelName       string     `gorm:"size:255;not null" json:"model_name"`
	Endpoint        string     `gorm:"size:100;not null" json:"endpoint"`
	PromptTokens    int        `gorm:"default:0" json:"prompt_tokens"`
	CompletionTokens int        `gorm:"default:0" json:"completion_tokens"`
	TotalTokens     int        `gorm:"default:0" json:"total_tokens"`
	LatencyMs       int        `gorm:"default:0" json:"latency_ms"`
	Status          string     `gorm:"size:50;not null" json:"status"`
	IPAddress       *string    `gorm:"size:45" json:"ip_address,omitempty"`
	UserAgent       *string    `gorm:"size:512" json:"user_agent,omitempty"`
	ErrorMessage    *string    `gorm:"type:text" json:"error_message,omitempty"`
	Streamed        bool       `gorm:"default:false" json:"streamed"`
	CreatedAt       time.Time  `gorm:"index" json:"created_at"`
}

func (UsageLog) TableName() string { return "usage_logs" }

func (u *UsageLog) BeforeCreate(tx *gorm.DB) error {
	if u.ID == uuid.Nil {
		u.ID = uuid.New()
	}
	return nil
}

type Quota struct {
	ID          uuid.UUID  `gorm:"type:uuid;primaryKey" json:"id"`
	ApiKeyID    uuid.UUID  `gorm:"type:uuid;not null;index" json:"api_key_id"`
	ModelScope  *string    `gorm:"size:255" json:"model_scope,omitempty"`
	PeriodStart time.Time  `gorm:"not null;index" json:"period_start"`
	PeriodEnd   time.Time  `gorm:"not null" json:"period_end"`
	TokensLimit int        `gorm:"default:0" json:"tokens_limit"`
	TokensUsed  int        `gorm:"default:0" json:"tokens_used"`
	CreatedAt   time.Time  `json:"created_at"`
	UpdatedAt   time.Time  `json:"updated_at"`
}

func (Quota) TableName() string { return "quotas" }

func (q *Quota) BeforeCreate(tx *gorm.DB) error {
	if q.ID == uuid.Nil {
		q.ID = uuid.New()
	}
	return nil
}

type Webhook struct {
	ID        uuid.UUID  `gorm:"type:uuid;primaryKey" json:"id"`
	ApiKeyID  uuid.UUID  `gorm:"type:uuid;not null;index" json:"api_key_id"`
	URL       string     `gorm:"size:2048;not null" json:"url"`
	Event     string     `gorm:"size:100;not null" json:"event"`
	Secret    *string    `gorm:"size:255" json:"secret,omitempty"`
	IsActive  bool       `gorm:"default:true" json:"is_active"`
	CreatedAt time.Time  `json:"created_at"`
}

func (Webhook) TableName() string { return "webhooks" }

func (w *Webhook) BeforeCreate(tx *gorm.DB) error {
	if w.ID == uuid.Nil {
		w.ID = uuid.New()
	}
	return nil
}