Jay Brown
112 lines
2.5 KiB
Go
112 lines
2.5 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"queryorchestration/internal/serviceconfig"
|
|
"queryorchestration/internal/serviceconfig/aws"
|
|
"queryorchestration/internal/serviceconfig/build"
|
|
"queryorchestration/internal/serviceconfig/objectstore"
|
|
"queryorchestration/internal/serviceconfig/textract"
|
|
|
|
documenttext "queryorchestration/internal/document/text"
|
|
documenttypes "queryorchestration/internal/document/types"
|
|
)
|
|
|
|
// Config provides the configuration for the text extractor tool.
|
|
// Note: Query-related configuration has been removed. See remove_query_plan.md.
|
|
type Config struct {
|
|
serviceconfig.BaseConfig
|
|
textract.TextractConfig
|
|
objectstore.ObjectStoreConfig
|
|
}
|
|
|
|
func main() {
|
|
// Print version information before any environment checks
|
|
build.PrintVersionInfo("textExtractor")
|
|
|
|
ctx := context.Background()
|
|
|
|
directory := "files/"
|
|
|
|
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)
|
|
}
|
|
|
|
entries, err := os.ReadDir(directory)
|
|
if err != nil {
|
|
slog.Error("error finding directory", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
for _, entry := range entries {
|
|
if entry.IsDir() {
|
|
continue
|
|
}
|
|
|
|
extractFile(ctx, cfg, directory, entry)
|
|
}
|
|
}
|
|
|
|
func extractFile(ctx context.Context, cfg *Config, directory string, entry os.DirEntry) {
|
|
filePath := filepath.Join(directory, entry.Name())
|
|
pdfFile, err := os.Open(filePath)
|
|
if err != nil {
|
|
slog.Error("error opening pdf file", "error", err)
|
|
return
|
|
}
|
|
|
|
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", entry.Name())
|
|
file, err := os.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))
|
|
}
|