refactor backend packages and global db service

This commit is contained in:
2026-07-18 21:58:15 +08:00
parent 62bd3d0455
commit 7800b07d42
51 changed files with 719 additions and 647 deletions

View File

@@ -0,0 +1,141 @@
package ai
import (
"crypto/aes"
"crypto/cipher"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"errors"
"io"
"time"
"gorm.io/gorm/clause"
"senlinai-agent/backend/internal/models"
)
type Gateway struct {
systemKey string
encryptionSecret string
}
type SelectedKey struct {
Provider string
APIKey string
KeyType string
}
func NewGateway(systemKey string) *Gateway {
return NewGatewayWithSecret(systemKey, "development-ai-key-secret-change-me")
}
func NewGatewayWithSecret(systemKey string, encryptionSecret string) *Gateway {
return &Gateway{systemKey: systemKey, encryptionSecret: encryptionSecret}
}
func (g *Gateway) SaveUserKey(userID uint, provider string, apiKey string) error {
encrypted, err := encryptAPIKey(apiKey, g.encryptionSecret)
if err != nil {
return err
}
key := models.SenlinAgentAIKey{UserID: userID, Provider: provider, EncryptedAPIKey: encrypted}
return models.DBService.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "user_id"}},
DoUpdates: clause.AssignmentColumns([]string{"provider", "encrypted_api_key", "updated_at"}),
}).Create(&key).Error
}
func (g *Gateway) SelectKey(userID uint) (SelectedKey, error) {
var userKey models.SenlinAgentAIKey
if err := models.DBService.Where("user_id = ?", userID).First(&userKey).Error; err == nil {
apiKey, err := decryptAPIKey(userKey.EncryptedAPIKey, g.encryptionSecret)
if err != nil {
return SelectedKey{}, err
}
return SelectedKey{Provider: userKey.Provider, APIKey: apiKey, KeyType: "user"}, nil
}
if g.systemKey == "" {
return SelectedKey{}, errors.New("no ai key available")
}
return SelectedKey{Provider: "openai", APIKey: g.systemKey, KeyType: "system"}, nil
}
func (g *Gateway) RecordCall(userID uint, provider string, usedKeyType string, action string, status string, errText string) error {
return models.DBService.Create(&models.SenlinAgentAICallLog{
UserID: userID,
Provider: provider,
UsedKeyType: usedKeyType,
Action: action,
Status: status,
Error: errText,
}).Error
}
func (g *Gateway) CheckRateLimit(userID uint, action string, limit int, window time.Duration) error {
if limit <= 0 {
return nil
}
var count int64
if err := models.DBService.Model(&models.SenlinAgentAICallLog{}).
Where("user_id = ? AND action = ? AND created_at >= ?", userID, action, time.Now().Add(-window)).
Count(&count).Error; err != nil {
return err
}
if count >= int64(limit) {
return errors.New("ai rate limit exceeded")
}
return nil
}
func encryptAPIKey(apiKey string, secret string) (string, error) {
if apiKey == "" {
return "", errors.New("api key is required")
}
block, err := aes.NewCipher(encryptionKey(secret))
if err != nil {
return "", err
}
gcm, err := cipher.NewGCM(block)
if err != nil {
return "", err
}
nonce := make([]byte, gcm.NonceSize())
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
return "", err
}
ciphertext := gcm.Seal(nil, nonce, []byte(apiKey), nil)
return "v1:" + base64.RawURLEncoding.EncodeToString(append(nonce, ciphertext...)), nil
}
func decryptAPIKey(encrypted string, secret string) (string, error) {
if len(encrypted) < 3 || encrypted[:3] != "v1:" {
return encrypted, nil
}
raw, err := base64.RawURLEncoding.DecodeString(encrypted[3:])
if err != nil {
return "", errors.New("invalid encrypted api key")
}
block, err := aes.NewCipher(encryptionKey(secret))
if err != nil {
return "", err
}
gcm, err := cipher.NewGCM(block)
if err != nil {
return "", err
}
if len(raw) < gcm.NonceSize() {
return "", errors.New("invalid encrypted api key")
}
nonce := raw[:gcm.NonceSize()]
ciphertext := raw[gcm.NonceSize():]
plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
if err != nil {
return "", errors.New("invalid encrypted api key")
}
return string(plaintext), nil
}
func encryptionKey(secret string) []byte {
sum := sha256.Sum256([]byte(secret))
return sum[:]
}

