feat(generation): migrate article generation and template-assist to RabbitMQ queues
Replace in-process generation with queue-based async execution via RabbitMQ. Add generation task runtime for lifecycle management, article generation payload/queue abstractions, and template-assist queue publisher. Refactor prompt generation service to support batch generation with per-item error tracking. Update stream hub to bridge queue events to SSE. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -94,7 +94,7 @@ func New(configPath string) (*App, error) {
|
|||||||
vectorStore := retrieval.NewQdrantStore(cfg.Qdrant, logger)
|
vectorStore := retrieval.NewQdrantStore(cfg.Qdrant, logger)
|
||||||
objectStorageClient := objectstorage.New(cfg.ObjectStorage, logger)
|
objectStorageClient := objectstorage.New(cfg.ObjectStorage, logger)
|
||||||
auditLogs := auditlog.NewAsyncWriter(pool, logger)
|
auditLogs := auditlog.NewAsyncWriter(pool, logger)
|
||||||
generationStreams := stream.NewGenerationHub()
|
generationStreams := stream.NewGenerationHub(mqClient)
|
||||||
appCache := cache.New(cfg.Cache.Driver, rdb)
|
appCache := cache.New(cfg.Cache.Driver, rdb)
|
||||||
|
|
||||||
if cfg.Server.Mode == "release" {
|
if cfg.Server.Mode == "release" {
|
||||||
|
|||||||
@@ -1,8 +1,14 @@
|
|||||||
package stream
|
package stream
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
|
||||||
)
|
)
|
||||||
|
|
||||||
type GenerationEvent struct {
|
type GenerationEvent struct {
|
||||||
@@ -18,6 +24,11 @@ type GenerationEvent struct {
|
|||||||
UpdatedAt time.Time `json:"updated_at"`
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type generationEventEnvelope struct {
|
||||||
|
Source string `json:"source"`
|
||||||
|
Event GenerationEvent `json:"event"`
|
||||||
|
}
|
||||||
|
|
||||||
type generationStream struct {
|
type generationStream struct {
|
||||||
snapshot GenerationEvent
|
snapshot GenerationEvent
|
||||||
subscribers map[int]chan GenerationEvent
|
subscribers map[int]chan GenerationEvent
|
||||||
@@ -26,17 +37,42 @@ type generationStream struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type GenerationHub struct {
|
type GenerationHub struct {
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
streams map[int64]*generationStream
|
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{
|
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 {
|
h.publish(articleID, func(state *generationStream) GenerationEvent {
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
state.completed = false
|
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) {
|
func (h *GenerationHub) AppendDelta(articleID, taskID int64, title, delta string) {
|
||||||
h.publish(articleID, func(state *generationStream) GenerationEvent {
|
h.publish(articleID, func(state *generationStream) GenerationEvent {
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
@@ -69,6 +109,7 @@ func (h *GenerationHub) AppendDelta(articleID, taskID int64, title, delta string
|
|||||||
state.snapshot.Status = "generating"
|
state.snapshot.Status = "generating"
|
||||||
state.snapshot.Content += delta
|
state.snapshot.Content += delta
|
||||||
state.snapshot.UpdatedAt = now
|
state.snapshot.UpdatedAt = now
|
||||||
|
state.completed = false
|
||||||
|
|
||||||
return GenerationEvent{
|
return GenerationEvent{
|
||||||
Type: "delta",
|
Type: "delta",
|
||||||
@@ -188,14 +229,83 @@ func (h *GenerationHub) Snapshot(articleID int64) (*GenerationEvent, bool) {
|
|||||||
return &snapshot, true
|
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) {
|
func (h *GenerationHub) publish(articleID int64, mutate func(*generationStream) GenerationEvent) {
|
||||||
h.mu.Lock()
|
h.mu.Lock()
|
||||||
state := h.ensureLocked(articleID)
|
state := h.ensureLocked(articleID)
|
||||||
event := mutate(state)
|
event := mutate(state)
|
||||||
|
subscribers := collectSubscribersLocked(state)
|
||||||
subscribers := make([]chan GenerationEvent, 0, len(state.subscribers))
|
shouldDelete := state.completed && len(state.subscribers) == 0
|
||||||
for _, ch := range state.subscribers {
|
if shouldDelete {
|
||||||
subscribers = append(subscribers, ch)
|
delete(h.streams, articleID)
|
||||||
}
|
}
|
||||||
h.mu.Unlock()
|
h.mu.Unlock()
|
||||||
|
|
||||||
@@ -205,6 +315,16 @@ func (h *GenerationHub) publish(articleID int64, mutate func(*generationStream)
|
|||||||
default:
|
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 {
|
func (h *GenerationHub) ensureLocked(articleID int64) *generationStream {
|
||||||
@@ -219,3 +339,78 @@ func (h *GenerationHub) ensureLocked(articleID int64) *generationStream {
|
|||||||
h.streams[articleID] = state
|
h.streams[articleID] = state
|
||||||
return 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,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -219,7 +219,6 @@ func NewKnowledgeService(
|
|||||||
go svc.runParseWorker()
|
go svc.runParseWorker()
|
||||||
}
|
}
|
||||||
go svc.recoverPendingParseTasks()
|
go svc.recoverPendingParseTasks()
|
||||||
go svc.runDeletedCleanupWorker()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return svc
|
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() {
|
func (s *KnowledgeService) scheduleDeletedKnowledgeCleanup() {
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
@@ -1778,6 +1766,13 @@ func (s *KnowledgeService) scheduleDeletedKnowledgeCleanup() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (s *KnowledgeService) RunDeletedCleanupSweep() {
|
||||||
|
if s == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.scheduleDeletedKnowledgeCleanup()
|
||||||
|
}
|
||||||
|
|
||||||
func (s *KnowledgeService) listKnowledgeCleanupCandidates(
|
func (s *KnowledgeService) listKnowledgeCleanupCandidates(
|
||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
limit int,
|
limit int,
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ import (
|
|||||||
"github.com/geo-platform/tenant-api/internal/shared/auth"
|
"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/config"
|
||||||
"github.com/geo-platform/tenant-api/internal/shared/llm"
|
"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/response"
|
||||||
"github.com/geo-platform/tenant-api/internal/shared/stream"
|
"github.com/geo-platform/tenant-api/internal/shared/stream"
|
||||||
"github.com/geo-platform/tenant-api/internal/tenant/prompts"
|
"github.com/geo-platform/tenant-api/internal/tenant/prompts"
|
||||||
@@ -22,12 +23,12 @@ import (
|
|||||||
type PromptRuleGenerationService struct {
|
type PromptRuleGenerationService struct {
|
||||||
pool *pgxpool.Pool
|
pool *pgxpool.Pool
|
||||||
llm llm.Client
|
llm llm.Client
|
||||||
|
rabbitMQ *rabbitmq.Client
|
||||||
knowledge *KnowledgeService
|
knowledge *KnowledgeService
|
||||||
streams *stream.GenerationHub
|
streams *stream.GenerationHub
|
||||||
streamEnabled bool
|
streamEnabled bool
|
||||||
articleTimeout time.Duration
|
articleTimeout time.Duration
|
||||||
maxOutputTokens int64
|
maxOutputTokens int64
|
||||||
jobs chan promptRuleGenerationJob
|
|
||||||
}
|
}
|
||||||
|
|
||||||
type promptRuleGenerationJob struct {
|
type promptRuleGenerationJob struct {
|
||||||
@@ -50,10 +51,34 @@ type GenerateFromRuleRequest struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type GenerateFromRuleResponse 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"`
|
ArticleID int64 `json:"article_id"`
|
||||||
TaskID int64 `json:"task_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 {
|
type promptRuleGenerationRecord struct {
|
||||||
ID int64
|
ID int64
|
||||||
Name string
|
Name string
|
||||||
@@ -68,24 +93,25 @@ type promptRuleGenerationRecord struct {
|
|||||||
|
|
||||||
const promptRuleGeneratingTitlePlaceholder = "标题生成中"
|
const promptRuleGeneratingTitlePlaceholder = "标题生成中"
|
||||||
|
|
||||||
|
type enqueuePromptRuleGenerationInput struct {
|
||||||
|
OperatorID *int64
|
||||||
|
TenantID int64
|
||||||
|
PromptRuleID int64
|
||||||
|
BrandID *int64
|
||||||
|
TargetPlatform *string
|
||||||
|
ExtraParams map[string]interface{}
|
||||||
|
FallbackName string
|
||||||
|
}
|
||||||
|
|
||||||
func NewPromptRuleGenerationService(
|
func NewPromptRuleGenerationService(
|
||||||
pool *pgxpool.Pool,
|
pool *pgxpool.Pool,
|
||||||
llmClient llm.Client,
|
llmClient llm.Client,
|
||||||
|
rabbitMQClient *rabbitmq.Client,
|
||||||
knowledge *KnowledgeService,
|
knowledge *KnowledgeService,
|
||||||
streams *stream.GenerationHub,
|
streams *stream.GenerationHub,
|
||||||
cfg config.GenerationConfig,
|
cfg config.GenerationConfig,
|
||||||
llmMaxOutputTokens int64,
|
llmMaxOutputTokens int64,
|
||||||
) *PromptRuleGenerationService {
|
) *PromptRuleGenerationService {
|
||||||
queueSize := cfg.QueueSize
|
|
||||||
if queueSize <= 0 {
|
|
||||||
queueSize = 128
|
|
||||||
}
|
|
||||||
|
|
||||||
workerCount := cfg.WorkerConcurrency
|
|
||||||
if workerCount <= 0 {
|
|
||||||
workerCount = 1
|
|
||||||
}
|
|
||||||
|
|
||||||
articleTimeout := cfg.ArticleTimeout
|
articleTimeout := cfg.ArticleTimeout
|
||||||
if articleTimeout <= 0 {
|
if articleTimeout <= 0 {
|
||||||
articleTimeout = 8 * time.Minute
|
articleTimeout = 8 * time.Minute
|
||||||
@@ -94,16 +120,12 @@ func NewPromptRuleGenerationService(
|
|||||||
svc := &PromptRuleGenerationService{
|
svc := &PromptRuleGenerationService{
|
||||||
pool: pool,
|
pool: pool,
|
||||||
llm: llmClient,
|
llm: llmClient,
|
||||||
|
rabbitMQ: rabbitMQClient,
|
||||||
knowledge: knowledge,
|
knowledge: knowledge,
|
||||||
streams: streams,
|
streams: streams,
|
||||||
streamEnabled: cfg.StreamEnabled,
|
streamEnabled: cfg.StreamEnabled,
|
||||||
articleTimeout: articleTimeout,
|
articleTimeout: articleTimeout,
|
||||||
maxOutputTokens: llmMaxOutputTokens,
|
maxOutputTokens: llmMaxOutputTokens,
|
||||||
jobs: make(chan promptRuleGenerationJob, queueSize),
|
|
||||||
}
|
|
||||||
|
|
||||||
for i := 0; i < workerCount; i++ {
|
|
||||||
go svc.runGenerationWorker()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return svc
|
return svc
|
||||||
@@ -111,12 +133,86 @@ func NewPromptRuleGenerationService(
|
|||||||
|
|
||||||
func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req GenerateFromRuleRequest) (*GenerateFromRuleResponse, error) {
|
func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req GenerateFromRuleRequest) (*GenerateFromRuleResponse, error) {
|
||||||
actor := auth.MustActor(ctx)
|
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 {
|
if err := s.llm.Validate(); err != nil {
|
||||||
return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error())
|
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 {
|
if err != nil {
|
||||||
return nil, err
|
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")
|
return nil, response.ErrBadRequest(40017, "prompt_rule_disabled", "prompt rule is disabled")
|
||||||
}
|
}
|
||||||
|
|
||||||
extra := req.ExtraParams
|
extra := input.ExtraParams
|
||||||
if extractGenerateCount(extra) > 1 {
|
generationMode := strings.TrimSpace(extractString(extra, "generation_mode"))
|
||||||
return nil, response.ErrBadRequest(40018, "generate_count_not_supported", "prompt rule instant generation currently supports generate_count = 1 only")
|
generateCount, err := normalizePromptRuleGenerateCount(extractGenerateCount(extra))
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
taskName := strings.TrimSpace(extractString(extra, "task_name"))
|
taskName := strings.TrimSpace(extractString(extra, "task_name"))
|
||||||
@@ -134,34 +232,46 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
|
|||||||
taskName = strings.TrimSpace(rule.Name)
|
taskName = strings.TrimSpace(rule.Name)
|
||||||
}
|
}
|
||||||
if taskName == "" {
|
if taskName == "" {
|
||||||
taskName = "即时任务"
|
taskName = strings.TrimSpace(input.FallbackName)
|
||||||
|
}
|
||||||
|
if taskName == "" {
|
||||||
|
taskName = "内容任务"
|
||||||
}
|
}
|
||||||
|
|
||||||
quotaRepo := repository.NewQuotaRepository(s.pool)
|
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 {
|
if err != nil || balance < 1 {
|
||||||
return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan")
|
return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan")
|
||||||
}
|
}
|
||||||
|
|
||||||
inputParams := map[string]interface{}{
|
inputParams := map[string]interface{}{
|
||||||
"prompt_rule_id": req.PromptRuleID,
|
"prompt_rule_id": input.PromptRuleID,
|
||||||
"task_name": taskName,
|
"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"),
|
"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
|
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 {
|
if len(platformIDs) > 0 {
|
||||||
inputParams["target_platforms"] = platformIDs
|
inputParams["target_platforms"] = platformIDs
|
||||||
}
|
}
|
||||||
if req.BrandID != nil {
|
if input.BrandID != nil {
|
||||||
inputParams["brand_id"] = *req.BrandID
|
inputParams["brand_id"] = *input.BrandID
|
||||||
}
|
}
|
||||||
if len(rule.KnowledgeGroupIDs) > 0 {
|
if len(rule.KnowledgeGroupIDs) > 0 {
|
||||||
inputParams["knowledge_group_ids"] = rule.KnowledgeGroupIDs
|
inputParams["knowledge_group_ids"] = rule.KnowledgeGroupIDs
|
||||||
}
|
}
|
||||||
|
inputParams = attachPromptRuleSnapshot(inputParams, rule)
|
||||||
|
|
||||||
knowledgePrompt := ""
|
knowledgePrompt := ""
|
||||||
if len(rule.KnowledgeGroupIDs) > 0 {
|
if len(rule.KnowledgeGroupIDs) > 0 {
|
||||||
@@ -170,7 +280,7 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
|
|||||||
}
|
}
|
||||||
knowledgeContext, resolveErr := s.knowledge.ResolveContext(
|
knowledgeContext, resolveErr := s.knowledge.ResolveContext(
|
||||||
ctx,
|
ctx,
|
||||||
actor.TenantID,
|
input.TenantID,
|
||||||
rule.KnowledgeGroupIDs,
|
rule.KnowledgeGroupIDs,
|
||||||
buildPromptRuleKnowledgeQuery(rule, inputParams),
|
buildPromptRuleKnowledgeQuery(rule, inputParams),
|
||||||
)
|
)
|
||||||
@@ -178,6 +288,9 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
|
|||||||
return nil, resolveErr
|
return nil, resolveErr
|
||||||
}
|
}
|
||||||
knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext)
|
knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext)
|
||||||
|
if knowledgePrompt != "" {
|
||||||
|
inputParams[generationPayloadKnowledgePrompt] = knowledgePrompt
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
inputJSON, _ := json.Marshal(inputParams)
|
inputJSON, _ := json.Marshal(inputParams)
|
||||||
@@ -193,13 +306,17 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
|
|||||||
|
|
||||||
quotaTx := repository.NewQuotaRepository(tx)
|
quotaTx := repository.NewQuotaRepository(tx)
|
||||||
auditTx := repository.NewAuditRepository(tx)
|
auditTx := repository.NewAuditRepository(tx)
|
||||||
|
operatorID := int64(0)
|
||||||
|
if input.OperatorID != nil {
|
||||||
|
operatorID = *input.OperatorID
|
||||||
|
}
|
||||||
|
|
||||||
newBalance := balance - 1
|
newBalance := balance - 1
|
||||||
reason := "generation_reserve"
|
reason := "generation_reserve"
|
||||||
referenceType := "generation_task"
|
referenceType := "generation_task"
|
||||||
_, err = quotaTx.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
_, err = quotaTx.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
||||||
TenantID: actor.TenantID,
|
TenantID: input.TenantID,
|
||||||
OperatorID: actor.UserID,
|
OperatorID: operatorID,
|
||||||
QuotaType: "article_generation",
|
QuotaType: "article_generation",
|
||||||
Delta: -1,
|
Delta: -1,
|
||||||
BalanceAfter: newBalance,
|
BalanceAfter: newBalance,
|
||||||
@@ -211,8 +328,8 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
|
|||||||
}
|
}
|
||||||
|
|
||||||
reservationID, err := quotaTx.CreateQuotaReservation(ctx, repository.QuotaReservationInput{
|
reservationID, err := quotaTx.CreateQuotaReservation(ctx, repository.QuotaReservationInput{
|
||||||
TenantID: actor.TenantID,
|
TenantID: input.TenantID,
|
||||||
OperatorID: actor.UserID,
|
OperatorID: operatorID,
|
||||||
QuotaType: "article_generation",
|
QuotaType: "article_generation",
|
||||||
ResourceType: "article",
|
ResourceType: "article",
|
||||||
ResourceID: nil,
|
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)
|
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)
|
VALUES ($1, 'custom_generation', $2, 'generating', 'unpublished', $3)
|
||||||
RETURNING id
|
RETURNING id
|
||||||
`, actor.TenantID, req.PromptRuleID, wizardStateJSON).Scan(&articleID)
|
`, input.TenantID, input.PromptRuleID, wizardStateJSON).Scan(&articleID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, response.ErrInternal(50012, "create_failed", "failed to create prompt rule article")
|
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")
|
return nil, response.ErrInternal(50012, "create_failed", "failed to update quota reservation resource")
|
||||||
}
|
}
|
||||||
|
|
||||||
taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{
|
taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{
|
||||||
OperatorID: actor.UserID,
|
OperatorID: operatorID,
|
||||||
TenantID: actor.TenantID,
|
TenantID: input.TenantID,
|
||||||
ArticleID: &articleID,
|
ArticleID: &articleID,
|
||||||
QuotaReservationID: &reservationID,
|
QuotaReservationID: &reservationID,
|
||||||
TaskType: "custom_generation",
|
TaskType: "custom_generation",
|
||||||
@@ -254,25 +371,26 @@ func (s *PromptRuleGenerationService) GenerateFromRule(ctx context.Context, req
|
|||||||
}
|
}
|
||||||
|
|
||||||
job := promptRuleGenerationJob{
|
job := promptRuleGenerationJob{
|
||||||
OperatorID: actor.UserID,
|
OperatorID: operatorID,
|
||||||
TenantID: actor.TenantID,
|
TenantID: input.TenantID,
|
||||||
ArticleID: articleID,
|
ArticleID: articleID,
|
||||||
TaskID: taskID,
|
TaskID: taskID,
|
||||||
ReservationID: reservationID,
|
ReservationID: reservationID,
|
||||||
Prompt: buildPromptRuleGenerationPrompt(rule, inputParams, knowledgePrompt),
|
|
||||||
InputParams: inputParams,
|
InputParams: inputParams,
|
||||||
InitialTitle: promptRuleGeneratingTitlePlaceholder,
|
InitialTitle: promptRuleGeneratingTitlePlaceholder,
|
||||||
WebSearch: llm.WebSearchOptions{
|
|
||||||
Enabled: extractBool(extra, "enable_web_search"),
|
|
||||||
},
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := s.enqueueGeneration(job); err != nil {
|
if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{
|
||||||
s.failGeneration(context.Background(), job, job.InitialTitle, "queue_enqueue", err)
|
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")
|
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)
|
s.streams.Start(articleID, taskID, promptRuleGeneratingTitlePlaceholder)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -322,21 +440,6 @@ func (s *PromptRuleGenerationService) loadPromptRule(ctx context.Context, tenant
|
|||||||
return record, nil
|
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) {
|
func (s *PromptRuleGenerationService) executeGeneration(ctx context.Context, job promptRuleGenerationJob) {
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
title := strings.TrimSpace(job.InitialTitle)
|
title := strings.TrimSpace(job.InitialTitle)
|
||||||
@@ -557,6 +660,50 @@ func extractGenerateCount(extra map[string]interface{}) int {
|
|||||||
return value
|
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 {
|
func extractBool(extra map[string]interface{}, key string) bool {
|
||||||
if len(extra) == 0 {
|
if len(extra) == 0 {
|
||||||
return false
|
return false
|
||||||
|
|||||||
@@ -120,6 +120,8 @@ type assistJob struct {
|
|||||||
OutlineRequest *OutlineTaskRequest
|
OutlineRequest *OutlineTaskRequest
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type AssistJob = assistJob
|
||||||
|
|
||||||
func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int64, req AnalyzeTaskRequest) (*AssistTaskCreateResponse, error) {
|
func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int64, req AnalyzeTaskRequest) (*AssistTaskCreateResponse, error) {
|
||||||
actor := auth.MustActor(ctx)
|
actor := auth.MustActor(ctx)
|
||||||
|
|
||||||
@@ -164,7 +166,7 @@ func (s *TemplateService) CreateAnalyzeTask(ctx context.Context, templateID int6
|
|||||||
AnalyzePrompt: extractAnalyzePromptTemplate(templateRecord.CardConfigJSON),
|
AnalyzePrompt: extractAnalyzePromptTemplate(templateRecord.CardConfigJSON),
|
||||||
AnalyzeRequest: &req,
|
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())
|
_ = 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")
|
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),
|
TitlePrompt: extractTitlePromptTemplate(templateRecord.CardConfigJSON),
|
||||||
TitleRequest: &req,
|
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())
|
_ = 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")
|
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),
|
OutlinePrompt: extractOutlinePromptTemplate(templateRecord.CardConfigJSON),
|
||||||
OutlineRequest: &req,
|
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())
|
_ = 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")
|
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
|
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) {
|
func (s *TemplateService) executeAssist(ctx context.Context, job assistJob) {
|
||||||
repo := repository.NewTemplateAssistTaskRepository(s.pool)
|
repo := repository.NewTemplateAssistTaskRepository(s.pool)
|
||||||
now := time.Now()
|
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) {
|
func (s *TemplateService) executeAnalyzeTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) {
|
||||||
if job.AnalyzeRequest == nil {
|
if job.AnalyzeRequest == nil {
|
||||||
s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, fmt.Errorf("missing analyze request"))
|
s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, fmt.Errorf("missing analyze request"))
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -13,6 +13,7 @@ import (
|
|||||||
"github.com/geo-platform/tenant-api/internal/shared/auth"
|
"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/config"
|
||||||
"github.com/geo-platform/tenant-api/internal/shared/llm"
|
"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/response"
|
||||||
"github.com/geo-platform/tenant-api/internal/shared/stream"
|
"github.com/geo-platform/tenant-api/internal/shared/stream"
|
||||||
"github.com/geo-platform/tenant-api/internal/tenant/repository"
|
"github.com/geo-platform/tenant-api/internal/tenant/repository"
|
||||||
@@ -22,13 +23,12 @@ type TemplateService struct {
|
|||||||
pool *pgxpool.Pool
|
pool *pgxpool.Pool
|
||||||
templates repository.TemplateRepository
|
templates repository.TemplateRepository
|
||||||
llm llm.Client
|
llm llm.Client
|
||||||
|
rabbitMQ *rabbitmq.Client
|
||||||
knowledge *KnowledgeService
|
knowledge *KnowledgeService
|
||||||
streams *stream.GenerationHub
|
streams *stream.GenerationHub
|
||||||
streamEnabled bool
|
streamEnabled bool
|
||||||
articleTimeout time.Duration
|
articleTimeout time.Duration
|
||||||
maxOutputTokens int64
|
maxOutputTokens int64
|
||||||
jobs chan generationJob
|
|
||||||
assistJobs chan assistJob
|
|
||||||
}
|
}
|
||||||
|
|
||||||
type generationJob struct {
|
type generationJob struct {
|
||||||
@@ -49,21 +49,12 @@ func NewTemplateService(
|
|||||||
pool *pgxpool.Pool,
|
pool *pgxpool.Pool,
|
||||||
templates repository.TemplateRepository,
|
templates repository.TemplateRepository,
|
||||||
llmClient llm.Client,
|
llmClient llm.Client,
|
||||||
|
rabbitMQClient *rabbitmq.Client,
|
||||||
knowledge *KnowledgeService,
|
knowledge *KnowledgeService,
|
||||||
streams *stream.GenerationHub,
|
streams *stream.GenerationHub,
|
||||||
cfg config.GenerationConfig,
|
cfg config.GenerationConfig,
|
||||||
llmMaxOutputTokens int64,
|
llmMaxOutputTokens int64,
|
||||||
) *TemplateService {
|
) *TemplateService {
|
||||||
queueSize := cfg.QueueSize
|
|
||||||
if queueSize <= 0 {
|
|
||||||
queueSize = 128
|
|
||||||
}
|
|
||||||
|
|
||||||
workerCount := cfg.WorkerConcurrency
|
|
||||||
if workerCount <= 0 {
|
|
||||||
workerCount = 1
|
|
||||||
}
|
|
||||||
|
|
||||||
articleTimeout := cfg.ArticleTimeout
|
articleTimeout := cfg.ArticleTimeout
|
||||||
if articleTimeout <= 0 {
|
if articleTimeout <= 0 {
|
||||||
articleTimeout = 8 * time.Minute
|
articleTimeout = 8 * time.Minute
|
||||||
@@ -73,18 +64,12 @@ func NewTemplateService(
|
|||||||
pool: pool,
|
pool: pool,
|
||||||
templates: templates,
|
templates: templates,
|
||||||
llm: llmClient,
|
llm: llmClient,
|
||||||
|
rabbitMQ: rabbitMQClient,
|
||||||
knowledge: knowledge,
|
knowledge: knowledge,
|
||||||
streams: streams,
|
streams: streams,
|
||||||
streamEnabled: cfg.StreamEnabled,
|
streamEnabled: cfg.StreamEnabled,
|
||||||
articleTimeout: articleTimeout,
|
articleTimeout: articleTimeout,
|
||||||
maxOutputTokens: llmMaxOutputTokens,
|
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
|
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")
|
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)
|
tx, err := s.pool.Begin(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -369,26 +360,22 @@ func (s *TemplateService) Generate(ctx context.Context, templateID int64, req Ge
|
|||||||
}
|
}
|
||||||
|
|
||||||
job := generationJob{
|
job := generationJob{
|
||||||
OperatorID: actor.UserID,
|
OperatorID: actor.UserID,
|
||||||
TenantID: actor.TenantID,
|
TenantID: actor.TenantID,
|
||||||
ArticleID: articleID,
|
ArticleID: articleID,
|
||||||
TaskID: taskID,
|
TaskID: taskID,
|
||||||
ReservationID: reservationID,
|
ReservationID: reservationID,
|
||||||
TemplateKey: templateRecord.TemplateKey,
|
Params: taskInputParams,
|
||||||
TemplateName: templateRecord.TemplateName,
|
InitialTitle: initialTitle,
|
||||||
PromptTemplate: templateRecord.PromptTemplate,
|
|
||||||
Params: cloneInputParams(req.InputParams),
|
|
||||||
InitialTitle: initialTitle,
|
|
||||||
WebSearch: llm.WebSearchOptions{
|
|
||||||
Enabled: req.EnableWebSearch,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
if req.WebSearchLimit != nil {
|
|
||||||
job.WebSearch.Limit = *req.WebSearchLimit
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := s.enqueueGeneration(job); err != nil {
|
if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{
|
||||||
s.failGeneration(context.Background(), job, job.InitialTitle, "queue_enqueue", err)
|
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")
|
return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry")
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -427,21 +414,6 @@ func updateGeneratedArticleState(
|
|||||||
return nil
|
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) {
|
func (s *TemplateService) executeGeneration(ctx context.Context, job generationJob) {
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
title := strings.TrimSpace(job.InitialTitle)
|
title := strings.TrimSpace(job.InitialTitle)
|
||||||
@@ -622,6 +594,8 @@ func formatGenerationFailure(stage string, genErr error) string {
|
|||||||
|
|
||||||
func generationFailureStageLabel(stage string) string {
|
func generationFailureStageLabel(stage string) string {
|
||||||
switch stage {
|
switch stage {
|
||||||
|
case "task_load":
|
||||||
|
return "加载生成任务失败"
|
||||||
case "queue_enqueue":
|
case "queue_enqueue":
|
||||||
return "任务入队失败"
|
return "任务入队失败"
|
||||||
case "llm_generate":
|
case "llm_generate":
|
||||||
|
|||||||
@@ -70,9 +70,12 @@ func (r *auditRepository) CreateAuditLog(ctx context.Context, input AuditLogInpu
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (r *auditRepository) CreateGenerationTask(ctx context.Context, input GenerationTaskInput) (int64, error) {
|
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{
|
return r.q.CreateGenerationTask(ctx, generated.CreateGenerationTaskParams{
|
||||||
OperatorID: pgInt8(&operatorID),
|
OperatorID: pgInt8(operatorID),
|
||||||
TenantID: input.TenantID,
|
TenantID: input.TenantID,
|
||||||
ArticleID: pgInt8(input.ArticleID),
|
ArticleID: pgInt8(input.ArticleID),
|
||||||
QuotaReservationID: pgInt8(input.QuotaReservationID),
|
QuotaReservationID: pgInt8(input.QuotaReservationID),
|
||||||
|
|||||||
@@ -57,10 +57,13 @@ func (r *quotaRepository) GetCurrentBalance(ctx context.Context, tenantID int64,
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (r *quotaRepository) InsertQuotaLedger(ctx context.Context, input QuotaLedgerInput) (int64, error) {
|
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{
|
return r.q.InsertQuotaLedger(ctx, generated.InsertQuotaLedgerParams{
|
||||||
TenantID: input.TenantID,
|
TenantID: input.TenantID,
|
||||||
OperatorID: pgInt8(&operatorID),
|
OperatorID: pgInt8(operatorID),
|
||||||
QuotaType: input.QuotaType,
|
QuotaType: input.QuotaType,
|
||||||
Delta: int32(input.Delta),
|
Delta: int32(input.Delta),
|
||||||
BalanceAfter: int32(input.BalanceAfter),
|
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) {
|
func (r *quotaRepository) CreateQuotaReservation(ctx context.Context, input QuotaReservationInput) (int64, error) {
|
||||||
expireAt := input.ExpireAt
|
expireAt := input.ExpireAt
|
||||||
operatorID := input.OperatorID
|
var operatorID *int64
|
||||||
|
if input.OperatorID > 0 {
|
||||||
|
operatorID = &input.OperatorID
|
||||||
|
}
|
||||||
return r.q.CreateQuotaReservation(ctx, generated.CreateQuotaReservationParams{
|
return r.q.CreateQuotaReservation(ctx, generated.CreateQuotaReservationParams{
|
||||||
TenantID: input.TenantID,
|
TenantID: input.TenantID,
|
||||||
OperatorID: pgInt8(&operatorID),
|
OperatorID: pgInt8(operatorID),
|
||||||
QuotaType: input.QuotaType,
|
QuotaType: input.QuotaType,
|
||||||
ResourceType: input.ResourceType,
|
ResourceType: input.ResourceType,
|
||||||
ResourceID: pgInt8(input.ResourceID),
|
ResourceID: pgInt8(input.ResourceID),
|
||||||
|
|||||||
@@ -51,6 +51,7 @@ func NewArticleHandler(a *bootstrap.App) *ArticleHandler {
|
|||||||
promptGenerateSvc: app.NewPromptRuleGenerationService(
|
promptGenerateSvc: app.NewPromptRuleGenerationService(
|
||||||
a.DB,
|
a.DB,
|
||||||
a.LLM,
|
a.LLM,
|
||||||
|
a.RabbitMQ,
|
||||||
knowledgeSvc,
|
knowledgeSvc,
|
||||||
a.GenerationStreams,
|
a.GenerationStreams,
|
||||||
a.Config.Generation,
|
a.Config.Generation,
|
||||||
@@ -369,6 +370,11 @@ func (h *ArticleHandler) Stream(c *gin.Context) {
|
|||||||
return
|
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)
|
events, snapshot, cancel := h.streams.Subscribe(id)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
|
||||||
@@ -377,20 +383,10 @@ func (h *ArticleHandler) Stream(c *gin.Context) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
flusher.Flush()
|
flusher.Flush()
|
||||||
|
if snapshot.Done {
|
||||||
|
return
|
||||||
|
}
|
||||||
} else {
|
} 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 {
|
if err := writeSSE(c, fallback.Type, fallback); err != nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
@@ -455,6 +451,32 @@ func writeSSE(c *gin.Context, event string, payload interface{}) error {
|
|||||||
return nil
|
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 {
|
func resolveOptimizeStreamErrorMessage(appErr *response.AppError) string {
|
||||||
if appErr == nil {
|
if appErr == nil {
|
||||||
return "optimize selection failed"
|
return "optimize selection failed"
|
||||||
|
|||||||
@@ -37,6 +37,7 @@ func NewTemplateHandler(a *bootstrap.App) *TemplateHandler {
|
|||||||
a.Cache,
|
a.Cache,
|
||||||
),
|
),
|
||||||
a.LLM,
|
a.LLM,
|
||||||
|
a.RabbitMQ,
|
||||||
knowledgeSvc,
|
knowledgeSvc,
|
||||||
a.GenerationStreams,
|
a.GenerationStreams,
|
||||||
a.Config.Generation,
|
a.Config.Generation,
|
||||||
|
|||||||
Reference in New Issue
Block a user