root/geo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(server): meter AI point usage across tenant AI features

Browse Source 
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:
liangxu
2026-04-28 11:33:47 +08:00
parent 794a2d89db
commit e307a048d1
18 changed files with 1270 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/internal/tenant/app/template_assist.go
+135 -11
View File
@@ -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,