Merged in feature/checkbox (pull request #145)

Checkbox Primary Edge Cases

* hascheckboxlogic

* gen

* preppinggen

* exploringcheckbox

* removeunselected

* tests

* failingtest

* failintests

* nomockqueryapi

* db

* name

* nocontainer

* deleterow

* cnc

* extradocsandupdatetextracts

* moretables

* appndedtables

* baseversion

internal/document/clean/create.go

@@ -2,6 +2,7 @@ package documentclean
 
 import (
 	"context"
+	"fmt"
 
 	doctextrunner "queryorchestration/api/docTextRunner"
 	"queryorchestration/internal/serviceconfig/queue"
@@ -12,13 +13,13 @@ import (
 func (s *Service) Clean(ctx context.Context, documentId uuid.UUID) error {
 	isclean, err := s.cfg.GetDBQueries().HasDocumentCleanEntry(ctx, documentId)
 	if err != nil {
-		return err
+		return fmt.Errorf("unable to verify if document has been cleaned: %s", err)
 	}
 
 	if !isclean {
 		err = s.clean(ctx, documentId)
 		if err != nil {
-			return err
+			return fmt.Errorf("unable to clean document: %s", err)
 		}
 	}
 

internal/document/text/checkbox.go

@@ -0,0 +1,54 @@
+package documenttext
+
+import (
+	"strings"
+)
+
+func (s *Service) hasCheckboxKeywords(text string) bool {
+	conditions := []func(string) bool{
+		s.hasExhibit,
+		s.hasBenefit,
+		s.hasPlans,
+		s.hasServices,
+	}
+
+	elementsContained := 0
+	for _, hasCondition := range conditions {
+		if hasCondition(text) {
+			elementsContained++
+		}
+	}
+
+	return elementsContained >= 3
+}
+
+func (s *Service) hasExhibit(text string) bool {
+	return strings.Contains(text, "exhibit")
+}
+
+func (s *Service) hasBenefit(text string) bool {
+	keywords := []string{"benefit plan", "applicable benefit", "benefits"}
+
+	return orKeywords(text, keywords)
+}
+
+func (s *Service) hasPlans(text string) bool {
+	keywords := []string{"plan/program:", "plan(s)", "plans"}
+
+	return orKeywords(text, keywords)
+}
+
+func (s *Service) hasServices(text string) bool {
+	keywords := []string{"service:", "services", "service(s)"}
+	return orKeywords(text, keywords)
+}
+
+func orKeywords(text string, keywords []string) bool {
+	for _, keyword := range keywords {
+		if strings.Contains(text, keyword) {
+			return true
+		}
+	}
+
+	return false
+}

internal/document/text/checkbox_test.go

@@ -0,0 +1,57 @@
+package documenttext
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHasCheckboxKeywords(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasCheckboxKeywords(""))
+	assert.False(t, svc.hasCheckboxKeywords("aaaaaa"))
+	assert.False(t, svc.hasCheckboxKeywords("benefit plan"))
+	assert.False(t, svc.hasCheckboxKeywords("benefit plans"))
+	assert.True(t, svc.hasCheckboxKeywords("exhibit benefit plans"))
+}
+
+func TestHasExhibit(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasExhibit(""))
+	assert.False(t, svc.hasExhibit("aaaaaa"))
+	assert.True(t, svc.hasExhibit("exhibit"))
+}
+
+func TestHasBenefit(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasBenefit(""))
+	assert.False(t, svc.hasBenefit("aaaaaa"))
+	assert.False(t, svc.hasBenefit("benefit"))
+	assert.True(t, svc.hasBenefit("benefit plan"))
+	assert.True(t, svc.hasBenefit("benefit plan(s)"))
+	assert.True(t, svc.hasBenefit("applicable benefit"))
+	assert.True(t, svc.hasBenefit("benefits"))
+}
+
+func TestHasPlans(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasPlans(""))
+	assert.False(t, svc.hasPlans("aaaaaa"))
+	assert.True(t, svc.hasPlans("plan(s):"))
+	assert.True(t, svc.hasPlans("plan(s)/program(s):"))
+	assert.True(t, svc.hasPlans("plans:"))
+	assert.True(t, svc.hasPlans("plan/program:"))
+	assert.True(t, svc.hasPlans("plan(s)"))
+	assert.True(t, svc.hasPlans("plans"))
+}
+
+func TestHasServices(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasServices(""))
+	assert.False(t, svc.hasServices("aaaaaa"))
+	assert.True(t, svc.hasServices("services:"))
+	assert.True(t, svc.hasServices("service:"))
+	assert.True(t, svc.hasServices("service(s):"))
+	assert.True(t, svc.hasServices("services"))
+	assert.True(t, svc.hasServices("service(s)"))
+}

