Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -103,10 +103,6 @@ def __init__(
**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
Expand All @@ -126,6 +122,10 @@ def hook(self) -> AnthropicHook:
return AnthropicHook(conn_id=self.conn_id)

def execute(self, context: Context) -> str | None:
if (self.message is None) == (self.outcome is None):
raise ValueError("Provide exactly one of 'message' or 'outcome'.")
if self.outcome is not None and not {"description", "rubric"} <= self.outcome.keys():
raise ValueError("'outcome' must include both 'description' and 'rubric'.")
create_kwargs: dict[str, Any] = dict(self.session_kwargs)
if self.vault_ids:
create_kwargs["vault_ids"] = self.vault_ids
Expand Down
44 changes: 31 additions & 13 deletions providers/anthropic/tests/unit/anthropic/operators/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,28 +38,46 @@ def _context():


def test_requires_exactly_one_of_message_or_outcome():
op = 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")
op.execute(_context())

op = AnthropicAgentSessionOperator(
task_id="a", agent_id="ag", environment_id="env", message="hi", outcome={"description": "x"}
)
with pytest.raises(ValueError, match="exactly one"):
AnthropicAgentSessionOperator(
task_id="a", agent_id="ag", environment_id="env", message="hi", outcome={"description": "x"}
)
op.execute(_context())


def test_outcome_requires_description_and_rubric():
# missing rubric
op = AnthropicAgentSessionOperator(
task_id="a", agent_id="ag", environment_id="env", outcome={"description": "x"}
)
with pytest.raises(ValueError, match="description.*rubric"):
AnthropicAgentSessionOperator(
task_id="a", agent_id="ag", environment_id="env", outcome={"description": "x"}
)
op.execute(_context())

# missing description
op = AnthropicAgentSessionOperator(
task_id="a",
agent_id="ag",
environment_id="env",
outcome={"rubric": {"type": "text", "content": "c"}},
)
with pytest.raises(ValueError, match="description.*rubric"):
AnthropicAgentSessionOperator(
task_id="a",
agent_id="ag",
environment_id="env",
outcome={"rubric": {"type": "text", "content": "c"}},
)
op.execute(_context())


def test_init_does_not_validate_message_or_outcome():
"""Regression test: __init__ must not read template-field values (see #70296)."""
op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env")
assert op.message is None
assert op.outcome is None

op = AnthropicAgentSessionOperator(
task_id="a", agent_id="ag", environment_id="env", outcome={"description": "x"}
)
assert op.outcome == {"description": "x"}


class TestExecute:
Expand Down
1 change: 0 additions & 1 deletion scripts/ci/prek/validate_operators_init_exemptions.txt
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ providers/amazon/src/airflow/providers/amazon/aws/operators/step_function.py::St
providers/amazon/src/airflow/providers/amazon/aws/transfers/base.py::AwsToAwsBaseOperator
providers/amazon/src/airflow/providers/amazon/aws/transfers/gcs_to_s3.py::GCSToS3Operator
providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py::S3ToRedshiftOperator
providers/anthropic/src/airflow/providers/anthropic/operators/agent.py::AnthropicAgentSessionOperator
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py::KubernetesPodOperator
providers/docker/src/airflow/providers/docker/operators/docker.py::DockerOperator
providers/google/src/airflow/providers/google/cloud/operators/bigquery.py::BigQueryInsertJobOperator
Expand Down