774 lines
23 KiB
Go
774 lines
23 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
"go.uber.org/zap"
|
|
|
|
"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"
|
|
"github.com/geo-platform/tenant-api/internal/shared/response"
|
|
"github.com/geo-platform/tenant-api/internal/shared/stream"
|
|
"github.com/geo-platform/tenant-api/internal/tenant/repository"
|
|
)
|
|
|
|
type TemplateService struct {
|
|
pool *pgxpool.Pool
|
|
templates repository.TemplateRepository
|
|
llm llm.Client
|
|
rabbitMQ *rabbitmq.Client
|
|
knowledge *KnowledgeService
|
|
streams *stream.GenerationHub
|
|
cfg generationRuntimeConfig
|
|
configProvider runtimeConfigProvider
|
|
cache sharedcache.Cache
|
|
logger *zap.Logger
|
|
}
|
|
|
|
type generationJob struct {
|
|
OperatorID int64
|
|
TenantID int64
|
|
ArticleID int64
|
|
TaskID int64
|
|
ReservationID int64
|
|
TemplateKey string
|
|
TemplateName string
|
|
PromptTemplate *string
|
|
Params map[string]interface{}
|
|
InitialTitle string
|
|
WebSearch llm.WebSearchOptions
|
|
}
|
|
|
|
func NewTemplateService(
|
|
pool *pgxpool.Pool,
|
|
templates repository.TemplateRepository,
|
|
llmClient llm.Client,
|
|
rabbitMQClient *rabbitmq.Client,
|
|
knowledge *KnowledgeService,
|
|
streams *stream.GenerationHub,
|
|
cfg config.GenerationConfig,
|
|
llmMaxOutputTokens int64,
|
|
) *TemplateService {
|
|
svc := &TemplateService{
|
|
pool: pool,
|
|
templates: templates,
|
|
llm: llmClient,
|
|
rabbitMQ: rabbitMQClient,
|
|
knowledge: knowledge,
|
|
streams: streams,
|
|
cfg: generationRuntimeFromConfig(cfg, config.LLMConfig{MaxOutputTokens: llmMaxOutputTokens}),
|
|
}
|
|
|
|
return svc
|
|
}
|
|
|
|
func (s *TemplateService) WithCache(c sharedcache.Cache) *TemplateService {
|
|
s.cache = c
|
|
return s
|
|
}
|
|
|
|
func (s *TemplateService) WithConfigProvider(provider runtimeConfigProvider) *TemplateService {
|
|
s.configProvider = provider
|
|
return s
|
|
}
|
|
|
|
func (s *TemplateService) WithLogger(logger *zap.Logger) *TemplateService {
|
|
s.logger = logger
|
|
return s
|
|
}
|
|
|
|
func (s *TemplateService) runtimeConfig() generationRuntimeConfig {
|
|
if s != nil && s.configProvider != nil {
|
|
if cfg := s.configProvider.Current(); cfg != nil {
|
|
return generationRuntimeFromConfig(cfg.Generation, cfg.LLM)
|
|
}
|
|
}
|
|
if s == nil {
|
|
return generationRuntimeFromConfig(config.GenerationConfig{}, config.LLMConfig{})
|
|
}
|
|
return s.cfg
|
|
}
|
|
|
|
type TemplateListItem struct {
|
|
ID int64 `json:"id"`
|
|
Scope string `json:"scope"`
|
|
OriginType string `json:"origin_type"`
|
|
TemplateKey string `json:"template_key"`
|
|
TemplateName string `json:"template_name"`
|
|
PromptVisibility string `json:"prompt_visibility"`
|
|
CardConfigJSON map[string]interface{} `json:"card_config"`
|
|
Status string `json:"status"`
|
|
VersionNo int `json:"version_no"`
|
|
CreatedAt time.Time `json:"created_at"`
|
|
}
|
|
|
|
func (s *TemplateService) List(ctx context.Context) ([]TemplateListItem, error) {
|
|
actor := auth.MustActor(ctx)
|
|
rows, err := s.templates.ListTemplates(ctx, actor.TenantID)
|
|
if err != nil {
|
|
return nil, response.ErrInternal(50010, "query_failed", "failed to list templates")
|
|
}
|
|
|
|
items := make([]TemplateListItem, 0, len(rows))
|
|
for _, row := range rows {
|
|
item := TemplateListItem{
|
|
ID: row.ID,
|
|
Scope: row.Scope,
|
|
OriginType: row.OriginType,
|
|
TemplateKey: row.TemplateKey,
|
|
TemplateName: row.TemplateName,
|
|
PromptVisibility: effectivePromptVisibility(&row, actor.TenantID),
|
|
CardConfigJSON: map[string]interface{}{},
|
|
Status: row.Status,
|
|
VersionNo: row.VersionNo,
|
|
CreatedAt: row.CreatedAt,
|
|
}
|
|
_ = json.Unmarshal(row.CardConfigJSON, &item.CardConfigJSON)
|
|
items = append(items, item)
|
|
}
|
|
|
|
return items, nil
|
|
}
|
|
|
|
type TemplateDetail struct {
|
|
TemplateListItem
|
|
PromptTemplate *string `json:"prompt_template,omitempty"`
|
|
}
|
|
|
|
func (s *TemplateService) Detail(ctx context.Context, id int64) (*TemplateDetail, error) {
|
|
actor := auth.MustActor(ctx)
|
|
record, err := s.templates.GetTemplateByID(ctx, id, actor.TenantID)
|
|
if err != nil {
|
|
return nil, mapTemplateLookupError(err)
|
|
}
|
|
|
|
detail := &TemplateDetail{
|
|
TemplateListItem: TemplateListItem{
|
|
ID: record.ID,
|
|
Scope: record.Scope,
|
|
OriginType: record.OriginType,
|
|
TemplateKey: record.TemplateKey,
|
|
TemplateName: record.TemplateName,
|
|
PromptVisibility: effectivePromptVisibility(record, actor.TenantID),
|
|
CardConfigJSON: map[string]interface{}{},
|
|
Status: record.Status,
|
|
VersionNo: record.VersionNo,
|
|
CreatedAt: record.CreatedAt,
|
|
},
|
|
}
|
|
_ = json.Unmarshal(record.CardConfigJSON, &detail.CardConfigJSON)
|
|
|
|
if canViewPromptTemplate(record, actor.TenantID) {
|
|
detail.PromptTemplate = record.PromptTemplate
|
|
}
|
|
|
|
return detail, nil
|
|
}
|
|
|
|
func canViewPromptTemplate(record *repository.TemplateRecord, actorTenantID int64) bool {
|
|
if record == nil || record.TenantID == nil {
|
|
return false
|
|
}
|
|
return *record.TenantID == actorTenantID
|
|
}
|
|
|
|
func mapTemplateLookupError(err error) error {
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return response.ErrNotFound(40410, "template_not_found", "template not found")
|
|
}
|
|
return err
|
|
}
|
|
|
|
func effectivePromptVisibility(record *repository.TemplateRecord, actorTenantID int64) string {
|
|
if canViewPromptTemplate(record, actorTenantID) {
|
|
return "visible"
|
|
}
|
|
return "sealed"
|
|
}
|
|
|
|
type GenerateRequest struct {
|
|
ArticleID *int64 `json:"article_id"`
|
|
InputParams map[string]interface{} `json:"input_params" binding:"required"`
|
|
WizardState map[string]interface{} `json:"wizard_state"`
|
|
EnableWebSearch bool `json:"enable_web_search"`
|
|
WebSearchLimit *int32 `json:"web_search_limit"`
|
|
}
|
|
|
|
type GenerateResponse struct {
|
|
ArticleID int64 `json:"article_id"`
|
|
TaskID int64 `json:"task_id"`
|
|
}
|
|
|
|
type SaveDraftRequest struct {
|
|
ArticleID *int64 `json:"article_id"`
|
|
CurrentStep int `json:"current_step"`
|
|
WizardState map[string]interface{} `json:"wizard_state" binding:"required"`
|
|
}
|
|
|
|
type SaveDraftResponse struct {
|
|
ArticleID int64 `json:"article_id"`
|
|
}
|
|
|
|
func (s *TemplateService) SaveDraft(ctx context.Context, templateID int64, req SaveDraftRequest) (*SaveDraftResponse, error) {
|
|
actor := auth.MustActor(ctx)
|
|
|
|
if _, err := s.templates.GetTemplateByID(ctx, templateID, actor.TenantID); err != nil {
|
|
return nil, mapTemplateLookupError(err)
|
|
}
|
|
|
|
stateJSON, err := json.Marshal(normalizeWizardState(req.WizardState, req.CurrentStep))
|
|
if err != nil {
|
|
return nil, response.ErrBadRequest(40016, "invalid_wizard_state", "wizard_state must be valid json")
|
|
}
|
|
|
|
if req.ArticleID == nil {
|
|
var articleID int64
|
|
err = s.pool.QueryRow(ctx, `
|
|
INSERT INTO articles (tenant_id, source_type, template_id, generate_status, publish_status, wizard_state_json)
|
|
VALUES ($1, 'template', $2, 'draft', 'unpublished', $3)
|
|
RETURNING id
|
|
`, actor.TenantID, templateID, stateJSON).Scan(&articleID)
|
|
if err != nil {
|
|
return nil, response.ErrInternal(50012, "draft_create_failed", "failed to create article draft")
|
|
}
|
|
invalidateArticleCaches(ctx, s.cache, actor.TenantID, &articleID)
|
|
return &SaveDraftResponse{ArticleID: articleID}, nil
|
|
}
|
|
|
|
var articleID int64
|
|
err = s.pool.QueryRow(ctx, `
|
|
UPDATE articles
|
|
SET wizard_state_json = $1,
|
|
generate_status = 'draft',
|
|
publish_status = 'unpublished',
|
|
updated_at = NOW()
|
|
WHERE id = $2
|
|
AND tenant_id = $3
|
|
AND template_id = $4
|
|
AND source_type = 'template'
|
|
AND deleted_at IS NULL
|
|
AND generate_status = 'draft'
|
|
RETURNING id
|
|
`, stateJSON, *req.ArticleID, actor.TenantID, templateID).Scan(&articleID)
|
|
if err != nil {
|
|
return nil, response.ErrConflict(40913, "draft_not_editable", "draft is not editable")
|
|
}
|
|
|
|
invalidateArticleCaches(ctx, s.cache, actor.TenantID, &articleID)
|
|
return &SaveDraftResponse{ArticleID: articleID}, nil
|
|
}
|
|
|
|
func (s *TemplateService) Generate(ctx context.Context, templateID int64, req GenerateRequest) (*GenerateResponse, error) {
|
|
actor := auth.MustActor(ctx)
|
|
|
|
templateRecord, err := s.templates.GetTemplateByID(ctx, templateID, actor.TenantID)
|
|
if err != nil {
|
|
return nil, mapTemplateLookupError(err)
|
|
}
|
|
if err := s.llm.Validate(); err != nil {
|
|
return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error())
|
|
}
|
|
|
|
initialTitle := resolveArticleTitle(req.InputParams, "")
|
|
wizardStateJSON, err := json.Marshal(normalizeWizardState(req.WizardState, wizardStateCurrentStep(req.WizardState)))
|
|
if err != nil {
|
|
return nil, response.ErrBadRequest(40016, "invalid_wizard_state", "wizard_state must be valid json")
|
|
}
|
|
quotaRepo := repository.NewQuotaRepository(s.pool)
|
|
balance, err := quotaRepo.GetCurrentBalance(ctx, actor.TenantID, "article_generation")
|
|
if err != nil || balance < 1 {
|
|
return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan")
|
|
}
|
|
|
|
taskInputParams := cloneInputParams(req.InputParams)
|
|
taskInputParams["enable_web_search"] = req.EnableWebSearch
|
|
if req.WebSearchLimit != nil {
|
|
taskInputParams["web_search_limit"] = *req.WebSearchLimit
|
|
}
|
|
taskInputParams = attachTemplateSnapshot(taskInputParams, templateRecord)
|
|
inputJSON, _ := json.Marshal(taskInputParams)
|
|
|
|
tx, err := s.pool.Begin(ctx)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("begin tx: %w", err)
|
|
}
|
|
defer func() {
|
|
_ = tx.Rollback(ctx)
|
|
}()
|
|
|
|
quotaTx := repository.NewQuotaRepository(tx)
|
|
articleTx := repository.NewArticleRepository(tx)
|
|
auditTx := repository.NewAuditRepository(tx)
|
|
|
|
newBalance := balance - 1
|
|
reason := "generation_reserve"
|
|
referenceType := "generation_task"
|
|
_, err = quotaTx.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
|
TenantID: actor.TenantID,
|
|
OperatorID: actor.UserID,
|
|
QuotaType: "article_generation",
|
|
Delta: -1,
|
|
BalanceAfter: newBalance,
|
|
Reason: &reason,
|
|
ReferenceType: &referenceType,
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("deduct quota: %w", err)
|
|
}
|
|
|
|
reservationID, err := quotaTx.CreateQuotaReservation(ctx, repository.QuotaReservationInput{
|
|
TenantID: actor.TenantID,
|
|
OperatorID: actor.UserID,
|
|
QuotaType: "article_generation",
|
|
ResourceType: "article",
|
|
ResourceID: nil,
|
|
ReservedAmount: 1,
|
|
ExpireAt: time.Now().Add(time.Hour),
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create reservation: %w", err)
|
|
}
|
|
|
|
var articleID int64
|
|
if req.ArticleID != nil {
|
|
err = tx.QueryRow(ctx, `
|
|
SELECT id
|
|
FROM articles
|
|
WHERE id = $1
|
|
AND tenant_id = $2
|
|
AND template_id = $3
|
|
AND source_type = 'template'
|
|
AND deleted_at IS NULL
|
|
AND generate_status = 'draft'
|
|
FOR UPDATE
|
|
`, *req.ArticleID, actor.TenantID, templateID).Scan(&articleID)
|
|
if err != nil {
|
|
return nil, response.ErrConflict(40913, "draft_not_editable", "draft is not editable")
|
|
}
|
|
if err := updateGeneratedArticleState(ctx, tx, articleID, actor.TenantID, initialTitle, "generating", wizardStateJSON); err != nil {
|
|
return nil, fmt.Errorf("mark draft generating: %w", err)
|
|
}
|
|
} else {
|
|
articleID, err = articleTx.CreateArticle(ctx, repository.CreateArticleInput{
|
|
TenantID: actor.TenantID,
|
|
SourceType: "template",
|
|
TemplateID: &templateID,
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create article: %w", err)
|
|
}
|
|
if err := updateGeneratedArticleState(ctx, tx, articleID, actor.TenantID, initialTitle, "generating", wizardStateJSON); err != nil {
|
|
return nil, fmt.Errorf("mark article generating: %w", err)
|
|
}
|
|
}
|
|
|
|
if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, actor.TenantID, articleID); err != nil {
|
|
return nil, fmt.Errorf("update reservation: %w", err)
|
|
}
|
|
|
|
taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{
|
|
OperatorID: actor.UserID,
|
|
TenantID: actor.TenantID,
|
|
ArticleID: &articleID,
|
|
QuotaReservationID: &reservationID,
|
|
TaskType: "template",
|
|
InputParamsJSON: inputJSON,
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create task: %w", err)
|
|
}
|
|
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return nil, fmt.Errorf("commit: %w", err)
|
|
}
|
|
invalidateArticleCaches(ctx, s.cache, actor.TenantID, &articleID)
|
|
|
|
job := generationJob{
|
|
OperatorID: actor.UserID,
|
|
TenantID: actor.TenantID,
|
|
ArticleID: articleID,
|
|
TaskID: taskID,
|
|
ReservationID: reservationID,
|
|
Params: taskInputParams,
|
|
InitialTitle: initialTitle,
|
|
}
|
|
|
|
if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{
|
|
TaskID: taskID,
|
|
TaskType: "template",
|
|
TenantID: actor.TenantID,
|
|
ArticleID: articleID,
|
|
}); err != nil {
|
|
s.failGeneration(context.Background(), job, initialTitle, "queue_enqueue", err)
|
|
return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry")
|
|
}
|
|
|
|
runtimeCfg := s.runtimeConfig()
|
|
if runtimeCfg.StreamEnabled && s.streams != nil {
|
|
s.streams.Start(articleID, taskID, initialTitle)
|
|
}
|
|
|
|
return &GenerateResponse{ArticleID: articleID, TaskID: taskID}, nil
|
|
}
|
|
|
|
func updateGeneratedArticleState(
|
|
ctx context.Context,
|
|
tx pgx.Tx,
|
|
articleID, tenantID int64,
|
|
title, status string,
|
|
wizardStateJSON []byte,
|
|
) error {
|
|
tag, err := tx.Exec(ctx, `
|
|
UPDATE articles
|
|
SET wizard_state_json = CASE
|
|
WHEN $1::jsonb IS NOT NULL THEN jsonb_set($1::jsonb, '{title}', to_jsonb($2::text), true)
|
|
WHEN NULLIF($2, '') IS NULL THEN COALESCE(wizard_state_json, '{}'::jsonb)
|
|
ELSE jsonb_set(COALESCE(wizard_state_json, '{}'::jsonb), '{title}', to_jsonb($2::text), true)
|
|
END,
|
|
generate_status = $3,
|
|
publish_status = 'unpublished',
|
|
updated_at = NOW()
|
|
WHERE id = $4 AND tenant_id = $5
|
|
`, nullableJSONString(wizardStateJSON), strings.TrimSpace(title), status, articleID, tenantID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if tag.RowsAffected() == 0 {
|
|
return fmt.Errorf("article not found")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *TemplateService) executeGeneration(ctx context.Context, job generationJob) {
|
|
now := time.Now()
|
|
title := strings.TrimSpace(job.InitialTitle)
|
|
if title == "" {
|
|
title = "Generated Article"
|
|
}
|
|
|
|
auditRepo := repository.NewAuditRepository(s.pool)
|
|
_ = auditRepo.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{
|
|
ID: job.TaskID,
|
|
TenantID: job.TenantID,
|
|
Status: "running",
|
|
StartedAt: &now,
|
|
})
|
|
|
|
knowledgePrompt := ""
|
|
if knowledgeGroupIDs := extractKnowledgeGroupIDs(job.Params["knowledge_group_ids"]); len(knowledgeGroupIDs) > 0 {
|
|
if s.knowledge == nil {
|
|
s.failGeneration(ctx, job, title, "knowledge_resolve", fmt.Errorf("knowledge service is not configured"))
|
|
return
|
|
}
|
|
knowledgeContext, resolveErr := s.knowledge.ResolveContext(
|
|
ctx,
|
|
job.TenantID,
|
|
knowledgeGroupIDs,
|
|
buildGenerationKnowledgeQuery(job.Params),
|
|
)
|
|
if resolveErr != nil {
|
|
s.failGeneration(ctx, job, title, "knowledge_resolve", resolveErr)
|
|
return
|
|
}
|
|
knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext)
|
|
}
|
|
|
|
prompt := buildGenerationPrompt(job.TemplateKey, job.TemplateName, job.PromptTemplate, job.Params, knowledgePrompt)
|
|
runtimeCfg := s.runtimeConfig()
|
|
generateReq := llm.GenerateRequest{
|
|
Prompt: prompt,
|
|
Timeout: runtimeCfg.ArticleTimeout,
|
|
MaxOutputTokens: runtimeCfg.MaxOutputTokens,
|
|
}
|
|
if job.WebSearch.Enabled {
|
|
generateReq.WebSearch = &job.WebSearch
|
|
}
|
|
result, err := s.llm.Generate(ctx, generateReq, func(delta string) {
|
|
if runtimeCfg.StreamEnabled && delta != "" {
|
|
s.streams.AppendDelta(job.ArticleID, job.TaskID, title, delta)
|
|
}
|
|
})
|
|
if err != nil {
|
|
s.failGeneration(ctx, job, title, "llm_generate", err)
|
|
return
|
|
}
|
|
|
|
content := strings.TrimSpace(result.Content)
|
|
if content == "" {
|
|
s.failGeneration(ctx, job, title, "llm_generate", fmt.Errorf("模型返回空内容"))
|
|
return
|
|
}
|
|
|
|
title = resolveArticleTitle(job.Params, content)
|
|
wordCount := estimateWordCount(content)
|
|
|
|
tx, err := s.pool.Begin(ctx)
|
|
if err != nil {
|
|
s.failGeneration(ctx, job, title, "persist_begin_tx", err)
|
|
return
|
|
}
|
|
defer rollbackGenerationTx(tx)
|
|
|
|
failWithRollback := func(stage string, err error) {
|
|
rollbackGenerationTx(tx)
|
|
s.failGeneration(ctx, job, title, stage, err)
|
|
}
|
|
|
|
articleRepo := repository.NewArticleRepository(tx)
|
|
quotaRepo := repository.NewQuotaRepository(tx)
|
|
auditTx := repository.NewAuditRepository(tx)
|
|
|
|
versionID, err := articleRepo.CreateArticleVersion(ctx, repository.CreateArticleVersionInput{
|
|
ArticleID: job.ArticleID,
|
|
VersionNo: 1,
|
|
Title: title,
|
|
HTMLContent: "",
|
|
MarkdownContent: content,
|
|
WordCount: wordCount,
|
|
SourceLabel: result.Model,
|
|
})
|
|
if err != nil {
|
|
failWithRollback("persist_create_version", err)
|
|
return
|
|
}
|
|
|
|
if err := articleRepo.UpdateArticleCurrentVersion(ctx, job.ArticleID, job.TenantID, versionID); err != nil {
|
|
failWithRollback("persist_current_version", err)
|
|
return
|
|
}
|
|
if err := articleRepo.UpdateArticleGenerateStatus(ctx, job.ArticleID, job.TenantID, "completed"); err != nil {
|
|
failWithRollback("persist_article_status", err)
|
|
return
|
|
}
|
|
|
|
completed := time.Now()
|
|
if err := auditTx.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{
|
|
ID: job.TaskID,
|
|
TenantID: job.TenantID,
|
|
Status: "completed",
|
|
StartedAt: &now,
|
|
CompletedAt: &completed,
|
|
}); err != nil {
|
|
failWithRollback("persist_task_status", err)
|
|
return
|
|
}
|
|
|
|
if err := quotaRepo.ConfirmReservation(ctx, job.ReservationID, job.TenantID); err != nil {
|
|
failWithRollback("confirm_reservation", err)
|
|
return
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
failWithRollback("persist_commit", err)
|
|
return
|
|
}
|
|
cleanupCtx, cancel := newGenerationCleanupContext()
|
|
defer cancel()
|
|
invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID)
|
|
|
|
if runtimeCfg.StreamEnabled && s.streams != nil {
|
|
s.streams.Complete(job.ArticleID, job.TaskID, title, content)
|
|
}
|
|
}
|
|
|
|
func (s *TemplateService) failGeneration(ctx context.Context, job generationJob, title, stage string, genErr error) {
|
|
cleanupCtx, cancel := newGenerationCleanupContext()
|
|
defer cancel()
|
|
|
|
errMsg := formatGenerationFailure(stage, genErr)
|
|
now := time.Now()
|
|
|
|
auditRepo := repository.NewAuditRepository(s.pool)
|
|
articleRepo := repository.NewArticleRepository(s.pool)
|
|
quotaRepo := repository.NewQuotaRepository(s.pool)
|
|
|
|
if err := auditRepo.UpdateGenerationTaskStatus(cleanupCtx, repository.GenerationTaskStatusInput{
|
|
ID: job.TaskID,
|
|
TenantID: job.TenantID,
|
|
Status: "failed",
|
|
ErrorMessage: &errMsg,
|
|
CompletedAt: &now,
|
|
}); err != nil && s.logger != nil {
|
|
s.logger.Warn("template generation task failed status update failed",
|
|
zap.Error(err),
|
|
zap.Int64("task_id", job.TaskID),
|
|
zap.Int64("article_id", job.ArticleID),
|
|
zap.String("stage", stage),
|
|
)
|
|
}
|
|
if err := articleRepo.UpdateArticleGenerateStatus(cleanupCtx, job.ArticleID, job.TenantID, "failed"); err != nil && s.logger != nil {
|
|
s.logger.Warn("template generation article failed status update failed",
|
|
zap.Error(err),
|
|
zap.Int64("task_id", job.TaskID),
|
|
zap.Int64("article_id", job.ArticleID),
|
|
zap.String("stage", stage),
|
|
)
|
|
}
|
|
|
|
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
|
if balanceErr == nil {
|
|
reason := "generation_refund"
|
|
referenceType := "generation_task"
|
|
taskReferenceID := job.TaskID
|
|
_, _ = quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
|
|
TenantID: job.TenantID,
|
|
OperatorID: job.OperatorID,
|
|
QuotaType: "article_generation",
|
|
Delta: 1,
|
|
BalanceAfter: balance + 1,
|
|
Reason: &reason,
|
|
ReferenceType: &referenceType,
|
|
ReferenceID: &taskReferenceID,
|
|
})
|
|
}
|
|
|
|
_ = quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID)
|
|
invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID)
|
|
|
|
runtimeCfg := s.runtimeConfig()
|
|
if runtimeCfg.StreamEnabled && s.streams != nil {
|
|
s.streams.Fail(job.ArticleID, job.TaskID, title, errMsg)
|
|
}
|
|
}
|
|
|
|
func formatGenerationFailure(stage string, genErr error) string {
|
|
cause := ""
|
|
if genErr != nil {
|
|
cause = strings.TrimSpace(genErr.Error())
|
|
}
|
|
|
|
label := generationFailureStageLabel(stage)
|
|
if cause == "" {
|
|
return fmt.Sprintf("%s [%s]", label, stage)
|
|
}
|
|
return fmt.Sprintf("%s [%s]: %s", label, stage, cause)
|
|
}
|
|
|
|
func generationFailureStageLabel(stage string) string {
|
|
switch stage {
|
|
case "task_load":
|
|
return "加载生成任务失败"
|
|
case "queue_enqueue":
|
|
return "任务入队失败"
|
|
case "llm_generate":
|
|
return "AI 正文生成失败"
|
|
case "knowledge_resolve":
|
|
return "知识库检索失败"
|
|
case "source_fetch":
|
|
return "源文章抓取失败"
|
|
case "persist_begin_tx":
|
|
return "写入生成结果事务启动失败"
|
|
case "persist_create_version":
|
|
return "保存文章版本失败"
|
|
case "persist_current_version":
|
|
return "更新文章当前版本失败"
|
|
case "persist_article_status":
|
|
return "更新文章生成状态失败"
|
|
case "persist_task_status":
|
|
return "更新生成任务状态失败"
|
|
case "confirm_reservation":
|
|
return "确认配额扣减失败"
|
|
case "persist_commit":
|
|
return "提交生成结果失败"
|
|
default:
|
|
return "文章生成失败"
|
|
}
|
|
}
|
|
|
|
func cloneInputParams(params map[string]interface{}) map[string]interface{} {
|
|
if len(params) == 0 {
|
|
return map[string]interface{}{}
|
|
}
|
|
|
|
cloned := make(map[string]interface{}, len(params))
|
|
for key, value := range params {
|
|
cloned[key] = value
|
|
}
|
|
return cloned
|
|
}
|
|
|
|
func normalizeWizardState(state map[string]interface{}, currentStep int) map[string]interface{} {
|
|
if state == nil {
|
|
state = map[string]interface{}{}
|
|
}
|
|
|
|
cloned := make(map[string]interface{}, len(state)+1)
|
|
for key, value := range state {
|
|
cloned[key] = value
|
|
}
|
|
|
|
if currentStep < 0 {
|
|
currentStep = 0
|
|
}
|
|
if currentStep > 2 {
|
|
currentStep = 2
|
|
}
|
|
cloned["current_step"] = currentStep
|
|
return cloned
|
|
}
|
|
|
|
func wizardStateCurrentStep(state map[string]interface{}) int {
|
|
if state == nil {
|
|
return 0
|
|
}
|
|
value, ok := state["current_step"]
|
|
if !ok {
|
|
return 0
|
|
}
|
|
switch typed := value.(type) {
|
|
case int:
|
|
return typed
|
|
case int32:
|
|
return int(typed)
|
|
case int64:
|
|
return int(typed)
|
|
case float64:
|
|
return int(typed)
|
|
default:
|
|
return 0
|
|
}
|
|
}
|
|
|
|
func nullableJSONString(value []byte) *string {
|
|
if len(value) == 0 {
|
|
return nil
|
|
}
|
|
text := string(value)
|
|
return &text
|
|
}
|
|
|
|
func extractKnowledgeGroupIDs(value interface{}) []int64 {
|
|
switch typed := value.(type) {
|
|
case []int64:
|
|
return normalizeInt64IDs(typed)
|
|
case []interface{}:
|
|
ids := make([]int64, 0, len(typed))
|
|
for _, item := range typed {
|
|
switch raw := item.(type) {
|
|
case int64:
|
|
ids = append(ids, raw)
|
|
case int32:
|
|
ids = append(ids, int64(raw))
|
|
case int:
|
|
ids = append(ids, int64(raw))
|
|
case float64:
|
|
ids = append(ids, int64(raw))
|
|
case float32:
|
|
ids = append(ids, int64(raw))
|
|
}
|
|
}
|
|
return normalizeInt64IDs(ids)
|
|
default:
|
|
return nil
|
|
}
|
|
}
|