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