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

feat(generation): migrate article generation and template-assist to RabbitMQ queues

Browse Source 
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:
liangxu
2026-04-15 14:20:26 +08:00
parent e8f48c6d37
commit ce13331e26
14 changed files with 951 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/internal/shared/stream/generation.go
+204 -9
View File
@@ -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,
})
}
}