Source code for operators.bedrock_pdf_extractor

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config

from airflow.models import BaseOperator
from util import Util




class BedrockPDFExtractorOperator(BaseOperator):
    ui_color = "#A6E6A6"

    @apply_defaults
    def __init__(
        self,
        model_id,
        output_s3_prefix,
        output_athena_dataset_name,
        output_schema: dict,
        input_files_s3_path,
        input_files_s3_bucket_name,
        mode="overwrite",
        inference_config={"maxTokens": 4096, "temperature": 0.7, "topP": 0.9},
        *args,
        **kwargs,
    ):
        
        self.map_name = f"{model_id}"
        self.model_id = model_id
        self.inference_config = inference_config
        self.output_schema = output_schema

        self.input_files_s3_path = input_files_s3_path
        self.input_files_s3_bucket_name = input_files_s3_bucket_name
        self.output_athena_dataset_name = output_athena_dataset_name
        self.output_s3_prefix = output_s3_prefix
        self.mode = mode
        super().__init__(*args, **kwargs)




    def get_llm(self):
        import boto3
        from botocore.config import Config as BotocoreConfig
        from langchain_aws import ChatBedrockConverse

        config = BotocoreConfig(read_timeout=600)
        llm = ChatBedrockConverse(
            model=self.model_id,
            temperature=self.inference_config.get("temperature"),
            top_p=self.inference_config.get("topP"),
            max_tokens=self.inference_config.get("maxTokens"),
            config=config,
        )
        structured_llm = llm.with_structured_output(self.output_schema)
        return structured_llm




    def invoke_llm(self, input, llm, sessionId):
        from langchain_core.messages import HumanMessage
        import os
        import boto3

        s3_client = boto3.client("s3")
        file_obj = s3_client.get_object(
            Bucket=self.input_files_s3_bucket_name, Key=input["Key"]
        )
        encoded_string = file_obj["Body"].read()
        file_name = os.path.basename(input["Key"])
        message = HumanMessage(
            content=[
                {
                    "text": f"I have a certificate document that I would like to process. useful information file_name:{file_name}, s3_key:{input['Key']}",
                },
                {
                    "document": {
                        "format": "pdf",
                        "name": "file_name",
                        "source": {"bytes": encoded_string},
                    },
                },
            ],
        )
        ai_msg = llm.invoke([message])
        return ai_msg




    def execute(self, context):

        import awswrangler as wr
        import uuid
        import pandas as pd

        # add map_index_template
        context["params"]["map_name"] = self.map_name
        context["map_index_template"] = "{{ params.map_name }}"

        structured_llm = self.get_llm()
        session_id = str(uuid.uuid4())

        list_documents = Util.get_existing_s3_files(
            self.s3_bucket_name, self.s3_prefix, file_type='pdf',return_raw_objects=True
        )
        processed_results = []
        for document in list_documents:
            ai_msg = self.invoke_llm(document, structured_llm, session_id)
            print(ai_msg)
            results=ai_msg.model_dump()
            results['s3_file']=document['Key']
            processed_results.append(results)

        df = pd.DataFrame(processed_results)
        save_path = f"s3://{Config.DATA_LAKE_BUCKET}/{Config.CURATED_PREFIX}/{self.output_s3_prefix}/{self.output_athena_dataset_name}"
        
        wr.s3.to_parquet(
            df=df,
            path=save_path,
            database=Config.DATA_LAKE_DB,
            dataset=True,
            index=False,
            mode=self.mode,
            table=self.dataset_name,
        )