Add tests for extraction converter

extraction_app.py

@@ -0,0 +1,4 @@
+from marker.scripts.run_streamlit_app import extraction_app_cli
+
+if __name__ == "__main__":
+    extraction_app_cli()

marker/converters/extraction.py

@@ -1,3 +1,4 @@
+import json
 import re
 
 from marker.builders.document import DocumentBuilder
@@ -8,7 +9,7 @@ from marker.converters.pdf import PdfConverter
 from marker.extractors.page import PageExtractor, json_schema_to_base_model
 from marker.providers.registry import provider_from_filepath
 
-from marker.renderers.extraction import ExtractionMerger
+from marker.renderers.extraction import ExtractionMerger, ExtractionOutput
 from marker.renderers.markdown import MarkdownRenderer
 
 from marker.logger import get_logger
@@ -36,13 +37,13 @@ class ExtractionConverter(PdfConverter):
 
         return document, provider
 
-    def __call__(self, filepath: str) -> str:
+    def __call__(self, filepath: str) -> ExtractionOutput:
         self.config["paginate_output"] = True  # Ensure we can split the output properly
         self.config["output_format"] = (
             "markdown"  # Output must be markdown for extraction
         )
         try:
-            json_schema_to_base_model(self.config["page_schema"])
+            json_schema_to_base_model(json.loads(self.config["page_schema"]))
         except Exception as e:
             logger.error(f"Could not parse page schema: {e}")
             raise ValueError(

marker/renderers/extraction.py

@@ -45,7 +45,7 @@ class ExtractionMerger:
     def __init__(self):
         pass
 
-    def __call__(self, outputs: Dict[int, ExtractionResult]):
+    def __call__(self, outputs: Dict[int, ExtractionResult]) -> ExtractionOutput:
         pnums = sorted(list(outputs.keys()))
         merged_result = outputs[pnums[0]].extracted_data.copy()
         confidence_exists = outputs[pnums[0]].existence_confidence

pyproject.toml

@@ -70,6 +70,7 @@ marker = "marker.scripts.convert:convert_cli"
 marker_single = "marker.scripts.convert_single:convert_single_cli"
 marker_chunk_convert = "marker.scripts.chunk_convert:chunk_convert_cli"
 marker_gui = "marker.scripts.run_streamlit_app:streamlit_app_cli"
+marker_extract = "marker.scripts.run_streamlit_app:extraction_app_cli"
 marker_server = "marker.scripts.server:server_cli"
 
 [build-system]

tests/converters/test_extraction_converter.py

@@ -0,0 +1,66 @@
+import json
+import pytest
+
+from marker.converters.extraction import ExtractionConverter
+from marker.extractors.page import PageExtractionSchema
+from marker.services import BaseService
+
+
+class MockLLMService(BaseService):
+    def __call__(self, prompt, image=None, page=None, response_schema=None, **kwargs):
+        assert response_schema == PageExtractionSchema
+        return {
+            "description": "Mock extraction description",
+            "extracted_json": json.dumps({"test_key": "test_value"}),
+            "existence_confidence": 5,
+            "value_confidence": 5,
+        }
+
+
+@pytest.fixture
+def mock_llm_service():
+    return MockLLMService
+
+
+@pytest.fixture
+def extraction_converter(config, model_dict, mock_llm_service):
+    test_schema = {
+        "title": "TestSchema",
+        "type": "object",
+        "properties": {"test_key": {"title": "Test Key", "type": "string"}},
+        "required": ["test_key"],
+    }
+
+    config["page_schema"] = json.dumps(test_schema)
+    config["output_format"] = "markdown"
+    model_dict["llm_service"] = mock_llm_service
+
+    converter = ExtractionConverter(
+        artifact_dict=model_dict, processor_list=None, config=config
+    )
+    converter.default_llm_service = MockLLMService
+    return converter
+
+
+@pytest.mark.config({"page_range": [0]})
+def test_extraction_converter_invalid_schema(
+    config, model_dict, mock_llm_service, temp_doc
+):
+    config["page_schema"] = "invalid json"
+
+    model_dict["llm_service"] = mock_llm_service
+    converter = ExtractionConverter(
+        artifact_dict=model_dict, processor_list=None, config=config
+    )
+
+    with pytest.raises(ValueError):
+        converter(temp_doc.name)
+
+
+@pytest.mark.config({"page_range": [0, 1]})
+def test_extraction_converter_multiple_pages(extraction_converter, temp_doc):
+    result = extraction_converter(temp_doc.name)
+
+    assert result is not None
+    assert result.json is not None
+    assert result.json == {"test_key": "test_value"}