Made refiner, referencer and builder less strict

structure_refiner.py

@@ -1,4 +1,5 @@
 import json
+import logging
 import traceback
 
 from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator
@@ -58,6 +59,25 @@ class StructureRefiner:
             if isinstance(element, ImageElement) and element.event_id not in input_image_ids:
                 raise ValueError("image.event_id must come from present")
 
+    @staticmethod
+    def _keep_present_images(
+        result: _RefineResult,
+        window: Window[StructureElement],
+    ) -> _RefineResult:
+        """Discard image elements invented from context windows."""
+        input_image_ids = {
+            element.event_id
+            for element in window.present
+            if isinstance(element, ImageElement)
+        }
+        elements: list[StructureElement] = []
+        for element in result.elements:
+            if isinstance(element, ImageElement) and element.event_id not in input_image_ids:
+                logging.warning("Ignoring image outside present window: id=%s", element.event_id)
+                continue
+            elements.append(element)
+        return result.model_copy(update={"elements": elements})
+
     def _refine_window(self, window: Window[StructureElement]) -> _RefineResult:
         messages = [
             self._system_prompt,
@@ -91,6 +111,7 @@ class StructureRefiner:
 
             try:
                 result = _RefineResult.model_validate_json(response)
+                result = self._keep_present_images(result, window)
                 self._validate_result(result, window)
                 return result
             except (ValidationError, ValueError, TypeError):