query-orchestration/internal/document/text/keyvalue.go

package documenttext

import (
	"fmt"
	"log/slog"
	"strings"

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

func (s *Service) getKeyValueSets(id uuid.UUID, elements PageElements) (map[uuid.UUID]uuid.UUID, error) {
	block := elements.BlockMap[id]
	keyValuePairs := map[uuid.UUID]bool{}
	lines := map[uuid.UUID]bool{}
	lineToKey := map[uuid.UUID]uuid.UUID{}

	for _, lineId := range s.getRelationshipIds(block) {
		for _, wordId := range s.getRelationshipIds(elements.BlockMap[lineId]) {
			for _, b := range elements.BlockMap {
				for _, childId := range s.getRelationshipIds(b) {
					if wordId != childId || b.BlockType != types.BlockTypeKeyValueSet {
						continue
					}

					id, err := uuid.Parse(*b.Id)
					if err != nil {
						return nil, err
					}

					lines[lineId] = true
					keyValuePairs[id] = true
					lineToKey[lineId] = id
				}
			}
		}
	}

	for _, lineId := range s.getRelationshipIds(block) {
		if lines[lineId] {
			continue
		}
		for _, wordId := range s.getRelationshipIds(elements.BlockMap[lineId]) {
			for _, b := range elements.BlockMap {
				for _, childId := range s.getRelationshipIds(b) {
					id, err := uuid.Parse(*b.Id)
					if err != nil {
						return nil, err
					}

					if wordId != childId || keyValuePairs[id] {
						continue
					}
					keyValuePairs[id] = true
					lineToKey[lineId] = id
				}
			}
		}
	}

	return lineToKey, nil
}

func (s *Service) addKeyValue(id uuid.UUID, elements PageElements, builder *strings.Builder) error {
	lineToKey, err := s.getKeyValueSets(id, elements)
	if err != nil {
		return err
	}

	for _, lineId := range s.getRelationshipIds(elements.BlockMap[id]) {
		keysBlock := lineToKey[lineId]
		needsSelection := false
		isSelected := false

		for _, id := range s.getRelationshipIdsByRel(elements.BlockMap[keysBlock], types.RelationshipTypeValue) {
			for _, id := range s.getRelationshipIdsByRel(elements.BlockMap[id], types.RelationshipTypeChild) {
				keyBlock := elements.BlockMap[id]
				if keyBlock.BlockType != types.BlockTypeSelectionElement {
					continue
				}
				needsSelection = true
				if keyBlock.SelectionStatus == types.SelectionStatusSelected {
					isSelected = true
					break
				}
			}
		}

		if !needsSelection {
			builder.WriteString(fmt.Sprintf("%s\n", *elements.BlockMap[lineId].Text))
			continue
		} else if !isSelected {
			continue
		}

		var key strings.Builder
		for _, id := range s.getRelationshipIdsByRel(elements.BlockMap[keysBlock], types.RelationshipTypeChild) {
			keyBlock := elements.BlockMap[id]
			if keyBlock.BlockType != types.BlockTypeWord {
				slog.Info("uncaught element in key value", "type", keyBlock.BlockType)
				continue
			}
			if key.String() != "" {
				key.WriteRune(' ')
			}
			key.WriteString(*keyBlock.Text)
		}

		builder.WriteString(fmt.Sprintf("%s\n", key.String()))
	}

	return nil
}