[SameerMesiah97]

Description

This change adds management operators for Snowflake Cortex Agents.

The Snowflake provider now includes operators to create, update and delete Cortex Agents, complementing the existing SnowflakeCortexAgentOperator used to execute them. Each operator is a thin wrapper around the corresponding hook method and follows the existing provider patterns.

Rationale

The existing hook already exposes the Snowflake Cortex Agent management APIs. Providing corresponding operators allows these operations to be used directly in DAGs without requiring users to invoke the hook from custom Python tasks.

This completes the operator support for the Cortex Agent lifecycle while remaining consistent with the abstractions provided throughout the Snowflake provider.

Tests

Added unit tests verifying that:

• SnowflakeCortexAgentCreateOperator delegates to create_agent() with the expected arguments.
• SnowflakeCortexAgentUpdateOperator delegates to update_agent() with the expected arguments.
• SnowflakeCortexAgentDeleteOperator delegates to delete_agent() with the expected arguments.
• Template fields are correctly configured for each new operator.

Notes

Kindly note that SnowflakeCortexAgentHook will be aligned with BaseManagedAgentHook introduced in PR #73532 in a follow-up PR.

Documentation

The example DAG and operator documentation have also been expanded to include usage examples for each management operator.

Backwards Compatibility

This change is fully backwards compatible. The existing SnowflakeCortexAgentOperator is unchanged. The new management operators are additive and do not modify the behaviour of existing hooks or operators.

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

Generated-by: [GPT 5.5] following the guidelines

[SameerMesiah97]

Depends on PR #70703 and PR #70101. Please do not review this until those 2 PRs are merged.

[vikramkoka]

@kaxil This is similar to the conversation we just had about the other PRs.

@SameerMesiah97 would like to chat about this.

[SameerMesiah97]

@kaxil This is similar to the conversation we just had about the other PRs.

@SameerMesiah97 would like to chat about this.

I would appreciate if you could provide more context here. Is there a blocking concern you wish for me to address? If you could link the existing conversation (if public) so that I can better understand the situation, that would be great.

[vikramkoka]

@kaxil This is similar to the conversation we just had about the other PRs.
@SameerMesiah97 would like to chat about this.

I would appreciate if you could provide more context here. Is there a blocking concern you wish for me to address? If you could link the existing conversation (if public) so that I can better understand the situation, that would be great.

Hey @SameerMesiah97 ,

Sorry, I was not being deliberately cryptic :)

I had seen a couple of very similar PRs to this from you (for Cortex) for other frameworks such as Google Vertex AI, etc.
They had a similar pattern of xxAgentCRUD operations.

I was verbally chatting with @kaxil about making the interfaces consistent across these, so that DAG authors don't have to learn multiple invocation mechanisms.
And then, starting flagging PRs for him to look at.

Hopefully, that clarifies?

[SameerMesiah97]

@kaxil This is similar to the conversation we just had about the other PRs.
@SameerMesiah97 would like to chat about this.

I would appreciate if you could provide more context here. Is there a blocking concern you wish for me to address? If you could link the existing conversation (if public) so that I can better understand the situation, that would be great.

Hey @SameerMesiah97 ,

Sorry, I was not being deliberately cryptic :)

I had seen a couple of very similar PRs to this from you (for Cortex) for other frameworks such as Google Vertex AI, etc. They had a similar pattern of xxAgentCRUD operations.

I was verbally chatting with @kaxil about making the interfaces consistent across these, so that DAG authors don't have to learn multiple invocation mechanisms. And then, starting flagging PRs for him to look at.

Hopefully, that clarifies?

Okay. I understand completely.

[gopidesupavan]

hey @SameerMesiah97 can you align to this interaface #71946

cc: @vikramkoka

[SameerMesiah97]

hey @SameerMesiah97 can you align to this interaface #71946

cc: @vikramkoka

Sure. But keep in mind that this is dependent on PR #70101 and PR #70703 being merged first.

[kaxil]

@SameerMesiah97 Could you post some screenshot of Airflow UI and Snowflake UI on how this all works e2e

[SameerMesiah97]

@SameerMesiah97 Could you post some screenshot of Airflow UI and Snowflake UI on how this all works e2e

I have ran the following DAG that uses all 4 Cortex Agent operators:

from __future__ import annotations

from airflow import DAG
from airflow.sdk import timezone

from airflow.providers.snowflake.hooks.snowflake_cortex_agent import CreateMode
from airflow.providers.snowflake.operators.snowflake_cortex_agent import (
    SnowflakeCortexAgentCreateOperator,
    SnowflakeCortexAgentDeleteOperator,
    SnowflakeCortexAgentOperator,
    SnowflakeCortexAgentUpdateOperator,
)

DATABASE = "<REDACTED>"
SCHEMA = "<REDACTED>"
AGENT_NAME = "AIRFLOW_TEST_AGENT"

with DAG(
    dag_id="test_snowflake_cortex_agent_operators",
    start_date=timezone.datetime(2025, 1, 1),
    schedule=None,
    catchup=False,
    tags=["snowflake", "cortex"],
) as dag:
    create_agent = SnowflakeCortexAgentCreateOperator(
        task_id="create_agent",
        snowflake_conn_id="snowflake_default",
        database=DATABASE,
        schema=SCHEMA,
        agent_name=AGENT_NAME,
        comment="Created by Airflow",
        instructions={
            "response": "Respond in a friendly and concise manner.",
        },
        create_mode=CreateMode.OR_REPLACE,
    )

    update_agent = SnowflakeCortexAgentUpdateOperator(
        task_id="update_agent",
        snowflake_conn_id="snowflake_default",
        database=DATABASE,
        schema=SCHEMA,
        agent_name=AGENT_NAME,
        comment="Updated by Airflow",
        instructions={
            "response": "Always answer in one sentence.",
        },
    )

    run_agent = SnowflakeCortexAgentOperator(
        task_id="run_agent",
        snowflake_conn_id="snowflake_default",
        database=DATABASE,
        schema=SCHEMA,
        agent_name=AGENT_NAME,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "What instructions are you currently using?",
                    }
                ],
            }
        ],
    )

    delete_agent = SnowflakeCortexAgentDeleteOperator(
        task_id="delete_agent",
        snowflake_conn_id="snowflake_default",
        database=DATABASE,
        schema=SCHEMA,
        agent_name=AGENT_NAME,
        if_exists=True,
    )

    create_agent >> update_agent >> run_agent >> delete_agent

All tasks executed sucessfully:

Because the DAG deletes the agent, we cannot see the actual thread in the Snowflake UI that is created when the agent is prompted after the DAG finishes execution. However, the return values (which are returned by the API) are pushed to XCom and we can see them in the Airflow UI. For example, this is what I see when for the create_agent task:

This is what I see for the run_agent task:

And this is what I see when for delete_agent task:

Keep in mind that update agent endpoint does not return a non-empty payload. But after running a modified version of the above DAG (with only the create_agent and update_agent tasks), I can see this in the Snowflake UI:

[SameerMesiah97]

CI failure appears unrelated.