Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce Amazon Bedrock service #38602

Merged
merged 2 commits into from
Mar 30, 2024

Conversation

ferruzzi
Copy link
Contributor

Includes the basic doc page, hook, operator, unit tests, and system test.

Currently only supports invoking one of the prebuilt "foundation" models. Custom model support is coming next.

Sample DAG:

from airflow.models import DAG
from airflow.providers.amazon.aws.operators.bedrock import BedrockInvokeModelOperator
from airflow.utils.dates import days_ago

PROMPT = "What are the first ten digits of pi?"

with DAG(
    dag_id="bedrock_testing",
    schedule_interval=None,
    start_date=days_ago(1)
) as dag:

    invoke_llama = BedrockInvokeModelOperator(
        task_id="invoke_llama",
        model_id="meta.llama2-13b-chat-v1",
        input_data={"prompt": PROMPT},
    )

Log output from the above sample DAG run:

image

Manually tested as a DAG in Breeze:
image

This is the generated doc page, for easier viewing:
image
The link at the bottom of that screenshot beings you here

Includes the basic doc page, hook, operator, unit tests, and system test.
@ferruzzi
Copy link
Contributor Author

@Taragolis - I added you to the review list since you saw the previous attempt at this, don't feel obligated though.

Copy link
Contributor

@Taragolis Taragolis left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some minor things

airflow/providers/amazon/aws/operators/bedrock.py Outdated Show resolved Hide resolved
airflow/providers/amazon/aws/operators/bedrock.py Outdated Show resolved Hide resolved
airflow/providers/amazon/aws/operators/bedrock.py Outdated Show resolved Hide resolved
Copy link
Contributor

@eladkal eladkal left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@eladkal eladkal merged commit 0723a8f into apache:main Mar 30, 2024
40 checks passed
@ephraimbuddy ephraimbuddy added the changelog:skip Changes that should be skipped from the changelog (CI, tests, etc..) label Mar 30, 2024
mathiaHT pushed a commit to mathiaHT/airflow that referenced this pull request Apr 4, 2024
* Introduce Amazon Bedrock service
utkarsharma2 pushed a commit to astronomer/airflow that referenced this pull request Apr 22, 2024
* Introduce Amazon Bedrock service
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
area:providers area:system-tests changelog:skip Changes that should be skipped from the changelog (CI, tests, etc..) kind:documentation provider:amazon-aws AWS/Amazon - related issues
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants