Commit
·
285f1d6
1
Parent(s):
8c63b8f
Add test
Browse files
tests/services/test_service_init.py
CHANGED
|
@@ -4,6 +4,7 @@ from marker.converters.pdf import PdfConverter
|
|
| 4 |
from marker.services.gemini import GoogleGeminiService
|
| 5 |
from marker.services.ollama import OllamaService
|
| 6 |
from marker.services.vertex import GoogleVertexService
|
|
|
|
| 7 |
|
| 8 |
|
| 9 |
@pytest.mark.output_format("markdown")
|
|
@@ -38,4 +39,10 @@ def test_llm_vertex(pdf_converter: PdfConverter, temp_doc):
|
|
| 38 |
@pytest.mark.config({"page_range": [0], "use_llm": True, "llm_service": "marker.services.ollama.OllamaService"})
|
| 39 |
def test_llm_ollama(pdf_converter: PdfConverter, temp_doc):
|
| 40 |
assert pdf_converter.artifact_dict["llm_service"] is not None
|
| 41 |
-
assert isinstance(pdf_converter.llm_service, OllamaService)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 4 |
from marker.services.gemini import GoogleGeminiService
|
| 5 |
from marker.services.ollama import OllamaService
|
| 6 |
from marker.services.vertex import GoogleVertexService
|
| 7 |
+
from marker.services.openai import OpenAIService
|
| 8 |
|
| 9 |
|
| 10 |
@pytest.mark.output_format("markdown")
|
|
|
|
| 39 |
@pytest.mark.config({"page_range": [0], "use_llm": True, "llm_service": "marker.services.ollama.OllamaService"})
|
| 40 |
def test_llm_ollama(pdf_converter: PdfConverter, temp_doc):
|
| 41 |
assert pdf_converter.artifact_dict["llm_service"] is not None
|
| 42 |
+
assert isinstance(pdf_converter.llm_service, OllamaService)
|
| 43 |
+
|
| 44 |
+
@pytest.mark.output_format("markdown")
|
| 45 |
+
@pytest.mark.config({"page_range": [0], "use_llm": True, "llm_service": "marker.services.openai.OpenAIService", "openai_api_key": "test"})
|
| 46 |
+
def test_llm_ollama(pdf_converter: PdfConverter, temp_doc):
|
| 47 |
+
assert pdf_converter.artifact_dict["llm_service"] is not None
|
| 48 |
+
assert isinstance(pdf_converter.llm_service, OpenAIService)
|