Initial project scaffold
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Cursor
@@ -0,0 +1,508 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type Store struct {
|
||||
pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
func Connect(ctx context.Context, databaseURL string) (*Store, error) {
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := pool.Ping(ctx); err != nil {
|
||||
pool.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &Store{pool: pool}, nil
|
||||
}
|
||||
|
||||
func (s *Store) Close() {
|
||||
s.pool.Close()
|
||||
}
|
||||
|
||||
func (s *Store) Ping(ctx context.Context) error {
|
||||
return s.pool.Ping(ctx)
|
||||
}
|
||||
|
||||
type Platform struct {
|
||||
ID string `json:"id"`
|
||||
Provider string `json:"provider"`
|
||||
PlatformKey string `json:"platformKey"`
|
||||
Name string `json:"name"`
|
||||
BaseURL string `json:"baseUrl,omitempty"`
|
||||
AuthType string `json:"authType"`
|
||||
Status string `json:"status"`
|
||||
Priority int `json:"priority"`
|
||||
DefaultPricingMode string `json:"defaultPricingMode"`
|
||||
DefaultDiscountFactor float64 `json:"defaultDiscountFactor"`
|
||||
Config map[string]any `json:"config,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type CreatePlatformInput struct {
|
||||
Provider string `json:"provider"`
|
||||
PlatformKey string `json:"platformKey"`
|
||||
Name string `json:"name"`
|
||||
BaseURL string `json:"baseUrl"`
|
||||
AuthType string `json:"authType"`
|
||||
Credentials map[string]any `json:"credentials"`
|
||||
Config map[string]any `json:"config"`
|
||||
DefaultPricingMode string `json:"defaultPricingMode"`
|
||||
DefaultDiscountFactor float64 `json:"defaultDiscountFactor"`
|
||||
Priority int `json:"priority"`
|
||||
}
|
||||
|
||||
type PlatformModel struct {
|
||||
ID string `json:"id"`
|
||||
PlatformID string `json:"platformId"`
|
||||
BaseModelID string `json:"baseModelId,omitempty"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
PlatformName string `json:"platformName,omitempty"`
|
||||
ModelName string `json:"modelName"`
|
||||
ModelAlias string `json:"modelAlias,omitempty"`
|
||||
ModelType string `json:"modelType"`
|
||||
DisplayName string `json:"displayName"`
|
||||
CapabilityOverride map[string]any `json:"capabilityOverride,omitempty"`
|
||||
Capabilities map[string]any `json:"capabilities,omitempty"`
|
||||
PricingMode string `json:"pricingMode"`
|
||||
DiscountFactor float64 `json:"discountFactor,omitempty"`
|
||||
BillingConfigOverride map[string]any `json:"billingConfigOverride,omitempty"`
|
||||
BillingConfig map[string]any `json:"billingConfig,omitempty"`
|
||||
Enabled bool `json:"enabled"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type CatalogProvider struct {
|
||||
ID string `json:"id"`
|
||||
ProviderKey string `json:"providerKey"`
|
||||
DisplayName string `json:"displayName"`
|
||||
ProviderType string `json:"providerType"`
|
||||
CapabilitySchema map[string]any `json:"capabilitySchema,omitempty"`
|
||||
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type BaseModel struct {
|
||||
ID string `json:"id"`
|
||||
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,omitempty"`
|
||||
BaseBillingConfig map[string]any `json:"baseBillingConfig,omitempty"`
|
||||
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"`
|
||||
PricingVersion int `json:"pricingVersion"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
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"`
|
||||
}
|
||||
|
||||
type RateLimitWindow struct {
|
||||
ScopeType string `json:"scopeType"`
|
||||
ScopeKey string `json:"scopeKey"`
|
||||
Metric string `json:"metric"`
|
||||
WindowStart time.Time `json:"windowStart"`
|
||||
LimitValue float64 `json:"limitValue"`
|
||||
UsedValue float64 `json:"usedValue"`
|
||||
ReservedValue float64 `json:"reservedValue"`
|
||||
ResetAt time.Time `json:"resetAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type CreateTaskInput struct {
|
||||
Kind string `json:"kind"`
|
||||
Model string `json:"model"`
|
||||
Request map[string]any `json:"request"`
|
||||
}
|
||||
|
||||
type GatewayTask struct {
|
||||
ID string `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
UserID string `json:"userId"`
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
Model string `json:"model"`
|
||||
Request map[string]any `json:"request,omitempty"`
|
||||
Status string `json:"status"`
|
||||
Result map[string]any `json:"result,omitempty"`
|
||||
Billings []any `json:"billings,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
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
|
||||
FROM integration_platforms
|
||||
ORDER BY priority ASC, created_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
platforms := make([]Platform, 0)
|
||||
for rows.Next() {
|
||||
var platform Platform
|
||||
var configBytes []byte
|
||||
if err := rows.Scan(
|
||||
&platform.ID,
|
||||
&platform.Provider,
|
||||
&platform.PlatformKey,
|
||||
&platform.Name,
|
||||
&platform.BaseURL,
|
||||
&platform.AuthType,
|
||||
&platform.Status,
|
||||
&platform.Priority,
|
||||
&platform.DefaultPricingMode,
|
||||
&platform.DefaultDiscountFactor,
|
||||
&configBytes,
|
||||
&platform.CreatedAt,
|
||||
&platform.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
platform.Config = decodeObject(configBytes)
|
||||
platforms = append(platforms, platform)
|
||||
}
|
||||
return platforms, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) CreatePlatform(ctx context.Context, input CreatePlatformInput) (Platform, error) {
|
||||
credentials, _ := json.Marshal(input.Credentials)
|
||||
config, _ := json.Marshal(input.Config)
|
||||
if input.DefaultPricingMode == "" {
|
||||
input.DefaultPricingMode = "inherit_discount"
|
||||
}
|
||||
if input.DefaultDiscountFactor == 0 {
|
||||
input.DefaultDiscountFactor = 1
|
||||
}
|
||||
if input.Priority == 0 {
|
||||
input.Priority = 100
|
||||
}
|
||||
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)
|
||||
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,
|
||||
).Scan(
|
||||
&platform.ID,
|
||||
&platform.Provider,
|
||||
&platform.PlatformKey,
|
||||
&platform.Name,
|
||||
&platform.BaseURL,
|
||||
&platform.AuthType,
|
||||
&platform.Status,
|
||||
&platform.Priority,
|
||||
&platform.DefaultPricingMode,
|
||||
&platform.DefaultDiscountFactor,
|
||||
&configBytes,
|
||||
&platform.CreatedAt,
|
||||
&platform.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return Platform{}, err
|
||||
}
|
||||
platform.Config = decodeObject(configBytes)
|
||||
return platform, nil
|
||||
}
|
||||
|
||||
func (s *Store) ListModels(ctx context.Context) ([]PlatformModel, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
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
|
||||
FROM platform_models m
|
||||
JOIN integration_platforms p ON p.id = m.platform_id
|
||||
ORDER BY m.model_type ASC, m.model_name ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
models := make([]PlatformModel, 0)
|
||||
for rows.Next() {
|
||||
var model PlatformModel
|
||||
var capabilityOverride []byte
|
||||
var capabilities []byte
|
||||
var billingConfigOverride []byte
|
||||
var billingConfig []byte
|
||||
if err := rows.Scan(
|
||||
&model.ID,
|
||||
&model.PlatformID,
|
||||
&model.BaseModelID,
|
||||
&model.Provider,
|
||||
&model.PlatformName,
|
||||
&model.ModelName,
|
||||
&model.ModelAlias,
|
||||
&model.ModelType,
|
||||
&model.DisplayName,
|
||||
&capabilityOverride,
|
||||
&capabilities,
|
||||
&model.PricingMode,
|
||||
&model.DiscountFactor,
|
||||
&billingConfigOverride,
|
||||
&billingConfig,
|
||||
&model.Enabled,
|
||||
&model.CreatedAt,
|
||||
&model.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
model.CapabilityOverride = decodeObject(capabilityOverride)
|
||||
model.Capabilities = decodeObject(capabilities)
|
||||
model.BillingConfigOverride = decodeObject(billingConfigOverride)
|
||||
model.BillingConfig = decodeObject(billingConfig)
|
||||
models = append(models, model)
|
||||
}
|
||||
return models, rows.Err()
|
||||
}
|
||||
|
||||
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
|
||||
FROM model_catalog_providers
|
||||
ORDER BY provider_key ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
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 {
|
||||
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()
|
||||
}
|
||||
|
||||
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
|
||||
FROM model_pricing_rules
|
||||
ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]PricingRule, 0)
|
||||
for rows.Next() {
|
||||
var item PricingRule
|
||||
var baseWeight []byte
|
||||
var dynamicWeight []byte
|
||||
if err := rows.Scan(
|
||||
&item.ID,
|
||||
&item.ScopeType,
|
||||
&item.ScopeID,
|
||||
&item.ResourceType,
|
||||
&item.Unit,
|
||||
&item.BasePrice,
|
||||
&item.Currency,
|
||||
&baseWeight,
|
||||
&dynamicWeight,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.BaseWeight = decodeObject(baseWeight)
|
||||
item.DynamicWeight = decodeObject(dynamicWeight)
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) ListRateLimitWindows(ctx context.Context) ([]RateLimitWindow, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT scope_type, scope_key, metric, window_start, limit_value::float8, used_value::float8,
|
||||
reserved_value::float8, reset_at, updated_at
|
||||
FROM gateway_rate_limit_counters
|
||||
WHERE reset_at >= now() - interval '5 minutes'
|
||||
ORDER BY window_start DESC, scope_type ASC, scope_key ASC, metric ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]RateLimitWindow, 0)
|
||||
for rows.Next() {
|
||||
var item RateLimitWindow
|
||||
if err := rows.Scan(
|
||||
&item.ScopeType,
|
||||
&item.ScopeKey,
|
||||
&item.Metric,
|
||||
&item.WindowStart,
|
||||
&item.LimitValue,
|
||||
&item.UsedValue,
|
||||
&item.ReservedValue,
|
||||
&item.ResetAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *auth.User) (GatewayTask, error) {
|
||||
requestBody, _ := json.Marshal(input.Request)
|
||||
var task GatewayTask
|
||||
var requestBytes []byte
|
||||
var resultBytes []byte
|
||||
var billingsBytes []byte
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
INSERT INTO gateway_tasks (kind, user_id, tenant_id, model, request, status)
|
||||
VALUES ($1, $2, NULLIF($3, ''), $4, $5, 'queued')
|
||||
RETURNING id::text, kind, user_id, COALESCE(tenant_id, ''), model, request, status, result, billings, COALESCE(error, ''), created_at, updated_at`,
|
||||
input.Kind, user.ID, user.TenantID, input.Model, requestBody,
|
||||
).Scan(&task.ID, &task.Kind, &task.UserID, &task.TenantID, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
task.Request = decodeObject(requestBytes)
|
||||
task.Result = decodeObject(resultBytes)
|
||||
task.Billings = decodeArray(billingsBytes)
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func (s *Store) GetTask(ctx context.Context, taskID string) (GatewayTask, error) {
|
||||
var task GatewayTask
|
||||
var requestBytes []byte
|
||||
var resultBytes []byte
|
||||
var billingsBytes []byte
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
SELECT id::text, kind, user_id, COALESCE(tenant_id, ''), model, request, status, result, billings, COALESCE(error, ''), created_at, updated_at
|
||||
FROM gateway_tasks
|
||||
WHERE id=$1`, taskID,
|
||||
).Scan(&task.ID, &task.Kind, &task.UserID, &task.TenantID, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
task.Request = decodeObject(requestBytes)
|
||||
task.Result = decodeObject(resultBytes)
|
||||
task.Billings = decodeArray(billingsBytes)
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func IsNotFound(err error) bool {
|
||||
return err == pgx.ErrNoRows
|
||||
}
|
||||
|
||||
func decodeObject(bytes []byte) map[string]any {
|
||||
if len(bytes) == 0 {
|
||||
return nil
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal(bytes, &out); err != nil {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decodeArray(bytes []byte) []any {
|
||||
if len(bytes) == 0 {
|
||||
return nil
|
||||
}
|
||||
var out []any
|
||||
if err := json.Unmarshal(bytes, &out); err != nil {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
Reference in New Issue
Block a user