Test services

README.md

@@ -165,7 +165,8 @@ converter = PdfConverter(
     config=config_parser.generate_config_dict(),
     artifact_dict=create_model_dict(),
     processor_list=config_parser.get_processors(),
-    renderer=config_parser.get_renderer()
+    renderer=config_parser.get_renderer(),
+    llm_service=config_parser.get_llm_service()
 )
 rendered = converter("FILEPATH")
 ```

marker/builders/layout.py

@@ -2,10 +2,8 @@ from typing import Annotated, List, Optional
 
 from surya.layout import LayoutPredictor
 from surya.layout.schema import LayoutResult, LayoutBox
-from surya.ocr_error.schema import OCRErrorDetectionResult
 
 from marker.builders import BaseBuilder
-from marker.providers import ProviderPageLines
 from marker.providers.pdf import PdfProvider
 from marker.schema import BlockTypes
 from marker.schema.document import Document

marker/builders/llm_layout.py

@@ -1,5 +1,5 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Annotated, Type
+from typing import Annotated
 
 from surya.layout import LayoutPredictor
 from tqdm import tqdm

marker/builders/ocr.py

@@ -5,7 +5,6 @@ from ftfy import fix_text
 from surya.recognition import RecognitionPredictor
 
 from marker.builders import BaseBuilder
-from marker.providers import ProviderPageLines
 from marker.providers.pdf import PdfProvider
 from marker.schema import BlockTypes
 from marker.schema.blocks import BlockId

marker/config/parser.py

@@ -83,7 +83,7 @@ class ConfigParser:
 
     def get_llm_service(self):
         # Only return an LLM service when use_llm is enabled
-        if not self.cli_options["use_llm"]:
+        if not self.cli_options.get("use_llm", False):
             return None
 
         service_cls = self.cli_options["llm_service"]

marker/converters/pdf.py

@@ -115,6 +115,8 @@ class PdfConverter(BaseConverter):
         if llm_service:
             llm_service_cls = strings_to_classes([llm_service])[0]
             llm_service = self.resolve_dependencies(llm_service_cls)
+        elif config.get("use_llm", False):
+            llm_service = self.resolve_dependencies(GoogleGeminiService)
 
         # Inject llm service into artifact_dict so it can be picked up by processors, etc.
         artifact_dict["llm_service"] = llm_service

marker/util.py

@@ -33,7 +33,7 @@ def verify_config_keys(obj):
             if value is None:
                 none_vals += f"{attr_name}, "
 
-    assert len(none_vals) == 0, f"Missing values for {none_vals} are not allowed in {obj.__class__.__name__}."
+    assert len(none_vals) == 0, f"In order to use {obj.__class__.__name__}, you must set the configuration values `{none_vals}`."
 
 
 def assign_config(cls, config: BaseModel | dict | None):

tests/conftest.py

@@ -19,7 +19,8 @@ from marker.renderers.markdown import MarkdownRenderer
 from marker.renderers.json import JSONRenderer
 from marker.schema.registry import register_block_class
 from marker.services.gemini import GoogleGeminiService
-from marker.util import classes_to_strings
+from marker.util import classes_to_strings, strings_to_classes
+
 
 @pytest.fixture(scope="session")
 def model_dict():
@@ -105,12 +106,15 @@ def pdf_document(request, config, pdf_provider, layout_model, ocr_error_model, r
 
 
 @pytest.fixture(scope="function")
-def pdf_converter(request, config, model_dict, renderer):
+def pdf_converter(request, config, model_dict, renderer, llm_service):
+    if llm_service:
+        llm_service = classes_to_strings([llm_service])[0]
     yield PdfConverter(
         artifact_dict=model_dict,
         processor_list=None,
         renderer=classes_to_strings([renderer])[0],
-        config=config
+        config=config,
+        llm_service=llm_service
     )
 
 
@@ -129,13 +133,12 @@ def renderer(request, config):
 
 
 @pytest.fixture(scope="function")
-def llm_service(request):
-    llm_service = GoogleGeminiService(
-        config={
-            "gemini_api_key": "test"
-        }
-    )
-    yield llm_service
+def llm_service(request, config):
+    llm_service = config.get("llm_service")
+    if not llm_service:
+        yield None
+    else:
+        yield strings_to_classes([llm_service])[0]
 
 
 @pytest.fixture(scope="function")

tests/processors/test_llm_processors.py

@@ -29,7 +29,7 @@ def test_llm_form_processor_no_config(pdf_document, llm_service):
 @pytest.mark.filename("form_1040.pdf")
 @pytest.mark.config({"page_range": [0]})
 def test_llm_form_processor_no_cells(pdf_document, llm_service):
-    config = {"use_llm": True, "google_api_key": "test"}
+    config = {"use_llm": True, "gemini_api_key": "test"}
     processor_lst = [LLMFormProcessor(config)]
     processor = LLMSimpleBlockMetaProcessor(processor_lst, llm_service, config)
     processor(pdf_document)
@@ -50,7 +50,7 @@ def test_llm_form_processor(pdf_document, detection_model, table_rec_model, reco
     cell_processor = TableProcessor(detection_model, recognition_model, table_rec_model)
     cell_processor(pdf_document)
 
-    config = {"use_llm": True, "google_api_key": "test"}
+    config = {"use_llm": True, "gemini_api_key": "test"}
     processor_lst = [LLMFormProcessor(config)]
     processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)
@@ -92,7 +92,7 @@ def test_llm_table_processor(pdf_document, detection_model, table_rec_model, rec
     cell_processor = TableProcessor(detection_model, recognition_model, table_rec_model)
     cell_processor(pdf_document)
 
-    processor = LLMTableProcessor(mock_cls, {"use_llm": True, "google_api_key": "test"})
+    processor = LLMTableProcessor(mock_cls, {"use_llm": True, "gemini_api_key": "test"})
     processor(pdf_document)
 
     tables = pdf_document.contained_blocks((BlockTypes.Table,))
@@ -106,7 +106,7 @@ def test_llm_table_processor(pdf_document, detection_model, table_rec_model, rec
 @pytest.mark.filename("A17_FlightPlan.pdf")
 @pytest.mark.config({"page_range": [0]})
 def test_llm_caption_processor_disabled(pdf_document):
-    config = {"use_llm": True, "google_api_key": "test"}
+    config = {"use_llm": True, "gemini_api_key": "test"}
     mock_cls = MagicMock()
     processor_lst = [LLMImageDescriptionProcessor(config)]
     processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
@@ -122,7 +122,7 @@ def test_llm_caption_processor(pdf_document):
     mock_cls = Mock()
     mock_cls.return_value = {"image_description": description}
 
-    config = {"use_llm": True, "google_api_key": "test", "extract_images": False}
+    config = {"use_llm": True, "gemini_api_key": "test", "extract_images": False}
     processor_lst = [LLMImageDescriptionProcessor(config)]
     processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)
@@ -152,7 +152,7 @@ def test_llm_complex_region_processor(pdf_document):
     pdf_document.pages[0].replace_block(old_block, new_block)
 
     # Test processor
-    config = {"use_llm": True, "google_api_key": "test"}
+    config = {"use_llm": True, "gemini_api_key": "test"}
     processor_lst = [LLMComplexRegionProcessor(config)]
     processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)
@@ -170,7 +170,7 @@ def test_multi_llm_processors(pdf_document):
     mock_cls = Mock()
     mock_cls.return_value = {"image_description": description, "html_equation": description}
 
-    config = {"use_llm": True, "google_api_key": "test", "extract_images": False, "min_equation_height": .001}
+    config = {"use_llm": True, "gemini_api_key": "test", "extract_images": False, "min_equation_height": .001}
     processor_lst = [LLMImageDescriptionProcessor(config), LLMEquationProcessor(config)]
     processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)

tests/processors/test_table_merge.py

@@ -21,7 +21,7 @@ def test_llm_table_processor_nomerge(pdf_document, detection_model, table_rec_mo
     tables = pdf_document.contained_blocks((BlockTypes.Table,))
     assert len(tables) == 3
 
-    processor = LLMTableMergeProcessor(mock_cls, {"use_llm": True, "google_api_key": "test"})
+    processor = LLMTableMergeProcessor(mock_cls, {"use_llm": True, "gemini_api_key": "test"})
     processor(pdf_document)
 
     tables = pdf_document.contained_blocks((BlockTypes.Table,))

tests/services/test_service_init.py

@@ -0,0 +1,41 @@
+import pytest
+
+from marker.converters.pdf import PdfConverter
+from marker.services.gemini import GoogleGeminiService
+from marker.services.ollama import OllamaService
+from marker.services.vertex import GoogleVertexService
+
+
+@pytest.mark.output_format("markdown")
+@pytest.mark.config({"page_range": [0]})
+def test_empty_llm(pdf_converter: PdfConverter, temp_pdf):
+    assert pdf_converter.artifact_dict["llm_service"] is None
+    assert pdf_converter.llm_service is None
+
+
+def test_llm_no_keys(model_dict, config):
+    with pytest.raises(AssertionError):
+        PdfConverter(
+            artifact_dict=model_dict,
+            config={"use_llm": True}
+        )
+
+@pytest.mark.output_format("markdown")
+@pytest.mark.config({"page_range": [0], "use_llm": True, "gemini_api_key": "test"})
+def test_llm_gemini(pdf_converter: PdfConverter, temp_pdf):
+    assert pdf_converter.artifact_dict["llm_service"] is not None
+    assert isinstance(pdf_converter.llm_service, GoogleGeminiService)
+
+
+@pytest.mark.output_format("markdown")
+@pytest.mark.config({"page_range": [0], "use_llm": True, "vertex_project_id": "test", "llm_service": "marker.services.vertex.GoogleVertexService"})
+def test_llm_vertex(pdf_converter: PdfConverter, temp_pdf):
+    assert pdf_converter.artifact_dict["llm_service"] is not None
+    assert isinstance(pdf_converter.llm_service, GoogleVertexService)
+
+
+@pytest.mark.output_format("markdown")
+@pytest.mark.config({"page_range": [0], "use_llm": True, "llm_service": "marker.services.ollama.OllamaService"})
+def test_llm_ollama(pdf_converter: PdfConverter, temp_pdf):
+    assert pdf_converter.artifact_dict["llm_service"] is not None
+    assert isinstance(pdf_converter.llm_service, OllamaService)