Initial commit: LlamaLink Go rewrite
Complete rewrite from Python/FastAPI to Go/Gin: - Go backend: auth (API keys + bcrypt), llama.cpp subprocess manager, hot-swap multi-model, rate limiting, quota system, webhooks - Vue 3 SPA admin panel (src/) with Tailwind CSS - Deployment: Docker multi-stage, docker-compose, nginx, systemd - GORM/SQLite models: ApiKey, Model, UsageLog, Quota, Webhook - REST API: /api/v1/admin/* (keys, models, chat, usage, health) - Embedded frontend via go:embed (build output at web/dist/) Removed legacy Python artifacts (app/, tests/, pyproject.toml, etc.)
This commit is contained in:
@@ -0,0 +1,119 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"log/slog"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/llamalink/llamalink/internal/db"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrQuotaExceeded = errors.New("quota exceeded for this period")
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewService(db *gorm.DB) *Service {
|
||||
return &Service{db: db}
|
||||
}
|
||||
|
||||
func (s *Service) CheckQuota(apiKeyID uuid.UUID, modelScope string) (bool, string, error) {
|
||||
now := time.Now().UTC()
|
||||
|
||||
var quota db.Quota
|
||||
err := s.db.Where(
|
||||
"api_key_id = ? AND period_start <= ? AND period_end > ?",
|
||||
apiKeyID, now, now,
|
||||
).First("a).Error
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return true, "", nil // No quota configured
|
||||
}
|
||||
return false, "", err
|
||||
}
|
||||
|
||||
// Check model-specific scope if set
|
||||
if quota.ModelScope != nil && *quota.ModelScope != "" && modelScope != "" {
|
||||
if *quota.ModelScope != modelScope {
|
||||
return true, "", nil // Quota doesn't apply to this model
|
||||
}
|
||||
}
|
||||
|
||||
if quota.TokensUsed >= quota.TokensLimit {
|
||||
return false, "quota exceeded", nil
|
||||
}
|
||||
|
||||
// Warning at 90%
|
||||
if quota.TokensLimit > 0 {
|
||||
usage := float64(quota.TokensUsed) / float64(quota.TokensLimit)
|
||||
if usage >= 0.9 {
|
||||
slog.Warn("quota usage warning",
|
||||
"api_key_id", apiKeyID,
|
||||
"usage_percent", usage*100,
|
||||
"tokens_used", quota.TokensUsed,
|
||||
"tokens_limit", quota.TokensLimit,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
return true, "", nil
|
||||
}
|
||||
|
||||
func (s *Service) ConsumeQuota(apiKeyID uuid.UUID, modelScope string, tokens int) error {
|
||||
now := time.Now().UTC()
|
||||
periodStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC)
|
||||
periodEnd := periodStart.AddDate(0, 1, 0)
|
||||
|
||||
return s.db.Transaction(func(tx *gorm.DB) error {
|
||||
var quota db.Quota
|
||||
err := tx.Where(
|
||||
"api_key_id = ? AND period_start = ? AND period_end = ?",
|
||||
apiKeyID, periodStart, periodEnd,
|
||||
).First("a).Error
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil // No quota configured
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// Check model scope
|
||||
if quota.ModelScope != nil && *quota.ModelScope != "" && modelScope != "" {
|
||||
if *quota.ModelScope != modelScope {
|
||||
return nil // Quota doesn't apply
|
||||
}
|
||||
}
|
||||
|
||||
newUsed := quota.TokensUsed + tokens
|
||||
if quota.TokensLimit > 0 && newUsed > quota.TokensLimit {
|
||||
return ErrQuotaExceeded
|
||||
}
|
||||
|
||||
return tx.Model("a).Update("tokens_used", newUsed).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) GetUsage(apiKeyID uuid.UUID, periodStart, periodEnd time.Time) (used int, limit int, err error) {
|
||||
var quota db.Quota
|
||||
err = s.db.Where(
|
||||
"api_key_id = ? AND period_start = ? AND period_end = ?",
|
||||
apiKeyID, periodStart, periodEnd,
|
||||
).First("a).Error
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return 0, 0, nil
|
||||
}
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
return quota.TokensUsed, quota.TokensLimit, nil
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/llamalink/llamalink/internal/db"
|
||||
)
|
||||
|
||||
type WebhookPayload struct {
|
||||
Event string `json:"event"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
KeyID uuid.UUID `json:"key_id"`
|
||||
KeyName string `json:"key_name"`
|
||||
Message string `json:"message"`
|
||||
Usage *struct {
|
||||
Used int `json:"tokens_used"`
|
||||
Limit int `json:"tokens_limit"`
|
||||
} `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
type WebhookService struct {
|
||||
db *gorm.DB
|
||||
client *http.Client
|
||||
webhookSecret string
|
||||
}
|
||||
|
||||
func NewWebhookService(db *gorm.DB) *WebhookService {
|
||||
return &WebhookService{
|
||||
db: db,
|
||||
client: &http.Client{Timeout: 10 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *WebhookService) Dispatch(event string, key *db.ApiKey, message string, usage *struct{ Used, Limit int }) {
|
||||
var webhooks []db.Webhook
|
||||
s.db.Where("api_key_id = ? AND event = ? AND is_active = ?", key.ID, event, true).Find(&webhooks)
|
||||
|
||||
for _, wh := range webhooks {
|
||||
go s.send(wh, event, key, message, usage)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *WebhookService) send(wh db.Webhook, event string, key *db.ApiKey, message string, usage *struct{ Used, Limit int }) {
|
||||
payload := WebhookPayload{
|
||||
Event: event,
|
||||
Timestamp: time.Now().UTC(),
|
||||
KeyID: key.ID,
|
||||
KeyName: key.Name,
|
||||
Message: message,
|
||||
}
|
||||
|
||||
if usage != nil {
|
||||
payload.Usage = &struct {
|
||||
Used int `json:"tokens_used"`
|
||||
Limit int `json:"tokens_limit"`
|
||||
}{Used: usage.Used, Limit: usage.Limit}
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
slog.Error("failed to marshal webhook payload", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("POST", wh.URL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
slog.Error("failed to create webhook request", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("X-LlamaLink-Event", event)
|
||||
req.Header.Set("X-LlamaLink-Timestamp", fmt.Sprintf("%d", time.Now().Unix()))
|
||||
|
||||
if wh.Secret != nil && *wh.Secret != "" {
|
||||
sig := s.sign(body, *wh.Secret)
|
||||
req.Header.Set("X-LlamaLink-Signature", sig)
|
||||
}
|
||||
|
||||
resp, err := s.client.Do(req)
|
||||
if err != nil {
|
||||
slog.Error("webhook delivery failed", "url", wh.URL, "error", err)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
slog.Warn("webhook returned error", "url", wh.URL, "status", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *WebhookService) sign(body []byte, secret string) string {
|
||||
h := hmac.New(sha256.New, []byte(secret))
|
||||
h.Write(body)
|
||||
return "sha256=" + hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
Reference in New Issue
Block a user