geo/server/internal/tenant/app/kol_assist_service.go

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
}