Add SnowflakeCortexAgentHook#68942
Draft
SameerMesiah97 wants to merge 1 commit into
Draft
Conversation
the Snowflake REST API and extracting text responses. Includes unit tests covering successful execution, threaded conversations, optional parameters, and error handling.
ce974bc to
4ec3220
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This change introduces a new hook for interacting with Snowflake Cortex Agents (
SnowflakeCortexAgentHook).The hook provides a
run_agent()method that executes a Cortex Agent using Snowflake's REST API and returns the JSON response payload. It supports conversation threading (thread_idandparent_message_id), model configuration, agent instructions, orchestration settings, tool selection, and tool resource configuration.A helper method,
get_text_response(), is also provided to extract text content from Cortex Agent responses.Authentication is delegated to
SnowflakeHookby reusing connection parameters and access tokens resolved by the existing Snowflake provider infrastructure.Rationale
Snowflake Cortex Agents provide a programmable interface for interacting with AI agents hosted within Snowflake. While the Snowflake provider currently includes hooks and operators for executing SQL workloads, it does not provide an integration for Cortex Agents.
Introducing a dedicated hook establishes the foundation for future Cortex Agent integrations while providing users with a reusable interface for invoking agents from Airflow tasks.
This initial implementation is scoped to OAuth authentication by reusing access tokens resolved by
SnowflakeHook. Support for additional authentication mechanisms can be added in future enhancements.Tests
Added unit tests verifying that:
run_agent()executes a Cortex Agent request and returns the expected response.ValueErroris raised whenthread_idis provided withoutparent_message_id.ValueErroris raised when an access token is unavailable.get_text_response()correctly extracts text content from agent responses.Documentation
Docstrings have been added describing the hook and its public methods, including supported request parameters and response behavior.
Backwards Compatibility
This change is additive only and does not modify existing Snowflake provider functionality or public APIs.