kaxil commented on code in PR #70930:
URL: https://github.com/apache/airflow/pull/70930#discussion_r4195056801
##########
providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py:
##########
@@ -24,13 +24,21 @@
from datetime import datetime
from airflow import DAG
+from airflow.providers.snowflake.hooks.snowflake_cortex_agent import CreateMode
from airflow.providers.snowflake.operators.snowflake_cortex_agent import (
+ SnowflakeCortexAgentCreateOperator,
+ SnowflakeCortexAgentDeleteOperator,
SnowflakeCortexAgentOperator,
+ SnowflakeCortexAgentUpdateOperator,
)
SNOWFLAKE_CONN_ID = "my_snowflake_conn"
DAG_ID = "example_snowflake_cortex_agent"
+DATABASE = "DEFAULT_DATABASE"
+SCHEMA = "DEFAULT_SCHEMA"
+AGENT_NAME = "default_agent"
Review Comment:
Now that this Dag creates the agent, could the name include
`SYSTEM_TESTS_ENV_ID` like `example_snowflake.py` does? Two concurrent system
test runs would otherwise collide on `default_agent`. The module docstring on
line 19 also still describes this as an example of
`SnowflakeCortexAgentOperator` only.
##########
providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py:
##########
@@ -59,6 +94,18 @@
)
# [END howto_operator_snowflake_cortex_agent]
+ # [START howto_operator_snowflake_cortex_agent_delete]
+ delete_agent = SnowflakeCortexAgentDeleteOperator(
+ task_id="delete_agent",
+ database=DATABASE,
+ schema=SCHEMA,
+ agent_name=AGENT_NAME,
+ if_exists=True,
+ )
+ # [END howto_operator_snowflake_cortex_agent_delete]
+
+ create_agent >> update_agent >> run_agent >> delete_agent
Review Comment:
If `update_agent` or `run_agent` fails, `delete_agent` ends up
`upstream_failed` and the agent stays in Snowflake. The next run's create then
hits `ERROR_IF_EXISTS` on the same fixed name and fails before it reaches
cleanup, so every rerun fails until someone drops the agent by hand.
`delete_agent.as_teardown(setups=create_agent)` would fix that: it runs only
when the create succeeded, and the Dag run still reports the upstream failure.
The Bedrock AgentCore and Vertex Agent Engine examples handle the same case
with `ALL_DONE` plus `watcher()`.
##########
providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py:
##########
@@ -104,3 +108,173 @@ def test_template_fields(self):
"messages",
"snowflake_conn_id",
)
+
+
+class TestSnowflakeCortexAgentCreateOperator:
+ @mock.patch(
+
"airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentHook.create_agent"
Review Comment:
These patches have no `autospec`, so the mock accepts any keyword. The
operators only forward arguments, so the hook signature is the main thing worth
pinning: `mock.patch.object(SnowflakeCortexAgentHook, "create_agent",
autospec=True)` (with `mock.ANY` for `self`) would catch drift. Every test also
passes every optional argument, so flipping a default like `ERROR_IF_EXISTS` or
`if_exists=False` leaves the suite green. A test per operator with only the
required arguments would cover that. Same for the update and delete tests below.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py:
##########
@@ -139,3 +139,275 @@ def execute(self, context: Context) -> dict[str, Any]:
tool_resources=self.tool_resources,
timeout=self.timeout,
)
+
+
+class SnowflakeCortexAgentCreateOperator(BaseOperator):
Review Comment:
On the open ask to align with #71946: `BaseManagedAgentToolset` there, and
`BaseManagedAgentHook` from #73532, only cover invoking an agent
(`resolve_agent`, `get_agent_capabilities`, `invoke_agent`). Neither has
create, update or delete, so there isn't a shared shape for these lifecycle
operators to follow yet. The alignment that does apply is
`SnowflakeCortexAgentHook` implementing `BaseManagedAgentHook`, the way
`BedrockAgentCoreHook` and `AgentEngineHook` do, and that can be a separate PR.
Could you say so in the description so the earlier thread is closed out?
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py:
##########
@@ -139,3 +139,275 @@ def execute(self, context: Context) -> dict[str, Any]:
tool_resources=self.tool_resources,
timeout=self.timeout,
)
+
+
+class SnowflakeCortexAgentCreateOperator(BaseOperator):
+ """
+ Create a Snowflake Cortex Agent.
+
+ :param database: Database in which to create the Cortex Agent.
+ :param schema: Schema in which to create the Cortex Agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param comment: Optional comment.
+ :param profile: Agent profile configuration. Optional.
+ :param models: Model configuration. Optional.
+ :param instructions: Agent instructions. Optional.
+ :param orchestration: Orchestration configuration. Optional.
+ :param tools: Agent tools. Optional.
+ :param tool_resources: Tool resource configuration. Optional.
+ :param create_mode: Resource creation mode. Defaults to
+ ``CreateMode.ERROR_IF_EXISTS``.
+ :param timeout: Maximum time in seconds to wait for the request to
+ complete. Defaults to ``600``.
+ :param snowflake_conn_id: Snowflake connection ID. Defaults to
+ ``snowflake_default``.
+ """
+
+ template_fields: Sequence[str] = (
+ "database",
+ "schema",
+ "agent_name",
+ "comment",
+ "profile",
+ "models",
+ "instructions",
+ "orchestration",
+ "tools",
+ "tool_resources",
+ "snowflake_conn_id",
+ )
+
+ ui_color = "#29B5E8"
+
+ def __init__(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ comment: str | None = None,
+ profile: dict[str, Any] | None = None,
+ models: dict[str, Any] | None = None,
+ instructions: dict[str, Any] | None = None,
+ orchestration: dict[str, Any] | None = None,
+ tools: list[dict[str, Any]] | None = None,
+ tool_resources: dict[str, Any] | None = None,
+ create_mode: CreateMode | str = CreateMode.ERROR_IF_EXISTS,
+ timeout: int | None = 600,
+ snowflake_conn_id: str = "snowflake_default",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+
+ self.database = database
+ self.schema = schema
+ self.agent_name = agent_name
+ self.comment = comment
+ self.profile = profile
+ self.models = models
+ self.instructions = instructions
+ self.orchestration = orchestration
+ self.tools = tools
+ self.tool_resources = tool_resources
+ self.create_mode = create_mode
Review Comment:
`create_mode` isn't templated, so it can be validated here with
`self.create_mode = CreateMode(create_mode)`. Today a value like `"OR_REPLACE"`
(the member name rather than the value) passes Dag parsing and only fails with
a `ValueError` when the task runs. The docstring could also list the accepted
values (`errorIfExists`, `orReplace`, `ifNotExists`) the way the hook's does.
##########
providers/snowflake/docs/operators/snowflake_cortex_agent.rst:
##########
@@ -64,3 +64,45 @@ An example usage of the ``SnowflakeCortexAgentOperator`` is
as follows:
Parameters passed to the operator take precedence over the corresponding
values configured in the Airflow connection metadata, such as ``database``,
``schema`` and ``role``.
+
+.. _howto/operator:SnowflakeCortexAgentCreateOperator:
+
+SnowflakeCortexAgentCreateOperator
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To create a Snowflake Cortex Agent you can use
Review Comment:
The Prerequisite Tasks section above still tells users to create the agent
outside Airflow, which this operator now does. These three sections also use
`^^^`, so they nest under the `SnowflakeCortexAgentOperator` title in the TOC.
Retitling the page (something like "Snowflake Cortex Agent operators") or
promoting these headings to `===` would fix the nesting.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
