3978622050
- Refactored AI gateway to utilize new cost management structures for usage tracking. - Replaced deprecated token extraction methods with a unified cost parsing approach. - Enhanced usage fallback mechanisms and introduced detailed usage metrics in responses. - Added new metering functionality to record AI usage and costs effectively. - Updated tests to reflect changes in usage parsing and cost calculations. - Introduced new API endpoints for retrieving AI usage summaries and pricing information.
97 lines
2.8 KiB
Go
97 lines
2.8 KiB
Go
package ai
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
|
|
"github.com/ultisuite/ulti-backend/internal/ai/cost"
|
|
"github.com/ultisuite/ulti-backend/internal/llm"
|
|
)
|
|
|
|
// QuotaService wraps cost policy and metering for backward compatibility.
|
|
type QuotaService struct {
|
|
db *pgxpool.Pool
|
|
policy *cost.PolicyService
|
|
meter *cost.Meter
|
|
}
|
|
|
|
func NewQuotaService(db *pgxpool.Pool) *QuotaService {
|
|
return &QuotaService{
|
|
db: db,
|
|
policy: cost.NewPolicyService(db),
|
|
meter: cost.NewMeter(db),
|
|
}
|
|
}
|
|
|
|
func (s *QuotaService) Check(ctx context.Context, externalUserID string) (SpendStatus, error) {
|
|
orgBilling, _ := s.usesOrgBilling(ctx, externalUserID)
|
|
status, err := s.policy.GetStatus(ctx, externalUserID, orgBilling)
|
|
if err != nil {
|
|
return SpendStatus{}, err
|
|
}
|
|
return status, nil
|
|
}
|
|
|
|
func (s *QuotaService) AssertAvailable(ctx context.Context, externalUserID string, provider llm.Provider, useOrgSettings bool) error {
|
|
scope := ResolveBillingScope(ctx, s.db, externalUserID, provider, useOrgSettings)
|
|
return s.policy.AssertAvailable(ctx, externalUserID, scope)
|
|
}
|
|
|
|
func (s *QuotaService) RecordUsage(ctx context.Context, in cost.RecordInput) error {
|
|
return s.meter.RecordUsage(ctx, in)
|
|
}
|
|
|
|
func (s *QuotaService) usesOrgBilling(ctx context.Context, externalUserID string) (bool, error) {
|
|
settings, err := LoadEffectiveLLMSettings(ctx, s.db, externalUserID)
|
|
if err != nil {
|
|
return true, err
|
|
}
|
|
org, err := loadOrgLLMPolicy(ctx, s.db)
|
|
if err != nil {
|
|
return true, err
|
|
}
|
|
if org.EnforceOrgProviders || len(settings.Providers) == 0 {
|
|
return true, nil
|
|
}
|
|
user, err := loadUserLLMSettings(ctx, s.db, externalUserID)
|
|
if err != nil {
|
|
return true, err
|
|
}
|
|
return len(user.Providers) == 0, nil
|
|
}
|
|
|
|
// ResolveBillingScope determines whether usage is billed to org or user's own key.
|
|
func ResolveBillingScope(ctx context.Context, db *pgxpool.Pool, externalUserID string, provider llm.Provider, useOrgSettings bool) string {
|
|
if useOrgSettings {
|
|
return cost.BillingScopeOrg
|
|
}
|
|
org, err := loadOrgLLMPolicy(ctx, db)
|
|
if err != nil || org.EnforceOrgProviders {
|
|
return cost.BillingScopeOrg
|
|
}
|
|
user, err := loadUserLLMSettings(ctx, db, externalUserID)
|
|
if err != nil || len(user.Providers) == 0 {
|
|
return cost.BillingScopeOrg
|
|
}
|
|
apiKey := strings.TrimSpace(provider.APIKey)
|
|
for _, op := range org.Providers {
|
|
if op.ID == provider.ID && strings.TrimSpace(op.APIKey) == apiKey && apiKey != "" {
|
|
return cost.BillingScopeOrg
|
|
}
|
|
}
|
|
for _, up := range user.Providers {
|
|
if up.ID == provider.ID && strings.TrimSpace(up.APIKey) != "" {
|
|
return cost.BillingScopeUser
|
|
}
|
|
}
|
|
return cost.BillingScopeOrg
|
|
}
|
|
|
|
// SpendStatus is the user-facing quota/spend response.
|
|
type SpendStatus = cost.SpendStatus
|
|
|
|
// ErrQuotaExceeded aliases cost limit error for backward compatibility.
|
|
var ErrQuotaExceeded = cost.ErrCostLimitExceeded
|