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

package app

import (
	"context"
	"encoding/json"
	"strings"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/geo-platform/tenant-api/internal/shared/auth"
	sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
	"github.com/geo-platform/tenant-api/internal/shared/config"
	"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"
	"go.uber.org/zap"
)

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 {
	pool        *pgxpool.Pool
	profileSvc  *KolProfileService
	repo        repository.KolAssistRepository
	llm         llm.Client
	messaging   *rabbitmq.Client
	logger      *zap.Logger
	urlResolver *kolAssistURLResolver
	cache       sharedcache.Cache
}

func NewKolAssistService(
	pool *pgxpool.Pool,
	profileSvc *KolProfileService,
	repo repository.KolAssistRepository,
	llmClient llm.Client,
	messagingClient *rabbitmq.Client,
	llmCfg config.LLMConfig,
	logger *zap.Logger,
) *KolAssistService {
	return &KolAssistService{
		pool:        pool,
		profileSvc:  profileSvc,
		repo:        repo,
		llm:         llmClient,
		messaging:   messagingClient,
		logger:      logger,
		urlResolver: newKolAssistURLResolver(llmCfg, logger),
	}
}

func (s *KolAssistService) WithCache(c sharedcache.Cache) *KolAssistService {
	s.cache = c
	return s
}

func (s *KolAssistService) WithConfigProvider(provider runtimeConfigProvider) *KolAssistService {
	if s != nil && s.urlResolver != nil {
		s.urlResolver.withConfigProvider(provider)
	}
	return s
}

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()
	reservation, err := s.reserveAIPoints(ctx, actor, taskID, req)
	if err != nil {
		return "", err
	}

	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 {
		s.refundReservedAIPoints(context.Background(), actor, reservation, err.Error())
		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())
		s.refundReservedAIPoints(context.Background(), actor, reservation, 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
}

func (s *KolAssistService) reserveAIPoints(
	ctx context.Context,
	actor auth.Actor,
	taskID string,
	req AssistRequest,
) (*AIPointReservation, error) {
	usageType := AIUsageTypeKolPromptGenerate
	if req.Mode == kolAssistModeOptimize {
		usageType = AIUsageTypeKolPromptOptimize
	}
	resourceType := "kol_assist_task"
	resourceUID := taskID
	var resourceID *int64
	if req.PromptID != nil && *req.PromptID > 0 {
		resourceID = req.PromptID
	}
	return ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{
		TenantID:     actor.TenantID,
		OperatorID:   actor.UserID,
		UsageType:    usageType,
		ResourceType: &resourceType,
		ResourceID:   resourceID,
		ResourceUID:  &resourceUID,
		MeteredText:  buildKolAssistAIPointMeteredText(req),
		Metadata: map[string]any{
			"task_id":   taskID,
			"mode":      req.Mode,
			"prompt_id": req.PromptID,
		},
	})
}

func (s *KolAssistService) reserveStreamingAIPoints(ctx context.Context, actor auth.Actor, req AssistRequest) (*AIPointReservation, error) {
	usageType := AIUsageTypeKolPromptGenerate
	if req.Mode == kolAssistModeOptimize {
		usageType = AIUsageTypeKolPromptOptimize
	}
	resourceType := "kol_prompt"
	var resourceID *int64
	if req.PromptID != nil && *req.PromptID > 0 {
		resourceID = req.PromptID
	}
	return ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{
		TenantID:     actor.TenantID,
		OperatorID:   actor.UserID,
		UsageType:    usageType,
		ResourceType: &resourceType,
		ResourceID:   resourceID,
		MeteredText:  buildKolAssistAIPointMeteredText(req),
		Metadata: map[string]any{
			"mode":      req.Mode,
			"prompt_id": req.PromptID,
			"stream":    true,
		},
	})
}

func (s *KolAssistService) refundReservedAIPoints(ctx context.Context, actor auth.Actor, reservation *AIPointReservation, errorMessage string) {
	if reservation == nil {
		return
	}
	_ = RefundAIPoints(ctx, s.pool, s.cache, actor.TenantID, actor.UserID, *reservation, errorMessage)
}

func buildKolAssistAIPointMeteredText(req AssistRequest) string {
	parts := []string{req.Description}
	if req.Mode == kolAssistModeOptimize {
		parts = append(parts, req.CurrentContent)
	}
	return strings.TrimSpace(strings.Join(parts, "\n"))
}