feat(kol): AI assist service + worker
This commit is contained in:
@@ -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...")
|
||||
}
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user