query-orchestration/cmd/textractGenerator_test/textdetection.go

//go:build aws

//go:generate go run -tags=aws ./textdetection.go
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	documenttext "queryorchestration/internal/document/text"
	documenttypes "queryorchestration/internal/document/types"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/aws"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/serviceconfig/queue/querysync"
	"queryorchestration/internal/serviceconfig/textract"
	"sync"
	"time"

	awstextract "github.com/aws/aws-sdk-go-v2/service/textract"
	"github.com/aws/aws-sdk-go-v2/service/textract/types"
)

type Config struct {
	serviceconfig.BaseConfig
	textract.TextractConfig
	querysync.QuerySyncConfig
	objectstore.ObjectStoreConfig
}

func main() {
	ctx := context.Background()

	if len(os.Args) != 2 {
		fmt.Println("Error: expected one argument")
		fmt.Println("Usage:", os.Args[0], "<file_base>")
		os.Exit(1)
	}

	name := os.Args[1]

	slog.Info("arguments", "file_base", name)

	cfg := &Config{}
	profile := aws.Profile(os.Getenv("AWS_PROFILE"))
	if profile == "" {
		slog.Error("profile not available")
		os.Exit(1)
	}

	err := cfg.SetTextractClientWithProfile(ctx, profile)
	if err != nil {
		slog.Error("error creating client", "error", err)
		os.Exit(1)
	}

	root, err := os.OpenRoot("assets")
	if err != nil {
		slog.Error("error opening root", "error", err)
		os.Exit(1)
	}

	pdfFile, err := root.Open(fmt.Sprintf("original/%s.pdf", name))
	if err != nil {
		slog.Error("error opening pdf file", "error", err)
		os.Exit(1)
	}

	pdf, err := documenttypes.NewPDFFromReader(pdfFile)
	if err != nil {
		slog.Error("error creating pdf file", "error", err)
		os.Exit(1)
	}

	textSvc := documenttext.New(cfg)

	count, err := pdf.GetPageCount(ctx)
	if err != nil {
		slog.Error("error getting page count", "error", err)
		os.Exit(1)
	}

	slog.Info("detecting text", "page_count", count)

	startTime := time.Now()
	allResults := []map[string]*awstextract.AnalyzeDocumentOutput{}
	allResults = make([]map[string]*awstextract.AnalyzeDocumentOutput, count)
	var wg sync.WaitGroup
	for i := range count {
		wg.Add(1)
		go func() {
			slog.Info("getting page", "index", i)

			baseResult, err := textSvc.GetBasePage(ctx, pdf, i)
			if err != nil {
				slog.Error("error detecting text", "error", err)
				os.Exit(1)
			}

			allResults[i] = map[string]*awstextract.AnalyzeDocumentOutput{
				"base": &baseResult,
			}

			hasTable := false
			for _, block := range baseResult.Blocks {
				if block.BlockType == types.BlockTypeLayoutTable {
					hasTable = true
				}
			}

			if hasTable {
				slog.Info("getting table", "index", i)
				pageBytes, err := pdf.GetPageAsPNG(ctx, i)
				if err != nil {
					slog.Error("error getting page", "error", err)
					os.Exit(1)
				}

				tableResult, err := cfg.GetTextractClient().AnalyzeDocument(ctx, &awstextract.AnalyzeDocumentInput{
					Document: &types.Document{
						Bytes: pageBytes,
					},
					FeatureTypes: []types.FeatureType{
						types.FeatureTypeTables,
					},
				})
				if err != nil {
					slog.Error("error detecting table", "error", err)
					os.Exit(1)
				}

				allResults[i]["table"] = tableResult
			}
			wg.Done()
		}()
	}
	wg.Wait()

	jsonContent, err := json.MarshalIndent(allResults, "", "\t")
	if err != nil {
		slog.Error("error marshaling json", "error", err)
		os.Exit(1)
	}

	outputFilename := fmt.Sprintf("textract/%s.gen", name)
	file, err := root.Create(outputFilename)
	if err != nil {
		slog.Error("error creating file", "error", err)
		os.Exit(1)
	}
	defer file.Close()

	_, err = file.Write(jsonContent)
	if err != nil {
		slog.Error("error writing file", "error", err)
		os.Exit(1)
	}

	slog.Info("extracted text", "output", outputFilename, "time", time.Since(startTime))
}