Merged in feature/assertions (pull request #158)

Assertions + Gen Text * assert
2025-05-28 12:36:47 +00:00
parent cc2278086f
commit db469e7491
7 changed files with 290 additions and 17 deletions
@@ -0,0 +1,104 @@
+//go:build aws
+
+//go:generate go run -tags=aws ./textdetection.go
+package main
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"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"
+	"time"
+)
+
+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)
+
+	startTime := time.Now()
+
+	text, err := textSvc.GetDocumentText(ctx, pdf)
+	if err != nil {
+		slog.Error("error getting page count", "error", err)
+		os.Exit(1)
+	}
+
+	outputFilename := fmt.Sprintf("text/%s.txt", name)
+	file, err := root.Create(outputFilename)
+	if err != nil {
+		slog.Error("error creating file", "error", err)
+		os.Exit(1)
+	}
+	defer file.Close()
+
+	s, err := io.ReadAll(text)
+	if err != nil {
+		slog.Error("error reading text", "error", err)
+		os.Exit(1)
+	}
+
+	_, err = file.Write(s)
+	if err != nil {
+		slog.Error("error writing file", "error", err)
+		os.Exit(1)
+	}
+
+	slog.Info("extracted text", "output", outputFilename, "time", time.Since(startTime))
+}
@@ -0,0 +1,148 @@
+//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"
+)
+
+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 index := range count {
+		i := index
+		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,
+			}
+
+			features, err := textSvc.ListRequiredFeatures(ctx, baseResult.Blocks, i)
+			if err != nil {
+				slog.Error("error listing features", "error", err)
+				os.Exit(1)
+			}
+
+			if len(features) > 0 {
+				fullFeaturesResult, err := textSvc.GetPageWithFeatures(ctx, pdf, i, features)
+				if err != nil {
+					slog.Error("error getting features", "error", err)
+					os.Exit(1)
+				}
+
+				allResults[i]["full_features"] = &fullFeaturesResult
+			}
+
+			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))
+}