Close #LGVISIUM-74: Update the API based on the new s3 specs & feedback

src/app/api/v1/endpoints/create_pngs.py

@@ -38,7 +38,7 @@ def create_pngs(aws_filename: Path):
             pix = page.get_pixmap(matrix=fitz.Matrix(3, 3))
             png_filename = f"{filename}-{page_number + 1}.png"
             png_path = f"/tmp/{png_filename}"  # Local path to save the PNG
-            s3_bucket_png_path = f"pngs/{png_filename}"
+            s3_bucket_png_path = f"dataextraction/{png_filename}"
 
             pix.save(png_path)
 

src/app/common/aws.py

@@ -26,7 +26,7 @@ def load_pdf_from_aws(filename: Path) -> fitz.Document:
     """
     # Load the PDF from the S3 object
     try:
-        data = load_data_from_aws(filename, "pdfs")
+        data = load_data_from_aws(filename)
         pdf_document = fitz.open(stream=data, filetype="pdf")
     except Exception:
         raise HTTPException(
@@ -45,7 +45,7 @@ def load_png_from_aws(filename: Path) -> np.ndarray:
     Returns:
         ndarray: The loaded PNG image.
     """
-    data = load_data_from_aws(filename, "pngs")
+    data = load_data_from_aws(filename, "dataextraction")
 
     # Convert the PNG data to an image using PIL
     image = Image.open(io.BytesIO(data))
@@ -54,21 +54,21 @@ def load_png_from_aws(filename: Path) -> np.ndarray:
     return np.array(image)
 
 
-def load_data_from_aws(filename: Path, format: str) -> bytes:
+def load_data_from_aws(filename: Path, prefix: str = "") -> bytes:
     """Load a document from AWS S3.
 
     Args:
         filename (str): The filename of the PNG image.
-        format (str): The format of the file.
+        prefix (str): The prefix of the file in the bucket.
 
     Returns:
         bytes: The loaded PNG image.
     """
     # Check if the PNG exists in S3
     try:
-        png_object = s3_client.get_object(Bucket=config.bucket_name, Key=str(format / filename))
+        png_object = s3_client.get_object(Bucket=config.bucket_name, Key=str(prefix / filename))
     except s3_client.exceptions.NoSuchKey:
-        raise HTTPException(status_code=404, detail=f"Document {format + filename} not found in S3 bucket.") from None
+        raise HTTPException(status_code=404, detail=f"Document {prefix + filename} not found in S3 bucket.") from None
 
     # Load the PNG from the S3 object
     try:

tests/test_create_pngs.py

@@ -13,9 +13,9 @@
 from botocore.exceptions import ClientError
 from fastapi.testclient import TestClient
 
-TEST_PDF_KEY = "pdfs/sample.pdf"
+TEST_PDF_KEY = "sample.pdf"
 TEST_PDF_PATH = Path(__file__).parent.parent / "example" / "example_borehole_profile.pdf"
-TEST_PNG_KEY = "pngs/sample-1.png"
+TEST_PNG_KEY = "dataextraction/sample-1.png"
 TEST_PNG_PATH = Path(__file__).parent.parent / "example" / "sample-1.png"
 
 

tests/test_data_extraction_from_bbox.py

@@ -16,9 +16,9 @@
 from app.common.schemas import ExtractDataRequest, FormatTypes
 from fastapi.testclient import TestClient
 
-TEST_PDF_KEY = Path("pdfs/sample.pdf")
+TEST_PDF_KEY = Path("sample.pdf")
 TEST_PDF_PATH = Path(__file__).parent.parent / "example" / "example_borehole_profile.pdf"
-TEST_PNG_KEY = Path("pngs/sample-1.png")
+TEST_PNG_KEY = Path("dataextraction/sample-1.png")
 TEST_PNG_PATH = Path(__file__).parent.parent / "example" / "sample-1.png"