267 lines
7.3 KiB
Go
267 lines
7.3 KiB
Go
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
|
|
}
|