From 8387b1f1b46b9ae0a79783c78794762049407cad Mon Sep 17 00:00:00 2001 From: Sameer Mesiah Date: Fri, 10 Jul 2026 21:55:50 +0100 Subject: [PATCH] Set stream=False for Snowflake Cortex Agent requests. --- .../providers/snowflake/hooks/snowflake_cortex_agent.py | 5 ++++- .../unit/snowflake/hooks/test_snowflake_cortex_agent.py | 3 +++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py index 7b7f37a249787..1413d7fee7135 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py +++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py @@ -130,7 +130,10 @@ def run_agent( if thread_id is not None and parent_message_id is None: raise ValueError("parent_message_id must be provided when thread_id is specified.") - payload: dict[str, Any] = {"messages": messages} + payload: dict[str, Any] = { + "messages": messages, + "stream": False, + } if thread_id is not None: payload["thread_id"] = thread_id diff --git a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_cortex_agent.py b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_cortex_agent.py index 8ec86d532f966..e686b698574ad 100644 --- a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_cortex_agent.py +++ b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_cortex_agent.py @@ -125,6 +125,7 @@ def test_run_agent( ], } ], + "stream": False, }, timeout=REQUEST_TIMEOUT, ) @@ -175,6 +176,7 @@ def test_run_agent_includes_thread_fields( assert payload["thread_id"] == 123 assert payload["parent_message_id"] == 456 + assert payload["stream"] is False @mock.patch(f"{MODULE_PATH}.requests.request") @mock.patch(f"{HOOK_PATH}._get_conn_params") @@ -215,6 +217,7 @@ def test_run_agent_includes_optional_fields( assert payload["orchestration"] == {"max_tokens": 1000} assert payload["tools"] == [{"name": "search_tool"}] assert payload["tool_resources"] == {"search_tool": {"config": "value"}} + assert payload["stream"] is False @mock.patch(f"{MODULE_PATH}.requests.request") @mock.patch(f"{HOOK_PATH}._get_conn_params")