feat: Add .env.example configuration file and integrate Redis service into Docker Compose for improved task management. Update desktop application to support document indexing and extraction workflows, enhancing user experience with a new extraction wizard and stepper components. Revise README and documentation for clarity on setup and usage.

2026-07-05 22:39:07 -05:00
parent 9aca926988
commit d95f8f6ad8
70 changed files with 9807 additions and 499 deletions
@@ -0,0 +1,74 @@
+package main
+
+import (
+	"context"
+	"log"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"github.com/pys/pdf-count/internal/api"
+	"github.com/pys/pdf-count/internal/config"
+	"github.com/pys/pdf-count/internal/model"
+	"github.com/pys/pdf-count/internal/ocr"
+	"github.com/pys/pdf-count/internal/queue"
+	"github.com/pys/pdf-count/internal/storage"
+	"github.com/pys/pdf-count/internal/store"
+)
+
+func main() {
+	cfg := config.Load()
+	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
+
+	predictor, err := model.NewPredictor(cfg.ModelPath, cfg.InferencePath)
+	if err != nil {
+		log.Fatalf("model: %v", err)
+	}
+	defer predictor.Close()
+
+	poolSize := cfg.OCRPoolSize
+	if poolSize <= 0 {
+		poolSize = 4
+	}
+	ocrEngine, err := ocr.NewEngine(cfg.OCRLang, poolSize)
+	if err != nil {
+		log.Fatalf("ocr: %v", err)
+	}
+	defer ocrEngine.Close()
+
+	st, err := store.New(cfg.RedisURL, cfg.JobResultTTL)
+	if err != nil {
+		log.Fatalf("redis: %v", err)
+	}
+	defer st.Close()
+
+	stg, err := storage.New(cfg.MinIOEndpoint, cfg.MinIOAccessKey, cfg.MinIOSecretKey, cfg.MinIOBucket, cfg.MinIOUseSSL)
+	if err != nil {
+		log.Fatalf("minio: %v", err)
+	}
+	initCtx, initCancel := context.WithTimeout(ctx, 10*time.Second)
+	if err := stg.EnsureBucket(initCtx); err != nil {
+		initCancel()
+		log.Fatalf("minio bucket: %v", err)
+	}
+	initCancel()
+
+	pub, err := queue.NewPublisher(cfg.NATSURL, cfg.StreamName, cfg.SubjectPrefix)
+	if err != nil {
+		log.Fatalf("nats: %v", err)
+	}
+
+	openAPIPath := os.Getenv("OPENAPI_PATH")
+	if openAPIPath == "" {
+		openAPIPath = "api/openapi.json"
+	}
+
+	srv := api.NewServer(cfg, predictor, ocrEngine, st, stg, pub, api.LoadOpenAPI(openAPIPath))
+	addr := ":" + cfg.Port
+	log.Printf("api listening on %s", addr)
+	if err := api.ListenAndServe(ctx, addr, srv.Router()); err != nil {
+		log.Fatalf("api: %v", err)
+	}
+}
@@ -0,0 +1,157 @@
+package main
+
+import (
+	"context"
+	"log"
+	"math"
+	"os"
+	"os/exec"
+	"os/signal"
+	"strconv"
+	"strings"
+	"syscall"
+	"time"
+
+	"github.com/pys/pdf-count/internal/store"
+)
+
+func main() {
+	cfg := loadConfig()
+	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
+
+	st, err := store.New(cfg.redisURL, 72*time.Hour)
+	if err != nil {
+		log.Fatalf("redis: %v", err)
+	}
+	defer st.Close()
+
+	log.Printf("autoscaler started min=%d max=%d slots=%d service=%s",
+		cfg.minWorkers, cfg.maxWorkers, cfg.workerSlots, cfg.composeService)
+
+	var idleSince time.Time
+
+	for {
+		select {
+		case <-ctx.Done():
+			log.Println("autoscaler stopped")
+			return
+		default:
+		}
+
+		pending, err := st.PendingCount(ctx)
+		if err != nil {
+			log.Printf("autoscaler pending count: %v", err)
+			time.Sleep(cfg.pollInterval)
+			continue
+		}
+
+		current, err := countWorkers(cfg)
+		if err != nil {
+			log.Printf("autoscaler count workers: %v", err)
+			time.Sleep(cfg.pollInterval)
+			continue
+		}
+
+		desired := cfg.minWorkers
+		if pending > 0 {
+			desired = int(math.Ceil(float64(pending) / float64(cfg.workerSlots)))
+			if desired < cfg.minWorkers {
+				desired = cfg.minWorkers
+			}
+			if desired > cfg.maxWorkers {
+				desired = cfg.maxWorkers
+			}
+			idleSince = time.Time{}
+		} else if current > cfg.minWorkers {
+			if idleSince.IsZero() {
+				idleSince = time.Now()
+			} else if time.Since(idleSince) >= cfg.scaleDownIdle {
+				desired = current - 1
+				idleSince = time.Time{}
+			} else {
+				desired = current
+			}
+		}
+
+		if desired != current {
+			log.Printf("autoscaler pending=%d current=%d desired=%d scaling", pending, current, desired)
+			if err := scaleWorkers(cfg, desired); err != nil {
+				log.Printf("autoscaler scale: %v", err)
+			}
+		}
+
+		time.Sleep(cfg.pollInterval)
+	}
+}
+
+type autoscalerConfig struct {
+	redisURL       string
+	composeFile    string
+	composeService string
+	minWorkers     int
+	maxWorkers     int
+	workerSlots    int
+	pollInterval   time.Duration
+	scaleDownIdle  time.Duration
+}
+
+func loadConfig() autoscalerConfig {
+	return autoscalerConfig{
+		redisURL:       getenv("REDIS_URL", "redis://redis:6379/0"),
+		composeFile:    getenv("COMPOSE_FILE", "/project/docker-compose.yml"),
+		composeService: getenv("COMPOSE_SERVICE", "worker"),
+		minWorkers:     getenvInt("MIN_WORKERS", 1),
+		maxWorkers:     getenvInt("MAX_WORKERS", 8),
+		workerSlots:    getenvInt("WORKER_SLOTS", 4),
+		pollInterval:   time.Duration(getenvInt("POLL_INTERVAL_SEC", 5)) * time.Second,
+		scaleDownIdle:  time.Duration(getenvInt("SCALE_DOWN_IDLE_SEC", 60)) * time.Second,
+	}
+}
+
+func countWorkers(cfg autoscalerConfig) (int, error) {
+	cmd := exec.Command("docker", "compose",
+		"-f", cfg.composeFile,
+		"ps", cfg.composeService,
+		"--status", "running",
+		"-q",
+	)
+	out, err := cmd.Output()
+	if err != nil {
+		return 0, err
+	}
+	trimmed := strings.TrimSpace(string(out))
+	if trimmed == "" {
+		return 0, nil
+	}
+	return len(strings.Split(trimmed, "\n")), nil
+}
+
+func scaleWorkers(cfg autoscalerConfig, n int) error {
+	scaleArg := cfg.composeService + "=" + strconv.Itoa(n)
+	cmd := exec.Command("docker", "compose",
+		"-f", cfg.composeFile,
+		"up", "-d", "--no-recreate",
+		"--scale", scaleArg,
+		cfg.composeService,
+	)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	return cmd.Run()
+}
+
+func getenv(key, fallback string) string {
+	if v := os.Getenv(key); v != "" {
+		return v
+	}
+	return fallback
+}
+
+func getenvInt(key string, fallback int) int {
+	if v := os.Getenv(key); v != "" {
+		if n, err := strconv.Atoi(v); err == nil {
+			return n
+		}
+	}
+	return fallback
+}
@@ -0,0 +1,7 @@
+package main
+
+import "github.com/pys/pdf-count/internal/worker"
+
+func main() {
+	worker.RunMain()
+}