save and load parsed pdf

Author: lsarra-ax

src/axiomatic/client.py

@@ -1,5 +1,6 @@
 import base64
 import requests
+import os
 from typing import Dict
 
 from .base_client import BaseClient, AsyncBaseClient
@@ -9,27 +10,31 @@
 class Axiomatic(BaseClient):
 
     def __init__(self, *args, **kwargs):
+        if "timeout" not in kwargs:
+            kwargs["timeout"] = 600
         super().__init__(*args, **kwargs)
 
         self.document_helper = DocumentHelper(self)
 
 
 class DocumentHelper:
 
+    _ax_client: Axiomatic
+
     def __init__(self, ax_client: Axiomatic):
-        self.ax_client = ax_client
+        self._ax_client = ax_client
 
     def pdf_from_url(self, url: str) -> MdResponse:
         """Download a PDF document from a URL and parse it into a Markdown response."""
         file = requests.get(url)
-        response = self.ax_client.document.parse(file=file.content)
+        response = self._ax_client.document.parse(file=file.content)
         return response.content
 
     def pdf_from_file(self, path: str) -> MdResponse:
         """Open a PDF document from a file path and parse it into a Markdown response."""
         with open(path, "rb") as f:
             file = f.read()
-        response = self.ax_client.document.parse(file=file)
+        response = self._ax_client.document.parse(file=file)
         return response.content
 
     def plot_b64_images(self, images: Dict[str, str]):
@@ -60,5 +65,28 @@ def navigate_image(change):
         display(layout)
         display_base64_image(current_index[0])
 
+    def save_parsed_pdf(self, response: MdResponse, path: str):
+        """Save a parsed PDF response to a file."""
+        os.makedirs(path, exist_ok=True)
+        for img_name, img in response.images.items():
+            with open(os.path.join(path, f"{img_name}.png"), "wb") as f:
+                f.write(base64.b64decode(img))
+
+        with open(os.path.join(path, "text.md"), "w") as f:
+            f.write(response.markdown)
+
+    def load_parsed_pdf(self, path: str) -> MdResponse:
+        """Load a parsed PDF response from a file."""
+        with open(os.path.join(path, "text.md"), "r") as f:
+            markdown = f.read()
+
+        images = {}
+        for img_name in os.listdir(path):
+            if img_name.endswith((".png")):
+                with open(os.path.join(path, img_name), "rb") as img_file:
+                    images[img_name] = base64.b64encode(img_file.read()).decode("utf-8")
+
+        return MdResponse(markdown=markdown, images=images)
+
 
 class AsyncAxiomatic(AsyncBaseClient): ...