Add runtime restore and temp asset cleanup
This commit is contained in:
@@ -0,0 +1,149 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type ConversationMessageInput struct {
|
||||
Hash string
|
||||
Role string
|
||||
Snapshot map[string]any
|
||||
AssetSHA256s []string
|
||||
}
|
||||
|
||||
type TaskMessageRefInput struct {
|
||||
MessageID string
|
||||
Position int
|
||||
}
|
||||
|
||||
type ConversationMessageRef struct {
|
||||
MessageID string `json:"messageId"`
|
||||
Position int `json:"position"`
|
||||
Message map[string]any `json:"message"`
|
||||
}
|
||||
|
||||
func (s *Store) EnsureConversation(ctx context.Context, user *auth.User, conversationKey string, metadata map[string]any) (string, error) {
|
||||
conversationKey = strings.TrimSpace(conversationKey)
|
||||
if conversationKey == "" {
|
||||
return "", nil
|
||||
}
|
||||
userID := ""
|
||||
gatewayUserID := ""
|
||||
if user != nil {
|
||||
userID = strings.TrimSpace(user.ID)
|
||||
gatewayUserID = strings.TrimSpace(user.GatewayUserID)
|
||||
}
|
||||
if userID == "" {
|
||||
userID = "anonymous"
|
||||
}
|
||||
metadataJSON, _ := json.Marshal(emptyObjectIfNil(metadata))
|
||||
var conversationID string
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
INSERT INTO gateway_conversations (user_id, gateway_user_id, conversation_key, metadata)
|
||||
VALUES ($1, NULLIF($2, '')::uuid, $3, $4::jsonb)
|
||||
ON CONFLICT (user_id, conversation_key) DO UPDATE
|
||||
SET gateway_user_id = COALESCE(gateway_conversations.gateway_user_id, EXCLUDED.gateway_user_id),
|
||||
metadata = gateway_conversations.metadata || EXCLUDED.metadata,
|
||||
updated_at = now()
|
||||
RETURNING id::text`, userID, gatewayUserID, conversationKey, string(metadataJSON)).Scan(&conversationID)
|
||||
return conversationID, err
|
||||
}
|
||||
|
||||
func (s *Store) UpsertConversationMessages(ctx context.Context, conversationID string, messages []ConversationMessageInput) ([]TaskMessageRefInput, int, error) {
|
||||
if strings.TrimSpace(conversationID) == "" || len(messages) == 0 {
|
||||
return nil, 0, nil
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
refs := make([]TaskMessageRefInput, 0, len(messages))
|
||||
newCount := 0
|
||||
for index, message := range messages {
|
||||
snapshotJSON, _ := json.Marshal(emptyObjectIfNil(message.Snapshot))
|
||||
var messageID string
|
||||
var inserted bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
INSERT INTO gateway_conversation_messages (
|
||||
conversation_id, message_hash, role, message_snapshot, asset_sha256s
|
||||
)
|
||||
VALUES ($1::uuid, $2, NULLIF($3, ''), $4::jsonb, $5)
|
||||
ON CONFLICT (conversation_id, message_hash) DO UPDATE
|
||||
SET updated_at = gateway_conversation_messages.updated_at
|
||||
RETURNING id::text, (xmax = 0) AS inserted`,
|
||||
conversationID,
|
||||
message.Hash,
|
||||
message.Role,
|
||||
string(snapshotJSON),
|
||||
message.AssetSHA256s,
|
||||
).Scan(&messageID, &inserted); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
if inserted {
|
||||
newCount++
|
||||
}
|
||||
refs = append(refs, TaskMessageRefInput{MessageID: messageID, Position: index})
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
return refs, newCount, nil
|
||||
}
|
||||
|
||||
func (s *Store) ListTaskConversationMessages(ctx context.Context, taskID string) ([]ConversationMessageRef, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT refs.message_id::text, refs.position, messages.message_snapshot
|
||||
FROM gateway_task_message_refs refs
|
||||
JOIN gateway_conversation_messages messages ON messages.id = refs.message_id
|
||||
WHERE refs.task_id = $1::uuid
|
||||
ORDER BY refs.position ASC`, taskID)
|
||||
if err != nil {
|
||||
if IsUndefinedDatabaseObject(err) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]ConversationMessageRef, 0)
|
||||
for rows.Next() {
|
||||
var item ConversationMessageRef
|
||||
var snapshot []byte
|
||||
if err := rows.Scan(&item.MessageID, &item.Position, &snapshot); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.Message = decodeObject(snapshot)
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func insertTaskMessageRefs(ctx context.Context, tx pgx.Tx, taskID string, refs []TaskMessageRefInput) error {
|
||||
if len(refs) == 0 {
|
||||
return nil
|
||||
}
|
||||
for _, ref := range refs {
|
||||
if strings.TrimSpace(ref.MessageID) == "" {
|
||||
continue
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_task_message_refs (task_id, message_id, position)
|
||||
VALUES ($1::uuid, $2::uuid, $3)
|
||||
ON CONFLICT (task_id, position) DO UPDATE
|
||||
SET message_id = EXCLUDED.message_id`,
|
||||
taskID,
|
||||
ref.MessageID,
|
||||
ref.Position,
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -12,8 +12,9 @@ import (
|
||||
const defaultServerMainUploadURL = "http://127.0.0.1:3001/v1/files/upload"
|
||||
|
||||
const (
|
||||
FileStorageSceneUpload = "upload"
|
||||
FileStorageSceneImageResult = "image_result"
|
||||
FileStorageSceneUpload = "upload"
|
||||
FileStorageSceneImageResult = "image_result"
|
||||
FileStorageSceneRequestAsset = "request_asset"
|
||||
)
|
||||
|
||||
const (
|
||||
|
||||
@@ -380,11 +380,14 @@ type RateLimitWindow struct {
|
||||
}
|
||||
|
||||
type CreateTaskInput struct {
|
||||
Kind string `json:"kind"`
|
||||
Model string `json:"model"`
|
||||
RunMode string `json:"runMode"`
|
||||
Async bool `json:"async"`
|
||||
Request map[string]any `json:"request"`
|
||||
Kind string `json:"kind"`
|
||||
Model string `json:"model"`
|
||||
RunMode string `json:"runMode"`
|
||||
Async bool `json:"async"`
|
||||
Request map[string]any `json:"request"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
NewMessageCount int `json:"newMessageCount"`
|
||||
MessageRefs []TaskMessageRefInput `json:"messageRefs"`
|
||||
}
|
||||
|
||||
type GatewayTask struct {
|
||||
@@ -407,6 +410,8 @@ type GatewayTask struct {
|
||||
RequestedModel string `json:"requestedModel,omitempty"`
|
||||
ResolvedModel string `json:"resolvedModel,omitempty"`
|
||||
RequestID string `json:"requestId,omitempty"`
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
NewMessageCount int `json:"newMessageCount,omitempty"`
|
||||
Request map[string]any `json:"request,omitempty"`
|
||||
AsyncMode bool `json:"asyncMode"`
|
||||
RiverJobID int64 `json:"riverJobId,omitempty"`
|
||||
@@ -438,6 +443,7 @@ COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_
|
||||
COALESCE(api_key_id, ''), COALESCE(api_key_name, ''), COALESCE(api_key_prefix, ''),
|
||||
COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model,
|
||||
COALESCE(model_type, ''), COALESCE(requested_model, ''), COALESCE(resolved_model, ''), COALESCE(request_id, ''),
|
||||
COALESCE(conversation_id::text, ''), COALESCE(new_message_count, 0),
|
||||
request, COALESCE(async_mode, false), COALESCE(river_job_id, 0), status, COALESCE(attempt_count, 0),
|
||||
COALESCE(remote_task_id, ''), COALESCE(remote_task_payload, '{}'::jsonb),
|
||||
COALESCE(result, '{}'::jsonb), COALESCE(billings, '[]'::jsonb),
|
||||
@@ -1746,15 +1752,18 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
|
||||
INSERT INTO gateway_tasks (
|
||||
kind, run_mode, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
|
||||
api_key_id, api_key_name, api_key_prefix, user_group_id, user_group_key,
|
||||
model, requested_model, request, async_mode, status, result, billings, finished_at
|
||||
model, requested_model, request, async_mode, status, result, billings, conversation_id, new_message_count, finished_at
|
||||
)
|
||||
VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, '')::uuid, NULLIF($13, ''), $14, $14, $15, $16, $17, $18::jsonb, $19::jsonb, CASE WHEN $20 THEN now() ELSE NULL END)
|
||||
VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, '')::uuid, NULLIF($13, ''), $14, $14, $15, $16, $17, $18::jsonb, $19::jsonb, NULLIF($20, '')::uuid, $21, CASE WHEN $22 THEN now() ELSE NULL END)
|
||||
RETURNING `+gatewayTaskColumns,
|
||||
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, input.Async, status, resultBody, billingsBody, false,
|
||||
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, input.Async, status, resultBody, billingsBody, input.ConversationID, input.NewMessageCount, false,
|
||||
))
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
if err := insertTaskMessageRefs(ctx, tx, task.ID, input.MessageRefs); err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
events := taskEventsForCreate(task.ID, runMode, status, nil)
|
||||
for _, event := range events {
|
||||
payload, _ := json.Marshal(event.Payload)
|
||||
@@ -1822,6 +1831,8 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
&task.RequestedModel,
|
||||
&task.ResolvedModel,
|
||||
&task.RequestID,
|
||||
&task.ConversationID,
|
||||
&task.NewMessageCount,
|
||||
&requestBytes,
|
||||
&task.AsyncMode,
|
||||
&task.RiverJobID,
|
||||
|
||||
@@ -7,6 +7,8 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type RateLimitMetricStatus struct {
|
||||
@@ -82,6 +84,59 @@ type PlatformPolicyEvent struct {
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
}
|
||||
|
||||
func (s *Store) RestorePlatformModelRuntimeStatus(ctx context.Context, platformModelID string) (ModelRateLimitStatus, error) {
|
||||
platformModelID = strings.TrimSpace(platformModelID)
|
||||
if platformModelID == "" {
|
||||
return ModelRateLimitStatus{}, pgx.ErrNoRows
|
||||
}
|
||||
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return ModelRateLimitStatus{}, err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
var restoredModelID string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
WITH restored_model AS (
|
||||
UPDATE platform_models
|
||||
SET enabled = true,
|
||||
cooldown_until = NULL,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
RETURNING id::text, platform_id
|
||||
),
|
||||
restored_platform AS (
|
||||
UPDATE integration_platforms
|
||||
SET status = 'enabled',
|
||||
disabled_reason = NULL,
|
||||
cooldown_until = NULL,
|
||||
updated_at = now()
|
||||
WHERE id = (SELECT platform_id FROM restored_model)
|
||||
AND deleted_at IS NULL
|
||||
RETURNING id
|
||||
)
|
||||
SELECT id
|
||||
FROM restored_model
|
||||
WHERE EXISTS (SELECT 1 FROM restored_platform)`, platformModelID).Scan(&restoredModelID); err != nil {
|
||||
return ModelRateLimitStatus{}, err
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return ModelRateLimitStatus{}, err
|
||||
}
|
||||
|
||||
items, err := s.ListModelRateLimitStatuses(ctx)
|
||||
if err != nil {
|
||||
return ModelRateLimitStatus{}, err
|
||||
}
|
||||
for _, item := range items {
|
||||
if item.PlatformModelID == restoredModelID {
|
||||
return item, nil
|
||||
}
|
||||
}
|
||||
return ModelRateLimitStatus{}, pgx.ErrNoRows
|
||||
}
|
||||
|
||||
func (s *Store) ListModelRateLimitStatuses(ctx context.Context) ([]ModelRateLimitStatus, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT m.id::text, m.platform_id::text, p.name, p.provider, p.status,
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type RequestAsset struct {
|
||||
ID string `json:"id"`
|
||||
SHA256 string `json:"sha256"`
|
||||
ContentType string `json:"contentType"`
|
||||
ByteSize int64 `json:"byteSize"`
|
||||
URL string `json:"url"`
|
||||
StorageProvider string `json:"storageProvider"`
|
||||
LocalPath string `json:"localPath,omitempty"`
|
||||
ExpiresAt *time.Time `json:"expiresAt,omitempty"`
|
||||
ExpiredAt *time.Time `json:"expiredAt,omitempty"`
|
||||
RefCount int `json:"refCount"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type RequestAssetInput struct {
|
||||
SHA256 string
|
||||
ContentType string
|
||||
ByteSize int64
|
||||
URL string
|
||||
StorageProvider string
|
||||
LocalPath string
|
||||
ExpiresAt *time.Time
|
||||
}
|
||||
|
||||
func (s *Store) FindRequestAsset(ctx context.Context, sha256 string, contentType string) (RequestAsset, bool, error) {
|
||||
asset, err := scanRequestAsset(s.pool.QueryRow(ctx, `
|
||||
SELECT id::text, sha256, content_type, byte_size, url, storage_provider,
|
||||
COALESCE(local_path, ''), expires_at, expired_at, ref_count, created_at, updated_at
|
||||
FROM gateway_request_assets
|
||||
WHERE sha256 = $1 AND content_type = $2`, sha256, contentType))
|
||||
if err != nil {
|
||||
if err == pgx.ErrNoRows {
|
||||
return RequestAsset{}, false, nil
|
||||
}
|
||||
return RequestAsset{}, false, err
|
||||
}
|
||||
return asset, true, nil
|
||||
}
|
||||
|
||||
func (s *Store) UpsertRequestAsset(ctx context.Context, input RequestAssetInput) (RequestAsset, error) {
|
||||
return scanRequestAsset(s.pool.QueryRow(ctx, `
|
||||
INSERT INTO gateway_request_assets (
|
||||
sha256, content_type, byte_size, url, storage_provider, local_path, expires_at, expired_at, ref_count
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, NULLIF($6, ''), $7, NULL, 1)
|
||||
ON CONFLICT (sha256, content_type) DO UPDATE
|
||||
SET byte_size = EXCLUDED.byte_size,
|
||||
url = EXCLUDED.url,
|
||||
storage_provider = EXCLUDED.storage_provider,
|
||||
local_path = EXCLUDED.local_path,
|
||||
expires_at = EXCLUDED.expires_at,
|
||||
expired_at = NULL,
|
||||
ref_count = gateway_request_assets.ref_count + 1,
|
||||
updated_at = now()
|
||||
RETURNING id::text, sha256, content_type, byte_size, url, storage_provider,
|
||||
COALESCE(local_path, ''), expires_at, expired_at, ref_count, created_at, updated_at`,
|
||||
input.SHA256,
|
||||
input.ContentType,
|
||||
input.ByteSize,
|
||||
input.URL,
|
||||
input.StorageProvider,
|
||||
input.LocalPath,
|
||||
input.ExpiresAt,
|
||||
))
|
||||
}
|
||||
|
||||
func (s *Store) IncrementRequestAssetRefCount(ctx context.Context, sha256 string, contentType string) error {
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_request_assets
|
||||
SET ref_count = ref_count + 1,
|
||||
updated_at = now()
|
||||
WHERE sha256 = $1 AND content_type = $2`, sha256, contentType)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) MarkRequestAssetExpiredByLocalPath(ctx context.Context, localPath string, expiredAt time.Time) error {
|
||||
if localPath == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_request_assets
|
||||
SET expired_at = COALESCE(expired_at, $2),
|
||||
updated_at = now()
|
||||
WHERE local_path = $1
|
||||
AND storage_provider = 'local_static'
|
||||
AND expired_at IS NULL`, localPath, expiredAt)
|
||||
return err
|
||||
}
|
||||
|
||||
func scanRequestAsset(scanner interface{ Scan(dest ...any) error }) (RequestAsset, error) {
|
||||
var asset RequestAsset
|
||||
var localPath string
|
||||
var expiresAt sql.NullTime
|
||||
var expiredAt sql.NullTime
|
||||
if err := scanner.Scan(
|
||||
&asset.ID,
|
||||
&asset.SHA256,
|
||||
&asset.ContentType,
|
||||
&asset.ByteSize,
|
||||
&asset.URL,
|
||||
&asset.StorageProvider,
|
||||
&localPath,
|
||||
&expiresAt,
|
||||
&expiredAt,
|
||||
&asset.RefCount,
|
||||
&asset.CreatedAt,
|
||||
&asset.UpdatedAt,
|
||||
); err != nil {
|
||||
return RequestAsset{}, err
|
||||
}
|
||||
asset.LocalPath = localPath
|
||||
if expiresAt.Valid {
|
||||
asset.ExpiresAt = &expiresAt.Time
|
||||
}
|
||||
if expiredAt.Valid {
|
||||
asset.ExpiredAt = &expiredAt.Time
|
||||
}
|
||||
return asset, nil
|
||||
}
|
||||
Reference in New Issue
Block a user