This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 9bf8cc0a6fb Add Snowflake Cortex Agent management operators (#70930)
9bf8cc0a6fb is described below
commit 9bf8cc0a6fb0dc19f7579031da4acf445ff769a3
Author: SameerMesiah97 <[email protected]>
AuthorDate: Sat Oct 10 13:05:35 2026 +0100
Add Snowflake Cortex Agent management operators (#70930)
* Add Snowflake Cortex Agent management operators
Add create, update and delete operators backed by existing hook methods.
Include unit tests for hook delegation and template fields, and expand
the example DAG and documentation with management operator usage.
* Harden Snowflake Cortex Agent management operators
Validate create modes during DAG parsing and document accepted values.
Use autospecced hook mocks to verify operator defaults and signatures.
Make system-test agents unique, guarantee teardown, and reorganize the
operator documentation.
* Expand Cortex Agent operator tests and update documentation
Parametrize operator tests to cover defaults and all optional arguments.
Document that unset update fields retain their existing Snowflake values.
---------
Co-authored-by: Sameer Mesiah <[email protected]>
---
.../docs/operators/snowflake_cortex_agent.rst | 108 ++++---
.../snowflake/operators/snowflake_cortex_agent.py | 278 +++++++++++++++++-
.../snowflake/example_snowflake_cortex_agent.py | 60 +++-
.../operators/test_snowflake_cortex_agent.py | 309 +++++++++++++++++----
4 files changed, 667 insertions(+), 88 deletions(-)
diff --git a/providers/snowflake/docs/operators/snowflake_cortex_agent.rst
b/providers/snowflake/docs/operators/snowflake_cortex_agent.rst
index a3eab48031f..4f7d182addf 100644
--- a/providers/snowflake/docs/operators/snowflake_cortex_agent.rst
+++ b/providers/snowflake/docs/operators/snowflake_cortex_agent.rst
@@ -15,22 +15,18 @@
specific language governing permissions and limitations
under the License.
-.. _howto/operator:SnowflakeCortexAgentOperator:
-
-SnowflakeCortexAgentOperator
-============================
+Snowflake Cortex Agent operators
+================================
-Use the
:class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentOperator`
-to execute `Snowflake Cortex Agents
<https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-agents>`__.
+The Snowflake provider includes operators for creating, updating, executing,
+and deleting `Snowflake Cortex Agents
+<https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-agents>`__.
-The operator wraps the Snowflake Cortex Agent Run API and executes an existing
-Cortex Agent. It returns the JSON response payload from the agent, allowing
-responses to be consumed by downstream Airflow tasks through XCom.
Prerequisite Tasks
^^^^^^^^^^^^^^^^^^
-To use this operator, you must do a few things:
+To use these operators, you must do the following:
* Install the provider package via **pip**.
@@ -40,30 +36,10 @@ To use this operator, you must do a few things:
Detailed information is available for :doc:`Installation
<apache-airflow:installation/index>`.
- * :doc:`Setup a Snowflake Connection </connections/snowflake>`.
-
- * Create a Snowflake Cortex Agent. See the
- `Snowflake Cortex Agents documentation
<https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-agents>`__.
+ * :doc:`Set up a Snowflake connection </connections/snowflake>`.
-Using the Operator
-^^^^^^^^^^^^^^^^^^
-
-Use the ``snowflake_conn_id`` argument to specify the connection used. If not
-specified, ``snowflake_default`` will be used.
-
-An example usage of the ``SnowflakeCortexAgentOperator`` is as follows:
-
-.. exampleinclude::
/../../snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
- :language: python
- :start-after: [START howto_operator_snowflake_cortex_agent]
- :end-before: [END howto_operator_snowflake_cortex_agent]
- :dedent: 4
-
-.. note::
-
- Parameters passed to the operator take precedence over the corresponding
- values configured in the Airflow connection metadata, such as ``database``,
- ``schema`` and ``role``.
+Use the ``snowflake_conn_id`` argument to specify the connection used by an
+operator. If not specified, ``snowflake_default`` is used.
Authentication
^^^^^^^^^^^^^^
@@ -83,3 +59,69 @@ See the :doc:`Snowflake connection </connections/snowflake>`
page for the full f
A hook instance keeps its resolved key-pair JWT and reuses it within its
renewal window instead of
signing a new one on every request. This benefits code that calls
``SnowflakeCortexAgentHook``
several times on the same instance; the operator makes one request per
``execute`` call.
+
+.. _howto/operator:SnowflakeCortexAgentCreateOperator:
+
+SnowflakeCortexAgentCreateOperator
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To create a Snowflake Cortex Agent you can use
+:class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentCreateOperator`.
+
+.. exampleinclude::
/../../snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
+ :language: python
+ :start-after: [START howto_operator_snowflake_cortex_agent_create]
+ :end-before: [END howto_operator_snowflake_cortex_agent_create]
+ :dedent: 4
+
+.. _howto/operator:SnowflakeCortexAgentUpdateOperator:
+
+SnowflakeCortexAgentUpdateOperator
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To update an existing Snowflake Cortex Agent you can use
+:class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentUpdateOperator`.
+
+Only fields explicitly provided are updated. Optional fields left as ``None``
+retain their existing values on the Cortex Agent.
+
+.. exampleinclude::
/../../snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
+ :language: python
+ :start-after: [START howto_operator_snowflake_cortex_agent_update]
+ :end-before: [END howto_operator_snowflake_cortex_agent_update]
+ :dedent: 4
+
+.. _howto/operator:SnowflakeCortexAgentOperator:
+
+SnowflakeCortexAgentOperator
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Use the
:class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentOperator`
+to execute an existing Snowflake Cortex Agent.
+
+The operator wraps the Snowflake Cortex Agent Run API and executes an existing
+Cortex Agent. It returns the JSON response payload from the agent, allowing
+responses to be consumed by downstream Airflow tasks through XCom.
+
+An example usage of the ``SnowflakeCortexAgentOperator`` is as follows:
+
+.. exampleinclude::
/../../snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
+ :language: python
+ :start-after: [START howto_operator_snowflake_cortex_agent]
+ :end-before: [END howto_operator_snowflake_cortex_agent]
+ :dedent: 4
+
+
+.. _howto/operator:SnowflakeCortexAgentDeleteOperator:
+
+SnowflakeCortexAgentDeleteOperator
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To delete a Snowflake Cortex Agent you can use
+:class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentDeleteOperator`.
+
+.. exampleinclude::
/../../snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
+ :language: python
+ :start-after: [START howto_operator_snowflake_cortex_agent_delete]
+ :end-before: [END howto_operator_snowflake_cortex_agent_delete]
+ :dedent: 4
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py
index 35a586f2c59..f631bbb79df 100644
---
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py
+++
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py
@@ -23,7 +23,7 @@ from functools import cached_property
from typing import TYPE_CHECKING, Any
from airflow.providers.common.compat.sdk import BaseOperator
-from airflow.providers.snowflake.hooks.snowflake_cortex_agent import
SnowflakeCortexAgentHook
+from airflow.providers.snowflake.hooks.snowflake_cortex_agent import
CreateMode, SnowflakeCortexAgentHook
if TYPE_CHECKING:
from airflow.providers.common.compat.sdk import Context
@@ -139,3 +139,279 @@ class SnowflakeCortexAgentOperator(BaseOperator):
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. Accepted values are
+ ``errorIfExists``, ``orReplace`` and ``ifNotExists``.
+ Defaults to ``errorIfExists``.
+ :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 = CreateMode(create_mode)
+ self.timeout = timeout
+ self.snowflake_conn_id = snowflake_conn_id
+
+ @cached_property
+ def hook(self) -> SnowflakeCortexAgentHook:
+ """Return the Snowflake Cortex Agent hook."""
+ return SnowflakeCortexAgentHook(
+ snowflake_conn_id=self.snowflake_conn_id,
+ )
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ """Create the Snowflake Cortex Agent."""
+ self.log.info(
+ "Creating Snowflake Cortex Agent '%s.%s.%s'.",
+ self.database,
+ self.schema,
+ self.agent_name,
+ )
+
+ return self.hook.create_agent(
+ database=self.database,
+ schema=self.schema,
+ agent_name=self.agent_name,
+ comment=self.comment,
+ profile=self.profile,
+ models=self.models,
+ instructions=self.instructions,
+ orchestration=self.orchestration,
+ tools=self.tools,
+ tool_resources=self.tool_resources,
+ create_mode=self.create_mode,
+ timeout=self.timeout,
+ )
+
+
+class SnowflakeCortexAgentUpdateOperator(BaseOperator):
+ """
+ Update a Snowflake Cortex Agent.
+
+ Only fields explicitly provided are included in the update request. Fields
left
+ as ``None`` retain their current values in Snowflake.
+
+ :param database: Database containing the Cortex Agent.
+ :param schema: Schema containing 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 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,
+ 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.timeout = timeout
+ self.snowflake_conn_id = snowflake_conn_id
+
+ @cached_property
+ def hook(self) -> SnowflakeCortexAgentHook:
+ """Return the Snowflake Cortex Agent hook."""
+ return SnowflakeCortexAgentHook(
+ snowflake_conn_id=self.snowflake_conn_id,
+ )
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ """Update the Snowflake Cortex Agent."""
+ self.log.info(
+ "Updating Snowflake Cortex Agent '%s.%s.%s'.",
+ self.database,
+ self.schema,
+ self.agent_name,
+ )
+
+ return self.hook.update_agent(
+ database=self.database,
+ schema=self.schema,
+ agent_name=self.agent_name,
+ comment=self.comment,
+ profile=self.profile,
+ models=self.models,
+ instructions=self.instructions,
+ orchestration=self.orchestration,
+ tools=self.tools,
+ tool_resources=self.tool_resources,
+ timeout=self.timeout,
+ )
+
+
+class SnowflakeCortexAgentDeleteOperator(BaseOperator):
+ """
+ Delete a Snowflake Cortex Agent.
+
+ :param database: Database containing the Cortex Agent.
+ :param schema: Schema containing the Cortex Agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param if_exists: Ignore the request if the agent does not exist.
+ Defaults to ``False``.
+ :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",
+ "snowflake_conn_id",
+ )
+
+ ui_color = "#29B5E8"
+
+ def __init__(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ if_exists: bool = False,
+ 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.if_exists = if_exists
+ self.timeout = timeout
+ self.snowflake_conn_id = snowflake_conn_id
+
+ @cached_property
+ def hook(self) -> SnowflakeCortexAgentHook:
+ """Return the Snowflake Cortex Agent hook."""
+ return SnowflakeCortexAgentHook(
+ snowflake_conn_id=self.snowflake_conn_id,
+ )
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ """Delete the Snowflake Cortex Agent."""
+ self.log.info(
+ "Deleting Snowflake Cortex Agent '%s.%s.%s'.",
+ self.database,
+ self.schema,
+ self.agent_name,
+ )
+
+ return self.hook.delete_agent(
+ database=self.database,
+ schema=self.schema,
+ agent_name=self.agent_name,
+ if_exists=self.if_exists,
+ timeout=self.timeout,
+ )
diff --git
a/providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
b/providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
index 35a93bfadfc..8b579e256cc 100644
---
a/providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
+++
b/providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
@@ -16,21 +16,34 @@
# specific language governing permissions and limitations
# under the License.
"""
-Example use of SnowflakeCortexAgentOperator.
+Example use of Snowflake Cortex Agent operators.
+
+This example creates, updates, executes, and deletes a Snowflake Cortex Agent.
"""
from __future__ import annotations
+import os
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"
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
+
+DATABASE = "DEFAULT_DATABASE"
+SCHEMA = "DEFAULT_SCHEMA"
+AGENT_NAME = f"airflow_cortex_agent_{ENV_ID}"
+
with DAG(
DAG_ID,
start_date=datetime(2024, 1, 1),
@@ -39,12 +52,39 @@ with DAG(
tags=["example"],
catchup=False,
) as dag:
+ # [START howto_operator_snowflake_cortex_agent_create]
+ create_agent = SnowflakeCortexAgentCreateOperator(
+ task_id="create_agent",
+ database=DATABASE,
+ schema=SCHEMA,
+ agent_name=AGENT_NAME,
+ comment="Created by Airflow",
+ instructions={
+ "response": "Respond in a friendly and concise manner.",
+ },
+ create_mode=CreateMode.ERROR_IF_EXISTS,
+ )
+ # [END howto_operator_snowflake_cortex_agent_create]
+
+ # [START howto_operator_snowflake_cortex_agent_update]
+ update_agent = SnowflakeCortexAgentUpdateOperator(
+ task_id="update_agent",
+ database=DATABASE,
+ schema=SCHEMA,
+ agent_name=AGENT_NAME,
+ comment="Updated by Airflow",
+ instructions={
+ "response": "Respond in one sentence.",
+ },
+ )
+ # [END howto_operator_snowflake_cortex_agent_update]
+
# [START howto_operator_snowflake_cortex_agent]
run_agent = SnowflakeCortexAgentOperator(
task_id="run_agent",
- database="DEFAULT_DATABASE",
- schema="DEFAULT_SCHEMA",
- agent_name="default_agent",
+ database=DATABASE,
+ schema=SCHEMA,
+ agent_name=AGENT_NAME,
messages=[
{
"role": "user",
@@ -59,6 +99,18 @@ with DAG(
)
# [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.as_teardown(setups=create_agent)
+
from tests_common.test_utils.system_tests import get_test_run # noqa: E402
diff --git
a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py
b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py
index 5777e1d436b..226e1a224e0 100644
---
a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py
+++
b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py
@@ -19,88 +19,297 @@ from __future__ import annotations
from unittest import mock
-from airflow.models.dag import DAG
-from airflow.providers.common.compat.sdk import timezone
+import pytest
+
+from airflow.providers.snowflake.hooks.snowflake_cortex_agent import (
+ CreateMode,
+ SnowflakeCortexAgentHook,
+)
from airflow.providers.snowflake.operators.snowflake_cortex_agent import (
+ SnowflakeCortexAgentCreateOperator,
+ SnowflakeCortexAgentDeleteOperator,
SnowflakeCortexAgentOperator,
+ SnowflakeCortexAgentUpdateOperator,
)
TASK_ID = "run_agent"
-CONN_ID = "snowflake_default"
class TestSnowflakeCortexAgentOperator:
- @mock.patch(
-
"airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentHook.run_agent"
+ @pytest.mark.parametrize(
+ ("operator_kwargs", "expected_kwargs"),
+ [
+ pytest.param(
+ {},
+ {
+ "thread_id": None,
+ "parent_message_id": None,
+ "tool_choice": None,
+ "models": None,
+ "instructions": None,
+ "orchestration": None,
+ "tools": None,
+ "tool_resources": None,
+ "timeout": 600,
+ },
+ id="defaults",
+ ),
+ pytest.param(
+ {
+ "thread_id": "thread-id",
+ "parent_message_id": "parent-message-id",
+ "tool_choice": {"type": "auto"},
+ "models": {"orchestration": "test-model"},
+ "instructions": {"response": "Answer concisely."},
+ "orchestration": {"budget": {"seconds": 30}},
+ "tools": [{"tool_spec": {"name": "search"}}],
+ "tool_resources": {"search": {"name": "search_service"}},
+ "timeout": 300,
+ },
+ {
+ "thread_id": "thread-id",
+ "parent_message_id": "parent-message-id",
+ "tool_choice": {"type": "auto"},
+ "models": {"orchestration": "test-model"},
+ "instructions": {"response": "Answer concisely."},
+ "orchestration": {"budget": {"seconds": 30}},
+ "tools": [{"tool_spec": {"name": "search"}}],
+ "tool_resources": {"search": {"name": "search_service"}},
+ "timeout": 300,
+ },
+ id="all-arguments",
+ ),
+ ],
+ )
+ @mock.patch.object(
+ SnowflakeCortexAgentHook,
+ "run_agent",
+ autospec=True,
)
- def test_execute(self, mock_run_agent):
- """Test that the operator delegates execution to the hook."""
+ def test_execute(self, mock_run_agent, operator_kwargs, expected_kwargs):
response = {"content": [{"type": "text", "text": "Hello"}]}
mock_run_agent.return_value = response
+ messages = [{"role": "user", "content": "Hello"}]
operator = SnowflakeCortexAgentOperator(
task_id=TASK_ID,
- snowflake_conn_id=CONN_ID,
database="MY_DATABASE",
schema="MY_SCHEMA",
agent_name="my_agent",
- messages=[
- {
- "role": "user",
- "content": "Hello",
- }
- ],
+ messages=messages,
+ **operator_kwargs,
)
result = operator.execute(context={})
mock_run_agent.assert_called_once_with(
+ mock.ANY,
database="MY_DATABASE",
schema="MY_SCHEMA",
agent_name="my_agent",
- messages=[
+ messages=messages,
+ **expected_kwargs,
+ )
+ assert result == response
+
+
+class TestSnowflakeCortexAgentCreateOperator:
+ @pytest.mark.parametrize(
+ ("operator_kwargs", "expected_kwargs"),
+ [
+ pytest.param(
+ {},
+ {
+ "comment": None,
+ "profile": None,
+ "models": None,
+ "instructions": None,
+ "orchestration": None,
+ "tools": None,
+ "tool_resources": None,
+ "create_mode": CreateMode.ERROR_IF_EXISTS,
+ "timeout": 600,
+ },
+ id="defaults",
+ ),
+ pytest.param(
{
- "role": "user",
- "content": "Hello",
- }
- ],
- thread_id=None,
- parent_message_id=None,
- tool_choice=None,
- models=None,
- instructions=None,
- orchestration=None,
- tools=None,
- tool_resources=None,
- timeout=600,
+ "comment": "Created by Airflow",
+ "profile": {"display_name": "Test Agent"},
+ "models": {"orchestration": "test-model"},
+ "instructions": {"response": "Answer concisely."},
+ "orchestration": {"budget": {"seconds": 30}},
+ "tools": [{"tool_spec": {"name": "search"}}],
+ "tool_resources": {"search": {"name": "search_service"}},
+ "create_mode": "orReplace",
+ "timeout": 300,
+ },
+ {
+ "comment": "Created by Airflow",
+ "profile": {"display_name": "Test Agent"},
+ "models": {"orchestration": "test-model"},
+ "instructions": {"response": "Answer concisely."},
+ "orchestration": {"budget": {"seconds": 30}},
+ "tools": [{"tool_spec": {"name": "search"}}],
+ "tool_resources": {"search": {"name": "search_service"}},
+ "create_mode": CreateMode.OR_REPLACE,
+ "timeout": 300,
+ },
+ id="all-arguments",
+ ),
+ ],
+ )
+ @mock.patch.object(
+ SnowflakeCortexAgentHook,
+ "create_agent",
+ autospec=True,
+ )
+ def test_execute(self, mock_create_agent, operator_kwargs,
expected_kwargs):
+ response = {"status": "created"}
+ mock_create_agent.return_value = response
+
+ operator = SnowflakeCortexAgentCreateOperator(
+ task_id="create_agent",
+ database="MY_DATABASE",
+ schema="MY_SCHEMA",
+ agent_name="my_agent",
+ **operator_kwargs,
)
- assert result == response
+ result = operator.execute(context={})
- def test_template_fields(self):
- dag = DAG(
- dag_id="test_template_fields",
- start_date=timezone.datetime(2024, 1, 1),
+ mock_create_agent.assert_called_once_with(
+ mock.ANY,
+ database="MY_DATABASE",
+ schema="MY_SCHEMA",
+ agent_name="my_agent",
+ **expected_kwargs,
)
+ assert result == response
- operator = SnowflakeCortexAgentOperator(
- task_id=TASK_ID,
- dag=dag,
- database="{{ var.value.database }}",
- schema="{{ params.schema }}",
- agent_name="{{ dag_run.conf['agent_name'] }}",
- messages=[
+ def test_invalid_create_mode(self):
+ with pytest.raises(ValueError, match="'OR_REPLACE' is not a valid
CreateMode"):
+ SnowflakeCortexAgentCreateOperator(
+ task_id="create_agent",
+ database="MY_DATABASE",
+ schema="MY_SCHEMA",
+ agent_name="my_agent",
+ create_mode="OR_REPLACE",
+ )
+
+
+class TestSnowflakeCortexAgentUpdateOperator:
+ @pytest.mark.parametrize(
+ ("operator_kwargs", "expected_kwargs"),
+ [
+ pytest.param(
+ {},
+ {
+ "comment": None,
+ "profile": None,
+ "models": None,
+ "instructions": None,
+ "orchestration": None,
+ "tools": None,
+ "tool_resources": None,
+ "timeout": 600,
+ },
+ id="defaults",
+ ),
+ pytest.param(
{
- "role": "user",
- "content": "{{ ds }}",
- }
- ],
+ "comment": "Updated by Airflow",
+ "profile": {"display_name": "Updated Agent"},
+ "models": {"orchestration": "updated-model"},
+ "instructions": {"response": "Answer briefly."},
+ "orchestration": {"budget": {"seconds": 60}},
+ "tools": [{"tool_spec": {"name": "analyst"}}],
+ "tool_resources": {"analyst": {"semantic_view":
"MY_VIEW"}},
+ "timeout": 300,
+ },
+ {
+ "comment": "Updated by Airflow",
+ "profile": {"display_name": "Updated Agent"},
+ "models": {"orchestration": "updated-model"},
+ "instructions": {"response": "Answer briefly."},
+ "orchestration": {"budget": {"seconds": 60}},
+ "tools": [{"tool_spec": {"name": "analyst"}}],
+ "tool_resources": {"analyst": {"semantic_view":
"MY_VIEW"}},
+ "timeout": 300,
+ },
+ id="all-arguments",
+ ),
+ ],
+ )
+ @mock.patch.object(
+ SnowflakeCortexAgentHook,
+ "update_agent",
+ autospec=True,
+ )
+ def test_execute(self, mock_update_agent, operator_kwargs,
expected_kwargs):
+ response = {}
+ mock_update_agent.return_value = response
+
+ operator = SnowflakeCortexAgentUpdateOperator(
+ task_id="update_agent",
+ database="MY_DATABASE",
+ schema="MY_SCHEMA",
+ agent_name="my_agent",
+ **operator_kwargs,
+ )
+
+ result = operator.execute(context={})
+
+ mock_update_agent.assert_called_once_with(
+ mock.ANY,
+ database="MY_DATABASE",
+ schema="MY_SCHEMA",
+ agent_name="my_agent",
+ **expected_kwargs,
+ )
+ assert result == response
+
+
+class TestSnowflakeCortexAgentDeleteOperator:
+ @pytest.mark.parametrize(
+ ("operator_kwargs", "expected_kwargs"),
+ [
+ pytest.param(
+ {},
+ {"if_exists": False, "timeout": 600},
+ id="defaults",
+ ),
+ pytest.param(
+ {"if_exists": True, "timeout": 300},
+ {"if_exists": True, "timeout": 300},
+ id="all-arguments",
+ ),
+ ],
+ )
+ @mock.patch.object(
+ SnowflakeCortexAgentHook,
+ "delete_agent",
+ autospec=True,
+ )
+ def test_execute(self, mock_delete_agent, operator_kwargs,
expected_kwargs):
+ response = {"status": "deleted"}
+ mock_delete_agent.return_value = response
+
+ operator = SnowflakeCortexAgentDeleteOperator(
+ task_id="delete_agent",
+ database="MY_DATABASE",
+ schema="MY_SCHEMA",
+ agent_name="my_agent",
+ **operator_kwargs,
)
- assert operator.template_fields == (
- "database",
- "schema",
- "agent_name",
- "messages",
- "snowflake_conn_id",
+ result = operator.execute(context={})
+
+ mock_delete_agent.assert_called_once_with(
+ mock.ANY,
+ database="MY_DATABASE",
+ schema="MY_SCHEMA",
+ agent_name="my_agent",
+ **expected_kwargs,
)
+ assert result == response
