feat(server): meter AI point usage across tenant AI features
Adds an ai_point_usage_logs ledger and a reserve/refund/complete pipeline that charges AI points for article selection optimize, template analyze /title/outline, and KOL prompt generate/optimize. Pending reservations are reconciled when the kol-assist worker and template-assist tasks finish or fail, so points refund automatically on errors. Workspace now exposes the AI quota status and a paginated usage ledger. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -27,10 +27,12 @@ func main() {
|
||||
defer cancel()
|
||||
|
||||
assistworker.NewKolAssistWorker(
|
||||
app.DB,
|
||||
app.RabbitMQ,
|
||||
app.KolAssists,
|
||||
app.LLM,
|
||||
app.Logger,
|
||||
app.Cache,
|
||||
app.Config.Generation.WorkerConcurrency,
|
||||
app.Config.Generation.ArticleTimeout,
|
||||
).Start(ctx)
|
||||
|
||||
@@ -0,0 +1,416 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/response"
|
||||
"github.com/geo-platform/tenant-api/internal/tenant/repository"
|
||||
)
|
||||
|
||||
const (
|
||||
AIUsageTypeArticleSelectionOptimize = "article_selection_optimize"
|
||||
AIUsageTypeTemplateAnalyze = "template_analyze"
|
||||
AIUsageTypeTemplateTitleGenerate = "template_title_generate"
|
||||
AIUsageTypeTemplateOutlineGenerate = "template_outline_generate"
|
||||
AIUsageTypeKolPromptGenerate = "kol_prompt_generate"
|
||||
AIUsageTypeKolPromptOptimize = "kol_prompt_optimize"
|
||||
|
||||
aiPointsQuotaType = "ai_points"
|
||||
)
|
||||
|
||||
type AIPointReserveInput struct {
|
||||
TenantID int64
|
||||
OperatorID int64
|
||||
UsageType string
|
||||
ResourceType *string
|
||||
ResourceID *int64
|
||||
ResourceUID *string
|
||||
MeteredText string
|
||||
FixedPoints int
|
||||
Metadata map[string]any
|
||||
}
|
||||
|
||||
type AIPointReservation struct {
|
||||
ReservationID int64
|
||||
UsageLogID int64
|
||||
Points int
|
||||
RequestChars int
|
||||
BaseChars int
|
||||
}
|
||||
|
||||
func ReserveAIPoints(
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
cache sharedcache.Cache,
|
||||
input AIPointReserveInput,
|
||||
) (*AIPointReservation, error) {
|
||||
if pool == nil {
|
||||
return nil, response.ErrServiceUnavailable(50344, "ai_points_unavailable", "ai points store is unavailable")
|
||||
}
|
||||
|
||||
input.UsageType = strings.TrimSpace(input.UsageType)
|
||||
if input.UsageType == "" {
|
||||
return nil, response.ErrBadRequest(40075, "ai_points_usage_type_required", "usage type is required")
|
||||
}
|
||||
|
||||
tx, err := pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50120, "ai_points_tx_failed", "failed to begin ai points transaction")
|
||||
}
|
||||
defer rollbackGenerationTx(tx)
|
||||
|
||||
if err := lockTenantAIPoints(ctx, tx, input.TenantID); err != nil {
|
||||
return nil, response.ErrInternal(50121, "ai_points_lock_failed", "failed to lock ai points quota")
|
||||
}
|
||||
|
||||
now := time.Now().UTC()
|
||||
quotaRepo := repository.NewQuotaRepository(tx)
|
||||
status, err := quotaRepo.GetAIQuotaStatus(ctx, input.TenantID, now)
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are not available for this plan")
|
||||
}
|
||||
return nil, response.ErrInternal(50122, "ai_points_query_failed", "failed to query ai points")
|
||||
}
|
||||
|
||||
requestChars := countAIPointChars(input.MeteredText)
|
||||
points := calculateAIPointCost(requestChars, status.BaseChars)
|
||||
if input.FixedPoints > 0 {
|
||||
points = input.FixedPoints
|
||||
}
|
||||
if status.Total <= 0 || status.Balance < points {
|
||||
return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are insufficient, please upgrade your plan")
|
||||
}
|
||||
|
||||
resourceType := "ai_usage"
|
||||
if input.ResourceType != nil && strings.TrimSpace(*input.ResourceType) != "" {
|
||||
resourceType = strings.TrimSpace(*input.ResourceType)
|
||||
}
|
||||
|
||||
reservationID, err := quotaRepo.CreateQuotaReservation(ctx, repository.QuotaReservationInput{
|
||||
TenantID: input.TenantID,
|
||||
OperatorID: input.OperatorID,
|
||||
QuotaType: aiPointsQuotaType,
|
||||
ResourceType: resourceType,
|
||||
ResourceID: input.ResourceID,
|
||||
ReservedAmount: points,
|
||||
ExpireAt: now.Add(time.Hour),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50123, "ai_points_reserve_failed", "failed to reserve ai points")
|
||||
}
|
||||
|
||||
metadata := input.Metadata
|
||||
if metadata == nil {
|
||||
metadata = map[string]any{}
|
||||
}
|
||||
metadataJSON, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50124, "ai_points_metadata_failed", "failed to encode ai points metadata")
|
||||
}
|
||||
usageID, err := repository.NewAIPointUsageRepository(tx).Create(ctx, repository.CreateAIPointUsageLogInput{
|
||||
TenantID: input.TenantID,
|
||||
OperatorID: input.OperatorID,
|
||||
QuotaReservationID: reservationID,
|
||||
UsageType: input.UsageType,
|
||||
ResourceType: &resourceType,
|
||||
ResourceID: input.ResourceID,
|
||||
ResourceUID: input.ResourceUID,
|
||||
RequestChars: requestChars,
|
||||
BaseChars: status.BaseChars,
|
||||
Points: points,
|
||||
MetadataJSON: metadataJSON,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50125, "ai_points_audit_failed", "failed to create ai points audit log")
|
||||
}
|
||||
|
||||
reason := "ai_points_reserve"
|
||||
referenceType := "ai_point_usage"
|
||||
balanceAfter := status.Balance - points
|
||||
if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
||||
TenantID: input.TenantID,
|
||||
OperatorID: input.OperatorID,
|
||||
QuotaType: aiPointsQuotaType,
|
||||
Delta: -points,
|
||||
BalanceAfter: balanceAfter,
|
||||
Reason: &reason,
|
||||
ReferenceType: &referenceType,
|
||||
ReferenceID: &usageID,
|
||||
}); err != nil {
|
||||
return nil, response.ErrInternal(50126, "ai_points_ledger_failed", "failed to write ai points ledger")
|
||||
}
|
||||
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return nil, response.ErrInternal(50127, "ai_points_commit_failed", "failed to commit ai points reservation")
|
||||
}
|
||||
deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(input.TenantID))
|
||||
|
||||
return &AIPointReservation{
|
||||
ReservationID: reservationID,
|
||||
UsageLogID: usageID,
|
||||
Points: points,
|
||||
RequestChars: requestChars,
|
||||
BaseChars: status.BaseChars,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func CompleteAIPoints(
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
tenantID int64,
|
||||
reservation AIPointReservation,
|
||||
model string,
|
||||
) error {
|
||||
if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
tx, err := pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollbackGenerationTx(tx)
|
||||
|
||||
pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID)
|
||||
if err != nil || !pending {
|
||||
return err
|
||||
}
|
||||
if err := confirmAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil {
|
||||
return err
|
||||
}
|
||||
var modelPtr *string
|
||||
if strings.TrimSpace(model) != "" {
|
||||
normalized := strings.TrimSpace(model)
|
||||
modelPtr = &normalized
|
||||
}
|
||||
if err := repository.NewAIPointUsageRepository(tx).MarkCompleted(ctx, tenantID, reservation.UsageLogID, modelPtr); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func RefundAIPoints(
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
cache sharedcache.Cache,
|
||||
tenantID int64,
|
||||
operatorID int64,
|
||||
reservation AIPointReservation,
|
||||
errorMessage string,
|
||||
) error {
|
||||
if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 || reservation.Points <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
tx, err := pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollbackGenerationTx(tx)
|
||||
|
||||
if err := lockTenantAIPoints(ctx, tx, tenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID)
|
||||
if err != nil || !pending {
|
||||
return err
|
||||
}
|
||||
|
||||
status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
quotaRepo := repository.NewQuotaRepository(tx)
|
||||
reason := "ai_points_refund"
|
||||
referenceType := "ai_point_usage"
|
||||
usageID := reservation.UsageLogID
|
||||
if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
||||
TenantID: tenantID,
|
||||
OperatorID: operatorID,
|
||||
QuotaType: aiPointsQuotaType,
|
||||
Delta: reservation.Points,
|
||||
BalanceAfter: status.Balance + reservation.Points,
|
||||
Reason: &reason,
|
||||
ReferenceType: &referenceType,
|
||||
ReferenceID: &usageID,
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := refundAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(tenantID))
|
||||
return nil
|
||||
}
|
||||
|
||||
func CompletePendingAIPointsByResource(
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
tenantID int64,
|
||||
resourceType string,
|
||||
resourceID *int64,
|
||||
resourceUID *string,
|
||||
model string,
|
||||
) error {
|
||||
reservation, err := findPendingAIPointReservation(ctx, pool, tenantID, resourceType, resourceID, resourceUID)
|
||||
if err != nil || reservation == nil {
|
||||
return err
|
||||
}
|
||||
return CompleteAIPoints(ctx, pool, tenantID, *reservation, model)
|
||||
}
|
||||
|
||||
func RefundPendingAIPointsByResource(
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
cache sharedcache.Cache,
|
||||
tenantID int64,
|
||||
resourceType string,
|
||||
resourceID *int64,
|
||||
resourceUID *string,
|
||||
errorMessage string,
|
||||
) error {
|
||||
reservation, operatorID, err := findPendingAIPointReservationWithOperator(ctx, pool, tenantID, resourceType, resourceID, resourceUID)
|
||||
if err != nil || reservation == nil {
|
||||
return err
|
||||
}
|
||||
return RefundAIPoints(ctx, pool, cache, tenantID, operatorID, *reservation, errorMessage)
|
||||
}
|
||||
|
||||
func findPendingAIPointReservation(
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
tenantID int64,
|
||||
resourceType string,
|
||||
resourceID *int64,
|
||||
resourceUID *string,
|
||||
) (*AIPointReservation, error) {
|
||||
reservation, _, err := findPendingAIPointReservationWithOperator(ctx, pool, tenantID, resourceType, resourceID, resourceUID)
|
||||
return reservation, err
|
||||
}
|
||||
|
||||
func findPendingAIPointReservationWithOperator(
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
tenantID int64,
|
||||
resourceType string,
|
||||
resourceID *int64,
|
||||
resourceUID *string,
|
||||
) (*AIPointReservation, int64, error) {
|
||||
if pool == nil {
|
||||
return nil, 0, nil
|
||||
}
|
||||
log, err := repository.NewAIPointUsageRepository(pool).FindPendingByResource(ctx, tenantID, resourceType, resourceID, resourceUID)
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, 0, nil
|
||||
}
|
||||
return nil, 0, err
|
||||
}
|
||||
if log == nil || log.QuotaReservationID == nil {
|
||||
return nil, 0, nil
|
||||
}
|
||||
operatorID := int64(0)
|
||||
if log.OperatorID != nil {
|
||||
operatorID = *log.OperatorID
|
||||
}
|
||||
return &AIPointReservation{
|
||||
ReservationID: *log.QuotaReservationID,
|
||||
UsageLogID: log.ID,
|
||||
Points: log.Points,
|
||||
RequestChars: log.RequestChars,
|
||||
BaseChars: log.BaseChars,
|
||||
}, operatorID, nil
|
||||
}
|
||||
|
||||
func countAIPointChars(text string) int {
|
||||
return utf8.RuneCountInString(strings.TrimSpace(text))
|
||||
}
|
||||
|
||||
func calculateAIPointCost(chars, baseChars int) int {
|
||||
if baseChars <= 0 {
|
||||
baseChars = 1000
|
||||
}
|
||||
if chars <= 0 {
|
||||
return 1
|
||||
}
|
||||
return chars/baseChars + 1
|
||||
}
|
||||
|
||||
func lockTenantAIPoints(ctx context.Context, tx pgx.Tx, tenantID int64) error {
|
||||
_, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(20260428, hashint8($1)::int)`, tenantID)
|
||||
return err
|
||||
}
|
||||
|
||||
func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogID int64) (bool, error) {
|
||||
var status string
|
||||
err := tx.QueryRow(ctx, `
|
||||
SELECT status
|
||||
FROM ai_point_usage_logs
|
||||
WHERE id = $1 AND tenant_id = $2
|
||||
FOR UPDATE
|
||||
`, usageLogID, tenantID).Scan(&status)
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return false, nil
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
return status == "pending", nil
|
||||
}
|
||||
|
||||
func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE quota_reservations
|
||||
SET status = 'confirmed',
|
||||
consumed_amount = reserved_amount,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND tenant_id = $2
|
||||
AND quota_type = $3
|
||||
AND status = 'pending'
|
||||
`, reservationID, tenantID, aiPointsQuotaType)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return errors.New("ai points reservation is not pending")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE quota_reservations
|
||||
SET status = 'refunded',
|
||||
refunded_amount = reserved_amount,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND tenant_id = $2
|
||||
AND quota_type = $3
|
||||
AND status = 'pending'
|
||||
`, reservationID, tenantID, aiPointsQuotaType)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return errors.New("ai points reservation is not pending")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package app
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestCalculateAIPointCost(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
chars int
|
||||
baseChars int
|
||||
want int
|
||||
}{
|
||||
{name: "empty still charges base usage", chars: 0, baseChars: 1000, want: 1},
|
||||
{name: "single char", chars: 1, baseChars: 1000, want: 1},
|
||||
{name: "first base boundary", chars: 1000, baseChars: 1000, want: 2},
|
||||
{name: "over first base", chars: 1001, baseChars: 1000, want: 2},
|
||||
{name: "second base boundary", chars: 2000, baseChars: 1000, want: 3},
|
||||
{name: "over second base", chars: 2001, baseChars: 1000, want: 3},
|
||||
{name: "default base chars", chars: 1000, baseChars: 0, want: 2},
|
||||
{name: "custom base chars", chars: 1500, baseChars: 500, want: 4},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := calculateAIPointCost(tt.chars, tt.baseChars); got != tt.want {
|
||||
t.Fatalf("calculateAIPointCost(%d, %d) = %d, want %d", tt.chars, tt.baseChars, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/auth"
|
||||
sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/llm"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/response"
|
||||
)
|
||||
@@ -38,6 +39,7 @@ type ArticleSelectionOptimizeService struct {
|
||||
pool *pgxpool.Pool
|
||||
llm llm.Client
|
||||
defaultMaxOutTokens int64
|
||||
cache sharedcache.Cache
|
||||
}
|
||||
|
||||
func NewArticleSelectionOptimizeService(
|
||||
@@ -52,6 +54,11 @@ func NewArticleSelectionOptimizeService(
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ArticleSelectionOptimizeService) WithCache(c sharedcache.Cache) *ArticleSelectionOptimizeService {
|
||||
s.cache = c
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *ArticleSelectionOptimizeService) ValidateRequest(
|
||||
ctx context.Context,
|
||||
articleID int64,
|
||||
@@ -70,9 +77,29 @@ func (s *ArticleSelectionOptimizeService) ValidateRequest(
|
||||
|
||||
func (s *ArticleSelectionOptimizeService) OptimizeSelection(
|
||||
ctx context.Context,
|
||||
articleID int64,
|
||||
req ArticleSelectionOptimizeRequest,
|
||||
onDelta func(string),
|
||||
) (*ArticleSelectionOptimizeResult, error) {
|
||||
actor := auth.MustActor(ctx)
|
||||
resourceType := "article"
|
||||
reservation, err := ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{
|
||||
TenantID: actor.TenantID,
|
||||
OperatorID: actor.UserID,
|
||||
UsageType: AIUsageTypeArticleSelectionOptimize,
|
||||
ResourceType: &resourceType,
|
||||
ResourceID: &articleID,
|
||||
MeteredText: req.SelectedText,
|
||||
Metadata: map[string]any{
|
||||
"article_id": articleID,
|
||||
"instruction": strings.TrimSpace(req.Instruction),
|
||||
"title_length": countAIPointChars(req.Title),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
result, err := s.llm.Generate(
|
||||
ctx,
|
||||
llm.GenerateRequest{
|
||||
@@ -83,14 +110,27 @@ func (s *ArticleSelectionOptimizeService) OptimizeSelection(
|
||||
onDelta,
|
||||
)
|
||||
if err != nil {
|
||||
refundCtx, cancel := newGenerationCleanupContext()
|
||||
_ = RefundAIPoints(refundCtx, s.pool, s.cache, actor.TenantID, actor.UserID, *reservation, err.Error())
|
||||
cancel()
|
||||
return nil, response.ErrInternal(50019, "article_selection_optimize_failed", err.Error())
|
||||
}
|
||||
|
||||
content := strings.TrimSpace(result.Content)
|
||||
if content == "" {
|
||||
refundCtx, cancel := newGenerationCleanupContext()
|
||||
_ = RefundAIPoints(refundCtx, s.pool, s.cache, actor.TenantID, actor.UserID, *reservation, "optimized content is empty")
|
||||
cancel()
|
||||
return nil, response.ErrInternal(50019, "article_selection_optimize_failed", "optimized content is empty")
|
||||
}
|
||||
|
||||
completeCtx, cancel := newGenerationCleanupContext()
|
||||
if err := CompleteAIPoints(completeCtx, s.pool, actor.TenantID, *reservation, result.Model); err != nil {
|
||||
cancel()
|
||||
return nil, response.ErrInternal(50019, "article_selection_optimize_failed", "failed to confirm ai points usage")
|
||||
}
|
||||
cancel()
|
||||
|
||||
return &ArticleSelectionOptimizeResult{
|
||||
Content: content,
|
||||
Model: result.Model,
|
||||
|
||||
@@ -47,26 +47,42 @@ func (s *KolAssistService) Run(
|
||||
req.Description = s.enrichGenerateDescription(ctx, req.Description)
|
||||
}
|
||||
|
||||
reservation, err := s.reserveStreamingAIPoints(ctx, actor, req)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
generateReq, err := BuildKolAssistGenerateRequest(req)
|
||||
if err != nil {
|
||||
s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
|
||||
return nil, "", response.ErrInternal(50087, "kol_assist_prompt_build_failed", err.Error())
|
||||
}
|
||||
|
||||
streamAccumulator := newKolAssistContentStream(onDelta)
|
||||
generated, err := s.llm.Generate(ctx, generateReq, streamAccumulator.ConsumeRawDelta)
|
||||
if err != nil {
|
||||
s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
result, err := ParseKolAssistResponse(generated.Content)
|
||||
if err != nil {
|
||||
s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
|
||||
return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", err.Error())
|
||||
}
|
||||
result, err = FinalizeKolAssistResult(req, result)
|
||||
if err != nil {
|
||||
s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
|
||||
return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", err.Error())
|
||||
}
|
||||
|
||||
completeCtx, cancel := newGenerationCleanupContext()
|
||||
if err := CompleteAIPoints(completeCtx, s.pool, actor.TenantID, *reservation, generated.Model); err != nil {
|
||||
cancel()
|
||||
return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", "failed to confirm ai points usage")
|
||||
}
|
||||
cancel()
|
||||
|
||||
return result, generated.Model, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -6,8 +6,10 @@ import (
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/auth"
|
||||
sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/config"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/llm"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
|
||||
@@ -48,15 +50,18 @@ type KolAssistTaskResponse struct {
|
||||
}
|
||||
|
||||
type KolAssistService struct {
|
||||
pool *pgxpool.Pool
|
||||
profileSvc *KolProfileService
|
||||
repo repository.KolAssistRepository
|
||||
llm llm.Client
|
||||
messaging *rabbitmq.Client
|
||||
logger *zap.Logger
|
||||
urlResolver *kolAssistURLResolver
|
||||
cache sharedcache.Cache
|
||||
}
|
||||
|
||||
func NewKolAssistService(
|
||||
pool *pgxpool.Pool,
|
||||
profileSvc *KolProfileService,
|
||||
repo repository.KolAssistRepository,
|
||||
llmClient llm.Client,
|
||||
@@ -65,6 +70,7 @@ func NewKolAssistService(
|
||||
logger *zap.Logger,
|
||||
) *KolAssistService {
|
||||
return &KolAssistService{
|
||||
pool: pool,
|
||||
profileSvc: profileSvc,
|
||||
repo: repo,
|
||||
llm: llmClient,
|
||||
@@ -74,6 +80,11 @@ func NewKolAssistService(
|
||||
}
|
||||
}
|
||||
|
||||
func (s *KolAssistService) WithCache(c sharedcache.Cache) *KolAssistService {
|
||||
s.cache = c
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *KolAssistService) Submit(ctx context.Context, actor auth.Actor, req AssistRequest) (string, error) {
|
||||
profile, err := s.profileSvc.RequireActiveKol(ctx, actor)
|
||||
if err != nil {
|
||||
@@ -94,6 +105,11 @@ func (s *KolAssistService) Submit(ctx context.Context, actor auth.Actor, req Ass
|
||||
}
|
||||
|
||||
taskID := uuid.NewString()
|
||||
reservation, err := s.reserveAIPoints(ctx, actor, taskID, req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if err := s.repo.Create(ctx, repository.CreateKolAssistInput{
|
||||
ID: taskID,
|
||||
TenantID: profile.TenantID,
|
||||
@@ -103,6 +119,7 @@ func (s *KolAssistService) Submit(ctx context.Context, actor auth.Actor, req Ass
|
||||
TaskType: taskType,
|
||||
RequestJSON: requestJSON,
|
||||
}); err != nil {
|
||||
s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
|
||||
return "", response.ErrInternal(50084, "kol_assist_create_failed", err.Error())
|
||||
}
|
||||
|
||||
@@ -111,6 +128,7 @@ func (s *KolAssistService) Submit(ctx context.Context, actor auth.Actor, req Ass
|
||||
TenantID: profile.TenantID,
|
||||
}); err != nil {
|
||||
_ = s.repo.MarkFailed(context.Background(), profile.TenantID, taskID, err.Error())
|
||||
s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
|
||||
return "", response.ErrServiceUnavailable(50342, "kol_assist_queue_unavailable", "assist queue is busy, please retry")
|
||||
}
|
||||
|
||||
@@ -184,3 +202,75 @@ func normalizeAssistRequest(req AssistRequest) (AssistRequest, string, error) {
|
||||
|
||||
return req, "ai_" + req.Mode, nil
|
||||
}
|
||||
|
||||
func (s *KolAssistService) reserveAIPoints(
|
||||
ctx context.Context,
|
||||
actor auth.Actor,
|
||||
taskID string,
|
||||
req AssistRequest,
|
||||
) (*AIPointReservation, error) {
|
||||
usageType := AIUsageTypeKolPromptGenerate
|
||||
if req.Mode == kolAssistModeOptimize {
|
||||
usageType = AIUsageTypeKolPromptOptimize
|
||||
}
|
||||
resourceType := "kol_assist_task"
|
||||
resourceUID := taskID
|
||||
var resourceID *int64
|
||||
if req.PromptID != nil && *req.PromptID > 0 {
|
||||
resourceID = req.PromptID
|
||||
}
|
||||
return ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{
|
||||
TenantID: actor.TenantID,
|
||||
OperatorID: actor.UserID,
|
||||
UsageType: usageType,
|
||||
ResourceType: &resourceType,
|
||||
ResourceID: resourceID,
|
||||
ResourceUID: &resourceUID,
|
||||
MeteredText: buildKolAssistAIPointMeteredText(req),
|
||||
Metadata: map[string]any{
|
||||
"task_id": taskID,
|
||||
"mode": req.Mode,
|
||||
"prompt_id": req.PromptID,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *KolAssistService) reserveStreamingAIPoints(ctx context.Context, actor auth.Actor, req AssistRequest) (*AIPointReservation, error) {
|
||||
usageType := AIUsageTypeKolPromptGenerate
|
||||
if req.Mode == kolAssistModeOptimize {
|
||||
usageType = AIUsageTypeKolPromptOptimize
|
||||
}
|
||||
resourceType := "kol_prompt"
|
||||
var resourceID *int64
|
||||
if req.PromptID != nil && *req.PromptID > 0 {
|
||||
resourceID = req.PromptID
|
||||
}
|
||||
return ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{
|
||||
TenantID: actor.TenantID,
|
||||
OperatorID: actor.UserID,
|
||||
UsageType: usageType,
|
||||
ResourceType: &resourceType,
|
||||
ResourceID: resourceID,
|
||||
MeteredText: buildKolAssistAIPointMeteredText(req),
|
||||
Metadata: map[string]any{
|
||||
"mode": req.Mode,
|
||||
"prompt_id": req.PromptID,
|
||||
"stream": true,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *KolAssistService) refundReservedAIPoints(ctx context.Context, actor auth.Actor, reservation *AIPointReservation, errorMessage string) {
|
||||
if reservation == nil {
|
||||
return
|
||||
}
|
||||
_ = RefundAIPoints(ctx, s.pool, s.cache, actor.TenantID, actor.UserID, *reservation, errorMessage)
|
||||
}
|
||||
|
||||
func buildKolAssistAIPointMeteredText(req AssistRequest) string {
|
||||
parts := []string{req.Description}
|
||||
if req.Mode == kolAssistModeOptimize {
|
||||
parts = append(parts, req.CurrentContent)
|
||||
}
|
||||
return strings.TrimSpace(strings.Join(parts, "\n"))
|
||||
}
|
||||
|
||||
@@ -144,6 +144,11 @@ func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int6
|
||||
}
|
||||
|
||||
taskID := uuid.NewString()
|
||||
reservation, err := s.reserveTemplateAssistAIPoints(ctx, actor.TenantID, actor.UserID, templateID, taskID, templateAssistTaskTypeAnalyze, buildAnalyzeAIPointMeteredText(req))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
repo := repository.NewTemplateAssistTaskRepository(s.pool)
|
||||
if err := repo.CreateTask(ctx, repository.CreateTemplateAssistTaskInput{
|
||||
ID: taskID,
|
||||
@@ -153,6 +158,7 @@ func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int6
|
||||
TemplateID: templateID,
|
||||
RequestJSON: payload,
|
||||
}); err != nil {
|
||||
s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error())
|
||||
return nil, response.ErrInternal(50021, "task_create_failed", err.Error())
|
||||
}
|
||||
|
||||
@@ -168,6 +174,7 @@ func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int6
|
||||
}
|
||||
if err := publishTemplateAssistTask(ctx, s.rabbitMQ, job); err != nil {
|
||||
_ = repo.FailTask(context.Background(), taskID, templateAssistTaskTypeAnalyze, actor.TenantID, err.Error(), time.Now())
|
||||
s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error())
|
||||
return nil, response.ErrServiceUnavailable(50312, "analyze_queue_unavailable", "analyze queue is busy, please retry")
|
||||
}
|
||||
|
||||
@@ -225,6 +232,11 @@ func (s *TemplateService) CreateTitleTask(ctx context.Context, templateID int64,
|
||||
}
|
||||
|
||||
taskID := uuid.NewString()
|
||||
reservation, err := s.reserveTemplateAssistAIPoints(ctx, actor.TenantID, actor.UserID, templateID, taskID, templateAssistTaskTypeTitle, buildTitleAIPointMeteredText(req))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
repo := repository.NewTemplateAssistTaskRepository(s.pool)
|
||||
if err := repo.CreateTask(ctx, repository.CreateTemplateAssistTaskInput{
|
||||
ID: taskID,
|
||||
@@ -234,6 +246,7 @@ func (s *TemplateService) CreateTitleTask(ctx context.Context, templateID int64,
|
||||
TemplateID: templateID,
|
||||
RequestJSON: payload,
|
||||
}); err != nil {
|
||||
s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error())
|
||||
return nil, response.ErrInternal(50025, "task_create_failed", err.Error())
|
||||
}
|
||||
|
||||
@@ -249,6 +262,7 @@ func (s *TemplateService) CreateTitleTask(ctx context.Context, templateID int64,
|
||||
}
|
||||
if err := publishTemplateAssistTask(ctx, s.rabbitMQ, job); err != nil {
|
||||
_ = repo.FailTask(context.Background(), taskID, templateAssistTaskTypeTitle, actor.TenantID, err.Error(), time.Now())
|
||||
s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error())
|
||||
return nil, response.ErrServiceUnavailable(50313, "title_queue_unavailable", "title queue is busy, please retry")
|
||||
}
|
||||
|
||||
@@ -306,6 +320,11 @@ func (s *TemplateService) CreateOutlineTask(ctx context.Context, templateID int6
|
||||
}
|
||||
|
||||
taskID := uuid.NewString()
|
||||
reservation, err := s.reserveTemplateAssistAIPoints(ctx, actor.TenantID, actor.UserID, templateID, taskID, templateAssistTaskTypeOutline, buildOutlineAIPointMeteredText(req))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
repo := repository.NewTemplateAssistTaskRepository(s.pool)
|
||||
if err := repo.CreateTask(ctx, repository.CreateTemplateAssistTaskInput{
|
||||
ID: taskID,
|
||||
@@ -315,6 +334,7 @@ func (s *TemplateService) CreateOutlineTask(ctx context.Context, templateID int6
|
||||
TemplateID: templateID,
|
||||
RequestJSON: payload,
|
||||
}); err != nil {
|
||||
s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error())
|
||||
return nil, response.ErrInternal(50029, "task_create_failed", err.Error())
|
||||
}
|
||||
|
||||
@@ -330,6 +350,7 @@ func (s *TemplateService) CreateOutlineTask(ctx context.Context, templateID int6
|
||||
}
|
||||
if err := publishTemplateAssistTask(ctx, s.rabbitMQ, job); err != nil {
|
||||
_ = repo.FailTask(context.Background(), taskID, templateAssistTaskTypeOutline, actor.TenantID, err.Error(), time.Now())
|
||||
s.refundReservedAIPoints(context.Background(), actor.TenantID, actor.UserID, reservation, err.Error())
|
||||
return nil, response.ErrServiceUnavailable(50314, "outline_queue_unavailable", "outline queue is busy, please retry")
|
||||
}
|
||||
|
||||
@@ -401,7 +422,24 @@ func (s *TemplateService) IsAssistTaskTerminal(ctx context.Context, job AssistJo
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return record.Status == "completed" || record.Status == "failed", nil
|
||||
switch record.Status {
|
||||
case "completed":
|
||||
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, "")
|
||||
return true, nil
|
||||
case "failed":
|
||||
if isTemplateAssistAIPointTask(job.TaskType) {
|
||||
resourceType := "template_assist_task"
|
||||
resourceUID := job.TaskID
|
||||
errorMessage := "template assist task failed"
|
||||
if record.ErrorMessage != nil && strings.TrimSpace(*record.ErrorMessage) != "" {
|
||||
errorMessage = strings.TrimSpace(*record.ErrorMessage)
|
||||
}
|
||||
_ = RefundPendingAIPointsByResource(ctx, s.pool, s.cache, job.TenantID, resourceType, nil, &resourceUID, errorMessage)
|
||||
}
|
||||
return true, nil
|
||||
default:
|
||||
return false, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TemplateService) executeAnalyzeTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) {
|
||||
@@ -433,7 +471,9 @@ func (s *TemplateService) executeAnalyzeTask(ctx context.Context, repo repositor
|
||||
|
||||
if err := repo.CompleteTask(ctx, job.TaskID, job.TaskType, job.TenantID, resultJSON, time.Now()); err != nil {
|
||||
s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err)
|
||||
return
|
||||
}
|
||||
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model)
|
||||
}
|
||||
|
||||
func (s *TemplateService) executeTitleTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) {
|
||||
@@ -470,7 +510,9 @@ func (s *TemplateService) executeTitleTask(ctx context.Context, repo repository.
|
||||
|
||||
if err := repo.CompleteTask(ctx, job.TaskID, job.TaskType, job.TenantID, resultJSON, time.Now()); err != nil {
|
||||
s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err)
|
||||
return
|
||||
}
|
||||
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model)
|
||||
}
|
||||
|
||||
func (s *TemplateService) executeOutlineTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) {
|
||||
@@ -512,12 +554,75 @@ func (s *TemplateService) executeOutlineTask(ctx context.Context, repo repositor
|
||||
|
||||
if err := repo.CompleteTask(ctx, job.TaskID, job.TaskType, job.TenantID, resultJSON, time.Now()); err != nil {
|
||||
s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err)
|
||||
return
|
||||
}
|
||||
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model)
|
||||
}
|
||||
|
||||
func (s *TemplateService) failAssist(ctx context.Context, taskID, taskType string, tenantID int64, err error) {
|
||||
repo := repository.NewTemplateAssistTaskRepository(s.pool)
|
||||
_ = repo.FailTask(ctx, taskID, taskType, tenantID, err.Error(), time.Now())
|
||||
if isTemplateAssistAIPointTask(taskType) {
|
||||
resourceType := "template_assist_task"
|
||||
resourceUID := taskID
|
||||
_ = RefundPendingAIPointsByResource(ctx, s.pool, s.cache, tenantID, resourceType, nil, &resourceUID, err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TemplateService) reserveTemplateAssistAIPoints(
|
||||
ctx context.Context,
|
||||
tenantID int64,
|
||||
operatorID int64,
|
||||
templateID int64,
|
||||
taskID string,
|
||||
taskType string,
|
||||
meteredText string,
|
||||
) (*AIPointReservation, error) {
|
||||
usageType := AIUsageTypeTemplateTitleGenerate
|
||||
fixedPoints := 0
|
||||
switch taskType {
|
||||
case templateAssistTaskTypeAnalyze:
|
||||
usageType = AIUsageTypeTemplateAnalyze
|
||||
fixedPoints = 1
|
||||
case templateAssistTaskTypeOutline:
|
||||
usageType = AIUsageTypeTemplateOutlineGenerate
|
||||
}
|
||||
resourceType := "template_assist_task"
|
||||
resourceUID := taskID
|
||||
return ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{
|
||||
TenantID: tenantID,
|
||||
OperatorID: operatorID,
|
||||
UsageType: usageType,
|
||||
ResourceType: &resourceType,
|
||||
ResourceID: &templateID,
|
||||
ResourceUID: &resourceUID,
|
||||
MeteredText: meteredText,
|
||||
FixedPoints: fixedPoints,
|
||||
Metadata: map[string]any{
|
||||
"template_id": templateID,
|
||||
"task_id": taskID,
|
||||
"task_type": taskType,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *TemplateService) refundReservedAIPoints(ctx context.Context, tenantID, operatorID int64, reservation *AIPointReservation, errorMessage string) {
|
||||
if reservation == nil {
|
||||
return
|
||||
}
|
||||
_ = RefundAIPoints(ctx, s.pool, s.cache, tenantID, operatorID, *reservation, errorMessage)
|
||||
}
|
||||
|
||||
func (s *TemplateService) completeTemplateAssistAIPoints(ctx context.Context, tenantID int64, taskID string, model string) {
|
||||
resourceType := "template_assist_task"
|
||||
resourceUID := taskID
|
||||
_ = CompletePendingAIPointsByResource(ctx, s.pool, tenantID, resourceType, nil, &resourceUID, model)
|
||||
}
|
||||
|
||||
func isTemplateAssistAIPointTask(taskType string) bool {
|
||||
return taskType == templateAssistTaskTypeAnalyze ||
|
||||
taskType == templateAssistTaskTypeTitle ||
|
||||
taskType == templateAssistTaskTypeOutline
|
||||
}
|
||||
|
||||
func normalizeAnalyzeTaskRequest(req AnalyzeTaskRequest) AnalyzeTaskRequest {
|
||||
@@ -661,16 +766,7 @@ func hasOutlineContext(req OutlineTaskRequest) bool {
|
||||
}
|
||||
|
||||
func buildAnalyzePrompt(templateKey, templateName string, req AnalyzeTaskRequest, analyzePromptTemplate *string) string {
|
||||
contextPayload := map[string]interface{}{
|
||||
"template_key": templateKey,
|
||||
"template_name": templateName,
|
||||
"locale": req.Locale,
|
||||
"brand_name": req.BrandName,
|
||||
"official_website": req.Website,
|
||||
"input_params": req.InputParams,
|
||||
"existing_keywords": req.ExistingKeywords,
|
||||
"existing_competitors": req.ExistingCompetitors,
|
||||
}
|
||||
contextPayload := analyzePromptParams(templateKey, templateName, req)
|
||||
|
||||
if analyzePromptTemplate != nil && strings.TrimSpace(*analyzePromptTemplate) != "" {
|
||||
basePrompt := strings.TrimSpace(renderPromptTemplate(analyzePromptTemplate, contextPayload))
|
||||
@@ -684,6 +780,11 @@ func buildAnalyzePrompt(templateKey, templateName string, req AnalyzeTaskRequest
|
||||
return prompts.AnalyzeFallbackPrompt(string(contextJSON))
|
||||
}
|
||||
|
||||
func buildAnalyzeAIPointMeteredText(req AnalyzeTaskRequest) string {
|
||||
payload, _ := json.Marshal(analyzePromptParams("", "", req))
|
||||
return string(payload)
|
||||
}
|
||||
|
||||
func buildTitlePrompt(templateKey, templateName string, req TitleTaskRequest, titlePromptTemplate *string) string {
|
||||
contextPayload := titlePromptParams(templateKey, templateName, req)
|
||||
if titlePromptTemplate != nil && strings.TrimSpace(*titlePromptTemplate) != "" {
|
||||
@@ -723,6 +824,29 @@ func buildOutlinePrompt(templateKey, templateName string, req OutlineTaskRequest
|
||||
return prompts.OutlineFallbackPrompt(string(contextJSON))
|
||||
}
|
||||
|
||||
func buildTitleAIPointMeteredText(req TitleTaskRequest) string {
|
||||
payload, _ := json.Marshal(titlePromptParams("", "", req))
|
||||
return string(payload)
|
||||
}
|
||||
|
||||
func buildOutlineAIPointMeteredText(req OutlineTaskRequest) string {
|
||||
payload, _ := json.Marshal(outlinePromptParams("", "", req))
|
||||
return string(payload)
|
||||
}
|
||||
|
||||
func analyzePromptParams(templateKey, templateName string, req AnalyzeTaskRequest) map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"template_key": templateKey,
|
||||
"template_name": templateName,
|
||||
"locale": req.Locale,
|
||||
"brand_name": req.BrandName,
|
||||
"official_website": req.Website,
|
||||
"input_params": req.InputParams,
|
||||
"existing_keywords": req.ExistingKeywords,
|
||||
"existing_competitors": req.ExistingCompetitors,
|
||||
}
|
||||
}
|
||||
|
||||
func titlePromptParams(templateKey, templateName string, req TitleTaskRequest) map[string]interface{} {
|
||||
params := map[string]interface{}{
|
||||
"template_key": templateKey,
|
||||
|
||||
@@ -20,6 +20,7 @@ type WorkspaceService struct {
|
||||
repo repository.WorkspaceRepository
|
||||
templates repository.TemplateRepository
|
||||
quota repository.QuotaRepository
|
||||
aiPointUsage repository.AIPointUsageRepository
|
||||
supportedPlatformCount int
|
||||
cache sharedcache.Cache
|
||||
cacheGroup singleflight.Group
|
||||
@@ -29,11 +30,13 @@ func NewWorkspaceService(
|
||||
repo repository.WorkspaceRepository,
|
||||
templates repository.TemplateRepository,
|
||||
quota repository.QuotaRepository,
|
||||
aiPointUsage repository.AIPointUsageRepository,
|
||||
) *WorkspaceService {
|
||||
return &WorkspaceService{
|
||||
repo: repo,
|
||||
templates: templates,
|
||||
quota: quota,
|
||||
aiPointUsage: aiPointUsage,
|
||||
supportedPlatformCount: 5,
|
||||
}
|
||||
}
|
||||
@@ -105,18 +108,80 @@ func (s *WorkspaceService) QuotaSummary(ctx context.Context) (*domain.QuotaSumma
|
||||
}
|
||||
return nil, response.ErrInternal(50010, "query_failed", "failed to get quota balance")
|
||||
}
|
||||
aiStatus, err := s.quota.GetAIQuotaStatus(loadCtx, actor.TenantID, time.Now().UTC())
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
aiStatus = &repository.AIQuotaStatus{}
|
||||
} else {
|
||||
return nil, response.ErrInternal(50010, "query_failed", "failed to get ai points balance")
|
||||
}
|
||||
}
|
||||
|
||||
return &domain.QuotaSummary{
|
||||
PlanCode: status.PlanCode,
|
||||
PlanName: status.PlanName,
|
||||
TotalQuota: status.Total,
|
||||
UsedQuota: status.Used,
|
||||
Balance: status.Balance,
|
||||
ResetAt: status.ResetAt,
|
||||
PlanCode: status.PlanCode,
|
||||
PlanName: status.PlanName,
|
||||
TotalQuota: status.Total,
|
||||
UsedQuota: status.Used,
|
||||
Balance: status.Balance,
|
||||
ResetAt: status.ResetAt,
|
||||
AIPointsTotal: aiStatus.Total,
|
||||
AIPointsUsed: aiStatus.Used,
|
||||
AIPointsBalance: aiStatus.Balance,
|
||||
AIPointBaseChars: aiStatus.BaseChars,
|
||||
AIPointsResetAt: aiStatus.ResetAt,
|
||||
}, nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *WorkspaceService) AIPointUsageLogs(ctx context.Context, page, pageSize int) (*domain.AIPointUsageListResponse, error) {
|
||||
actor := auth.MustActor(ctx)
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
if pageSize < 1 || pageSize > 100 {
|
||||
pageSize = 20
|
||||
}
|
||||
offset := (page - 1) * pageSize
|
||||
|
||||
rows, err := s.aiPointUsage.List(ctx, repository.AIPointUsageListParams{
|
||||
TenantID: actor.TenantID,
|
||||
Limit: pageSize,
|
||||
Offset: offset,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50010, "query_failed", "failed to list ai point usage")
|
||||
}
|
||||
total, err := s.aiPointUsage.Count(ctx, actor.TenantID)
|
||||
if err != nil {
|
||||
return nil, response.ErrInternal(50010, "query_failed", "failed to count ai point usage")
|
||||
}
|
||||
|
||||
items := make([]domain.AIPointUsageLog, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
items = append(items, domain.AIPointUsageLog{
|
||||
ID: row.ID,
|
||||
UsageType: row.UsageType,
|
||||
ResourceType: row.ResourceType,
|
||||
ResourceID: row.ResourceID,
|
||||
ResourceUID: row.ResourceUID,
|
||||
RequestChars: row.RequestChars,
|
||||
BaseChars: row.BaseChars,
|
||||
Points: row.Points,
|
||||
Status: row.Status,
|
||||
Model: row.Model,
|
||||
ErrorMessage: row.ErrorMessage,
|
||||
CompletedAt: row.CompletedAt,
|
||||
CreatedAt: row.CreatedAt,
|
||||
})
|
||||
}
|
||||
return &domain.AIPointUsageListResponse{
|
||||
Items: items,
|
||||
Total: total,
|
||||
Page: page,
|
||||
PageSize: pageSize,
|
||||
}, nil
|
||||
}
|
||||
|
||||
type TemplateCard struct {
|
||||
ID int64 `json:"id"`
|
||||
Scope string `json:"scope"`
|
||||
|
||||
@@ -23,12 +23,17 @@ type RecentArticle struct {
|
||||
}
|
||||
|
||||
type QuotaSummary struct {
|
||||
PlanCode string `json:"plan_code"`
|
||||
PlanName string `json:"plan_name"`
|
||||
TotalQuota int `json:"total_quota"`
|
||||
UsedQuota int `json:"used_quota"`
|
||||
Balance int `json:"balance"`
|
||||
ResetAt *time.Time `json:"reset_at"`
|
||||
PlanCode string `json:"plan_code"`
|
||||
PlanName string `json:"plan_name"`
|
||||
TotalQuota int `json:"total_quota"`
|
||||
UsedQuota int `json:"used_quota"`
|
||||
Balance int `json:"balance"`
|
||||
ResetAt *time.Time `json:"reset_at"`
|
||||
AIPointsTotal int `json:"ai_points_total"`
|
||||
AIPointsUsed int `json:"ai_points_used"`
|
||||
AIPointsBalance int `json:"ai_points_balance"`
|
||||
AIPointBaseChars int `json:"ai_point_base_chars"`
|
||||
AIPointsResetAt *time.Time `json:"ai_points_reset_at"`
|
||||
}
|
||||
|
||||
type KolWorkspaceCard struct {
|
||||
@@ -41,3 +46,26 @@ type KolWorkspaceCard struct {
|
||||
PackageCover *string `json:"package_cover"`
|
||||
KolDisplayName string `json:"kol_display_name"`
|
||||
}
|
||||
|
||||
type AIPointUsageLog struct {
|
||||
ID int64 `json:"id"`
|
||||
UsageType string `json:"usage_type"`
|
||||
ResourceType *string `json:"resource_type"`
|
||||
ResourceID *int64 `json:"resource_id"`
|
||||
ResourceUID *string `json:"resource_uid"`
|
||||
RequestChars int `json:"request_chars"`
|
||||
BaseChars int `json:"base_chars"`
|
||||
Points int `json:"points"`
|
||||
Status string `json:"status"`
|
||||
Model *string `json:"model"`
|
||||
ErrorMessage *string `json:"error_message"`
|
||||
CompletedAt *time.Time `json:"completed_at"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type AIPointUsageListResponse struct {
|
||||
Items []AIPointUsageLog `json:"items"`
|
||||
Total int64 `json:"total"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"page_size"`
|
||||
}
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/tenant/repository/generated"
|
||||
)
|
||||
|
||||
type AIPointUsageLog struct {
|
||||
ID int64
|
||||
TenantID int64
|
||||
OperatorID *int64
|
||||
QuotaReservationID *int64
|
||||
UsageType string
|
||||
ResourceType *string
|
||||
ResourceID *int64
|
||||
ResourceUID *string
|
||||
RequestChars int
|
||||
BaseChars int
|
||||
Points int
|
||||
Status string
|
||||
Model *string
|
||||
MetadataJSON []byte
|
||||
ErrorMessage *string
|
||||
CompletedAt *time.Time
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type CreateAIPointUsageLogInput struct {
|
||||
TenantID int64
|
||||
OperatorID int64
|
||||
QuotaReservationID int64
|
||||
UsageType string
|
||||
ResourceType *string
|
||||
ResourceID *int64
|
||||
ResourceUID *string
|
||||
RequestChars int
|
||||
BaseChars int
|
||||
Points int
|
||||
MetadataJSON []byte
|
||||
}
|
||||
|
||||
type AIPointUsageListParams struct {
|
||||
TenantID int64
|
||||
Limit int
|
||||
Offset int
|
||||
}
|
||||
|
||||
type AIPointUsageRepository interface {
|
||||
Create(ctx context.Context, input CreateAIPointUsageLogInput) (int64, error)
|
||||
MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error
|
||||
MarkRefunded(ctx context.Context, tenantID, id int64, errorMessage string) error
|
||||
FindPendingByResource(ctx context.Context, tenantID int64, resourceType string, resourceID *int64, resourceUID *string) (*AIPointUsageLog, error)
|
||||
List(ctx context.Context, params AIPointUsageListParams) ([]AIPointUsageLog, error)
|
||||
Count(ctx context.Context, tenantID int64) (int64, error)
|
||||
}
|
||||
|
||||
type aiPointUsageRepository struct {
|
||||
db generated.DBTX
|
||||
}
|
||||
|
||||
func NewAIPointUsageRepository(db generated.DBTX) AIPointUsageRepository {
|
||||
return &aiPointUsageRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *aiPointUsageRepository) Create(ctx context.Context, input CreateAIPointUsageLogInput) (int64, error) {
|
||||
metadataJSON := input.MetadataJSON
|
||||
if len(metadataJSON) == 0 {
|
||||
metadataJSON = []byte("{}")
|
||||
}
|
||||
var operatorID *int64
|
||||
if input.OperatorID > 0 {
|
||||
operatorID = &input.OperatorID
|
||||
}
|
||||
var reservationID *int64
|
||||
if input.QuotaReservationID > 0 {
|
||||
reservationID = &input.QuotaReservationID
|
||||
}
|
||||
|
||||
var id int64
|
||||
err := r.db.QueryRow(ctx, `
|
||||
INSERT INTO ai_point_usage_logs (
|
||||
tenant_id, operator_id, quota_reservation_id, usage_type,
|
||||
resource_type, resource_id, resource_uid,
|
||||
request_chars, base_chars, points, metadata_json
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::jsonb)
|
||||
RETURNING id
|
||||
`, input.TenantID, operatorID, reservationID, input.UsageType,
|
||||
input.ResourceType, input.ResourceID, input.ResourceUID,
|
||||
input.RequestChars, input.BaseChars, input.Points, metadataJSON,
|
||||
).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
func (r *aiPointUsageRepository) MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error {
|
||||
_, err := r.db.Exec(ctx, `
|
||||
UPDATE ai_point_usage_logs
|
||||
SET status = 'completed',
|
||||
model = $1,
|
||||
error_message = NULL,
|
||||
completed_at = COALESCE(completed_at, NOW()),
|
||||
updated_at = NOW()
|
||||
WHERE id = $2
|
||||
AND tenant_id = $3
|
||||
AND status = 'pending'
|
||||
`, model, id, tenantID)
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *aiPointUsageRepository) MarkRefunded(ctx context.Context, tenantID, id int64, errorMessage string) error {
|
||||
_, err := r.db.Exec(ctx, `
|
||||
UPDATE ai_point_usage_logs
|
||||
SET status = 'refunded',
|
||||
error_message = $1,
|
||||
completed_at = COALESCE(completed_at, NOW()),
|
||||
updated_at = NOW()
|
||||
WHERE id = $2
|
||||
AND tenant_id = $3
|
||||
AND status = 'pending'
|
||||
`, errorMessage, id, tenantID)
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *aiPointUsageRepository) FindPendingByResource(
|
||||
ctx context.Context,
|
||||
tenantID int64,
|
||||
resourceType string,
|
||||
resourceID *int64,
|
||||
resourceUID *string,
|
||||
) (*AIPointUsageLog, error) {
|
||||
row := r.db.QueryRow(ctx, `
|
||||
SELECT id, tenant_id, operator_id, quota_reservation_id, usage_type,
|
||||
resource_type, resource_id, resource_uid, request_chars, base_chars,
|
||||
points, status, model, metadata_json, error_message, completed_at,
|
||||
created_at, updated_at
|
||||
FROM ai_point_usage_logs
|
||||
WHERE tenant_id = $1
|
||||
AND resource_type = $2
|
||||
AND ($3::bigint IS NULL OR resource_id = $3)
|
||||
AND ($4::varchar IS NULL OR resource_uid = $4)
|
||||
AND status = 'pending'
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT 1
|
||||
`, tenantID, resourceType, resourceID, resourceUID)
|
||||
|
||||
var item AIPointUsageLog
|
||||
err := scanAIPointUsageLog(row, &item)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (r *aiPointUsageRepository) List(ctx context.Context, params AIPointUsageListParams) ([]AIPointUsageLog, error) {
|
||||
limit := params.Limit
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 50
|
||||
}
|
||||
offset := params.Offset
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
|
||||
rows, err := r.db.Query(ctx, `
|
||||
SELECT id, tenant_id, operator_id, quota_reservation_id, usage_type,
|
||||
resource_type, resource_id, resource_uid, request_chars, base_chars,
|
||||
points, status, model, metadata_json, error_message, completed_at,
|
||||
created_at, updated_at
|
||||
FROM ai_point_usage_logs
|
||||
WHERE tenant_id = $1
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT $2 OFFSET $3
|
||||
`, params.TenantID, limit, offset)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]AIPointUsageLog, 0, limit)
|
||||
for rows.Next() {
|
||||
var item AIPointUsageLog
|
||||
if err := scanAIPointUsageLog(rows, &item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (r *aiPointUsageRepository) Count(ctx context.Context, tenantID int64) (int64, error) {
|
||||
var total int64
|
||||
err := r.db.QueryRow(ctx, `
|
||||
SELECT COUNT(*)
|
||||
FROM ai_point_usage_logs
|
||||
WHERE tenant_id = $1
|
||||
`, tenantID).Scan(&total)
|
||||
return total, err
|
||||
}
|
||||
|
||||
type aiPointUsageScanner interface {
|
||||
Scan(dest ...any) error
|
||||
}
|
||||
|
||||
func scanAIPointUsageLog(scanner aiPointUsageScanner, item *AIPointUsageLog) error {
|
||||
var operatorID pgtype.Int8
|
||||
var quotaReservationID pgtype.Int8
|
||||
var resourceType pgtype.Text
|
||||
var resourceID pgtype.Int8
|
||||
var resourceUID pgtype.Text
|
||||
var model pgtype.Text
|
||||
var errorMessage pgtype.Text
|
||||
var completedAt pgtype.Timestamptz
|
||||
|
||||
if err := scanner.Scan(
|
||||
&item.ID,
|
||||
&item.TenantID,
|
||||
&operatorID,
|
||||
"aReservationID,
|
||||
&item.UsageType,
|
||||
&resourceType,
|
||||
&resourceID,
|
||||
&resourceUID,
|
||||
&item.RequestChars,
|
||||
&item.BaseChars,
|
||||
&item.Points,
|
||||
&item.Status,
|
||||
&model,
|
||||
&item.MetadataJSON,
|
||||
&errorMessage,
|
||||
&completedAt,
|
||||
&item.CreatedAt,
|
||||
&item.UpdatedAt,
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
item.OperatorID = nullableInt64(operatorID)
|
||||
item.QuotaReservationID = nullableInt64(quotaReservationID)
|
||||
item.ResourceType = nullableText(resourceType)
|
||||
item.ResourceID = nullableInt64(resourceID)
|
||||
item.ResourceUID = nullableText(resourceUID)
|
||||
item.Model = nullableText(model)
|
||||
item.ErrorMessage = nullableText(errorMessage)
|
||||
item.CompletedAt = optionalTime(completedAt)
|
||||
return nil
|
||||
}
|
||||
@@ -19,6 +19,16 @@ type ArticleQuotaStatus struct {
|
||||
ResetAt *time.Time
|
||||
}
|
||||
|
||||
type AIQuotaStatus struct {
|
||||
PlanCode string
|
||||
PlanName string
|
||||
Total int
|
||||
Used int
|
||||
Balance int
|
||||
BaseChars int
|
||||
ResetAt *time.Time
|
||||
}
|
||||
|
||||
type QuotaLedgerInput struct {
|
||||
TenantID int64
|
||||
OperatorID int64
|
||||
@@ -43,6 +53,7 @@ type QuotaReservationInput struct {
|
||||
type QuotaRepository interface {
|
||||
GetCurrentBalance(ctx context.Context, tenantID int64, quotaType string) (int, error)
|
||||
GetArticleQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*ArticleQuotaStatus, error)
|
||||
GetAIQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*AIQuotaStatus, error)
|
||||
InsertQuotaLedger(ctx context.Context, input QuotaLedgerInput) (int64, error)
|
||||
CreateQuotaReservation(ctx context.Context, input QuotaReservationInput) (int64, error)
|
||||
UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error
|
||||
@@ -78,6 +89,58 @@ func (r *quotaRepository) GetCurrentBalance(ctx context.Context, tenantID int64,
|
||||
return status.Balance, nil
|
||||
}
|
||||
|
||||
func (r *quotaRepository) GetAIQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*AIQuotaStatus, error) {
|
||||
access, err := loadTenantPlanAccess(ctx, r.db, tenantID, now)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if access == nil {
|
||||
return nil, pgx.ErrNoRows
|
||||
}
|
||||
|
||||
baseChars := access.AIPointBaseChars()
|
||||
if !access.HasActiveAccess(now) {
|
||||
return &AIQuotaStatus{
|
||||
PlanCode: access.PlanCode,
|
||||
PlanName: access.PlanName,
|
||||
BaseChars: baseChars,
|
||||
}, nil
|
||||
}
|
||||
|
||||
total := access.AIPointsMonthlyLimit()
|
||||
if total <= 0 {
|
||||
return &AIQuotaStatus{
|
||||
PlanCode: access.PlanCode,
|
||||
PlanName: access.PlanName,
|
||||
BaseChars: baseChars,
|
||||
}, nil
|
||||
}
|
||||
|
||||
windowStart, windowEnd, resetAt := access.AIPointsWindow(now)
|
||||
if windowEnd.IsZero() {
|
||||
return nil, errors.New("ai points quota window end is missing")
|
||||
}
|
||||
|
||||
used, err := CountEffectiveQuotaReservations(ctx, r.db, tenantID, "ai_points", windowStart, windowEnd)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
balance := total - used
|
||||
if balance < 0 {
|
||||
balance = 0
|
||||
}
|
||||
|
||||
return &AIQuotaStatus{
|
||||
PlanCode: access.PlanCode,
|
||||
PlanName: access.PlanName,
|
||||
Total: total,
|
||||
Used: used,
|
||||
Balance: balance,
|
||||
BaseChars: baseChars,
|
||||
ResetAt: resetAt,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *quotaRepository) GetArticleQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*ArticleQuotaStatus, error) {
|
||||
access, err := loadTenantPlanAccess(ctx, r.db, tenantID, now)
|
||||
if err != nil {
|
||||
|
||||
@@ -48,7 +48,7 @@ func NewArticleHandler(a *bootstrap.App) *ArticleHandler {
|
||||
a.DB,
|
||||
a.LLM,
|
||||
a.Config.LLM.MaxOutputTokens,
|
||||
),
|
||||
).WithCache(a.Cache),
|
||||
promptGenerateSvc: app.NewPromptRuleGenerationService(
|
||||
a.DB,
|
||||
a.LLM,
|
||||
@@ -344,7 +344,7 @@ func (h *ArticleHandler) OptimizeSelectionStream(c *gin.Context) {
|
||||
var streamed strings.Builder
|
||||
var streamWriteErr error
|
||||
|
||||
result, err := h.selectionOptimize.OptimizeSelection(streamCtx, req, func(delta string) {
|
||||
result, err := h.selectionOptimize.OptimizeSelection(streamCtx, id, req, func(delta string) {
|
||||
if delta == "" || streamWriteErr != nil {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -42,7 +42,7 @@ func NewKolManageHandler(a *bootstrap.App, assets *AssetHandler) *KolManageHandl
|
||||
profileSvc: profileSvc,
|
||||
pkgSvc: app.NewKolPackageService(profileSvc, a.KolPackages, a.KolPrompts).WithCache(a.Cache),
|
||||
promptSvc: app.NewKolPromptService(a.DB, profileSvc, a.KolPackages, a.KolPrompts, a.KolSubscriptions, a.KolPromptAsset).WithCache(a.Cache),
|
||||
assistSvc: app.NewKolAssistService(profileSvc, a.KolAssists, a.LLM, a.RabbitMQ, a.Config.LLM, a.Logger),
|
||||
assistSvc: app.NewKolAssistService(a.DB, profileSvc, a.KolAssists, a.LLM, a.RabbitMQ, a.Config.LLM, a.Logger).WithCache(a.Cache),
|
||||
assets: assets,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -80,6 +80,7 @@ func RegisterRoutes(app *bootstrap.App) {
|
||||
workspace.GET("/overview", wsHandler.Overview)
|
||||
workspace.GET("/recent-articles", wsHandler.RecentArticles)
|
||||
workspace.GET("/quota-summary", wsHandler.QuotaSummary)
|
||||
workspace.GET("/ai-point-usage", wsHandler.AIPointUsageLogs)
|
||||
workspace.GET("/template-cards", wsHandler.TemplateCards)
|
||||
workspace.GET("/kol-cards", wsHandler.KolCards)
|
||||
|
||||
|
||||
@@ -22,6 +22,7 @@ func NewWorkspaceHandler(a *bootstrap.App) *WorkspaceHandler {
|
||||
a.Cache,
|
||||
),
|
||||
repository.NewQuotaRepository(a.DB),
|
||||
repository.NewAIPointUsageRepository(a.DB),
|
||||
).WithCache(a.Cache),
|
||||
}
|
||||
}
|
||||
@@ -53,6 +54,30 @@ func (h *WorkspaceHandler) QuotaSummary(c *gin.Context) {
|
||||
response.Success(c, data)
|
||||
}
|
||||
|
||||
func (h *WorkspaceHandler) AIPointUsageLogs(c *gin.Context) {
|
||||
page := parseIntQuery(c.Query("page"), 1)
|
||||
pageSize := parseIntQuery(c.Query("page_size"), 20)
|
||||
if c.Query("page") == "" && c.Query("page_size") == "" && c.Query("limit") != "" {
|
||||
limit := parseIntQuery(c.Query("limit"), 20)
|
||||
offset := parseIntQuery(c.Query("offset"), 0)
|
||||
if limit > 0 {
|
||||
pageSize = limit
|
||||
page = offset/limit + 1
|
||||
}
|
||||
}
|
||||
|
||||
data, err := h.svc.AIPointUsageLogs(
|
||||
c.Request.Context(),
|
||||
page,
|
||||
pageSize,
|
||||
)
|
||||
if err != nil {
|
||||
response.Error(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, data)
|
||||
}
|
||||
|
||||
func (h *WorkspaceHandler) TemplateCards(c *gin.Context) {
|
||||
data, err := h.svc.TemplateCards(c.Request.Context())
|
||||
if err != nil {
|
||||
|
||||
@@ -8,9 +8,11 @@ import (
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
amqp "github.com/rabbitmq/amqp091-go"
|
||||
"go.uber.org/zap"
|
||||
|
||||
sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/llm"
|
||||
"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
|
||||
tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
|
||||
@@ -20,10 +22,12 @@ import (
|
||||
var errKolAssistConsumerClosed = errors.New("kol assist consumer channel closed")
|
||||
|
||||
type KolAssistWorker struct {
|
||||
pool *pgxpool.Pool
|
||||
rabbitMQ *rabbitmq.Client
|
||||
repo repository.KolAssistRepository
|
||||
llm llm.Client
|
||||
logger *zap.Logger
|
||||
cache sharedcache.Cache
|
||||
retryInterval time.Duration
|
||||
processTimeout time.Duration
|
||||
consumerPrefix string
|
||||
@@ -31,10 +35,12 @@ type KolAssistWorker struct {
|
||||
}
|
||||
|
||||
func NewKolAssistWorker(
|
||||
pool *pgxpool.Pool,
|
||||
rabbitMQClient *rabbitmq.Client,
|
||||
repo repository.KolAssistRepository,
|
||||
llmClient llm.Client,
|
||||
logger *zap.Logger,
|
||||
cache sharedcache.Cache,
|
||||
workerConcurrency int,
|
||||
processTimeout time.Duration,
|
||||
) *KolAssistWorker {
|
||||
@@ -46,10 +52,12 @@ func NewKolAssistWorker(
|
||||
}
|
||||
|
||||
return &KolAssistWorker{
|
||||
pool: pool,
|
||||
rabbitMQ: rabbitMQClient,
|
||||
repo: repo,
|
||||
llm: llmClient,
|
||||
logger: logger,
|
||||
cache: cache,
|
||||
retryInterval: 5 * time.Second,
|
||||
processTimeout: processTimeout,
|
||||
consumerPrefix: fmt.Sprintf("kol-assist-%d", os.Getpid()),
|
||||
@@ -73,7 +81,19 @@ func (w *KolAssistWorker) Process(ctx context.Context, job tenantapp.KolAssistJo
|
||||
if err != nil {
|
||||
return fmt.Errorf("load kol assist task: %w", err)
|
||||
}
|
||||
if task == nil || task.Status == "completed" || task.Status == "failed" {
|
||||
if task == nil {
|
||||
return nil
|
||||
}
|
||||
switch task.Status {
|
||||
case "completed":
|
||||
w.completeAIPoints(ctx, job, "")
|
||||
return nil
|
||||
case "failed":
|
||||
errorMessage := "kol assist task failed"
|
||||
if task.ErrorMessage != nil && *task.ErrorMessage != "" {
|
||||
errorMessage = *task.ErrorMessage
|
||||
}
|
||||
w.refundAIPoints(ctx, job, errorMessage)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -83,6 +103,7 @@ func (w *KolAssistWorker) Process(ctx context.Context, job tenantapp.KolAssistJo
|
||||
|
||||
if err := w.processTask(ctx, job, task); err != nil {
|
||||
_ = w.repo.MarkFailed(context.Background(), job.TenantID, job.TaskID, err.Error())
|
||||
w.refundAIPoints(context.Background(), job, err.Error())
|
||||
if w.logger != nil {
|
||||
w.logger.Warn("kol assist task failed",
|
||||
zap.Error(err),
|
||||
@@ -215,5 +236,30 @@ func (w *KolAssistWorker) processTask(ctx context.Context, job tenantapp.KolAssi
|
||||
if err := w.repo.MarkCompleted(ctx, job.TenantID, job.TaskID, resultJSON); err != nil {
|
||||
return fmt.Errorf("mark kol assist completed: %w", err)
|
||||
}
|
||||
w.completeAIPoints(ctx, job, result.Model)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *KolAssistWorker) completeAIPoints(ctx context.Context, job tenantapp.KolAssistJob, model string) {
|
||||
resourceType := "kol_assist_task"
|
||||
resourceUID := job.TaskID
|
||||
if err := tenantapp.CompletePendingAIPointsByResource(ctx, w.pool, job.TenantID, resourceType, nil, &resourceUID, model); err != nil && w.logger != nil {
|
||||
w.logger.Warn("kol assist ai points confirmation failed",
|
||||
zap.Error(err),
|
||||
zap.String("task_id", job.TaskID),
|
||||
zap.Int64("tenant_id", job.TenantID),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *KolAssistWorker) refundAIPoints(ctx context.Context, job tenantapp.KolAssistJob, errorMessage string) {
|
||||
resourceType := "kol_assist_task"
|
||||
resourceUID := job.TaskID
|
||||
if err := tenantapp.RefundPendingAIPointsByResource(ctx, w.pool, w.cache, job.TenantID, resourceType, nil, &resourceUID, errorMessage); err != nil && w.logger != nil {
|
||||
w.logger.Warn("kol assist ai points refund failed",
|
||||
zap.Error(err),
|
||||
zap.String("task_id", job.TaskID),
|
||||
zap.Int64("tenant_id", job.TenantID),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
DROP INDEX IF EXISTS idx_quota_reservations_tenant_type_created_status;
|
||||
|
||||
DROP TABLE IF EXISTS ai_point_usage_logs;
|
||||
@@ -0,0 +1,37 @@
|
||||
CREATE TABLE ai_point_usage_logs (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
tenant_id BIGINT NOT NULL REFERENCES tenants(id),
|
||||
operator_id BIGINT REFERENCES users(id),
|
||||
quota_reservation_id BIGINT REFERENCES quota_reservations(id),
|
||||
usage_type VARCHAR(64) NOT NULL,
|
||||
resource_type VARCHAR(64),
|
||||
resource_id BIGINT,
|
||||
resource_uid VARCHAR(128),
|
||||
request_chars INT NOT NULL DEFAULT 0,
|
||||
base_chars INT NOT NULL DEFAULT 1000,
|
||||
points INT NOT NULL,
|
||||
status VARCHAR(20) NOT NULL DEFAULT 'pending',
|
||||
model VARCHAR(128),
|
||||
metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb,
|
||||
error_message TEXT,
|
||||
completed_at TIMESTAMPTZ,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
CONSTRAINT chk_ai_point_usage_status
|
||||
CHECK (status IN ('pending', 'completed', 'refunded', 'failed'))
|
||||
);
|
||||
|
||||
CREATE INDEX idx_ai_point_usage_tenant_created
|
||||
ON ai_point_usage_logs(tenant_id, created_at DESC);
|
||||
|
||||
CREATE INDEX idx_ai_point_usage_tenant_operator_created
|
||||
ON ai_point_usage_logs(tenant_id, operator_id, created_at DESC);
|
||||
|
||||
CREATE INDEX idx_ai_point_usage_reservation
|
||||
ON ai_point_usage_logs(quota_reservation_id);
|
||||
|
||||
CREATE INDEX idx_ai_point_usage_resource
|
||||
ON ai_point_usage_logs(tenant_id, resource_type, resource_id, resource_uid);
|
||||
|
||||
CREATE INDEX idx_quota_reservations_tenant_type_created_status
|
||||
ON quota_reservations(tenant_id, quota_type, created_at, status);
|
||||
Reference in New Issue
Block a user