my-aws-helpers 2.6.0__tar.gz → 2.6.2__tar.gz

{my_aws_helpers-2.6.0 → my_aws_helpers-2.6.2}/PKG-INFO

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

my_aws_helpers-2.6.2/my_aws_helpers/bedrock.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+import boto3
+from botocore.config import Config
+import json
+import time
+import os
+import io
+from typing import Optional, List, Dict
+from enum import Enum
+import pymupdf
+import concurrent.futures
+from dataclasses import dataclass
+from my_aws_helpers.logging import select_powertools_logger
+
+
+logger = select_powertools_logger('bedrock-boy')
+
+logger.info("Got logger")
+
+class PromptType(str, Enum):
+    json = "json_system_prompt.txt"
+    markdown = "markdown_system_prompt.txt"
+
+
+@dataclass
+class TokenUsage:
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, int]) -> TokenUsage:
+        return cls(
+            input_tokens=data.get('inputTokens', 0),
+            output_tokens=data.get('outputTokens', 0),
+            total_tokens=data.get('totalTokens', 0),
+        )
+
+
+@dataclass
+class OCRResult:
+    content: List[Dict[str, str]]
+    token_usage: TokenUsage
+    page_number: int
+
+    @classmethod
+    def from_dict(cls, data: Dict) -> OCRResult:
+        return cls(
+            content = data.get("content", []),
+            token_usage = TokenUsage.from_dict(data.get("token_usage", {})),
+            page_number = data.get("page_number", 0)
+        )
+
+
+class Bedrock:
+    def __init__(
+        self,
+        model_id: str = "apac.anthropic.claude-3-5-sonnet-20241022-v2:0",
+    ):
+
+        self.session = Bedrock._set_session_params()
+        region_name = "ap-southeast-2"
+        if self.session is None:
+            self.session = boto3.Session(region_name = region_name)
+
+        
+        custom_config = Config(
+            retries={
+                'max_attempts': 2,  # Total attempts = 1 initial + 1 retry
+                'mode': 'standard',  # or 'adaptive'
+            }
+        )      
+        self.client = self.session.client("bedrock-runtime", region_name=region_name, config = custom_config)
+        self.model_id = model_id
+
+    @staticmethod
+    def _set_session_params():
+        try:
+            aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"]
+            aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"]
+            aws_session_token = os.environ["AWS_SESSION_TOKEN"]
+            region_name = os.environ["AWS_DEFAULT_REGION"]
+            return boto3.Session(
+                aws_access_key_id=aws_access_key_id,
+                aws_secret_access_key=aws_secret_access_key,
+                aws_session_token=aws_session_token,
+                region_name=region_name
+            )
+        except Exception as e:
+            logger.exception(e)
+            return None
+
+    def _get_prompt(self, prompt_type: str) -> Optional[str]:
+        if prompt_type not in list(PromptType):
+            raise Exception(f"Error: Invalid prompt type")
+        
+        path = os.path.join(os.path.dirname(__file__), "prompts", prompt_type)
+        try:
+            with open(path, "r") as f:
+                prompt = f.read()
+                return prompt
+        except Exception as e:
+            logger.exception(f"Failed to get {prompt_type} prompt due to {e}")
+        return None
+
+    def _ocr(
+        self, 
+        prompt: str, 
+        image_bytes: bytes,
+        page_number: Optional[int] = 0
+    ) -> Optional[OCRResult]:
+        system_prompt = [{"text": prompt}]
+        message = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "image": {
+                            "format": "png",
+                            "source": {
+                                "bytes": image_bytes,
+                            }
+                        }
+                    }
+                ]
+            }
+        ]
+        retries = 3
+        for i in range(retries):
+            logger.info(f"Attempt number {i} for {self.model_id} converse")
+            try:
+                response = self.client.converse(modelId = self.model_id, messages = message, system = system_prompt)
+                if response['ResponseMetadata']['HTTPStatusCode'] == 200:
+                    break
+            except Exception as e:
+                logger.exception(f"Error during conversation due to {e}")
+                if i >= len(retries) - 1: raise Exception(e)
+                continue
+        
+        result = {}
+        result["content"] = json.loads(response["output"]["message"]["content"][0]["text"])
+        result["token_usage"] = response["usage"]
+        result["page_number"] = page_number
+        return OCRResult.from_dict(data = result)
+    
+    def _parallel_ocr(
+            self, 
+            image_bytes_list: List[bytes], 
+            prompt_type: str,
+            max_workers: int = 10,
+    ):
+        results = list()
+        prompt = self._get_prompt(prompt_type=prompt_type)
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+            logger.info("Some log")
+            execution_futures = {
+                executor.submit(
+                    self._ocr, 
+                    prompt = prompt,
+                    image_bytes = img,
+                ): img for img in image_bytes_list
+            }
+            for future in concurrent.futures.as_completed(execution_futures):
+                result = future.result()
+                if result:
+                    results.append(result)
+        return results
+    
+    def get_ocr_result(
+        self,
+        pdf_bytes: io.BytesIO, 
+        prompt_type: str,
+        zoom: int = 7,
+    ) -> List[OCRResult]:
+        logger.info("Getting OCR Results")
+        try:
+            document = pymupdf.open(stream=pdf_bytes, filetype="pdf")
+        except Exception as e:
+            logger.exception(f"Failed to open pdf due to {e}")
+            return []
+
+        pages: List[pymupdf.Page] = [p for p in document]
+
+        image_bytes_list: List[bytes] = list()
+        for i, p in enumerate(pages):
+            try:
+                image_bytes: bytes = p.get_pixmap(matrix = pymupdf.Matrix(zoom, zoom)).tobytes("png")
+                image_bytes_list.append(image_bytes)
+            except Exception as e:
+                logger.error(f"Could not get pix map for page {i}")
+                continue
+        prompt = self._get_prompt(prompt_type=prompt_type)
+        logger.info("Got Prompt")
+        results = list()
+        for i, image_bytes in enumerate(image_bytes_list):
+            logger.info(f"Starting OCR for page: {i}")
+            results.append(self._ocr(image_bytes=image_bytes, prompt=prompt))
+        return results
+            

{my_aws_helpers-2.6.0 → my_aws_helpers-2.6.2}/my_aws_helpers.egg-info/PKG-INFO

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

{my_aws_helpers-2.6.0 → my_aws_helpers-2.6.2}/my_aws_helpers.egg-info/SOURCES.txt

@@ -3,6 +3,7 @@ setup.cfg
 setup.py
 my_aws_helpers/api.py
 my_aws_helpers/auth.py
+my_aws_helpers/bedrock.py
 my_aws_helpers/cognito.py
 my_aws_helpers/dynamo.py
 my_aws_helpers/errors.py
@@ -16,4 +17,5 @@ my_aws_helpers.egg-info/dependency_links.txt
 my_aws_helpers.egg-info/requires.txt
 my_aws_helpers.egg-info/top_level.txt
 my_aws_helpers.egg-info/zip-safe
+my_aws_helpers/prompts/__init__.py
 tests/test_event.py

{my_aws_helpers-2.6.0 → my_aws_helpers-2.6.2}/setup.py

@@ -3,7 +3,7 @@ from setuptools import find_namespace_packages, setup
 
 base_path = os.path.abspath(os.path.dirname(__file__))
 
-version = "2.6.0"
+version = "2.6.2"
 
 setup(
     name="my_aws_helpers",
@@ -23,5 +23,6 @@ setup(
     zip_safe = True,
     install_requires = [
         "boto3==1.34.36"
-    ]
+    ],
+    include_package_data=True,
 )