View File

@@ -0,0 +1,106 @@
package ai
import (
"fmt"
"testing"
"time"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"senlinai-agent/backend/internal/models"
)
func TestSelectKeyPrefersUserKey(t *testing.T) {
database := newTestDB(t)
require.NoError(t, database.Create(&models.SenlinAgentAIKey{UserID: 3, Provider: "openai", EncryptedAPIKey: "user-key"}).Error)
gateway := NewGateway("system-key")
selected, err := gateway.SelectKey(3)
require.NoError(t, err)
require.Equal(t, "user", selected.KeyType)
require.Equal(t, "user-key", selected.APIKey)
}
func TestSaveUserKeyEncryptsStoredKey(t *testing.T) {
database := newTestDB(t)
gateway := NewGatewayWithSecret("system-key", "test-encryption-secret")
require.NoError(t, gateway.SaveUserKey(3, "openai", "user-key"))
var stored models.SenlinAgentAIKey
require.NoError(t, database.Where("user_id = ?", 3).First(&stored).Error)
require.NotEqual(t, "user-key", stored.EncryptedAPIKey)
require.Contains(t, stored.EncryptedAPIKey, "v1:")
selected, err := gateway.SelectKey(3)
require.NoError(t, err)
require.Equal(t, "user-key", selected.APIKey)
}
func TestSelectKeyFallsBackToSystemKey(t *testing.T) {
newTestDB(t)
gateway := NewGateway("system-key")
selected, err := gateway.SelectKey(3)
require.NoError(t, err)
require.Equal(t, "system", selected.KeyType)
require.Equal(t, "system-key", selected.APIKey)
}
func TestSelectKeyReturnsErrorWhenNoKeyAvailable(t *testing.T) {
newTestDB(t)
gateway := NewGateway("")
_, err := gateway.SelectKey(3)
require.ErrorContains(t, err, "no ai key available")
}
func TestRecordCallStoresAuditFields(t *testing.T) {
database := newTestDB(t)
gateway := NewGateway("system-key")
require.NoError(t, gateway.RecordCall(3, "openai", "system", "inbox_analyze", "failed", "rate limited"))
var log models.SenlinAgentAICallLog
require.NoError(t, database.First(&log).Error)
require.Equal(t, uint(3), log.UserID)
require.Equal(t, "openai", log.Provider)
require.Equal(t, "system", log.UsedKeyType)
require.Equal(t, "inbox_analyze", log.Action)
require.Equal(t, "failed", log.Status)
require.Equal(t, "rate limited", log.Error)
}
func TestCheckRateLimitRejectsCallsOverWindow(t *testing.T) {
newTestDB(t)
gateway := NewGateway("system-key")
require.NoError(t, gateway.RecordCall(3, "openai", "system", "inbox_analyze", "succeeded", ""))
err := gateway.CheckRateLimit(3, "inbox_analyze", 1, time.Hour)
require.ErrorContains(t, err, "ai rate limit exceeded")
}
func TestCreateAISession(t *testing.T) {
newTestDB(t)
service := NewSessionService()
session, err := service.Create(7, 3, "报价分析")
require.NoError(t, err)
require.Equal(t, uint(7), session.ProjectID)
require.Equal(t, uint(3), session.CreatedBy)
require.Equal(t, "报价分析", session.Title)
}
func newTestDB(t *testing.T) *gorm.DB {
t.Helper()
database, err := gorm.Open(sqlite.Open(fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name())), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, models.AutoMigrate(database))
models.DBService = database
return database
}

View File

@@ -0,0 +1,24 @@
package ai
import (
"errors"
"strings"
"senlinai-agent/backend/internal/models"
)
type SessionService struct {
}
func NewSessionService() *SessionService {
return &SessionService{}
}
func (s *SessionService) Create(projectID uint, userID uint, title string) (*models.SenlinAgentAISession, error) {
title = strings.TrimSpace(title)
if title == "" {
return nil, errors.New("session title is required")
}
session := &models.SenlinAgentAISession{ProjectID: projectID, CreatedBy: userID, Title: title}
return session, models.DBService.Create(session).Error
}