SameerMesiah97 commented on PR #70930:
URL: https://github.com/apache/airflow/pull/70930#issuecomment-6027191494
> @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:
<img width="1800" height="687" alt="image"
src="https://github.com/user-attachments/assets/1a078edb-d391-400c-80b6-774ef6be8a24";
/>
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:
<img width="1478" height="137" alt="image"
src="https://github.com/user-attachments/assets/512c2eec-6936-4c17-a88f-079c75581084";
/>
This is what I see for the `run_agent` task:
<img width="1030" height="416" alt="image"
src="https://github.com/user-attachments/assets/57266ed2-7fd0-4bcc-b384-f5ae53e42c50";
/>
And this is what I see when for `delete_agent` task:
<img width="933" height="127" alt="image"
src="https://github.com/user-attachments/assets/a50bfea1-57cf-463a-8509-ef61ee85c116";
/>
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:
<img width="967" height="572" alt="image"
src="https://github.com/user-attachments/assets/54ce9320-d966-42c4-917b-6fe7385449e4";
/>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
