diff --git a/.github/ISSUE_TEMPLATE/1-airflow_bug_report.yml b/.github/ISSUE_TEMPLATE/1-airflow_bug_report.yml index 60c117f437d32..724f5a0b87ce3 100644 --- a/.github/ISSUE_TEMPLATE/1-airflow_bug_report.yml +++ b/.github/ISSUE_TEMPLATE/1-airflow_bug_report.yml @@ -115,6 +115,7 @@ body: - akeyless - alibaba - amazon + - anthropic - apache-beam - apache-cassandra - apache-drill diff --git a/.github/boring-cyborg.yml b/.github/boring-cyborg.yml index f83c78cb39ee2..98e38346e98a3 100644 --- a/.github/boring-cyborg.yml +++ b/.github/boring-cyborg.yml @@ -30,6 +30,9 @@ labelPRBasedOnFilePath: provider:amazon: - providers/amazon/** + provider:anthropic: + - providers/anthropic/** + provider:apache-beam: - providers/apache/beam/** diff --git a/airflow-core/docs/extra-packages-ref.rst b/airflow-core/docs/extra-packages-ref.rst index eb0acd31df1cb..c2bb02ba55dbd 100644 --- a/airflow-core/docs/extra-packages-ref.rst +++ b/airflow-core/docs/extra-packages-ref.rst @@ -243,6 +243,8 @@ These are extras that add dependencies needed for integration with external serv +---------------------+-----------------------------------------------------+-----------------------------------------------------+ | amazon | ``pip install 'apache-airflow[amazon]'`` | Amazon Web Services | +---------------------+-----------------------------------------------------+-----------------------------------------------------+ +| anthropic | ``pip install 'apache-airflow[anthropic]'`` | Anthropic hooks and operators | ++---------------------+-----------------------------------------------------+-----------------------------------------------------+ | asana | ``pip install 'apache-airflow[asana]'`` | Asana hooks and operators | +---------------------+-----------------------------------------------------+-----------------------------------------------------+ | atlassian-jira | ``pip install 'apache-airflow[atlassian-jira]'`` | Jira hooks and operators | diff --git a/airflow-core/src/airflow/provider.yaml.schema.json b/airflow-core/src/airflow/provider.yaml.schema.json index 50dacbc79634c..b9d12bc29754d 100644 --- a/airflow-core/src/airflow/provider.yaml.schema.json +++ b/airflow-core/src/airflow/provider.yaml.schema.json @@ -99,9 +99,11 @@ "ai", "ai-inference", "alibaba", + "anthropic", "apache", "aws", "azure", + "claude", "dbt", "embeddings", "gcp", diff --git a/dev/breeze/doc/images/output-commands.svg b/dev/breeze/doc/images/output-commands.svg index 6dfe7e09c52c4..4891c30ed5124 100644 --- a/dev/breeze/doc/images/output-commands.svg +++ b/dev/breeze/doc/images/output-commands.svg @@ -367,7 +367,7 @@ -Usage:breeze[OPTIONS] [COMMAND] [ARGS]... +Usage:breeze[OPTIONS] COMMAND [ARGS]... ╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ │--answer -aForce answer to questions. (y | n | q | yes | no | quit)│ diff --git a/dev/breeze/doc/images/output_build-docs.svg b/dev/breeze/doc/images/output_build-docs.svg index 172f0f6576eb2..21272cc62296c 100644 --- a/dev/breeze/doc/images/output_build-docs.svg +++ b/dev/breeze/doc/images/output_build-docs.svg @@ -240,7 +240,7 @@ Usage:                                                                                                                 breeze build-docs                                                                                                      -[OPTIONS] [airbyte | akeyless | alibaba | all-providers | amazon | apache-airflow | apache-airflow-ctl |               +[OPTIONS] [airbyte | akeyless | alibaba | all-providers | amazon | anthropic | apache-airflow | apache-airflow-ctl |   apache-airflow-providers | apache.cassandra | apache.drill | apache.druid | apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | cloudant | cncf.kubernetes diff --git a/dev/breeze/doc/images/output_build-docs.txt b/dev/breeze/doc/images/output_build-docs.txt index a734b81b0c766..48b11286965d5 100644 --- a/dev/breeze/doc/images/output_build-docs.txt +++ b/dev/breeze/doc/images/output_build-docs.txt @@ -1 +1 @@ -de29604dc4d7031b9a26a1e0350d7806 +28c9a42f6270c2fa37c15cca78a8c38b diff --git a/dev/breeze/doc/images/output_release-management_add-back-references.svg b/dev/breeze/doc/images/output_release-management_add-back-references.svg index 88ba1696907ae..814f0d28ea328 100644 --- a/dev/breeze/doc/images/output_release-management_add-back-references.svg +++ b/dev/breeze/doc/images/output_release-management_add-back-references.svg @@ -149,7 +149,7 @@ Usage:                                                                                                                 breeze release-management add-back-references                                                                          -[OPTIONS] [airbyte | akeyless | alibaba | all-providers | amazon | apache-airflow | apache-airflow-ctl |               +[OPTIONS] [airbyte | akeyless | alibaba | all-providers | amazon | anthropic | apache-airflow | apache-airflow-ctl |   apache-airflow-providers | apache.cassandra | apache.drill | apache.druid | apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | cloudant | cncf.kubernetes diff --git a/dev/breeze/doc/images/output_release-management_add-back-references.txt b/dev/breeze/doc/images/output_release-management_add-back-references.txt index 9d5b0b1d2fef4..8562a83f6d21c 100644 --- a/dev/breeze/doc/images/output_release-management_add-back-references.txt +++ b/dev/breeze/doc/images/output_release-management_add-back-references.txt @@ -1 +1 @@ -4b44b7c52911adf1c1bdfb87e10e873b +a100217703deb9fddf97c53ef7e959a5 diff --git a/dev/breeze/doc/images/output_release-management_classify-provider-changes.svg b/dev/breeze/doc/images/output_release-management_classify-provider-changes.svg index e7b78f7153ee1..cef6e4d3b1d14 100644 --- a/dev/breeze/doc/images/output_release-management_classify-provider-changes.svg +++ b/dev/breeze/doc/images/output_release-management_classify-provider-changes.svg @@ -162,17 +162,17 @@ Usage:                                                                                                                 breeze release-management classify-provider-changes                                                                    -[OPTIONS] [airbyte | akeyless | alibaba | amazon | apache.cassandra | apache.drill | apache.druid | apache.flink |     -apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig |  -apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | -cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging | common.sql |          -databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp -| git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins | keycloak |        -microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc | openai |         -openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres |     -presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake |    -sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb |          -zendesk]...                                                                                                            +[OPTIONS] [airbyte | akeyless | alibaba | amazon | anthropic | apache.cassandra | apache.drill | apache.druid |        +apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy +| apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | +clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging |        +common.sql | databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab |   +facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins |    +keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc |       +openai | openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone |       +postgres | presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp |     +snowflake | sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex |    +ydb | zendesk]...                                                                                                      Classify each provider's unreleased changes with hard-coded, high-confidence rules, flagging ambiguous commits as  'needs_llm' for an agent/skill to assess. Outputs JSON - a deterministic alternative to the random '--non-interactive' diff --git a/dev/breeze/doc/images/output_release-management_classify-provider-changes.txt b/dev/breeze/doc/images/output_release-management_classify-provider-changes.txt index 226f5e3c0feab..6e02c700afdf5 100644 --- a/dev/breeze/doc/images/output_release-management_classify-provider-changes.txt +++ b/dev/breeze/doc/images/output_release-management_classify-provider-changes.txt @@ -1 +1 @@ -57466e5cb9470b0d4398724fda2a1825 +4cac13b21eee8b732a46c5a15aec7a4b diff --git a/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.svg b/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.svg index 1ce32a4d94152..2f9692be54fa0 100644 --- a/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.svg +++ b/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.svg @@ -151,17 +151,17 @@ Usage:                                                                                                                 breeze release-management generate-issue-content-providers                                                             -[OPTIONS] [airbyte | akeyless | alibaba | amazon | apache.cassandra | apache.drill | apache.druid | apache.flink |     -apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig |  -apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | -cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging | common.sql |          -databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp -| git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins | keycloak |        -microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc | openai |         -openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres |     -presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake |    -sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb |          -zendesk]...                                                                                                            +[OPTIONS] [airbyte | akeyless | alibaba | amazon | anthropic | apache.cassandra | apache.drill | apache.druid |        +apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy +| apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | +clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging |        +common.sql | databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab |   +facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins |    +keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc |       +openai | openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone |       +postgres | presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp |     +snowflake | sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex |    +ydb | zendesk]...                                                                                                      Generates content for issue to test the release. diff --git a/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.txt b/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.txt index 5231c239ad495..6afb289b5854f 100644 --- a/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.txt +++ b/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.txt @@ -1 +1 @@ -e2d5118e3339b53ce6e87dc6be6ae053 +bd80bf2dd63a27111b8c5745df5dfb86 diff --git a/dev/breeze/doc/images/output_release-management_generate-providers-metadata.svg b/dev/breeze/doc/images/output_release-management_generate-providers-metadata.svg index 40358ab62d9ca..64efb0af6b65b 100644 --- a/dev/breeze/doc/images/output_release-management_generate-providers-metadata.svg +++ b/dev/breeze/doc/images/output_release-management_generate-providers-metadata.svg @@ -1,4 +1,4 @@ - + │--provider-id     Provider_id to generate metadata for. If not specified, metadata for all packages will be       │ │generated. This is debug-only option. When this option is specified, the metadata produced for  │ │the distribution will not be written to the file, but printed via console. (airbyte | akeyless |│ -│alibaba | amazon | apache.beam | apache.cassandra | apache.drill | apache.druid | apache.flink |│ -│apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | │ -│apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb |│ -│asana | atlassian.jira | celery | clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai│ -│| common.compat | common.io | common.messaging | common.sql | databricks | datadog | dbt.cloud |│ -│dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp | git | │ -│github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins | │ -│keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql │ -│| neo4j | odbc | openai | openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | │ -│papermill | pgvector | pinecone | postgres | presto | qdrant | redis | salesforce | samba | │ -│segment | sendgrid | sftp | singularity | slack | smtp | snowflake | sqlite | ssh | standard | │ -│tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb | zendesk)│ -│--provider-versionProvider version to generate metadata for. Only used when --provider-id is specified. Limits    │ -│running metadata generation to only this version of the provider. (TEXT)│ -╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ -╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--dry-run-DIf dry-run is set, commands are only printed, not executed.│ -│--verbose-vPrint verbose information about performed steps.│ -│--help   -hShow this message and exit.│ -╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +│alibaba | amazon | anthropic | apache.beam | apache.cassandra | apache.drill | apache.druid | │ +│apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | │ +│apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | │ +│apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | cloudant | cncf.kubernetes│ +│| cohere | common.ai | common.compat | common.io | common.messaging | common.sql | databricks | │ +│datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | │ +│facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica│ +│| jdbc | jenkins | keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | │ +│microsoft.winrm | mongo | mysql | neo4j | odbc | openai | openfaas | openlineage | opensearch | │ +│opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres | presto | qdrant | │ +│redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake │ +│| sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate |│ +│yandex | ydb | zendesk)│ +│--provider-versionProvider version to generate metadata for. Only used when --provider-id is specified. Limits    │ +│running metadata generation to only this version of the provider. (TEXT)│ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ +│--dry-run-DIf dry-run is set, commands are only printed, not executed.│ +│--verbose-vPrint verbose information about performed steps.│ +│--help   -hShow this message and exit.│ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ diff --git a/dev/breeze/doc/images/output_release-management_generate-providers-metadata.txt b/dev/breeze/doc/images/output_release-management_generate-providers-metadata.txt index 088e7af08dc1f..7c14e04e63151 100644 --- a/dev/breeze/doc/images/output_release-management_generate-providers-metadata.txt +++ b/dev/breeze/doc/images/output_release-management_generate-providers-metadata.txt @@ -1 +1 @@ -93f9b27ce19f188a7d7b70cfcabf8c2b +5865810f5a4b12c9199267b7dd6459b3 diff --git a/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.svg b/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.svg index 3baea7f764abe..27a1d92de9a4b 100644 --- a/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.svg +++ b/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.svg @@ -189,17 +189,17 @@ Usage:                                                                                                                 breeze release-management prepare-provider-distributions                                                               -[OPTIONS] [airbyte | akeyless | alibaba | amazon | apache.cassandra | apache.drill | apache.druid | apache.flink |     -apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig |  -apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | -cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging | common.sql |          -databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp -| git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins | keycloak |        -microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc | openai |         -openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres |     -presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake |    -sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb |          -zendesk]...                                                                                                            +[OPTIONS] [airbyte | akeyless | alibaba | amazon | anthropic | apache.cassandra | apache.drill | apache.druid |        +apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy +| apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | +clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging |        +common.sql | databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab |   +facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins |    +keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc |       +openai | openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone |       +postgres | presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp |     +snowflake | sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex |    +ydb | zendesk]...                                                                                                      Prepare sdist/whl distributions of Airflow Providers. Each provider directory is wiped with `git clean -fdx`  (preserving .venv, .idea, .vscode) before build to keep in-tree generated files out of the artifact. See dev/breeze  diff --git a/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.txt b/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.txt index a4104f3c3a4d2..f054c03806e05 100644 --- a/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.txt +++ b/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.txt @@ -1 +1 @@ -3581abc65c55d4b8a096005c312df783 +a27c1726f5902e5fdb501ecdee226476 diff --git a/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.svg b/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.svg index 571cc1a3c80e1..b9c71f9c81aec 100644 --- a/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.svg +++ b/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.svg @@ -210,17 +210,17 @@ Usage:                                                                                                                 breeze release-management prepare-provider-documentation                                                               -[OPTIONS] [airbyte | akeyless | alibaba | amazon | apache.cassandra | apache.drill | apache.druid | apache.flink |     -apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig |  -apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | -cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging | common.sql |          -databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp -| git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins | keycloak |        -microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc | openai |         -openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres |     -presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake |    -sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb |          -zendesk]...                                                                                                            +[OPTIONS] [airbyte | akeyless | alibaba | amazon | anthropic | apache.cassandra | apache.drill | apache.druid |        +apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy +| apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | +clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging |        +common.sql | databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab |   +facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins |    +keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc |       +openai | openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone |       +postgres | presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp |     +snowflake | sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex |    +ydb | zendesk]...                                                                                                      Prepare CHANGELOG, README and COMMITS information for providers. diff --git a/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.txt b/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.txt index 9ada882dbcbf1..a54080b6bfbb5 100644 --- a/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.txt +++ b/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.txt @@ -1 +1 @@ -0a45b898de7677eed8aa19c7e4f07345 +c33a2f6d00a3a8dbec8b56c1c2d88d54 diff --git a/dev/breeze/doc/images/output_release-management_publish-docs.svg b/dev/breeze/doc/images/output_release-management_publish-docs.svg index b01450b3c1f6b..6c08dc6c566d0 100644 --- a/dev/breeze/doc/images/output_release-management_publish-docs.svg +++ b/dev/breeze/doc/images/output_release-management_publish-docs.svg @@ -188,7 +188,7 @@ Usage:                                                                                                                 breeze release-management publish-docs                                                                                 -[OPTIONS] [airbyte | akeyless | alibaba | all-providers | amazon | apache-airflow | apache-airflow-ctl |               +[OPTIONS] [airbyte | akeyless | alibaba | all-providers | amazon | anthropic | apache-airflow | apache-airflow-ctl |   apache-airflow-providers | apache.cassandra | apache.drill | apache.druid | apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | cloudant | cncf.kubernetes diff --git a/dev/breeze/doc/images/output_release-management_publish-docs.txt b/dev/breeze/doc/images/output_release-management_publish-docs.txt index ca7d9b2eebe9f..9c2563d46b87b 100644 --- a/dev/breeze/doc/images/output_release-management_publish-docs.txt +++ b/dev/breeze/doc/images/output_release-management_publish-docs.txt @@ -1 +1 @@ -30b0621c66436e215e0d34ec9bfd407b +71cd889146297a3b78c71773b6bb5bcd diff --git a/dev/breeze/doc/images/output_sbom_generate-providers-requirements.svg b/dev/breeze/doc/images/output_sbom_generate-providers-requirements.svg index 71bce1509344e..a9d2e173ef2b8 100644 --- a/dev/breeze/doc/images/output_sbom_generate-providers-requirements.svg +++ b/dev/breeze/doc/images/output_sbom_generate-providers-requirements.svg @@ -182,8 +182,8 @@ ╭─ Generate provider requirements flags ───────────────────────────────────────────────────────────────────────────────╮ │--python-versions Comma separate list of Python versions to update sbom from (defaults to all historical python   │ │versions) (3.6 | 3.7 | 3.8 | 3.9 | 3.10 | 3.11 | 3.12 | 3.13 | 3.14)│ -│--provider-id     Provider id to generate the requirements for (airbyte | akeyless | alibaba | amazon | │ -│apache.beam | apache.cassandra | apache.drill | apache.druid | apache.flink | apache.hdfs | │ +│--provider-id     Provider id to generate the requirements for (airbyte | akeyless | alibaba | amazon | anthropic │ +│| apache.beam | apache.cassandra | apache.drill | apache.druid | apache.flink | apache.hdfs | │ │apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | │ │apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | │ │atlassian.jira | celery | clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai | │ diff --git a/dev/breeze/doc/images/output_sbom_generate-providers-requirements.txt b/dev/breeze/doc/images/output_sbom_generate-providers-requirements.txt index ff491c1b8cf35..bb3262922cec4 100644 --- a/dev/breeze/doc/images/output_sbom_generate-providers-requirements.txt +++ b/dev/breeze/doc/images/output_sbom_generate-providers-requirements.txt @@ -1 +1 @@ -4e0a8dab83f1d257bce1cb90b778b009 +84f322f2ffa9ded046c2d3678135a8f7 diff --git a/dev/breeze/doc/images/output_workflow-run_publish-docs.svg b/dev/breeze/doc/images/output_workflow-run_publish-docs.svg index 48033b138d1af..ff1db47a5f239 100644 --- a/dev/breeze/doc/images/output_workflow-run_publish-docs.svg +++ b/dev/breeze/doc/images/output_workflow-run_publish-docs.svg @@ -203,7 +203,7 @@ Usage:                                                                                                                 breeze workflow-run publish-docs                                                                                       -[OPTIONS] [airbyte | akeyless | alibaba | all-providers | amazon | apache-airflow | apache-airflow-ctl |               +[OPTIONS] [airbyte | akeyless | alibaba | all-providers | amazon | anthropic | apache-airflow | apache-airflow-ctl |   apache-airflow-providers | apache.cassandra | apache.drill | apache.druid | apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | cloudant | cncf.kubernetes diff --git a/dev/breeze/doc/images/output_workflow-run_publish-docs.txt b/dev/breeze/doc/images/output_workflow-run_publish-docs.txt index bcdbd541efa08..01d247c731835 100644 --- a/dev/breeze/doc/images/output_workflow-run_publish-docs.txt +++ b/dev/breeze/doc/images/output_workflow-run_publish-docs.txt @@ -1 +1 @@ -01f2b301d8965951b49e2c8cd114700a +6a419bdaebbe36070dfea872c87a2a23 diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py index 8d6bd32e6ceb4..38ccf343cedee 100644 --- a/dev/breeze/src/airflow_breeze/global_constants.py +++ b/dev/breeze/src/airflow_breeze/global_constants.py @@ -828,7 +828,7 @@ def get_airflow_extras(): { "python-version": "3.10", "airflow-version": "2.11.1", - "remove-providers": "common.messaging edge3 fab git keycloak informatica common.ai opensearch", + "remove-providers": "anthropic common.messaging edge3 fab git keycloak informatica common.ai opensearch", "run-unit-tests": "true", }, { diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 0d4ccbfffdd60..eb25e39ebefe9 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -940,6 +940,7 @@ KerberosClient keycloak Keyfile keyfile +keyless KeyManagementServiceClient keyring keyspace @@ -1038,6 +1039,7 @@ maxcompute Maxime MaxRuntimeInSeconds mb +MCP md mem memcached @@ -1046,6 +1048,7 @@ memorystore memray Mesos mesos +MessageBatch metaclass metadatabase metadataStores diff --git a/providers/anthropic/.gitignore b/providers/anthropic/.gitignore new file mode 100644 index 0000000000000..bff2d7629604d --- /dev/null +++ b/providers/anthropic/.gitignore @@ -0,0 +1 @@ +*.iml diff --git a/providers/anthropic/LICENSE b/providers/anthropic/LICENSE new file mode 100644 index 0000000000000..11069edd79019 --- /dev/null +++ b/providers/anthropic/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright [yyyy] [name of copyright owner] + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/providers/anthropic/NOTICE b/providers/anthropic/NOTICE new file mode 100644 index 0000000000000..a51bd9390d030 --- /dev/null +++ b/providers/anthropic/NOTICE @@ -0,0 +1,5 @@ +Apache Airflow +Copyright 2016-2026 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). diff --git a/providers/anthropic/README.rst b/providers/anthropic/README.rst new file mode 100644 index 0000000000000..a42bbbd2c662b --- /dev/null +++ b/providers/anthropic/README.rst @@ -0,0 +1,83 @@ + +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN! + +.. IF YOU WANT TO MODIFY TEMPLATE FOR THIS FILE, YOU SHOULD MODIFY THE TEMPLATE + ``PROVIDER_README_TEMPLATE.rst.jinja2`` IN the ``dev/breeze/src/airflow_breeze/templates`` DIRECTORY + +Package ``apache-airflow-providers-anthropic`` + +Release: ``0.1.0`` + + +`Anthropic `__ provider for Apache Airflow. +Wraps the official Anthropic Python SDK to run the Claude Message Batches API +asynchronously from Airflow, plus direct message and token-counting helpers. + + +Provider package +---------------- + +This is a provider package for ``anthropic`` provider. All classes for this provider package +are in ``airflow.providers.anthropic`` python package. + +You can find package information and changelog for the provider +in the `documentation `_. + +Installation +------------ + +You can install this package on top of an existing Airflow installation (see ``Requirements`` below +for the minimum Airflow version supported) via +``pip install apache-airflow-providers-anthropic`` + +The package supports the following python versions: 3.10,3.11,3.12,3.13,3.14 + +Requirements +------------ + +========================================== ================== +PIP package Version required +========================================== ================== +``apache-airflow`` ``>=3.0.0`` +``apache-airflow-providers-common-compat`` ``>=1.12.0`` +``anthropic`` ``>=0.101.0`` +========================================== ================== + +Cross provider package dependencies +----------------------------------- + +Those are dependencies that might be needed in order to use all the features of the package. +You need to install the specified providers in order to use them. + +You can install such cross-provider dependencies when installing from PyPI. For example: + +.. code-block:: bash + + pip install apache-airflow-providers-anthropic[common.compat] + + +================================================================================================================== ================= +Dependent package Extra +================================================================================================================== ================= +`apache-airflow-providers-common-compat `_ ``common.compat`` +================================================================================================================== ================= + +The changelog for the provider package can be found in the +`changelog `_. diff --git a/providers/anthropic/docs/.latest-doc-only-change.txt b/providers/anthropic/docs/.latest-doc-only-change.txt new file mode 100644 index 0000000000000..2c1ab461a9c8e --- /dev/null +++ b/providers/anthropic/docs/.latest-doc-only-change.txt @@ -0,0 +1 @@ +da9caffdbbeab1917e1cec5726e50af5f14a5206 diff --git a/providers/anthropic/docs/changelog.rst b/providers/anthropic/docs/changelog.rst new file mode 100644 index 0000000000000..57383492e8320 --- /dev/null +++ b/providers/anthropic/docs/changelog.rst @@ -0,0 +1,26 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +``apache-airflow-providers-anthropic`` + +Changelog +--------- + +0.1.0 +..... + +Initial version of the provider. diff --git a/providers/anthropic/docs/commits.rst b/providers/anthropic/docs/commits.rst new file mode 100644 index 0000000000000..2c44b2b758579 --- /dev/null +++ b/providers/anthropic/docs/commits.rst @@ -0,0 +1,34 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + .. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN! + + .. IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE + `PROVIDER_COMMITS_TEMPLATE.rst.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY + + .. THE REMAINDER OF THE FILE IS AUTOMATICALLY GENERATED. IT WILL BE OVERWRITTEN! + +Package apache-airflow-providers-anthropic +------------------------------------------------------ + +`Anthropic `__ + + +This is detailed commit list of changes for versions provider package: ``anthropic``. +For high-level changelog, see :doc:`package information including changelog `. + +.. airflow-providers-commits:: diff --git a/providers/anthropic/docs/conf.py b/providers/anthropic/docs/conf.py new file mode 100644 index 0000000000000..e1ccf1ad0f465 --- /dev/null +++ b/providers/anthropic/docs/conf.py @@ -0,0 +1,27 @@ +# Disable Flake8 because of all the sphinx imports +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Configuration of Providers docs building.""" + +from __future__ import annotations + +import os + +os.environ["AIRFLOW_PACKAGE_NAME"] = "apache-airflow-providers-anthropic" + +from docs.provider_conf import * # noqa: F403 diff --git a/providers/anthropic/docs/connections.rst b/providers/anthropic/docs/connections.rst new file mode 100644 index 0000000000000..26937763900cf --- /dev/null +++ b/providers/anthropic/docs/connections.rst @@ -0,0 +1,106 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/connection:anthropic: + +Anthropic Connection +==================== + +The `Anthropic `__ connection type enables access to the Claude API +through the official Anthropic Python SDK. + +Default Connection IDs +---------------------- + +The Anthropic hook points to the ``anthropic_default`` connection by default. + +Configuring the Connection +-------------------------- + +API Key + Your Anthropic API key. Required for the first-party API (``platform="anthropic"``), + and also used as the Microsoft Foundry API key (``platform="foundry"``, together with + the ``resource`` extra). Not required for the ``bedrock``, ``vertex`` or ``aws`` + platforms, which authenticate through the respective cloud provider's credential chain. + +Base URL (optional) + A custom base URL for the first-party API (for example, an LLM gateway or proxy). + +Extra (optional) + A JSON dictionary of additional parameters. All keys are optional: + + * ``platform`` — which client to build: ``anthropic`` (default), ``bedrock``, ``vertex``, + ``aws`` or ``foundry``. + * ``model`` — default model id used whenever an operator or hook call does not pass + ``model`` (e.g. ``hook.create_message(...)``, ``hook.create_agent(...)``). Set it here + to change the model without editing Dags; falls back to the provider default + (``claude-opus-4-8``). + * ``aws_region`` — AWS region for the ``bedrock`` platform. + * ``project_id`` / ``region`` — GCP project and region for the ``vertex`` platform. + * ``resource`` — Azure resource name for the ``foundry`` platform. + * ``anthropic_client_kwargs`` — a nested dictionary forwarded verbatim to the client + constructor (for example ``timeout``, ``max_retries`` or ``default_headers``). + + For example, to set the client timeout: + + .. code-block:: json + + { + "anthropic_client_kwargs": { + "timeout": 30, + "max_retries": 5 + } + } + +Workload Identity Federation (keyless auth) +------------------------------------------- + +For the first-party ``anthropic`` platform you can authenticate with short-lived OIDC +tokens instead of a static API key, via `Workload Identity Federation +`__. +Two ways: + +* **Configured on the connection** — leave the API Key empty and set a ``workload_identity`` + block in ``extra``: + + .. code-block:: json + + { + "workload_identity": { + "identity_token_file": "/var/run/secrets/anthropic.com/token", + "federation_rule_id": "fdrl_...", + "organization_id": "00000000-0000-0000-0000-000000000000", + "service_account_id": "svac_...", + "workspace_id": "wrkspc_..." + } + } + +* **From the environment** — leave both the API Key and ``extra`` empty; the SDK resolves + credentials from the standard ``ANTHROPIC_FEDERATION_RULE_ID`` / ``ANTHROPIC_ORGANIZATION_ID`` + / ``ANTHROPIC_SERVICE_ACCOUNT_ID`` / ``ANTHROPIC_IDENTITY_TOKEN_FILE`` environment variables. + +.. note:: + + ``ANTHROPIC_API_KEY`` in the worker environment **shadows** federation — unset it where + the worker runs if you intend to use WIF. + +.. note:: + + The Message Batches API, token counting and the Models API are served only by the + first-party Anthropic API (``platform="anthropic"``) and Claude Platform on AWS + (``platform="aws"``). They are **not** available on Amazon Bedrock, Google Vertex AI + or Microsoft Foundry; the hook raises a clear error if you call them on those platforms. diff --git a/providers/anthropic/docs/index.rst b/providers/anthropic/docs/index.rst new file mode 100644 index 0000000000000..57fc40b4e4c39 --- /dev/null +++ b/providers/anthropic/docs/index.rst @@ -0,0 +1,130 @@ + +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +``apache-airflow-providers-anthropic`` +====================================== + + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Basics + + Home + Changelog + Security + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Guides + + Connection types + Operators + + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Resources + + Python API <_api/airflow/providers/anthropic/index> + PyPI Repository + Installing from sources + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: System tests + + System Tests <_api/tests/system/anthropic/index> + +.. THE REMAINDER OF THE FILE IS AUTOMATICALLY GENERATED. IT WILL BE OVERWRITTEN AT RELEASE TIME! + + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Commits + + Detailed list of commits + + +apache-airflow-providers-anthropic package +------------------------------------------------------ + +`Anthropic `__ provider for Apache Airflow. +Wraps the official Anthropic Python SDK to run the Claude Message Batches API +asynchronously from Airflow, plus direct message and token-counting helpers. + + +Release: 0.1.0 + +Provider package +---------------- + +This package is for the ``anthropic`` provider. +All classes for this package are included in the ``airflow.providers.anthropic`` python package. + +Installation +------------ + +You can install this package on top of an existing Airflow installation via +``pip install apache-airflow-providers-anthropic``. +For the minimum Airflow version supported, see ``Requirements`` below. + +Requirements +------------ + +The minimum Apache Airflow version supported by this provider distribution is ``3.0.0``. + +========================================== ================== +PIP package Version required +========================================== ================== +``apache-airflow`` ``>=3.0.0`` +``apache-airflow-providers-common-compat`` ``>=1.12.0`` +``anthropic`` ``>=0.101.0`` +========================================== ================== + +Optional cross provider package dependencies +-------------------------------------------- + +Those are dependencies that might be needed in order to use all the features of the package. +You need to install the specified provider distributions in order to use them. + +You can install such cross-provider dependencies when installing from PyPI. For example: + +.. code-block:: bash + + pip install apache-airflow-providers-anthropic[common.compat] + + +================================================================================================================== ================= +Dependent package Extra +================================================================================================================== ================= +`apache-airflow-providers-common-compat `_ ``common.compat`` +================================================================================================================== ================= + +Downloading official packages +----------------------------- + +You can download officially released packages and verify their checksums and signatures from the +`Official Apache Download site `_ + +* `The apache-airflow-providers-anthropic 0.1.0 sdist package `_ (`asc `__, `sha512 `__) +* `The apache-airflow-providers-anthropic 0.1.0 wheel package `_ (`asc `__, `sha512 `__) diff --git a/providers/anthropic/docs/installing-providers-from-sources.rst b/providers/anthropic/docs/installing-providers-from-sources.rst new file mode 100644 index 0000000000000..a72b45ffaa6e8 --- /dev/null +++ b/providers/anthropic/docs/installing-providers-from-sources.rst @@ -0,0 +1,18 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. include:: /../../../devel-common/src/sphinx_exts/includes/installing-providers-from-sources.rst diff --git a/providers/anthropic/docs/integration-logos/Anthropic.png b/providers/anthropic/docs/integration-logos/Anthropic.png new file mode 100644 index 0000000000000..2c66ccf7963fb Binary files /dev/null and b/providers/anthropic/docs/integration-logos/Anthropic.png differ diff --git a/providers/anthropic/docs/operators/anthropic.rst b/providers/anthropic/docs/operators/anthropic.rst new file mode 100644 index 0000000000000..f22d100ee1910 --- /dev/null +++ b/providers/anthropic/docs/operators/anthropic.rst @@ -0,0 +1,163 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +Anthropic Operators +=================== + +The Anthropic provider runs the `Claude Message Batches API +`__ from Airflow. +Message Batches process many ``messages.create`` requests asynchronously at 50% of +standard cost; most complete within an hour, with a 24-hour SLA — a good fit for +Airflow's deferrable execution model. + +.. note:: + + For interactive, single-call or agentic LLM workloads, prefer the vendor-agnostic + ``apache-airflow-providers-common-ai`` provider with ``model="anthropic:claude-opus-4-8"``. + This provider focuses on the batch/async surface and direct SDK access that the agent + abstraction does not model. + +.. _howto/operator:AnthropicBatchOperator: + +AnthropicBatchOperator +---------------------- + +:class:`~airflow.providers.anthropic.operators.batch.AnthropicBatchOperator` submits a +Message Batch and waits for it to reach the terminal ``ended`` status. In deferrable mode it +releases the worker slot while an +:class:`~airflow.providers.anthropic.triggers.batch.AnthropicBatchTrigger` polls for +completion. + +The operator returns the **batch ID only**. Pull the per-request results with +:meth:`~airflow.providers.anthropic.hooks.anthropic.AnthropicHook.stream_batch_results` and +persist them to object storage — results can be very large and must not be pushed to XCom. +Results are retained for 29 days after the batch is created. + +Parameters +"""""""""" + +* ``requests`` — a list of ``{"custom_id": str, "params": {...}}`` dicts, where ``params`` is a + ``messages.create`` payload (``model``, ``max_tokens``, ``messages``, ...). +* ``conn_id`` — the Anthropic connection ID (default ``anthropic_default``). +* ``deferrable`` — run in deferrable mode (defaults to the ``operators.default_deferrable`` config). +* ``poll_interval`` — seconds between status checks, in both the synchronous and deferrable paths. +* ``timeout`` — seconds to wait for a terminal status; defaults to 24 hours (the batch SLA). +* ``wait_for_completion`` — if ``False``, return the batch ID immediately after submission. +* ``fail_on_partial_error`` — if ``True``, fail the task when any request errored or expired. + Defaults to ``False`` (succeed and log a warning so successful results are not discarded). + +.. warning:: + + A task retry re-submits a **new** batch. Prefer ``retries=0`` on this task. The submitted + ``batch_id`` is pushed to XCom under key ``batch_id`` immediately after submission, so a + crashed run never loses track of an in-flight batch. + +Example +""""""" + +.. exampleinclude:: /../tests/system/anthropic/example_anthropic_batch.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_anthropic_batch] + :end-before: [END howto_operator_anthropic_batch] + +.. _howto/sensor:AnthropicBatchSensor: + +AnthropicBatchSensor +-------------------- + +:class:`~airflow.providers.anthropic.sensors.batch.AnthropicBatchSensor` waits for an +already-submitted batch (by ``batch_id``) to reach a terminal status. Pair it with +``AnthropicBatchOperator(wait_for_completion=False)`` for a fire-and-forget submit followed +by a re-entrant await — because the sensor only polls an existing batch, retrying it never +re-submits, which sidesteps the "retry creates a new batch" hazard of a waiting submit task. + +It applies the same terminal-status policy as the operator (skip on full cancellation, +``fail_on_partial_error`` to fail on errored/expired requests) and supports ``deferrable`` +mode via the shared trigger. + +.. code-block:: python + + from airflow.providers.anthropic.operators.batch import AnthropicBatchOperator + from airflow.providers.anthropic.sensors.batch import AnthropicBatchSensor + + submit = AnthropicBatchOperator( + task_id="submit", + requests=requests, + wait_for_completion=False, # fire-and-forget; recommend retries=0 + ) + wait = AnthropicBatchSensor( + task_id="wait", + batch_id="{{ ti.xcom_pull(task_ids='submit') }}", + deferrable=True, + ) + submit >> wait + +.. _howto/operator:AnthropicAgentSessionOperator: + +AnthropicAgentSessionOperator +----------------------------- + +:class:`~airflow.providers.anthropic.operators.agent.AnthropicAgentSessionOperator` runs a +`Managed Agents `__ session: +Anthropic runs the agent loop server-side while the worker drives a session and waits for it +to finish. Unlike the ``common.ai`` provider (a *local* pydantic-ai loop), the loop and its +tool-execution sandbox run on Anthropic's infrastructure; the worker only orchestrates. + +**Agents and environments are created once** (via +:meth:`~airflow.providers.anthropic.hooks.anthropic.AnthropicHook.create_agent` / +:meth:`~airflow.providers.anthropic.hooks.anthropic.AnthropicHook.create_environment`, the +``ant`` CLI, or the Console) and referenced by ID on every run — the operator never creates +an agent per task. Configure the agent for **autonomous** operation (no client-side custom +tools or ``always_ask`` permission) so the session reaches ``idle`` (turn complete) rather +than blocking on input the operator cannot supply. + +Provide exactly one of ``message`` (a single user turn) or ``outcome`` (a +``user.define_outcome`` rubric the agent iterates against until satisfied). The operator +returns the **session ID only**; pull artifacts the agent wrote to ``/mnt/session/outputs/`` +afterwards via the Files API (``scope_id=``). + +Parameters +"""""""""" + +* ``agent_id`` / ``environment_id`` — IDs of a pre-created agent and environment. +* ``message`` — a single user message to start the session (mutually exclusive with ``outcome``). +* ``outcome`` — a ``user.define_outcome`` payload (``description`` + required ``rubric``, + optional ``max_iterations``); mutually exclusive with ``message``. +* ``conn_id`` — the Anthropic connection ID (default ``anthropic_default``). +* ``deferrable`` — run in deferrable mode (defaults to ``operators.default_deferrable``). +* ``poll_interval`` — seconds between session status checks. +* ``timeout`` — seconds to wait for a terminal status; defaults to 24 hours. +* ``vault_ids`` — vault IDs providing MCP/credential access to the session. +* ``session_resources`` — files, GitHub repos, or memory stores to mount (forwarded to + ``sessions.create`` as ``resources``; renamed to avoid the reserved ``BaseOperator.resources``). +* ``session_kwargs`` — extra keyword arguments forwarded to ``sessions.create``. + +.. note:: + + Completion is detected accurately for both modes. A ``message`` run inspects the + terminal ``session.status_idle`` event's ``stop_reason`` (correlated against the + kickoff event): ``end_turn`` succeeds; ``requires_action`` and ``retries_exhausted`` + raise an error. An ``outcome`` run is judged from the ``outcome_evaluations`` verdict. + The agent must still be configured for autonomous operation (no client-side custom + tools / ``always_ask``). + +.. exampleinclude:: /../tests/system/anthropic/example_anthropic_agent.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_anthropic_agent_session] + :end-before: [END howto_operator_anthropic_agent_session] diff --git a/providers/anthropic/docs/security.rst b/providers/anthropic/docs/security.rst new file mode 100644 index 0000000000000..66c6f79a4ecfc --- /dev/null +++ b/providers/anthropic/docs/security.rst @@ -0,0 +1,38 @@ + + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +Releasing security patches +-------------------------- + +Airflow providers are released independently from Airflow itself and the information about vulnerabilities +is published separately. You can upgrade providers independently from Airflow itself, following the +instructions found in :doc:`apache-airflow:installation/installing-from-pypi`. + +When we release Provider version, the development is always done from the ``main`` branch where we prepare +the next version. The provider uses strict `SemVer `_ versioning policy. Depending on +the scope of the change, Provider will get ''MAJOR'' version upgrade when there are +breaking changes, ``MINOR`` version upgrade when there are new features or ``PATCHLEVEL`` version upgrade +when there are only bug fixes (including security bugfixes) - and this is the only version that receives +security fixes by default, so you should upgrade to latest version of the provider if you want to receive +all released security fixes. + +The only exception to that rule is when we have a critical security fix and good reason to provide an +out-of-band release for the provider, in which case stakeholders in the provider might decide to cherry-pick +and prepare a branch for an older version of the provider following the +`mixed governance model `_ +and requires interested parties to cherry-pick and test the fixes. diff --git a/providers/anthropic/provider.yaml b/providers/anthropic/provider.yaml new file mode 100644 index 0000000000000..6b2e9a39db8f5 --- /dev/null +++ b/providers/anthropic/provider.yaml @@ -0,0 +1,139 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +--- +package-name: apache-airflow-providers-anthropic + +name: Anthropic + +description: | + `Anthropic `__ provider for Apache Airflow. + Wraps the official Anthropic Python SDK to run the Claude Message Batches API + asynchronously from Airflow, plus direct message and token-counting helpers. + +state: ready +lifecycle: incubation +source-date-epoch: 1780426378 + +# Note that those versions are maintained by release manager - do not update them manually +versions: + - 0.1.0 + +integrations: + - integration-name: Anthropic + external-doc-url: https://docs.claude.com/ + logo: /docs/integration-logos/Anthropic.png + how-to-guide: + - /docs/apache-airflow-providers-anthropic/operators/anthropic.rst + tags: + - ai + +hooks: + - integration-name: Anthropic + python-modules: + - airflow.providers.anthropic.hooks.anthropic + +operators: + - integration-name: Anthropic + python-modules: + - airflow.providers.anthropic.operators.batch + - airflow.providers.anthropic.operators.agent + +sensors: + - integration-name: Anthropic + python-modules: + - airflow.providers.anthropic.sensors.batch + +triggers: + - integration-name: Anthropic + python-modules: + - airflow.providers.anthropic.triggers.batch + - airflow.providers.anthropic.triggers.agent + +connection-types: + - hook-class-name: airflow.providers.anthropic.hooks.anthropic.AnthropicHook + hook-name: "Anthropic" + connection-type: anthropic + conn-fields: + platform: + label: Platform + schema: + type: + - string + - 'null' + enum: + - anthropic + - bedrock + - vertex + - aws + - foundry + default: anthropic + description: >- + Which client to build: anthropic (first-party API, default), bedrock (Amazon + Bedrock), vertex (Google Vertex AI), aws (Claude Platform on AWS) or foundry + (Microsoft Foundry). + model: + label: Default Model + schema: + type: + - string + - 'null' + description: >- + Default model id used whenever an operator or hook call does not pass model + (for example hook.create_message(...)). Falls back to claude-opus-4-8. + aws_region: + label: AWS Region + schema: + type: + - string + - 'null' + description: AWS region for the bedrock platform (for example us-east-1). + project_id: + label: GCP Project ID + schema: + type: + - string + - 'null' + description: Google Cloud project id for the vertex platform. + region: + label: GCP Region + schema: + type: + - string + - 'null' + description: Google Cloud region for the vertex platform (for example us-east5). + resource: + label: Azure Resource + schema: + type: + - string + - 'null' + description: Azure resource name for the foundry platform. + ui-field-behaviour: + hidden-fields: + - schema + - port + - login + relabeling: + password: API Key + host: Base URL + placeholders: + extra: >- + {"anthropic_client_kwargs": {"timeout": 30, "max_retries": 5}, + "workload_identity": {"federation_rule_id": "fdrl_...", "organization_id": "...", + "service_account_id": "svac_...", "workspace_id": "wrkspc_...", + "identity_token_file": "/var/run/secrets/anthropic.com/token"}} diff --git a/providers/anthropic/pyproject.toml b/providers/anthropic/pyproject.toml new file mode 100644 index 0000000000000..74154f78bb805 --- /dev/null +++ b/providers/anthropic/pyproject.toml @@ -0,0 +1,134 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN! + +# IF YOU WANT TO MODIFY THIS FILE EXCEPT DEPENDENCIES, YOU SHOULD MODIFY THE TEMPLATE +# `pyproject_TEMPLATE.toml.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY +[build-system] +requires = ["flit_core==3.12.0"] +build-backend = "flit_core.buildapi" + +[project] +name = "apache-airflow-providers-anthropic" +version = "0.1.0" +description = "Provider package apache-airflow-providers-anthropic for Apache Airflow" +readme = "README.rst" +license = "Apache-2.0" +license-files = ['LICENSE', 'NOTICE'] +authors = [ + {name="Apache Software Foundation", email="dev@airflow.apache.org"}, +] +maintainers = [ + {name="Apache Software Foundation", email="dev@airflow.apache.org"}, +] +keywords = [ "airflow-provider", "anthropic", "airflow", "integration" ] +classifiers = [ + "Development Status :: 5 - Production/Stable", + "Environment :: Console", + "Environment :: Web Environment", + "Intended Audience :: Developers", + "Intended Audience :: System Administrators", + "Framework :: Apache Airflow", + "Framework :: Apache Airflow :: Provider", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: System :: Monitoring", +] +requires-python = ">=3.10" + +# The dependencies should be modified in place in the generated file. +# Any change in the dependencies is preserved when the file is regenerated +# Make sure to run ``prek update-providers-dependencies --all-files`` +# After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` +dependencies = [ + "apache-airflow>=3.0.0", + "apache-airflow-providers-common-compat>=1.12.0", + # 0.101.0 is the first release that ships AnthropicAWS, the newest of the platform + # client classes the hook imports at module top. + "anthropic>=0.101.0", +] + +# The optional dependencies should be modified in place in the generated file +# Any change in the dependencies is preserved when the file is regenerated +[project.optional-dependencies] +"bedrock" = ["anthropic[bedrock]>=0.101.0"] +"vertex" = ["anthropic[vertex]>=0.101.0"] +"aws" = ["anthropic[aws]>=0.101.0"] + +[dependency-groups] +dev = [ + "apache-airflow", + "apache-airflow-task-sdk", + "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", + # Additional devel dependencies (do not remove this line and add extra development dependencies) +] + +# To build docs: +# +# uv run --group docs build-docs +# +# To enable auto-refreshing build with server: +# +# uv run --group docs build-docs --autobuild +# +# To see more options: +# +# uv run --group docs build-docs --help +# +docs = [ + "apache-airflow-devel-common[docs]" +] + +[tool.uv.sources] +# These names must match the names as defined in the pyproject.toml of the workspace items, +# *not* the workspace folder paths +apache-airflow = {workspace = true} +apache-airflow-devel-common = {workspace = true} +apache-airflow-task-sdk = {workspace = true} +apache-airflow-providers-common-sql = {workspace = true} +apache-airflow-providers-standard = {workspace = true} + +[project.urls] +"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-anthropic/0.1.0" +"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-anthropic/0.1.0/changelog.html" +"Bug Tracker" = "https://github.com/apache/airflow/issues" +"Source Code" = "https://github.com/apache/airflow" +"Slack Chat" = "https://s.apache.org/airflow-slack" +"Mastodon" = "https://fosstodon.org/@airflow" +"YouTube" = "https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/" + +[project.entry-points."apache_airflow_provider"] +provider_info = "airflow.providers.anthropic.get_provider_info:get_provider_info" + +[tool.flit.module] +name = "airflow.providers.anthropic" + +# Explicit sdist contents so the build does not rely on VCS information +# (flit 4.0 makes --no-use-vcs the default — see https://github.com/pypa/flit/pull/782). +[tool.flit.sdist] +include = [ + "docs/", + "provider.yaml", + "src/airflow/__init__.py", + "src/airflow/providers/__init__.py", + "tests/", +] diff --git a/providers/anthropic/src/airflow/__init__.py b/providers/anthropic/src/airflow/__init__.py new file mode 100644 index 0000000000000..5966d6b1d5261 --- /dev/null +++ b/providers/anthropic/src/airflow/__init__.py @@ -0,0 +1,17 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/providers/anthropic/src/airflow/providers/__init__.py b/providers/anthropic/src/airflow/providers/__init__.py new file mode 100644 index 0000000000000..5966d6b1d5261 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/__init__.py @@ -0,0 +1,17 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/providers/anthropic/src/airflow/providers/anthropic/__init__.py b/providers/anthropic/src/airflow/providers/anthropic/__init__.py new file mode 100644 index 0000000000000..117734e0ae5e9 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/__init__.py @@ -0,0 +1,39 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE +# OVERWRITTEN WHEN PREPARING DOCUMENTATION FOR THE PACKAGES. +# +# IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE +# `PROVIDER__INIT__PY_TEMPLATE.py.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY +# +from __future__ import annotations + +import packaging.version + +from airflow import __version__ as airflow_version + +__all__ = ["__version__"] + +__version__ = "0.1.0" + +if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse( + "3.0.0" +): + raise RuntimeError( + f"The package `apache-airflow-providers-anthropic:{__version__}` needs Apache Airflow 3.0.0+" + ) diff --git a/providers/anthropic/src/airflow/providers/anthropic/exceptions.py b/providers/anthropic/src/airflow/providers/anthropic/exceptions.py new file mode 100644 index 0000000000000..f4e81b868a185 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/exceptions.py @@ -0,0 +1,37 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + + +class AnthropicError(Exception): + """Base class for all Anthropic provider errors.""" + + +class AnthropicBatchJobError(AnthropicError): + """Raised when an Anthropic Message Batch fails or finishes with failed requests.""" + + +class AnthropicBatchTimeout(AnthropicError): + """Raised when an Anthropic Message Batch does not reach a terminal status in time.""" + + +class AnthropicAgentSessionError(AnthropicError): + """Raised when a Managed Agents session terminates or fails.""" + + +class AnthropicAgentSessionTimeout(AnthropicError): + """Raised when a Managed Agents session does not reach a terminal status in time.""" diff --git a/providers/anthropic/src/airflow/providers/anthropic/get_provider_info.py b/providers/anthropic/src/airflow/providers/anthropic/get_provider_info.py new file mode 100644 index 0000000000000..839e94a8626f1 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/get_provider_info.py @@ -0,0 +1,115 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN! +# +# IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE +# `get_provider_info_TEMPLATE.py.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY + + +def get_provider_info(): + return { + "package-name": "apache-airflow-providers-anthropic", + "name": "Anthropic", + "description": "`Anthropic `__ provider for Apache Airflow.\nWraps the official Anthropic Python SDK to run the Claude Message Batches API\nasynchronously from Airflow, plus direct message and token-counting helpers.\n", + "integrations": [ + { + "integration-name": "Anthropic", + "external-doc-url": "https://docs.claude.com/", + "logo": "/docs/integration-logos/Anthropic.png", + "how-to-guide": ["/docs/apache-airflow-providers-anthropic/operators/anthropic.rst"], + "tags": ["ai"], + } + ], + "hooks": [ + { + "integration-name": "Anthropic", + "python-modules": ["airflow.providers.anthropic.hooks.anthropic"], + } + ], + "operators": [ + { + "integration-name": "Anthropic", + "python-modules": [ + "airflow.providers.anthropic.operators.batch", + "airflow.providers.anthropic.operators.agent", + ], + } + ], + "sensors": [ + {"integration-name": "Anthropic", "python-modules": ["airflow.providers.anthropic.sensors.batch"]} + ], + "triggers": [ + { + "integration-name": "Anthropic", + "python-modules": [ + "airflow.providers.anthropic.triggers.batch", + "airflow.providers.anthropic.triggers.agent", + ], + } + ], + "connection-types": [ + { + "hook-class-name": "airflow.providers.anthropic.hooks.anthropic.AnthropicHook", + "hook-name": "Anthropic", + "connection-type": "anthropic", + "conn-fields": { + "platform": { + "label": "Platform", + "schema": { + "type": ["string", "null"], + "enum": ["anthropic", "bedrock", "vertex", "aws", "foundry"], + "default": "anthropic", + }, + "description": "Which client to build: anthropic (first-party API, default), bedrock (Amazon Bedrock), vertex (Google Vertex AI), aws (Claude Platform on AWS) or foundry (Microsoft Foundry).", + }, + "model": { + "label": "Default Model", + "schema": {"type": ["string", "null"]}, + "description": "Default model id used whenever an operator or hook call does not pass model (for example hook.create_message(...)). Falls back to claude-opus-4-8.", + }, + "aws_region": { + "label": "AWS Region", + "schema": {"type": ["string", "null"]}, + "description": "AWS region for the bedrock platform (for example us-east-1).", + }, + "project_id": { + "label": "GCP Project ID", + "schema": {"type": ["string", "null"]}, + "description": "Google Cloud project id for the vertex platform.", + }, + "region": { + "label": "GCP Region", + "schema": {"type": ["string", "null"]}, + "description": "Google Cloud region for the vertex platform (for example us-east5).", + }, + "resource": { + "label": "Azure Resource", + "schema": {"type": ["string", "null"]}, + "description": "Azure resource name for the foundry platform.", + }, + }, + "ui-field-behaviour": { + "hidden-fields": ["schema", "port", "login"], + "relabeling": {"password": "API Key", "host": "Base URL"}, + "placeholders": { + "extra": '{"anthropic_client_kwargs": {"timeout": 30, "max_retries": 5}, "workload_identity": {"federation_rule_id": "fdrl_...", "organization_id": "...", "service_account_id": "svac_...", "workspace_id": "wrkspc_...", "identity_token_file": "/var/run/secrets/anthropic.com/token"}}' + }, + }, + } + ], + } diff --git a/providers/anthropic/src/airflow/providers/anthropic/hooks/__init__.py b/providers/anthropic/src/airflow/providers/anthropic/hooks/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/hooks/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py new file mode 100644 index 0000000000000..dbefdf84d5df4 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py @@ -0,0 +1,612 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import logging +import time +from enum import Enum +from functools import cached_property +from typing import TYPE_CHECKING, Any, cast + +from anthropic import ( + Anthropic, + AnthropicAWS, + AnthropicBedrock, + AnthropicFoundry, + AnthropicVertex, + IdentityTokenFile, + WorkloadIdentityCredentials, +) + +from airflow.providers.anthropic.exceptions import ( + AnthropicAgentSessionError, + AnthropicAgentSessionTimeout, + AnthropicBatchJobError, + AnthropicBatchTimeout, + AnthropicError, +) +from airflow.providers.common.compat.sdk import AirflowSkipException, BaseHook + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + + from anthropic.types import Message + from anthropic.types.beta import ( + BetaEnvironment, + BetaManagedAgentsAgent, + BetaManagedAgentsSession, + environment_create_params, + ) + from anthropic.types.beta.sessions import BetaManagedAgentsEventParams + from anthropic.types.messages import MessageBatch, MessageBatchIndividualResponse + from anthropic.types.messages.batch_create_params import Request + +#: Default model used when an operator or hook caller does not specify one. +#: Prefer configuring the model on the connection so it can be updated without +#: a provider release when this model ID is retired. +DEFAULT_MODEL = "claude-opus-4-8" + +#: Platforms that serve the first-party-only endpoints (Message Batches, token +#: counting, the Models API). Amazon Bedrock, Google Vertex AI and Microsoft +#: Foundry do not serve these, so the hook fails fast rather than surfacing a +#: raw ``404`` from the SDK. +FIRST_PARTY_PLATFORMS = frozenset({"anthropic", "aws"}) + +AnthropicClient = Anthropic | AnthropicBedrock | AnthropicVertex | AnthropicAWS | AnthropicFoundry + +#: Consecutive failed polls tolerated in the synchronous wait helpers before giving up +#: (transient errors). Mirrors the deferrable triggers' tolerance so a single blip does +#: not fail (and cancel/archive) a still-healthy batch or session. +MAX_CONSECUTIVE_POLL_FAILURES = 5 + + +class BatchStatus(str, Enum): + """Top-level ``processing_status`` of an Anthropic Message Batch.""" + + IN_PROGRESS = "in_progress" + CANCELING = "canceling" + ENDED = "ended" + + @classmethod + def is_in_progress(cls, status: str) -> bool: + """ + Return ``True`` while the batch has not reached the terminal ``ended`` status. + + This is broader than the ``in_progress`` value: a ``canceling`` batch is also + non-terminal (cancellation is in flight but the batch has not ended yet), so it + returns ``True`` too. Read the name as "not yet terminal", not "equals the + ``in_progress`` status". + """ + return status != cls.ENDED + + +class SessionStatus(str, Enum): + """Status of a Managed Agents session.""" + + RESCHEDULING = "rescheduling" + RUNNING = "running" + IDLE = "idle" + TERMINATED = "terminated" + + @classmethod + def is_terminal(cls, status: str) -> bool: + """ + Return ``True`` once the session has stopped working. + + ``idle`` means the agent finished its turn (done, for an autonomous run); + ``terminated`` is an unrecoverable failure. Both stop the wait. + """ + return status in (cls.IDLE, cls.TERMINATED) + + +#: ``outcome_evaluations[].result`` values that mean the outcome did NOT succeed. +OUTCOME_FAILURE_RESULTS = frozenset({"failed", "max_iterations_reached", "interrupted"}) + + +def evaluate_session_state( + session: BetaManagedAgentsSession, *, expect_outcome: bool +) -> tuple[bool, str | None, bool]: + """ + Judge a polled session from its object fields alone. + + Returns ``(done, error_message, needs_event_check)``. ``done=False`` means keep + polling. ``needs_event_check=True`` means the session is ``idle`` on a ``message`` + run and the object can't say *why* — the caller must inspect the event log (see + :meth:`AnthropicHook.poll_session_completion`). + + The ``status`` field can't distinguish a genuine ``end_turn`` from ``requires_action`` + or ``retries_exhausted``, nor a just-created ``idle``. For an outcome run the true + verdict is in ``outcome_evaluations`` (judged here, which also defeats the start race). + """ + if session.status == SessionStatus.TERMINATED: + return True, f"Session {session.id} terminated.", False + if session.status != SessionStatus.IDLE: + return False, None, False + if not expect_outcome: + return False, None, True + for evaluation in session.outcome_evaluations: + if evaluation.result == "satisfied": + return True, None, False + if evaluation.result in OUTCOME_FAILURE_RESULTS: + return True, f"Outcome not satisfied for session {session.id}: {evaluation.result}.", False + # idle but no terminal outcome verdict yet (e.g. the run has not started) + return False, None, False + + +def evaluate_batch_counts( + *, + batch_id: str | None, + canceled: int, + errored: int, + expired: int, + succeeded: int, + fail_on_partial_error: bool, +) -> None: + """ + Apply the success/skip/fail policy for a terminal batch's request counts. + + Lives in the hook module so both :class:`AnthropicBatchOperator` and + :class:`~airflow.providers.anthropic.sensors.batch.AnthropicBatchSensor` share it + without an operator/sensor cross-import. Raises ``AirflowSkipException`` for a + fully-cancelled batch, ``AnthropicBatchJobError`` when ``fail_on_partial_error`` and any + request failed, otherwise returns (logging a warning for partial failures). + """ + total = canceled + errored + expired + succeeded + if total and canceled == total: + raise AirflowSkipException(f"Batch {batch_id} was fully cancelled.") + failed = errored + expired + if failed: + message = ( + f"Batch {batch_id} ended with {failed} failed request(s) " + f"(errored={errored}, expired={expired}, succeeded={succeeded})." + ) + if fail_on_partial_error: + raise AnthropicBatchJobError(message) + logger.warning("%s Successful results are still available.", message) + + +class AnthropicHook(BaseHook): + """ + Use the Anthropic SDK to interact with the Claude API. + + The connection's ``password`` is used as the API key and ``host`` as an optional + base URL (for gateways/proxies). The ``extra`` field selects the platform client + and passes platform-specific configuration: + + - ``platform``: one of ``anthropic`` (default), ``bedrock``, ``vertex``, ``aws``, ``foundry``. + - ``model``: default model id used when an operator/hook call omits ``model`` (lets you + change the model without editing Dags); falls back to :data:`DEFAULT_MODEL`. + - ``aws_region``: region for the ``bedrock`` and ``aws`` platforms. + - ``project_id`` / ``region``: project and region for the ``vertex`` platform. + - ``resource``: Azure resource name for the ``foundry`` platform. + - ``anthropic_client_kwargs``: extra keyword arguments forwarded to the client + constructor (e.g. ``timeout``, ``max_retries``, ``default_headers``). + - ``workload_identity``: configure `Workload Identity Federation + `__ + (keyless OIDC auth) with ``identity_token_file``, ``federation_rule_id``, + ``organization_id``, ``service_account_id`` and optional ``workspace_id`` / ``scope``. + + When the ``anthropic`` platform has no API Key and no ``workload_identity`` block, the + client is built with no static credential so the SDK resolves them from the environment + — supporting env-driven Workload Identity Federation and ``ant`` profiles. + + .. seealso:: https://docs.claude.com/en/api/client-sdks + + :param conn_id: :ref:`Anthropic connection id `. + """ + + conn_name_attr = "conn_id" + default_conn_name = "anthropic_default" + conn_type = "anthropic" + hook_name = "Anthropic" + + def __init__(self, conn_id: str = default_conn_name, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.conn_id = conn_id + + @cached_property + def _connection(self): + return self.get_connection(self.conn_id) + + @cached_property + def platform(self) -> str: + """Return the configured platform (defaults to ``anthropic``).""" + return (self._connection.extra_dejson.get("platform") or "anthropic").lower() + + @cached_property + def default_model(self) -> str: + """Default model id — connection ``extra['model']`` if set, else :data:`DEFAULT_MODEL`.""" + return self._connection.extra_dejson.get("model") or DEFAULT_MODEL + + @cached_property + def conn(self) -> AnthropicClient: + """Return the Anthropic client for the configured platform.""" + return self.get_conn() + + def get_conn(self) -> AnthropicClient: + """Build and return the Anthropic client for the configured platform.""" + conn = self._connection + extras = conn.extra_dejson + client_kwargs = dict(extras.get("anthropic_client_kwargs", {})) + platform = self.platform + self.log.debug("Building Anthropic client for platform %r (conn_id=%s)", platform, self.conn_id) + if platform == "bedrock": + return AnthropicBedrock(aws_region=extras.get("aws_region"), **client_kwargs) + if platform == "vertex": + return AnthropicVertex( + project_id=extras.get("project_id"), region=extras.get("region"), **client_kwargs + ) + if platform == "aws": + return AnthropicAWS(aws_region=extras.get("aws_region"), **client_kwargs) + if platform == "foundry": + api_key = client_kwargs.pop("api_key", None) or conn.password + return AnthropicFoundry(api_key=api_key, resource=extras.get("resource"), **client_kwargs) + if platform != "anthropic": + raise AnthropicError( + f"Unknown Anthropic platform {platform!r}. " + "Expected one of: anthropic, bedrock, vertex, aws, foundry." + ) + base_url = client_kwargs.pop("base_url", None) or conn.host or None + wif = extras.get("workload_identity") + if wif: + return Anthropic( + credentials=self._workload_identity_credentials(wif), base_url=base_url, **client_kwargs + ) + api_key = client_kwargs.pop("api_key", None) or conn.password + if api_key: + return Anthropic(api_key=api_key, base_url=base_url, **client_kwargs) + # No static key and no explicit federation config: let the SDK resolve credentials + # from the environment, which supports env-driven Workload Identity Federation + # (ANTHROPIC_FEDERATION_RULE_ID etc.) and ``ant`` profiles. + return Anthropic(base_url=base_url, **client_kwargs) + + @staticmethod + def _workload_identity_credentials(wif: dict[str, Any]) -> WorkloadIdentityCredentials: + """ + Build a WIF credential from the connection ``extra['workload_identity']`` mapping. + + Exchanges a short-lived OIDC token (read from ``identity_token_file``) for an + Anthropic access token. See + https://platform.claude.com/docs/en/manage-claude/workload-identity-federation. + """ + kwargs: dict[str, Any] = { + "identity_token_provider": IdentityTokenFile(wif["identity_token_file"]), + "federation_rule_id": wif["federation_rule_id"], + "organization_id": wif["organization_id"], + "service_account_id": wif["service_account_id"], + } + if wif.get("workspace_id"): + kwargs["workspace_id"] = wif["workspace_id"] + if wif.get("scope"): + kwargs["scope"] = wif["scope"] + return WorkloadIdentityCredentials(**kwargs) + + def _require_first_party(self, feature: str) -> None: + if self.platform not in FIRST_PARTY_PLATFORMS: + raise AnthropicError( + f"{feature} is not available on the {self.platform!r} platform. " + "Use the first-party Anthropic API (platform='anthropic') or " + "Claude Platform on AWS (platform='aws')." + ) + + @property + def _first_party_conn(self) -> Anthropic: + """ + Client cast to the first-party type for endpoints only it exposes. + + Callers must guard with :meth:`_require_first_party` first; the Bedrock/Vertex/ + Foundry clients don't expose ``beta.agents``/``beta.sessions``/``models``. + """ + return cast("Anthropic", self.conn) + + def test_connection(self) -> tuple[bool, str]: + """Test the Anthropic connection.""" + try: + if self.platform in FIRST_PARTY_PLATFORMS: + # Narrowed by the platform guard: only the first-party / AWS clients, + # which expose the Models API, reach this branch. + self._first_party_conn.models.list() + return True, "Connection established!" + # models.list() is not served on bedrock/vertex/foundry; building the + # client validates the configuration without a paid request. + self.get_conn() + return True, f"Connection configured for platform {self.platform!r} (no live check available)." + except Exception as e: + return False, str(e) + + def create_message( + self, + messages: list[dict[str, Any]], + model: str | None = None, + max_tokens: int = 1024, + system: str | None = None, + **kwargs: Any, + ) -> Message: + """ + Create a single message response (one-shot ``messages.create``). + + :param messages: The conversation so far, as a list of message dicts. + :param model: Model ID to use. Defaults to :attr:`default_model` (the connection's + ``extra['model']`` or :data:`DEFAULT_MODEL`). + :param max_tokens: Maximum number of tokens to generate. + :param system: Optional system prompt. + """ + params: dict[str, Any] = { + "model": model or self.default_model, + "max_tokens": max_tokens, + "messages": messages, + **kwargs, + } + if system is not None: + params["system"] = system + return self.conn.messages.create(**params) + + def count_tokens( + self, + messages: list[dict[str, Any]], + model: str | None = None, + system: str | None = None, + **kwargs: Any, + ) -> int: + """Return the number of input tokens the given request would consume.""" + self._require_first_party("Token counting") + params: dict[str, Any] = {"model": model or self.default_model, "messages": messages, **kwargs} + if system is not None: + params["system"] = system + return self.conn.messages.count_tokens(**params).input_tokens + + def create_batch(self, requests: list[dict[str, Any]]) -> MessageBatch: + """ + Submit a Message Batch. + + :param requests: A list of ``{"custom_id": str, "params": {...}}`` dicts, where + ``params`` is a ``messages.create`` payload (``model``, ``max_tokens``, + ``messages``, ...). + """ + self._require_first_party("The Message Batches API") + # ``Request`` is a TypedDict, so the plain dicts callers build match structurally. + return self.conn.messages.batches.create(requests=cast("Iterable[Request]", requests)) + + def get_batch(self, batch_id: str) -> MessageBatch: + """Retrieve a Message Batch by ID.""" + self._require_first_party("The Message Batches API") + return self.conn.messages.batches.retrieve(batch_id) + + def cancel_batch(self, batch_id: str) -> MessageBatch: + """Request cancellation of a Message Batch.""" + self._require_first_party("The Message Batches API") + return self.conn.messages.batches.cancel(batch_id) + + def list_batches(self, **kwargs: Any) -> Any: + """Return a (paginated) list of Message Batches.""" + self._require_first_party("The Message Batches API") + return self.conn.messages.batches.list(**kwargs) + + def stream_batch_results(self, batch_id: str) -> Iterator[MessageBatchIndividualResponse]: + """ + Return a streaming iterator of per-request results, keyed by ``custom_id``. + + Results stream from the API and arrive in **arbitrary order** — key them by + ``result.custom_id``, never by position. Results are available for 29 days + after the batch is created. The result set can be very large: iterate and + persist to object storage; do not materialize it into XCom. + """ + # Return (don't ``yield``) so the platform guard fails fast at call time + # rather than only when the caller starts iterating. + self._require_first_party("The Message Batches API") + return self.conn.messages.batches.results(batch_id) + + def wait_for_batch( + self, batch_id: str, wait_seconds: float = 3, timeout: float = 24 * 60 * 60 + ) -> MessageBatch: + """ + Poll a batch synchronously until it reaches the terminal ``ended`` status. + + :param batch_id: The batch to wait for. + :param wait_seconds: Seconds to sleep between polls. + :param timeout: Maximum seconds to wait before raising :class:`AnthropicBatchTimeout`. + :return: The terminal :class:`~anthropic.types.messages.MessageBatch`. + """ + start = time.monotonic() + consecutive_failures = 0 + while True: + try: + batch = self.get_batch(batch_id) + except Exception as e: + # Tolerate transient poll errors (as the deferrable trigger does) so a + # single blip does not fail — and cancel — a still-running batch whose + # results remain recoverable for 29 days. + consecutive_failures += 1 + if ( + consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES + or time.monotonic() - start > timeout + ): + raise + self.log.warning("Polling batch %s failed (%s); retrying.", batch_id, e) + time.sleep(wait_seconds) + continue + consecutive_failures = 0 + self.log.debug("Batch %s status=%s", batch_id, batch.processing_status) + if not BatchStatus.is_in_progress(batch.processing_status): + return batch + if time.monotonic() - start > timeout: + raise AnthropicBatchTimeout( + f"Batch {batch_id} did not reach a terminal status within {timeout} seconds." + ) + time.sleep(wait_seconds) + + # --- Managed Agents ------------------------------------------------------- + # Agents and environments are persisted, reusable resources: create them once + # (these helpers, the ``ant`` CLI, or a setup script) and store the IDs. The + # operator references those IDs; it never creates an agent per run. + + def create_agent(self, name: str, model: str | None = None, **kwargs: Any) -> BetaManagedAgentsAgent: + """ + Create a (reusable, versioned) Managed Agents agent. One-time setup. + + ``model`` defaults to :attr:`default_model` (the connection's ``extra['model']`` + or :data:`DEFAULT_MODEL`). + """ + self._require_first_party("Managed Agents") + agent = self._first_party_conn.beta.agents.create( + name=name, model=model or self.default_model, **kwargs + ) + self.log.debug("Created agent %s (name=%r, model=%s)", agent.id, name, model or self.default_model) + return agent + + def create_environment( + self, name: str, config: dict[str, Any] | None = None, **kwargs: Any + ) -> BetaEnvironment: + """Create a (reusable) environment for agent sessions. One-time setup.""" + self._require_first_party("Managed Agents") + if config is None: + config = {"type": "cloud", "networking": {"type": "unrestricted"}} + environment = self._first_party_conn.beta.environments.create( + name=name, config=cast("environment_create_params.Config", config), **kwargs + ) + self.log.debug("Created environment %s (name=%r)", environment.id, name) + return environment + + def create_session(self, agent: str, environment_id: str, **kwargs: Any) -> BetaManagedAgentsSession: + """Start a session against a pre-created agent + environment.""" + self._require_first_party("Managed Agents") + return self._first_party_conn.beta.sessions.create( + agent=agent, environment_id=environment_id, **kwargs + ) + + def get_session(self, session_id: str) -> BetaManagedAgentsSession: + """Retrieve a session (carries its current ``status``).""" + self._require_first_party("Managed Agents") + return self._first_party_conn.beta.sessions.retrieve(session_id) + + def send_event(self, session_id: str, event: dict[str, Any]) -> Any: + """Send a single event (e.g. a ``user.message`` or ``user.define_outcome``).""" + self._require_first_party("Managed Agents") + # Event dicts callers build match the SDK's TypedDict union structurally. + return self._first_party_conn.beta.sessions.events.send( + session_id, events=cast("list[BetaManagedAgentsEventParams]", [event]) + ) + + def archive_session(self, session_id: str) -> Any: + """Archive a session (frees the server-side container). Best-effort teardown.""" + self._require_first_party("Managed Agents") + return self._first_party_conn.beta.sessions.archive(session_id) + + def _latest_idle_reason(self, session_id: str, kickoff_event_id: str | None) -> str | None: + """ + Return the ``stop_reason`` of the newest ``session.status_idle`` event, or ``None``. + + Walks the event log newest-first. Returns ``None`` if the kickoff event is the most + recent event (the agent has not responded yet — defeats the start race) or no idle + event is found in the scan window. + """ + # The SDK cursor auto-paginates (page size 20); cap the walk at 100 events so a + # long event log can't make one poll iterate unboundedly. + examined = 0 + for event in self._first_party_conn.beta.sessions.events.list(session_id, order="desc", limit=20): + if kickoff_event_id is not None and event.id == kickoff_event_id: + return None + if event.type == "session.status_idle": + return event.stop_reason.type + examined += 1 + if examined >= 100: + break + return None + + def poll_session_completion( + self, session_id: str, *, expect_outcome: bool = False, kickoff_event_id: str | None = None + ) -> tuple[bool, str | None]: + """ + Return ``(done, error_message)`` for one poll of a session. + + Combines the session object (status / outcome verdict) with the event log + (``stop_reason`` of the latest idle) so a ``message`` run distinguishes genuine + ``end_turn`` completion from ``requires_action`` / ``retries_exhausted``. + """ + session = self.get_session(session_id) + done, error_message, needs_event_check = evaluate_session_state( + session, expect_outcome=expect_outcome + ) + self.log.debug( + "Session %s status=%s done=%s needs_event_check=%s", + session_id, + session.status, + done, + needs_event_check, + ) + if not needs_event_check: + return done, error_message + reason = self._latest_idle_reason(session_id, kickoff_event_id) + if reason is None: + return False, None + if reason == "end_turn": + return True, None + return True, ( + f"Session {session_id} is idle but did not complete ({reason}); " + "configure an autonomous agent or use an outcome run." + ) + + def wait_for_session( + self, + session_id: str, + expect_outcome: bool = False, + kickoff_event_id: str | None = None, + poll_interval: float = 30, + timeout: float = 24 * 60 * 60, + ) -> None: + """ + Poll a session synchronously until it completes. + + :param session_id: The session to wait for. + :param expect_outcome: Whether the session is running a ``user.define_outcome`` loop + (completion judged from ``outcome_evaluations``). + :param kickoff_event_id: ID of the kickoff event, used to correlate the terminal + idle event on a ``message`` run (defeats the start race). + :param poll_interval: Seconds to sleep between polls. + :param timeout: Maximum seconds to wait before raising :class:`AnthropicAgentSessionTimeout`. + """ + start = time.monotonic() + consecutive_failures = 0 + while True: + try: + done, error_message = self.poll_session_completion( + session_id, expect_outcome=expect_outcome, kickoff_event_id=kickoff_event_id + ) + except Exception as e: + # Tolerate transient poll errors (as the deferrable trigger does) so a + # single blip does not fail — and archive — a still-running session. + consecutive_failures += 1 + if ( + consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES + or time.monotonic() - start > timeout + ): + raise + self.log.warning("Polling session %s failed (%s); retrying.", session_id, e) + time.sleep(poll_interval) + continue + consecutive_failures = 0 + if done: + if error_message: + raise AnthropicAgentSessionError(error_message) + return + if time.monotonic() - start > timeout: + raise AnthropicAgentSessionTimeout( + f"Session {session_id} did not reach a terminal status within {timeout} seconds." + ) + time.sleep(poll_interval) diff --git a/providers/anthropic/src/airflow/providers/anthropic/operators/__init__.py b/providers/anthropic/src/airflow/providers/anthropic/operators/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/operators/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py new file mode 100644 index 0000000000000..7b1f5f86e14c1 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py @@ -0,0 +1,228 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from collections.abc import Sequence +from datetime import timedelta +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.anthropic.exceptions import AnthropicAgentSessionError, AnthropicAgentSessionTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook +from airflow.providers.anthropic.triggers.agent import AnthropicAgentSessionTrigger +from airflow.providers.common.compat.sdk import BaseOperator, conf + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class AnthropicAgentSessionOperator(BaseOperator): + """ + Run a Managed Agents session against a pre-created agent and environment. + + Anthropic runs the agent loop server-side; the worker creates a session, sends the + initial instruction, and waits for the session to reach a terminal status. In + deferrable mode it releases the worker slot while a trigger polls the session status. + + Provide exactly one of ``message`` (a single user turn) or ``outcome`` (a + ``user.define_outcome`` rubric that the agent iterates against until satisfied). + + .. important:: + This operator is for **autonomous** agents — configure the agent without + client-side custom tools or an ``always_ask`` permission policy. + + Completion is detected accurately for both modes. A ``message`` run reads the + terminal ``session.status_idle`` event's ``stop_reason`` (correlated against the + kickoff event to avoid a start-race false positive): ``end_turn`` succeeds, while + ``requires_action`` (the agent is blocked on input) and ``retries_exhausted`` + raise an error rather than silently passing. An ``outcome`` run is judged from the + session's ``outcome_evaluations`` verdict (``satisfied`` vs. + ``failed``/``max_iterations_reached``/``interrupted``). + + Agents and environments are created once (see + :meth:`~airflow.providers.anthropic.hooks.anthropic.AnthropicHook.create_agent`), + not per task run. + + Outputs the agent writes to ``/mnt/session/outputs/`` are retrieved afterwards via the + Files API (``scope_id=``); the operator returns the **session ID only**. + + .. seealso:: + For more information, take a look at the guide: + :ref:`howto/operator:AnthropicAgentSessionOperator` + + :param agent_id: ID of a pre-created agent. + :param environment_id: ID of a pre-created environment. + :param message: A single user message to start the session. Mutually exclusive with ``outcome``. + :param outcome: A ``user.define_outcome`` payload (``description``, ``rubric``, + optional ``max_iterations``). Mutually exclusive with ``message``. + :param conn_id: The Anthropic connection ID to use. + :param deferrable: Run the operator in deferrable mode. + :param poll_interval: Seconds between session status checks (both paths). + :param timeout: Seconds to wait for a terminal status. Defaults to 24 hours. In + deferrable mode the trigger enforces this and tears the session down on timeout; + set ``execution_timeout`` only for a shorter hard cap (which preempts that + graceful teardown). + :param vault_ids: Vault IDs providing MCP/credential access to the session. + :param session_resources: Session resources (files, GitHub repos, memory stores). Named + ``session_resources`` to avoid colliding with the reserved ``BaseOperator.resources``; + forwarded to ``sessions.create`` as ``resources``. + :param session_kwargs: Extra keyword arguments forwarded to ``sessions.create``. + """ + + template_fields: Sequence[str] = ("agent_id", "environment_id", "message", "outcome") + + def __init__( + self, + *, + agent_id: str, + environment_id: str, + message: str | None = None, + outcome: dict[str, Any] | None = None, + conn_id: str = AnthropicHook.default_conn_name, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + poll_interval: float = 30, + timeout: float = 24 * 60 * 60, + vault_ids: list[str] | None = None, + session_resources: list[dict[str, Any]] | None = None, + session_kwargs: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + if (message is None) == (outcome is None): + raise ValueError("Provide exactly one of 'message' or 'outcome'.") + if outcome is not None and not {"description", "rubric"} <= outcome.keys(): + raise ValueError("'outcome' must include both 'description' and 'rubric'.") + self.agent_id = agent_id + self.environment_id = environment_id + self.message = message + self.outcome = outcome + self.conn_id = conn_id + self.deferrable = deferrable + self.poll_interval = poll_interval + self.timeout = timeout + self.vault_ids = vault_ids + self.session_resources = session_resources + self.session_kwargs = session_kwargs or {} + self.session_id: str | None = None + + @cached_property + def hook(self) -> AnthropicHook: + """Return an instance of the AnthropicHook.""" + return AnthropicHook(conn_id=self.conn_id) + + def execute(self, context: Context) -> str | None: + create_kwargs: dict[str, Any] = dict(self.session_kwargs) + if self.vault_ids: + create_kwargs["vault_ids"] = self.vault_ids + if self.session_resources: + create_kwargs["resources"] = self.session_resources + session = self.hook.create_session( + agent=self.agent_id, environment_id=self.environment_id, **create_kwargs + ) + self.session_id = session.id + context["ti"].xcom_push(key="session_id", value=session.id) + self.log.info("Started Anthropic session %s for agent %s", session.id, self.agent_id) + + try: + if self.outcome is not None: + response = self.hook.send_event(session.id, {"type": "user.define_outcome", **self.outcome}) + else: + response = self.hook.send_event( + session.id, + {"type": "user.message", "content": [{"type": "text", "text": self.message}]}, + ) + except Exception: + # send_event failed after create_session allocated the container; tear it down. + self._archive_session(session.id) + raise + # Correlate completion against the kickoff event so a message run is not fooled by a + # just-created idle session (the start race). + sent = response.data + kickoff_event_id = sent[-1].id if sent else None + + expect_outcome = self.outcome is not None + if self.deferrable: + # Backstop the deferral slightly beyond the trigger's own end_time so the + # trigger's clean "timeout" event (which tears the session down) wins the race + # rather than a generic AirflowTaskTimeout. A user-set execution_timeout still + # applies as a shorter hard cap (and then teardown is skipped — documented). + self.defer( + timeout=self.execution_timeout or timedelta(seconds=self.timeout + self.poll_interval + 60), + trigger=AnthropicAgentSessionTrigger( + conn_id=self.conn_id, + session_id=session.id, + poll_interval=self.poll_interval, + end_time=time.time() + self.timeout, + expect_outcome=expect_outcome, + kickoff_event_id=kickoff_event_id, + ), + method_name="execute_complete", + ) + + self.log.info("Waiting for session %s to complete", session.id) + try: + self.hook.wait_for_session( + session.id, + expect_outcome=expect_outcome, + kickoff_event_id=kickoff_event_id, + poll_interval=self.poll_interval, + timeout=self.timeout, + ) + except Exception: + # Any failure after the session starts (timeout, SDK 5xx, auth expiry) leaves + # the server-side container running; archive it best-effort before failing. + self._archive_session(session.id) + raise + return session.id + + def execute_complete(self, context: Context, event: Any = None) -> str: + if not event: + raise AnthropicAgentSessionError("Trigger resumed without an event payload.") + # The deferred task is a fresh instance; restore the session id from the event. + self.session_id = event["session_id"] + status = event["status"] + if status == "timeout": + self._archive_session(self.session_id) + raise AnthropicAgentSessionTimeout(event["message"]) + if status == "error": + # The trigger yields "error" when polling gives up while the session may still + # be running; archive it best-effort so its container does not linger. + self._archive_session(self.session_id) + raise AnthropicAgentSessionError(event["message"]) + self.log.info("Session %s completed.", self.session_id) + return self.session_id + + def _archive_session(self, session_id: str | None) -> None: + """Best-effort teardown of the server-side session (frees its container).""" + if not session_id: + return + try: + self.hook.archive_session(session_id) + except Exception as e: + self.log.warning("Failed to archive session %s: %s", session_id, e) + + def on_kill(self) -> None: + """ + Archive the session if the (non-deferred) task is killed. + + Only fires while the worker process is alive, i.e. the synchronous path + (``deferrable=False``). On Airflow 3.3+ a killed deferred task is archived by the + trigger's ``on_kill``. On older Airflow the session of a killed deferred task is not + archived automatically; archive it manually via the hook. + """ + self._archive_session(self.session_id) diff --git a/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py b/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py new file mode 100644 index 0000000000000..7fd6f8f6a88c7 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py @@ -0,0 +1,206 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from collections.abc import Sequence +from datetime import timedelta +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.anthropic.exceptions import AnthropicBatchJobError, AnthropicBatchTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook, evaluate_batch_counts +from airflow.providers.anthropic.triggers.batch import AnthropicBatchTrigger +from airflow.providers.common.compat.sdk import BaseOperator, conf + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class AnthropicBatchOperator(BaseOperator): + """ + Submit an Anthropic Message Batch and wait for it to complete. + + Message Batches process many ``messages.create`` requests asynchronously at 50% of + standard cost; most complete within an hour (24h SLA). This operator submits the + batch and, in deferrable mode, releases the worker slot while a trigger polls for + completion. + + The operator returns the **batch ID only** — never the results. Pull results with + :meth:`~airflow.providers.anthropic.hooks.anthropic.AnthropicHook.stream_batch_results` + and persist them to object storage; results can be very large and must not be pushed + to XCom. Results are retained for 29 days after the batch is created. + + .. note:: + A retry re-submits a brand-new batch. Prefer ``retries=0`` on this task (the + submitted ``batch_id`` is pushed to XCom under key ``batch_id`` immediately, so + a crashed run never loses track of an in-flight batch). + + .. seealso:: + For more information, take a look at the guide: + :ref:`howto/operator:AnthropicBatchOperator` + + :param requests: A list of ``{"custom_id": str, "params": {...}}`` dicts, where + ``params`` is a ``messages.create`` payload (``model``, ``max_tokens``, ``messages``, ...). + :param conn_id: The Anthropic connection ID to use. + :param deferrable: Run the operator in deferrable mode. + :param poll_interval: Seconds between status checks, in both the synchronous and + deferrable paths. + :param timeout: Seconds to wait for the batch to reach a terminal status. Defaults to + 24 hours (the Message Batches SLA). In deferrable mode this also bounds the + deferral; set ``execution_timeout`` only if you want a shorter hard cap (note a + shorter ``execution_timeout`` preempts the graceful cancel-on-timeout path). + :param wait_for_completion: Whether to wait for the batch to complete. If ``False``, + the operator returns the batch ID immediately after submission. + :param fail_on_partial_error: If ``True``, fail the task when any request errored or + expired. Defaults to ``False`` (succeed and log a warning so the successful + results are not discarded). + """ + + template_fields: Sequence[str] = ("requests",) + + def __init__( + self, + requests: list[dict[str, Any]], + conn_id: str = AnthropicHook.default_conn_name, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + poll_interval: float = 60, + timeout: float = 24 * 60 * 60, + wait_for_completion: bool = True, + fail_on_partial_error: bool = False, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.requests = requests + self.conn_id = conn_id + self.deferrable = deferrable + self.poll_interval = poll_interval + self.timeout = timeout + self.wait_for_completion = wait_for_completion + self.fail_on_partial_error = fail_on_partial_error + self.batch_id: str | None = None + + @cached_property + def hook(self) -> AnthropicHook: + """Return an instance of the AnthropicHook.""" + return AnthropicHook(conn_id=self.conn_id) + + def execute(self, context: Context) -> str | None: + if not self.requests: + raise ValueError("AnthropicBatchOperator requires at least one request; got an empty list.") + batch = self.hook.create_batch(self.requests) + self.batch_id = batch.id + # Push immediately so a crash between submit and completion never loses the batch. + context["ti"].xcom_push(key="batch_id", value=batch.id) + self.log.info("Submitted Anthropic Message Batch %s (%d requests)", batch.id, len(self.requests)) + + if not self.wait_for_completion: + return self.batch_id + + if self.deferrable: + self.defer( + # Backstop the deferral slightly beyond the trigger's own end_time so the + # trigger's clean "timeout" event (which cancels the batch) wins over a + # generic AirflowTaskTimeout. A user-set execution_timeout still applies + # as a shorter hard cap. + timeout=self.execution_timeout or timedelta(seconds=self.timeout + self.poll_interval + 60), + trigger=AnthropicBatchTrigger( + conn_id=self.conn_id, + batch_id=self.batch_id, + poll_interval=self.poll_interval, + end_time=time.time() + self.timeout, + ), + method_name="execute_complete", + ) + + self.log.info("Waiting for batch %s to complete", self.batch_id) + try: + batch = self.hook.wait_for_batch( + self.batch_id, wait_seconds=self.poll_interval, timeout=self.timeout + ) + except Exception: + # Any failure after submission (timeout, SDK 5xx, auth expiry) leaves the batch + # running and billing; cancel it best-effort before the task fails. + self.log.warning("Batch %s failed while waiting; requesting cancellation.", self.batch_id) + self._cancel_batch_quietly() + raise + counts = batch.request_counts + self._apply_policy(counts.canceled, counts.errored, counts.expired, counts.succeeded) + return self.batch_id + + def execute_complete(self, context: Context, event: Any = None) -> str: + """ + Resume after the trigger fires. + + The deferred task is a fresh instance, so the batch ID is read from the event, + not ``self.batch_id``. + """ + self.batch_id = event["batch_id"] + status = event["status"] + if status == "timeout": + self.log.warning("Batch %s timed out; requesting cancellation.", self.batch_id) + self._cancel_batch_quietly() + raise AnthropicBatchTimeout(event["message"]) + if status == "error": + # The trigger yields "error" when polling gives up (transient failures + # exhausted or the deadline passed mid-poll) while the batch may still be + # running; cancel it best-effort so it does not keep billing. + self.log.warning("Batch %s errored while polling; requesting cancellation.", self.batch_id) + self._cancel_batch_quietly() + raise AnthropicBatchJobError(event["message"]) + + counts = event.get("request_counts") or {} + self._apply_policy( + counts.get("canceled", 0), + counts.get("errored", 0), + counts.get("expired", 0), + counts.get("succeeded", 0), + ) + self.log.info("%s completed successfully.", self.task_id) + return self.batch_id + + def _apply_policy(self, canceled: int, errored: int, expired: int, succeeded: int) -> None: + evaluate_batch_counts( + batch_id=self.batch_id, + canceled=canceled, + errored=errored, + expired=expired, + succeeded=succeeded, + fail_on_partial_error=self.fail_on_partial_error, + ) + + def on_kill(self) -> None: + """ + Cancel the batch if the (non-deferred) task is killed. + + This only fires while the worker process is alive, i.e. the synchronous path + (``deferrable=False``). On Airflow 3.3+ a killed deferred task is cancelled by the + trigger's ``on_kill``. On older Airflow the batch of a killed deferred task is not + cancelled automatically; cancel it manually via the hook. + """ + if self.batch_id: + self.log.info("on_kill: cancelling Anthropic batch %s", self.batch_id) + self._cancel_batch_quietly() + + def _cancel_batch_quietly(self) -> None: + """Best-effort batch cancellation for the timeout and kill paths.""" + if not self.batch_id: + return + try: + self.hook.cancel_batch(self.batch_id) + except Exception as e: + self.log.warning("Failed to cancel batch %s: %s", self.batch_id, e) diff --git a/providers/anthropic/src/airflow/providers/anthropic/sensors/__init__.py b/providers/anthropic/src/airflow/providers/anthropic/sensors/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/sensors/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/src/airflow/providers/anthropic/sensors/batch.py b/providers/anthropic/src/airflow/providers/anthropic/sensors/batch.py new file mode 100644 index 0000000000000..a5128b5c1b764 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/sensors/batch.py @@ -0,0 +1,124 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from collections.abc import Sequence +from datetime import timedelta +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.anthropic.exceptions import AnthropicBatchJobError, AnthropicBatchTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook, BatchStatus, evaluate_batch_counts +from airflow.providers.anthropic.triggers.batch import AnthropicBatchTrigger +from airflow.providers.common.compat.sdk import BaseSensorOperator, conf + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class AnthropicBatchSensor(BaseSensorOperator): + """ + Wait for an already-submitted Anthropic Message Batch to reach a terminal status. + + Pairs with ``AnthropicBatchOperator(wait_for_completion=False)`` (or any out-of-band + submission) for a fire-and-forget submit + re-entrant await. Because the sensor only + polls an existing ``batch_id``, it is naturally idempotent across retries — unlike a + submit step, retrying it never creates a new batch. + + On a terminal batch it applies the same policy as the operator: a fully-cancelled + batch skips the task, and ``fail_on_partial_error`` controls whether errored/expired + requests fail it. + + .. seealso:: + For more information, take a look at the guide: + :ref:`howto/sensor:AnthropicBatchSensor` + + :param batch_id: The ID of the batch to wait for. + :param conn_id: The Anthropic connection ID to use. + :param deferrable: Run the sensor in deferrable mode (polls via a trigger). + :param fail_on_partial_error: If ``True``, fail when any request errored or expired. + Defaults to ``False`` (succeed and log a warning). + """ + + template_fields: Sequence[str] = ("batch_id",) + + def __init__( + self, + *, + batch_id: str, + conn_id: str = AnthropicHook.default_conn_name, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + fail_on_partial_error: bool = False, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.batch_id = batch_id + self.conn_id = conn_id + self.deferrable = deferrable + self.fail_on_partial_error = fail_on_partial_error + + @cached_property + def hook(self) -> AnthropicHook: + """Return an instance of the AnthropicHook.""" + return AnthropicHook(conn_id=self.conn_id) + + def poke(self, context: Context) -> bool: + batch = self.hook.get_batch(self.batch_id) + if BatchStatus.is_in_progress(batch.processing_status): + return False + counts = batch.request_counts + evaluate_batch_counts( + batch_id=self.batch_id, + canceled=counts.canceled, + errored=counts.errored, + expired=counts.expired, + succeeded=counts.succeeded, + fail_on_partial_error=self.fail_on_partial_error, + ) + return True + + def execute(self, context: Context) -> None: + if self.deferrable: + self.defer( + timeout=timedelta(seconds=self.timeout), + trigger=AnthropicBatchTrigger( + conn_id=self.conn_id, + batch_id=self.batch_id, + poll_interval=self.poke_interval, + end_time=time.time() + self.timeout, + ), + method_name="execute_complete", + ) + super().execute(context) + + def execute_complete(self, context: Context, event: Any = None) -> None: + status = event["status"] + if status == "timeout": + raise AnthropicBatchTimeout(event["message"]) + if status == "error": + raise AnthropicBatchJobError(event["message"]) + counts = event.get("request_counts") or {} + evaluate_batch_counts( + batch_id=event["batch_id"], + canceled=counts.get("canceled", 0), + errored=counts.get("errored", 0), + expired=counts.get("expired", 0), + succeeded=counts.get("succeeded", 0), + fail_on_partial_error=self.fail_on_partial_error, + ) + self.log.info("Batch %s reached a terminal status.", event["batch_id"]) diff --git a/providers/anthropic/src/airflow/providers/anthropic/triggers/__init__.py b/providers/anthropic/src/airflow/providers/anthropic/triggers/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/triggers/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/src/airflow/providers/anthropic/triggers/agent.py b/providers/anthropic/src/airflow/providers/anthropic/triggers/agent.py new file mode 100644 index 0000000000000..2045f4cc2d003 --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/triggers/agent.py @@ -0,0 +1,140 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +import time +from collections.abc import AsyncIterator +from typing import Any + +from airflow.providers.anthropic.hooks.anthropic import MAX_CONSECUTIVE_POLL_FAILURES, AnthropicHook +from airflow.triggers.base import BaseTrigger, TriggerEvent + + +class AnthropicAgentSessionTrigger(BaseTrigger): + """ + Poll a Managed Agents session until it reaches a terminal status. + + :param conn_id: The Anthropic connection ID. + :param session_id: The session to poll. + :param poll_interval: Seconds to sleep between polls. + :param end_time: Wall-clock deadline (``time.time()`` epoch seconds). Wall-clock is + used deliberately: the trigger is serialized and may resume in a different + triggerer process, where a per-process ``time.monotonic()`` value is meaningless. + :param expect_outcome: Whether the session is running a ``user.define_outcome`` loop + (completion is then judged from ``outcome_evaluations``, not from ``idle`` alone). + :param kickoff_event_id: ID of the kickoff event, used to correlate the terminal idle + event on a ``message`` run (defeats the start race). + """ + + def __init__( + self, + conn_id: str, + session_id: str, + poll_interval: float, + end_time: float, + expect_outcome: bool = False, + kickoff_event_id: str | None = None, + ) -> None: + super().__init__() + self.conn_id = conn_id + self.session_id = session_id + self.poll_interval = poll_interval + self.end_time = end_time + self.expect_outcome = expect_outcome + self.kickoff_event_id = kickoff_event_id + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize AnthropicAgentSessionTrigger arguments and class path.""" + return ( + "airflow.providers.anthropic.triggers.agent.AnthropicAgentSessionTrigger", + { + "conn_id": self.conn_id, + "session_id": self.session_id, + "poll_interval": self.poll_interval, + "end_time": self.end_time, + "expect_outcome": self.expect_outcome, + "kickoff_event_id": self.kickoff_event_id, + }, + ) + + async def on_kill(self) -> None: + """ + Archive the session when a user kills the deferred task. + + Runs in the triggerer event loop on Airflow 3.3+ (a no-op override on older + versions, which never call it). Closes the gap the operator's ``on_kill`` leaves + for deferred tasks, which have released their worker slot. + """ + hook = AnthropicHook(conn_id=self.conn_id) + try: + await asyncio.to_thread(hook.archive_session, self.session_id) + self.log.info("on_kill: archived Anthropic session %s", self.session_id) + except Exception as e: + self.log.warning("on_kill: failed to archive session %s: %s", self.session_id, e) + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Poll the session and yield exactly one terminal event.""" + hook = AnthropicHook(conn_id=self.conn_id) + consecutive_failures = 0 + while True: + try: + # poll_session_completion does blocking SDK HTTP calls; run off the event loop. + done, error_message = await asyncio.to_thread( + hook.poll_session_completion, + self.session_id, + expect_outcome=self.expect_outcome, + kickoff_event_id=self.kickoff_event_id, + ) + except Exception as e: + # Tolerate transient polling errors rather than failing the whole wait. + consecutive_failures += 1 + if consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES or time.time() > self.end_time: + yield TriggerEvent({"status": "error", "session_id": self.session_id, "message": str(e)}) + return + self.log.warning("Polling session %s failed (%s); retrying.", self.session_id, e) + await asyncio.sleep(self.poll_interval) + continue + + consecutive_failures = 0 + if done: + if error_message: + yield TriggerEvent( + {"status": "error", "session_id": self.session_id, "message": error_message} + ) + else: + yield TriggerEvent( + { + "status": "success", + "session_id": self.session_id, + "message": f"Session {self.session_id} completed.", + } + ) + return + if time.time() > self.end_time: + yield TriggerEvent( + { + "status": "timeout", + "session_id": self.session_id, + "message": ( + f"Session {self.session_id} did not reach a terminal status " + "before the configured timeout." + ), + } + ) + return + await asyncio.sleep(self.poll_interval) diff --git a/providers/anthropic/src/airflow/providers/anthropic/triggers/batch.py b/providers/anthropic/src/airflow/providers/anthropic/triggers/batch.py new file mode 100644 index 0000000000000..cdb6973b7d9eb --- /dev/null +++ b/providers/anthropic/src/airflow/providers/anthropic/triggers/batch.py @@ -0,0 +1,129 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +import time +from collections.abc import AsyncIterator +from typing import Any + +from airflow.providers.anthropic.hooks.anthropic import ( + MAX_CONSECUTIVE_POLL_FAILURES, + AnthropicHook, + BatchStatus, +) +from airflow.triggers.base import BaseTrigger, TriggerEvent + + +class AnthropicBatchTrigger(BaseTrigger): + """ + Poll an Anthropic Message Batch until it reaches the terminal ``ended`` status. + + :param conn_id: The Anthropic connection ID. + :param batch_id: The batch to poll. + :param poll_interval: Seconds to sleep between polls. + :param end_time: Wall-clock deadline (``time.time()`` epoch seconds) after which a + ``timeout`` event is emitted. Wall-clock is used deliberately: the trigger is + serialized to the metadata DB and may resume in a different triggerer process, + so a per-process ``time.monotonic()`` value would not survive serialization. + """ + + def __init__(self, conn_id: str, batch_id: str, poll_interval: float, end_time: float) -> None: + super().__init__() + self.conn_id = conn_id + self.batch_id = batch_id + self.poll_interval = poll_interval + self.end_time = end_time + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize AnthropicBatchTrigger arguments and class path.""" + return ( + "airflow.providers.anthropic.triggers.batch.AnthropicBatchTrigger", + { + "conn_id": self.conn_id, + "batch_id": self.batch_id, + "poll_interval": self.poll_interval, + "end_time": self.end_time, + }, + ) + + async def on_kill(self) -> None: + """ + Cancel the batch when a user kills the deferred task. + + Runs in the triggerer event loop on Airflow 3.3+ (a no-op override on older + versions, which never call it). Closes the gap the operator's ``on_kill`` leaves + for deferred tasks, which have released their worker slot. + """ + hook = AnthropicHook(conn_id=self.conn_id) + try: + await asyncio.to_thread(hook.cancel_batch, self.batch_id) + self.log.info("on_kill: cancelled Anthropic batch %s", self.batch_id) + except Exception as e: + self.log.warning("on_kill: failed to cancel batch %s: %s", self.batch_id, e) + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Poll the batch status and yield exactly one terminal event.""" + hook = AnthropicHook(conn_id=self.conn_id) + consecutive_failures = 0 + while True: + try: + # get_batch is a blocking SDK HTTP call; run it off the event loop so a + # single poll does not stall every other trigger on this triggerer. + batch = await asyncio.to_thread(hook.get_batch, self.batch_id) + except Exception as e: + # Tolerate transient polling errors rather than failing a (up to 24h) wait. + consecutive_failures += 1 + if consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES or time.time() > self.end_time: + yield TriggerEvent({"status": "error", "batch_id": self.batch_id, "message": str(e)}) + return + self.log.warning("Polling batch %s failed (%s); retrying.", self.batch_id, e) + await asyncio.sleep(self.poll_interval) + continue + + consecutive_failures = 0 + self.log.debug("Batch %s status=%s", self.batch_id, batch.processing_status) + if not BatchStatus.is_in_progress(batch.processing_status): + counts = batch.request_counts + yield TriggerEvent( + { + "status": "success", + "batch_id": self.batch_id, + "message": f"Batch {self.batch_id} has ended.", + "request_counts": { + "succeeded": counts.succeeded, + "errored": counts.errored, + "canceled": counts.canceled, + "expired": counts.expired, + "processing": counts.processing, + }, + } + ) + return + if time.time() > self.end_time: + yield TriggerEvent( + { + "status": "timeout", + "batch_id": self.batch_id, + "message": ( + f"Batch {self.batch_id} did not reach a terminal status " + "before the configured timeout." + ), + } + ) + return + await asyncio.sleep(self.poll_interval) diff --git a/providers/anthropic/tests/conftest.py b/providers/anthropic/tests/conftest.py new file mode 100644 index 0000000000000..f56ccce0a3f69 --- /dev/null +++ b/providers/anthropic/tests/conftest.py @@ -0,0 +1,19 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +pytest_plugins = "tests_common.pytest_plugin" diff --git a/providers/anthropic/tests/system/__init__.py b/providers/anthropic/tests/system/__init__.py new file mode 100644 index 0000000000000..5966d6b1d5261 --- /dev/null +++ b/providers/anthropic/tests/system/__init__.py @@ -0,0 +1,17 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/providers/anthropic/tests/system/anthropic/__init__.py b/providers/anthropic/tests/system/anthropic/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/tests/system/anthropic/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/tests/system/anthropic/example_anthropic_agent.py b/providers/anthropic/tests/system/anthropic/example_anthropic_agent.py new file mode 100644 index 0000000000000..88b8b2cfee79e --- /dev/null +++ b/providers/anthropic/tests/system/anthropic/example_anthropic_agent.py @@ -0,0 +1,65 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from airflow.sdk import dag, task + +ANTHROPIC_CONN_ID = "anthropic_default" + + +@dag(schedule=None, catchup=False) +def anthropic_managed_agent(): + @task + def setup_agent_and_environment() -> dict[str, str]: + # One-time setup helper — in production run this once and store the IDs, do not + # create a fresh agent every Dag run. Here we create them and pass the IDs via XCom. + from airflow.providers.anthropic.hooks.anthropic import AnthropicHook + + hook = AnthropicHook(conn_id=ANTHROPIC_CONN_ID) + agent = hook.create_agent( + name="airflow-research-agent", + system="You are a concise research assistant.", + tools=[{"type": "agent_toolset_20260401"}], + ) + environment = hook.create_environment(name="airflow-agent-env") + return {"agent_id": agent.id, "environment_id": environment.id} + + setup = setup_agent_and_environment() + + # [START howto_operator_anthropic_agent_session] + from airflow.providers.anthropic.operators.agent import AnthropicAgentSessionOperator + + run_agent = AnthropicAgentSessionOperator( + task_id="run_agent", + conn_id=ANTHROPIC_CONN_ID, + agent_id=setup["agent_id"], + environment_id=setup["environment_id"], + message="Summarize the latest stable Apache Airflow release in two sentences.", + deferrable=True, + ) + # [END howto_operator_anthropic_agent_session] + + run_agent + + +anthropic_managed_agent() + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example Dag with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/anthropic/tests/system/anthropic/example_anthropic_batch.py b/providers/anthropic/tests/system/anthropic/example_anthropic_batch.py new file mode 100644 index 0000000000000..117a9209ea886 --- /dev/null +++ b/providers/anthropic/tests/system/anthropic/example_anthropic_batch.py @@ -0,0 +1,81 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import Any + +from airflow.sdk import dag, task + +ANTHROPIC_CONN_ID = "anthropic_default" +MODEL = "claude-opus-4-8" + +POKEMON = ["pikachu", "charmander", "bulbasaur"] + + +@dag(schedule=None, catchup=False) +def anthropic_batch_messages(): + @task + def build_requests(names: list[str]) -> list[dict[str, Any]]: + return [ + { + "custom_id": name, + "params": { + "model": MODEL, + "max_tokens": 256, + "messages": [{"role": "user", "content": f"Describe {name} in one sentence."}], + }, + } + for name in names + ] + + @task + def collect_results(batch_id: str) -> dict[str, str]: + # Results stream from the API unordered; key them by custom_id. For large + # batches, persist to object storage instead of returning via XCom. + from airflow.providers.anthropic.hooks.anthropic import AnthropicHook + + hook = AnthropicHook(conn_id=ANTHROPIC_CONN_ID) + summaries: dict[str, str] = {} + for entry in hook.stream_batch_results(batch_id): + if entry.result.type == "succeeded": + text = next((b.text for b in entry.result.message.content if b.type == "text"), "") + summaries[entry.custom_id] = text + return summaries + + requests = build_requests(POKEMON) + + # [START howto_operator_anthropic_batch] + from airflow.providers.anthropic.operators.batch import AnthropicBatchOperator + + run_batch = AnthropicBatchOperator( + task_id="run_batch", + conn_id=ANTHROPIC_CONN_ID, + requests=requests, + deferrable=True, + ) + # [END howto_operator_anthropic_batch] + + collect_results(batch_id=run_batch.output) + + +anthropic_batch_messages() + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example Dag with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/anthropic/tests/unit/__init__.py b/providers/anthropic/tests/unit/__init__.py new file mode 100644 index 0000000000000..5966d6b1d5261 --- /dev/null +++ b/providers/anthropic/tests/unit/__init__.py @@ -0,0 +1,17 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/providers/anthropic/tests/unit/anthropic/__init__.py b/providers/anthropic/tests/unit/anthropic/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/tests/unit/anthropic/hooks/__init__.py b/providers/anthropic/tests/unit/anthropic/hooks/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/hooks/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py new file mode 100644 index 0000000000000..1f0974d540cc1 --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py @@ -0,0 +1,492 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.providers.anthropic.exceptions import ( + AnthropicAgentSessionError, + AnthropicAgentSessionTimeout, + AnthropicBatchTimeout, + AnthropicError, +) +from airflow.providers.anthropic.hooks.anthropic import ( + DEFAULT_MODEL, + MAX_CONSECUTIVE_POLL_FAILURES, + AnthropicHook, + BatchStatus, + SessionStatus, + evaluate_session_state, +) + +pytest.importorskip("anthropic") + +HOOK_PATH = "airflow.providers.anthropic.hooks.anthropic" + + +def _conn(password="sk-ant-test", host=None, extra=None): + conn = mock.MagicMock() + conn.password = password + conn.host = host + conn.extra_dejson = extra or {} + return conn + + +def _make_hook(extra=None): + """Build a hook with the cached connection + client pre-injected (no real lookup).""" + hook = AnthropicHook() + hook.__dict__["_connection"] = _conn(extra=extra) + hook.__dict__["platform"] = (extra or {}).get("platform", "anthropic").lower() + client = mock.MagicMock() + hook.__dict__["conn"] = client + return hook, client + + +class TestBatchStatus: + @pytest.mark.parametrize( + ("status", "expected"), + [("in_progress", True), ("canceling", True), ("ended", False)], + ) + def test_is_in_progress(self, status, expected): + assert BatchStatus.is_in_progress(status) is expected + + +class TestSessionStatus: + @pytest.mark.parametrize( + ("status", "expected"), + [("running", False), ("rescheduling", False), ("idle", True), ("terminated", True)], + ) + def test_is_terminal(self, status, expected): + assert SessionStatus.is_terminal(status) is expected + + +def _session(status, outcome_results=None): + s = mock.MagicMock() + s.status = status + s.id = "sess_1" + s.outcome_evaluations = [mock.MagicMock(result=r) for r in (outcome_results or [])] + return s + + +def _idle_event(reason, event_id="evt_idle"): + return mock.MagicMock(type="session.status_idle", id=event_id, stop_reason=mock.MagicMock(type=reason)) + + +class TestEvaluateSessionState: + def test_terminated_is_done_with_error(self): + done, err, check = evaluate_session_state(_session("terminated"), expect_outcome=False) + assert done is True + assert err is not None + assert "terminated" in err + assert check is False + + @pytest.mark.parametrize("status", ["running", "rescheduling"]) + def test_non_idle_not_done(self, status): + assert evaluate_session_state(_session(status), expect_outcome=False) == (False, None, False) + + def test_message_idle_needs_event_check(self): + # message run: the session object can't say why it's idle -> defer to the event log + assert evaluate_session_state(_session("idle"), expect_outcome=False) == (False, None, True) + + def test_outcome_idle_without_verdict_not_done(self): + # just-created idle: no verdict yet -> keep waiting (fixes the start race) + assert evaluate_session_state(_session("idle"), expect_outcome=True) == (False, None, False) + + def test_outcome_satisfied_is_done_success(self): + assert evaluate_session_state(_session("idle", ["satisfied"]), expect_outcome=True) == ( + True, + None, + False, + ) + + @pytest.mark.parametrize("result", ["failed", "max_iterations_reached", "interrupted"]) + def test_outcome_failure_is_done_with_error(self, result): + done, err, check = evaluate_session_state(_session("idle", [result]), expect_outcome=True) + assert done is True + assert err is not None + assert result in err + + +class TestPollSessionCompletion: + def test_terminated_is_error(self): + hook, client = _make_hook() + client.beta.sessions.retrieve.return_value = _session("terminated") + done, err = hook.poll_session_completion("s") + assert done is True + assert err is not None + + def test_message_end_turn_success(self): + hook, client = _make_hook() + client.beta.sessions.retrieve.return_value = _session("idle") + client.beta.sessions.events.list.return_value = [_idle_event("end_turn")] + assert hook.poll_session_completion("s", kickoff_event_id="evt_kick") == (True, None) + + @pytest.mark.parametrize("reason", ["requires_action", "retries_exhausted"]) + def test_message_blocked_is_error(self, reason): + hook, client = _make_hook() + client.beta.sessions.retrieve.return_value = _session("idle") + client.beta.sessions.events.list.return_value = [_idle_event(reason)] + done, err = hook.poll_session_completion("s", kickoff_event_id="evt_kick") + assert done is True + assert err is not None + assert reason in err + + def test_message_no_response_yet_not_done(self): + # newest event is our kickoff (agent hasn't responded) -> keep waiting (start race) + hook, client = _make_hook() + client.beta.sessions.retrieve.return_value = _session("idle") + client.beta.sessions.events.list.return_value = [mock.MagicMock(type="user.message", id="evt_kick")] + assert hook.poll_session_completion("s", kickoff_event_id="evt_kick") == (False, None) + + def test_outcome_satisfied_skips_event_check(self): + hook, client = _make_hook() + client.beta.sessions.retrieve.return_value = _session("idle", ["satisfied"]) + assert hook.poll_session_completion("s", expect_outcome=True) == (True, None) + client.beta.sessions.events.list.assert_not_called() + + +class TestDefaultModel: + def test_defaults_to_constant(self): + hook, _ = _make_hook() + assert hook.default_model == DEFAULT_MODEL + + def test_from_connection_extra(self): + hook, _ = _make_hook(extra={"model": "claude-sonnet-4-6"}) + assert hook.default_model == "claude-sonnet-4-6" + + def test_create_message_uses_connection_model(self): + hook, client = _make_hook(extra={"model": "claude-sonnet-4-6"}) + hook.create_message([{"role": "user", "content": "hi"}]) + assert client.messages.create.call_args.kwargs["model"] == "claude-sonnet-4-6" + + def test_explicit_model_overrides_connection(self): + hook, client = _make_hook(extra={"model": "claude-sonnet-4-6"}) + hook.create_message([{"role": "user", "content": "hi"}], model="claude-haiku-4-5") + assert client.messages.create.call_args.kwargs["model"] == "claude-haiku-4-5" + + def test_create_agent_uses_connection_model(self): + hook, client = _make_hook(extra={"model": "claude-sonnet-4-6"}) + hook.create_agent(name="a") + assert client.beta.agents.create.call_args.kwargs["model"] == "claude-sonnet-4-6" + + +class TestManagedAgentsHook: + def _hook_with_client(self, extra=None): + hook = AnthropicHook() + hook.__dict__["_connection"] = _conn(extra=extra) + hook.__dict__["platform"] = (extra or {}).get("platform", "anthropic").lower() + client = mock.MagicMock() + hook.__dict__["conn"] = client + return hook, client + + def test_create_agent(self): + hook, client = self._hook_with_client() + hook.create_agent(name="ag1", system="be brief") + client.beta.agents.create.assert_called_once_with(name="ag1", model=DEFAULT_MODEL, system="be brief") + + def test_archive_session(self): + hook, client = self._hook_with_client() + hook.archive_session("sess_1") + client.beta.sessions.archive.assert_called_once_with("sess_1") + + def test_create_session(self): + hook, client = self._hook_with_client() + hook.create_session(agent="ag", environment_id="env") + client.beta.sessions.create.assert_called_once_with(agent="ag", environment_id="env") + + def test_create_environment_defaults_cloud(self): + hook, client = self._hook_with_client() + hook.create_environment(name="env1") + client.beta.environments.create.assert_called_once_with( + name="env1", config={"type": "cloud", "networking": {"type": "unrestricted"}} + ) + + def test_send_event_wraps_in_events_list(self): + hook, client = self._hook_with_client() + event = {"type": "user.message", "content": [{"type": "text", "text": "hi"}]} + hook.send_event("sess_1", event) + client.beta.sessions.events.send.assert_called_once_with("sess_1", events=[event]) + + def test_get_session(self): + hook, client = self._hook_with_client() + hook.get_session("sess_1") + client.beta.sessions.retrieve.assert_called_once_with("sess_1") + + @pytest.mark.parametrize("platform", ["bedrock", "vertex", "foundry"]) + def test_managed_agents_unavailable_on_non_first_party(self, platform): + hook, _ = self._hook_with_client(extra={"platform": platform}) + with pytest.raises(AnthropicError, match="Managed Agents is not available"): + hook.create_session(agent="ag", environment_id="env") + + +class TestWaitForSession: + @mock.patch(f"{HOOK_PATH}.time.sleep") + @mock.patch.object(AnthropicHook, "poll_session_completion") + @mock.patch.object(AnthropicHook, "get_connection") + def test_returns_when_done(self, mock_get_connection, mock_poll, mock_sleep): + mock_get_connection.return_value = _conn() + mock_poll.side_effect = [(False, None), (True, None)] + AnthropicHook().wait_for_session("sess_1", poll_interval=0.01) + assert mock_poll.call_count == 2 + + @mock.patch(f"{HOOK_PATH}.time.monotonic") + @mock.patch(f"{HOOK_PATH}.time.sleep") + @mock.patch.object(AnthropicHook, "poll_session_completion") + @mock.patch.object(AnthropicHook, "get_connection") + def test_raises_on_timeout(self, mock_get_connection, mock_poll, mock_sleep, mock_monotonic): + mock_get_connection.return_value = _conn() + mock_poll.return_value = (False, None) + mock_monotonic.side_effect = [0, 100] + with pytest.raises(AnthropicAgentSessionTimeout, match="did not reach a terminal status"): + AnthropicHook().wait_for_session("sess_1", poll_interval=0.01, timeout=10) + + @mock.patch(f"{HOOK_PATH}.time.sleep") + @mock.patch.object(AnthropicHook, "poll_session_completion") + @mock.patch.object(AnthropicHook, "get_connection") + def test_failure_raises(self, mock_get_connection, mock_poll, mock_sleep): + mock_get_connection.return_value = _conn() + mock_poll.return_value = (True, "Outcome not satisfied for session sess_1: failed.") + with pytest.raises(AnthropicAgentSessionError, match="not satisfied"): + AnthropicHook().wait_for_session("sess_1", expect_outcome=True, poll_interval=0.01) + + +class TestAnthropicHookGetConn: + @mock.patch(f"{HOOK_PATH}.Anthropic") + @mock.patch.object(AnthropicHook, "get_connection") + def test_first_party_uses_password_and_host(self, mock_get_connection, mock_anthropic): + mock_get_connection.return_value = _conn(password="sk-ant", host="https://gw.example") + client = AnthropicHook().get_conn() + mock_anthropic.assert_called_once_with(api_key="sk-ant", base_url="https://gw.example") + assert client is mock_anthropic.return_value + + @mock.patch(f"{HOOK_PATH}.Anthropic") + @mock.patch.object(AnthropicHook, "get_connection") + def test_client_kwargs_override_api_key_and_base_url(self, mock_get_connection, mock_anthropic): + mock_get_connection.return_value = _conn( + password="from-password", + extra={"anthropic_client_kwargs": {"api_key": "from-extra", "max_retries": 5}}, + ) + AnthropicHook().get_conn() + mock_anthropic.assert_called_once_with(api_key="from-extra", base_url=None, max_retries=5) + + @mock.patch(f"{HOOK_PATH}.AnthropicBedrock") + @mock.patch.object(AnthropicHook, "get_connection") + def test_bedrock_platform(self, mock_get_connection, mock_bedrock): + mock_get_connection.return_value = _conn(extra={"platform": "bedrock", "aws_region": "us-east-1"}) + client = AnthropicHook().get_conn() + mock_bedrock.assert_called_once_with(aws_region="us-east-1") + assert client is mock_bedrock.return_value + + @mock.patch(f"{HOOK_PATH}.AnthropicVertex") + @mock.patch.object(AnthropicHook, "get_connection") + def test_vertex_platform(self, mock_get_connection, mock_vertex): + mock_get_connection.return_value = _conn( + extra={"platform": "vertex", "project_id": "p1", "region": "us-central1"} + ) + AnthropicHook().get_conn() + mock_vertex.assert_called_once_with(project_id="p1", region="us-central1") + + @mock.patch(f"{HOOK_PATH}.AnthropicAWS", autospec=True) + @mock.patch.object(AnthropicHook, "get_connection") + def test_aws_platform(self, mock_get_connection, mock_aws): + mock_get_connection.return_value = _conn(extra={"platform": "aws", "aws_region": "us-east-1"}) + AnthropicHook().get_conn() + mock_aws.assert_called_once_with(aws_region="us-east-1") + + @mock.patch(f"{HOOK_PATH}.AnthropicAWS", autospec=True) + @mock.patch.object(AnthropicHook, "get_connection") + def test_aws_platform_without_region(self, mock_get_connection, mock_aws): + # No aws_region configured -> pass None so the SDK falls back to its own resolution. + mock_get_connection.return_value = _conn(extra={"platform": "aws"}) + AnthropicHook().get_conn() + mock_aws.assert_called_once_with(aws_region=None) + + @mock.patch(f"{HOOK_PATH}.AnthropicFoundry") + @mock.patch.object(AnthropicHook, "get_connection") + def test_foundry_platform(self, mock_get_connection, mock_foundry): + mock_get_connection.return_value = _conn( + password="azkey", extra={"platform": "foundry", "resource": "r1"} + ) + AnthropicHook().get_conn() + mock_foundry.assert_called_once_with(api_key="azkey", resource="r1") + + @mock.patch.object(AnthropicHook, "get_connection") + def test_unknown_platform_raises(self, mock_get_connection): + mock_get_connection.return_value = _conn(extra={"platform": "bogus"}) + with pytest.raises(AnthropicError, match="Unknown Anthropic platform"): + AnthropicHook().get_conn() + + @mock.patch(f"{HOOK_PATH}.Anthropic") + @mock.patch(f"{HOOK_PATH}.IdentityTokenFile") + @mock.patch(f"{HOOK_PATH}.WorkloadIdentityCredentials") + @mock.patch.object(AnthropicHook, "get_connection") + def test_workload_identity_federation_explicit( + self, mock_get_connection, mock_wic, mock_itf, mock_anthropic + ): + mock_get_connection.return_value = _conn( + password=None, + extra={ + "workload_identity": { + "identity_token_file": "/var/run/secrets/anthropic.com/token", + "federation_rule_id": "fdrl_x", + "organization_id": "org_x", + "service_account_id": "svac_x", + "workspace_id": "wrkspc_x", + } + }, + ) + AnthropicHook().get_conn() + mock_itf.assert_called_once_with("/var/run/secrets/anthropic.com/token") + mock_wic.assert_called_once_with( + identity_token_provider=mock_itf.return_value, + federation_rule_id="fdrl_x", + organization_id="org_x", + service_account_id="svac_x", + workspace_id="wrkspc_x", + ) + mock_anthropic.assert_called_once_with(credentials=mock_wic.return_value, base_url=None) + + @mock.patch(f"{HOOK_PATH}.Anthropic") + @mock.patch.object(AnthropicHook, "get_connection") + def test_no_key_lets_sdk_resolve_credentials(self, mock_get_connection, mock_anthropic): + # No static key and no explicit WIF config: construct without api_key so the SDK + # resolves env-driven Workload Identity Federation / profile credentials. + mock_get_connection.return_value = _conn(password=None) + AnthropicHook().get_conn() + mock_anthropic.assert_called_once_with(base_url=None) + + +class TestAnthropicHookFeatures: + def _hook_with_client(self, extra=None): + hook = AnthropicHook() + # Pre-populate the cached_property values so no real connection lookup or + # client construction happens. + hook.__dict__["_connection"] = _conn(extra=extra) + hook.__dict__["platform"] = (extra or {}).get("platform", "anthropic").lower() + client = mock.MagicMock() + hook.__dict__["conn"] = client + return hook, client + + def test_create_message_passes_params(self): + hook, client = self._hook_with_client() + hook.create_message([{"role": "user", "content": "hi"}], system="be brief") + client.messages.create.assert_called_once_with( + model=DEFAULT_MODEL, + max_tokens=1024, + messages=[{"role": "user", "content": "hi"}], + system="be brief", + ) + + def test_count_tokens_returns_input_tokens(self): + hook, client = self._hook_with_client() + client.messages.count_tokens.return_value.input_tokens = 42 + assert hook.count_tokens([{"role": "user", "content": "hi"}]) == 42 + + def test_create_batch_passes_requests(self): + hook, client = self._hook_with_client() + reqs = [{"custom_id": "a", "params": {"model": DEFAULT_MODEL, "max_tokens": 1, "messages": []}}] + hook.create_batch(reqs) + client.messages.batches.create.assert_called_once_with(requests=reqs) + + def test_get_and_cancel_batch(self): + hook, client = self._hook_with_client() + hook.get_batch("batch_1") + client.messages.batches.retrieve.assert_called_once_with("batch_1") + hook.cancel_batch("batch_1") + client.messages.batches.cancel.assert_called_once_with("batch_1") + + def test_stream_batch_results_is_generator(self): + hook, client = self._hook_with_client() + client.messages.batches.results.return_value = iter([1, 2, 3]) + assert list(hook.stream_batch_results("batch_1")) == [1, 2, 3] + + @pytest.mark.parametrize("platform", ["bedrock", "vertex", "foundry"]) + def test_batch_features_unavailable_on_non_first_party(self, platform): + hook, _ = self._hook_with_client(extra={"platform": platform}) + with pytest.raises(AnthropicError, match="not available on the"): + hook.create_batch([]) + with pytest.raises(AnthropicError, match="not available on the"): + hook.count_tokens([{"role": "user", "content": "x"}]) + + +class TestWaitForBatch: + def _batch(self, status): + batch = mock.MagicMock() + batch.processing_status = status + return batch + + @mock.patch(f"{HOOK_PATH}.time.sleep") + @mock.patch.object(AnthropicHook, "get_batch") + @mock.patch.object(AnthropicHook, "get_connection") + def test_returns_when_ended(self, mock_get_connection, mock_get_batch, mock_sleep): + mock_get_connection.return_value = _conn() + mock_get_batch.side_effect = [self._batch("in_progress"), self._batch("ended")] + AnthropicHook().wait_for_batch("batch_1", wait_seconds=0.01) + assert mock_get_batch.call_count == 2 + + @mock.patch(f"{HOOK_PATH}.time.monotonic") + @mock.patch(f"{HOOK_PATH}.time.sleep") + @mock.patch.object(AnthropicHook, "get_batch") + @mock.patch.object(AnthropicHook, "get_connection") + def test_raises_on_timeout(self, mock_get_connection, mock_get_batch, mock_sleep, mock_monotonic): + mock_get_connection.return_value = _conn() + mock_get_batch.return_value = self._batch("in_progress") + mock_monotonic.side_effect = [0, 100] # start=0, elapsed=100 > timeout + with pytest.raises(AnthropicBatchTimeout, match="did not reach a terminal status"): + AnthropicHook().wait_for_batch("batch_1", wait_seconds=0.01, timeout=10) + + @mock.patch(f"{HOOK_PATH}.time.sleep") + @mock.patch.object(AnthropicHook, "get_batch") + @mock.patch.object(AnthropicHook, "get_connection") + def test_tolerates_transient_poll_error(self, mock_get_connection, mock_get_batch, mock_sleep): + mock_get_connection.return_value = _conn() + mock_get_batch.side_effect = [RuntimeError("blip"), self._batch("ended")] + AnthropicHook().wait_for_batch("batch_1", wait_seconds=0.01) + assert mock_get_batch.call_count == 2 + + @mock.patch(f"{HOOK_PATH}.time.sleep") + @mock.patch.object(AnthropicHook, "get_batch") + @mock.patch.object(AnthropicHook, "get_connection") + def test_gives_up_after_consecutive_poll_failures(self, mock_get_connection, mock_get_batch, mock_sleep): + mock_get_connection.return_value = _conn() + mock_get_batch.side_effect = RuntimeError("persistent") + with pytest.raises(RuntimeError, match="persistent"): + AnthropicHook().wait_for_batch("batch_1", wait_seconds=0.01) + assert mock_get_batch.call_count == MAX_CONSECUTIVE_POLL_FAILURES + + +class TestTestConnection: + @mock.patch.object(AnthropicHook, "get_connection") + def test_first_party_calls_models_list(self, mock_get_connection): + mock_get_connection.return_value = _conn() + hook = AnthropicHook() + client = mock.MagicMock() + hook.__dict__["conn"] = client + ok, msg = hook.test_connection() + assert ok is True + client.models.list.assert_called_once() + + @mock.patch.object(AnthropicHook, "get_conn") + @mock.patch.object(AnthropicHook, "get_connection") + def test_non_first_party_skips_live_check(self, mock_get_connection, mock_get_conn): + mock_get_connection.return_value = _conn(extra={"platform": "bedrock", "aws_region": "us-east-1"}) + ok, msg = AnthropicHook().test_connection() + assert ok is True + assert "no live check" in msg + mock_get_conn.assert_called_once() diff --git a/providers/anthropic/tests/unit/anthropic/operators/__init__.py b/providers/anthropic/tests/unit/anthropic/operators/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/operators/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/tests/unit/anthropic/operators/test_agent.py b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py new file mode 100644 index 0000000000000..d461f060f340e --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py @@ -0,0 +1,226 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.exceptions import TaskDeferred +from airflow.providers.anthropic.exceptions import AnthropicAgentSessionError, AnthropicAgentSessionTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook +from airflow.providers.anthropic.operators.agent import AnthropicAgentSessionOperator +from airflow.providers.anthropic.triggers.agent import AnthropicAgentSessionTrigger + +pytest.importorskip("anthropic") + + +def _context(): + return {"ti": mock.MagicMock()} + + +def test_requires_exactly_one_of_message_or_outcome(): + with pytest.raises(ValueError, match="exactly one"): + AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env") + with pytest.raises(ValueError, match="exactly one"): + AnthropicAgentSessionOperator( + task_id="a", agent_id="ag", environment_id="env", message="hi", outcome={"description": "x"} + ) + + +def test_outcome_requires_description_and_rubric(): + # missing rubric + with pytest.raises(ValueError, match="description.*rubric"): + AnthropicAgentSessionOperator( + task_id="a", agent_id="ag", environment_id="env", outcome={"description": "x"} + ) + # missing description + with pytest.raises(ValueError, match="description.*rubric"): + AnthropicAgentSessionOperator( + task_id="a", + agent_id="ag", + environment_id="env", + outcome={"rubric": {"type": "text", "content": "c"}}, + ) + + +class TestExecute: + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_message_sends_user_message_and_waits(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + mock_hook_prop.return_value = hook + + op = AnthropicAgentSessionOperator( + task_id="a", agent_id="ag", environment_id="env", message="summarize", deferrable=False + ) + context = _context() + assert op.execute(context) == "sess_1" + hook.create_session.assert_called_once_with(agent="ag", environment_id="env") + hook.send_event.assert_called_once_with( + "sess_1", {"type": "user.message", "content": [{"type": "text", "text": "summarize"}]} + ) + hook.wait_for_session.assert_called_once() + context["ti"].xcom_push.assert_called_once_with(key="session_id", value="sess_1") + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_outcome_sends_define_outcome(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + mock_hook_prop.return_value = hook + + outcome = {"description": "build a CSV", "rubric": {"type": "text", "content": "has a price column"}} + op = AnthropicAgentSessionOperator( + task_id="a", agent_id="ag", environment_id="env", outcome=outcome, deferrable=False + ) + op.execute(_context()) + hook.send_event.assert_called_once_with("sess_1", {"type": "user.define_outcome", **outcome}) + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_passes_vault_ids_and_resources(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + mock_hook_prop.return_value = hook + + op = AnthropicAgentSessionOperator( + task_id="a", + agent_id="ag", + environment_id="env", + message="hi", + deferrable=False, + vault_ids=["vlt_1"], + session_resources=[{"type": "file", "file_id": "f1", "mount_path": "/workspace/f"}], + ) + op.execute(_context()) + hook.create_session.assert_called_once_with( + agent="ag", + environment_id="env", + vault_ids=["vlt_1"], + resources=[{"type": "file", "file_id": "f1", "mount_path": "/workspace/f"}], + ) + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_sync_timeout_archives_and_raises(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + hook.wait_for_session.side_effect = AnthropicAgentSessionTimeout("too slow") + mock_hook_prop.return_value = hook + + op = AnthropicAgentSessionOperator( + task_id="a", agent_id="ag", environment_id="env", message="hi", deferrable=False + ) + with pytest.raises(AnthropicAgentSessionTimeout, match="too slow"): + op.execute(_context()) + hook.archive_session.assert_called_once_with("sess_1") + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_sync_non_timeout_error_archives_and_raises(self, mock_hook_prop): + # A non-timeout failure while waiting (SDK 5xx, auth expiry) also leaves the session + # container running, so the broadened except archives it best-effort before re-raising. + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + hook.wait_for_session.side_effect = RuntimeError("api 5xx") + mock_hook_prop.return_value = hook + + op = AnthropicAgentSessionOperator( + task_id="a", agent_id="ag", environment_id="env", message="hi", deferrable=False + ) + with pytest.raises(RuntimeError, match="api 5xx"): + op.execute(_context()) + hook.archive_session.assert_called_once_with("sess_1") + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_send_event_failure_archives_session(self, mock_hook_prop): + # send_event fails after create_session allocated the container; it must be archived. + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + hook.send_event.side_effect = RuntimeError("send boom") + mock_hook_prop.return_value = hook + + op = AnthropicAgentSessionOperator( + task_id="a", agent_id="ag", environment_id="env", message="hi", deferrable=False + ) + with pytest.raises(RuntimeError, match="send boom"): + op.execute(_context()) + hook.archive_session.assert_called_once_with("sess_1") + hook.wait_for_session.assert_not_called() + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_deferrable_defers_with_trigger(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + mock_hook_prop.return_value = hook + + op = AnthropicAgentSessionOperator( + task_id="a", agent_id="ag", environment_id="env", message="hi", deferrable=True + ) + with pytest.raises(TaskDeferred) as exc: + op.execute(_context()) + assert isinstance(exc.value.trigger, AnthropicAgentSessionTrigger) + assert exc.value.trigger.session_id == "sess_1" + assert exc.value.method_name == "execute_complete" + hook.wait_for_session.assert_not_called() + + +class TestExecuteComplete: + def test_success_returns_session_id(self): + op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") + assert op.execute_complete({}, {"status": "success", "session_id": "sess_1"}) == "sess_1" + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_error_archives_and_raises(self, mock_hook_prop): + # The trigger's "error" event means polling gave up while the session may still be + # running, so the operator archives it best-effort before failing. + hook = mock.MagicMock(spec=AnthropicHook) + mock_hook_prop.return_value = hook + op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") + with pytest.raises(AnthropicAgentSessionError, match="boom"): + op.execute_complete({}, {"status": "error", "session_id": "s", "message": "boom"}) + hook.archive_session.assert_called_once_with("s") + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_timeout_archives_and_raises(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + mock_hook_prop.return_value = hook + op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") + with pytest.raises(AnthropicAgentSessionTimeout): + op.execute_complete({}, {"status": "timeout", "session_id": "s", "message": "slow"}) + hook.archive_session.assert_called_once_with("s") + + def test_none_event_raises(self): + op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") + with pytest.raises(AnthropicAgentSessionError, match="without an event"): + op.execute_complete({}, None) + + +class TestOnKill: + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_on_kill_archives_session(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + mock_hook_prop.return_value = hook + op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") + op.session_id = "sess_1" + op.on_kill() + hook.archive_session.assert_called_once_with("sess_1") + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_on_kill_noop_without_session(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + mock_hook_prop.return_value = hook + op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") + op.on_kill() + hook.archive_session.assert_not_called() diff --git a/providers/anthropic/tests/unit/anthropic/operators/test_batch.py b/providers/anthropic/tests/unit/anthropic/operators/test_batch.py new file mode 100644 index 0000000000000..b79d1e5b03710 --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/operators/test_batch.py @@ -0,0 +1,208 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.exceptions import TaskDeferred +from airflow.providers.anthropic.exceptions import AnthropicBatchJobError, AnthropicBatchTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook +from airflow.providers.anthropic.operators.batch import AnthropicBatchOperator +from airflow.providers.anthropic.triggers.batch import AnthropicBatchTrigger +from airflow.providers.common.compat.sdk import AirflowSkipException + +pytest.importorskip("anthropic") + +REQUESTS = [{"custom_id": "a", "params": {"model": "claude-opus-4-8", "max_tokens": 8, "messages": []}}] + + +def _counts(succeeded=0, errored=0, canceled=0, expired=0, processing=0): + counts = mock.MagicMock() + counts.succeeded = succeeded + counts.errored = errored + counts.canceled = canceled + counts.expired = expired + counts.processing = processing + return counts + + +def _context(): + return {"ti": mock.MagicMock()} + + +class TestAnthropicBatchOperatorExecute: + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_sync_waits_and_returns_batch_id(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_batch.return_value.id = "batch_1" + hook.wait_for_batch.return_value.request_counts = _counts(succeeded=1) + mock_hook_prop.return_value = hook + + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS, deferrable=False) + context = _context() + result = op.execute(context) + + assert result == "batch_1" + hook.wait_for_batch.assert_called_once() + context["ti"].xcom_push.assert_called_once_with(key="batch_id", value="batch_1") + + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_deferrable_defers_with_trigger(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_batch.return_value.id = "batch_1" + mock_hook_prop.return_value = hook + + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS, deferrable=True) + with pytest.raises(TaskDeferred) as exc: + op.execute(_context()) + assert isinstance(exc.value.trigger, AnthropicBatchTrigger) + assert exc.value.trigger.batch_id == "batch_1" + assert exc.value.method_name == "execute_complete" + hook.wait_for_batch.assert_not_called() + + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_no_wait_returns_immediately(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_batch.return_value.id = "batch_1" + mock_hook_prop.return_value = hook + + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS, wait_for_completion=False) + assert op.execute(_context()) == "batch_1" + hook.wait_for_batch.assert_not_called() + + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_sync_applies_policy_skip_on_full_cancel(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_batch.return_value.id = "batch_1" + hook.wait_for_batch.return_value.request_counts = _counts(canceled=2) + mock_hook_prop.return_value = hook + + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS, deferrable=False) + with pytest.raises(AirflowSkipException): + op.execute(_context()) + + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_sync_applies_policy_fail_on_partial_error_when_strict(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_batch.return_value.id = "batch_1" + hook.wait_for_batch.return_value.request_counts = _counts(succeeded=4, errored=1) + mock_hook_prop.return_value = hook + + op = AnthropicBatchOperator( + task_id="t", requests=REQUESTS, deferrable=False, fail_on_partial_error=True + ) + with pytest.raises(AnthropicBatchJobError): + op.execute(_context()) + + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_sync_timeout_cancels_and_raises(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_batch.return_value.id = "batch_1" + hook.wait_for_batch.side_effect = AnthropicBatchTimeout("too slow") + mock_hook_prop.return_value = hook + + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS, deferrable=False) + with pytest.raises(AnthropicBatchTimeout, match="too slow"): + op.execute(_context()) + hook.cancel_batch.assert_called_once_with("batch_1") + + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_sync_non_timeout_error_cancels_and_raises(self, mock_hook_prop): + # A non-timeout failure while waiting (SDK 5xx, auth expiry) also leaves the batch + # running, so the broadened except cancels it best-effort before re-raising. + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_batch.return_value.id = "batch_1" + hook.wait_for_batch.side_effect = RuntimeError("api 5xx") + mock_hook_prop.return_value = hook + + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS, deferrable=False) + with pytest.raises(RuntimeError, match="api 5xx"): + op.execute(_context()) + hook.cancel_batch.assert_called_once_with("batch_1") + + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_empty_requests_raises_before_any_api_call(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + mock_hook_prop.return_value = hook + + op = AnthropicBatchOperator(task_id="t", requests=[]) + with pytest.raises(ValueError, match="at least one request"): + op.execute(_context()) + hook.create_batch.assert_not_called() + + +class TestExecuteComplete: + def test_success_returns_batch_id(self): + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS) + event = {"status": "success", "batch_id": "batch_1", "request_counts": {"succeeded": 3}} + assert op.execute_complete(_context(), event) == "batch_1" + + @mock.patch("airflow.providers.anthropic.operators.batch.AnthropicHook", autospec=True) + def test_error_cancels_and_raises_job_error(self, mock_hook_cls): + # The trigger's "error" event means polling gave up while the batch may still be + # running, so the operator cancels it best-effort before failing. + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS) + event = {"status": "error", "batch_id": "batch_1", "message": "boom"} + with pytest.raises(AnthropicBatchJobError, match="boom"): + op.execute_complete(_context(), event) + mock_hook_cls.return_value.cancel_batch.assert_called_once_with("batch_1") + + @mock.patch("airflow.providers.anthropic.operators.batch.AnthropicHook", autospec=True) + def test_timeout_cancels_and_raises(self, mock_hook_cls): + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS) + event = {"status": "timeout", "batch_id": "batch_1", "message": "too slow"} + with pytest.raises(AnthropicBatchTimeout, match="too slow"): + op.execute_complete(_context(), event) + mock_hook_cls.return_value.cancel_batch.assert_called_once_with("batch_1") + + def test_fully_cancelled_skips(self): + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS) + event = {"status": "success", "batch_id": "batch_1", "request_counts": {"canceled": 4}} + with pytest.raises(AirflowSkipException): + op.execute_complete(_context(), event) + + def test_partial_error_warns_by_default(self): + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS) + event = {"status": "success", "batch_id": "b", "request_counts": {"succeeded": 9, "errored": 1}} + assert op.execute_complete(_context(), event) == "b" + + def test_partial_error_fails_when_strict(self): + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS, fail_on_partial_error=True) + event = {"status": "success", "batch_id": "b", "request_counts": {"succeeded": 9, "errored": 1}} + with pytest.raises(AnthropicBatchJobError, match="failed request"): + op.execute_complete(_context(), event) + + +class TestOnKill: + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_on_kill_cancels_known_batch(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + mock_hook_prop.return_value = hook + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS) + op.batch_id = "batch_1" + op.on_kill() + hook.cancel_batch.assert_called_once_with("batch_1") + + @mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock) + def test_on_kill_noop_without_batch(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + mock_hook_prop.return_value = hook + op = AnthropicBatchOperator(task_id="t", requests=REQUESTS) + op.on_kill() + hook.cancel_batch.assert_not_called() diff --git a/providers/anthropic/tests/unit/anthropic/sensors/__init__.py b/providers/anthropic/tests/unit/anthropic/sensors/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/sensors/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/tests/unit/anthropic/sensors/test_batch.py b/providers/anthropic/tests/unit/anthropic/sensors/test_batch.py new file mode 100644 index 0000000000000..02d071feb0634 --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/sensors/test_batch.py @@ -0,0 +1,104 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.exceptions import TaskDeferred +from airflow.providers.anthropic.exceptions import AnthropicBatchJobError, AnthropicBatchTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook +from airflow.providers.anthropic.sensors.batch import AnthropicBatchSensor +from airflow.providers.anthropic.triggers.batch import AnthropicBatchTrigger +from airflow.providers.common.compat.sdk import AirflowSkipException + +pytest.importorskip("anthropic") + + +def _batch(status, succeeded=0, errored=0, canceled=0, expired=0): + batch = mock.MagicMock() + batch.processing_status = status + counts = batch.request_counts + counts.succeeded, counts.errored, counts.canceled, counts.expired = ( + succeeded, + errored, + canceled, + expired, + ) + return batch + + +class TestAnthropicBatchSensorPoke: + @mock.patch.object(AnthropicBatchSensor, "hook", new_callable=mock.PropertyMock) + def test_poke_false_while_in_progress(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.get_batch.return_value = _batch("in_progress") + mock_hook_prop.return_value = hook + sensor = AnthropicBatchSensor(task_id="s", batch_id="b1") + assert sensor.poke({}) is False + + @mock.patch.object(AnthropicBatchSensor, "hook", new_callable=mock.PropertyMock) + def test_poke_true_when_ended(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.get_batch.return_value = _batch("ended", succeeded=5) + mock_hook_prop.return_value = hook + sensor = AnthropicBatchSensor(task_id="s", batch_id="b1") + assert sensor.poke({}) is True + + @mock.patch.object(AnthropicBatchSensor, "hook", new_callable=mock.PropertyMock) + def test_poke_skips_on_full_cancel(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.get_batch.return_value = _batch("ended", canceled=3) + mock_hook_prop.return_value = hook + sensor = AnthropicBatchSensor(task_id="s", batch_id="b1") + with pytest.raises(AirflowSkipException): + sensor.poke({}) + + @mock.patch.object(AnthropicBatchSensor, "hook", new_callable=mock.PropertyMock) + def test_poke_fails_on_partial_error_when_strict(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.get_batch.return_value = _batch("ended", succeeded=4, errored=1) + mock_hook_prop.return_value = hook + sensor = AnthropicBatchSensor(task_id="s", batch_id="b1", fail_on_partial_error=True) + with pytest.raises(AnthropicBatchJobError): + sensor.poke({}) + + +class TestAnthropicBatchSensorDeferrable: + def test_execute_defers(self): + sensor = AnthropicBatchSensor(task_id="s", batch_id="b1", deferrable=True) + with pytest.raises(TaskDeferred) as exc: + sensor.execute({}) + assert isinstance(exc.value.trigger, AnthropicBatchTrigger) + assert exc.value.trigger.batch_id == "b1" + assert exc.value.method_name == "execute_complete" + + def test_execute_complete_error_raises(self): + sensor = AnthropicBatchSensor(task_id="s", batch_id="b1") + with pytest.raises(AnthropicBatchJobError, match="boom"): + sensor.execute_complete({}, {"status": "error", "batch_id": "b1", "message": "boom"}) + + def test_execute_complete_timeout_raises(self): + sensor = AnthropicBatchSensor(task_id="s", batch_id="b1") + with pytest.raises(AnthropicBatchTimeout): + sensor.execute_complete({}, {"status": "timeout", "batch_id": "b1", "message": "slow"}) + + def test_execute_complete_success(self): + sensor = AnthropicBatchSensor(task_id="s", batch_id="b1") + event = {"status": "success", "batch_id": "b1", "request_counts": {"succeeded": 2}} + assert sensor.execute_complete({}, event) is None diff --git a/providers/anthropic/tests/unit/anthropic/test_exceptions.py b/providers/anthropic/tests/unit/anthropic/test_exceptions.py new file mode 100644 index 0000000000000..365ba875eb4cb --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/test_exceptions.py @@ -0,0 +1,44 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import pytest + +from airflow.providers.anthropic.exceptions import ( + AnthropicAgentSessionError, + AnthropicAgentSessionTimeout, + AnthropicBatchJobError, + AnthropicBatchTimeout, + AnthropicError, +) + + +@pytest.mark.parametrize( + "exc", + [ + AnthropicError, + AnthropicBatchJobError, + AnthropicBatchTimeout, + AnthropicAgentSessionError, + AnthropicAgentSessionTimeout, + ], +) +def test_provider_errors_share_base_and_are_not_airflow_exceptions(exc): + # Every provider error is catchable via the AnthropicError base and is deliberately + # NOT an AirflowException subclass (the no-new-AirflowException direction). + assert issubclass(exc, AnthropicError) + assert not any(base.__name__ == "AirflowException" for base in exc.__mro__) diff --git a/providers/anthropic/tests/unit/anthropic/triggers/__init__.py b/providers/anthropic/tests/unit/anthropic/triggers/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/triggers/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/providers/anthropic/tests/unit/anthropic/triggers/test_agent.py b/providers/anthropic/tests/unit/anthropic/triggers/test_agent.py new file mode 100644 index 0000000000000..6e01b4259b80e --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/triggers/test_agent.py @@ -0,0 +1,128 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from unittest import mock + +import pytest + +from airflow.providers.anthropic.triggers.agent import AnthropicAgentSessionTrigger +from airflow.triggers.base import TriggerEvent + +pytest.importorskip("anthropic") + +TRIGGER_PATH = "airflow.providers.anthropic.triggers.agent" +POLL = "airflow.providers.anthropic.hooks.anthropic.AnthropicHook.poll_session_completion" + + +def _trigger(end_time=None, expect_outcome=False): + return AnthropicAgentSessionTrigger( + conn_id="anthropic_default", + session_id="sess_1", + poll_interval=1.0, + end_time=end_time if end_time is not None else time.time() + 3600, + expect_outcome=expect_outcome, + kickoff_event_id="evt_kick", + ) + + +def test_serialization(): + end_time = time.time() + 3600 + path, kwargs = _trigger(end_time).serialize() + assert path == "airflow.providers.anthropic.triggers.agent.AnthropicAgentSessionTrigger" + assert kwargs == { + "conn_id": "anthropic_default", + "session_id": "sess_1", + "poll_interval": 1.0, + "end_time": end_time, + "expect_outcome": False, + "kickoff_event_id": "evt_kick", + } + + +@pytest.mark.asyncio +@mock.patch(f"{TRIGGER_PATH}.AnthropicHook", autospec=True) +async def test_on_kill_archives_session(mock_hook_cls): + # A user killing the deferred task archives the still-running session (Airflow 3.3+). + await _trigger().on_kill() + mock_hook_cls.return_value.archive_session.assert_called_once_with("sess_1") + + +@pytest.mark.asyncio +@mock.patch(POLL) +async def test_done_success_yields_success(mock_poll): + mock_poll.return_value = (True, None) + event = await _trigger().run().__anext__() + assert event.payload["status"] == "success" + assert event.payload["session_id"] == "sess_1" + + +@pytest.mark.asyncio +@mock.patch(POLL) +async def test_done_error_yields_error(mock_poll): + mock_poll.return_value = (True, "Session sess_1 terminated.") + event = await _trigger().run().__anext__() + assert event.payload["status"] == "error" + assert "terminated" in event.payload["message"] + + +@pytest.mark.asyncio +@mock.patch(POLL) +async def test_timeout_yields_timeout(mock_poll): + mock_poll.return_value = (False, None) + event = await _trigger(end_time=time.time() - 1).run().__anext__() + assert event.payload["status"] == "timeout" + + +@pytest.mark.asyncio +@mock.patch(f"{TRIGGER_PATH}.asyncio.sleep") +@mock.patch(POLL) +async def test_polls_until_done(mock_poll, mock_sleep): + mock_poll.side_effect = [(False, None), (False, None), (True, None)] + event = await _trigger().run().__anext__() + assert event.payload["status"] == "success" + assert mock_poll.call_count == 3 + assert mock_sleep.await_count == 2 + + +@pytest.mark.asyncio +@mock.patch(f"{TRIGGER_PATH}.asyncio.sleep") +@mock.patch(POLL) +async def test_persistent_error_yields_error_after_retries(mock_poll, mock_sleep): + mock_poll.side_effect = RuntimeError("kaboom") + event = await _trigger().run().__anext__() + assert event == TriggerEvent({"status": "error", "session_id": "sess_1", "message": "kaboom"}) + assert mock_poll.call_count == 5 + + +@pytest.mark.asyncio +@mock.patch(f"{TRIGGER_PATH}.asyncio.sleep") +@mock.patch(POLL) +async def test_transient_error_then_success(mock_poll, mock_sleep): + mock_poll.side_effect = [RuntimeError("blip"), (True, None)] + event = await _trigger().run().__anext__() + assert event.payload["status"] == "success" + + +@pytest.mark.asyncio +@mock.patch(POLL) +async def test_outcome_failure_yields_error(mock_poll): + mock_poll.return_value = (True, "Outcome not satisfied for session sess_1: max_iterations_reached.") + event = await _trigger(expect_outcome=True).run().__anext__() + assert event.payload["status"] == "error" + assert "max_iterations_reached" in event.payload["message"] diff --git a/providers/anthropic/tests/unit/anthropic/triggers/test_batch.py b/providers/anthropic/tests/unit/anthropic/triggers/test_batch.py new file mode 100644 index 0000000000000..564681f545df7 --- /dev/null +++ b/providers/anthropic/tests/unit/anthropic/triggers/test_batch.py @@ -0,0 +1,132 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from unittest import mock + +import pytest + +from airflow.providers.anthropic.triggers.batch import AnthropicBatchTrigger +from airflow.triggers.base import TriggerEvent + +pytest.importorskip("anthropic") + +TRIGGER_PATH = "airflow.providers.anthropic.triggers.batch" +HOOK_PATH = "airflow.providers.anthropic.hooks.anthropic.AnthropicHook.get_batch" + + +def _batch(status, succeeded=0, errored=0, canceled=0, expired=0, processing=0): + batch = mock.MagicMock() + batch.processing_status = status + counts = mock.MagicMock() + counts.succeeded, counts.errored = succeeded, errored + counts.canceled, counts.expired, counts.processing = canceled, expired, processing + batch.request_counts = counts + return batch + + +class TestAnthropicBatchTrigger: + CONN_ID = "anthropic_default" + BATCH_ID = "batch_1" + POLL = 1.0 + + def _trigger(self, end_time=None): + return AnthropicBatchTrigger( + conn_id=self.CONN_ID, + batch_id=self.BATCH_ID, + poll_interval=self.POLL, + end_time=end_time if end_time is not None else time.time() + 3600, + ) + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_PATH}.AnthropicHook", autospec=True) + async def test_on_kill_cancels_batch(self, mock_hook_cls): + # A user killing the deferred task cancels the still-running batch (Airflow 3.3+). + await self._trigger().on_kill() + mock_hook_cls.return_value.cancel_batch.assert_called_once_with(self.BATCH_ID) + + def test_serialization(self): + end_time = time.time() + 3600 + class_path, kwargs = self._trigger(end_time).serialize() + assert class_path == "airflow.providers.anthropic.triggers.batch.AnthropicBatchTrigger" + assert kwargs == { + "conn_id": self.CONN_ID, + "batch_id": self.BATCH_ID, + "poll_interval": self.POLL, + "end_time": end_time, + } + + @pytest.mark.asyncio + @mock.patch(HOOK_PATH) + async def test_ended_yields_success_with_counts(self, mock_get_batch): + mock_get_batch.return_value = _batch("ended", succeeded=3, errored=1) + event = await self._trigger().run().__anext__() + assert event == TriggerEvent( + { + "status": "success", + "batch_id": self.BATCH_ID, + "message": f"Batch {self.BATCH_ID} has ended.", + "request_counts": { + "succeeded": 3, + "errored": 1, + "canceled": 0, + "expired": 0, + "processing": 0, + }, + } + ) + + @pytest.mark.asyncio + @mock.patch(HOOK_PATH) + async def test_timeout_yields_timeout_event(self, mock_get_batch): + mock_get_batch.return_value = _batch("in_progress") + # end_time already in the past -> timeout on first poll. + event = await self._trigger(end_time=time.time() - 1).run().__anext__() + assert event.payload["status"] == "timeout" + assert event.payload["batch_id"] == self.BATCH_ID + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_PATH}.asyncio.sleep") + @mock.patch(HOOK_PATH) + async def test_persistent_error_yields_error_after_retries(self, mock_get_batch, mock_sleep): + mock_get_batch.side_effect = RuntimeError("kaboom") + event = await self._trigger().run().__anext__() + assert event == TriggerEvent({"status": "error", "batch_id": self.BATCH_ID, "message": "kaboom"}) + assert mock_get_batch.call_count == 5 + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_PATH}.asyncio.sleep") + @mock.patch(HOOK_PATH) + async def test_transient_error_then_success(self, mock_get_batch, mock_sleep): + mock_get_batch.side_effect = [RuntimeError("blip"), _batch("ended", succeeded=1)] + event = await self._trigger().run().__anext__() + assert event.payload["status"] == "success" + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_PATH}.asyncio.sleep") + @mock.patch(HOOK_PATH) + async def test_polls_until_ended(self, mock_get_batch, mock_sleep): + mock_get_batch.side_effect = [ + _batch("in_progress"), + _batch("canceling"), + _batch("ended", succeeded=1), + ] + event = await self._trigger().run().__anext__() + assert event.payload["status"] == "success" + assert mock_get_batch.call_count == 3 + assert mock_sleep.await_count == 2 diff --git a/pyproject.toml b/pyproject.toml index 6cd884588ba6e..633a208e4dcfe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,6 +110,9 @@ apache-airflow = "airflow.__main__:main" "amazon" = [ "apache-airflow-providers-amazon>=9.0.0" ] +"anthropic" = [ + "apache-airflow-providers-anthropic>=0.1.0" +] "apache.cassandra" = [ "apache-airflow-providers-apache-cassandra>=3.7.0; python_version !=\"3.14\"" ] @@ -411,6 +414,7 @@ apache-airflow = "airflow.__main__:main" "apache-airflow-providers-akeyless>=0.1.0", "apache-airflow-providers-alibaba>=3.0.0", "apache-airflow-providers-amazon>=9.0.0", + "apache-airflow-providers-anthropic>=0.1.0", "apache-airflow-providers-apache-cassandra>=3.7.0; python_version !=\"3.14\"", "apache-airflow-providers-apache-drill>=2.8.1", "apache-airflow-providers-apache-druid>=3.12.0", @@ -1090,6 +1094,8 @@ mypy_path = [ "$MYPY_CONFIG_FILE_DIR/providers/alibaba/tests", "$MYPY_CONFIG_FILE_DIR/providers/amazon/src", "$MYPY_CONFIG_FILE_DIR/providers/amazon/tests", + "$MYPY_CONFIG_FILE_DIR/providers/anthropic/src", + "$MYPY_CONFIG_FILE_DIR/providers/anthropic/tests", "$MYPY_CONFIG_FILE_DIR/providers/apache/cassandra/src", "$MYPY_CONFIG_FILE_DIR/providers/apache/cassandra/tests", "$MYPY_CONFIG_FILE_DIR/providers/apache/drill/src", @@ -1427,6 +1433,7 @@ apache-airflow-providers-airbyte = false apache-airflow-providers-akeyless = false apache-airflow-providers-alibaba = false apache-airflow-providers-amazon = false +apache-airflow-providers-anthropic = false apache-airflow-providers-apache-cassandra = false apache-airflow-providers-apache-drill = false apache-airflow-providers-apache-druid = false @@ -1571,6 +1578,7 @@ apache-airflow-providers-airbyte = false apache-airflow-providers-akeyless = false apache-airflow-providers-alibaba = false apache-airflow-providers-amazon = false +apache-airflow-providers-anthropic = false apache-airflow-providers-apache-cassandra = false apache-airflow-providers-apache-drill = false apache-airflow-providers-apache-druid = false @@ -1730,6 +1738,7 @@ apache-airflow-providers-airbyte = { workspace = true } apache-airflow-providers-akeyless = { workspace = true } apache-airflow-providers-alibaba = { workspace = true } apache-airflow-providers-amazon = { workspace = true } +apache-airflow-providers-anthropic = { workspace = true } apache-airflow-providers-apache-cassandra = { workspace = true } apache-airflow-providers-apache-drill = { workspace = true } apache-airflow-providers-apache-druid = { workspace = true } @@ -1869,6 +1878,7 @@ members = [ "providers/akeyless", "providers/alibaba", "providers/amazon", + "providers/anthropic", "providers/apache/cassandra", "providers/apache/drill", "providers/apache/druid", diff --git a/scripts/ci/docker-compose/remove-sources.yml b/scripts/ci/docker-compose/remove-sources.yml index f10ae1472e9d6..ed2e4b1bed789 100644 --- a/scripts/ci/docker-compose/remove-sources.yml +++ b/scripts/ci/docker-compose/remove-sources.yml @@ -30,6 +30,7 @@ services: - ../../../empty:/opt/airflow/providers/akeyless/src - ../../../empty:/opt/airflow/providers/alibaba/src - ../../../empty:/opt/airflow/providers/amazon/src + - ../../../empty:/opt/airflow/providers/anthropic/src - ../../../empty:/opt/airflow/providers/apache/beam/src - ../../../empty:/opt/airflow/providers/apache/cassandra/src - ../../../empty:/opt/airflow/providers/apache/drill/src diff --git a/scripts/ci/docker-compose/tests-sources.yml b/scripts/ci/docker-compose/tests-sources.yml index d33deafc970fa..1f6e47ba961c2 100644 --- a/scripts/ci/docker-compose/tests-sources.yml +++ b/scripts/ci/docker-compose/tests-sources.yml @@ -43,6 +43,7 @@ services: - ../../../providers/akeyless/tests:/opt/airflow/providers/akeyless/tests - ../../../providers/alibaba/tests:/opt/airflow/providers/alibaba/tests - ../../../providers/amazon/tests:/opt/airflow/providers/amazon/tests + - ../../../providers/anthropic/tests:/opt/airflow/providers/anthropic/tests - ../../../providers/apache/beam/tests:/opt/airflow/providers/apache/beam/tests - ../../../providers/apache/cassandra/tests:/opt/airflow/providers/apache/cassandra/tests - ../../../providers/apache/drill/tests:/opt/airflow/providers/apache/drill/tests diff --git a/uv.lock b/uv.lock index 939b97ad4cc26..9df1b9bf7f592 100644 --- a/uv.lock +++ b/uv.lock @@ -68,6 +68,7 @@ apache-airflow-providers-microsoft-mssql = false apache-airflow-providers-teradata = false apache-airflow-providers-jdbc = false apache-airflow-helm-chart = false +apache-airflow-providers-anthropic = false apache-airflow-providers-common-io = false apache-airflow-providers-cohere = false apache-airflow-providers-pinecone = false @@ -175,6 +176,7 @@ members = [ "apache-airflow-providers-akeyless", "apache-airflow-providers-alibaba", "apache-airflow-providers-amazon", + "apache-airflow-providers-anthropic", "apache-airflow-providers-apache-cassandra", "apache-airflow-providers-apache-drill", "apache-airflow-providers-apache-druid", @@ -958,6 +960,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a9/26/ea71185027956325be1903d4fcaf7461d5ef40ca8f0e64f992e24ea9db0e/anthropic-0.112.0-py3-none-any.whl", hash = "sha256:bcc6268612c716dbb77133dd60fc41d26016d1b81dee9a52314d210193638751", size = 931954, upload-time = "2026-06-24T18:45:58.205Z" }, ] +[package.optional-dependencies] +aws = [ + { name = "boto3" }, + { name = "botocore" }, +] +bedrock = [ + { name = "boto3" }, + { name = "botocore" }, +] +vertex = [ + { name = "google-auth", extra = ["requests"] }, +] + [[package]] name = "anyio" version = "4.14.1" @@ -1001,6 +1016,7 @@ all = [ { name = "apache-airflow-providers-akeyless" }, { name = "apache-airflow-providers-alibaba" }, { name = "apache-airflow-providers-amazon", extra = ["aiobotocore", "python3-saml", "s3fs"] }, + { name = "apache-airflow-providers-anthropic" }, { name = "apache-airflow-providers-apache-cassandra", marker = "python_full_version != '3.14.*'" }, { name = "apache-airflow-providers-apache-drill" }, { name = "apache-airflow-providers-apache-druid" }, @@ -1119,6 +1135,9 @@ amazon = [ amazon-aws-auth = [ { name = "apache-airflow-providers-amazon", extra = ["python3-saml"] }, ] +anthropic = [ + { name = "apache-airflow-providers-anthropic" }, +] apache-atlas = [ { name = "atlasclient" }, ] @@ -1567,6 +1586,8 @@ requires-dist = [ { name = "apache-airflow-providers-amazon", extras = ["aiobotocore"], marker = "extra == 'aiobotocore'", editable = "providers/amazon" }, { name = "apache-airflow-providers-amazon", extras = ["python3-saml"], marker = "extra == 'amazon-aws-auth'", editable = "providers/amazon" }, { name = "apache-airflow-providers-amazon", extras = ["s3fs"], marker = "extra == 's3fs'", editable = "providers/amazon" }, + { name = "apache-airflow-providers-anthropic", marker = "extra == 'all'", editable = "providers/anthropic" }, + { name = "apache-airflow-providers-anthropic", marker = "extra == 'anthropic'", editable = "providers/anthropic" }, { name = "apache-airflow-providers-apache-cassandra", marker = "python_full_version != '3.14.*' and extra == 'all'", editable = "providers/apache/cassandra" }, { name = "apache-airflow-providers-apache-cassandra", marker = "python_full_version != '3.14.*' and extra == 'apache-cassandra'", editable = "providers/apache/cassandra" }, { name = "apache-airflow-providers-apache-drill", marker = "extra == 'all'", editable = "providers/apache/drill" }, @@ -1778,7 +1799,7 @@ requires-dist = [ { name = "sentry-sdk", marker = "extra == 'sentry'", specifier = ">=2.30.0" }, { name = "uv", marker = "extra == 'uv'", specifier = ">=0.11.25" }, ] -provides-extras = ["all-core", "async", "graphviz", "gunicorn", "kerberos", "memray", "otel", "statsd", "all-task-sdk", "airbyte", "akeyless", "alibaba", "amazon", "apache-cassandra", "apache-drill", "apache-druid", "apache-flink", "apache-hdfs", "apache-hive", "apache-iceberg", "apache-impala", "apache-kafka", "apache-kylin", "apache-livy", "apache-pig", "apache-pinot", "apache-spark", "apache-tinkerpop", "apprise", "arangodb", "asana", "atlassian-jira", "celery", "clickhousedb", "cloudant", "cncf-kubernetes", "cohere", "common-ai", "common-compat", "common-io", "common-messaging", "common-sql", "databricks", "datadog", "dbt-cloud", "dingding", "discord", "docker", "edge3", "elasticsearch", "exasol", "fab", "facebook", "ftp", "git", "github", "google", "grpc", "hashicorp", "http", "ibm-mq", "imap", "influxdb", "informatica", "jdbc", "jenkins", "keycloak", "microsoft-azure", "microsoft-mssql", "microsoft-psrp", "microsoft-winrm", "mongo", "mysql", "neo4j", "odbc", "openai", "openfaas", "openlineage", "opensearch", "opsgenie", "oracle", "pagerduty", "papermill", "pgvector", "pinecone", "postgres", "presto", "qdrant", "redis", "salesforce", "samba", "segment", "sendgrid", "sftp", "singularity", "slack", "smtp", "snowflake", "sqlite", "ssh", "standard", "tableau", "telegram", "teradata", "trino", "vertica", "vespa", "weaviate", "yandex", "ydb", "zendesk", "all", "aiobotocore", "apache-atlas", "apache-webhdfs", "amazon-aws-auth", "cloudpickle", "github-enterprise", "google-auth", "ldap", "pandas", "polars", "rabbitmq", "sentry", "s3fs", "uv"] +provides-extras = ["all-core", "async", "graphviz", "gunicorn", "kerberos", "memray", "otel", "statsd", "all-task-sdk", "airbyte", "akeyless", "alibaba", "amazon", "anthropic", "apache-cassandra", "apache-drill", "apache-druid", "apache-flink", "apache-hdfs", "apache-hive", "apache-iceberg", "apache-impala", "apache-kafka", "apache-kylin", "apache-livy", "apache-pig", "apache-pinot", "apache-spark", "apache-tinkerpop", "apprise", "arangodb", "asana", "atlassian-jira", "celery", "clickhousedb", "cloudant", "cncf-kubernetes", "cohere", "common-ai", "common-compat", "common-io", "common-messaging", "common-sql", "databricks", "datadog", "dbt-cloud", "dingding", "discord", "docker", "edge3", "elasticsearch", "exasol", "fab", "facebook", "ftp", "git", "github", "google", "grpc", "hashicorp", "http", "ibm-mq", "imap", "influxdb", "informatica", "jdbc", "jenkins", "keycloak", "microsoft-azure", "microsoft-mssql", "microsoft-psrp", "microsoft-winrm", "mongo", "mysql", "neo4j", "odbc", "openai", "openfaas", "openlineage", "opensearch", "opsgenie", "oracle", "pagerduty", "papermill", "pgvector", "pinecone", "postgres", "presto", "qdrant", "redis", "salesforce", "samba", "segment", "sendgrid", "sftp", "singularity", "slack", "smtp", "snowflake", "sqlite", "ssh", "standard", "tableau", "telegram", "teradata", "trino", "vertica", "vespa", "weaviate", "yandex", "ydb", "zendesk", "all", "aiobotocore", "apache-atlas", "apache-webhdfs", "amazon-aws-auth", "cloudpickle", "github-enterprise", "google-auth", "ldap", "pandas", "polars", "rabbitmq", "sentry", "s3fs", "uv"] [package.metadata.requires-dev] ci-image = [ @@ -3219,6 +3240,58 @@ dev = [ ] docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "devel-common" }] +[[package]] +name = "apache-airflow-providers-anthropic" +version = "0.1.0" +source = { editable = "providers/anthropic" } +dependencies = [ + { name = "anthropic" }, + { name = "apache-airflow" }, + { name = "apache-airflow-providers-common-compat" }, +] + +[package.optional-dependencies] +aws = [ + { name = "anthropic", extra = ["aws"] }, +] +bedrock = [ + { name = "anthropic", extra = ["bedrock"] }, +] +vertex = [ + { name = "anthropic", extra = ["vertex"] }, +] + +[package.dev-dependencies] +dev = [ + { name = "apache-airflow" }, + { name = "apache-airflow-devel-common" }, + { name = "apache-airflow-providers-common-compat" }, + { name = "apache-airflow-task-sdk" }, +] +docs = [ + { name = "apache-airflow-devel-common", extra = ["docs"] }, +] + +[package.metadata] +requires-dist = [ + { name = "anthropic", specifier = ">=0.101.0" }, + { name = "anthropic", extras = ["aws"], marker = "extra == 'aws'", specifier = ">=0.101.0" }, + { name = "anthropic", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=0.101.0" }, + { name = "anthropic", extras = ["vertex"], marker = "extra == 'vertex'", specifier = ">=0.101.0" }, + { name = "apache-airflow", editable = "." }, + { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, +] +provides-extras = ["bedrock", "vertex", "aws"] + +[package.metadata.requires-dev] +dev = [ + { name = "apache-airflow", editable = "." }, + { name = "apache-airflow-devel-common", editable = "devel-common" }, + { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, + { name = "apache-airflow-task-sdk", editable = "task-sdk" }, +] +docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "devel-common" }] + [[package]] name = "apache-airflow-providers-apache-cassandra" version = "3.9.5"