from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
from airflow.models import BaseOperator
[docs]
class BedrockInvokeOperator(BaseOperator):
ui_color = "#A6E6A6"
[docs]
@staticmethod
def get_table_name(dataset_name):
return f"c_ai_{dataset_name}"
@apply_defaults
def __init__(
self,
s3_prefix,
dataset_name,
input_dataset,
mode="overwrite",
*args,
**kwargs,
):
super(BedrockInvokeOperator, self).__init__(
*args,
**kwargs,
)
self.input_dataset = input_dataset
self.s3_prefix = s3_prefix
self.dataset_name = BedrockInvokeOperator.get_table_name(dataset_name)
self.mode = mode
self.map_name = None
[docs]
def get_client(self):
raise NotImplementedError
[docs]
def invoke_agent(self, input_text, client, sessionId):
raise NotImplementedError
[docs]
def execute(self, context):
import boto3
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 }}"
client = self.get_client()
session_id = str(uuid.uuid4())
data_to_save = []
if isinstance(self.input_dataset, str):
dataset = wr.s3.read_parquet(self.input_dataset)
for idx, row in dataset.iterrows():
response = self.invoke_agent(row, client, session_id)
data_to_save.append(
{"question": row["question"], "response": response, "id": row["id"]}
)
elif isinstance(self.input_dataset, dict):
response = self.invoke_agent(self.input_dataset, client, session_id)
data_to_save.append({"question": "single", "response": response, "id": 0})
df = pd.DataFrame(data_to_save)
save_path = f"s3://{Config.DATA_LAKE_BUCKET}/{Config.CURATED_PREFIX}/{self.s3_prefix}/{self.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,
dtype={"question": "string", "response": "string", "id": "string"},
)
class BedrockAgentOperator(BedrockInvokeOperator):
ui_color = "#A6E6A6"
@apply_defaults
def __init__(
self,
agent_id,
alias_id,
s3_prefix,
dataset_name,
input_dataset=None,
*args,
**kwargs,
):
self.map_name = f"{agent_id}-{alias_id}"
self.agent_id = agent_id
self.alias_id = alias_id
self.input_dataset = input_dataset
self.s3_prefix = s3_prefix
super(BedrockAgentOperator, self).__init__(
s3_prefix=s3_prefix,
input_dataset=input_dataset,
dataset_name=dataset_name,
*args,
**kwargs,
)
print(f"agent_id: {self.agent_id}, alias_id: {self.alias_id}")
def get_client(self):
import boto3
from botocore.config import Config as BotocoreConfig
config = BotocoreConfig(read_timeout=600)
return boto3.client(
service_name="bedrock-agent-runtime",
config=config,
)
def invoke_agent(self, input, client, sessionId):
response = client.invoke_agent(
agentId=self.agent_id,
agentAliasId=self.alias_id,
inputText=input["question"],
sessionId=sessionId,
)
# completion = ""
# for event in response.get("completion"):
# chunk = event["chunk"]
# completion = completion + chunk["bytes"].decode()
# return completion
events = [event for event in response.get("completion")]
chunk = [event for event in events if "chunk" in event] # probably just one
# Process response
completion_text = "".join(
event_chunk["chunk"]["bytes"].decode() for event_chunk in chunk
)
print(f"Bedrock response: {completion_text}")
return completion_text
class BedrockFileSummaryOperator(BaseOperator):
ui_color = "#A6E6A6"
@apply_defaults
def __init__(
self,
model_id,
file_s3_path,
s3_bucket_name,
user_prompt="""
summarize the following document. Each summary should be in this format: Title, Summery, Key points, Published date, Author, Source (link). Only return the results WITHOUT saying anything else ".
""",
inference_config={"maxTokens": 4096, "temperature": 0.7, "topP": 0.9},
*args,
**kwargs,
):
self.map_name = f"{file_s3_path}"
self.model_id = model_id
self.inference_config = inference_config
self.file_s3_path = file_s3_path
self.s3_bucket_name = s3_bucket_name
self.user_prompt = user_prompt
super(BedrockFileSummaryOperator, self).__init__(
*args,
**kwargs,
)
def get_client(self):
import boto3
from botocore.config import Config as BotocoreConfig
config = BotocoreConfig(read_timeout=600)
return boto3.client(
service_name="bedrock-runtime",
config=config,
)
def invoke_agent(self, conversation, client):
response = client.converse(
modelId=self.model_id,
messages=conversation,
inferenceConfig=self.inference_config,
)
return response["output"]["message"]["content"][0]["text"]
def execute(self, context):
import boto3
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 }}"
bedrock_client = self.get_client()
s3_client = boto3.client("s3")
response = s3_client.get_object(
Bucket=self.s3_bucket_name, Key=self.file_s3_path
)
object_metadata = response.get("Metadata", {})
source_url = object_metadata.get("source_url", "")
print("Source URL:", source_url)
bin = response["Body"].read()
modified_prompt = (
self.user_prompt
if not source_url
else f"{self.user_prompt}. Use the following url as the source link if applicable: {source_url}"
)
try:
conversation = [
{
"role": "user",
"content": [
{
"text": modified_prompt,
},
{
"document": {
"format": (
"pdf" if self.file_s3_path.endswith(".pdf") else "html"
),
"name": "string",
"source": {"bytes": bin},
},
},
],
}
]
response_text = self.invoke_agent(conversation, bedrock_client)
#do something with the response
except Exception as e:
print(f"ERROR: Can't invoke model. Reason: {e}")
class BedrockModelOperator(BedrockInvokeOperator):
ui_color = "#A6E6A6"
@apply_defaults
def __init__(
self,
model_id,
s3_prefix,
dataset_name,
input_dataset=None,
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
super(BedrockModelOperator, self).__init__(
s3_prefix=s3_prefix,
dataset_name=dataset_name,
input_dataset_path=input_dataset,
*args,
**kwargs,
)
print(f"Model ID: {self.model_id}")
def get_client(self):
import boto3
from botocore.config import Config as BotocoreConfig
config = BotocoreConfig(read_timeout=600)
return boto3.client(
service_name="bedrock-runtime",
config=config,
)
def invoke_agent(self, input, client, sessionId):
conversation = [
{
"role": "user",
"content": [
{
"text": input["question"],
}
],
}
]
response = client.converse(
modelId=self.model_id,
messages=conversation,
inferenceConfig=self.inference_config,
)
return response["output"]["message"]["content"][0]["text"]