tests: add tests for partitioned parse

Author: pmrowla

tests/parse/test_llama_parse.py

@@ -1,6 +1,7 @@
 import os
 import pytest
 import shutil
+from typing import Optional, cast
 from fsspec.implementations.local import LocalFileSystem
 from httpx import AsyncClient
 
@@ -20,11 +21,15 @@ def test_simple_page_text() -> None:
     assert len(result[0].text) > 0
 
 
-@pytest.fixture
-def markdown_parser() -> LlamaParse:
+@pytest.fixture(params=[None, 2])
+def markdown_parser(request: pytest.FixtureRequest) -> LlamaParse:
     if os.environ.get("LLAMA_CLOUD_API_KEY", "") == "":
         pytest.skip("LLAMA_CLOUD_API_KEY not set")
-    return LlamaParse(result_type="markdown", ignore_errors=False)
+    return LlamaParse(
+        result_type="markdown",
+        ignore_errors=False,
+        partition_pages=cast(Optional[int], request.param),
+    )
 
 
 def test_simple_page_markdown(markdown_parser: LlamaParse) -> None:
@@ -35,8 +40,6 @@ def test_simple_page_markdown(markdown_parser: LlamaParse) -> None:
 
 
 def test_simple_page_markdown_bytes(markdown_parser: LlamaParse) -> None:
-    markdown_parser = LlamaParse(result_type="markdown", ignore_errors=False)
-
     filepath = "tests/test_files/attention_is_all_you_need.pdf"
     with open(filepath, "rb") as f:
         file_bytes = f.read()
@@ -51,8 +54,6 @@ def test_simple_page_markdown_bytes(markdown_parser: LlamaParse) -> None:
 
 
 def test_simple_page_markdown_buffer(markdown_parser: LlamaParse) -> None:
-    markdown_parser = LlamaParse(result_type="markdown", ignore_errors=False)
-
     filepath = "tests/test_files/attention_is_all_you_need.pdf"
     with open(filepath, "rb") as f:
         # client must provide extra_info with file_name
@@ -161,9 +162,12 @@ async def test_mixing_input_types() -> None:
     os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
     reason="LLAMA_CLOUD_API_KEY not set",
 )
+@pytest.mark.parametrize("partition_pages", [None, 2])
 @pytest.mark.asyncio
-async def test_download_images() -> None:
-    parser = LlamaParse(result_type="markdown", take_screenshot=True)
+async def test_download_images(partition_pages: Optional[int]) -> None:
+    parser = LlamaParse(
+        result_type="markdown", take_screenshot=True, partition_pages=partition_pages
+    )
     filepath = "tests/test_files/attention_is_all_you_need.pdf"
     json_result = await parser.aget_json([filepath])
 
@@ -175,3 +179,17 @@ async def test_download_images() -> None:
 
     await parser.aget_images(json_result, download_path)
     assert len(os.listdir(download_path)) == len(json_result[0]["pages"][0]["images"])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("split_by_page,expected", [(True, 4), (False, 1)])
+async def test_multiple_page_markdown(
+    markdown_parser: LlamaParse,
+    split_by_page: bool,
+    expected: int,
+) -> None:
+    markdown_parser.split_by_page = split_by_page
+    filepath = "tests/test_files/TOS.pdf"
+    result = await markdown_parser.aload_data(filepath)
+    assert len(result) == expected
+    assert all(len(doc.text) > 0 for doc in result)

tests/parse/test_llama_parse_result.py

@@ -1,6 +1,7 @@
 import tempfile
 import os
 import pytest
+from typing import Optional
 from llama_cloud_services import LlamaParse
 from llama_cloud_services.parse.types import JobResult
 
@@ -15,16 +16,23 @@ def chart_file_path() -> str:
     return "tests/test_files/attention_is_all_you_need_chart.pdf"
 
 
+@pytest.fixture
+def multiple_page_path() -> str:
+    return "tests/test_files/TOS.pdf"
+
+
 @pytest.mark.asyncio
 @pytest.mark.skipif(
     os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
     reason="LLAMA_CLOUD_API_KEY not set",
 )
-async def test_basic_parse_result(file_path: str):
+@pytest.mark.parametrize("partition_pages", [None, 2])
+async def test_basic_parse_result(file_path: str, partition_pages: Optional[int]):
     parser = LlamaParse(
         take_screenshot=True,
         auto_mode=True,
         fast_mode=False,
+        partition_pages=partition_pages,
     )
     result = await parser.aparse(file_path)
 
@@ -142,8 +150,11 @@ async def test_parse_layout(file_path: str):
     os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
     reason="LLAMA_CLOUD_API_KEY not set",
 )
-def test_parse_multiple_files(file_path: str, chart_file_path: str):
-    parser = LlamaParse()
+@pytest.mark.parametrize("partition_pages", [None, 2])
+def test_parse_multiple_files(
+    file_path: str, chart_file_path: str, partition_pages: Optional[int]
+):
+    parser = LlamaParse(partition_pages=partition_pages)
     result = parser.parse([file_path, chart_file_path])
 
     assert isinstance(result, list)
@@ -152,3 +163,40 @@ def test_parse_multiple_files(file_path: str, chart_file_path: str):
     assert isinstance(result[1], JobResult)
     assert result[0].file_name == file_path
     assert result[1].file_name == chart_file_path
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
+    reason="LLAMA_CLOUD_API_KEY not set",
+)
+@pytest.mark.parametrize("partition_pages", [None, 2])
+async def test_multiple_page_parse_result(
+    multiple_page_path: str, partition_pages: Optional[int]
+):
+    parser = LlamaParse(
+        take_screenshot=True,
+        auto_mode=True,
+        fast_mode=False,
+        partition_pages=partition_pages,
+    )
+    results = await parser.aparse(multiple_page_path)
+    if partition_pages is None:
+        assert isinstance(results, JobResult)
+        results = [results]
+    else:
+        assert isinstance(results, list)
+
+    for result in results:
+        assert isinstance(result, JobResult)
+        assert result.job_id is not None
+        assert result.file_name == multiple_page_path
+        assert len(result.pages) > 0
+
+        assert result.pages[0].text is not None
+        assert len(result.pages[0].text) > 0
+
+        assert result.pages[0].md is not None
+        assert len(result.pages[0].md) > 0
+
+        assert result.pages[0].md != result.pages[0].text