snowflake.snowpark.DataFrameAIFunctions.count_tokens

DataFrameAIFunctions.count_tokens(function_name: str, prompt: Union[snowflake.snowpark.column.Column, str], *, model: Optional[str] = None, output_column: Optional[str] = None, options: Optional[Dict[str, Any]] = None, return_error_details: Optional[bool] = None) → snowflake.snowpark.DataFrame[source]

Count the number of tokens in text for a specified AI function.

This method returns an estimate of the number of tokens that would be consumed by a call to the given AI function. This is useful for estimating costs and ensuring inputs fit within model token limits.

Parameters:

  • function_name – The AI function whose tokenizer should be used for counting (e.g. 'ai_complete', 'ai_classify', 'ai_embed', 'ai_sentiment'). Must begin with 'ai_' and use only lowercase letters.

  • prompt – The column (Column object or column name as string) containing the text to count tokens for.

  • model –

    The model name, required for functions that accept a model parameter (e.g. ai_complete, ai_embed). Supported models include:

    • deepseek-r1, e5-base-v2, e5-large-v2

    • gemma-7b, jamba-1.5-large, jamba-1.5-mini, jamba-instruct

    • llama2-70b-chat, llama3-70b, llama3-8b

    • llama3.1-405b, llama3.1-70b, llama3.1-8b

    • llama3.2-1b, llama3.2-3b, llama3.3-70b

    • llama4-maverick, llama4-scout

    • mistral-7b, mistral-large, mistral-large2, mixtral-8x7b

    • nv-embed-qa-4, reka-core, reka-flash

    • snowflake-arctic-embed-l-v2.0, snowflake-arctic-embed-m-v1.5

    • snowflake-arctic-embed-m, snowflake-arctic

    • snowflake-llama-3.1-405b, snowflake-llama-3.3-70b

    • voyage-multilingual-2

  • output_column – The name of the output column to be appended. If not provided, a column named COUNT_TOKENS_OUTPUT is appended.

  • options – Optional dict specifying additional processing parameters for the underlying AI_COUNT_TOKENS call.

  • return_error_details – When True, returns an OBJECT with value and error fields instead of returning NULL on failure.

Returns:

A new DataFrame with an appended output column containing the token count as an integer. When return_error_details=True, the column contains an OBJECT with value and error fields instead.

Examples:

>>> # Count tokens for ai_complete with a specific model
>>> df = session.create_dataframe([
...     ["What is a large language model?"],
...     ["Explain quantum computing in simple terms."],
... ], schema=["text"])
>>> result_df = df.ai.count_tokens(
...     "ai_complete",
...     prompt="text",
...     model="llama3.1-70b",
...     output_column="token_count"
... )
>>> result_df.collect()[0]["TOKEN_COUNT"] > 0
True

>>> # Count tokens for ai_sentiment (no model required)
>>> result_df = df.ai.count_tokens(
...     "ai_sentiment",
...     prompt="text",
...     output_column="token_count"
... )
>>> result_df.collect()[0]["TOKEN_COUNT"] > 0
True

Note

The token count does not account for any managed system prompt that may be automatically added when using other Cortex AI functions. The actual token usage may be higher when using those functions.