From afdb76f55dbe2c8bdb6da2df8d9a4afdfcb0ec92 Mon Sep 17 00:00:00 2001 From: Sameer Mesiah Date: Wed, 15 Jul 2026 22:20:54 +0100 Subject: [PATCH 1/2] Introduce SnowflakeCortexAgentOperator to execute Snowflake Cortex Agents using SnowflakeCortexAgentHook. Add unit tests, documentation, and an example DAG demonstrating how to invoke a Cortex Agent from an Airflow DAG. --- .../docs/operators/snowflake_cortex_agent.rst | 66 +++++++++ providers/snowflake/provider.yaml | 2 + .../providers/snowflake/get_provider_info.py | 2 + .../operators/snowflake_cortex_agent.py | 135 ++++++++++++++++++ .../example_snowflake_cortex_agent.py | 66 +++++++++ .../operators/test_snowflake_cortex_agent.py | 76 ++++++++++ 6 files changed, 347 insertions(+) create mode 100644 providers/snowflake/docs/operators/snowflake_cortex_agent.rst create mode 100644 providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py create mode 100644 providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py create mode 100644 providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py diff --git a/providers/snowflake/docs/operators/snowflake_cortex_agent.rst b/providers/snowflake/docs/operators/snowflake_cortex_agent.rst new file mode 100644 index 0000000000000..438519e6870c1 --- /dev/null +++ b/providers/snowflake/docs/operators/snowflake_cortex_agent.rst @@ -0,0 +1,66 @@ +.. 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/operator:SnowflakeCortexAgentOperator: + +SnowflakeCortexAgentOperator +============================ + +Use the :class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentOperator` +to execute `Snowflake 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: + + * Install the provider package via **pip**. + + .. code-block:: bash + + pip install 'apache-airflow-providers-snowflake' + + Detailed information is available for :doc:`Installation `. + + * :doc:`Setup a Snowflake Connection `. + + * Create a Snowflake Cortex Agent. See the + `Snowflake Cortex Agents documentation `__. + +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``. diff --git a/providers/snowflake/provider.yaml b/providers/snowflake/provider.yaml index 557f0246febea..0d1742567c99b 100644 --- a/providers/snowflake/provider.yaml +++ b/providers/snowflake/provider.yaml @@ -120,6 +120,7 @@ integrations: - /docs/apache-airflow-providers-snowflake/operators/snowflake.rst - /docs/apache-airflow-providers-snowflake/operators/snowpark.rst - /docs/apache-airflow-providers-snowflake/operators/snowpark_containers.rst + - /docs/apache-airflow-providers-snowflake/operators/snowflake_cortex_agent.rst logo: /docs/integration-logos/Snowflake.png tags: [service] @@ -129,6 +130,7 @@ operators: - airflow.providers.snowflake.operators.snowflake - airflow.providers.snowflake.operators.snowpark - airflow.providers.snowflake.operators.snowpark_containers + - airflow.providers.snowflake.operators.snowflake_cortex_agent task-decorators: - class-name: airflow.providers.snowflake.decorators.snowpark.snowpark_task diff --git a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py index 7e20af8004ee2..8c93b0f9e49bc 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py +++ b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py @@ -34,6 +34,7 @@ def get_provider_info(): "/docs/apache-airflow-providers-snowflake/operators/snowflake.rst", "/docs/apache-airflow-providers-snowflake/operators/snowpark.rst", "/docs/apache-airflow-providers-snowflake/operators/snowpark_containers.rst", + "/docs/apache-airflow-providers-snowflake/operators/snowflake_cortex_agent.rst", ], "logo": "/docs/integration-logos/Snowflake.png", "tags": ["service"], @@ -46,6 +47,7 @@ def get_provider_info(): "airflow.providers.snowflake.operators.snowflake", "airflow.providers.snowflake.operators.snowpark", "airflow.providers.snowflake.operators.snowpark_containers", + "airflow.providers.snowflake.operators.snowflake_cortex_agent", ], } ], 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 new file mode 100644 index 0000000000000..5b8822b73612c --- /dev/null +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py @@ -0,0 +1,135 @@ +# +# 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 collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.models import BaseOperator +from airflow.providers.snowflake.hooks.snowflake_cortex_agent import SnowflakeCortexAgentHook + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class SnowflakeCortexAgentOperator(BaseOperator): + """ + Execute 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 messages: Conversation messages to send to the agent. + :param thread_id: Existing conversation thread identifier. Optional. + Defaults to ``None``. + :param parent_message_id: Parent message identifier within the specified + thread. Required when ``thread_id`` is provided. Defaults to ``None``. + :param tool_choice: Tool selection configuration. Optional. Defaults to + ``None``. + :param models: Model configuration. Optional. Defaults to ``None``. + :param instructions: Agent instruction overrides. Optional. Defaults to + ``None``. + :param orchestration: Orchestration configuration. Optional. Defaults to + ``None``. + :param tools: Additional tools available to the agent. Optional. Defaults + to ``None``. + :param tool_resources: Configuration for tools specified in ``tools``. + Optional. Defaults to ``None``. + :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] = ("messages",) + + template_fields_renderers = { + "messages": "json", + } + + ui_color = "#29B5E8" + + def __init__( + self, + *, + database: str, + schema: str, + agent_name: str, + messages: list[dict[str, Any]], + thread_id: int | None = None, + parent_message_id: int | None = None, + tool_choice: 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.messages = messages + self.thread_id = thread_id + self.parent_message_id = parent_message_id + self.tool_choice = tool_choice + 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]: + """Execute the Snowflake Cortex Agent.""" + self.log.info( + "Executing Snowflake Cortex Agent '%s.%s.%s'.", + self.database, + self.schema, + self.agent_name, + ) + + return self.hook.run_agent( + database=self.database, + schema=self.schema, + agent_name=self.agent_name, + messages=self.messages, + thread_id=self.thread_id, + parent_message_id=self.parent_message_id, + tool_choice=self.tool_choice, + models=self.models, + instructions=self.instructions, + orchestration=self.orchestration, + tools=self.tools, + tool_resources=self.tool_resources, + 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 new file mode 100644 index 0000000000000..35a93bfadfca1 --- /dev/null +++ b/providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py @@ -0,0 +1,66 @@ +# +# 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. +""" +Example use of SnowflakeCortexAgentOperator. +""" + +from __future__ import annotations + +from datetime import datetime + +from airflow import DAG +from airflow.providers.snowflake.operators.snowflake_cortex_agent import ( + SnowflakeCortexAgentOperator, +) + +SNOWFLAKE_CONN_ID = "my_snowflake_conn" +DAG_ID = "example_snowflake_cortex_agent" + +with DAG( + DAG_ID, + start_date=datetime(2024, 1, 1), + schedule="@once", + default_args={"snowflake_conn_id": SNOWFLAKE_CONN_ID}, + tags=["example"], + catchup=False, +) as dag: + # [START howto_operator_snowflake_cortex_agent] + run_agent = SnowflakeCortexAgentOperator( + task_id="run_agent", + database="DEFAULT_DATABASE", + schema="DEFAULT_SCHEMA", + agent_name="default_agent", + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What can you help me with?", + } + ], + } + ], + ) + # [END howto_operator_snowflake_cortex_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/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py new file mode 100644 index 0000000000000..52a724fd847e1 --- /dev/null +++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py @@ -0,0 +1,76 @@ +# +# 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 + +from airflow.providers.snowflake.operators.snowflake_cortex_agent import ( + SnowflakeCortexAgentOperator, +) + +TASK_ID = "run_agent" +CONN_ID = "snowflake_default" + + +class TestSnowflakeCortexAgentOperator: + @mock.patch( + "airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentHook.run_agent" + ) + def test_execute(self, mock_run_agent): + """Test that the operator delegates execution to the hook.""" + response = {"content": [{"type": "text", "text": "Hello"}]} + mock_run_agent.return_value = response + + 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", + } + ], + ) + + result = operator.execute(context={}) + + mock_run_agent.assert_called_once_with( + database="MY_DATABASE", + schema="MY_SCHEMA", + agent_name="my_agent", + messages=[ + { + "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, + ) + + assert result == response From 550298bb4228d782be3910b73f064c67a7d8d138 Mon Sep 17 00:00:00 2001 From: Sameer Mesiah Date: Thu, 23 Jul 2026 22:40:12 +0100 Subject: [PATCH 2/2] Add templated fields and correct import for BaseOperator. --- .../operators/snowflake_cortex_agent.py | 9 ++++-- .../operators/test_snowflake_cortex_agent.py | 29 +++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) 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 5b8822b73612c..f00cde01a9a7a 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 @@ -22,7 +22,7 @@ from functools import cached_property from typing import TYPE_CHECKING, Any -from airflow.models import BaseOperator +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.snowflake.hooks.snowflake_cortex_agent import SnowflakeCortexAgentHook if TYPE_CHECKING: @@ -58,7 +58,12 @@ class SnowflakeCortexAgentOperator(BaseOperator): ``snowflake_default``. """ - template_fields: Sequence[str] = ("messages",) + template_fields: Sequence[str] = ( + "database", + "schema", + "agent_name", + "messages", + ) template_fields_renderers = { "messages": "json", 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 52a724fd847e1..130a2878f52b2 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,9 +19,11 @@ from unittest import mock +from airflow.models.dag import DAG from airflow.providers.snowflake.operators.snowflake_cortex_agent import ( SnowflakeCortexAgentOperator, ) +from airflow.utils import timezone TASK_ID = "run_agent" CONN_ID = "snowflake_default" @@ -74,3 +76,30 @@ def test_execute(self, mock_run_agent): ) assert result == response + + def test_template_fields(self): + dag = DAG( + dag_id="test_template_fields", + start_date=timezone.datetime(2024, 1, 1), + ) + + operator = SnowflakeCortexAgentOperator( + task_id=TASK_ID, + dag=dag, + database="{{ var.value.database }}", + schema="{{ params.schema }}", + agent_name="{{ dag_run.conf['agent_name'] }}", + messages=[ + { + "role": "user", + "content": "{{ ds }}", + } + ], + ) + + assert operator.template_fields == ( + "database", + "schema", + "agent_name", + "messages", + )