chore: implement suggestions

llama_cloud_services/parse/base.py

@@ -37,6 +37,7 @@ from llama_cloud_services.parse.utils import (
     nest_asyncio_msg,
     make_api_request,
     partition_pages,
+    extract_tables_from_json_results,
 )
 
 # can put in a path to the file or the file bytes itself
@@ -1620,6 +1621,16 @@ class LlamaParse(BasePydanticReader):
             else:
                 raise e
 
+    def get_tables(self, json_results: List[dict], download_path: str) -> List[str]:
+        if not os.path.exists(download_path):
+            os.makedirs(download_path)
+        return extract_tables_from_json_results(json_results, download_path)
+
+    async def aget_tables(
+        self, json_result: List[dict], download_path: str
+    ) -> List[str]:
+        return await asyncio.to_thread(self.get_tables, json_result, download_path)
+
     async def aget_xlsx(
         self, json_result: List[dict], download_path: str
     ) -> List[dict]:

llama_cloud_services/parse/utils.py

@@ -1,7 +1,10 @@
 import httpx
 import itertools
 import logging
+import os
 from enum import Enum
+from pathlib import Path
+from datetime import datetime
 from tenacity import (
     retry,
     stop_after_attempt,
@@ -9,7 +12,7 @@ from tenacity import (
     retry_if_exception,
     before_sleep_log,
 )
-from typing import Any, Iterable, Iterator, Optional
+from typing import Any, Iterable, Iterator, Optional, List, cast
 
 logger = logging.getLogger(__name__)
 
@@ -351,3 +354,30 @@ def partition_pages(
             yield ",".join(targets)
         else:
             return
+
+
+def extract_tables_from_json_results(
+    json_results: List[dict], download_path: str
+) -> List[str]:
+    tables = []
+    for json_result in json_results:
+        pages = json_result["pages"]
+        for page in pages:
+            page = cast(dict, page)
+            items = page.get("items", [])
+            if items:
+                for i, item in enumerate(items):
+                    item = cast(dict, item)
+                    if item.get("type", "") == "table" and item.get("csv", ""):
+                        savepath = os.path.join(
+                            download_path,
+                            f"table_{datetime.now().strftime('%Y_%d_%m_%H_%M_%S_%f')[:-3]}.csv",
+                        )
+                        if Path(savepath).exists():
+                            savepath = (
+                                savepath.replace(".csv", "_")[0] + str(i) + ".csv"
+                            )
+                        with open(savepath, "w") as f:
+                            f.write(item["csv"])
+                        tables.append(savepath)
+    return tables

pyproject.toml

@@ -8,7 +8,7 @@ python_version = "3.10"
 
 [tool.poetry]
 name = "llama-cloud-services"
-version = "0.6.46"
+version = "0.6.47"
 description = "Tailored SDK clients for LlamaCloud services."
 authors = ["Logan Markewich <logan@runllama.ai>"]
 license = "MIT"

tests/parse/test_utils.py

@@ -1,7 +1,30 @@
 import pytest
 
+from pathlib import Path
+from llama_cloud_services.parse.utils import (
+    expand_target_pages,
+    partition_pages,
+    extract_tables_from_json_results,
+)
+from typing import List
 
-from llama_cloud_services.parse.utils import expand_target_pages, partition_pages
+
+@pytest.fixture()
+def pseudo_json_results() -> List[dict]:
+    return [
+        {
+            "pages": [
+                {
+                    "items": [
+                        {
+                            "type": "table",
+                            "csv": "Name,Age,Height (cm)\nAnna,12,140\nBob,22,175\nClaire,33,173\nDenis,44,185\n",
+                        }
+                    ]
+                }
+            ]
+        }
+    ]
 
 
 def test_expand_target_pages() -> None:
@@ -28,3 +51,12 @@ def test_partion_pages() -> None:
     assert result == ["0,2-3", "5,8"]
     result = list(partition_pages(pages, 3, max_pages=10))
     assert result == ["0,2-3", "5,8-9", "10"]
+
+
+def test_table_extraction(pseudo_json_results: List[dict], tmpdir: Path) -> None:
+    tables = extract_tables_from_json_results(pseudo_json_results, tmpdir)
+    assert len(tables) == 1
+    for table in tables:
+        assert Path(table).exists()
+        with open(table) as t:
+            assert t.read() == pseudo_json_results[0]["pages"][0]["items"][0]["csv"]