410 lines
18 KiB
Go
410 lines
18 KiB
Go
package contentmodel
|
|
|
|
import (
|
|
"context"
|
|
"crypto/aes"
|
|
"crypto/cipher"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"data-terminal-cloud/internal/audit"
|
|
"data-terminal-cloud/internal/database"
|
|
"data-terminal-cloud/internal/extraction"
|
|
"data-terminal-cloud/internal/identity"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
var (
|
|
ErrAccessDenied = errors.New("content model access denied")
|
|
ErrNotFound = errors.New("content model not found")
|
|
ErrInvalid = errors.New("invalid content model configuration")
|
|
ErrVersion = errors.New("content model version conflict")
|
|
ErrTestRequired = errors.New("content model must pass a connection test before activation")
|
|
ErrActiveRevoke = errors.New("active content model must be replaced before revocation")
|
|
)
|
|
|
|
type Model struct {
|
|
ID string `json:"id"`
|
|
TenantID string `json:"tenant_id"`
|
|
DisplayName string `json:"display_name"`
|
|
Provider string `json:"provider"`
|
|
BaseURL string `json:"base_url"`
|
|
ModelID string `json:"model_id"`
|
|
CredentialHint string `json:"credential_hint"`
|
|
Temperature float64 `json:"temperature"`
|
|
MaxTokens int `json:"max_tokens"`
|
|
Status string `json:"status"`
|
|
TestStatus string `json:"test_status"`
|
|
LastTestedAt *time.Time `json:"last_tested_at,omitempty"`
|
|
LastErrorCode string `json:"last_error_code,omitempty"`
|
|
ActivatedAt *time.Time `json:"activated_at,omitempty"`
|
|
Version int64 `json:"version"`
|
|
CreatedAt time.Time `json:"created_at"`
|
|
UpdatedAt time.Time `json:"updated_at"`
|
|
}
|
|
|
|
type CreateInput struct {
|
|
TenantID string `json:"tenant_id"`
|
|
DisplayName string `json:"display_name"`
|
|
Provider string `json:"provider"`
|
|
BaseURL string `json:"base_url"`
|
|
ModelID string `json:"model_id"`
|
|
APIKey string `json:"api_key"`
|
|
Temperature float64 `json:"temperature"`
|
|
MaxTokens int `json:"max_tokens"`
|
|
}
|
|
|
|
type Registry struct {
|
|
pool *pgxpool.Pool
|
|
aead cipher.AEAD
|
|
fallback extraction.Generator
|
|
}
|
|
|
|
type Service struct {
|
|
registry *Registry
|
|
audit *audit.Repository
|
|
}
|
|
|
|
type storedModel struct {
|
|
Model
|
|
tenantInternalID string
|
|
encrypted []byte
|
|
nonce []byte
|
|
keyVersion int64
|
|
}
|
|
|
|
type actor struct {
|
|
tenantID, tenantPublicID, userID, userPublicID string
|
|
canManage bool
|
|
}
|
|
|
|
func NewRegistry(pool *pgxpool.Pool, masterKeyHex string, fallback extraction.Generator) (*Registry, error) {
|
|
if pool == nil || fallback == nil {
|
|
return nil, errors.New("content model database and fallback generator are required")
|
|
}
|
|
key, err := hex.DecodeString(strings.TrimSpace(masterKeyHex))
|
|
if err != nil || len(key) != 32 {
|
|
return nil, errors.New("content model master key must be 32 bytes encoded as 64 hexadecimal characters")
|
|
}
|
|
block, err := aes.NewCipher(key)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create content model cipher: %w", err)
|
|
}
|
|
aead, err := cipher.NewGCM(block)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create content model AEAD: %w", err)
|
|
}
|
|
return &Registry{pool: pool, aead: aead, fallback: fallback}, nil
|
|
}
|
|
|
|
func NewService(registry *Registry, auditRepository *audit.Repository) (*Service, error) {
|
|
if registry == nil || auditRepository == nil {
|
|
return nil, errors.New("content model registry and audit repository are required")
|
|
}
|
|
return &Service{registry: registry, audit: auditRepository}, nil
|
|
}
|
|
|
|
// Generate resolves the active tenant model at execution time. This makes an
|
|
// activation change effective for the next summary job without restarting the worker.
|
|
func (r *Registry) Generate(ctx context.Context, input extraction.GenerationInput) (extraction.GenerationResult, error) {
|
|
if strings.TrimSpace(input.TenantID) == "" {
|
|
return r.fallback.Generate(ctx, input)
|
|
}
|
|
item, err := r.active(ctx, input.TenantID)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return r.fallback.Generate(ctx, input)
|
|
}
|
|
if err != nil {
|
|
return extraction.GenerationResult{}, err
|
|
}
|
|
generator, err := r.generator(item)
|
|
if err != nil {
|
|
return extraction.GenerationResult{}, err
|
|
}
|
|
return generator.Generate(ctx, input)
|
|
}
|
|
|
|
func (s *Service) List(ctx context.Context, principal identity.Principal, tenantPublicID string) ([]Model, error) {
|
|
current, err := s.resolveActor(ctx, principal, tenantPublicID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
rows, err := s.registry.pool.Query(ctx, modelSelect+` WHERE cm.tenant_id=$1::uuid AND cm.status<>'revoked' ORDER BY (cm.status='active') DESC,cm.updated_at DESC`, current.tenantID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
items := []Model{}
|
|
for rows.Next() {
|
|
item, scanErr := scanModel(rows)
|
|
if scanErr != nil {
|
|
return nil, scanErr
|
|
}
|
|
items = append(items, item.Model)
|
|
}
|
|
return items, rows.Err()
|
|
}
|
|
|
|
func (s *Service) Create(ctx context.Context, principal identity.Principal, input CreateInput, requestID string) (Model, error) {
|
|
normalizeInput(&input)
|
|
if err := validateInput(input); err != nil {
|
|
return Model{}, err
|
|
}
|
|
current, err := s.resolveActor(ctx, principal, input.TenantID)
|
|
if err != nil {
|
|
return Model{}, err
|
|
}
|
|
if !current.canManage {
|
|
return Model{}, ErrAccessDenied
|
|
}
|
|
var result Model
|
|
err = database.Transact(ctx, s.registry.pool, func(tx pgx.Tx) error {
|
|
publicID := uuid.NewString()
|
|
nonce := make([]byte, s.registry.aead.NonceSize())
|
|
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
|
|
return err
|
|
}
|
|
const keyVersion int64 = 1
|
|
encrypted := s.registry.aead.Seal(nil, nonce, []byte(input.APIKey), modelAAD(current.tenantID, publicID, keyVersion))
|
|
row := tx.QueryRow(ctx, `
|
|
INSERT INTO content_model_config(public_id,tenant_id,display_name,provider,base_url,model_id,encrypted_api_key,nonce,key_version,credential_hint,temperature,max_tokens,created_by,updated_by)
|
|
VALUES($1::uuid,$2::uuid,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13::uuid,$13::uuid)
|
|
RETURNING public_id::text,$14::text,display_name,provider,base_url,model_id,credential_hint,temperature::float8,max_tokens,status,test_status,last_tested_at,COALESCE(last_error_code,''),activated_at,version,created_at,updated_at,tenant_id::text,encrypted_api_key,nonce,key_version`,
|
|
publicID, current.tenantID, input.DisplayName, input.Provider, input.BaseURL, input.ModelID, encrypted, nonce, keyVersion, credentialHint(input.APIKey), input.Temperature, input.MaxTokens, current.userID, current.tenantPublicID)
|
|
item, scanErr := scanModel(row)
|
|
if scanErr != nil {
|
|
return scanErr
|
|
}
|
|
result = item.Model
|
|
resourceID := result.ID
|
|
return s.audit.Record(ctx, tx, audit.Event{TenantID: ¤t.tenantID, ActorType: "user", ActorPublicID: ¤t.userPublicID, Action: "content_model.create", ResourceType: "content_model", ResourceID: &resourceID, RequestID: requestID, Purpose: "配置聊天内容总结模型", Result: "success", Detail: map[string]any{"provider": result.Provider, "model_id": result.ModelID}})
|
|
})
|
|
return result, err
|
|
}
|
|
|
|
func (s *Service) Test(ctx context.Context, principal identity.Principal, tenantPublicID, publicID, requestID string) (Model, error) {
|
|
current, err := s.resolveActor(ctx, principal, tenantPublicID)
|
|
if err != nil || !current.canManage {
|
|
return Model{}, ErrAccessDenied
|
|
}
|
|
item, err := s.get(ctx, current.tenantID, publicID)
|
|
if err != nil {
|
|
return Model{}, err
|
|
}
|
|
generator, err := s.registry.generator(item)
|
|
testStatus, errorCode := "passed", ""
|
|
if err == nil {
|
|
_, err = generator.Generate(ctx, extraction.GenerationInput{TenantID: current.tenantID, Content: "模型连通性测试，不包含长者或家庭数据。"})
|
|
}
|
|
if err != nil {
|
|
testStatus, errorCode = "failed", normalizedErrorCode(err)
|
|
}
|
|
var result Model
|
|
updateErr := database.Transact(ctx, s.registry.pool, func(tx pgx.Tx) error {
|
|
row := tx.QueryRow(ctx, modelUpdateReturning+` WHERE cm.tenant_id=$1::uuid AND cm.public_id=$2::uuid`, current.tenantID, publicID, testStatus, errorCode, current.userID, current.tenantPublicID)
|
|
updated, scanErr := scanModel(row)
|
|
if scanErr != nil {
|
|
return scanErr
|
|
}
|
|
result = updated.Model
|
|
resourceID := result.ID
|
|
return s.audit.Record(ctx, tx, audit.Event{TenantID: ¤t.tenantID, ActorType: "user", ActorPublicID: ¤t.userPublicID, Action: "content_model.test", ResourceType: "content_model", ResourceID: &resourceID, RequestID: requestID, Purpose: "验证总结模型连通性", Result: testStatus, Detail: map[string]any{"error_code": errorCode}})
|
|
})
|
|
if updateErr != nil {
|
|
return Model{}, updateErr
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (s *Service) Activate(ctx context.Context, principal identity.Principal, tenantPublicID, publicID, requestID string, expectedVersion int64) (Model, error) {
|
|
current, err := s.resolveActor(ctx, principal, tenantPublicID)
|
|
if err != nil || !current.canManage {
|
|
return Model{}, ErrAccessDenied
|
|
}
|
|
var result Model
|
|
err = database.Transact(ctx, s.registry.pool, func(tx pgx.Tx) error {
|
|
var internalID string
|
|
var version int64
|
|
var testStatus, status string
|
|
if err := tx.QueryRow(ctx, `SELECT id::text,version,test_status,status FROM content_model_config WHERE tenant_id=$1::uuid AND public_id=$2::uuid FOR UPDATE`, current.tenantID, publicID).Scan(&internalID, &version, &testStatus, &status); errors.Is(err, pgx.ErrNoRows) {
|
|
return ErrNotFound
|
|
} else if err != nil {
|
|
return err
|
|
}
|
|
if version != expectedVersion {
|
|
return ErrVersion
|
|
}
|
|
if status == "revoked" {
|
|
return ErrNotFound
|
|
}
|
|
if testStatus != "passed" {
|
|
return ErrTestRequired
|
|
}
|
|
if _, err := tx.Exec(ctx, `UPDATE content_model_config SET status='inactive',activated_at=NULL,updated_at=now(),version=version+1,updated_by=$2::uuid WHERE tenant_id=$1::uuid AND status='active' AND id<>$3::uuid`, current.tenantID, current.userID, internalID); err != nil {
|
|
return err
|
|
}
|
|
row := tx.QueryRow(ctx, `UPDATE content_model_config cm SET status='active',activated_at=now(),updated_at=now(),version=version+1,updated_by=$3::uuid WHERE tenant_id=$1::uuid AND public_id=$2::uuid RETURNING `+modelReturning, current.tenantID, publicID, current.userID, current.tenantPublicID)
|
|
updated, scanErr := scanModel(row)
|
|
if scanErr != nil {
|
|
return scanErr
|
|
}
|
|
result = updated.Model
|
|
resourceID := result.ID
|
|
return s.audit.Record(ctx, tx, audit.Event{TenantID: ¤t.tenantID, ActorType: "user", ActorPublicID: ¤t.userPublicID, Action: "content_model.activate", ResourceType: "content_model", ResourceID: &resourceID, RequestID: requestID, Purpose: "激活聊天总结模型", Result: "success", Detail: map[string]any{"model_id": result.ModelID}})
|
|
})
|
|
return result, err
|
|
}
|
|
|
|
func (s *Service) Revoke(ctx context.Context, principal identity.Principal, tenantPublicID, publicID, requestID string, expectedVersion int64) (Model, error) {
|
|
current, err := s.resolveActor(ctx, principal, tenantPublicID)
|
|
if err != nil || !current.canManage {
|
|
return Model{}, ErrAccessDenied
|
|
}
|
|
var result Model
|
|
err = database.Transact(ctx, s.registry.pool, func(tx pgx.Tx) error {
|
|
item, lockErr := s.getTx(ctx, tx, current.tenantID, publicID, true)
|
|
if lockErr != nil {
|
|
return lockErr
|
|
}
|
|
if item.Version != expectedVersion {
|
|
return ErrVersion
|
|
}
|
|
if item.Status == "active" {
|
|
return ErrActiveRevoke
|
|
}
|
|
row := tx.QueryRow(ctx, `UPDATE content_model_config cm SET status='revoked',encrypted_api_key=decode('','hex'),nonce=decode('','hex'),updated_at=now(),version=version+1,updated_by=$3::uuid WHERE tenant_id=$1::uuid AND public_id=$2::uuid RETURNING `+modelReturning, current.tenantID, publicID, current.userID, current.tenantPublicID)
|
|
updated, scanErr := scanModel(row)
|
|
if scanErr != nil {
|
|
return scanErr
|
|
}
|
|
result = updated.Model
|
|
resourceID := result.ID
|
|
return s.audit.Record(ctx, tx, audit.Event{TenantID: ¤t.tenantID, ActorType: "user", ActorPublicID: ¤t.userPublicID, Action: "content_model.revoke", ResourceType: "content_model", ResourceID: &resourceID, RequestID: requestID, Purpose: "移除未激活的总结模型", Result: "success"})
|
|
})
|
|
return result, err
|
|
}
|
|
|
|
func (r *Registry) active(ctx context.Context, tenantID string) (storedModel, error) {
|
|
row := r.pool.QueryRow(ctx, modelSelect+` WHERE cm.tenant_id=$1::uuid AND cm.status='active'`, tenantID)
|
|
return scanModel(row)
|
|
}
|
|
|
|
func (r *Registry) generator(item storedModel) (extraction.Generator, error) {
|
|
key, err := r.aead.Open(nil, item.nonce, item.encrypted, modelAAD(item.tenantInternalID, item.ID, item.keyVersion))
|
|
if err != nil || len(key) < 8 {
|
|
return nil, errors.New("content model credential cannot be decrypted")
|
|
}
|
|
return extraction.NewCompatibleGenerator(item.Provider, item.BaseURL, string(key), item.ModelID, item.Temperature, item.MaxTokens, 30*time.Second)
|
|
}
|
|
|
|
func (s *Service) get(ctx context.Context, tenantID, publicID string) (storedModel, error) {
|
|
return s.getTx(ctx, s.registry.pool, tenantID, publicID, false)
|
|
}
|
|
|
|
type rowQueryer interface {
|
|
QueryRow(context.Context, string, ...any) pgx.Row
|
|
}
|
|
|
|
func (s *Service) getTx(ctx context.Context, q rowQueryer, tenantID, publicID string, lock bool) (storedModel, error) {
|
|
lockClause := ""
|
|
if lock {
|
|
lockClause = " FOR UPDATE OF cm"
|
|
}
|
|
item, err := scanModel(q.QueryRow(ctx, modelSelect+` WHERE cm.tenant_id=$1::uuid AND cm.public_id=$2::uuid`+lockClause, tenantID, publicID))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return storedModel{}, ErrNotFound
|
|
}
|
|
return item, err
|
|
}
|
|
|
|
func (s *Service) resolveActor(ctx context.Context, principal identity.Principal, tenantPublicID string) (actor, error) {
|
|
var result actor
|
|
err := s.registry.pool.QueryRow(ctx, `
|
|
SELECT t.id::text,t.public_id::text,ua.id::text,ua.public_id::text,
|
|
COALESCE(bool_or(hm.role IN ('owner','admin') AND hm.status='active'),false)
|
|
FROM tenant t JOIN user_account ua ON ua.tenant_id=t.id AND ua.status='active'
|
|
LEFT JOIN household_member hm ON hm.user_id=ua.id AND hm.tenant_id=t.id
|
|
WHERE t.public_id=$1::uuid AND t.status='active' AND ua.oidc_issuer=$2 AND ua.oidc_subject=$3
|
|
GROUP BY t.id,t.public_id,ua.id,ua.public_id`, tenantPublicID, principal.Issuer, principal.Subject).Scan(&result.tenantID, &result.tenantPublicID, &result.userID, &result.userPublicID, &result.canManage)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return actor{}, ErrAccessDenied
|
|
}
|
|
return result, err
|
|
}
|
|
|
|
func normalizeInput(input *CreateInput) {
|
|
input.TenantID = strings.TrimSpace(input.TenantID)
|
|
input.DisplayName = strings.TrimSpace(input.DisplayName)
|
|
input.Provider = strings.ToLower(strings.TrimSpace(input.Provider))
|
|
input.BaseURL = strings.TrimRight(strings.TrimSpace(input.BaseURL), "/")
|
|
input.ModelID = strings.TrimSpace(input.ModelID)
|
|
input.APIKey = strings.TrimSpace(input.APIKey)
|
|
if input.Temperature == 0 {
|
|
input.Temperature = 0.2
|
|
}
|
|
if input.MaxTokens == 0 {
|
|
input.MaxTokens = 1200
|
|
}
|
|
}
|
|
|
|
func validateInput(input CreateInput) error {
|
|
parsed, err := url.Parse(input.BaseURL)
|
|
validProvider := input.Provider == "ark" || input.Provider == "openai_compatible"
|
|
validURL := err == nil && parsed.Host != "" && (parsed.Scheme == "https" || (parsed.Scheme == "http" && (parsed.Hostname() == "localhost" || parsed.Hostname() == "127.0.0.1")))
|
|
if !validUUID(input.TenantID) || len([]rune(input.DisplayName)) < 1 || len([]rune(input.DisplayName)) > 120 || !validProvider || !validURL || input.ModelID == "" || len(input.APIKey) < 8 || input.Temperature < 0 || input.Temperature > 2 || input.MaxTokens < 256 || input.MaxTokens > 8192 {
|
|
return ErrInvalid
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func credentialHint(key string) string {
|
|
runes := []rune(strings.TrimSpace(key))
|
|
if len(runes) <= 4 {
|
|
return "••••"
|
|
}
|
|
return "••••" + string(runes[len(runes)-4:])
|
|
}
|
|
|
|
func normalizedErrorCode(err error) string {
|
|
value := strings.ToLower(err.Error())
|
|
switch {
|
|
case strings.Contains(value, "401"), strings.Contains(value, "403"):
|
|
return "AUTHENTICATION_FAILED"
|
|
case strings.Contains(value, "429"):
|
|
return "RATE_LIMITED"
|
|
case strings.Contains(value, "timeout"), strings.Contains(value, "deadline"):
|
|
return "TIMEOUT"
|
|
case strings.Contains(value, "validation"), strings.Contains(value, "decode"):
|
|
return "INVALID_RESPONSE"
|
|
default:
|
|
return "CONNECTION_FAILED"
|
|
}
|
|
}
|
|
|
|
func modelAAD(tenantID, publicID string, version int64) []byte {
|
|
return []byte(fmt.Sprintf("content-model-credential-v1\n%s\n%s\n%d", tenantID, publicID, version))
|
|
}
|
|
|
|
func validUUID(value string) bool { _, err := uuid.Parse(value); return err == nil }
|
|
|
|
const modelSelect = `SELECT cm.public_id::text,t.public_id::text,cm.display_name,cm.provider,cm.base_url,cm.model_id,cm.credential_hint,cm.temperature::float8,cm.max_tokens,cm.status,cm.test_status,cm.last_tested_at,COALESCE(cm.last_error_code,''),cm.activated_at,cm.version,cm.created_at,cm.updated_at,cm.tenant_id::text,cm.encrypted_api_key,cm.nonce,cm.key_version FROM content_model_config cm JOIN tenant t ON t.id=cm.tenant_id`
|
|
const modelReturning = `cm.public_id::text,$4::text,cm.display_name,cm.provider,cm.base_url,cm.model_id,cm.credential_hint,cm.temperature::float8,cm.max_tokens,cm.status,cm.test_status,cm.last_tested_at,COALESCE(cm.last_error_code,''),cm.activated_at,cm.version,cm.created_at,cm.updated_at,cm.tenant_id::text,cm.encrypted_api_key,cm.nonce,cm.key_version`
|
|
const modelUpdateReturning = `UPDATE content_model_config cm SET test_status=$3,last_tested_at=now(),last_error_code=NULLIF($4,''),updated_at=now(),version=version+1,updated_by=$5::uuid RETURNING cm.public_id::text,$6::text,cm.display_name,cm.provider,cm.base_url,cm.model_id,cm.credential_hint,cm.temperature::float8,cm.max_tokens,cm.status,cm.test_status,cm.last_tested_at,COALESCE(cm.last_error_code,''),cm.activated_at,cm.version,cm.created_at,cm.updated_at,cm.tenant_id::text,cm.encrypted_api_key,cm.nonce,cm.key_version`
|
|
|
|
type rowScanner interface{ Scan(...any) error }
|
|
|
|
func scanModel(row rowScanner) (storedModel, error) {
|
|
var item storedModel
|
|
err := row.Scan(&item.ID, &item.TenantID, &item.DisplayName, &item.Provider, &item.BaseURL, &item.ModelID, &item.CredentialHint, &item.Temperature, &item.MaxTokens, &item.Status, &item.TestStatus, &item.LastTestedAt, &item.LastErrorCode, &item.ActivatedAt, &item.Version, &item.CreatedAt, &item.UpdatedAt, &item.tenantInternalID, &item.encrypted, &item.nonce, &item.keyVersion)
|
|
return item, err
|
|
}
|