my-aws-helpers 5.2.0__tar.gz → 5.4.0__tar.gz

{my_aws_helpers-5.2.0 → my_aws_helpers-5.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: my_aws_helpers
-Version: 5.2.0
+Version: 5.4.0
 Summary: AWS Helpers
 Home-page: https://github.com/JarrodMccarthy/aws_helpers.git
 Author: Jarrod McCarthy

{my_aws_helpers-5.2.0 → my_aws_helpers-5.4.0}/my_aws_helpers/bedrock.py

@@ -11,7 +11,12 @@ from enum import Enum
 import pymupdf
 import concurrent.futures
 from dataclasses import dataclass
-from my_aws_helpers.s3 import S3Serialiser
+from my_aws_helpers.s3 import S3Serialiser, BaseS3Object, BaseS3Queries, S3, S3Location
+
+from my_aws_helpers.logging import select_powertools_logger
+
+
+logger = select_powertools_logger("aws-helpers-bedrock")
 
 
 class ImageType(str, Enum):
@@ -44,7 +49,7 @@ class TokenUsage:
 
 
 @dataclass
-class OCRResult:
+class OCRResult(BaseS3Object):
     content: List[Dict[str, str]]
     token_usage: TokenUsage
     page_number: int
@@ -59,12 +64,60 @@ class OCRResult:
 
     @classmethod
     def from_s3_representation(cls, obj: dict) -> OCRResult:
+        obj['token_usage'] = TokenUsage.from_dict(obj.get("token_usage", {})),
         return cls(**obj)
 
     def to_s3_representation(self) -> dict:
         obj = copy(vars(self))
+        obj['token_usage'] = S3Serialiser.object_serialiser(obj=obj['token_usage'])
         return S3Serialiser.object_serialiser(obj=obj)
 
+    def get_save_location(self, bucket_name: str) -> S3Location:
+        pass
+
+
+class OCRResultQueries(BaseS3Queries):
+    def __init__(self, s3_client: S3, bucket_name: str):
+        super().__init__(s3_client=s3_client, bucket_name=bucket_name)
+
+    def save_ocr_result_to_s3(
+        self, ocr_result: OCRResult, save_location: S3Location
+    ) -> Optional[S3Location]:
+        try:
+            obj = ocr_result.to_s3_representation()
+            return self.s3_client.save_dict_to_s3(
+                content=obj,
+                s3_location=save_location,
+            )
+        except Exception as e:
+            logger.exception(f"Failed to save ocr result to s3 due to {e}")
+        return None
+
+    def _concurrent_s3_read(
+        self, locations: List[S3Location], max_workers: int = 10
+    ) -> List[OCRResult]:
+        results: List[OCRResult] = list()
+        futures = list()
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+            for loc in locations:
+                future = executor.submit(
+                    self.s3_client.read_dict_from_s3,
+                    s3_location=loc,
+                )
+                futures.append(future)
+        for f in futures:
+            results.append(f.result())
+        results = [r for r in results if r is not None]
+        return results
+
+    def get_listing_by_key_prefix(self, prefix: str) -> List[OCRResult]:
+        locations = self.s3_client.list_objects_by_prefix(
+            bucket_name=self.bucket_name, prefix=prefix
+        )
+        objects = self._concurrent_s3_read(listing_locations=locations)
+        ocr_results = [OCRResult.from_s3_representation(obj=obj) for obj in objects]
+        return ocr_results
+
 
 class Bedrock:
     def __init__(

{my_aws_helpers-5.2.0 → my_aws_helpers-5.4.0}/my_aws_helpers.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: my-aws-helpers
-Version: 5.2.0
+Version: 5.4.0
 Summary: AWS Helpers
 Home-page: https://github.com/JarrodMccarthy/aws_helpers.git
 Author: Jarrod McCarthy

{my_aws_helpers-5.2.0 → my_aws_helpers-5.4.0}/setup.py

@@ -3,7 +3,7 @@ from setuptools import find_namespace_packages, setup
 
 base_path = os.path.abspath(os.path.dirname(__file__))
 
-version = "5.2.0"
+version = "5.4.0"
 
 core = [
     "boto3==1.34.36",