from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
import boto3
import json
import time
from botocore.exceptions import ClientError
[docs]
class BedrockAutomationBaseOperator(BaseOperator):
@apply_defaults
def __init__(
self,
input_pdf_key,
output_uri,
blueprint,
blueprint_version,
blueprint_stage,
max_wait_time=120,
check_interval=10,
**kwargs,
):
super().__init__(**kwargs)
self.input_pdf_key = input_pdf_key
self.output_uri = output_uri
self.bedrock_client = None
self.blueprint = blueprint
self.blueprint_version = blueprint_version
self.blueprint_stage = blueprint_stage
# Subclasses handling larger documents raise these; the defaults keep
# the original behaviour for existing callers.
self.max_wait_time = max_wait_time
self.check_interval = check_interval
[docs]
def get_data_automation_profile_arn(self):
sts_client = boto3.client("sts")
account_id = sts_client.get_caller_identity()["Account"]
region = boto3.Session().region_name
return f"arn:aws:bedrock:{region}:{account_id}:data-automation-profile/us.data-automation-v1"
[docs]
def execute(self, context):
self.create_bedrock_client()
job_id = self.submit_automation_job()
if not job_id:
return
response = self.wait_for_job_completion(
job_id,
max_wait_time=self.max_wait_time,
check_interval=self.check_interval,
)
if not response or response.get("status") != "Success":
self.log.error(f"Job failed: {response.get('status')}")
return
process_model_output = self.process_model_output(response)
if not process_model_output:
self.log.error("No results found in the model output")
return
return self.post_process(context, process_model_output, response)
[docs]
def post_process(self, context, process_model_output, response):
raise NotImplementedError(
"Post-processing method is not implemented. Please implement this method in your subclass."
)
[docs]
def process_model_output(self, response):
output_location = response.get("outputConfiguration", {}).get("s3Uri", "")
result_data = self.uri_to_json(output_location)
if not result_data:
self.log.error("No results found")
return
output_metadata = result_data.get("output_metadata", [])
if not output_metadata:
self.log.error("No output metadata found")
return
segment_metadata = output_metadata[0].get("segment_metadata", [])
custom_output_path = next(
(
segment["custom_output_path"]
for segment in segment_metadata
if segment.get("custom_output_status") == "MATCH"
),
None,
)
if not custom_output_path:
self.log.error("No custom output path found")
return
document_results = self.uri_to_json(custom_output_path)
if not document_results:
self.log.error("No results found")
return
return document_results
[docs]
def invoke_data_automation_async(self,payload, region="us-east-1"):
"""
Call AWS Bedrock InvokeDataAutomationAsync API using the correct headers. Remove this after boto3 is updated
"""
import requests
import json
import boto3
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
# Get AWS credentials
session = boto3.Session()
credentials = session.get_credentials()
if not credentials:
raise Exception("No AWS credentials found. Configure your credentials first.")
url = f"https://bedrock-data-automation-runtime.{region}.amazonaws.com/"
headers = {
"Content-Type": "application/x-amz-json-1.1",
"X-Amz-Target": "AmazonBedrockKeystoneRuntimeService.InvokeDataAutomationAsync",
"X-Amz-User-Agent": "aws-sdk-python/custom",
"X-Bedrock-Data-Insights-Consumer": "Python-Client"
}
request = AWSRequest(
method="POST",
url=url,
headers=headers,
data=json.dumps(payload)
)
SigV4Auth(credentials, 'bedrock', region).add_auth(request)
prepared = request.prepare()
print(f"Making request to: {prepared.url}")
try:
response = requests.post(
prepared.url,
headers=dict(prepared.headers),
data=prepared.body,
timeout=300
)
print(f"Response Status: {response.status_code}")
if response.status_code == 200:
return response.json()
else:
print(f"Error Response: {response.text}")
response.raise_for_status()
except requests.exceptions.RequestException as e:
print(f"Request failed: {e}")
if hasattr(e, 'response') and e.response is not None:
print(f"Error response body: {e.response.text}")
raise
[docs]
def submit_automation_job(self):
try:
data_automation_profile_arn = self.get_data_automation_profile_arn()
payload = {
"blueprints": [
{
"blueprintArn": self.blueprint,
"version": self.blueprint_version,
"stage": self.blueprint_stage,
}
],
"dataAutomationProfileArn": data_automation_profile_arn,
"inputConfiguration": {
"s3Uri": self.input_pdf_key
},
"notificationConfiguration": {
"eventBridgeConfiguration": {
"eventBridgeEnabled": True
}
},
"outputConfiguration": {
"s3Uri": self.output_uri
},
"tags": [
{
"key": "Name",
"value": "test"
}
]
}
response=self.invoke_data_automation_async(payload)
return response["invocationArn"]
except ClientError as e:
self.log.error(f"Error processing document: {e}")
return None
[docs]
def create_bedrock_client(self):
import boto3
self.bedrock_client = boto3.client(
"bedrock-data-automation-runtime",
config=boto3.session.Config(parameter_validation=False),
)
return self.bedrock_client
[docs]
def check_job_status(self, job_id):
try:
self.create_bedrock_client()
response = self.bedrock_client.get_data_automation_status(
invocationArn=job_id
)
return response["status"], response
except ClientError as e:
self.log.error(f"Error checking job status: {e}")
return None, None
[docs]
def wait_for_job_completion(self, job_id, max_wait_time=120, check_interval=10):
start_time = time.time()
while (time.time() - start_time) < max_wait_time:
status, response = self.check_job_status(job_id)
if status == "Success":
return response
elif status in ["Created", "InProgress"]:
self.log.info(f"Job status: {status}. Waiting for completion...")
elif status in ["Created", "InProgress", "ServiceError", "ClientError"]:
self.log.error(f"Job failed with status: {status}")
return None
time.sleep(check_interval)
self.log.error("Max wait time exceeded")
return None
[docs]
def uri_to_json(self, uri):
import boto3
import json
s3_client = boto3.client("s3")
bucket = uri.split("/")[2] if uri else None
prefix = "/".join(uri.split("/")[3:]) if uri else None
if not bucket or not prefix:
self.log.error("Output location information not available")
return None
try:
objects = s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)
for obj in objects.get("Contents", []):
key = obj["Key"]
if key.endswith(".json"):
response = s3_client.get_object(Bucket=bucket, Key=key)
return json.loads(response["Body"].read().decode("utf-8"))
except Exception as e:
self.log.error(f"Error retrieving results: {e}")
return None