Merge remote-tracking branch 'origin/RED-10200-2' into RED-10200-2

redaction-service-v1/document/src/main/java/com/iqser/red/service/redaction/v1/server/model/document/entity/EntityEventListener.java

@@ -0,0 +1,18 @@
+package com.iqser.red.service.redaction.v1.server.model.document.entity;
+
+public interface EntityEventListener {
+
+    /**
+     * Invoked when an entity is updated.
+     *
+     * @param entity The entity that was updated.
+     */
+    void onEntityUpdated(IEntity entity);
+
+    /**
+     * Invoked when an entity is removed.
+     *
+     * @param entity The entity that was removed.
+     */
+    void onEntityRemoved(IEntity entity);
+}

redaction-service-v1/document/src/main/java/com/iqser/red/service/redaction/v1/server/model/document/entity/EntityType.java

@@ -6,6 +6,5 @@ public enum EntityType {
     RECOMMENDATION,
     FALSE_POSITIVE,
     FALSE_RECOMMENDATION,
-    DICTIONARY_REMOVAL,
+    DICTIONARY_REMOVAL
-    TEMPORARY
 }

redaction-service-v1/document/src/main/java/com/iqser/red/service/redaction/v1/server/model/document/entity/IEntity.java

@@ -51,18 +51,6 @@ public interface IEntity {
     String type();
 
 
-    /**
-     * Marks this entity and all its intersecting nodes as updated
-     */
-    void update();
-
-
-    /**
-     * Marks this entity and all its intersecting nodes as removed
-     */
-    void remove();
-
-
     /**
      * An Entity is valid, when it active and not a false recommendation, a false positive or a dictionary removal.
      *
@@ -357,19 +345,6 @@ public interface IEntity {
     }
 
 
-    default void handleStateChange(boolean wasValid) {
-
-        if (valid() == wasValid) {
-            return;
-        }
-        if (!removed()) {
-            update();
-        } else {
-            remove();
-        }
-    }
-
-
     /**
      * Adds a collection of matched rules to this entity.
      *
@@ -380,12 +355,46 @@ public interface IEntity {
         if (getMatchedRuleList().equals(matchedRules)) {
             return;
         }
-        boolean valid = valid();
+        boolean wasValid = valid();
         getMatchedRuleList().addAll(matchedRules);
-        if (valid() == valid) {
+        handleStateChange(wasValid);
+    }
+
+
+    void addEntityEventListener(EntityEventListener listener);
+
+
+    void removeEntityEventListener(EntityEventListener listener);
+
+
+    default void notifyEntityUpdated() {
+
+        for (EntityEventListener listener : getEntityEventListeners()) {
+            listener.onEntityUpdated(this);
+        }
+    }
+
+    default void notifyEntityRemoved() {
+
+        for (EntityEventListener listener : getEntityEventListeners()) {
+            listener.onEntityRemoved(this);
+        }
+    }
+
+
+    Collection<EntityEventListener> getEntityEventListeners();
+
+
+    default void handleStateChange(boolean wasValid) {
+
+        if (valid() == wasValid) {
             return;
         }
-        update();
+        if (!removed()) {
+            notifyEntityUpdated();
+        } else {
+            notifyEntityRemoved();
+        }
     }
 
 

redaction-service-v1/document/src/main/java/com/iqser/red/service/redaction/v1/server/model/document/entity/IKieSessionUpdater.java

@@ -1,21 +0,0 @@
-package com.iqser.red.service.redaction.v1.server.model.document.entity;
-
-import com.iqser.red.service.redaction.v1.server.model.document.nodes.Image;
-
-public interface IKieSessionUpdater {
-
-    void insert(TextEntity textEntity);
-
-
-    void update(TextEntity textEntity);
-
-
-    void update(Image image);
-
-
-    void remove(TextEntity textEntity);
-
-
-    void remove(Image image);
-
-}

redaction-service-v1/document/src/main/java/com/iqser/red/service/redaction/v1/server/model/document/entity/TextEntity.java

@@ -78,6 +78,9 @@ public class TextEntity implements IEntity {
     @Builder.Default
     HashedMap<TextEntity, Relation> relations = new HashedMap<>();
 
+    @Builder.Default
+    Collection<EntityEventListener> entityEventListeners = new ArrayList<>();
+
 
     public static TextEntity initialEntityNode(TextRange textRange, String type, EntityType entityType, SemanticNode node) {
 
@@ -267,14 +270,14 @@ public class TextEntity implements IEntity {
     public void addManualChange(BaseAnnotation manualChange) {
 
         manualOverwrite.addChange(manualChange);
-        update();
+        notifyEntityUpdated();
     }
 
 
     public void addManualChanges(List<BaseAnnotation> manualChanges) {
 
         manualOverwrite.addChanges(manualChanges);
-        update();
+        notifyEntityUpdated();
     }
 
 
@@ -338,31 +341,18 @@ public class TextEntity implements IEntity {
                 .orElse(getMatchedRule().isWriteValueWithLineBreaks() ? getValueWithLineBreaks() : value);
     }
 
+    @Override
+    public void addEntityEventListener(EntityEventListener listener) {
 
-    public void update() {
+        entityEventListeners.add(listener);
-
-        getKieSessionUpdater().ifPresent(updater -> updater.update(this));
     }
 
 
-    public void remove() {
+    @Override
+    public void removeEntityEventListener(EntityEventListener listener) {
 
-        getKieSessionUpdater().ifPresent(updater -> updater.remove(this));
+        entityEventListeners.remove(listener);
-    }
 
-
-    private @NonNull Optional<IKieSessionUpdater> getKieSessionUpdater() {
-
-        if (intersectingNodes.isEmpty()) {
-            return Optional.empty();
-        }
-        if (intersectingNodes.get(0) instanceof Document document) {
-            if (document.getKieSessionUpdater() == null) {
-                return Optional.empty();
-            }
-            return Optional.of(document.getKieSessionUpdater());
-        }
-        return Optional.empty();
     }
 
 }

redaction-service-v1/document/src/main/java/com/iqser/red/service/redaction/v1/server/model/document/nodes/Document.java

@@ -11,7 +11,6 @@ import java.util.stream.Stream;
 
 import com.iqser.red.service.redaction.v1.server.model.document.DocumentTree;
 import com.iqser.red.service.redaction.v1.server.model.document.NodeVisitor;
-import com.iqser.red.service.redaction.v1.server.model.document.entity.IKieSessionUpdater;
 import com.iqser.red.service.redaction.v1.server.model.document.textblock.TextBlock;
 
 import lombok.AccessLevel;
@@ -40,8 +39,6 @@ public class Document extends AbstractSemanticNode {
     @Builder.Default
     static final SectionIdentifier sectionIdentifier = SectionIdentifier.document();
 
-    IKieSessionUpdater kieSessionUpdater;
-
     @Override
     public NodeType getType() {
 

redaction-service-v1/document/src/main/java/com/iqser/red/service/redaction/v1/server/model/document/nodes/Image.java

@@ -1,18 +1,19 @@
 package com.iqser.red.service.redaction.v1.server.model.document.nodes;
 
 import java.awt.geom.Rectangle2D;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Optional;
 import java.util.PriorityQueue;
 import java.util.Set;
 
 import com.iqser.red.service.redaction.v1.server.model.document.NodeVisitor;
 import com.iqser.red.service.redaction.v1.server.model.document.TextRange;
+import com.iqser.red.service.redaction.v1.server.model.document.entity.EntityEventListener;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.IEntity;
-import com.iqser.red.service.redaction.v1.server.model.document.entity.IKieSessionUpdater;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.ManualChangeOverwrite;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.MatchedRule;
 import com.iqser.red.service.redaction.v1.server.model.document.textblock.TextBlock;
@@ -23,7 +24,6 @@ import lombok.Builder;
 import lombok.Data;
 import lombok.EqualsAndHashCode;
 import lombok.NoArgsConstructor;
-import lombok.NonNull;
 import lombok.experimental.FieldDefaults;
 import lombok.experimental.SuperBuilder;
 
@@ -55,6 +55,9 @@ public class Image extends AbstractSemanticNode implements IEntity {
 
     Page page;
 
+    @Builder.Default
+    Collection<EntityEventListener> entityEventListeners = new ArrayList<>();
+
 
     @Override
     public NodeType getType() {
@@ -103,6 +106,21 @@ public class Image extends AbstractSemanticNode implements IEntity {
     }
 
 
+    @Override
+    public void addEntityEventListener(EntityEventListener listener) {
+
+        entityEventListeners.add(listener);
+    }
+
+
+    @Override
+    public void removeEntityEventListener(EntityEventListener listener) {
+
+        entityEventListeners.remove(listener);
+
+    }
+
+
     @Override
     public String type() {
 
@@ -111,19 +129,6 @@ public class Image extends AbstractSemanticNode implements IEntity {
     }
 
 
-    @Override
-    public void update() {
-
-        getKieSessionUpdater().ifPresent(updater -> updater.update(this));
-    }
-
-    @Override
-    public void remove() {
-
-        getKieSessionUpdater().ifPresent(updater -> updater.remove(this));
-    }
-
-
     @Override
     public String toString() {
 
@@ -192,18 +197,4 @@ public class Image extends AbstractSemanticNode implements IEntity {
         return true;
     }
 
-    private @NonNull Optional<IKieSessionUpdater> getKieSessionUpdater() {
-
-        if (getDocumentTree() == null) {
-            return Optional.empty();
-        }
-        if (getDocumentTree().getRoot().getNode() instanceof Document document) {
-            if (document.getKieSessionUpdater() == null) {
-                return Optional.empty();
-            }
-            return Optional.of(document.getKieSessionUpdater());
-        }
-        return Optional.empty();
-    }
-
 }

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/RedactionServiceSettings.java

@@ -40,7 +40,7 @@ public class RedactionServiceSettings {
 
     private boolean annotationMode;
 
-    private boolean droolsDebug = true;
+    private boolean droolsDebug;
 
     private boolean protobufJsonFallback = true;
 

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/model/PrecursorEntity.java

@@ -2,6 +2,7 @@ package com.iqser.red.service.redaction.v1.server.model;
 
 import static com.iqser.red.service.redaction.v1.server.service.NotFoundImportedEntitiesService.IMPORTED_REDACTION_TYPE;
 
+import java.util.Collection;
 import java.util.List;
 import java.util.Optional;
 import java.util.PriorityQueue;
@@ -15,6 +16,7 @@ import com.iqser.red.service.persistence.service.v1.api.shared.model.analysislog
 import com.iqser.red.service.persistence.service.v1.api.shared.model.annotations.entitymapped.ManualRedactionEntry;
 import com.iqser.red.service.persistence.service.v1.api.shared.model.annotations.entitymapped.ManualResizeRedaction;
 import com.iqser.red.service.redaction.v1.server.model.document.TextRange;
+import com.iqser.red.service.redaction.v1.server.model.document.entity.EntityEventListener;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.EntityType;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.IEntity;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.ManualChangeOverwrite;
@@ -183,18 +185,6 @@ public class PrecursorEntity implements IEntity {
     }
 
 
-    @Override
-    public void update() {
-
-        // not in KieSession, do nothing
-    }
-    @Override
-    public void remove() {
-
-        // not in KieSession, do nothing
-    }
-
-
     /**
      * @return true when this entity is of EntityType ENTITY or HINT
      */
@@ -211,6 +201,28 @@ public class PrecursorEntity implements IEntity {
     }
 
 
+    @Override
+    public void addEntityEventListener(EntityEventListener listener) {
+
+        throw new UnsupportedOperationException("PrecursorEntity does not support entityEventListeners");
+    }
+
+
+    @Override
+    public void removeEntityEventListener(EntityEventListener listener) {
+
+        throw new UnsupportedOperationException("PrecursorEntity does not support entityEventListeners");
+
+    }
+
+
+    @Override
+    public Collection<EntityEventListener> getEntityEventListeners() {
+
+        throw new UnsupportedOperationException("PrecursorEntity does not support entityEventListeners");
+    }
+
+
     private static EntityType getEntityType(EntryType entryType) {
 
         switch (entryType) {

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/service/EntityFindingUtility.java

@@ -181,7 +181,7 @@ public class EntityFindingUtility {
 
         return textBlocks.stream()
                 .flatMap(searchImplementation::getBoundaries)
-                .map(boundary -> entityCreationService.byTextRangeWithEngine(boundary, "temp", EntityType.TEMPORARY, node, Collections.emptySet()))
+                .map(boundary -> entityCreationService.byTextRangeWithEngine(boundary, "temp", EntityType.ENTITY, node, Collections.emptySet()))
                 .filter(Optional::isPresent)
                 .map(Optional::get)
                 .distinct()
@@ -208,7 +208,7 @@ public class EntityFindingUtility {
         return textBlocks.stream()
                 .flatMap(tb -> searchImplementation.getBoundaries(tb)
                         .filter(textRange -> entityCreationService.isValidEntityTextRange(tb, textRange)))
-                .map(boundary -> entityCreationService.byTextRangeWithEngine(boundary, "temp", EntityType.TEMPORARY, document, Collections.emptySet()))
+                .map(boundary -> entityCreationService.byTextRangeWithEngine(boundary, "temp", EntityType.ENTITY, document, Collections.emptySet()))
                 .filter(Optional::isPresent)
                 .map(Optional::get)
                 .distinct()
@@ -222,7 +222,7 @@ public class EntityFindingUtility {
 
         return searchImplementation.getBoundaries(document.getTextBlock())
                 .filter(textRange -> entityCreationService.isValidEntityTextRange(document.getTextBlock(), textRange))
-                .map(boundary -> entityCreationService.byTextRangeWithEngine(boundary, "temp", EntityType.TEMPORARY, document, Collections.emptySet()))
+                .map(boundary -> entityCreationService.byTextRangeWithEngine(boundary, "temp", EntityType.ENTITY, document, Collections.emptySet()))
                 .filter(Optional::isPresent)
                 .map(Optional::get)
                 .distinct()

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/service/EntityLogCreatorService.java

@@ -431,7 +431,7 @@ public class EntityLogCreatorService {
     private static EntryType getEntryType(EntityType entityType) {
 
         return switch (entityType) {
-            case ENTITY, TEMPORARY -> EntryType.ENTITY;
+            case ENTITY -> EntryType.ENTITY;
             case HINT -> EntryType.HINT;
             case FALSE_POSITIVE, DICTIONARY_REMOVAL -> EntryType.FALSE_POSITIVE;
             case RECOMMENDATION -> EntryType.RECOMMENDATION;

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/service/document/EntityCreationService.java

@@ -1,6 +1,10 @@
 package com.iqser.red.service.redaction.v1.server.service.document;
 
-import static com.iqser.red.service.redaction.v1.server.service.document.EntityCreationUtility.*;
+import static com.iqser.red.service.redaction.v1.server.service.document.EntityCreationUtility.allEntitiesIntersectAndHaveSameTypes;
+import static com.iqser.red.service.redaction.v1.server.service.document.EntityCreationUtility.checkIfBothStartAndEndAreEmpty;
+import static com.iqser.red.service.redaction.v1.server.service.document.EntityCreationUtility.findIntersectingSubNodes;
+import static com.iqser.red.service.redaction.v1.server.service.document.EntityCreationUtility.toLineAfterTextRange;
+import static com.iqser.red.service.redaction.v1.server.service.document.EntityCreationUtility.truncateEndIfLineBreakIsBetween;
 import static com.iqser.red.service.redaction.v1.server.utils.SeparatorUtils.boundaryIsSurroundedBySeparators;
 
 import java.util.Arrays;
@@ -23,25 +27,34 @@ import com.iqser.red.service.redaction.v1.server.model.document.ConsecutiveBound
 import com.iqser.red.service.redaction.v1.server.model.document.DocumentTree;
 import com.iqser.red.service.redaction.v1.server.model.document.TextRange;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.EntityType;
-import com.iqser.red.service.redaction.v1.server.model.document.entity.IKieSessionUpdater;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.ManualChangeOverwrite;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.TextEntity;
-import com.iqser.red.service.redaction.v1.server.model.document.nodes.Document;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.NodeType;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.SemanticNode;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.Table;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.TableCell;
 import com.iqser.red.service.redaction.v1.server.model.document.textblock.TextBlock;
+import com.iqser.red.service.redaction.v1.server.service.drools.KieSessionUpdater;
 import com.iqser.red.service.redaction.v1.server.utils.EntityEnrichmentService;
 import com.iqser.red.service.redaction.v1.server.utils.RectangleTransformations;
 import com.iqser.red.service.redaction.v1.server.utils.RedactionSearchUtility;
 import com.iqser.red.service.redaction.v1.server.utils.exception.NotFoundException;
 
+import lombok.NoArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
+@NoArgsConstructor
 public class EntityCreationService {
 
+    KieSessionUpdater kieSessionUpdater = null;
+
+
+    public EntityCreationService(KieSessionUpdater kieSessionUpdater) {
+
+        this.kieSessionUpdater = kieSessionUpdater;
+    }
+
 
     /**
      * Creates entities found between specified start and stop strings, case-sensitive.
@@ -990,9 +1003,7 @@ public class EntityCreationService {
             return Optional.empty();
         }
         entity.addEngines(engines);
-        if(!entityType.equals(EntityType.TEMPORARY)) {
+        insertToKieSession(entity);
-            insertToKieSession(entity, node);
-        }
         return Optional.of(entity);
     }
 
@@ -1067,7 +1078,7 @@ public class EntityCreationService {
         EntityEnrichmentService.enrichEntity(mergedEntity, node.getTextBlock());
 
         addEntityToGraph(mergedEntity, node);
-        insertToKieSession(mergedEntity, node);
+        insertToKieSession(mergedEntity);
 
         entitiesToMerge.stream()
                 .filter(e -> !e.equals(mergedEntity))
@@ -1134,14 +1145,11 @@ public class EntityCreationService {
      *
      * @param textEntity The merged text entity to insert.
      */
-    public void insertToKieSession(TextEntity textEntity, SemanticNode node) {
+    public void insertToKieSession(TextEntity textEntity) {
 
-        if (node.getDocumentTree().getRoot().getNode() instanceof Document document) {
+        if(kieSessionUpdater != null) {
-            IKieSessionUpdater updater = document.getKieSessionUpdater();
+            textEntity.addEntityEventListener(kieSessionUpdater);
-            if (updater == null) {
+            kieSessionUpdater.insert(textEntity);
-                return;
-            }
-            updater.insert(textEntity);
         }
     }
 

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/service/drools/EntityDroolsExecutionService.java

@@ -28,6 +28,7 @@ import com.iqser.red.service.redaction.v1.server.logger.RulesLogger;
 import com.iqser.red.service.redaction.v1.server.logger.TrackingAgendaEventListener;
 import com.iqser.red.service.redaction.v1.server.model.NerEntities;
 import com.iqser.red.service.redaction.v1.server.model.dictionary.Dictionary;
+import com.iqser.red.service.redaction.v1.server.model.document.entity.TextEntity;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.Document;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.SemanticNode;
 import com.iqser.red.service.redaction.v1.server.service.ManualChangesApplicationService;
@@ -94,9 +95,7 @@ public class EntityDroolsExecutionService {
 
         try {
             KieSessionUpdater kieSessionUpdater = new KieSessionUpdater(kieSession);
-            document.setKieSessionUpdater(kieSessionUpdater);
+            EntityCreationService entityCreationService = new EntityCreationService(kieSessionUpdater);
-
-            EntityCreationService entityCreationService = new EntityCreationService();
             RulesLogger logger = new RulesLogger(webSocketService, context);
             if (settings.isDroolsDebug()) {
                 logger.enableAgendaTracking();
@@ -120,8 +119,10 @@ public class EntityDroolsExecutionService {
                     .flatMap(SemanticNode::streamAllSubNodes)
                     .forEach(kieSession::insert);
 
-            document.getEntities()
+            for (TextEntity entity : document.getEntities()) {
-                    .forEach(kieSession::insert);
+                entity.addEntityEventListener(kieSessionUpdater);
+                kieSession.insert(entity);
+            }
             document.getEntities()
                     .forEach(textEntity -> textEntity.getRelations().values()
                             .forEach(kieSession::insert));

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/service/drools/KieSessionUpdater.java

@@ -3,7 +3,8 @@ package com.iqser.red.service.redaction.v1.server.service.drools;
 import org.kie.api.runtime.KieSession;
 import org.kie.api.runtime.rule.FactHandle;
 
-import com.iqser.red.service.redaction.v1.server.model.document.entity.IKieSessionUpdater;
+import com.iqser.red.service.redaction.v1.server.model.document.entity.EntityEventListener;
+import com.iqser.red.service.redaction.v1.server.model.document.entity.IEntity;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.TextEntity;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.Image;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.SemanticNode;
@@ -14,7 +15,7 @@ import lombok.experimental.FieldDefaults;
 
 @RequiredArgsConstructor
 @FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
-public class KieSessionUpdater implements IKieSessionUpdater {
+public class KieSessionUpdater implements EntityEventListener {
 
     KieSession kieSession;
 
@@ -31,51 +32,52 @@ public class KieSessionUpdater implements IKieSessionUpdater {
     }
 
 
-    public void update(TextEntity textEntity) {
+    @Override
+    public void onEntityUpdated(IEntity entity) {
 
-        kieSession.update(kieSession.getFactHandle(textEntity), textEntity);
+        if (entity instanceof TextEntity textEntity) {
-        updateIntersectingNodes(textEntity);
+            kieSession.update(kieSession.getFactHandle(textEntity), textEntity);
-        textEntity.getRelations().values()
+            updateIntersectingNodes(textEntity);
-                .forEach(this::updateFactIfPresent);
+            textEntity.getRelations().values()
-        textEntity.getRelations().keySet()
+                    .forEach(this::updateFactIfPresent);
-                .forEach(k -> updateFactIfPresent(k.getRelations()
+            textEntity.getRelations().keySet()
-                                                          .get(textEntity)));
+                    .forEach(k -> updateFactIfPresent(k.getRelations()
-    }
+                                                              .get(textEntity)));
-
+        }
-
+        if (entity instanceof Image image) {
-
+            kieSession.update(kieSession.getFactHandle(image), image);
-    public void update(Image image) {
+            SemanticNode parent = image;
-
+            while (parent.hasParent()) {
-        kieSession.update(kieSession.getFactHandle(image), image);
+                parent = parent.getParent();
-        SemanticNode parent = image;
+                kieSession.update(kieSession.getFactHandle(parent), parent);
-        while (parent.hasParent()) {
+            }
-            parent = parent.getParent();
-            kieSession.update(kieSession.getFactHandle(parent), parent);
         }
     }
 
 
-    public void remove(TextEntity textEntity) {
+    @Override
+    public void onEntityRemoved(IEntity entity) {
 
-        //test replace all deletes with updates
+        if (entity instanceof TextEntity textEntity) {
-        kieSession.delete(kieSession.getFactHandle(textEntity));
+            //test replace all deletes with updates
-        updateIntersectingNodes(textEntity);
+            kieSession.delete(kieSession.getFactHandle(textEntity));
-        textEntity.getRelations().values()
+            updateIntersectingNodes(textEntity);
-                .forEach(this::deleteFactIfPresent);
+            textEntity.getRelations().values()
-        textEntity.getRelations().keySet()
+                    .forEach(this::deleteFactIfPresent);
-                .forEach(k -> deleteFactIfPresent(k.getRelations()
+            textEntity.getRelations().keySet()
-                                                          .get(textEntity)));
+                    .forEach(k -> deleteFactIfPresent(k.getRelations()
-    }
+                                                              .get(textEntity)));
-
-
-    public void remove(Image image) {
-
-        kieSession.delete(kieSession.getFactHandle(image));
-        SemanticNode parent = image;
-        while (parent.hasParent()) {
-            parent = parent.getParent();
-            kieSession.update(kieSession.getFactHandle(parent), parent);
         }
+        if (entity instanceof Image image) {
+
+            kieSession.delete(kieSession.getFactHandle(image));
+            SemanticNode parent = image;
+            while (parent.hasParent()) {
+                parent = parent.getParent();
+                kieSession.update(kieSession.getFactHandle(parent), parent);
+            }
+        }
+
     }
 
 
@@ -86,7 +88,6 @@ public class KieSessionUpdater implements IKieSessionUpdater {
     }
 
 
-
     private void updateFactIfPresent(Object o) {
 
         FactHandle factHandle = kieSession.getFactHandle(o);
@@ -95,6 +96,7 @@ public class KieSessionUpdater implements IKieSessionUpdater {
         }
     }
 
+
     private void deleteFactIfPresent(Object o) {
 
         FactHandle factHandle = kieSession.getFactHandle(o);

redaction-service-v1/redaction-service-server-v1/src/test/java/com/iqser/red/service/redaction/v1/server/document/graph/DocumentIEntityInsertionIntegrationTest.java

@@ -7,7 +7,6 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.wildfly.common.Assert.assertTrue;
 
-import java.util.HashSet;
 import java.util.List;
 
 import org.junit.jupiter.api.BeforeEach;
@@ -44,7 +43,7 @@ public class DocumentIEntityInsertionIntegrationTest extends BuildDocumentIntegr
     public void createEntityCreationService() {
 
         MockitoAnnotations.initMocks(this);
-        entityCreationService = new EntityCreationService();
+        entityCreationService = new EntityCreationService(new KieSessionUpdater(kieSession));
     }
 
 
@@ -88,7 +87,6 @@ public class DocumentIEntityInsertionIntegrationTest extends BuildDocumentIntegr
     public void assertSameEntitiesCantBeCreatedTwice() {
 
         Document document = buildGraph("files/new/crafted document.pdf");
-        document.setKieSessionUpdater(new KieSessionUpdater(kieSession));
         String type = "CBI_author";
         assertTrue(entityCreationService.byTextRange(new TextRange(0, 10), type, EntityType.ENTITY, document)
                            .isPresent());

redaction-service-v1/redaction-service-server-v1/src/test/java/com/iqser/red/service/redaction/v1/server/manualchanges/ManualChangesIntegrationTest.java

@@ -15,22 +15,16 @@ import java.util.stream.Collectors;
 import org.junit.jupiter.api.Test;
 
 import com.google.common.collect.Sets;
-import com.iqser.red.service.persistence.service.v1.api.shared.model.AnalyzeRequest;
 import com.iqser.red.service.persistence.service.v1.api.shared.model.annotations.Rectangle;
 import com.iqser.red.service.persistence.service.v1.api.shared.model.annotations.entitymapped.IdRemoval;
 import com.iqser.red.service.persistence.service.v1.api.shared.model.annotations.entitymapped.ManualForceRedaction;
 import com.iqser.red.service.persistence.service.v1.api.shared.model.annotations.entitymapped.ManualResizeRedaction;
-import com.iqser.red.service.redaction.v1.server.mapper.DocumentGraphMapper;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.EntityType;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.PositionOnPage;
 import com.iqser.red.service.redaction.v1.server.model.document.entity.TextEntity;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.Document;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.Paragraph;
 import com.iqser.red.service.redaction.v1.server.rules.RulesIntegrationTest;
-import com.iqser.red.service.redaction.v1.server.service.drools.KieSessionUpdater;
-import com.knecon.fforesight.service.layoutparser.internal.api.queue.LayoutParsingType;
-
-import lombok.SneakyThrows;
 
 public class ManualChangesIntegrationTest extends RulesIntegrationTest {
 
@@ -231,14 +225,4 @@ public class ManualChangesIntegrationTest extends RulesIntegrationTest {
         return new Rectangle((float) rectangle2D.getMinX(), (float) rectangle2D.getMinY(), (float) rectangle2D.getWidth(), (float) rectangle2D.getHeight(), number);
     }
 
-
-    @Override
-    @SneakyThrows
-    protected Document buildGraph(String filename) {
-
-        Document document = super.buildGraph(filename);
-        document.setKieSessionUpdater(new KieSessionUpdater(kieSession));
-        return document;
-    }
-
 }

redaction-service-v1/redaction-service-server-v1/src/test/java/com/iqser/red/service/redaction/v1/server/rules/RulesIntegrationTest.java

@@ -25,6 +25,7 @@ import com.iqser.red.service.redaction.v1.server.model.dictionary.Dictionary;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.Document;
 import com.iqser.red.service.redaction.v1.server.service.ManualChangesApplicationService;
 import com.iqser.red.service.redaction.v1.server.service.document.EntityCreationService;
+import com.iqser.red.service.redaction.v1.server.service.drools.KieSessionUpdater;
 
 @ExtendWith(MockitoExtension.class)
 @Import(RulesIntegrationTest.TestConfiguration.class)
@@ -79,7 +80,7 @@ public class RulesIntegrationTest extends BuildDocumentIntegrationTest {
 
         Dictionary dict = Mockito.mock(Dictionary.class);
         kieSession = kieContainer.newKieSession();
-        entityCreationService = new EntityCreationService();
+        entityCreationService = new EntityCreationService(new KieSessionUpdater(kieSession));
         kieSession.setGlobal("manualChangesApplicationService", manualChangesApplicationService);
         kieSession.setGlobal("entityCreationService", entityCreationService);
         kieSession.setGlobal("dictionary", dict);