This is an automated email from the ASF dual-hosted git repository.
shahar1 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 6df6cb32ea5 Add SnowflakeCortexAgentOperator to Snowflake Provider
(#69939)
6df6cb32ea5 is described below
commit 6df6cb32ea5d2ae262ee342be4532039023e7a0d
Author: SameerMesiah97 <[email protected]>
AuthorDate: Mon Jul 27 13:57:03 2026 +0100
Add SnowflakeCortexAgentOperator to Snowflake Provider (#69939)
* 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.
* Add templated fields and correct import for BaseOperator.
---
.../docs/operators/snowflake_cortex_agent.rst | 66 ++++++++++
providers/snowflake/provider.yaml | 2 +
.../providers/snowflake/get_provider_info.py | 2 +
.../snowflake/operators/snowflake_cortex_agent.py | 140 +++++++++++++++++++++
.../snowflake/example_snowflake_cortex_agent.py | 66 ++++++++++
.../operators/test_snowflake_cortex_agent.py | 105 ++++++++++++++++
6 files changed, 381 insertions(+)
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 00000000000..438519e6870
--- /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
<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:
+
+ * Install the provider package via **pip**.
+
+ .. code-block:: bash
+
+ pip install 'apache-airflow-providers-snowflake'
+
+ 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>`__.
+
+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 557f0246feb..0d1742567c9 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 7e20af8004e..8c93b0f9e49 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 00000000000..f00cde01a9a
--- /dev/null
+++
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_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
+
+from collections.abc import Sequence
+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
+
+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] = (
+ "database",
+ "schema",
+ "agent_name",
+ "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 00000000000..35a93bfadfc
--- /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 00000000000..130a2878f52
--- /dev/null
+++
b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py
@@ -0,0 +1,105 @@
+#
+# 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.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"
+
+
+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
+
+ 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",
+ )