internal/document/text/childElements.go

@@ -7,11 +7,11 @@ import (
 	"github.com/google/uuid"
 )
 
-func (s *Service) getDirectChildren(block types.Block, blockMap map[uuid.UUID]types.Block) []uuid.UUID {
+func (s *Service) GetDirectChildren(id uuid.UUID, blockMap map[uuid.UUID]types.Block) []uuid.UUID {
 	lines := []uuid.UUID{}
 	found := map[uuid.UUID]*bool{}
 
-	for _, id := range s.getChildIds(block) {
+	for _, id := range s.getRelationshipIds(blockMap[id]) {
 		if found[id] != nil && !*found[id] {
 			continue
 		}
@@ -20,7 +20,7 @@ func (s *Service) getDirectChildren(block types.Block, blockMap map[uuid.UUID]ty
 		f := true
 		found[id] = &f
 
-		s.removeChildren(blockMap[id], blockMap, &found)
+		s.removeChildren(id, blockMap, &found)
 	}
 
 	output := []uuid.UUID{}
@@ -35,15 +35,15 @@ func (s *Service) getDirectChildren(block types.Block, blockMap map[uuid.UUID]ty
 	return output
 }
 
-func (s *Service) removeChildren(block types.Block, blockMap map[uuid.UUID]types.Block, found *map[uuid.UUID]*bool) {
-	for _, childId := range s.getChildIds(block) {
+func (s *Service) removeChildren(id uuid.UUID, blockMap map[uuid.UUID]types.Block, found *map[uuid.UUID]*bool) {
+	for _, childId := range s.getRelationshipIds(blockMap[id]) {
 		f := false
 		(*found)[childId] = &f
-		s.removeChildren(blockMap[childId], blockMap, found)
+		s.removeChildren(childId, blockMap, found)
 	}
 }
 
-func (s *Service) getBlockMap(blocksList []types.Block) map[uuid.UUID]types.Block {
+func (s *Service) GetBlockMap(blocksList []types.Block) map[uuid.UUID]types.Block {
 	blocks := make(map[uuid.UUID]types.Block, len(blocksList))
 	for _, block := range blocksList {
 		if block.Id == nil {
@@ -61,7 +61,7 @@ func (s *Service) getBlockMap(blocksList []types.Block) map[uuid.UUID]types.Bloc
 	return blocks
 }
 
-func (s *Service) getChildIds(block types.Block) []uuid.UUID {
+func (s *Service) getRelationshipIds(block types.Block) []uuid.UUID {
 	childIds := []uuid.UUID{}
 	for _, relationship := range block.Relationships {
 		for _, idStr := range relationship.Ids {
@@ -77,3 +77,23 @@ func (s *Service) getChildIds(block types.Block) []uuid.UUID {
 
 	return childIds
 }
+
+func (s *Service) getRelationshipIdsByRel(block types.Block, relType types.RelationshipType) []uuid.UUID {
+	childIds := []uuid.UUID{}
+	for _, relationship := range block.Relationships {
+		if relationship.Type != relType {
+			continue
+		}
+		for _, idStr := range relationship.Ids {
+			id, err := uuid.Parse(idStr)
+			if err != nil {
+				slog.Error("invalid child uuid", "id", id, "error", err)
+				continue
+			}
+
+			childIds = append(childIds, id)
+		}
+	}
+
+	return childIds
+}

internal/document/text/documentText_test.go

@@ -66,6 +66,82 @@ func TestGetDocumentText(t *testing.T) {
 		_, err := svc.getDocumentText(ctx, pdf)
 		require.EqualError(t, err, "no textract response")
 	})
+	t.Run("merged_cell_table", func(t *testing.T) {
+		t.Parallel()
+		cfg := DocTextConfig{}
+		svc := Service{
+			cfg: &cfg,
+		}
+		mockTextract := textractmock.NewMockTextractClient(t)
+		cfg.TextractClient = mockTextract
+
+		textractBaseOut, pdf, txtFile, close := getFiles(t, "merged_cell_table")
+		defer close()
+
+		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
+
+		text, err := svc.getDocumentText(ctx, pdf)
+		require.NoError(t, err)
+
+		assertReaders(t, txtFile, text)
+	})
+	t.Run("table_check", func(t *testing.T) {
+		t.Parallel()
+		cfg := DocTextConfig{}
+		svc := Service{
+			cfg: &cfg,
+		}
+		mockTextract := textractmock.NewMockTextractClient(t)
+		cfg.TextractClient = mockTextract
+
+		textractBaseOut, pdf, txtFile, close := getFiles(t, "table_check")
+		defer close()
+
+		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
+
+		text, err := svc.getDocumentText(ctx, pdf)
+		require.NoError(t, err)
+
+		assertReaders(t, txtFile, text)
+	})
+	t.Run("multi_column", func(t *testing.T) {
+		t.Parallel()
+		cfg := DocTextConfig{}
+		svc := Service{
+			cfg: &cfg,
+		}
+		mockTextract := textractmock.NewMockTextractClient(t)
+		cfg.TextractClient = mockTextract
+
+		textractBaseOut, pdf, txtFile, close := getFiles(t, "multi_column")
+		defer close()
+
+		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
+
+		text, err := svc.getDocumentText(ctx, pdf)
+		require.NoError(t, err)
+
+		assertReaders(t, txtFile, text)
+	})
+	t.Run("table_check_row", func(t *testing.T) {
+		t.Parallel()
+		cfg := DocTextConfig{}
+		svc := Service{
+			cfg: &cfg,
+		}
+		mockTextract := textractmock.NewMockTextractClient(t)
+		cfg.TextractClient = mockTextract
+
+		textractBaseOut, pdf, txtFile, close := getFiles(t, "table_check_row")
+		defer close()
+
+		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
+
+		text, err := svc.getDocumentText(ctx, pdf)
+		require.NoError(t, err)
+
+		assertReaders(t, txtFile, text)
+	})
 
 	t.Run("cnc_01-06", func(t *testing.T) {
 		t.Parallel()
@@ -141,25 +217,6 @@ func TestGetDocumentText(t *testing.T) {
 		text, err := svc.getDocumentText(ctx, pdf)
 		require.NoError(t, err)
 
-		assertReaders(t, txtFile, text)
-	})
-	t.Run("merged_cell_table", func(t *testing.T) {
-		t.Parallel()
-		cfg := DocTextConfig{}
-		svc := Service{
-			cfg: &cfg,
-		}
-		mockTextract := textractmock.NewMockTextractClient(t)
-		cfg.TextractClient = mockTextract
-
-		textractBaseOut, pdf, txtFile, close := getFiles(t, "merged_cell_table")
-		defer close()
-
-		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
-
-		text, err := svc.getDocumentText(ctx, pdf)
-		require.NoError(t, err)
-
 		assertReaders(t, txtFile, text)
 	})
 }
@@ -208,13 +265,20 @@ func textExpectations(t testing.TB, ctx context.Context, mockTextract *textractm
 			).
 			Return(base, nil)
 
-		table := pageResult["table"]
+		table := pageResult["full_features"]
 		if table != nil {
 			mockTextract.EXPECT().
 				AnalyzeDocument(
 					mock.Anything,
 					mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
-						return string(in.Document.Bytes) == string(page) && in.FeatureTypes[0] == types.FeatureTypeTables
+						found := false
+						for _, feature := range in.FeatureTypes {
+							if feature == types.FeatureTypeForms || feature == types.FeatureTypeTables {
+								found = true
+								break
+							}
+						}
+						return string(in.Document.Bytes) == string(page) && found
 					}),
 					mock.Anything,
 				).
@@ -223,17 +287,21 @@ func textExpectations(t testing.TB, ctx context.Context, mockTextract *textractm
 	}
 }
 
+func assertReaderAndString(t testing.TB, expected io.Reader, actual string) {
+	expectedBytes, err := io.ReadAll(expected)
+	require.NoError(t, err)
+	expectedStr := string(expectedBytes)
+
+	if !assert.Equal(t, expectedStr, actual) {
+		diff := cmp.Diff(expectedStr, actual)
+		t.Logf("Detailed diff (-expected +actual):\n%s", diff)
+	}
+}
+
 func assertReaders(t testing.TB, expected io.Reader, actual io.Reader) {
 	actualBytes, err := io.ReadAll(actual)
 	require.NoError(t, err)
 	actualStr := string(actualBytes)
 
-	expectedBytes, err := io.ReadAll(expected)
-	require.NoError(t, err)
-	expectedStr := string(expectedBytes)
-
-	if !assert.Equal(t, expectedStr, actualStr) {
-		diff := cmp.Diff(expectedStr, actualStr)
-		t.Logf("Detailed diff (-expected +actual):\n%s", diff)
-	}
+	assertReaderAndString(t, expected, actualStr)
 }

internal/document/text/keyvalue.go

@@ -0,0 +1,113 @@
+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
+}

internal/document/text/pageText.go

@@ -2,6 +2,7 @@ package documenttext
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"strings"
@@ -19,102 +20,126 @@ func (s *Service) getPageText(ctx context.Context, document documenttypes.File,
 		return "", err
 	}
 
-	var pageBlock *types.Block
-	for _, block := range res.Blocks {
-		if block.BlockType == types.BlockTypePage {
-			pageBlock = &block
-			break
-		}
-	}
-	if pageBlock == nil {
-		return "", fmt.Errorf("no page block found for index: %d", index)
-	}
-
-	blockMap := s.getBlockMap(res.Blocks)
-
-	lines, specialFeatures, err := s.getPageElements(ctx, document, index, *pageBlock, blockMap)
+	features, err := s.ListRequiredFeatures(ctx, res.Blocks, index)
 	if err != nil {
 		return "", err
 	}
 
-	return s.buildPageText(ctx, lines, blockMap, index, specialFeatures)
+	blocks := res.Blocks
+	if len(features) > 0 {
+		res, err := s.GetPageWithFeatures(ctx, document, index, features)
+		if err != nil {
+			return "", err
+		}
+
+		blocks = res.Blocks
+	}
+
+	elements, err := s.getPageElements(ctx, document, index, blocks)
+	if err != nil {
+		return "", err
+	}
+
+	return s.buildPageText(ctx, elements, index)
 }
 
-func (s *Service) getPageElements(ctx context.Context, document documenttypes.File, index int, pageBlock types.Block, blockMap map[uuid.UUID]types.Block) ([]uuid.UUID, specialisedFeatures, error) {
-	directChildren := s.getDirectChildren(pageBlock, blockMap)
+type PageElements struct {
+	Page           uuid.UUID
+	BlockMap       map[uuid.UUID]types.Block
+	DirectChildren []uuid.UUID
+	Tables         []uuid.UUID
+}
 
-	hasTable := false
-	for _, id := range directChildren {
-		if blockMap[id].BlockType == types.BlockTypeLayoutTable {
-			hasTable = true
+func (s *Service) getPageElements(ctx context.Context, document documenttypes.File, index int, blocks []types.Block) (PageElements, error) {
+	elements := PageElements{}
+
+	pageBlock, err := s.GetPageBlock(blocks)
+	if err != nil {
+		return PageElements{}, err
+	}
+
+	elements.Page = pageBlock
+	elements.BlockMap = s.GetBlockMap(blocks)
+	elements.DirectChildren = s.GetDirectChildren(pageBlock, elements.BlockMap)
+
+	for _, block := range blocks {
+		if block.BlockType != types.BlockTypeTable {
+			continue
+		}
+		id, err := uuid.Parse(*block.Id)
+		if err != nil {
+			return PageElements{}, err
+		}
+
+		elements.Tables = append(elements.Tables, id)
+	}
+
+	return elements, nil
+}
+
+func (s *Service) ListRequiredFeatures(ctx context.Context, blocks []types.Block, index int) ([]types.FeatureType, error) {
+	features := []types.FeatureType{}
+	for _, block := range blocks {
+		switch block.BlockType {
+		case types.BlockTypeLayoutTable:
+			features = append(features, types.FeatureTypeTables)
+		case types.BlockTypeLayoutKeyValue:
+			features = append(features, types.FeatureTypeForms)
 		}
 	}
 
-	features := []types.FeatureType{}
-	if hasTable {
-		features = append(features, types.FeatureTypeTables)
-		// features = append(features, types.FeatureTypeForms)
+	if len(features) > 0 {
+		features = append(features, types.FeatureTypeLayout)
+		features = append(features, types.FeatureTypeSignatures)
 	}
 
-	// Detecting Checkboxes -
-	// if has tables
-	// For a page with no tables, we look for 3 out of 4 of the following conditions
-	//     has_exhibit→ 'exhibit' is on the page
-	//     has_benefit→ one of ['benefit plan', 'benefit plan(s)', 'applicable benefit', 'benefits'] is on the page
-	//     has_plans→ one of ['plan(s):', 'plan(s)/program(s):', 'plans:', 'plan/program:', 'plan(s)', 'plans'] is on the page
-	//     has_services → one of ['services:', 'service:', 'service(s):', 'services', 'service(s)'] is on the page
-
-	if len(features) == 0 {
-		return directChildren, specialisedFeatures{}, nil
-	}
-
-	featuresOut, err := s.getSpecialisedFeatures(ctx, document, index, features)
-	if err != nil {
-		return nil, specialisedFeatures{}, err
-	}
-
-	return directChildren, featuresOut, nil
+	return features, nil
 }
 
 type buildParams struct {
-	numOfSignatures int
-	currentTable    int
+	tableIndex int
 }
 
-func (s *Service) buildPageText(ctx context.Context, lines []uuid.UUID, blockMap map[uuid.UUID]types.Block, index int, specialFeatures specialisedFeatures) (string, error) {
+func (s *Service) buildPageText(ctx context.Context, elements PageElements, index int) (string, error) {
 	var builder strings.Builder
-	params := buildParams{
-		numOfSignatures: 0,
-		currentTable:    0,
-	}
+	params := buildParams{}
 
-	for _, id := range lines {
-		err := s.addComponentText(ctx, blockMap[id], blockMap, specialFeatures, &params, &builder)
+	for _, id := range elements.DirectChildren {
+		err := s.addComponentText(ctx, id, elements, &params, &builder)
 		if err != nil {
 			return "", err
 		}
 	}
 
-	for params.currentTable < len(specialFeatures.Tables) {
-		id := specialFeatures.Tables[params.currentTable]
-		err := s.addComponentText(ctx, specialFeatures.BlockMap[id], blockMap, specialFeatures, &params, &builder)
+	for params.tableIndex < len(elements.Tables) {
+		err := s.addTable(ctx, elements.Tables[params.tableIndex], elements.BlockMap, &builder)
 		if err != nil {
 			return "", err
 		}
+
+		params.tableIndex++
 	}
 
-	if params.numOfSignatures > 0 {
-		builder.WriteString(fmt.Sprintf("\nThis page has %d signature.\n", params.numOfSignatures))
+	numOfSignatures := 0
+	for _, block := range elements.BlockMap {
+		if block.BlockType == types.BlockTypeSignature {
+			numOfSignatures++
+		}
+	}
+
+	if numOfSignatures > 0 {
+		builder.WriteString(fmt.Sprintf("\nThis page has %d signature.\n", numOfSignatures))
 	}
 
 	return builder.String(), nil
 }
 
-func (s *Service) addComponentText(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, specialFeatures specialisedFeatures, params *buildParams, builder *strings.Builder) error {
+func (s *Service) addComponentText(ctx context.Context, id uuid.UUID, elements PageElements, params *buildParams, builder *strings.Builder) error {
+	block := elements.BlockMap[id]
 	switch block.BlockType {
-	case types.BlockTypeLayoutText, types.BlockTypeLayoutTitle, types.BlockTypeLayoutHeader, types.BlockTypeLayoutFooter, types.BlockTypeLayoutSectionHeader, types.BlockTypeLayoutPageNumber, types.BlockTypeLayoutList, types.BlockTypeLayoutFigure, types.BlockTypeLayoutKeyValue:
-		for _, id := range s.getChildIds(block) {
-			err := s.addComponentText(ctx, blockMap[id], blockMap, specialFeatures, params, builder)
+	case types.BlockTypeLayoutText, types.BlockTypeLayoutTitle, types.BlockTypeLayoutHeader, types.BlockTypeLayoutFooter, types.BlockTypeLayoutSectionHeader, types.BlockTypeLayoutPageNumber, types.BlockTypeLayoutList, types.BlockTypeLayoutFigure:
+		for _, id := range s.getRelationshipIds(block) {
+			err := s.addComponentText(ctx, id, elements, params, builder)
 			if err != nil {
 				return err
 			}
@@ -131,23 +156,24 @@ func (s *Service) addComponentText(ctx context.Context, block types.Block, block
 			slog.Error("unable to write string", "error", err)
 			return err
 		}
-	case types.BlockTypeSignature:
-		params.numOfSignatures++
-	case types.BlockTypeLayoutTable, types.BlockTypeTable:
-		if block.BlockType == types.BlockTypeLayoutTable && len(specialFeatures.Tables) <= params.currentTable {
-			for _, child := range s.getChildIds(block) {
-				err := s.addComponentText(ctx, blockMap[child], blockMap, specialFeatures, params, builder)
+	case types.BlockTypeLayoutTable:
+		if params.tableIndex < len(elements.Tables) {
+			err := s.addTable(ctx, elements.Tables[params.tableIndex], elements.BlockMap, builder)
+			if err != nil {
+				return err
+			}
+		} else {
+			for _, child := range s.getRelationshipIds(block) {
+				err := s.addComponentText(ctx, child, elements, params, builder)
 				if err != nil {
 					return err
 				}
 			}
-			params.currentTable++
-			break
 		}
-		tableId := specialFeatures.Tables[params.currentTable]
-		params.currentTable++
-		tableBlock := specialFeatures.BlockMap[tableId]
-		err := s.addTable(ctx, tableBlock, specialFeatures.BlockMap, builder)
+
+		params.tableIndex++
+	case types.BlockTypeLayoutKeyValue:
+		err := s.addKeyValue(id, elements, builder)
 		if err != nil {
 			return err
 		}
@@ -159,6 +185,13 @@ func (s *Service) addComponentText(ctx context.Context, block types.Block, block
 }
 
 func (s *Service) GetBasePage(ctx context.Context, doc documenttypes.File, index int) (textract.AnalyzeDocumentOutput, error) {
+	return s.GetPageWithFeatures(ctx, doc, index, []types.FeatureType{
+		types.FeatureTypeLayout,
+		types.FeatureTypeSignatures,
+	})
+}
+
+func (s *Service) GetPageWithFeatures(ctx context.Context, doc documenttypes.File, index int, features []types.FeatureType) (textract.AnalyzeDocumentOutput, error) {
 	pageBytes, err := doc.GetPageAsPNG(ctx, index)
 	if err != nil {
 		return textract.AnalyzeDocumentOutput{}, err
@@ -168,10 +201,7 @@ func (s *Service) GetBasePage(ctx context.Context, doc documenttypes.File, index
 		Document: &types.Document{
 			Bytes: pageBytes,
 		},
-		FeatureTypes: []types.FeatureType{
-			types.FeatureTypeLayout,
-			types.FeatureTypeSignatures,
-		},
+		FeatureTypes: features,
 	})
 	if err != nil {
 		return textract.AnalyzeDocumentOutput{}, err
@@ -179,3 +209,18 @@ func (s *Service) GetBasePage(ctx context.Context, doc documenttypes.File, index
 
 	return *res, err
 }
+
+func (s *Service) GetPageBlock(blocks []types.Block) (uuid.UUID, error) {
+	var pageBlock *types.Block
+	for _, block := range blocks {
+		if block.BlockType == types.BlockTypePage {
+			pageBlock = &block
+			break
+		}
+	}
+	if pageBlock == nil {
+		return uuid.Nil, errors.New("no page block found")
+	}
+
+	return uuid.Parse(*pageBlock.Id)
+}

internal/document/text/pageText_test.go

@@ -51,14 +51,14 @@ func TestGetBlockMap(t *testing.T) {
 	svc := Service{}
 	t.Run("0 blocks", func(t *testing.T) {
 		blocks := []types.Block{}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
 	})
 	t.Run("nil id blocks", func(t *testing.T) {
 		blocks := []types.Block{
 			{},
 		}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
 	})
 	t.Run("1 block", func(t *testing.T) {
@@ -69,7 +69,7 @@ func TestGetBlockMap(t *testing.T) {
 				Id: &blockIdStr,
 			},
 		}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{
 			blockId: {
 				Id: &blockIdStr,
@@ -83,7 +83,7 @@ func TestGetBlockMap(t *testing.T) {
 				Id: &blockIdStr,
 			},
 		}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
 	})
 	t.Run("2 blocks", func(t *testing.T) {
@@ -99,7 +99,7 @@ func TestGetBlockMap(t *testing.T) {
 				Id: &blockTwoIdStr,
 			},
 		}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{
 			blockOneId: {
 				Id: &blockOneIdStr,

internal/document/text/tables.go

@@ -6,30 +6,33 @@ import (
 	"log/slog"
 	"strings"
 
-	documenttypes "queryorchestration/internal/document/types"
-
-	"github.com/aws/aws-sdk-go-v2/service/textract"
 	"github.com/aws/aws-sdk-go-v2/service/textract/types"
 	"github.com/google/uuid"
 )
 
 type cell struct {
-	value  string
-	row    int
-	column int
+	value      string
+	row        int
+	column     int
+	isSelected bool
 }
 
-func (s *Service) addTable(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, builder *strings.Builder) error {
+func (s *Service) addTable(ctx context.Context, id uuid.UUID, blockMap map[uuid.UUID]types.Block, builder *strings.Builder) error {
 	params := tableParams{
 		maxRow:    0,
 		maxColumn: 0,
 		table:     []cell{},
 	}
 
-	directChildren := s.getDirectChildren(block, blockMap)
+	for _, uid := range s.GetDirectChildren(id, blockMap) {
+		s.addComponentTable(ctx, uid, blockMap, &params)
+	}
 
-	for _, uid := range directChildren {
-		s.addComponentTable(ctx, blockMap[uid], blockMap, &params)
+	invalidRows := make([]bool, params.maxRow+1)
+	for _, cell := range params.table {
+		if !cell.isSelected && cell.column == 0 {
+			invalidRows[cell.row] = true
+		}
 	}
 
 	tableArr := make([][]string, params.maxRow+1)
@@ -41,7 +44,14 @@ func (s *Service) addTable(ctx context.Context, block types.Block, blockMap map[
 		tableArr[cell.row][cell.column] = cell.value
 	}
 
-	jsonData, err := json.Marshal(tableArr)
+	finalTable := [][]string{}
+	for i, row := range tableArr {
+		if !invalidRows[i] {
+			finalTable = append(finalTable, row)
+		}
+	}
+
+	jsonData, err := json.Marshal(finalTable)
 	if err != nil {
 		return err
 	}
@@ -59,16 +69,19 @@ type tableParams struct {
 	table     []cell
 }
 
-func (s *Service) addComponentTable(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, params *tableParams) {
+func (s *Service) addComponentTable(ctx context.Context, id uuid.UUID, blockMap map[uuid.UUID]types.Block, params *tableParams) {
+	block := blockMap[id]
+
 	switch block.BlockType {
 	case types.BlockTypeMergedCell:
 		content := ""
-		for _, uid := range s.getChildIds(block) {
+		addRow := true
+		for _, uid := range s.getRelationshipIds(block) {
 			if blockMap[uid].BlockType != types.BlockTypeCell {
 				slog.Warn("not a cell", "block_type", blockMap[uid].BlockType)
 				continue
 			}
-			content = s.getCellContent(blockMap[uid], blockMap)
+			content, addRow = s.getCellContent(uid, blockMap)
 			if content != "" {
 				break
 			}
@@ -83,16 +96,18 @@ func (s *Service) addComponentTable(ctx context.Context, block types.Block, bloc
 					row,
 					column,
 					content,
+					addRow,
 					params,
 				)
 			}
 		}
 	case types.BlockTypeCell:
-		content := s.getCellContent(block, blockMap)
+		content, addRow := s.getCellContent(id, blockMap)
 		s.addCellToTable(
 			int(*block.RowIndex)-1,
 			int(*block.ColumnIndex)-1,
 			content,
+			addRow,
 			params,
 		)
 	default:
@@ -100,21 +115,24 @@ func (s *Service) addComponentTable(ctx context.Context, block types.Block, bloc
 	}
 }
 
-func (s *Service) getCellContent(block types.Block, blockMap map[uuid.UUID]types.Block) string {
+func (s *Service) getCellContent(id uuid.UUID, blockMap map[uuid.UUID]types.Block) (string, bool) {
 	var str strings.Builder
-	for _, uid := range s.getChildIds(block) {
+	isSelected := true
+	for _, uid := range s.getRelationshipIds(blockMap[id]) {
 		if blockMap[uid].BlockType == types.BlockTypeWord {
 			if str.String() != "" {
 				str.WriteRune(' ')
 			}
 			str.WriteString(*blockMap[uid].Text)
+		} else if blockMap[uid].BlockType == types.BlockTypeSelectionElement && blockMap[uid].SelectionStatus == types.SelectionStatusNotSelected {
+			isSelected = false
 		}
 	}
 
-	return str.String()
+	return str.String(), isSelected
 }
 
-func (s *Service) addCellToTable(rowIndex int, columnIndex int, content string, params *tableParams) {
+func (s *Service) addCellToTable(rowIndex int, columnIndex int, content string, isSelected bool, params *tableParams) {
 	if params.maxRow < rowIndex {
 		params.maxRow = rowIndex
 	}
@@ -122,49 +140,9 @@ func (s *Service) addCellToTable(rowIndex int, columnIndex int, content string,
 		params.maxColumn = columnIndex
 	}
 	params.table = append(params.table, cell{
-		column: columnIndex,
-		row:    rowIndex,
-		value:  content,
+		column:     columnIndex,
+		row:        rowIndex,
+		value:      content,
+		isSelected: isSelected,
 	})
 }
-
-type specialisedFeatures struct {
-	BlockMap map[uuid.UUID]types.Block
-	Tables   []uuid.UUID
-}
-
-func (s *Service) getSpecialisedFeatures(ctx context.Context, document documenttypes.File, index int, features []types.FeatureType) (specialisedFeatures, error) {
-	pageBytes, err := document.GetPageAsPNG(ctx, index)
-	if err != nil {
-		return specialisedFeatures{}, err
-	}
-
-	res, err := s.cfg.GetTextractClient().AnalyzeDocument(ctx, &textract.AnalyzeDocumentInput{
-		Document: &types.Document{
-			Bytes: pageBytes,
-		},
-		FeatureTypes: features,
-	})
-	if err != nil {
-		return specialisedFeatures{}, err
-	}
-
-	tables := []uuid.UUID{}
-	for _, block := range res.Blocks {
-		if block.BlockType == types.BlockTypeTable {
-			id, err := uuid.Parse(*block.Id)
-			if err != nil {
-				slog.Warn("error parsing block id for tables", "error", err, "id", *block.Id)
-				continue
-			}
-
-			tables = append(tables, id)
-		}
-	}
-	blocks := s.getBlockMap(res.Blocks)
-
-	return specialisedFeatures{
-		BlockMap: blocks,
-		Tables:   tables,
-	}, nil
-}