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,591 @@
+package api
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"mime/multipart"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/go-chi/chi/v5/middleware"
+	"github.com/google/uuid"
+	"github.com/pys/pdf-count/internal/config"
+	"github.com/pys/pdf-count/internal/domain"
+	"github.com/pys/pdf-count/internal/model"
+	"github.com/pys/pdf-count/internal/ocr"
+	"github.com/pys/pdf-count/internal/pipeline"
+	"github.com/pys/pdf-count/internal/queue"
+	"github.com/pys/pdf-count/internal/storage"
+	"github.com/pys/pdf-count/internal/store"
+	"github.com/pys/pdf-count/web"
+	"gopkg.in/yaml.v3"
+)
+
+type Server struct {
+	cfg       config.Config
+	predictor *model.Predictor
+	ocrEngine *ocr.Engine
+	store     *store.Store
+	storage   *storage.Client
+	publisher *queue.Publisher
+	jobRunner *JobRunner
+	openAPI   []byte
+	ready     bool
+	readyMu   sync.RWMutex
+}
+
+func NewServer(cfg config.Config, predictor *model.Predictor, ocrEngine *ocr.Engine, st *store.Store, stg *storage.Client, pub *queue.Publisher, openAPI []byte) *Server {
+	srv := &Server{
+		cfg:       cfg,
+		predictor: predictor,
+		ocrEngine: ocrEngine,
+		store:     st,
+		storage:   stg,
+		publisher: pub,
+		openAPI:   openAPI,
+		ready:     true,
+	}
+	srv.jobRunner = NewJobRunner(st, srv.newProcessor, srv.ocrConcurrency())
+	return srv
+}
+
+func (s *Server) processConcurrency() int {
+	c := s.cfg.SyncProcessConcurrency
+	if c <= 0 {
+		c = s.cfg.OCRPoolSize
+	}
+	if c <= 0 {
+		c = 4
+	}
+	return c
+}
+
+func (s *Server) ocrConcurrency() int {
+	c := s.processConcurrency()
+	pool := s.cfg.OCRPoolSize
+	if pool <= 0 {
+		pool = 4
+	}
+	if c > pool {
+		c = pool
+	}
+	return c
+}
+
+func (s *Server) SetReady(v bool) {
+	s.readyMu.Lock()
+	s.ready = v
+	s.readyMu.Unlock()
+}
+
+func (s *Server) Router() http.Handler {
+	r := chi.NewRouter()
+	r.Use(middleware.RequestID)
+	r.Use(middleware.RealIP)
+	r.Use(middleware.Recoverer)
+
+	r.Get("/health", s.health)
+	r.Get("/ready", s.readyHandler)
+	r.Get("/openapi.json", s.openAPIHandler)
+
+	r.Get("/demo", func(w http.ResponseWriter, r *http.Request) {
+		http.Redirect(w, r, "/demo/", http.StatusMovedPermanently)
+	})
+	r.Handle("/demo/*", http.StripPrefix("/demo/", http.FileServer(http.FS(web.FS))))
+
+	r.Route("/v1", func(r chi.Router) {
+		r.Post("/count", s.countSync)
+		r.Post("/jobs", s.createJob)
+		r.Post("/jobs/{jobID}/upload", s.uploadChunk)
+		r.Get("/jobs/{jobID}", s.getJob)
+		r.Get("/jobs/{jobID}/results", s.getJobResults)
+		r.Delete("/jobs/{jobID}", s.cancelJob)
+	})
+	return r
+}
+
+func (s *Server) health(w http.ResponseWriter, r *http.Request) {
+	writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
+}
+
+func (s *Server) readyHandler(w http.ResponseWriter, r *http.Request) {
+	s.readyMu.RLock()
+	ok := s.ready
+	s.readyMu.RUnlock()
+	if !ok {
+		writeJSON(w, http.StatusServiceUnavailable, map[string]string{"status": "not_ready"})
+		return
+	}
+	writeJSON(w, http.StatusOK, map[string]string{"status": "ready"})
+}
+
+func (s *Server) openAPIHandler(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusOK)
+	_, _ = w.Write(s.openAPI)
+}
+
+func (s *Server) countSync(w http.ResponseWriter, r *http.Request) {
+	includeRejected := r.URL.Query().Get("include_rejected") == "true"
+	maxBytes := int64(s.cfg.MaxPDFSizeMB) * 1024 * 1024
+
+	if err := r.ParseMultipartForm(maxBytes * int64(s.cfg.SyncMaxFiles)); err != nil {
+		httpError(w, http.StatusBadRequest, "invalid multipart form")
+		return
+	}
+
+	files := r.MultipartForm.File["files"]
+	if len(files) == 0 {
+		httpError(w, http.StatusBadRequest, "field 'files' required")
+		return
+	}
+	if len(files) > s.cfg.SyncMaxFiles {
+		httpError(w, http.StatusBadRequest, "too many files")
+		return
+	}
+
+	proc := s.newProcessor(includeRejected)
+	type fileInput struct {
+		filename string
+		data     []byte
+		errRes   *domain.PDFResult
+	}
+	inputs := make([]fileInput, 0, len(files))
+
+	for _, fh := range files {
+		f, err := fh.Open()
+		if err != nil {
+			errRes := domain.PDFResult{
+				Filename: filepath.Base(fh.Filename),
+				Status:   domain.StatusError,
+				Error:    err.Error(),
+			}
+			inputs = append(inputs, fileInput{errRes: &errRes})
+			continue
+		}
+		data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
+		_ = f.Close()
+		if err != nil {
+			errRes := domain.PDFResult{
+				Filename: filepath.Base(fh.Filename),
+				Status:   domain.StatusError,
+				Error:    err.Error(),
+			}
+			inputs = append(inputs, fileInput{errRes: &errRes})
+			continue
+		}
+		if int64(len(data)) > maxBytes {
+			errRes := domain.PDFResult{
+				Filename: filepath.Base(fh.Filename),
+				Status:   domain.StatusError,
+				Error:    "file too large",
+			}
+			inputs = append(inputs, fileInput{errRes: &errRes})
+			continue
+		}
+		inputs = append(inputs, fileInput{
+			filename: filepath.Base(fh.Filename),
+			data:     data,
+		})
+	}
+
+	results := make([]domain.PDFResult, len(inputs))
+	failed := 0
+	sem := make(chan struct{}, s.ocrConcurrency())
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+
+	for i, in := range inputs {
+		if in.errRes != nil {
+			results[i] = *in.errRes
+			failed++
+			continue
+		}
+		wg.Add(1)
+		go func(idx int, in fileInput) {
+			defer wg.Done()
+			sem <- struct{}{}
+			defer func() { <-sem }()
+			res := proc.ProcessPDF(in.filename, in.data)
+			mu.Lock()
+			results[idx] = res
+			if res.Status == domain.StatusError {
+				failed++
+			}
+			mu.Unlock()
+		}(i, in)
+	}
+	wg.Wait()
+
+	writeJSON(w, http.StatusOK, domain.CountResponse{
+		Results:   results,
+		Processed: len(results) - failed,
+		Failed:    failed,
+	})
+}
+
+type createJobRequest struct {
+	Files           []string `json:"files"`
+	Prefix          string   `json:"prefix"`
+	IncludeRejected bool     `json:"include_rejected"`
+	CallbackURL     string   `json:"callback_url"`
+}
+
+func (s *Server) createJob(w http.ResponseWriter, r *http.Request) {
+	var req createJobRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		httpError(w, http.StatusBadRequest, "invalid json body")
+		return
+	}
+
+	ctx := r.Context()
+	pending, err := s.store.PendingCount(ctx)
+	if err == nil && int(pending) > s.cfg.QueueMaxPending {
+		httpError(w, http.StatusTooManyRequests, "queue full")
+		return
+	}
+
+	jobID := uuid.NewString()
+	var objectKeys []string
+	var filenames []string
+
+	if req.Prefix != "" {
+		prefix := strings.TrimSuffix(req.Prefix, "/")
+		keys, err := s.storage.ListByPrefix(ctx, prefix)
+		if err != nil {
+			httpError(w, http.StatusBadRequest, err.Error())
+			return
+		}
+		objectKeys = keys
+		for _, k := range keys {
+			filenames = append(filenames, filepath.Base(k))
+		}
+	} else if len(req.Files) > 0 {
+		for _, f := range req.Files {
+			name := filepath.Base(f)
+			objectKeys = append(objectKeys, storage.JobObjectKey(jobID, name))
+			filenames = append(filenames, name)
+		}
+	} else {
+		// Empty job shell for chunked upload flow.
+		if err := s.store.CreateJob(ctx, jobID, 0, req.IncludeRejected); err != nil {
+			httpError(w, http.StatusInternalServerError, err.Error())
+			return
+		}
+		writeJSON(w, http.StatusAccepted, domain.JobStatus{
+			JobID:  jobID,
+			Status: domain.JobQueued,
+			Total:  0,
+		})
+		return
+	}
+
+	if len(objectKeys) > s.cfg.JobMaxFiles {
+		httpError(w, http.StatusBadRequest, "too many files in job")
+		return
+	}
+
+	if err := s.store.CreateJob(ctx, jobID, len(objectKeys), req.IncludeRejected); err != nil {
+		httpError(w, http.StatusInternalServerError, err.Error())
+		return
+	}
+	if err := s.store.AddJobFiles(ctx, jobID, filenames); err != nil {
+		httpError(w, http.StatusInternalServerError, err.Error())
+		return
+	}
+	_ = s.store.MarkRunning(ctx, jobID)
+
+	for i, key := range objectKeys {
+		msg := domain.QueueMessage{
+			JobID:           jobID,
+			ObjectKey:       key,
+			Filename:        filenames[i],
+			IncludeRejected: req.IncludeRejected,
+		}
+		if err := s.publisher.Publish(ctx, msg); err != nil {
+			httpError(w, http.StatusInternalServerError, err.Error())
+			return
+		}
+	}
+
+	writeJSON(w, http.StatusAccepted, domain.JobStatus{
+		JobID:  jobID,
+		Status: domain.JobRunning,
+		Total:  len(objectKeys),
+	})
+}
+
+func (s *Server) uploadChunk(w http.ResponseWriter, r *http.Request) {
+	jobID := chi.URLParam(r, "jobID")
+	includeRejected := r.URL.Query().Get("include_rejected") == "true"
+	maxBytes := int64(s.cfg.MaxPDFSizeMB) * 1024 * 1024
+
+	if err := r.ParseMultipartForm(maxBytes * 500); err != nil {
+		httpError(w, http.StatusBadRequest, "invalid multipart form")
+		return
+	}
+	files := r.MultipartForm.File["files"]
+	if len(files) == 0 {
+		httpError(w, http.StatusBadRequest, "field 'files' required")
+		return
+	}
+	if len(files) > 500 {
+		httpError(w, http.StatusBadRequest, "max 500 files per chunk")
+		return
+	}
+
+	ctx := r.Context()
+	status, err := s.store.GetJob(ctx, jobID)
+	if err != nil {
+		_ = s.store.CreateJob(ctx, jobID, 0, includeRejected)
+		status = domain.JobStatus{JobID: jobID, Status: domain.JobQueued}
+	}
+	if status.Status == domain.JobCancelled {
+		httpError(w, http.StatusConflict, "job cancelled")
+		return
+	}
+
+	var filenames []string
+
+	if s.cfg.JobInlineProcessing {
+		_ = s.store.MarkRunning(ctx, jobID)
+		for _, fh := range files {
+			f, err := fh.Open()
+			if err != nil {
+				continue
+			}
+			data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
+			_ = f.Close()
+			if err != nil || int64(len(data)) > maxBytes {
+				continue
+			}
+			name := filepath.Base(fh.Filename)
+			if err := s.store.AddJobFiles(ctx, jobID, []string{name}); err != nil {
+				httpError(w, http.StatusInternalServerError, err.Error())
+				return
+			}
+			if err := s.store.IncrementTotal(ctx, jobID, 1); err != nil {
+				httpError(w, http.StatusInternalServerError, err.Error())
+				return
+			}
+			filenames = append(filenames, name)
+			s.jobRunner.ProcessOne(jobID, FilePayload{Filename: name, Data: data}, includeRejected)
+		}
+	} else {
+		var err error
+		filenames, err = s.uploadChunkMinIO(ctx, jobID, files, includeRejected, maxBytes)
+		if err != nil {
+			httpError(w, http.StatusInternalServerError, err.Error())
+			return
+		}
+	}
+
+	if len(filenames) == 0 {
+		httpError(w, http.StatusBadRequest, "no valid files uploaded")
+		return
+	}
+
+	writeJSON(w, http.StatusAccepted, map[string]interface{}{
+		"job_id":   jobID,
+		"uploaded": len(filenames),
+	})
+}
+
+func (s *Server) uploadChunkMinIO(ctx context.Context, jobID string, files []*multipart.FileHeader, includeRejected bool, maxBytes int64) ([]string, error) {
+	filenames := make([]string, 0, len(files))
+	for _, fh := range files {
+		filenames = append(filenames, filepath.Base(fh.Filename))
+	}
+	if len(filenames) == 0 {
+		return nil, nil
+	}
+
+	if err := s.store.AddJobFiles(ctx, jobID, filenames); err != nil {
+		return nil, err
+	}
+	if err := s.store.IncrementTotal(ctx, jobID, len(filenames)); err != nil {
+		return nil, err
+	}
+	if err := s.store.MarkRunning(ctx, jobID); err != nil {
+		return nil, err
+	}
+
+	ioLimit := 8
+	if len(files) < ioLimit {
+		ioLimit = len(files)
+	}
+	ioSem := make(chan struct{}, ioLimit)
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	uploadedNames := make([]string, 0, len(files))
+
+	for _, fh := range files {
+		wg.Add(1)
+		go func(fh *multipart.FileHeader) {
+			defer wg.Done()
+			ioSem <- struct{}{}
+			defer func() { <-ioSem }()
+
+			name := filepath.Base(fh.Filename)
+			f, err := fh.Open()
+			if err != nil {
+				s.recordUploadError(ctx, jobID, name, err.Error())
+				return
+			}
+			data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
+			_ = f.Close()
+			if err != nil {
+				s.recordUploadError(ctx, jobID, name, err.Error())
+				return
+			}
+			if int64(len(data)) > maxBytes {
+				s.recordUploadError(ctx, jobID, name, "file too large")
+				return
+			}
+
+			key := storage.JobObjectKey(jobID, name)
+			if err := s.storage.Put(ctx, key, data, "application/pdf"); err != nil {
+				s.recordUploadError(ctx, jobID, name, err.Error())
+				return
+			}
+
+			msg := domain.QueueMessage{
+				JobID:           jobID,
+				ObjectKey:       key,
+				Filename:        name,
+				IncludeRejected: includeRejected,
+			}
+			if err := s.publisher.Publish(ctx, msg); err != nil {
+				s.recordUploadError(ctx, jobID, name, err.Error())
+				return
+			}
+
+			mu.Lock()
+			uploadedNames = append(uploadedNames, name)
+			mu.Unlock()
+		}(fh)
+	}
+	wg.Wait()
+
+	if len(uploadedNames) == 0 {
+		return nil, nil
+	}
+	return uploadedNames, nil
+}
+
+func (s *Server) recordUploadError(ctx context.Context, jobID, filename, errMsg string) {
+	_ = s.store.SaveResult(ctx, jobID, domain.PDFResult{
+		Filename: filename,
+		Status:   domain.StatusError,
+		Error:    errMsg,
+	})
+	_ = s.store.TryCompleteJob(ctx, jobID)
+}
+
+func (s *Server) getJob(w http.ResponseWriter, r *http.Request) {
+	jobID := chi.URLParam(r, "jobID")
+	waitMs, _ := strconv.Atoi(r.URL.Query().Get("wait_ms"))
+	if waitMs > 30000 {
+		waitMs = 30000
+	}
+
+	deadline := time.Now()
+	if waitMs > 0 {
+		deadline = deadline.Add(time.Duration(waitMs) * time.Millisecond)
+	}
+
+	for {
+		status, err := s.store.GetJob(r.Context(), jobID)
+		if err != nil {
+			httpError(w, http.StatusNotFound, "job not found")
+			return
+		}
+		if isTerminalJobStatus(status.Status) || waitMs <= 0 || !time.Now().Before(deadline) {
+			writeJSON(w, http.StatusOK, status)
+			return
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+}
+
+func isTerminalJobStatus(status string) bool {
+	return status == domain.JobCompleted || status == domain.JobCancelled
+}
+
+func (s *Server) getJobResults(w http.ResponseWriter, r *http.Request) {
+	jobID := chi.URLParam(r, "jobID")
+	cursor := r.URL.Query().Get("cursor")
+	limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
+	page, err := s.store.GetResultsPage(r.Context(), jobID, cursor, limit)
+	if err != nil {
+		httpError(w, http.StatusNotFound, err.Error())
+		return
+	}
+	writeJSON(w, http.StatusOK, page)
+}
+
+func (s *Server) cancelJob(w http.ResponseWriter, r *http.Request) {
+	jobID := chi.URLParam(r, "jobID")
+	if err := s.store.CancelJob(r.Context(), jobID); err != nil {
+		httpError(w, http.StatusInternalServerError, err.Error())
+		return
+	}
+	writeJSON(w, http.StatusOK, map[string]string{"job_id": jobID, "status": domain.JobCancelled})
+}
+
+func (s *Server) newProcessor(includeRejected bool) *pipeline.Processor {
+	return pipeline.NewProcessor(s.predictor, s.ocrEngine, pipeline.Options{
+		DPI:               s.cfg.OCRDPI,
+		Lang:              s.cfg.OCRLang,
+		DBSCANEps:         s.cfg.DBSCANEps,
+		DBSCANMinSamples:  s.cfg.DBSCANMinSamples,
+		IncludeRejected:   includeRejected,
+		LogPipelineTiming: s.cfg.LogPipelineTiming,
+	})
+}
+
+func writeJSON(w http.ResponseWriter, status int, v interface{}) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(status)
+	_ = json.NewEncoder(w).Encode(v)
+}
+
+func httpError(w http.ResponseWriter, status int, msg string) {
+	writeJSON(w, status, map[string]string{"error": msg})
+}
+
+func LoadOpenAPI(path string) []byte {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return []byte(`{"openapi":"3.1.0","info":{"title":"pdf-count","version":"1.0.0"}}`)
+	}
+	if strings.HasSuffix(strings.ToLower(path), ".yaml") || strings.HasSuffix(strings.ToLower(path), ".yml") {
+		var doc interface{}
+		if err := yaml.Unmarshal(data, &doc); err == nil {
+			if j, err := json.Marshal(doc); err == nil {
+				return j
+			}
+		}
+	}
+	return data
+}
+
+func ListenAndServe(ctx context.Context, addr string, handler http.Handler) error {
+	srv := &http.Server{Addr: addr, Handler: handler}
+	go func() {
+		<-ctx.Done()
+		_ = srv.Shutdown(context.Background())
+	}()
+	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		return err
+	}
+	return nil
+}