diff --git a/server/internal/bootstrap/bootstrap.go b/server/internal/bootstrap/bootstrap.go index d90500b..05f31d0 100644 --- a/server/internal/bootstrap/bootstrap.go +++ b/server/internal/bootstrap/bootstrap.go @@ -94,7 +94,7 @@ func New(configPath string) (*App, error) { vectorStore := retrieval.NewQdrantStore(cfg.Qdrant, logger) objectStorageClient := objectstorage.New(cfg.ObjectStorage, logger) auditLogs := auditlog.NewAsyncWriter(pool, logger) - generationStreams := stream.NewGenerationHub() + generationStreams := stream.NewGenerationHub(mqClient) appCache := cache.New(cfg.Cache.Driver, rdb) if cfg.Server.Mode == "release" { diff --git a/server/internal/shared/stream/generation.go b/server/internal/shared/stream/generation.go index 5797f4c..0b55ecc 100644 --- a/server/internal/shared/stream/generation.go +++ b/server/internal/shared/stream/generation.go @@ -1,8 +1,14 @@ package stream import ( + "context" + "encoding/json" + "fmt" + "os" "sync" "time" + + "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" ) type GenerationEvent struct { @@ -18,6 +24,11 @@ type GenerationEvent struct { UpdatedAt time.Time `json:"updated_at"` } +type generationEventEnvelope struct { + Source string `json:"source"` + Event GenerationEvent `json:"event"` +} + type generationStream struct { snapshot GenerationEvent subscribers map[int]chan GenerationEvent @@ -26,17 +37,42 @@ type generationStream struct { } type GenerationHub struct { - mu sync.Mutex - streams map[int64]*generationStream + mu sync.Mutex + streams map[int64]*generationStream + rabbitMQ *rabbitmq.Client + instanceID string + consumerName string } -func NewGenerationHub() *GenerationHub { +func NewGenerationHub(rabbitMQClient *rabbitmq.Client) *GenerationHub { + instanceID := fmt.Sprintf("%d-%d", os.Getpid(), time.Now().UnixNano()) return &GenerationHub{ - streams: make(map[int64]*generationStream), + streams: make(map[int64]*generationStream), + rabbitMQ: rabbitMQClient, + instanceID: instanceID, + consumerName: "generation-stream-" + instanceID, } } -func (h *GenerationHub) Start(articleID, taskID int64, title string) { +func (h *GenerationHub) Run(ctx context.Context) { + if h == nil || h.rabbitMQ == nil { + return + } + go h.run(ctx) +} + +func (h *GenerationHub) Seed(event GenerationEvent) { + h.mu.Lock() + defer h.mu.Unlock() + + state := h.ensureLocked(event.ArticleID) + applySnapshotLocked(state, event) + if state.completed && len(state.subscribers) == 0 { + delete(h.streams, event.ArticleID) + } +} + +func (h *GenerationHub) StartEvent(articleID, taskID int64, title string) { h.publish(articleID, func(state *generationStream) GenerationEvent { now := time.Now() state.completed = false @@ -59,6 +95,10 @@ func (h *GenerationHub) Start(articleID, taskID int64, title string) { }) } +func (h *GenerationHub) Start(articleID, taskID int64, title string) { + h.StartEvent(articleID, taskID, title) +} + func (h *GenerationHub) AppendDelta(articleID, taskID int64, title, delta string) { h.publish(articleID, func(state *generationStream) GenerationEvent { now := time.Now() @@ -69,6 +109,7 @@ func (h *GenerationHub) AppendDelta(articleID, taskID int64, title, delta string state.snapshot.Status = "generating" state.snapshot.Content += delta state.snapshot.UpdatedAt = now + state.completed = false return GenerationEvent{ Type: "delta", @@ -188,14 +229,83 @@ func (h *GenerationHub) Snapshot(articleID int64) (*GenerationEvent, bool) { return &snapshot, true } +func (h *GenerationHub) run(ctx context.Context) { + for { + deliveries, ch, err := h.rabbitMQ.ConsumeGenerationEvents(h.consumerName) + if err != nil { + select { + case <-ctx.Done(): + return + case <-time.After(2 * time.Second): + continue + } + } + + closed := false + for !closed { + select { + case <-ctx.Done(): + _ = ch.Close() + return + case delivery, ok := <-deliveries: + if !ok { + closed = true + continue + } + h.handleDelivery(delivery.Body) + _ = delivery.Ack(false) + } + } + _ = ch.Close() + + select { + case <-ctx.Done(): + return + case <-time.After(2 * time.Second): + } + } +} + +func (h *GenerationHub) handleDelivery(body []byte) { + var envelope generationEventEnvelope + if err := json.Unmarshal(body, &envelope); err != nil { + return + } + if envelope.Source == h.instanceID { + return + } + + h.mu.Lock() + state := h.streams[envelope.Event.ArticleID] + if state == nil { + h.mu.Unlock() + return + } + + applyRemoteEventLocked(state, envelope.Event) + subscribers := collectSubscribersLocked(state) + shouldDelete := state.completed && len(state.subscribers) == 0 + if shouldDelete { + delete(h.streams, envelope.Event.ArticleID) + } + h.mu.Unlock() + + for _, ch := range subscribers { + select { + case ch <- envelope.Event: + default: + } + } +} + func (h *GenerationHub) publish(articleID int64, mutate func(*generationStream) GenerationEvent) { h.mu.Lock() state := h.ensureLocked(articleID) event := mutate(state) - - subscribers := make([]chan GenerationEvent, 0, len(state.subscribers)) - for _, ch := range state.subscribers { - subscribers = append(subscribers, ch) + subscribers := collectSubscribersLocked(state) + shouldDelete := state.completed && len(state.subscribers) == 0 + if shouldDelete { + delete(h.streams, articleID) } h.mu.Unlock() @@ -205,6 +315,16 @@ func (h *GenerationHub) publish(articleID int64, mutate func(*generationStream) default: } } + + if h.rabbitMQ != nil { + payload, err := json.Marshal(generationEventEnvelope{ + Source: h.instanceID, + Event: event, + }) + if err == nil { + _ = h.rabbitMQ.PublishGenerationEvent(context.Background(), payload) + } + } } func (h *GenerationHub) ensureLocked(articleID int64) *generationStream { @@ -219,3 +339,78 @@ func (h *GenerationHub) ensureLocked(articleID int64) *generationStream { h.streams[articleID] = state return state } + +func collectSubscribersLocked(state *generationStream) []chan GenerationEvent { + subscribers := make([]chan GenerationEvent, 0, len(state.subscribers)) + for _, ch := range state.subscribers { + subscribers = append(subscribers, ch) + } + return subscribers +} + +func applySnapshotLocked(state *generationStream, event GenerationEvent) { + state.snapshot = GenerationEvent{ + Type: "snapshot", + ArticleID: event.ArticleID, + TaskID: event.TaskID, + Title: event.Title, + Status: event.Status, + Content: event.Content, + Error: event.Error, + Done: event.Done, + UpdatedAt: event.UpdatedAt, + } + state.completed = event.Done +} + +func applyRemoteEventLocked(state *generationStream, event GenerationEvent) { + switch event.Type { + case "delta": + state.completed = false + state.snapshot.Type = "snapshot" + state.snapshot.ArticleID = event.ArticleID + state.snapshot.TaskID = event.TaskID + state.snapshot.Title = event.Title + state.snapshot.Status = event.Status + if event.Content != "" { + state.snapshot.Content = event.Content + } else { + state.snapshot.Content += event.Delta + } + state.snapshot.UpdatedAt = event.UpdatedAt + case "completed": + applySnapshotLocked(state, GenerationEvent{ + Type: "snapshot", + ArticleID: event.ArticleID, + TaskID: event.TaskID, + Title: event.Title, + Status: "completed", + Content: event.Content, + Done: true, + UpdatedAt: event.UpdatedAt, + }) + case "error": + applySnapshotLocked(state, GenerationEvent{ + Type: "snapshot", + ArticleID: event.ArticleID, + TaskID: event.TaskID, + Title: event.Title, + Status: "failed", + Error: event.Error, + Done: true, + UpdatedAt: event.UpdatedAt, + }) + default: + applySnapshotLocked(state, GenerationEvent{ + Type: "snapshot", + ArticleID: event.ArticleID, + TaskID: event.TaskID, + Title: event.Title, + Status: event.Status, + Content: event.Content, + Error: event.Error, + Done: event.Done, + UpdatedAt: event.UpdatedAt, + }) + } +} diff --git a/server/internal/tenant/app/article_generation_payload.go b/server/internal/tenant/app/article_generation_payload.go new file mode 100644 index 0000000..58159c7 --- /dev/null +++ b/server/internal/tenant/app/article_generation_payload.go @@ -0,0 +1,97 @@ +package app + +import "github.com/geo-platform/tenant-api/internal/tenant/repository" + +const ( + generationPayloadTemplateKey = "__snapshot_template_key" + generationPayloadTemplateName = "__snapshot_template_name" + generationPayloadTemplatePromptTemplate = "__snapshot_template_prompt_template" + generationPayloadPromptRuleID = "__snapshot_prompt_rule_id" + generationPayloadPromptRuleName = "__snapshot_prompt_rule_name" + generationPayloadPromptRuleContent = "__snapshot_prompt_rule_content" + generationPayloadPromptRuleScene = "__snapshot_prompt_rule_scene" + generationPayloadPromptRuleTone = "__snapshot_prompt_rule_tone" + generationPayloadPromptRuleWordCount = "__snapshot_prompt_rule_word_count" + generationPayloadKnowledgePrompt = "__snapshot_knowledge_prompt" +) + +func attachTemplateSnapshot(payload map[string]interface{}, record *repository.TemplateRecord) map[string]interface{} { + if payload == nil { + payload = map[string]interface{}{} + } + if record == nil { + return payload + } + + payload[generationPayloadTemplateKey] = record.TemplateKey + payload[generationPayloadTemplateName] = record.TemplateName + if record.PromptTemplate != nil { + payload[generationPayloadTemplatePromptTemplate] = *record.PromptTemplate + } + + return payload +} + +func extractTemplateSnapshot(payload map[string]interface{}) (string, string, *string) { + templateKey := extractString(payload, generationPayloadTemplateKey) + templateName := extractString(payload, generationPayloadTemplateName) + promptTemplateText := extractString(payload, generationPayloadTemplatePromptTemplate) + if templateKey == "" || templateName == "" { + return "", "", nil + } + if promptTemplateText == "" { + return templateKey, templateName, nil + } + return templateKey, templateName, &promptTemplateText +} + +func attachPromptRuleSnapshot(payload map[string]interface{}, rule *promptRuleGenerationRecord) map[string]interface{} { + if payload == nil { + payload = map[string]interface{}{} + } + if rule == nil { + return payload + } + + payload[generationPayloadPromptRuleID] = rule.ID + payload[generationPayloadPromptRuleName] = rule.Name + payload[generationPayloadPromptRuleContent] = rule.PromptContent + if rule.Scene != nil { + payload[generationPayloadPromptRuleScene] = *rule.Scene + } + if rule.DefaultTone != nil { + payload[generationPayloadPromptRuleTone] = *rule.DefaultTone + } + if rule.DefaultWordCount != nil { + payload[generationPayloadPromptRuleWordCount] = *rule.DefaultWordCount + } + + return payload +} + +func extractPromptRuleSnapshot(payload map[string]interface{}) *promptRuleGenerationRecord { + ruleID := int64(extractInt(payload, generationPayloadPromptRuleID)) + name := extractString(payload, generationPayloadPromptRuleName) + content := extractString(payload, generationPayloadPromptRuleContent) + if ruleID <= 0 || content == "" { + return nil + } + + record := &promptRuleGenerationRecord{ + ID: ruleID, + Name: name, + PromptContent: content, + Status: "enabled", + } + if scene := extractString(payload, generationPayloadPromptRuleScene); scene != "" { + record.Scene = &scene + } + if tone := extractString(payload, generationPayloadPromptRuleTone); tone != "" { + record.DefaultTone = &tone + } + if wordCount := extractInt(payload, generationPayloadPromptRuleWordCount); wordCount > 0 { + value := wordCount + record.DefaultWordCount = &value + } + return record +} diff --git a/server/internal/tenant/app/article_generation_queue.go b/server/internal/tenant/app/article_generation_queue.go new file mode 100644 index 0000000..446074b --- /dev/null +++ b/server/internal/tenant/app/article_generation_queue.go @@ -0,0 +1,63 @@ +package app + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" +) + +type articleGenerationTaskEnvelope struct { + TaskID int64 `json:"task_id"` + TaskType string `json:"task_type"` + TenantID int64 `json:"tenant_id"` + ArticleID int64 `json:"article_id"` +} + +type ArticleGenerationTaskEnvelope struct { + TaskID int64 `json:"task_id"` + TaskType string `json:"task_type"` + TenantID int64 `json:"tenant_id"` + ArticleID int64 `json:"article_id"` +} + +func publishArticleGenerationTask(ctx context.Context, client *rabbitmq.Client, envelope articleGenerationTaskEnvelope) error { + if client == nil { + return fmt.Errorf("rabbitmq client is not configured") + } + + payload, err := json.Marshal(envelope) + if err != nil { + return fmt.Errorf("marshal generation task envelope: %w", err) + } + + if err := client.PublishGenerationTask(ctx, payload); err != nil { + return fmt.Errorf("publish generation task: %w", err) + } + return nil +} + +func decodeArticleGenerationTaskEnvelope(payload []byte) (*articleGenerationTaskEnvelope, error) { + var envelope articleGenerationTaskEnvelope + if err := json.Unmarshal(payload, &envelope); err != nil { + return nil, fmt.Errorf("decode generation task envelope: %w", err) + } + if envelope.TaskID <= 0 { + return nil, fmt.Errorf("generation task envelope task_id is required") + } + return &envelope, nil +} + +func DecodeArticleGenerationTaskEnvelope(payload []byte) (*ArticleGenerationTaskEnvelope, error) { + envelope, err := decodeArticleGenerationTaskEnvelope(payload) + if err != nil { + return nil, err + } + return &ArticleGenerationTaskEnvelope{ + TaskID: envelope.TaskID, + TaskType: envelope.TaskType, + TenantID: envelope.TenantID, + ArticleID: envelope.ArticleID, + }, nil +} diff --git a/server/internal/tenant/app/generation_task_runtime.go b/server/internal/tenant/app/generation_task_runtime.go new file mode 100644 index 0000000..ce286ef --- /dev/null +++ b/server/internal/tenant/app/generation_task_runtime.go @@ -0,0 +1,222 @@ +package app + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/geo-platform/tenant-api/internal/shared/llm" + "github.com/geo-platform/tenant-api/internal/tenant/repository" +) + +type ClaimedGenerationTask struct { + ID int64 + TenantID int64 + OperatorID int64 + ArticleID int64 + QuotaReservationID int64 + TaskType string + InputParams map[string]interface{} +} + +func HandleClaimedGenerationTaskFailure( + ctx context.Context, + pool *pgxpool.Pool, + templateService *TemplateService, + promptRuleService *PromptRuleGenerationService, + task ClaimedGenerationTask, + taskErr error, +) error { + errMsg := formatGenerationFailure("task_load", taskErr) + now := time.Now() + + auditRepo := repository.NewAuditRepository(pool) + articleRepo := repository.NewArticleRepository(pool) + quotaRepo := repository.NewQuotaRepository(pool) + + _ = auditRepo.UpdateGenerationTaskStatus(ctx, repository.GenerationTaskStatusInput{ + ID: task.ID, + TenantID: task.TenantID, + Status: "failed", + ErrorMessage: &errMsg, + CompletedAt: &now, + }) + if task.ArticleID > 0 { + _ = articleRepo.UpdateArticleGenerateStatus(ctx, task.ArticleID, task.TenantID, "failed") + } + + if task.QuotaReservationID > 0 { + balance, balanceErr := quotaRepo.GetCurrentBalance(ctx, task.TenantID, "article_generation") + if balanceErr == nil && task.OperatorID > 0 { + reason := "generation_refund" + referenceType := "generation_task" + taskReferenceID := task.ID + _, _ = quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ + TenantID: task.TenantID, + OperatorID: task.OperatorID, + QuotaType: "article_generation", + Delta: 1, + BalanceAfter: balance + 1, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &taskReferenceID, + }) + } + + _ = quotaRepo.RefundReservation(ctx, task.QuotaReservationID, task.TenantID) + } + + title := resolveArticleTitle(task.InputParams, "") + if strings.TrimSpace(extractString(task.InputParams, "task_name")) != "" { + title = strings.TrimSpace(extractString(task.InputParams, "task_name")) + } + if title == "" { + title = promptRuleGeneratingTitlePlaceholder + } + + if task.TaskType == "template" && templateService != nil && templateService.streamEnabled { + templateService.streams.Fail(task.ArticleID, task.ID, title, errMsg) + } + if task.TaskType == "custom_generation" && promptRuleService != nil && promptRuleService.streamEnabled { + promptRuleService.streams.Fail(task.ArticleID, task.ID, title, errMsg) + } + + return nil +} + +func (s *TemplateService) ProcessQueuedTask(ctx context.Context, task ClaimedGenerationTask) { + job := generationJob{ + OperatorID: task.OperatorID, + TenantID: task.TenantID, + ArticleID: task.ArticleID, + TaskID: task.ID, + ReservationID: task.QuotaReservationID, + Params: cloneInputParams(task.InputParams), + InitialTitle: resolveArticleTitle(task.InputParams, ""), + WebSearch: llm.WebSearchOptions{ + Enabled: extractBool(task.InputParams, "enable_web_search"), + }, + } + if limit := extractInt(task.InputParams, "web_search_limit"); limit > 0 { + job.WebSearch.Limit = int32(limit) + } + + templateKey, templateName, promptTemplate := extractTemplateSnapshot(task.InputParams) + if templateKey == "" || templateName == "" { + record, err := s.loadTemplateRecordForTask(ctx, task) + if err != nil { + s.failGeneration(context.Background(), job, job.InitialTitle, "task_load", err) + return + } + templateKey = record.TemplateKey + templateName = record.TemplateName + promptTemplate = record.PromptTemplate + } + job.TemplateKey = templateKey + job.TemplateName = templateName + job.PromptTemplate = promptTemplate + + s.executeGeneration(ctx, job) +} + +func (s *TemplateService) loadTemplateRecordForTask(ctx context.Context, task ClaimedGenerationTask) (*repository.TemplateRecord, error) { + var templateID int64 + if err := s.pool.QueryRow(ctx, ` + SELECT template_id + FROM articles + WHERE id = $1 + AND tenant_id = $2 + AND source_type = 'template' + AND deleted_at IS NULL + `, task.ArticleID, task.TenantID).Scan(&templateID); err != nil { + return nil, fmt.Errorf("load task template id: %w", err) + } + + record, err := s.templates.GetTemplateByID(ctx, templateID, task.TenantID) + if err != nil { + return nil, fmt.Errorf("load task template: %w", err) + } + return record, nil +} + +func (s *PromptRuleGenerationService) ProcessQueuedTask(ctx context.Context, task ClaimedGenerationTask) { + inputParams := cloneInputParams(task.InputParams) + initialTitle := strings.TrimSpace(extractString(inputParams, "task_name")) + if initialTitle == "" { + initialTitle = promptRuleGeneratingTitlePlaceholder + } + + job := promptRuleGenerationJob{ + OperatorID: task.OperatorID, + TenantID: task.TenantID, + ArticleID: task.ArticleID, + TaskID: task.ID, + ReservationID: task.QuotaReservationID, + InputParams: inputParams, + InitialTitle: initialTitle, + WebSearch: llm.WebSearchOptions{ + Enabled: extractBool(inputParams, "enable_web_search"), + }, + } + + rule := extractPromptRuleSnapshot(inputParams) + if rule == nil { + promptRuleID, err := s.resolvePromptRuleIDForTask(ctx, task) + if err != nil { + s.failGeneration(context.Background(), job, initialTitle, "task_load", err) + return + } + rule, err = s.loadPromptRule(ctx, task.TenantID, promptRuleID) + if err != nil { + s.failGeneration(context.Background(), job, initialTitle, "task_load", err) + return + } + } + + knowledgePrompt := strings.TrimSpace(extractString(inputParams, generationPayloadKnowledgePrompt)) + if knowledgePrompt == "" { + if knowledgeGroupIDs := extractKnowledgeGroupIDs(inputParams["knowledge_group_ids"]); len(knowledgeGroupIDs) > 0 { + if s.knowledge == nil { + s.failGeneration(context.Background(), job, initialTitle, "knowledge_resolve", fmt.Errorf("knowledge service is not configured")) + return + } + + knowledgeContext, err := s.knowledge.ResolveContext( + ctx, + task.TenantID, + knowledgeGroupIDs, + buildPromptRuleKnowledgeQuery(rule, inputParams), + ) + if err != nil { + s.failGeneration(context.Background(), job, initialTitle, "knowledge_resolve", err) + return + } + knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext) + } + } + + job.Prompt = buildPromptRuleGenerationPrompt(rule, inputParams, knowledgePrompt) + s.executeGeneration(ctx, job) +} + +func (s *PromptRuleGenerationService) resolvePromptRuleIDForTask(ctx context.Context, task ClaimedGenerationTask) (int64, error) { + if promptRuleID := int64(extractInt(task.InputParams, "prompt_rule_id")); promptRuleID > 0 { + return promptRuleID, nil + } + + var promptRuleID int64 + if err := s.pool.QueryRow(ctx, ` + SELECT prompt_rule_id + FROM articles + WHERE id = $1 + AND tenant_id = $2 + AND source_type = 'custom_generation' + AND deleted_at IS NULL + `, task.ArticleID, task.TenantID).Scan(&promptRuleID); err != nil { + return 0, fmt.Errorf("load task prompt rule id: %w", err) + } + return promptRuleID, nil +} diff --git a/server/internal/tenant/app/knowledge_service.go b/server/internal/tenant/app/knowledge_service.go index 25a2e28..e6fda28 100644 --- a/server/internal/tenant/app/knowledge_service.go +++ b/server/internal/tenant/app/knowledge_service.go @@ -219,7 +219,6 @@ func NewKnowledgeService( go svc.runParseWorker() } go svc.recoverPendingParseTasks() - go svc.runDeletedCleanupWorker() } return svc @@ -1753,17 +1752,6 @@ func (s *KnowledgeService) cleanupDeletedKnowledgeItemAsync(tenantID, itemID int }() } -func (s *KnowledgeService) runDeletedCleanupWorker() { - s.scheduleDeletedKnowledgeCleanup() - - ticker := time.NewTicker(defaultKnowledgeCleanupPoll) - defer ticker.Stop() - - for range ticker.C { - s.scheduleDeletedKnowledgeCleanup() - } -} - func (s *KnowledgeService) scheduleDeletedKnowledgeCleanup() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() @@ -1778,6 +1766,13 @@ func (s *KnowledgeService) scheduleDeletedKnowledgeCleanup() { } } +func (s *KnowledgeService) RunDeletedCleanupSweep() { + if s == nil { + return + } + s.scheduleDeletedKnowledgeCleanup() +} + func (s *KnowledgeService) listKnowledgeCleanupCandidates( ctx context.Context, limit int, diff --git a/server/internal/tenant/app/prompt_generate_service.go b/server/internal/tenant/app/prompt_generate_service.go index 4ef53f6..7cbccb0 100644 --- a/server/internal/tenant/app/prompt_generate_service.go +++ b/server/internal/tenant/app/prompt_generate_service.go @@ -13,6 +13,7 @@ import ( "github.com/geo-platform/tenant-api/internal/shared/auth" "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/prompts" @@ -22,12 +23,12 @@ import ( type PromptRuleGenerationService struct { pool *pgxpool.Pool llm llm.Client + rabbitMQ *rabbitmq.Client knowledge *KnowledgeService streams *stream.GenerationHub streamEnabled bool articleTimeout time.Duration maxOutputTokens int64 - jobs chan promptRuleGenerationJob } type promptRuleGenerationJob struct { @@ -50,10 +51,34 @@ type GenerateFromRuleRequest struct { } type GenerateFromRuleResponse struct { + ArticleID int64 `json:"article_id"` + TaskID int64 `json:"task_id"` + ArticleIDs []int64 `json:"article_ids"` + TaskIDs []int64 `json:"task_ids"` + Items []GenerateFromRuleItem `json:"items"` + RequestedGenerateCount int `json:"requested_generate_count"` + SuccessCount int `json:"success_count"` + FailedCount int `json:"failed_count"` +} + +type GenerateFromRuleItem struct { ArticleID int64 `json:"article_id"` TaskID int64 `json:"task_id"` } +type SchedulePromptRuleGenerationInput struct { + ScheduleTaskID int64 + OperatorID *int64 + TenantID int64 + PromptRuleID int64 + BrandID *int64 + Name string + TargetPlatform *string + EnableWebSearch bool + GenerateCount int + ScheduledFor time.Time +} + type promptRuleGenerationRecord struct { ID int64 Name string @@ -68,24 +93,25 @@ type promptRuleGenerationRecord struct { const promptRuleGeneratingTitlePlaceholder = "标题生成中" +type enqueuePromptRuleGenerationInput struct { + OperatorID *int64 + TenantID int64 + PromptRuleID int64 + BrandID *int64 + TargetPlatform *string + ExtraParams map[string]interface{} + FallbackName string +} + func NewPromptRuleGenerationService( pool *pgxpool.Pool, llmClient llm.Client, + rabbitMQClient *rabbitmq.Client, knowledge *KnowledgeService, streams *stream.GenerationHub, cfg config.GenerationConfig, llmMaxOutputTokens int64, ) *PromptRuleGenerationService { - queueSize := cfg.QueueSize - if queueSize <= 0 { - queueSize = 128 - } - - workerCount := cfg.WorkerConcurrency - if workerCount <= 0 { - workerCount = 1 - } - articleTimeout := cfg.ArticleTimeout if articleTimeout <= 0 { articleTimeout = 8 * time.Minute @@ -94,16 +120,12 @@ func NewPromptRuleGenerationService( svc := &PromptRuleGenerationService{ pool: pool, llm: llmClient, + rabbitMQ: rabbitMQClient, knowledge: knowledge, streams: streams, streamEnabled: cfg.StreamEnabled, articleTimeout: articleTimeout, maxOutputTokens: llmMaxOutputTokens, - jobs: make(chan promptRuleGenerationJob, queueSize), - } - - for i := 0; i < workerCount; i++ { - go svc.runGenerationWorker() } return svc @@ -111,12 +133,86 @@ func NewPromptRuleGenerationService( func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req GenerateFromRuleRequest) (*GenerateFromRuleResponse, error) { actor := auth.MustActor(ctx) + extra := clonePromptRuleExtraParams(req.ExtraParams) + if strings.TrimSpace(extractString(extra, "generation_mode")) == "" { + extra["generation_mode"] = "instant" + } + generateCount, err := normalizePromptRuleGenerateCount(extractGenerateCount(extra)) + if err != nil { + return nil, err + } + quotaRepo := repository.NewQuotaRepository(s.pool) + balance, err := quotaRepo.GetCurrentBalance(ctx, actor.TenantID, "article_generation") + if err != nil || balance < generateCount { + return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan") + } + + results := make([]GenerateFromRuleItem, 0, generateCount) + var lastErr error + for idx := 0; idx < generateCount; idx++ { + result, enqueueErr := s.enqueuePromptRuleGeneration(ctx, enqueuePromptRuleGenerationInput{ + OperatorID: &actor.UserID, + TenantID: actor.TenantID, + PromptRuleID: req.PromptRuleID, + BrandID: req.BrandID, + TargetPlatform: req.TargetPlatform, + ExtraParams: clonePromptRuleExtraParams(extra), + FallbackName: "即时任务", + }) + if enqueueErr != nil { + lastErr = enqueueErr + continue + } + results = append(results, GenerateFromRuleItem{ + ArticleID: result.ArticleID, + TaskID: result.TaskID, + }) + } + + if len(results) == 0 { + if lastErr != nil { + return nil, lastErr + } + return nil, response.ErrInternal(50012, "create_failed", "failed to create generation task") + } + + return buildGenerateFromRuleResponse(results, generateCount), nil +} + +func (s *PromptRuleGenerationService) EnqueueScheduledGeneration(ctx context.Context, input SchedulePromptRuleGenerationInput) (*GenerateFromRuleResponse, error) { + extra := map[string]interface{}{ + "generation_mode": "schedule", + "schedule_task_id": input.ScheduleTaskID, + "enable_web_search": input.EnableWebSearch, + } + if input.GenerateCount > 0 { + extra["generate_count"] = input.GenerateCount + } + if name := strings.TrimSpace(input.Name); name != "" { + extra["task_name"] = name + } + if !input.ScheduledFor.IsZero() { + extra["scheduled_for"] = input.ScheduledFor.UTC().Format(time.RFC3339) + } + + return s.enqueuePromptRuleGeneration(ctx, enqueuePromptRuleGenerationInput{ + OperatorID: input.OperatorID, + TenantID: input.TenantID, + PromptRuleID: input.PromptRuleID, + BrandID: input.BrandID, + TargetPlatform: input.TargetPlatform, + ExtraParams: extra, + FallbackName: "定时任务", + }) +} + +func (s *PromptRuleGenerationService) enqueuePromptRuleGeneration(ctx context.Context, input enqueuePromptRuleGenerationInput) (*GenerateFromRuleResponse, error) { if err := s.llm.Validate(); err != nil { return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error()) } - rule, err := s.loadPromptRule(ctx, actor.TenantID, req.PromptRuleID) + rule, err := s.loadPromptRule(ctx, input.TenantID, input.PromptRuleID) if err != nil { return nil, err } @@ -124,9 +220,11 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req return nil, response.ErrBadRequest(40017, "prompt_rule_disabled", "prompt rule is disabled") } - extra := req.ExtraParams - if extractGenerateCount(extra) > 1 { - return nil, response.ErrBadRequest(40018, "generate_count_not_supported", "prompt rule instant generation currently supports generate_count = 1 only") + extra := input.ExtraParams + generationMode := strings.TrimSpace(extractString(extra, "generation_mode")) + generateCount, err := normalizePromptRuleGenerateCount(extractGenerateCount(extra)) + if err != nil { + return nil, err } taskName := strings.TrimSpace(extractString(extra, "task_name")) @@ -134,34 +232,46 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req taskName = strings.TrimSpace(rule.Name) } if taskName == "" { - taskName = "即时任务" + taskName = strings.TrimSpace(input.FallbackName) + } + if taskName == "" { + taskName = "内容任务" } quotaRepo := repository.NewQuotaRepository(s.pool) - balance, err := quotaRepo.GetCurrentBalance(ctx, actor.TenantID, "article_generation") + balance, err := quotaRepo.GetCurrentBalance(ctx, input.TenantID, "article_generation") if err != nil || balance < 1 { return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan") } inputParams := map[string]interface{}{ - "prompt_rule_id": req.PromptRuleID, + "prompt_rule_id": input.PromptRuleID, "task_name": taskName, - "target_platform": strings.TrimSpace(derefString(req.TargetPlatform)), + "target_platform": strings.TrimSpace(derefString(input.TargetPlatform)), "enable_web_search": extractBool(extra, "enable_web_search"), + "generate_count": generateCount, } - if generationMode := strings.TrimSpace(extractString(extra, "generation_mode")); generationMode != "" { + if generationMode != "" { inputParams["generation_mode"] = generationMode } - platformIDs := parsePlatformIDs(derefString(req.TargetPlatform)) + if scheduleTaskID := extractInt(extra, "schedule_task_id"); scheduleTaskID > 0 { + inputParams["schedule_task_id"] = scheduleTaskID + } + if scheduledFor := strings.TrimSpace(extractString(extra, "scheduled_for")); scheduledFor != "" { + inputParams["scheduled_for"] = scheduledFor + } + + platformIDs := parsePlatformIDs(derefString(input.TargetPlatform)) if len(platformIDs) > 0 { inputParams["target_platforms"] = platformIDs } - if req.BrandID != nil { - inputParams["brand_id"] = *req.BrandID + if input.BrandID != nil { + inputParams["brand_id"] = *input.BrandID } if len(rule.KnowledgeGroupIDs) > 0 { inputParams["knowledge_group_ids"] = rule.KnowledgeGroupIDs } + inputParams = attachPromptRuleSnapshot(inputParams, rule) knowledgePrompt := "" if len(rule.KnowledgeGroupIDs) > 0 { @@ -170,7 +280,7 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req } knowledgeContext, resolveErr := s.knowledge.ResolveContext( ctx, - actor.TenantID, + input.TenantID, rule.KnowledgeGroupIDs, buildPromptRuleKnowledgeQuery(rule, inputParams), ) @@ -178,6 +288,9 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req return nil, resolveErr } knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext) + if knowledgePrompt != "" { + inputParams[generationPayloadKnowledgePrompt] = knowledgePrompt + } } inputJSON, _ := json.Marshal(inputParams) @@ -193,13 +306,17 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req quotaTx := repository.NewQuotaRepository(tx) auditTx := repository.NewAuditRepository(tx) + operatorID := int64(0) + if input.OperatorID != nil { + operatorID = *input.OperatorID + } newBalance := balance - 1 reason := "generation_reserve" referenceType := "generation_task" _, err = quotaTx.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ - TenantID: actor.TenantID, - OperatorID: actor.UserID, + TenantID: input.TenantID, + OperatorID: operatorID, QuotaType: "article_generation", Delta: -1, BalanceAfter: newBalance, @@ -211,8 +328,8 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req } reservationID, err := quotaTx.CreateQuotaReservation(ctx, repository.QuotaReservationInput{ - TenantID: actor.TenantID, - OperatorID: actor.UserID, + TenantID: input.TenantID, + OperatorID: operatorID, QuotaType: "article_generation", ResourceType: "article", ResourceID: nil, @@ -228,18 +345,18 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req INSERT INTO articles (tenant_id, source_type, prompt_rule_id, generate_status, publish_status, wizard_state_json) VALUES ($1, 'custom_generation', $2, 'generating', 'unpublished', $3) RETURNING id - `, actor.TenantID, req.PromptRuleID, wizardStateJSON).Scan(&articleID) + `, input.TenantID, input.PromptRuleID, wizardStateJSON).Scan(&articleID) if err != nil { return nil, response.ErrInternal(50012, "create_failed", "failed to create prompt rule article") } - if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, actor.TenantID, articleID); err != nil { + if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, input.TenantID, articleID); err != nil { return nil, response.ErrInternal(50012, "create_failed", "failed to update quota reservation resource") } taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{ - OperatorID: actor.UserID, - TenantID: actor.TenantID, + OperatorID: operatorID, + TenantID: input.TenantID, ArticleID: &articleID, QuotaReservationID: &reservationID, TaskType: "custom_generation", @@ -254,25 +371,26 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req } job := promptRuleGenerationJob{ - OperatorID: actor.UserID, - TenantID: actor.TenantID, + OperatorID: operatorID, + TenantID: input.TenantID, ArticleID: articleID, TaskID: taskID, ReservationID: reservationID, - Prompt: buildPromptRuleGenerationPrompt(rule, inputParams, knowledgePrompt), InputParams: inputParams, InitialTitle: promptRuleGeneratingTitlePlaceholder, - WebSearch: llm.WebSearchOptions{ - Enabled: extractBool(extra, "enable_web_search"), - }, } - if err := s.enqueueGeneration(job); err != nil { - s.failGeneration(context.Background(), job, job.InitialTitle, "queue_enqueue", err) + if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{ + TaskID: taskID, + TaskType: "custom_generation", + TenantID: input.TenantID, + ArticleID: articleID, + }); err != nil { + s.failGeneration(context.Background(), job, promptRuleGeneratingTitlePlaceholder, "queue_enqueue", err) return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry") } - if s.streamEnabled { + if s.streamEnabled && s.streams != nil { s.streams.Start(articleID, taskID, promptRuleGeneratingTitlePlaceholder) } @@ -322,21 +440,6 @@ func (s *PromptRuleGenerationService) loadPromptRule(ctx context.Context, tenant return record, nil } -func (s *PromptRuleGenerationService) enqueueGeneration(job promptRuleGenerationJob) error { - select { - case s.jobs <- job: - return nil - default: - return fmt.Errorf("generation queue is full") - } -} - -func (s *PromptRuleGenerationService) runGenerationWorker() { - for job := range s.jobs { - s.executeGeneration(context.Background(), job) - } -} - func (s *PromptRuleGenerationService) executeGeneration(ctx context.Context, job promptRuleGenerationJob) { now := time.Now() title := strings.TrimSpace(job.InitialTitle) @@ -557,6 +660,50 @@ func extractGenerateCount(extra map[string]interface{}) int { return value } +func normalizePromptRuleGenerateCount(value int) (int, error) { + if value < 1 || value > 20 { + return 0, response.ErrBadRequest(40022, "invalid_generate_count", "generate_count must be between 1 and 20") + } + return value, nil +} + +func clonePromptRuleExtraParams(source map[string]interface{}) map[string]interface{} { + if len(source) == 0 { + return map[string]interface{}{} + } + + cloned := make(map[string]interface{}, len(source)) + for key, value := range source { + cloned[key] = value + } + return cloned +} + +func buildGenerateFromRuleResponse(items []GenerateFromRuleItem, requestedCount int) *GenerateFromRuleResponse { + responseItems := make([]GenerateFromRuleItem, 0, len(items)) + articleIDs := make([]int64, 0, len(items)) + taskIDs := make([]int64, 0, len(items)) + for _, item := range items { + responseItems = append(responseItems, item) + articleIDs = append(articleIDs, item.ArticleID) + taskIDs = append(taskIDs, item.TaskID) + } + + result := &GenerateFromRuleResponse{ + ArticleIDs: articleIDs, + TaskIDs: taskIDs, + Items: responseItems, + RequestedGenerateCount: requestedCount, + SuccessCount: len(items), + FailedCount: requestedCount - len(items), + } + if len(items) > 0 { + result.ArticleID = items[0].ArticleID + result.TaskID = items[0].TaskID + } + return result +} + func extractBool(extra map[string]interface{}, key string) bool { if len(extra) == 0 { return false diff --git a/server/internal/tenant/app/template_assist.go b/server/internal/tenant/app/template_assist.go index 228049d..59249c1 100644 --- a/server/internal/tenant/app/template_assist.go +++ b/server/internal/tenant/app/template_assist.go @@ -120,6 +120,8 @@ type assistJob struct { OutlineRequest *OutlineTaskRequest } +type AssistJob = assistJob + func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int64, req AnalyzeTaskRequest) (*AssistTaskCreateResponse, error) { actor := auth.MustActor(ctx) @@ -164,7 +166,7 @@ func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int6 AnalyzePrompt: extractAnalyzePromptTemplate(templateRecord.CardConfigJSON), AnalyzeRequest: &req, } - if err := s.enqueueAssist(job); err != nil { + if err := publishTemplateAssistTask(ctx, s.rabbitMQ, job); err != nil { _ = repo.FailTask(context.Background(), taskID, templateAssistTaskTypeAnalyze, actor.TenantID, err.Error(), time.Now()) return nil, response.ErrServiceUnavailable(50312, "analyze_queue_unavailable", "analyze queue is busy, please retry") } @@ -245,7 +247,7 @@ func (s *TemplateService) CreateTitleTask(ctx context.Context, templateID int64, TitlePrompt: extractTitlePromptTemplate(templateRecord.CardConfigJSON), TitleRequest: &req, } - if err := s.enqueueAssist(job); err != nil { + if err := publishTemplateAssistTask(ctx, s.rabbitMQ, job); err != nil { _ = repo.FailTask(context.Background(), taskID, templateAssistTaskTypeTitle, actor.TenantID, err.Error(), time.Now()) return nil, response.ErrServiceUnavailable(50313, "title_queue_unavailable", "title queue is busy, please retry") } @@ -326,7 +328,7 @@ func (s *TemplateService) CreateOutlineTask(ctx context.Context, templateID int6 OutlinePrompt: extractOutlinePromptTemplate(templateRecord.CardConfigJSON), OutlineRequest: &req, } - if err := s.enqueueAssist(job); err != nil { + if err := publishTemplateAssistTask(ctx, s.rabbitMQ, job); err != nil { _ = repo.FailTask(context.Background(), taskID, templateAssistTaskTypeOutline, actor.TenantID, err.Error(), time.Now()) return nil, response.ErrServiceUnavailable(50314, "outline_queue_unavailable", "outline queue is busy, please retry") } @@ -363,21 +365,6 @@ func (s *TemplateService) GetOutlineTask(ctx context.Context, templateID int64, return resp, nil } -func (s *TemplateService) enqueueAssist(job assistJob) error { - select { - case s.assistJobs <- job: - return nil - default: - return fmt.Errorf("assist queue is full") - } -} - -func (s *TemplateService) runAssistWorker() { - for job := range s.assistJobs { - s.executeAssist(context.Background(), job) - } -} - func (s *TemplateService) executeAssist(ctx context.Context, job assistJob) { repo := repository.NewTemplateAssistTaskRepository(s.pool) now := time.Now() @@ -395,6 +382,28 @@ func (s *TemplateService) executeAssist(ctx context.Context, job assistJob) { } } +func (s *TemplateService) ProcessQueuedAssist(ctx context.Context, job assistJob) { + s.executeAssist(ctx, job) +} + +func (s *TemplateService) ProcessQueuedAssistJob(ctx context.Context, job AssistJob) { + s.ProcessQueuedAssist(ctx, job) +} + +func (s *TemplateService) IsAssistTaskTerminal(ctx context.Context, job AssistJob) (bool, error) { + record, err := repository.NewTemplateAssistTaskRepository(s.pool).GetTask( + ctx, + job.TaskID, + job.TaskType, + job.TenantID, + job.TemplateID, + ) + if err != nil { + return false, err + } + return record.Status == "completed" || record.Status == "failed", nil +} + func (s *TemplateService) executeAnalyzeTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) { if job.AnalyzeRequest == nil { s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, fmt.Errorf("missing analyze request")) diff --git a/server/internal/tenant/app/template_assist_queue.go b/server/internal/tenant/app/template_assist_queue.go new file mode 100644 index 0000000..ca73d66 --- /dev/null +++ b/server/internal/tenant/app/template_assist_queue.go @@ -0,0 +1,44 @@ +package app + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" +) + +func publishTemplateAssistTask(ctx context.Context, client *rabbitmq.Client, job assistJob) error { + if client == nil { + return fmt.Errorf("rabbitmq client is not configured") + } + + payload, err := json.Marshal(job) + if err != nil { + return fmt.Errorf("marshal template assist job: %w", err) + } + + if err := client.PublishTemplateAssistTask(ctx, payload); err != nil { + return fmt.Errorf("publish template assist task: %w", err) + } + return nil +} + +func decodeTemplateAssistTask(payload []byte) (*assistJob, error) { + var job assistJob + if err := json.Unmarshal(payload, &job); err != nil { + return nil, fmt.Errorf("decode template assist job: %w", err) + } + if job.TaskID == "" || job.TaskType == "" || job.TenantID <= 0 || job.TemplateID <= 0 { + return nil, fmt.Errorf("template assist job is incomplete") + } + return &job, nil +} + +func DecodeTemplateAssistTask(payload []byte) (*AssistJob, error) { + job, err := decodeTemplateAssistTask(payload) + if err != nil { + return nil, err + } + return job, nil +} diff --git a/server/internal/tenant/app/template_service.go b/server/internal/tenant/app/template_service.go index 61d4b0c..9e8db80 100644 --- a/server/internal/tenant/app/template_service.go +++ b/server/internal/tenant/app/template_service.go @@ -13,6 +13,7 @@ import ( "github.com/geo-platform/tenant-api/internal/shared/auth" "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" @@ -22,13 +23,12 @@ type TemplateService struct { pool *pgxpool.Pool templates repository.TemplateRepository llm llm.Client + rabbitMQ *rabbitmq.Client knowledge *KnowledgeService streams *stream.GenerationHub streamEnabled bool articleTimeout time.Duration maxOutputTokens int64 - jobs chan generationJob - assistJobs chan assistJob } type generationJob struct { @@ -49,21 +49,12 @@ 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 { - queueSize := cfg.QueueSize - if queueSize <= 0 { - queueSize = 128 - } - - workerCount := cfg.WorkerConcurrency - if workerCount <= 0 { - workerCount = 1 - } - articleTimeout := cfg.ArticleTimeout if articleTimeout <= 0 { articleTimeout = 8 * time.Minute @@ -73,18 +64,12 @@ func NewTemplateService( pool: pool, templates: templates, llm: llmClient, + rabbitMQ: rabbitMQClient, knowledge: knowledge, streams: streams, streamEnabled: cfg.StreamEnabled, articleTimeout: articleTimeout, maxOutputTokens: llmMaxOutputTokens, - jobs: make(chan generationJob, queueSize), - assistJobs: make(chan assistJob, queueSize), - } - - for i := 0; i < workerCount; i++ { - go svc.runGenerationWorker() - go svc.runAssistWorker() } return svc @@ -272,7 +257,13 @@ func (s *TemplateService) Generate(ctx context.Context, templateID int64, req Ge return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan") } - inputJSON, _ := json.Marshal(req.InputParams) + 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 { @@ -369,26 +360,22 @@ func (s *TemplateService) Generate(ctx context.Context, templateID int64, req Ge } job := generationJob{ - OperatorID: actor.UserID, - TenantID: actor.TenantID, - ArticleID: articleID, - TaskID: taskID, - ReservationID: reservationID, - TemplateKey: templateRecord.TemplateKey, - TemplateName: templateRecord.TemplateName, - PromptTemplate: templateRecord.PromptTemplate, - Params: cloneInputParams(req.InputParams), - InitialTitle: initialTitle, - WebSearch: llm.WebSearchOptions{ - Enabled: req.EnableWebSearch, - }, - } - if req.WebSearchLimit != nil { - job.WebSearch.Limit = *req.WebSearchLimit + OperatorID: actor.UserID, + TenantID: actor.TenantID, + ArticleID: articleID, + TaskID: taskID, + ReservationID: reservationID, + Params: taskInputParams, + InitialTitle: initialTitle, } - if err := s.enqueueGeneration(job); err != nil { - s.failGeneration(context.Background(), job, job.InitialTitle, "queue_enqueue", err) + 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") } @@ -427,21 +414,6 @@ func updateGeneratedArticleState( return nil } -func (s *TemplateService) enqueueGeneration(job generationJob) error { - select { - case s.jobs <- job: - return nil - default: - return fmt.Errorf("generation queue is full") - } -} - -func (s *TemplateService) runGenerationWorker() { - for job := range s.jobs { - s.executeGeneration(context.Background(), job) - } -} - func (s *TemplateService) executeGeneration(ctx context.Context, job generationJob) { now := time.Now() title := strings.TrimSpace(job.InitialTitle) @@ -622,6 +594,8 @@ func formatGenerationFailure(stage string, genErr error) string { func generationFailureStageLabel(stage string) string { switch stage { + case "task_load": + return "加载生成任务失败" case "queue_enqueue": return "任务入队失败" case "llm_generate": diff --git a/server/internal/tenant/repository/audit_repo.go b/server/internal/tenant/repository/audit_repo.go index ccb281a..7dac2c9 100644 --- a/server/internal/tenant/repository/audit_repo.go +++ b/server/internal/tenant/repository/audit_repo.go @@ -70,9 +70,12 @@ func (r *auditRepository) CreateAuditLog(ctx context.Context, input AuditLogInpu } func (r *auditRepository) CreateGenerationTask(ctx context.Context, input GenerationTaskInput) (int64, error) { - operatorID := input.OperatorID + var operatorID *int64 + if input.OperatorID > 0 { + operatorID = &input.OperatorID + } return r.q.CreateGenerationTask(ctx, generated.CreateGenerationTaskParams{ - OperatorID: pgInt8(&operatorID), + OperatorID: pgInt8(operatorID), TenantID: input.TenantID, ArticleID: pgInt8(input.ArticleID), QuotaReservationID: pgInt8(input.QuotaReservationID), diff --git a/server/internal/tenant/repository/quota_repo.go b/server/internal/tenant/repository/quota_repo.go index 81cbf23..d2ec8af 100644 --- a/server/internal/tenant/repository/quota_repo.go +++ b/server/internal/tenant/repository/quota_repo.go @@ -57,10 +57,13 @@ func (r *quotaRepository) GetCurrentBalance(ctx context.Context, tenantID int64, } func (r *quotaRepository) InsertQuotaLedger(ctx context.Context, input QuotaLedgerInput) (int64, error) { - operatorID := input.OperatorID + var operatorID *int64 + if input.OperatorID > 0 { + operatorID = &input.OperatorID + } return r.q.InsertQuotaLedger(ctx, generated.InsertQuotaLedgerParams{ TenantID: input.TenantID, - OperatorID: pgInt8(&operatorID), + OperatorID: pgInt8(operatorID), QuotaType: input.QuotaType, Delta: int32(input.Delta), BalanceAfter: int32(input.BalanceAfter), @@ -72,10 +75,13 @@ func (r *quotaRepository) InsertQuotaLedger(ctx context.Context, input QuotaLedg func (r *quotaRepository) CreateQuotaReservation(ctx context.Context, input QuotaReservationInput) (int64, error) { expireAt := input.ExpireAt - operatorID := input.OperatorID + var operatorID *int64 + if input.OperatorID > 0 { + operatorID = &input.OperatorID + } return r.q.CreateQuotaReservation(ctx, generated.CreateQuotaReservationParams{ TenantID: input.TenantID, - OperatorID: pgInt8(&operatorID), + OperatorID: pgInt8(operatorID), QuotaType: input.QuotaType, ResourceType: input.ResourceType, ResourceID: pgInt8(input.ResourceID), diff --git a/server/internal/tenant/transport/article_handler.go b/server/internal/tenant/transport/article_handler.go index c8d06d2..59172c2 100644 --- a/server/internal/tenant/transport/article_handler.go +++ b/server/internal/tenant/transport/article_handler.go @@ -51,6 +51,7 @@ func NewArticleHandler(a *bootstrap.App) *ArticleHandler { promptGenerateSvc: app.NewPromptRuleGenerationService( a.DB, a.LLM, + a.RabbitMQ, knowledgeSvc, a.GenerationStreams, a.Config.Generation, @@ -369,6 +370,11 @@ func (h *ArticleHandler) Stream(c *gin.Context) { return } + fallback := buildGenerationSnapshot(detail) + if snapshot, ok := h.streams.Snapshot(id); !ok || (fallback.Done && snapshot.Status != fallback.Status) { + h.streams.Seed(fallback) + } + events, snapshot, cancel := h.streams.Subscribe(id) defer cancel() @@ -377,20 +383,10 @@ func (h *ArticleHandler) Stream(c *gin.Context) { return } flusher.Flush() + if snapshot.Done { + return + } } else { - fallback := stream.GenerationEvent{ - Type: "snapshot", - ArticleID: detail.ID, - Status: detail.GenerateStatus, - Done: detail.GenerateStatus == "completed" || detail.GenerateStatus == "failed", - UpdatedAt: time.Now().UTC(), - } - if detail.Title != nil { - fallback.Title = *detail.Title - } - if detail.MarkdownContent != nil { - fallback.Content = *detail.MarkdownContent - } if err := writeSSE(c, fallback.Type, fallback); err != nil { return } @@ -455,6 +451,32 @@ func writeSSE(c *gin.Context, event string, payload interface{}) error { return nil } +func buildGenerationSnapshot(detail *app.ArticleDetailResponse) stream.GenerationEvent { + event := stream.GenerationEvent{ + Type: "snapshot", + ArticleID: detail.ID, + Status: detail.GenerateStatus, + Done: detail.GenerateStatus == "completed" || detail.GenerateStatus == "failed", + UpdatedAt: time.Now().UTC(), + } + if detail.Title != nil { + event.Title = *detail.Title + } + if detail.MarkdownContent != nil { + event.Content = *detail.MarkdownContent + } + if detail.GenerateStatus == "completed" { + event.Type = "completed" + } + if detail.GenerateStatus == "failed" { + event.Type = "error" + if detail.GenerationErrorMessage != nil { + event.Error = *detail.GenerationErrorMessage + } + } + return event +} + func resolveOptimizeStreamErrorMessage(appErr *response.AppError) string { if appErr == nil { return "optimize selection failed" diff --git a/server/internal/tenant/transport/template_handler.go b/server/internal/tenant/transport/template_handler.go index 9b48d54..9ab3f3b 100644 --- a/server/internal/tenant/transport/template_handler.go +++ b/server/internal/tenant/transport/template_handler.go @@ -37,6 +37,7 @@ func NewTemplateHandler(a *bootstrap.App) *TemplateHandler { a.Cache, ), a.LLM, + a.RabbitMQ, knowledgeSvc, a.GenerationStreams, a.Config.Generation,