Snowflake Cortex Agent operators¶

The Snowflake provider includes operators for creating, updating, executing, and deleting Snowflake Cortex Agents.

Prerequisite Tasks¶

To use these operators, you must do the following:

Install the provider package via pip.
pip install 'apache-airflow-providers-snowflake'
Detailed information is available for Installation.
Set up a Snowflake connection.

Use the snowflake_conn_id argument to specify the connection used by an operator. If not specified, snowflake_default is used.

Authentication¶

SnowflakeCortexAgentHook (and the operator built on it) authenticate according to the connection’s authenticator extra:

OAuth: set authenticator to oauth and configure a refresh token, a client credentials grant, or azure_conn_id, as on the Snowflake connection generally.
PAT (Programmatic Access Token): set authenticator to programmatic_access_token and put the PAT value in the connection’s Password field.
Key-pair JWT: the default when neither of the above is set. Configure private_key_file or private_key_content (optionally with a passphrase in Password), and set Login to the Snowflake user name the key is registered to.

See the Snowflake connection page for the full field reference. 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.

SnowflakeCortexAgentCreateOperator¶

To create a Snowflake Cortex Agent you can use SnowflakeCortexAgentCreateOperator.

tests/system/snowflake/example_snowflake_cortex_agent.py[source]

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,
)

SnowflakeCortexAgentUpdateOperator¶

To update an existing Snowflake Cortex Agent you can use SnowflakeCortexAgentUpdateOperator.

Only fields explicitly provided are updated. Optional fields left as None retain their existing values on the Cortex Agent.

tests/system/snowflake/example_snowflake_cortex_agent.py[source]

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.",
    },
)

SnowflakeCortexAgentOperator¶

Use the 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:

tests/system/snowflake/example_snowflake_cortex_agent.py[source]

run_agent = SnowflakeCortexAgentOperator(
    task_id="run_agent",
    database=DATABASE,
    schema=SCHEMA,
    agent_name=AGENT_NAME,
    messages=[
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "What can you help me with?",
                }
            ],
        }
    ],
)

SnowflakeCortexAgentDeleteOperator¶

To delete a Snowflake Cortex Agent you can use SnowflakeCortexAgentDeleteOperator.

tests/system/snowflake/example_snowflake_cortex_agent.py[source]

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