Add missing AISQL function wrappers: ai_count_tokens, ai_multi_embed, ai_redact#4261
Open
sfc-gh-yixie wants to merge 17 commits into
Open
Add missing AISQL function wrappers: ai_count_tokens, ai_multi_embed, ai_redact#4261sfc-gh-yixie wants to merge 17 commits into
sfc-gh-yixie wants to merge 17 commits into
Conversation
…ti_embed, ai_redact) - Add standalone `ai_count_tokens`, `ai_multi_embed`, and `ai_redact` to `functions.py` - Add `DataFrame.ai.multi_embed` and `DataFrame.ai.redact` DataFrame methods - Fix `DataFrame.ai.count_tokens` to use new `AI_COUNT_TOKENS` instead of deprecated `SNOWFLAKE.CORTEX.COUNT_TOKENS` - Add `DataframeAiMultiEmbed` and `DataframeAiRedact` proto messages; regenerate `ast_pb2.py` - Add integration tests for all new functions Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Comment on lines
+7
to
+18
| #### New Features | ||
|
|
||
| - Added `ai_count_tokens` function to `snowflake.snowpark.functions` to estimate token counts for AI function calls. This is the updated replacement for the deprecated `SNOWFLAKE.CORTEX.COUNT_TOKENS`. | ||
| - Added `ai_multi_embed` function to `snowflake.snowpark.functions` to generate multimodal embeddings from text, images, audio, or video files. | ||
| - Added `ai_redact` function to `snowflake.snowpark.functions` to detect and redact personally identifiable information (PII) from text. | ||
| - Added `DataFrame.ai.multi_embed` method to generate multimodal embeddings via the DataFrame API. | ||
| - Added `DataFrame.ai.redact` method to detect and redact PII from text columns via the DataFrame API. | ||
|
|
||
| #### Bug Fixes | ||
|
|
||
| - Fixed `DataFrame.ai.count_tokens` to use `AI_COUNT_TOKENS` instead of the deprecated `SNOWFLAKE.CORTEX.COUNT_TOKENS`. | ||
|
|
Collaborator
There was a problem hiding this comment.
nit: this shall be moved to 1.54.0 after rebasing the main
Collaborator
Author
There was a problem hiding this comment.
Moved to the latest version
| df._ast_id = stmt.uid | ||
| return df | ||
|
|
||
| @experimental(version="1.52.0") |
Collaborator
There was a problem hiding this comment.
do we need experimental flag -- are these functions in pupr?
Collaborator
|
there are some lint error https://github.com/snowflakedb/snowpark-python/actions/runs/27932786054/job/82647949888?pr=4261 |
…ature/aisql-function-wrappers
Collaborator
Author
fixed lint |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #4261 +/- ##
===========================================
- Coverage 95.21% 84.59% -10.62%
===========================================
Files 171 170 -1
Lines 44538 44623 +85
Branches 7617 7674 +57
===========================================
- Hits 42406 37748 -4658
- Misses 1320 5646 +4326
- Partials 812 1229 +417 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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
3 participants
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.
Which Jira issue is this PR addressing? Make sure that there is an accompanying issue to your PR.
Fixes NO-SNOW
Fill out the following pre-review checklist:
Please describe how your code solves the related issue.
Summary
snowflake.snowpark.functions:ai_count_tokens,ai_multi_embed, andai_redactDataFrame.ai.multi_embedandDataFrame.ai.redacthigh-level DataFrame methodsDataFrame.ai.count_tokensto call the newAI_COUNT_TOKENSSQL function instead of the deprecatedSNOWFLAKE.CORTEX.COUNT_TOKENSDataframeAiMultiEmbedandDataframeAiRedactproto messages toast.protoand regeneratedast_pb2.pyMotivation
Snowflake now exposes 15
AI_*functions. Prior to this PR, the Snowpark Python client was missing three of them:AI_COUNT_TOKENSAI_MULTI_EMBEDAI_REDACTTest plan
test_ai_count_tokens_basic— verifiesai_count_tokensreturns a positive integertest_ai_count_tokens_no_model— verifies functions that don't require a modeltest_ai_count_tokens_column_input— verifies column references and relative token countstest_ai_redact_basic— verifies PII placeholders appear in outputtest_ai_redact_with_categories— verifies category-scoped redactiontest_ai_redact_detect_mode— verifies span metadata returned in detect modetest_ai_redact_column_input— verifies column referencestest_dataframe_ai_redact_basic/_detect_mode/_with_categories/_default_output_columntest_ai_multi_embed_basic— verifies embedding object structuretest_dataframe_ai_multi_embed_basic/_default_output_column🤖 Generated with Claude Code