From d55f1276f30eac039fe3b13fff8a84d19d7eedef Mon Sep 17 00:00:00 2001 From: liangxu Date: Fri, 17 Apr 2026 13:52:01 +0800 Subject: [PATCH] feat(kol): AI assist service + worker --- server/cmd/kol-assist-worker/main.go | 45 +++ server/configs/config.yaml | 6 + server/internal/bootstrap/bootstrap.go | 3 + server/internal/shared/config/config.go | 24 ++ .../shared/messaging/rabbitmq/client.go | 20 ++ .../internal/tenant/app/kol_assist_queue.go | 51 ++++ .../internal/tenant/app/kol_assist_service.go | 178 ++++++++++++ .../worker/assist/kol_assist_worker.go | 266 ++++++++++++++++++ 8 files changed, 593 insertions(+) create mode 100644 server/cmd/kol-assist-worker/main.go create mode 100644 server/internal/tenant/app/kol_assist_queue.go create mode 100644 server/internal/tenant/app/kol_assist_service.go create mode 100644 server/internal/worker/assist/kol_assist_worker.go diff --git a/server/cmd/kol-assist-worker/main.go b/server/cmd/kol-assist-worker/main.go new file mode 100644 index 0000000..630a91d --- /dev/null +++ b/server/cmd/kol-assist-worker/main.go @@ -0,0 +1,45 @@ +package main + +import ( + "context" + "log" + "os" + "os/signal" + "syscall" + + "github.com/geo-platform/tenant-api/internal/bootstrap" + assistworker "github.com/geo-platform/tenant-api/internal/worker/assist" +) + +func main() { + configPath := "configs/config.yaml" + if p := os.Getenv("CONFIG_PATH"); p != "" { + configPath = p + } + + app, err := bootstrap.New(configPath) + if err != nil { + log.Fatalf("bootstrap: %v", err) + } + defer app.Close() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + assistworker.NewKolAssistWorker( + app.RabbitMQ, + app.KolAssists, + app.LLM, + app.Logger, + app.Config.Generation.WorkerConcurrency, + app.Config.Generation.ArticleTimeout, + ).Start(ctx) + + app.Logger.Info("kol-assist-worker started") + + quit := make(chan os.Signal, 1) + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + <-quit + cancel() + app.Logger.Info("kol-assist-worker shutting down...") +} diff --git a/server/configs/config.yaml b/server/configs/config.yaml index e2d7a1a..5981dff 100644 --- a/server/configs/config.yaml +++ b/server/configs/config.yaml @@ -49,6 +49,12 @@ rabbitmq: template_assist_dlx: template.assist.dlx template_assist_dlq: template.assist.run.dlq template_assist_dlq_routing_key: template.assist.run.dlq + kol_assist_exchange: kol.assist + kol_assist_routing_key: kol.assist.run + kol_assist_queue: kol.assist.run + kol_assist_dlx: kol.assist.dlx + kol_assist_dlq: kol.assist.run.dlq + kol_assist_dlq_routing_key: kol.assist.run.dlq consumer_prefetch: 10 publish_channel_pool_size: 32 diff --git a/server/internal/bootstrap/bootstrap.go b/server/internal/bootstrap/bootstrap.go index 2fb176e..e914318 100644 --- a/server/internal/bootstrap/bootstrap.go +++ b/server/internal/bootstrap/bootstrap.go @@ -50,6 +50,7 @@ type App struct { KolProfiles repository.KolProfileRepository KolPackages repository.KolPackageRepository KolPrompts repository.KolPromptRepository + KolAssists repository.KolAssistRepository KolPromptAsset *tenantapp.KolPromptAsset } @@ -105,6 +106,7 @@ func New(configPath string) (*App, error) { kolProfiles := repository.NewKolProfileRepository(pool) kolPackages := repository.NewKolPackageRepository(pool) kolPrompts := repository.NewKolPromptRepository(pool) + kolAssists := repository.NewKolAssistRepository(pool) kolPromptAsset := tenantapp.NewKolPromptAsset(objectStorageClient) if cfg.Server.Mode == "release" { @@ -162,6 +164,7 @@ func New(configPath string) (*App, error) { KolProfiles: kolProfiles, KolPackages: kolPackages, KolPrompts: kolPrompts, + KolAssists: kolAssists, KolPromptAsset: kolPromptAsset, }, nil } diff --git a/server/internal/shared/config/config.go b/server/internal/shared/config/config.go index 5d1704d..08a9e87 100644 --- a/server/internal/shared/config/config.go +++ b/server/internal/shared/config/config.go @@ -81,6 +81,12 @@ type RabbitMQConfig struct { TemplateAssistDLX string `mapstructure:"template_assist_dlx"` TemplateAssistDLQ string `mapstructure:"template_assist_dlq"` TemplateAssistDLQRouteKey string `mapstructure:"template_assist_dlq_routing_key"` + KolAssistExchange string `mapstructure:"kol_assist_exchange"` + KolAssistRoutingKey string `mapstructure:"kol_assist_routing_key"` + KolAssistQueue string `mapstructure:"kol_assist_queue"` + KolAssistDLX string `mapstructure:"kol_assist_dlx"` + KolAssistDLQ string `mapstructure:"kol_assist_dlq"` + KolAssistDLQRouteKey string `mapstructure:"kol_assist_dlq_routing_key"` ImageAssetExchange string `mapstructure:"image_asset_exchange"` ImageAssetRoutingKey string `mapstructure:"image_asset_routing_key"` ImageAssetQueue string `mapstructure:"image_asset_queue"` @@ -359,6 +365,24 @@ func normalizeRabbitMQConfig(cfg *RabbitMQConfig) { if strings.TrimSpace(cfg.TemplateAssistDLQRouteKey) == "" { cfg.TemplateAssistDLQRouteKey = "template.assist.run.dlq" } + if strings.TrimSpace(cfg.KolAssistExchange) == "" { + cfg.KolAssistExchange = "kol.assist" + } + if strings.TrimSpace(cfg.KolAssistRoutingKey) == "" { + cfg.KolAssistRoutingKey = "kol.assist.run" + } + if strings.TrimSpace(cfg.KolAssistQueue) == "" { + cfg.KolAssistQueue = "kol.assist.run" + } + if strings.TrimSpace(cfg.KolAssistDLX) == "" { + cfg.KolAssistDLX = "kol.assist.dlx" + } + if strings.TrimSpace(cfg.KolAssistDLQ) == "" { + cfg.KolAssistDLQ = "kol.assist.run.dlq" + } + if strings.TrimSpace(cfg.KolAssistDLQRouteKey) == "" { + cfg.KolAssistDLQRouteKey = "kol.assist.run.dlq" + } if strings.TrimSpace(cfg.ImageAssetExchange) == "" { cfg.ImageAssetExchange = "image.asset" } diff --git a/server/internal/shared/messaging/rabbitmq/client.go b/server/internal/shared/messaging/rabbitmq/client.go index c1c95d1..d456078 100644 --- a/server/internal/shared/messaging/rabbitmq/client.go +++ b/server/internal/shared/messaging/rabbitmq/client.go @@ -102,6 +102,10 @@ func (c *Client) PublishTemplateAssistTask(ctx context.Context, body []byte) err return c.publish(ctx, c.cfg.TemplateAssistExchange, c.cfg.TemplateAssistRoutingKey, body) } +func (c *Client) PublishKolAssistTask(ctx context.Context, body []byte) error { + return c.publish(ctx, c.cfg.KolAssistExchange, c.cfg.KolAssistRoutingKey, body) +} + func (c *Client) PublishImageAssetTask(ctx context.Context, body []byte) error { return c.publish(ctx, c.cfg.ImageAssetExchange, c.cfg.ImageAssetRoutingKey, body) } @@ -180,6 +184,10 @@ func (c *Client) ConsumeTemplateAssistTask(consumerName string) (<-chan amqp.Del return c.consume(consumerName, c.cfg.TemplateAssistQueue) } +func (c *Client) ConsumeKolAssistTask(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) { + return c.consume(consumerName, c.cfg.KolAssistQueue) +} + func (c *Client) ConsumeImageAssetTask(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) { return c.consume(consumerName, c.cfg.ImageAssetQueue) } @@ -562,6 +570,18 @@ func (c *Client) ensureTopology(conn *amqp.Connection) error { return err } + if err := c.ensureWorkQueue( + ch, + c.cfg.KolAssistExchange, + c.cfg.KolAssistRoutingKey, + c.cfg.KolAssistQueue, + c.cfg.KolAssistDLX, + c.cfg.KolAssistDLQ, + c.cfg.KolAssistDLQRouteKey, + ); err != nil { + return err + } + if err := c.ensureWorkQueue( ch, c.cfg.ImageAssetExchange, diff --git a/server/internal/tenant/app/kol_assist_queue.go b/server/internal/tenant/app/kol_assist_queue.go new file mode 100644 index 0000000..10ce8c2 --- /dev/null +++ b/server/internal/tenant/app/kol_assist_queue.go @@ -0,0 +1,51 @@ +package app + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" +) + +type kolAssistJob struct { + TaskID string `json:"task_id"` + TenantID int64 `json:"tenant_id"` +} + +type KolAssistJob = kolAssistJob + +func publishKolAssistTask(ctx context.Context, client *rabbitmq.Client, job kolAssistJob) error { + if client == nil { + return fmt.Errorf("rabbitmq client is not configured") + } + + payload, err := json.Marshal(job) + if err != nil { + return fmt.Errorf("marshal kol assist job: %w", err) + } + + if err := client.PublishKolAssistTask(ctx, payload); err != nil { + return fmt.Errorf("publish kol assist task: %w", err) + } + return nil +} + +func decodeKolAssistTask(payload []byte) (*kolAssistJob, error) { + var job kolAssistJob + if err := json.Unmarshal(payload, &job); err != nil { + return nil, fmt.Errorf("decode kol assist job: %w", err) + } + if job.TaskID == "" || job.TenantID <= 0 { + return nil, fmt.Errorf("kol assist job is incomplete") + } + return &job, nil +} + +func DecodeKolAssistTask(payload []byte) (*KolAssistJob, error) { + job, err := decodeKolAssistTask(payload) + if err != nil { + return nil, err + } + return job, nil +} diff --git a/server/internal/tenant/app/kol_assist_service.go b/server/internal/tenant/app/kol_assist_service.go new file mode 100644 index 0000000..efe8505 --- /dev/null +++ b/server/internal/tenant/app/kol_assist_service.go @@ -0,0 +1,178 @@ +package app + +import ( + "context" + "encoding/json" + "strings" + + "github.com/google/uuid" + + "github.com/geo-platform/tenant-api/internal/shared/auth" + "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/tenant/repository" +) + +const ( + kolAssistModeGenerate = "generate" + kolAssistModeOptimize = "optimize" +) + +var ErrKolAssistTaskNotFound = response.ErrNotFound(40463, "kol_assist_task_not_found", "assist task not found") + +type AssistRequest struct { + Mode string `json:"mode" binding:"required"` + Description string `json:"description"` + CurrentContent string `json:"current_content"` + PromptID *int64 `json:"prompt_id"` + Schema *KolSchemaJSON `json:"schema"` +} + +type KolAssistSubmitResponse struct { + TaskID string `json:"task_id"` +} + +type KolAssistResult struct { + Content string `json:"content"` + Schema KolSchemaJSON `json:"schema"` +} + +type KolAssistTaskResponse struct { + ID string `json:"id"` + Status string `json:"status"` + Result *KolAssistResult `json:"result,omitempty"` + ErrorMessage *string `json:"error_message,omitempty"` +} + +type KolAssistService struct { + profileSvc *KolProfileService + repo repository.KolAssistRepository + llm llm.Client + messaging *rabbitmq.Client +} + +func NewKolAssistService( + profileSvc *KolProfileService, + repo repository.KolAssistRepository, + llmClient llm.Client, + messagingClient *rabbitmq.Client, +) *KolAssistService { + return &KolAssistService{ + profileSvc: profileSvc, + repo: repo, + llm: llmClient, + messaging: messagingClient, + } +} + +func (s *KolAssistService) Submit(ctx context.Context, actor auth.Actor, req AssistRequest) (string, error) { + profile, err := s.profileSvc.RequireActiveKol(ctx, actor) + if err != nil { + return "", err + } + if err := s.llm.Validate(); err != nil { + return "", response.ErrServiceUnavailable(50341, "llm_unavailable", err.Error()) + } + + req, taskType, err := normalizeAssistRequest(req) + if err != nil { + return "", err + } + + requestJSON, err := json.Marshal(req) + if err != nil { + return "", response.ErrInternal(50083, "kol_assist_marshal_failed", "failed to encode assist request") + } + + taskID := uuid.NewString() + if err := s.repo.Create(ctx, repository.CreateKolAssistInput{ + ID: taskID, + TenantID: profile.TenantID, + KolProfileID: profile.ID, + PromptID: req.PromptID, + OperatorID: actor.UserID, + TaskType: taskType, + RequestJSON: requestJSON, + }); err != nil { + return "", response.ErrInternal(50084, "kol_assist_create_failed", err.Error()) + } + + if err := publishKolAssistTask(ctx, s.messaging, kolAssistJob{ + TaskID: taskID, + TenantID: profile.TenantID, + }); err != nil { + _ = s.repo.MarkFailed(context.Background(), profile.TenantID, taskID, err.Error()) + return "", response.ErrServiceUnavailable(50342, "kol_assist_queue_unavailable", "assist queue is busy, please retry") + } + + return taskID, nil +} + +func (s *KolAssistService) Get(ctx context.Context, actor auth.Actor, id string) (*repository.KolAssistTask, error) { + profile, err := s.profileSvc.RequireActiveKol(ctx, actor) + if err != nil { + return nil, err + } + + task, err := s.repo.Get(ctx, profile.TenantID, strings.TrimSpace(id)) + if err != nil { + return nil, response.ErrInternal(50085, "kol_assist_query_failed", err.Error()) + } + if task == nil { + return nil, ErrKolAssistTaskNotFound + } + return task, nil +} + +func NewKolAssistTaskResponse(task *repository.KolAssistTask) (*KolAssistTaskResponse, error) { + if task == nil { + return nil, nil + } + + resp := &KolAssistTaskResponse{ + ID: task.ID, + Status: task.Status, + ErrorMessage: task.ErrorMessage, + } + if len(task.ResultJSON) == 0 { + return resp, nil + } + + var result KolAssistResult + if err := json.Unmarshal(task.ResultJSON, &result); err != nil { + return nil, response.ErrInternal(50086, "kol_assist_result_invalid", err.Error()) + } + if result.Schema.Variables == nil { + result.Schema.Variables = []KolVariableDefinition{} + } + resp.Result = &result + return resp, nil +} + +func normalizeAssistRequest(req AssistRequest) (AssistRequest, string, error) { + req.Mode = strings.ToLower(strings.TrimSpace(req.Mode)) + req.Description = strings.TrimSpace(req.Description) + req.CurrentContent = strings.TrimSpace(req.CurrentContent) + + switch req.Mode { + case kolAssistModeGenerate: + if req.Description == "" { + return AssistRequest{}, "", response.ErrBadRequest(40067, "kol_assist_description_required", "description is required") + } + case kolAssistModeOptimize: + if req.CurrentContent == "" { + return AssistRequest{}, "", response.ErrBadRequest(40068, "kol_assist_current_content_required", "current_content is required") + } + default: + return AssistRequest{}, "", response.ErrBadRequest(40069, "kol_assist_mode_invalid", "mode must be generate or optimize") + } + + if req.Schema != nil { + if err := ValidateSchema(*req.Schema); err != nil { + return AssistRequest{}, "", response.ErrBadRequest(40070, "kol_assist_schema_invalid", err.Error()) + } + } + + return req, "ai_" + req.Mode, nil +} diff --git a/server/internal/worker/assist/kol_assist_worker.go b/server/internal/worker/assist/kol_assist_worker.go new file mode 100644 index 0000000..6e7cde4 --- /dev/null +++ b/server/internal/worker/assist/kol_assist_worker.go @@ -0,0 +1,266 @@ +package assist + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "strings" + "time" + + amqp "github.com/rabbitmq/amqp091-go" + "go.uber.org/zap" + + "github.com/geo-platform/tenant-api/internal/shared/llm" + "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" + "github.com/geo-platform/tenant-api/internal/tenant/repository" +) + +const kolAssistSystemPrompt = "You are a prompt designer for marketing content. Respond ONLY with JSON: {content: string, schema: {variables: [...]}}" + +var errKolAssistConsumerClosed = errors.New("kol assist consumer channel closed") + +type KolAssistWorker struct { + rabbitMQ *rabbitmq.Client + repo repository.KolAssistRepository + llm llm.Client + logger *zap.Logger + retryInterval time.Duration + processTimeout time.Duration + consumerPrefix string + workerConcurrency int +} + +func NewKolAssistWorker( + rabbitMQClient *rabbitmq.Client, + repo repository.KolAssistRepository, + llmClient llm.Client, + logger *zap.Logger, + workerConcurrency int, + processTimeout time.Duration, +) *KolAssistWorker { + if workerConcurrency <= 0 { + workerConcurrency = 1 + } + if processTimeout <= 0 { + processTimeout = 5 * time.Minute + } + + return &KolAssistWorker{ + rabbitMQ: rabbitMQClient, + repo: repo, + llm: llmClient, + logger: logger, + retryInterval: 5 * time.Second, + processTimeout: processTimeout, + consumerPrefix: fmt.Sprintf("kol-assist-%d", os.Getpid()), + workerConcurrency: workerConcurrency, + } +} + +func (w *KolAssistWorker) Start(ctx context.Context) { + if w == nil || w.rabbitMQ == nil || w.repo == nil || w.llm == nil { + return + } + + for i := 0; i < w.workerConcurrency; i++ { + consumerName := fmt.Sprintf("%s-%d", w.consumerPrefix, i+1) + go w.run(ctx, consumerName) + } +} + +func (w *KolAssistWorker) Process(ctx context.Context, job tenantapp.KolAssistJob) error { + task, err := w.repo.Get(ctx, job.TenantID, job.TaskID) + if err != nil { + return fmt.Errorf("load kol assist task: %w", err) + } + if task == nil || task.Status == "completed" || task.Status == "failed" { + return nil + } + + if err := w.repo.MarkStarted(ctx, job.TenantID, job.TaskID); err != nil { + return fmt.Errorf("mark kol assist started: %w", err) + } + + if err := w.processTask(ctx, job, task); err != nil { + _ = w.repo.MarkFailed(context.Background(), job.TenantID, job.TaskID, err.Error()) + if w.logger != nil { + w.logger.Warn("kol assist task failed", + zap.Error(err), + zap.String("task_id", job.TaskID), + zap.Int64("tenant_id", job.TenantID), + ) + } + } + + return nil +} + +func (w *KolAssistWorker) run(ctx context.Context, consumerName string) { + for { + if err := w.consumeOnce(ctx, consumerName); err != nil && ctx.Err() == nil && w.logger != nil { + if errors.Is(err, errKolAssistConsumerClosed) { + w.logger.Info("kol assist consumer channel closed, retrying", + zap.String("consumer", consumerName), + ) + } else { + w.logger.Warn("kol assist worker stopped unexpectedly", + zap.Error(err), + zap.String("consumer", consumerName), + ) + } + } + + select { + case <-ctx.Done(): + return + case <-time.After(w.retryInterval): + } + } +} + +func (w *KolAssistWorker) consumeOnce(ctx context.Context, consumerName string) error { + deliveries, ch, err := w.rabbitMQ.ConsumeKolAssistTask(consumerName) + if err != nil { + return err + } + defer ch.Close() + + for { + select { + case <-ctx.Done(): + return nil + case delivery, ok := <-deliveries: + if !ok { + return errKolAssistConsumerClosed + } + w.handleDelivery(ctx, delivery) + } + } +} + +func (w *KolAssistWorker) handleDelivery(parent context.Context, delivery amqp.Delivery) { + job, err := tenantapp.DecodeKolAssistTask(delivery.Body) + if err != nil { + w.rejectDelivery(delivery, false, err, "", 0) + return + } + + ctx, cancel := context.WithTimeout(parent, w.processTimeout) + defer cancel() + + if err := w.Process(ctx, *job); err != nil { + requeue := !delivery.Redelivered + w.rejectDelivery(delivery, requeue, err, job.TaskID, job.TenantID) + return + } + + if err := delivery.Ack(false); err != nil && w.logger != nil { + w.logger.Warn("kol assist ack failed", + zap.Error(err), + zap.String("task_id", job.TaskID), + zap.Int64("tenant_id", job.TenantID), + ) + } +} + +func (w *KolAssistWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, taskID string, tenantID int64) { + if nackErr := delivery.Nack(false, requeue); nackErr != nil && w.logger != nil { + w.logger.Warn("kol assist nack failed", + zap.Error(nackErr), + zap.String("task_id", taskID), + zap.Int64("tenant_id", tenantID), + ) + } + if w.logger != nil { + w.logger.Warn("kol assist processing failed", + zap.Error(err), + zap.Bool("requeue", requeue), + zap.Bool("redelivered", delivery.Redelivered), + zap.String("task_id", taskID), + zap.Int64("tenant_id", tenantID), + ) + } +} + +func (w *KolAssistWorker) processTask(ctx context.Context, job tenantapp.KolAssistJob, task *repository.KolAssistTask) error { + var req tenantapp.AssistRequest + if err := json.Unmarshal(task.RequestJSON, &req); err != nil { + return fmt.Errorf("decode kol assist request: %w", err) + } + + userPrompt, err := buildKolAssistUserPrompt(req) + if err != nil { + return err + } + + result, err := w.llm.Generate(ctx, llm.GenerateRequest{ + Prompt: fmt.Sprintf("SYSTEM:\n%s\n\nUSER:\n%s", kolAssistSystemPrompt, userPrompt), + ResponseFormat: &llm.ResponseFormat{ + Type: llm.ResponseFormatTypeJSONObject, + Name: "kol_assist_response", + Description: "JSON object with content:string and schema:{variables:[]}", + }, + }, nil) + if err != nil { + return fmt.Errorf("generate kol assist content: %w", err) + } + + parsed, err := parseKolAssistResponse(result.Content) + if err != nil { + return err + } + + resultJSON, err := json.Marshal(parsed) + if err != nil { + return fmt.Errorf("marshal kol assist result: %w", err) + } + + if err := w.repo.MarkCompleted(ctx, job.TenantID, job.TaskID, resultJSON); err != nil { + return fmt.Errorf("mark kol assist completed: %w", err) + } + return nil +} + +func buildKolAssistUserPrompt(req tenantapp.AssistRequest) (string, error) { + switch strings.ToLower(strings.TrimSpace(req.Mode)) { + case "generate": + return req.Description, nil + case "optimize": + schemaJSON := "null" + if req.Schema != nil { + payload, err := json.Marshal(req.Schema) + if err != nil { + return "", fmt.Errorf("encode schema: %w", err) + } + schemaJSON = string(payload) + } + return fmt.Sprintf( + "Optimize the following prompt. Keep variable placeholders intact. Current schema: %s\n\nPROMPT:\n%s", + schemaJSON, + req.CurrentContent, + ), nil + default: + return "", fmt.Errorf("unsupported assist mode %q", req.Mode) + } +} + +func parseKolAssistResponse(content string) (*tenantapp.KolAssistResult, error) { + var result tenantapp.KolAssistResult + if err := json.Unmarshal([]byte(content), &result); err != nil { + return nil, fmt.Errorf("decode kol assist result: %w", err) + } + result.Content = strings.TrimSpace(result.Content) + if result.Content == "" { + return nil, fmt.Errorf("kol assist result content is empty") + } + if result.Schema.Variables == nil { + result.Schema.Variables = []tenantapp.KolVariableDefinition{} + } + if err := tenantapp.ValidateSchema(result.Schema); err != nil { + return nil, fmt.Errorf("invalid kol assist schema: %w", err) + } + return &result, nil +}