feat: implement AI gateway phase one runtime
This commit is contained in:
@@ -0,0 +1,190 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const baseModelColumns = `
|
||||
id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
|
||||
capabilities, base_billing_config, default_rate_limit_policy, metadata, pricing_version,
|
||||
status, created_at, updated_at`
|
||||
|
||||
type BaseModelInput struct {
|
||||
ProviderKey string `json:"providerKey"`
|
||||
CanonicalModelKey string `json:"canonicalModelKey"`
|
||||
ProviderModelName string `json:"providerModelName"`
|
||||
ModelType string `json:"modelType"`
|
||||
DisplayName string `json:"displayName"`
|
||||
Capabilities map[string]any `json:"capabilities"`
|
||||
BaseBillingConfig map[string]any `json:"baseBillingConfig"`
|
||||
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy"`
|
||||
Metadata map[string]any `json:"metadata"`
|
||||
PricingVersion int `json:"pricingVersion"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
type baseModelScanner interface {
|
||||
Scan(dest ...any) error
|
||||
}
|
||||
|
||||
func (s *Store) ListBaseModels(ctx context.Context) ([]BaseModel, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT `+baseModelColumns+`
|
||||
FROM base_model_catalog
|
||||
ORDER BY provider_key ASC, model_type ASC, canonical_model_key ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]BaseModel, 0)
|
||||
for rows.Next() {
|
||||
item, err := scanBaseModel(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) CreateBaseModel(ctx context.Context, input BaseModelInput) (BaseModel, error) {
|
||||
input = normalizeBaseModelInput(input)
|
||||
capabilities, _ := json.Marshal(emptyObjectIfNil(input.Capabilities))
|
||||
billingConfig, _ := json.Marshal(emptyObjectIfNil(input.BaseBillingConfig))
|
||||
rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.DefaultRateLimitPolicy))
|
||||
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
|
||||
|
||||
return scanBaseModel(s.pool.QueryRow(ctx, `
|
||||
INSERT INTO base_model_catalog (
|
||||
provider_id, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
|
||||
capabilities, base_billing_config, default_rate_limit_policy, metadata, pricing_version, status
|
||||
)
|
||||
VALUES (
|
||||
(SELECT id FROM model_catalog_providers WHERE provider_key = $1 OR provider_code = $1 LIMIT 1),
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11
|
||||
)
|
||||
RETURNING `+baseModelColumns,
|
||||
input.ProviderKey,
|
||||
input.CanonicalModelKey,
|
||||
input.ProviderModelName,
|
||||
input.ModelType,
|
||||
input.DisplayName,
|
||||
capabilities,
|
||||
billingConfig,
|
||||
rateLimitPolicy,
|
||||
metadata,
|
||||
input.PricingVersion,
|
||||
input.Status,
|
||||
))
|
||||
}
|
||||
|
||||
func (s *Store) UpdateBaseModel(ctx context.Context, id string, input BaseModelInput) (BaseModel, error) {
|
||||
input = normalizeBaseModelInput(input)
|
||||
capabilities, _ := json.Marshal(emptyObjectIfNil(input.Capabilities))
|
||||
billingConfig, _ := json.Marshal(emptyObjectIfNil(input.BaseBillingConfig))
|
||||
rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.DefaultRateLimitPolicy))
|
||||
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
|
||||
|
||||
return scanBaseModel(s.pool.QueryRow(ctx, `
|
||||
UPDATE base_model_catalog
|
||||
SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = $2 OR provider_code = $2 LIMIT 1),
|
||||
provider_key = $2,
|
||||
canonical_model_key = $3,
|
||||
provider_model_name = $4,
|
||||
model_type = $5,
|
||||
display_name = $6,
|
||||
capabilities = $7,
|
||||
base_billing_config = $8,
|
||||
default_rate_limit_policy = $9,
|
||||
metadata = $10,
|
||||
pricing_version = $11,
|
||||
status = $12,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
RETURNING `+baseModelColumns,
|
||||
id,
|
||||
input.ProviderKey,
|
||||
input.CanonicalModelKey,
|
||||
input.ProviderModelName,
|
||||
input.ModelType,
|
||||
input.DisplayName,
|
||||
capabilities,
|
||||
billingConfig,
|
||||
rateLimitPolicy,
|
||||
metadata,
|
||||
input.PricingVersion,
|
||||
input.Status,
|
||||
))
|
||||
}
|
||||
|
||||
func (s *Store) DeleteBaseModel(ctx context.Context, id string) error {
|
||||
result, err := s.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if result.RowsAffected() == 0 {
|
||||
return pgx.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
|
||||
var item BaseModel
|
||||
var capabilities []byte
|
||||
var billingConfig []byte
|
||||
var rateLimitPolicy []byte
|
||||
var metadata []byte
|
||||
if err := scanner.Scan(
|
||||
&item.ID,
|
||||
&item.ProviderKey,
|
||||
&item.CanonicalModelKey,
|
||||
&item.ProviderModelName,
|
||||
&item.ModelType,
|
||||
&item.DisplayName,
|
||||
&capabilities,
|
||||
&billingConfig,
|
||||
&rateLimitPolicy,
|
||||
&metadata,
|
||||
&item.PricingVersion,
|
||||
&item.Status,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return BaseModel{}, err
|
||||
}
|
||||
item.Capabilities = decodeObject(capabilities)
|
||||
item.BaseBillingConfig = decodeObject(billingConfig)
|
||||
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
|
||||
item.Metadata = decodeObject(metadata)
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func normalizeBaseModelInput(input BaseModelInput) BaseModelInput {
|
||||
input.ProviderKey = strings.TrimSpace(input.ProviderKey)
|
||||
input.CanonicalModelKey = strings.TrimSpace(input.CanonicalModelKey)
|
||||
input.ProviderModelName = strings.TrimSpace(input.ProviderModelName)
|
||||
input.ModelType = strings.TrimSpace(input.ModelType)
|
||||
input.DisplayName = strings.TrimSpace(input.DisplayName)
|
||||
input.Status = strings.TrimSpace(input.Status)
|
||||
if input.CanonicalModelKey == "" && input.ProviderKey != "" && input.ProviderModelName != "" {
|
||||
input.CanonicalModelKey = input.ProviderKey + ":" + input.ProviderModelName
|
||||
}
|
||||
if input.DisplayName == "" {
|
||||
input.DisplayName = input.ProviderModelName
|
||||
}
|
||||
if input.ModelType == "" {
|
||||
input.ModelType = "text_generate"
|
||||
}
|
||||
if input.PricingVersion <= 0 {
|
||||
input.PricingVersion = 1
|
||||
}
|
||||
if input.Status == "" {
|
||||
input.Status = "active"
|
||||
}
|
||||
return input
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string) ([]RuntimeModelCandidate, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT p.id::text, p.platform_key, p.name, p.provider, COALESCE(p.base_url, ''),
|
||||
p.auth_type, p.credentials, p.config, p.default_pricing_mode,
|
||||
p.default_discount_factor::float8, COALESCE(p.pricing_rule_set_id::text, ''),
|
||||
p.retry_policy, p.rate_limit_policy,
|
||||
COALESCE(p.dynamic_priority, p.priority) AS effective_priority,
|
||||
m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''),
|
||||
COALESCE(b.provider_model_name, ''), m.model_name, COALESCE(m.model_alias, ''),
|
||||
m.model_type, m.display_name, m.capabilities, m.capability_override,
|
||||
COALESCE(b.base_billing_config, '{}'::jsonb), m.billing_config, m.billing_config_override,
|
||||
m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''),
|
||||
m.retry_policy, m.rate_limit_policy
|
||||
FROM platform_models m
|
||||
JOIN integration_platforms p ON p.id = m.platform_id
|
||||
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
|
||||
LEFT JOIN runtime_client_states s
|
||||
ON s.client_id = p.platform_key || ':' || m.model_type || ':' || m.model_name
|
||||
WHERE p.status = 'enabled'
|
||||
AND p.deleted_at IS NULL
|
||||
AND m.enabled = true
|
||||
AND m.model_type = $2
|
||||
AND (p.cooldown_until IS NULL OR p.cooldown_until <= now())
|
||||
AND (
|
||||
m.model_name = $1
|
||||
OR m.model_alias = $1
|
||||
OR b.canonical_model_key = $1
|
||||
OR b.provider_model_name = $1
|
||||
)
|
||||
ORDER BY effective_priority ASC,
|
||||
COALESCE(s.limiter_ratio, 0) ASC,
|
||||
COALESCE(s.running_count, 0) ASC,
|
||||
COALESCE(s.waiting_count, 0) ASC,
|
||||
COALESCE(s.last_assigned_at, to_timestamp(0)) ASC,
|
||||
m.created_at ASC`, model, modelType)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]RuntimeModelCandidate, 0)
|
||||
for rows.Next() {
|
||||
var item RuntimeModelCandidate
|
||||
var credentials []byte
|
||||
var platformConfig []byte
|
||||
var platformRetryPolicy []byte
|
||||
var platformRateLimitPolicy []byte
|
||||
var capabilities []byte
|
||||
var capabilityOverride []byte
|
||||
var baseBilling []byte
|
||||
var billing []byte
|
||||
var billingOverride []byte
|
||||
var modelRetryPolicy []byte
|
||||
var modelRateLimitPolicy []byte
|
||||
if err := rows.Scan(
|
||||
&item.PlatformID,
|
||||
&item.PlatformKey,
|
||||
&item.PlatformName,
|
||||
&item.Provider,
|
||||
&item.BaseURL,
|
||||
&item.AuthType,
|
||||
&credentials,
|
||||
&platformConfig,
|
||||
&item.DefaultPricingMode,
|
||||
&item.DefaultDiscountFactor,
|
||||
&item.PlatformPricingRuleSetID,
|
||||
&platformRetryPolicy,
|
||||
&platformRateLimitPolicy,
|
||||
&item.PlatformPriority,
|
||||
&item.PlatformModelID,
|
||||
&item.BaseModelID,
|
||||
&item.CanonicalModelKey,
|
||||
&item.ProviderModelName,
|
||||
&item.ModelName,
|
||||
&item.ModelAlias,
|
||||
&item.ModelType,
|
||||
&item.DisplayName,
|
||||
&capabilities,
|
||||
&capabilityOverride,
|
||||
&baseBilling,
|
||||
&billing,
|
||||
&billingOverride,
|
||||
&item.PricingMode,
|
||||
&item.DiscountFactor,
|
||||
&item.ModelPricingRuleSetID,
|
||||
&modelRetryPolicy,
|
||||
&modelRateLimitPolicy,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.Credentials = decodeObject(credentials)
|
||||
item.PlatformConfig = decodeObject(platformConfig)
|
||||
item.PlatformRetryPolicy = decodeObject(platformRetryPolicy)
|
||||
item.PlatformRateLimitPolicy = decodeObject(platformRateLimitPolicy)
|
||||
item.Capabilities = decodeObject(capabilities)
|
||||
item.CapabilityOverride = decodeObject(capabilityOverride)
|
||||
item.BaseBillingConfig = decodeObject(baseBilling)
|
||||
item.BillingConfig = decodeObject(billing)
|
||||
item.BillingConfigOverride = decodeObject(billingOverride)
|
||||
item.ModelRetryPolicy = decodeObject(modelRetryPolicy)
|
||||
item.ModelRateLimitPolicy = decodeObject(modelRateLimitPolicy)
|
||||
item.ClientID = fmt.Sprintf("%s:%s:%s", item.PlatformKey, item.ModelType, item.ModelName)
|
||||
item.QueueKey = item.ClientID
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(items) == 0 {
|
||||
return nil, ErrNoModelCandidate
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const catalogProviderColumns = `
|
||||
id::text, provider_key, COALESCE(NULLIF(provider_code, ''), provider_key) AS provider_code,
|
||||
display_name, provider_type, COALESCE(icon_path, '') AS icon_path,
|
||||
COALESCE(source, '') AS source, capability_schema, default_rate_limit_policy,
|
||||
metadata, status, created_at, updated_at`
|
||||
|
||||
type CatalogProviderInput struct {
|
||||
ProviderKey string `json:"providerKey"`
|
||||
Code string `json:"code"`
|
||||
DisplayName string `json:"displayName"`
|
||||
ProviderType string `json:"providerType"`
|
||||
IconPath string `json:"iconPath"`
|
||||
Source string `json:"source"`
|
||||
CapabilitySchema map[string]any `json:"capabilitySchema"`
|
||||
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy"`
|
||||
Metadata map[string]any `json:"metadata"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
type catalogProviderScanner interface {
|
||||
Scan(dest ...any) error
|
||||
}
|
||||
|
||||
func (s *Store) CreateCatalogProvider(ctx context.Context, input CatalogProviderInput) (CatalogProvider, error) {
|
||||
input = normalizeCatalogProviderInput(input)
|
||||
capabilitySchema, _ := json.Marshal(emptyObjectIfNil(input.CapabilitySchema))
|
||||
rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.DefaultRateLimitPolicy))
|
||||
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
|
||||
|
||||
return scanCatalogProvider(s.pool.QueryRow(ctx, `
|
||||
INSERT INTO model_catalog_providers (
|
||||
provider_key, provider_code, display_name, provider_type, icon_path, source,
|
||||
capability_schema, default_rate_limit_policy, metadata, status
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, NULLIF($5, ''), $6, $7, $8, $9, $10)
|
||||
RETURNING `+catalogProviderColumns,
|
||||
input.ProviderKey,
|
||||
input.Code,
|
||||
input.DisplayName,
|
||||
input.ProviderType,
|
||||
input.IconPath,
|
||||
input.Source,
|
||||
capabilitySchema,
|
||||
rateLimitPolicy,
|
||||
metadata,
|
||||
input.Status,
|
||||
))
|
||||
}
|
||||
|
||||
func (s *Store) UpdateCatalogProvider(ctx context.Context, id string, input CatalogProviderInput) (CatalogProvider, error) {
|
||||
input = normalizeCatalogProviderInput(input)
|
||||
capabilitySchema, _ := json.Marshal(emptyObjectIfNil(input.CapabilitySchema))
|
||||
rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.DefaultRateLimitPolicy))
|
||||
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
|
||||
|
||||
return scanCatalogProvider(s.pool.QueryRow(ctx, `
|
||||
UPDATE model_catalog_providers
|
||||
SET provider_key = $2,
|
||||
provider_code = $3,
|
||||
display_name = $4,
|
||||
provider_type = $5,
|
||||
icon_path = NULLIF($6, ''),
|
||||
source = $7,
|
||||
capability_schema = $8,
|
||||
default_rate_limit_policy = $9,
|
||||
metadata = $10,
|
||||
status = $11,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
RETURNING `+catalogProviderColumns,
|
||||
id,
|
||||
input.ProviderKey,
|
||||
input.Code,
|
||||
input.DisplayName,
|
||||
input.ProviderType,
|
||||
input.IconPath,
|
||||
input.Source,
|
||||
capabilitySchema,
|
||||
rateLimitPolicy,
|
||||
metadata,
|
||||
input.Status,
|
||||
))
|
||||
}
|
||||
|
||||
func (s *Store) DeleteCatalogProvider(ctx context.Context, id string) error {
|
||||
result, err := s.pool.Exec(ctx, `DELETE FROM model_catalog_providers WHERE id = $1::uuid`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if result.RowsAffected() == 0 {
|
||||
return pgx.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func scanCatalogProvider(scanner catalogProviderScanner) (CatalogProvider, error) {
|
||||
var item CatalogProvider
|
||||
var capabilitySchema []byte
|
||||
var rateLimitPolicy []byte
|
||||
var metadata []byte
|
||||
if err := scanner.Scan(
|
||||
&item.ID,
|
||||
&item.ProviderKey,
|
||||
&item.Code,
|
||||
&item.DisplayName,
|
||||
&item.ProviderType,
|
||||
&item.IconPath,
|
||||
&item.Source,
|
||||
&capabilitySchema,
|
||||
&rateLimitPolicy,
|
||||
&metadata,
|
||||
&item.Status,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return CatalogProvider{}, err
|
||||
}
|
||||
item.CapabilitySchema = decodeObject(capabilitySchema)
|
||||
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
|
||||
item.Metadata = decodeObject(metadata)
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func normalizeCatalogProviderInput(input CatalogProviderInput) CatalogProviderInput {
|
||||
input.ProviderKey = strings.TrimSpace(input.ProviderKey)
|
||||
input.Code = strings.TrimSpace(input.Code)
|
||||
input.DisplayName = strings.TrimSpace(input.DisplayName)
|
||||
input.ProviderType = strings.TrimSpace(input.ProviderType)
|
||||
input.IconPath = strings.TrimSpace(input.IconPath)
|
||||
input.Source = strings.TrimSpace(input.Source)
|
||||
input.Status = strings.TrimSpace(input.Status)
|
||||
if input.Code == "" {
|
||||
input.Code = input.ProviderKey
|
||||
}
|
||||
if input.ProviderType == "" {
|
||||
input.ProviderType = "openai"
|
||||
}
|
||||
if input.Source == "" {
|
||||
input.Source = "gateway"
|
||||
}
|
||||
if input.Status == "" {
|
||||
input.Status = "active"
|
||||
}
|
||||
return input
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type modelCatalogSnapshot struct {
|
||||
ID string
|
||||
ProviderKey string
|
||||
CanonicalModelKey string
|
||||
ProviderModelName string
|
||||
ModelType string
|
||||
DisplayName string
|
||||
Capabilities map[string]any
|
||||
BaseBillingConfig map[string]any
|
||||
DefaultRateLimitPolicy map[string]any
|
||||
}
|
||||
|
||||
func (s *Store) CreatePlatformModel(ctx context.Context, input CreatePlatformModelInput) (PlatformModel, error) {
|
||||
base, err := s.lookupBaseModel(ctx, input.BaseModelID, input.CanonicalModelKey, input.ModelName)
|
||||
if err != nil && !IsNotFound(err) {
|
||||
return PlatformModel{}, err
|
||||
}
|
||||
if input.ModelType == "" {
|
||||
input.ModelType = base.ModelType
|
||||
}
|
||||
if input.ModelName == "" {
|
||||
input.ModelName = base.ProviderModelName
|
||||
}
|
||||
if input.DisplayName == "" {
|
||||
input.DisplayName = firstNonEmpty(base.DisplayName, input.ModelName)
|
||||
}
|
||||
if input.PricingMode == "" {
|
||||
input.PricingMode = "inherit_discount"
|
||||
}
|
||||
capabilities := input.Capabilities
|
||||
if len(capabilities) == 0 {
|
||||
capabilities = mergeObjects(base.Capabilities, input.CapabilityOverride)
|
||||
}
|
||||
billingConfig := input.BillingConfig
|
||||
if len(billingConfig) == 0 {
|
||||
billingConfig = mergeObjects(base.BaseBillingConfig, input.BillingConfigOverride)
|
||||
}
|
||||
rateLimitPolicy := input.RateLimitPolicy
|
||||
if len(rateLimitPolicy) == 0 {
|
||||
rateLimitPolicy = base.DefaultRateLimitPolicy
|
||||
}
|
||||
|
||||
capabilityOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.CapabilityOverride))
|
||||
capabilitiesJSON, _ := json.Marshal(emptyObjectIfNil(capabilities))
|
||||
billingOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingConfigOverride))
|
||||
billingJSON, _ := json.Marshal(emptyObjectIfNil(billingConfig))
|
||||
permissionJSON, _ := json.Marshal(emptyObjectIfNil(input.PermissionConfig))
|
||||
retryJSON, _ := json.Marshal(emptyObjectIfNil(input.RetryPolicy))
|
||||
rateLimitJSON, _ := json.Marshal(emptyObjectIfNil(rateLimitPolicy))
|
||||
|
||||
discount := any(nil)
|
||||
if input.DiscountFactor > 0 {
|
||||
discount = input.DiscountFactor
|
||||
}
|
||||
baseID := any(nil)
|
||||
if base.ID != "" {
|
||||
baseID = base.ID
|
||||
}
|
||||
|
||||
var model PlatformModel
|
||||
var capabilityOverrideBytes []byte
|
||||
var capabilitiesBytes []byte
|
||||
var billingOverrideBytes []byte
|
||||
var billingBytes []byte
|
||||
err = s.pool.QueryRow(ctx, `
|
||||
INSERT INTO platform_models (
|
||||
platform_id, base_model_id, model_name, model_alias, model_type, display_name,
|
||||
capability_override, capabilities, pricing_mode, discount_factor,
|
||||
pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy, enabled
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, $2::uuid, $3, NULLIF($4, ''), $5, $6,
|
||||
$7::jsonb, $8::jsonb, $9, $10::numeric,
|
||||
NULLIF($11, '')::uuid, $12::jsonb, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, true
|
||||
)
|
||||
ON CONFLICT (platform_id, model_name, model_type) DO UPDATE
|
||||
SET base_model_id = EXCLUDED.base_model_id,
|
||||
model_alias = EXCLUDED.model_alias,
|
||||
display_name = EXCLUDED.display_name,
|
||||
capability_override = EXCLUDED.capability_override,
|
||||
capabilities = EXCLUDED.capabilities,
|
||||
pricing_mode = EXCLUDED.pricing_mode,
|
||||
discount_factor = EXCLUDED.discount_factor,
|
||||
pricing_rule_set_id = EXCLUDED.pricing_rule_set_id,
|
||||
billing_config_override = EXCLUDED.billing_config_override,
|
||||
billing_config = EXCLUDED.billing_config,
|
||||
permission_config = EXCLUDED.permission_config,
|
||||
retry_policy = EXCLUDED.retry_policy,
|
||||
rate_limit_policy = EXCLUDED.rate_limit_policy,
|
||||
enabled = true,
|
||||
updated_at = now()
|
||||
RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_name,
|
||||
COALESCE(model_alias, ''), model_type, display_name, capability_override,
|
||||
capabilities, pricing_mode, COALESCE(discount_factor, 0)::float8,
|
||||
COALESCE(pricing_rule_set_id::text, ''), billing_config_override, billing_config, enabled, created_at, updated_at`,
|
||||
input.PlatformID,
|
||||
baseID,
|
||||
input.ModelName,
|
||||
input.ModelAlias,
|
||||
input.ModelType,
|
||||
input.DisplayName,
|
||||
string(capabilityOverrideJSON),
|
||||
string(capabilitiesJSON),
|
||||
input.PricingMode,
|
||||
discount,
|
||||
input.PricingRuleSetID,
|
||||
string(billingOverrideJSON),
|
||||
string(billingJSON),
|
||||
string(permissionJSON),
|
||||
string(retryJSON),
|
||||
string(rateLimitJSON),
|
||||
).Scan(
|
||||
&model.ID,
|
||||
&model.PlatformID,
|
||||
&model.BaseModelID,
|
||||
&model.ModelName,
|
||||
&model.ModelAlias,
|
||||
&model.ModelType,
|
||||
&model.DisplayName,
|
||||
&capabilityOverrideBytes,
|
||||
&capabilitiesBytes,
|
||||
&model.PricingMode,
|
||||
&model.DiscountFactor,
|
||||
&model.PricingRuleSetID,
|
||||
&billingOverrideBytes,
|
||||
&billingBytes,
|
||||
&model.Enabled,
|
||||
&model.CreatedAt,
|
||||
&model.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return PlatformModel{}, err
|
||||
}
|
||||
model.CapabilityOverride = decodeObject(capabilityOverrideBytes)
|
||||
model.Capabilities = decodeObject(capabilitiesBytes)
|
||||
model.BillingConfigOverride = decodeObject(billingOverrideBytes)
|
||||
model.BillingConfig = decodeObject(billingBytes)
|
||||
return model, nil
|
||||
}
|
||||
|
||||
func (s *Store) lookupBaseModel(ctx context.Context, id string, canonicalKey string, modelName string) (modelCatalogSnapshot, error) {
|
||||
var item modelCatalogSnapshot
|
||||
var capabilities []byte
|
||||
var billingConfig []byte
|
||||
var rateLimitPolicy []byte
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
SELECT id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
|
||||
capabilities, base_billing_config, default_rate_limit_policy
|
||||
FROM base_model_catalog
|
||||
WHERE ($1 <> '' AND id = NULLIF($1, '')::uuid)
|
||||
OR ($2 <> '' AND canonical_model_key = $2)
|
||||
OR ($3 <> '' AND provider_model_name = $3)
|
||||
ORDER BY CASE WHEN id::text = $1 THEN 0 WHEN canonical_model_key = $2 THEN 1 ELSE 2 END
|
||||
LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSpace(modelName)).Scan(
|
||||
&item.ID,
|
||||
&item.ProviderKey,
|
||||
&item.CanonicalModelKey,
|
||||
&item.ProviderModelName,
|
||||
&item.ModelType,
|
||||
&item.DisplayName,
|
||||
&capabilities,
|
||||
&billingConfig,
|
||||
&rateLimitPolicy,
|
||||
)
|
||||
if err != nil {
|
||||
if err == pgx.ErrNoRows {
|
||||
return modelCatalogSnapshot{}, err
|
||||
}
|
||||
return modelCatalogSnapshot{}, err
|
||||
}
|
||||
item.Capabilities = decodeObject(capabilities)
|
||||
item.BaseBillingConfig = decodeObject(billingConfig)
|
||||
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func mergeObjects(base map[string]any, override map[string]any) map[string]any {
|
||||
out := map[string]any{}
|
||||
for key, value := range base {
|
||||
out[key] = value
|
||||
}
|
||||
for key, value := range override {
|
||||
out[key] = value
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func emptyObjectIfNil(value map[string]any) map[string]any {
|
||||
if value == nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
return value
|
||||
}
|
||||
@@ -60,6 +60,7 @@ type Platform struct {
|
||||
Priority int `json:"priority"`
|
||||
DefaultPricingMode string `json:"defaultPricingMode"`
|
||||
DefaultDiscountFactor float64 `json:"defaultDiscountFactor"`
|
||||
PricingRuleSetID string `json:"pricingRuleSetId,omitempty"`
|
||||
Config map[string]any `json:"config,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
@@ -75,6 +76,7 @@ type CreatePlatformInput struct {
|
||||
Config map[string]any `json:"config"`
|
||||
DefaultPricingMode string `json:"defaultPricingMode"`
|
||||
DefaultDiscountFactor float64 `json:"defaultDiscountFactor"`
|
||||
PricingRuleSetID string `json:"pricingRuleSetId"`
|
||||
Priority int `json:"priority"`
|
||||
}
|
||||
|
||||
@@ -123,6 +125,7 @@ type PlatformModel struct {
|
||||
Capabilities map[string]any `json:"capabilities,omitempty"`
|
||||
PricingMode string `json:"pricingMode"`
|
||||
DiscountFactor float64 `json:"discountFactor,omitempty"`
|
||||
PricingRuleSetID string `json:"pricingRuleSetId,omitempty"`
|
||||
BillingConfigOverride map[string]any `json:"billingConfigOverride,omitempty"`
|
||||
BillingConfig map[string]any `json:"billingConfig,omitempty"`
|
||||
Enabled bool `json:"enabled"`
|
||||
@@ -133,10 +136,14 @@ type PlatformModel struct {
|
||||
type CatalogProvider struct {
|
||||
ID string `json:"id"`
|
||||
ProviderKey string `json:"providerKey"`
|
||||
Code string `json:"code"`
|
||||
DisplayName string `json:"displayName"`
|
||||
ProviderType string `json:"providerType"`
|
||||
IconPath string `json:"iconPath,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
CapabilitySchema map[string]any `json:"capabilitySchema,omitempty"`
|
||||
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
@@ -152,6 +159,7 @@ type BaseModel struct {
|
||||
Capabilities map[string]any `json:"capabilities,omitempty"`
|
||||
BaseBillingConfig map[string]any `json:"baseBillingConfig,omitempty"`
|
||||
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
PricingVersion int `json:"pricingVersion"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
@@ -159,17 +167,40 @@ type BaseModel struct {
|
||||
}
|
||||
|
||||
type PricingRule struct {
|
||||
ID string `json:"id"`
|
||||
ScopeType string `json:"scopeType"`
|
||||
ScopeID string `json:"scopeId,omitempty"`
|
||||
ResourceType string `json:"resourceType"`
|
||||
Unit string `json:"unit"`
|
||||
BasePrice float64 `json:"basePrice"`
|
||||
Currency string `json:"currency"`
|
||||
BaseWeight map[string]any `json:"baseWeight,omitempty"`
|
||||
DynamicWeight map[string]any `json:"dynamicWeight,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
ID string `json:"id"`
|
||||
RuleSetID string `json:"ruleSetId,omitempty"`
|
||||
RuleKey string `json:"ruleKey"`
|
||||
DisplayName string `json:"displayName"`
|
||||
ScopeType string `json:"scopeType"`
|
||||
ScopeID string `json:"scopeId,omitempty"`
|
||||
ResourceType string `json:"resourceType"`
|
||||
Unit string `json:"unit"`
|
||||
BasePrice float64 `json:"basePrice"`
|
||||
Currency string `json:"currency"`
|
||||
BaseWeight map[string]any `json:"baseWeight,omitempty"`
|
||||
DynamicWeight map[string]any `json:"dynamicWeight,omitempty"`
|
||||
CalculatorType string `json:"calculatorType"`
|
||||
DimensionSchema map[string]any `json:"dimensionSchema,omitempty"`
|
||||
FormulaConfig map[string]any `json:"formulaConfig,omitempty"`
|
||||
Priority int `json:"priority"`
|
||||
Status string `json:"status"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type PricingRuleSet struct {
|
||||
ID string `json:"id"`
|
||||
RuleSetKey string `json:"ruleSetKey"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Category string `json:"category"`
|
||||
Currency string `json:"currency"`
|
||||
Status string `json:"status"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
Rules []PricingRule `json:"rules,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type GatewayTenant struct {
|
||||
@@ -305,7 +336,8 @@ type TaskEvent struct {
|
||||
func (s *Store) ListPlatforms(ctx context.Context) ([]Platform, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, provider, platform_key, name, COALESCE(base_url, ''), auth_type, status, priority,
|
||||
default_pricing_mode, default_discount_factor::float8, config, created_at, updated_at
|
||||
default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''),
|
||||
config, created_at, updated_at
|
||||
FROM integration_platforms
|
||||
ORDER BY priority ASC, created_at DESC`)
|
||||
if err != nil {
|
||||
@@ -328,6 +360,7 @@ ORDER BY priority ASC, created_at DESC`)
|
||||
&platform.Priority,
|
||||
&platform.DefaultPricingMode,
|
||||
&platform.DefaultDiscountFactor,
|
||||
&platform.PricingRuleSetID,
|
||||
&configBytes,
|
||||
&platform.CreatedAt,
|
||||
&platform.UpdatedAt,
|
||||
@@ -355,11 +388,11 @@ func (s *Store) CreatePlatform(ctx context.Context, input CreatePlatformInput) (
|
||||
var platform Platform
|
||||
var configBytes []byte
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
INSERT INTO integration_platforms (provider, platform_key, name, base_url, auth_type, credentials, config, default_pricing_mode, default_discount_factor, priority)
|
||||
VALUES ($1, COALESCE(NULLIF($2, ''), 'platform_' || replace(gen_random_uuid()::text, '-', '')), $3, $4, $5, $6, $7, $8, $9, $10)
|
||||
INSERT INTO integration_platforms (provider, platform_key, name, base_url, auth_type, credentials, config, default_pricing_mode, default_discount_factor, pricing_rule_set_id, priority)
|
||||
VALUES ($1, COALESCE(NULLIF($2, ''), 'platform_' || replace(gen_random_uuid()::text, '-', '')), $3, $4, $5, $6, $7, $8, $9, NULLIF($10, '')::uuid, $11)
|
||||
RETURNING id::text, provider, platform_key, name, COALESCE(base_url, ''), auth_type, status, priority,
|
||||
default_pricing_mode, default_discount_factor::float8, config, created_at, updated_at`,
|
||||
input.Provider, input.PlatformKey, input.Name, input.BaseURL, input.AuthType, credentials, config, input.DefaultPricingMode, input.DefaultDiscountFactor, input.Priority,
|
||||
default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''), config, created_at, updated_at`,
|
||||
input.Provider, input.PlatformKey, input.Name, input.BaseURL, input.AuthType, credentials, config, input.DefaultPricingMode, input.DefaultDiscountFactor, input.PricingRuleSetID, input.Priority,
|
||||
).Scan(
|
||||
&platform.ID,
|
||||
&platform.Provider,
|
||||
@@ -371,6 +404,7 @@ RETURNING id::text, provider, platform_key, name, COALESCE(base_url, ''), auth_t
|
||||
&platform.Priority,
|
||||
&platform.DefaultPricingMode,
|
||||
&platform.DefaultDiscountFactor,
|
||||
&platform.PricingRuleSetID,
|
||||
&configBytes,
|
||||
&platform.CreatedAt,
|
||||
&platform.UpdatedAt,
|
||||
@@ -387,7 +421,8 @@ func (s *Store) ListModels(ctx context.Context) ([]PlatformModel, error) {
|
||||
SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.provider, p.name,
|
||||
m.model_name, COALESCE(m.model_alias, ''), m.model_type, m.display_name,
|
||||
m.capability_override, m.capabilities, m.pricing_mode, COALESCE(m.discount_factor, 0)::float8,
|
||||
m.billing_config_override, m.billing_config, m.enabled, m.created_at, m.updated_at
|
||||
COALESCE(m.pricing_rule_set_id::text, ''), m.billing_config_override, m.billing_config,
|
||||
m.enabled, m.created_at, m.updated_at
|
||||
FROM platform_models m
|
||||
JOIN integration_platforms p ON p.id = m.platform_id
|
||||
ORDER BY m.model_type ASC, m.model_name ASC`)
|
||||
@@ -417,6 +452,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`)
|
||||
&capabilities,
|
||||
&model.PricingMode,
|
||||
&model.DiscountFactor,
|
||||
&model.PricingRuleSetID,
|
||||
&billingConfigOverride,
|
||||
&billingConfig,
|
||||
&model.Enabled,
|
||||
@@ -436,8 +472,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`)
|
||||
|
||||
func (s *Store) ListCatalogProviders(ctx context.Context) ([]CatalogProvider, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, provider_key, display_name, provider_type, capability_schema,
|
||||
default_rate_limit_policy, status, created_at, updated_at
|
||||
SELECT `+catalogProviderColumns+`
|
||||
FROM model_catalog_providers
|
||||
ORDER BY provider_key ASC`)
|
||||
if err != nil {
|
||||
@@ -447,67 +482,10 @@ ORDER BY provider_key ASC`)
|
||||
|
||||
items := make([]CatalogProvider, 0)
|
||||
for rows.Next() {
|
||||
var item CatalogProvider
|
||||
var capabilitySchema []byte
|
||||
var rateLimitPolicy []byte
|
||||
if err := rows.Scan(
|
||||
&item.ID,
|
||||
&item.ProviderKey,
|
||||
&item.DisplayName,
|
||||
&item.ProviderType,
|
||||
&capabilitySchema,
|
||||
&rateLimitPolicy,
|
||||
&item.Status,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
item, err := scanCatalogProvider(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.CapabilitySchema = decodeObject(capabilitySchema)
|
||||
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) ListBaseModels(ctx context.Context) ([]BaseModel, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
|
||||
capabilities, base_billing_config, default_rate_limit_policy, pricing_version,
|
||||
status, created_at, updated_at
|
||||
FROM base_model_catalog
|
||||
ORDER BY provider_key ASC, model_type ASC, canonical_model_key ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]BaseModel, 0)
|
||||
for rows.Next() {
|
||||
var item BaseModel
|
||||
var capabilities []byte
|
||||
var billingConfig []byte
|
||||
var rateLimitPolicy []byte
|
||||
if err := rows.Scan(
|
||||
&item.ID,
|
||||
&item.ProviderKey,
|
||||
&item.CanonicalModelKey,
|
||||
&item.ProviderModelName,
|
||||
&item.ModelType,
|
||||
&item.DisplayName,
|
||||
&capabilities,
|
||||
&billingConfig,
|
||||
&rateLimitPolicy,
|
||||
&item.PricingVersion,
|
||||
&item.Status,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.Capabilities = decodeObject(capabilities)
|
||||
item.BaseBillingConfig = decodeObject(billingConfig)
|
||||
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
@@ -516,9 +494,11 @@ ORDER BY provider_key ASC, model_type ASC, canonical_model_key ASC`)
|
||||
func (s *Store) ListPricingRules(ctx context.Context) ([]PricingRule, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, scope_type, COALESCE(scope_id::text, ''), resource_type, unit,
|
||||
base_price::float8, currency, base_weight, dynamic_weight, created_at, updated_at
|
||||
base_price::float8, currency, base_weight, dynamic_weight,
|
||||
COALESCE(rule_set_id::text, ''), rule_key, display_name, calculator_type,
|
||||
dimension_schema, formula_config, priority, status, metadata, created_at, updated_at
|
||||
FROM model_pricing_rules
|
||||
ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
|
||||
ORDER BY COALESCE(rule_set_id::text, ''), priority ASC, resource_type ASC, created_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -529,6 +509,9 @@ ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
|
||||
var item PricingRule
|
||||
var baseWeight []byte
|
||||
var dynamicWeight []byte
|
||||
var dimensionSchema []byte
|
||||
var formulaConfig []byte
|
||||
var metadata []byte
|
||||
if err := rows.Scan(
|
||||
&item.ID,
|
||||
&item.ScopeType,
|
||||
@@ -539,6 +522,15 @@ ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
|
||||
&item.Currency,
|
||||
&baseWeight,
|
||||
&dynamicWeight,
|
||||
&item.RuleSetID,
|
||||
&item.RuleKey,
|
||||
&item.DisplayName,
|
||||
&item.CalculatorType,
|
||||
&dimensionSchema,
|
||||
&formulaConfig,
|
||||
&item.Priority,
|
||||
&item.Status,
|
||||
&metadata,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
@@ -546,6 +538,9 @@ ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
|
||||
}
|
||||
item.BaseWeight = decodeObject(baseWeight)
|
||||
item.DynamicWeight = decodeObject(dynamicWeight)
|
||||
item.DimensionSchema = decodeObject(dimensionSchema)
|
||||
item.FormulaConfig = decodeObject(formulaConfig)
|
||||
item.Metadata = decodeObject(metadata)
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
@@ -969,6 +964,17 @@ SELECT NOT EXISTS (
|
||||
).Scan(&tenantID); err != nil {
|
||||
return GatewayUser{}, err
|
||||
}
|
||||
_ = tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(default_user_group_id::text, '')
|
||||
FROM gateway_tenants
|
||||
WHERE id = $1::uuid`, tenantID).Scan(&userGroupID)
|
||||
if userGroupID == "" {
|
||||
_ = tx.QueryRow(ctx, `
|
||||
SELECT id::text
|
||||
FROM gateway_user_groups
|
||||
WHERE group_key = 'default' AND status = 'active'
|
||||
LIMIT 1`).Scan(&userGroupID)
|
||||
}
|
||||
if invitationCode != "" {
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT i.id::text, COALESCE(i.created_by::text, '')
|
||||
@@ -1182,17 +1188,8 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
|
||||
requestBody, _ := json.Marshal(input.Request)
|
||||
runMode := normalizeRunMode(input.RunMode, input.Request)
|
||||
status := "queued"
|
||||
result := map[string]any(nil)
|
||||
billings := []any(nil)
|
||||
finished := false
|
||||
if runMode == "simulation" {
|
||||
status = "succeeded"
|
||||
result = simulationResult(input.Kind, input.Model)
|
||||
billings = simulationBillings(input.Kind, input.Model)
|
||||
finished = true
|
||||
}
|
||||
resultBody, _ := json.Marshal(result)
|
||||
billingsBody, _ := json.Marshal(billings)
|
||||
resultBody, _ := json.Marshal(map[string]any(nil))
|
||||
billingsBody, _ := json.Marshal([]any(nil))
|
||||
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
@@ -1213,12 +1210,12 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
|
||||
RETURNING id::text, kind, run_mode, user_id, COALESCE(gateway_user_id::text, ''), user_source,
|
||||
COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_key, ''),
|
||||
COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model, request, status, result, billings, COALESCE(error, ''), created_at, updated_at`,
|
||||
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, status, resultBody, billingsBody, finished,
|
||||
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, status, resultBody, billingsBody, false,
|
||||
).Scan(&task.ID, &task.Kind, &task.RunMode, &task.UserID, &task.GatewayUserID, &task.UserSource, &task.GatewayTenantID, &task.TenantID, &task.TenantKey, &task.UserGroupID, &task.UserGroupKey, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
events := taskEventsForCreate(task.ID, runMode, status, result)
|
||||
events := taskEventsForCreate(task.ID, runMode, status, nil)
|
||||
for _, event := range events {
|
||||
payload, _ := json.Marshal(event.Payload)
|
||||
if _, err := tx.Exec(ctx, `
|
||||
@@ -1300,6 +1297,10 @@ func IsNotFound(err error) bool {
|
||||
return err == pgx.ErrNoRows
|
||||
}
|
||||
|
||||
func IsUniqueViolation(err error) bool {
|
||||
return isUniqueViolation(err)
|
||||
}
|
||||
|
||||
func isUniqueViolation(err error) bool {
|
||||
var pgErr *pgconn.PgError
|
||||
return errors.As(err, &pgErr) && pgErr.Code == "23505"
|
||||
@@ -1452,7 +1453,7 @@ func simulationBillings(kind string, model string) []any {
|
||||
}
|
||||
|
||||
func taskEventsForCreate(taskID string, runMode string, status string, result map[string]any) []TaskEvent {
|
||||
events := []TaskEvent{{
|
||||
return []TaskEvent{{
|
||||
TaskID: taskID,
|
||||
Seq: 1,
|
||||
EventType: "task.accepted",
|
||||
@@ -1463,33 +1464,6 @@ func taskEventsForCreate(taskID string, runMode string, status string, result ma
|
||||
Payload: map[string]any{"taskId": taskID},
|
||||
Simulated: runMode == "simulation",
|
||||
}}
|
||||
if runMode != "simulation" {
|
||||
return events
|
||||
}
|
||||
return append(events,
|
||||
TaskEvent{
|
||||
TaskID: taskID,
|
||||
Seq: 2,
|
||||
EventType: "task.progress",
|
||||
Status: "running",
|
||||
Phase: "simulation",
|
||||
Progress: 0.5,
|
||||
Message: "simulation client running",
|
||||
Payload: map[string]any{"taskId": taskID},
|
||||
Simulated: true,
|
||||
},
|
||||
TaskEvent{
|
||||
TaskID: taskID,
|
||||
Seq: 3,
|
||||
EventType: "task.completed",
|
||||
Status: status,
|
||||
Phase: "completed",
|
||||
Progress: 1,
|
||||
Message: "simulation completed",
|
||||
Payload: map[string]any{"taskId": taskID, "result": result},
|
||||
Simulated: true,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func decodeObject(bytes []byte) map[string]any {
|
||||
|
||||
@@ -0,0 +1,334 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const pricingRuleSetColumns = `
|
||||
id::text, rule_set_key, name, COALESCE(description, ''), category, currency,
|
||||
status, metadata, created_at, updated_at`
|
||||
|
||||
const pricingRuleColumns = `
|
||||
id::text, COALESCE(rule_set_id::text, ''), rule_key, display_name, scope_type,
|
||||
COALESCE(scope_id::text, ''), resource_type, unit, base_price::float8, currency,
|
||||
base_weight, dynamic_weight, calculator_type, dimension_schema, formula_config,
|
||||
priority, status, metadata, created_at, updated_at`
|
||||
|
||||
type PricingRuleInput struct {
|
||||
RuleKey string `json:"ruleKey"`
|
||||
DisplayName string `json:"displayName"`
|
||||
ResourceType string `json:"resourceType"`
|
||||
Unit string `json:"unit"`
|
||||
BasePrice float64 `json:"basePrice"`
|
||||
Currency string `json:"currency"`
|
||||
BaseWeight map[string]any `json:"baseWeight"`
|
||||
DynamicWeight map[string]any `json:"dynamicWeight"`
|
||||
CalculatorType string `json:"calculatorType"`
|
||||
DimensionSchema map[string]any `json:"dimensionSchema"`
|
||||
FormulaConfig map[string]any `json:"formulaConfig"`
|
||||
Priority int `json:"priority"`
|
||||
Status string `json:"status"`
|
||||
Metadata map[string]any `json:"metadata"`
|
||||
}
|
||||
|
||||
type PricingRuleSetInput struct {
|
||||
RuleSetKey string `json:"ruleSetKey"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Category string `json:"category"`
|
||||
Currency string `json:"currency"`
|
||||
Status string `json:"status"`
|
||||
Metadata map[string]any `json:"metadata"`
|
||||
Rules []PricingRuleInput `json:"rules"`
|
||||
}
|
||||
|
||||
type pricingScanner interface {
|
||||
Scan(dest ...any) error
|
||||
}
|
||||
|
||||
func (s *Store) ListPricingRuleSets(ctx context.Context) ([]PricingRuleSet, error) {
|
||||
rows, err := s.pool.Query(ctx, `SELECT `+pricingRuleSetColumns+` FROM model_pricing_rule_sets ORDER BY category ASC, name ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]PricingRuleSet, 0)
|
||||
byID := map[string]int{}
|
||||
for rows.Next() {
|
||||
item, err := scanPricingRuleSet(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
byID[item.ID] = len(items)
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ruleRows, err := s.pool.Query(ctx, `
|
||||
SELECT `+pricingRuleColumns+`
|
||||
FROM model_pricing_rules
|
||||
WHERE rule_set_id IS NOT NULL
|
||||
ORDER BY rule_set_id, priority ASC, resource_type ASC, rule_key ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer ruleRows.Close()
|
||||
for ruleRows.Next() {
|
||||
rule, err := scanPricingRule(ruleRows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if index, ok := byID[rule.RuleSetID]; ok {
|
||||
items[index].Rules = append(items[index].Rules, rule)
|
||||
}
|
||||
}
|
||||
return items, ruleRows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) CreatePricingRuleSet(ctx context.Context, input PricingRuleSetInput) (PricingRuleSet, error) {
|
||||
input = normalizePricingRuleSet(input)
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
|
||||
item, err := scanPricingRuleSet(tx.QueryRow(ctx, `
|
||||
INSERT INTO model_pricing_rule_sets (rule_set_key, name, description, category, currency, status, metadata)
|
||||
VALUES ($1, $2, NULLIF($3, ''), $4, $5, $6, $7)
|
||||
RETURNING `+pricingRuleSetColumns,
|
||||
input.RuleSetKey, input.Name, input.Description, input.Category, input.Currency, input.Status, metadata,
|
||||
))
|
||||
if err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
if err := insertPricingRules(ctx, tx, item.ID, input.Currency, input.Rules); err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
item.Rules = pricingInputsToRules(item.ID, input.Currency, input.Rules)
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func (s *Store) UpdatePricingRuleSet(ctx context.Context, id string, input PricingRuleSetInput) (PricingRuleSet, error) {
|
||||
input = normalizePricingRuleSet(input)
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
|
||||
item, err := scanPricingRuleSet(tx.QueryRow(ctx, `
|
||||
UPDATE model_pricing_rule_sets
|
||||
SET rule_set_key = $2,
|
||||
name = $3,
|
||||
description = NULLIF($4, ''),
|
||||
category = $5,
|
||||
currency = $6,
|
||||
status = $7,
|
||||
metadata = $8,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
RETURNING `+pricingRuleSetColumns,
|
||||
id, input.RuleSetKey, input.Name, input.Description, input.Category, input.Currency, input.Status, metadata,
|
||||
))
|
||||
if err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `DELETE FROM model_pricing_rules WHERE rule_set_id = $1::uuid`, id); err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
if err := insertPricingRules(ctx, tx, item.ID, input.Currency, input.Rules); err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
item.Rules = pricingInputsToRules(item.ID, input.Currency, input.Rules)
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func (s *Store) DeletePricingRuleSet(ctx context.Context, id string) error {
|
||||
result, err := s.pool.Exec(ctx, `DELETE FROM model_pricing_rule_sets WHERE id = $1::uuid`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if result.RowsAffected() == 0 {
|
||||
return pgx.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func insertPricingRules(ctx context.Context, tx pgx.Tx, ruleSetID string, defaultCurrency string, rules []PricingRuleInput) error {
|
||||
for index, rule := range rules {
|
||||
rule = normalizePricingRule(rule, index, defaultCurrency)
|
||||
baseWeight, _ := json.Marshal(emptyObjectIfNil(rule.BaseWeight))
|
||||
dynamicWeight, _ := json.Marshal(emptyObjectIfNil(rule.DynamicWeight))
|
||||
dimensionSchema, _ := json.Marshal(emptyObjectIfNil(rule.DimensionSchema))
|
||||
formulaConfig, _ := json.Marshal(emptyObjectIfNil(rule.FormulaConfig))
|
||||
metadata, _ := json.Marshal(emptyObjectIfNil(rule.Metadata))
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO model_pricing_rules (
|
||||
rule_set_id, rule_key, display_name, scope_type, scope_id, resource_type,
|
||||
unit, base_price, currency, base_weight, dynamic_weight, calculator_type,
|
||||
dimension_schema, formula_config, priority, status, metadata
|
||||
)
|
||||
VALUES ($1::uuid, $2, $3, 'rule_set', $1::uuid, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)`,
|
||||
ruleSetID, rule.RuleKey, rule.DisplayName, rule.ResourceType, rule.Unit,
|
||||
rule.BasePrice, rule.Currency, baseWeight, dynamicWeight, rule.CalculatorType,
|
||||
dimensionSchema, formulaConfig, rule.Priority, rule.Status, metadata,
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func scanPricingRuleSet(scanner pricingScanner) (PricingRuleSet, error) {
|
||||
var item PricingRuleSet
|
||||
var metadata []byte
|
||||
if err := scanner.Scan(
|
||||
&item.ID,
|
||||
&item.RuleSetKey,
|
||||
&item.Name,
|
||||
&item.Description,
|
||||
&item.Category,
|
||||
&item.Currency,
|
||||
&item.Status,
|
||||
&metadata,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return PricingRuleSet{}, err
|
||||
}
|
||||
item.Metadata = decodeObject(metadata)
|
||||
item.Rules = []PricingRule{}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func scanPricingRule(scanner pricingScanner) (PricingRule, error) {
|
||||
var item PricingRule
|
||||
var baseWeight []byte
|
||||
var dynamicWeight []byte
|
||||
var dimensionSchema []byte
|
||||
var formulaConfig []byte
|
||||
var metadata []byte
|
||||
if err := scanner.Scan(
|
||||
&item.ID,
|
||||
&item.RuleSetID,
|
||||
&item.RuleKey,
|
||||
&item.DisplayName,
|
||||
&item.ScopeType,
|
||||
&item.ScopeID,
|
||||
&item.ResourceType,
|
||||
&item.Unit,
|
||||
&item.BasePrice,
|
||||
&item.Currency,
|
||||
&baseWeight,
|
||||
&dynamicWeight,
|
||||
&item.CalculatorType,
|
||||
&dimensionSchema,
|
||||
&formulaConfig,
|
||||
&item.Priority,
|
||||
&item.Status,
|
||||
&metadata,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return PricingRule{}, err
|
||||
}
|
||||
item.BaseWeight = decodeObject(baseWeight)
|
||||
item.DynamicWeight = decodeObject(dynamicWeight)
|
||||
item.DimensionSchema = decodeObject(dimensionSchema)
|
||||
item.FormulaConfig = decodeObject(formulaConfig)
|
||||
item.Metadata = decodeObject(metadata)
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func normalizePricingRuleSet(input PricingRuleSetInput) PricingRuleSetInput {
|
||||
input.RuleSetKey = strings.TrimSpace(input.RuleSetKey)
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.Description = strings.TrimSpace(input.Description)
|
||||
input.Category = strings.TrimSpace(input.Category)
|
||||
input.Currency = strings.TrimSpace(input.Currency)
|
||||
input.Status = strings.TrimSpace(input.Status)
|
||||
if input.Category == "" {
|
||||
input.Category = "custom"
|
||||
}
|
||||
if input.Currency == "" {
|
||||
input.Currency = "resource"
|
||||
}
|
||||
if input.Status == "" {
|
||||
input.Status = "active"
|
||||
}
|
||||
return input
|
||||
}
|
||||
|
||||
func normalizePricingRule(input PricingRuleInput, index int, defaultCurrency string) PricingRuleInput {
|
||||
input.RuleKey = strings.TrimSpace(input.RuleKey)
|
||||
input.DisplayName = strings.TrimSpace(input.DisplayName)
|
||||
input.ResourceType = strings.TrimSpace(input.ResourceType)
|
||||
input.Unit = strings.TrimSpace(input.Unit)
|
||||
input.Currency = strings.TrimSpace(input.Currency)
|
||||
input.CalculatorType = strings.TrimSpace(input.CalculatorType)
|
||||
input.Status = strings.TrimSpace(input.Status)
|
||||
if input.RuleKey == "" {
|
||||
input.RuleKey = "rule_" + strings.ReplaceAll(input.ResourceType+"_"+input.Unit, " ", "_")
|
||||
}
|
||||
if input.DisplayName == "" {
|
||||
input.DisplayName = input.ResourceType
|
||||
}
|
||||
if input.Unit == "" {
|
||||
input.Unit = "item"
|
||||
}
|
||||
if input.Currency == "" {
|
||||
input.Currency = defaultCurrency
|
||||
}
|
||||
if input.CalculatorType == "" {
|
||||
input.CalculatorType = "unit_weight"
|
||||
}
|
||||
if input.Priority == 0 {
|
||||
input.Priority = (index + 1) * 10
|
||||
}
|
||||
if input.Status == "" {
|
||||
input.Status = "active"
|
||||
}
|
||||
return input
|
||||
}
|
||||
|
||||
func pricingInputsToRules(ruleSetID string, defaultCurrency string, rules []PricingRuleInput) []PricingRule {
|
||||
items := make([]PricingRule, 0, len(rules))
|
||||
for index, input := range rules {
|
||||
input = normalizePricingRule(input, index, defaultCurrency)
|
||||
items = append(items, PricingRule{
|
||||
RuleSetID: ruleSetID,
|
||||
RuleKey: input.RuleKey,
|
||||
DisplayName: input.DisplayName,
|
||||
ScopeType: "rule_set",
|
||||
ScopeID: ruleSetID,
|
||||
ResourceType: input.ResourceType,
|
||||
Unit: input.Unit,
|
||||
BasePrice: input.BasePrice,
|
||||
Currency: input.Currency,
|
||||
BaseWeight: emptyObjectIfNil(input.BaseWeight),
|
||||
DynamicWeight: emptyObjectIfNil(input.DynamicWeight),
|
||||
CalculatorType: input.CalculatorType,
|
||||
DimensionSchema: emptyObjectIfNil(input.DimensionSchema),
|
||||
FormulaConfig: emptyObjectIfNil(input.FormulaConfig),
|
||||
Priority: input.Priority,
|
||||
Status: input.Status,
|
||||
Metadata: emptyObjectIfNil(input.Metadata),
|
||||
})
|
||||
}
|
||||
return items
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
)
|
||||
|
||||
func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, reservations []RateLimitReservation) (RateLimitResult, error) {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return RateLimitResult{}, err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
result := RateLimitResult{}
|
||||
for _, reservation := range reservations {
|
||||
if reservation.Limit <= 0 || reservation.Amount <= 0 {
|
||||
continue
|
||||
}
|
||||
if reservation.Metric == "" || reservation.Amount > reservation.Limit {
|
||||
return RateLimitResult{}, ErrRateLimited
|
||||
}
|
||||
if reservation.WindowSeconds <= 0 {
|
||||
reservation.WindowSeconds = 60
|
||||
}
|
||||
if reservation.Metric == "concurrent" {
|
||||
if reservation.LeaseTTLSeconds <= 0 {
|
||||
reservation.LeaseTTLSeconds = 120
|
||||
}
|
||||
var active float64
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(SUM(lease_value), 0)::float8
|
||||
FROM gateway_concurrency_leases
|
||||
WHERE scope_type = $1
|
||||
AND scope_key = $2
|
||||
AND released_at IS NULL
|
||||
AND expires_at > now()`,
|
||||
reservation.ScopeType,
|
||||
reservation.ScopeKey,
|
||||
).Scan(&active); err != nil {
|
||||
return RateLimitResult{}, err
|
||||
}
|
||||
if active+reservation.Amount > reservation.Limit {
|
||||
return RateLimitResult{}, ErrRateLimited
|
||||
}
|
||||
var leaseID string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
INSERT INTO gateway_concurrency_leases (task_id, scope_type, scope_key, lease_value, expires_at)
|
||||
VALUES ($1::uuid, $2, $3, $4, now() + ($5::int * interval '1 second'))
|
||||
RETURNING id::text`,
|
||||
taskID,
|
||||
reservation.ScopeType,
|
||||
reservation.ScopeKey,
|
||||
reservation.Amount,
|
||||
reservation.LeaseTTLSeconds,
|
||||
).Scan(&leaseID); err != nil {
|
||||
return RateLimitResult{}, err
|
||||
}
|
||||
result.LeaseIDs = append(result.LeaseIDs, leaseID)
|
||||
continue
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_rate_limit_counters (
|
||||
scope_type, scope_key, metric, window_start, limit_value, used_value, reserved_value, reset_at
|
||||
)
|
||||
VALUES (
|
||||
$1, $2, $3, date_trunc('minute', now()), $4, $5, 0,
|
||||
date_trunc('minute', now()) + ($6::int * interval '1 second')
|
||||
)
|
||||
ON CONFLICT (scope_type, scope_key, metric, window_start) DO UPDATE
|
||||
SET limit_value = EXCLUDED.limit_value,
|
||||
used_value = gateway_rate_limit_counters.used_value + EXCLUDED.used_value,
|
||||
reset_at = EXCLUDED.reset_at,
|
||||
updated_at = now()
|
||||
WHERE gateway_rate_limit_counters.used_value + EXCLUDED.used_value <= EXCLUDED.limit_value`,
|
||||
reservation.ScopeType,
|
||||
reservation.ScopeKey,
|
||||
reservation.Metric,
|
||||
reservation.Limit,
|
||||
reservation.Amount,
|
||||
reservation.WindowSeconds,
|
||||
)
|
||||
if err != nil {
|
||||
return RateLimitResult{}, err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return RateLimitResult{}, ErrRateLimited
|
||||
}
|
||||
}
|
||||
return result, tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leaseIDs []string) error {
|
||||
if len(leaseIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
for _, leaseID := range leaseIDs {
|
||||
if leaseID == "" {
|
||||
continue
|
||||
}
|
||||
if _, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_concurrency_leases
|
||||
SET released_at = now()
|
||||
WHERE id = $1::uuid AND released_at IS NULL`, leaseID); err != nil && !errors.Is(err, ErrRateLimited) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package store
|
||||
|
||||
import "errors"
|
||||
|
||||
var (
|
||||
ErrNoModelCandidate = errors.New("no enabled platform model matches request")
|
||||
ErrRateLimited = errors.New("rate limit exceeded")
|
||||
)
|
||||
|
||||
type CreatePlatformModelInput struct {
|
||||
PlatformID string `json:"platformId"`
|
||||
BaseModelID string `json:"baseModelId"`
|
||||
CanonicalModelKey string `json:"canonicalModelKey"`
|
||||
ModelName string `json:"modelName"`
|
||||
ModelAlias string `json:"modelAlias"`
|
||||
ModelType string `json:"modelType"`
|
||||
DisplayName string `json:"displayName"`
|
||||
CapabilityOverride map[string]any `json:"capabilityOverride"`
|
||||
Capabilities map[string]any `json:"capabilities"`
|
||||
PricingMode string `json:"pricingMode"`
|
||||
DiscountFactor float64 `json:"discountFactor"`
|
||||
PricingRuleSetID string `json:"pricingRuleSetId"`
|
||||
BillingConfigOverride map[string]any `json:"billingConfigOverride"`
|
||||
BillingConfig map[string]any `json:"billingConfig"`
|
||||
PermissionConfig map[string]any `json:"permissionConfig"`
|
||||
RetryPolicy map[string]any `json:"retryPolicy"`
|
||||
RateLimitPolicy map[string]any `json:"rateLimitPolicy"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
type RuntimeModelCandidate struct {
|
||||
PlatformID string
|
||||
PlatformKey string
|
||||
PlatformName string
|
||||
Provider string
|
||||
BaseURL string
|
||||
AuthType string
|
||||
Credentials map[string]any
|
||||
PlatformConfig map[string]any
|
||||
DefaultPricingMode string
|
||||
DefaultDiscountFactor float64
|
||||
PlatformRetryPolicy map[string]any
|
||||
PlatformRateLimitPolicy map[string]any
|
||||
PlatformPriority int
|
||||
PlatformModelID string
|
||||
BaseModelID string
|
||||
CanonicalModelKey string
|
||||
ProviderModelName string
|
||||
ModelName string
|
||||
ModelAlias string
|
||||
ModelType string
|
||||
DisplayName string
|
||||
Capabilities map[string]any
|
||||
CapabilityOverride map[string]any
|
||||
BaseBillingConfig map[string]any
|
||||
BillingConfig map[string]any
|
||||
BillingConfigOverride map[string]any
|
||||
PricingMode string
|
||||
DiscountFactor float64
|
||||
PlatformPricingRuleSetID string
|
||||
ModelPricingRuleSetID string
|
||||
ModelRetryPolicy map[string]any
|
||||
ModelRateLimitPolicy map[string]any
|
||||
ClientID string
|
||||
QueueKey string
|
||||
}
|
||||
|
||||
type RateLimitReservation struct {
|
||||
ScopeType string
|
||||
ScopeKey string
|
||||
Metric string
|
||||
Limit float64
|
||||
Amount float64
|
||||
WindowSeconds int
|
||||
LeaseTTLSeconds int
|
||||
}
|
||||
|
||||
type RateLimitResult struct {
|
||||
LeaseIDs []string
|
||||
}
|
||||
|
||||
type CreateTaskAttemptInput struct {
|
||||
TaskID string
|
||||
AttemptNo int
|
||||
PlatformID string
|
||||
PlatformModelID string
|
||||
ClientID string
|
||||
QueueKey string
|
||||
Status string
|
||||
Simulated bool
|
||||
RequestSnapshot map[string]any
|
||||
}
|
||||
|
||||
type FinishTaskAttemptInput struct {
|
||||
AttemptID string
|
||||
Status string
|
||||
Retryable bool
|
||||
ResponseSnapshot map[string]any
|
||||
ErrorCode string
|
||||
ErrorMessage string
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType string, normalizedRequest map[string]any) error {
|
||||
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'running',
|
||||
model_type = NULLIF($2, ''),
|
||||
normalized_request = $3::jsonb,
|
||||
locked_at = now(),
|
||||
heartbeat_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID, modelType, string(normalizedJSON))
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) {
|
||||
requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot))
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
var attemptID string
|
||||
err = tx.QueryRow(ctx, `
|
||||
INSERT INTO gateway_task_attempts (
|
||||
task_id, attempt_no, platform_id, platform_model_id, client_id, queue_key,
|
||||
status, simulated, request_snapshot
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, $2, NULLIF($3, '')::uuid, NULLIF($4, '')::uuid, NULLIF($5, ''), $6,
|
||||
$7, $8, $9::jsonb
|
||||
)
|
||||
RETURNING id::text`,
|
||||
input.TaskID,
|
||||
input.AttemptNo,
|
||||
input.PlatformID,
|
||||
input.PlatformModelID,
|
||||
input.ClientID,
|
||||
input.QueueKey,
|
||||
firstNonEmpty(input.Status, "running"),
|
||||
input.Simulated,
|
||||
string(requestJSON),
|
||||
).Scan(&attemptID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET attempt_count = GREATEST(attempt_count, $2), updated_at = now()
|
||||
WHERE id = $1::uuid`, input.TaskID, input.AttemptNo); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return attemptID, tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptInput) error {
|
||||
responseJSON, _ := json.Marshal(emptyObjectIfNil(input.ResponseSnapshot))
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET status = $2,
|
||||
retryable = $3,
|
||||
response_snapshot = $4::jsonb,
|
||||
error_code = NULLIF($5, ''),
|
||||
error_message = NULLIF($6, ''),
|
||||
finished_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
input.AttemptID,
|
||||
input.Status,
|
||||
input.Retryable,
|
||||
string(responseJSON),
|
||||
input.ErrorCode,
|
||||
input.ErrorMessage,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) FinishTaskSuccess(ctx context.Context, taskID string, result map[string]any, billings []any) (GatewayTask, error) {
|
||||
resultJSON, _ := json.Marshal(emptyObjectIfNil(result))
|
||||
billingsJSON, _ := json.Marshal(billings)
|
||||
if _, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'succeeded',
|
||||
result = $2::jsonb,
|
||||
billings = $3::jsonb,
|
||||
error = NULL,
|
||||
error_code = NULL,
|
||||
error_message = NULL,
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID, string(resultJSON), string(billingsJSON)); err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
return s.GetTask(ctx, taskID)
|
||||
}
|
||||
|
||||
func (s *Store) FinishTaskFailure(ctx context.Context, taskID string, code string, message string) (GatewayTask, error) {
|
||||
if _, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'failed',
|
||||
error = NULLIF($2, ''),
|
||||
error_code = NULLIF($3, ''),
|
||||
error_message = NULLIF($2, ''),
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID, message, code); err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
return s.GetTask(ctx, taskID)
|
||||
}
|
||||
|
||||
func (s *Store) AddTaskEvent(ctx context.Context, taskID string, eventType string, status string, phase string, progress float64, message string, payload map[string]any, simulated bool) (TaskEvent, error) {
|
||||
payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload))
|
||||
var event TaskEvent
|
||||
var payloadBytes []byte
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
WITH next_seq AS (
|
||||
SELECT COALESCE(MAX(seq), 0) + 1 AS seq
|
||||
FROM gateway_task_events
|
||||
WHERE task_id = $1::uuid
|
||||
)
|
||||
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
|
||||
SELECT $1::uuid, next_seq.seq, $2, NULLIF($3, ''), NULLIF($4, ''), $5, NULLIF($6, ''), $7::jsonb, $8
|
||||
FROM next_seq
|
||||
RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''),
|
||||
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at`,
|
||||
taskID,
|
||||
eventType,
|
||||
status,
|
||||
phase,
|
||||
progress,
|
||||
message,
|
||||
string(payloadJSON),
|
||||
simulated,
|
||||
).Scan(
|
||||
&event.ID,
|
||||
&event.TaskID,
|
||||
&event.Seq,
|
||||
&event.EventType,
|
||||
&event.Status,
|
||||
&event.Phase,
|
||||
&event.Progress,
|
||||
&event.Message,
|
||||
&payloadBytes,
|
||||
&event.Simulated,
|
||||
&event.CreatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return TaskEvent{}, err
|
||||
}
|
||||
event.Payload = decodeObject(payloadBytes)
|
||||
return event, nil
|
||||
}
|
||||
|
||||
func (s *Store) QueueTaskCallback(ctx context.Context, event TaskEvent, callbackURL string) error {
|
||||
if callbackURL == "" {
|
||||
return nil
|
||||
}
|
||||
payloadJSON, _ := json.Marshal(map[string]any{
|
||||
"taskId": event.TaskID,
|
||||
"seq": event.Seq,
|
||||
"eventType": event.EventType,
|
||||
"status": event.Status,
|
||||
"phase": event.Phase,
|
||||
"progress": event.Progress,
|
||||
"message": event.Message,
|
||||
"payload": event.Payload,
|
||||
"simulated": event.Simulated,
|
||||
"createdAt": event.CreatedAt,
|
||||
})
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload)
|
||||
VALUES ($1::uuid, $2::uuid, $3, $4, $5::jsonb)
|
||||
ON CONFLICT (task_id, seq, callback_url) DO NOTHING`,
|
||||
event.TaskID,
|
||||
event.ID,
|
||||
event.Seq,
|
||||
callbackURL,
|
||||
string(payloadJSON),
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) RecordClientAssignment(ctx context.Context, candidate RuntimeModelCandidate) error {
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
INSERT INTO runtime_client_states (
|
||||
client_id, platform_id, provider, method_name, queue_key, running_count, last_assigned_at
|
||||
)
|
||||
VALUES ($1, $2::uuid, $3, $4, $5, 1, now())
|
||||
ON CONFLICT (client_id) DO UPDATE
|
||||
SET running_count = runtime_client_states.running_count + 1,
|
||||
last_assigned_at = now(),
|
||||
updated_at = now()`,
|
||||
candidate.ClientID,
|
||||
candidate.PlatformID,
|
||||
candidate.Provider,
|
||||
candidate.ModelType,
|
||||
candidate.QueueKey,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) RecordClientRelease(ctx context.Context, clientID string, lastError string) error {
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE runtime_client_states
|
||||
SET running_count = GREATEST(running_count - 1, 0),
|
||||
last_error = NULLIF($2, ''),
|
||||
updated_at = now()
|
||||
WHERE client_id = $1`, clientID, lastError)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type UserGroupPolicy struct {
|
||||
ID string
|
||||
GroupKey string
|
||||
RateLimitPolicy map[string]any
|
||||
BillingDiscountPolicy map[string]any
|
||||
}
|
||||
|
||||
func (s *Store) ResolveUserGroupPolicy(ctx context.Context, user *auth.User) (UserGroupPolicy, error) {
|
||||
userGroupID := ""
|
||||
if user != nil {
|
||||
userGroupID = user.UserGroupID
|
||||
}
|
||||
var item UserGroupPolicy
|
||||
var rateLimit []byte
|
||||
var billing []byte
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
SELECT id::text, group_key, rate_limit_policy, billing_discount_policy
|
||||
FROM gateway_user_groups
|
||||
WHERE status = 'active'
|
||||
AND (($1 <> '' AND id = NULLIF($1, '')::uuid) OR ($1 = '' AND group_key = 'default'))
|
||||
ORDER BY CASE WHEN id::text = $1 THEN 0 ELSE 1 END, priority ASC
|
||||
LIMIT 1`, userGroupID).Scan(&item.ID, &item.GroupKey, &rateLimit, &billing)
|
||||
if err != nil {
|
||||
if err == pgx.ErrNoRows {
|
||||
return UserGroupPolicy{}, nil
|
||||
}
|
||||
return UserGroupPolicy{}, err
|
||||
}
|
||||
item.RateLimitPolicy = decodeObject(rateLimit)
|
||||
item.BillingDiscountPolicy = decodeObject(billing)
|
||||
return item, nil
|
||||
}
|
||||
Reference in New Issue
Block a user