geo/server/cmd/worker-generate/main.go

package main

import (
	"context"
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/geo-platform/tenant-api/internal/bootstrap"
	tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
	tenantcompliance "github.com/geo-platform/tenant-api/internal/tenant/compliance"
	"github.com/geo-platform/tenant-api/internal/tenant/repository"
	generateworker "github.com/geo-platform/tenant-api/internal/worker/generate"
)

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()

	generationProvider := tenantapp.NewGenerationStreamDisabledConfigProvider(app.ConfigStore)
	generationCfg := app.Config().Generation
	generationCfg.StreamEnabled = false

	knowledgeSvc := tenantapp.NewKnowledgeService(
		app.DB,
		app.RetrievalProvider,
		app.VectorStore,
		app.ObjectStorage,
		app.LLM,
		app.Logger,
		app.Config().Retrieval,
		app.Config().LLM,
		0,
		0,
	).WithConfigProvider(app.ConfigStore)

	templateSvc := tenantapp.NewTemplateService(
		app.DB,
		repository.NewCachedTemplateRepository(
			repository.NewTemplateRepository(app.DB),
			app.Cache,
		),
		app.LLM,
		nil,
		knowledgeSvc,
		app.GenerationStreams,
		generationCfg,
		app.Config().LLM.MaxOutputTokens,
	).WithCache(app.Cache).WithLogger(app.Logger).WithConfigProvider(generationProvider)

	promptRuleSvc := tenantapp.NewPromptRuleGenerationService(
		app.DB,
		app.LLM,
		nil,
		knowledgeSvc,
		app.GenerationStreams,
		generationCfg,
		app.Config().LLM.MaxOutputTokens,
	).WithCache(app.Cache).WithLogger(app.Logger).WithConfigProvider(generationProvider)

	publishJobSvc := tenantapp.NewPublishJobServiceWithConfig(
		app.DB,
		app.RabbitMQ,
		app.Redis,
		app.Logger,
		app.ConfigStore,
	).WithCache(app.Cache)
	enterpriseSiteSvc := tenantapp.NewEnterpriseSiteService(app.DB, app.ConfigStore).
		WithCache(app.Cache).
		WithObjectStorage(app.ObjectStorage)
	promptRuleSvc.WithPublishJobService(publishJobSvc)
	promptRuleSvc.WithEnterpriseSiteService(enterpriseSiteSvc)

	imitationSvc := tenantapp.NewArticleImitationService(
		app.DB,
		app.LLM,
		nil,
		knowledgeSvc,
		app.GenerationStreams,
		generationCfg,
		app.Config().LLM.MaxOutputTokens,
	).WithCache(app.Cache).WithLogger(app.Logger).WithConfigProvider(generationProvider)

	kolWorker := generateworker.NewKolGenerationWorker(
		app.DB,
		app.LLM,
		knowledgeSvc,
		app.Logger,
		generationCfg,
		app.Config().LLM.MaxOutputTokens,
	).WithCache(app.Cache).
		WithPublishJobService(publishJobSvc).
		WithEnterpriseSiteService(enterpriseSiteSvc).
		WithConfigProvider(app.ConfigStore)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	app.StartConfigWatcher(ctx)

	generateworker.NewArticleGenerationWorker(
		app.DB,
		app.RabbitMQ,
		templateSvc,
		promptRuleSvc,
		imitationSvc,
		kolWorker,
		app.Logger,
		app.Config().Generation,
	).WithCache(app.Cache).WithConfigProvider(app.ConfigStore).Start(ctx)

	generateworker.NewTemplateAssistWorker(
		app.RabbitMQ,
		templateSvc,
		app.Logger,
		app.Config().Generation.WorkerConcurrency,
		app.Config().Generation.ArticleTimeout,
	).WithConfigProvider(app.ConfigStore).Start(ctx)

	generateworker.NewComplianceReviewWorker(
		app.RabbitMQ,
		tenantcompliance.NewService(app.DB, app.ConfigStore, app.Logger).WithRabbitMQ(app.RabbitMQ).WithLLM(app.LLM),
		app.Logger,
		app.Config().Compliance,
	).WithConfigProvider(app.ConfigStore).Start(ctx)

	app.Logger.Info("worker-generate started")

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	cancel()
	app.Logger.Info("worker-generate shutting down...")
}