snowflake-ml-python 2.3.0__py3-none-any.whl → 2.5.0__py3-none-any.whl

snowflake/cortex/_classify_text.py

@@ -1,10 +1,43 @@
-from typing import Union, cast
+import json
+from typing import Any, cast
 
 from typing_extensions import deprecated
 
 from snowflake import snowpark
-from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    call_sql_function,
+    warn_cortex_deprecated,
+)
 from snowflake.ml._internal import telemetry
+from snowflake.snowpark import functions
+
+_AI_CLASSIFY_FUNCTION_NAME = "AI_CLASSIFY"
+# Legacy CLASSIFY_TEXT always returned exactly one label. AI_CLASSIFY defaults to single-label mode, but
+# we set output_mode explicitly so the single-label contract can't drift and multi-label results can
+# never leak through.
+_AI_CLASSIFY_SINGLE_LABEL_CONFIG = {"output_mode": "single"}
+
+
+def _adapt_ai_classify_result(result: Any) -> str:
+    """Map an AI_CLASSIFY result to the legacy CLASSIFY_TEXT {"label": ...} JSON string."""
+    # Legacy classify_text returned the raw CLASSIFY_TEXT OBJECT, which the immediate (non-Column) path
+    # collects as a JSON string such as '{"label": "positive"}'. We reproduce that exact shape so
+    # existing callers that json.loads the result and read ["label"] keep working.
+    parsed: Any = result
+    if isinstance(result, str):
+        try:
+            parsed = json.loads(result)
+        except json.JSONDecodeError:
+            return result
+    label: Any = None
+    if isinstance(parsed, dict):
+        labels = parsed.get("labels")
+        if isinstance(labels, list) and labels:
+            label = labels[0]
+        elif parsed.get("label") is not None:
+            label = parsed.get("label")
+    return json.dumps({"label": label})
 
 
 @telemetry.send_api_usage_telemetry(
@@ -25,17 +58,30 @@ def classify_text(
     Returns:
         A column of classification responses.
     """
-
-    return _classify_text_impl("snowflake.cortex.classify_text", str_input, categories, session=session)
+    warn_cortex_deprecated("classify_text", "ai_classify")
+    result = _classify_text_impl(
+        _AI_CLASSIFY_FUNCTION_NAME, str_input, categories, _AI_CLASSIFY_SINGLE_LABEL_CONFIG, session=session
+    )
+    if isinstance(result, snowpark.Column):
+        # AI_CLASSIFY returns {"labels": [...]}; re-wrap as the legacy CLASSIFY_TEXT {"label": ...} OBJECT
+        # so downstream SQL that reads col["label"] is unchanged. Subscript access (not functions.get with
+        # a bare string, which Snowpark treats as a column identifier) extracts the field by literal key.
+        label = functions.coalesce(result["labels"][0], result["label"])
+        return functions.object_construct(functions.lit("label"), label)
+    return _adapt_ai_classify_result(result)
 
 
 def _classify_text_impl(
     function: str,
     str_input: str | snowpark.Column,
     categories: list[str] | snowpark.Column,
+    config: dict[str, Any] | None = None,
     session: snowpark.Session | None = None,
 ) -> str | snowpark.Column:
-    return cast(Union[str, snowpark.Column], call_sql_function(function, session, str_input, categories))
+    args: list[Any] = [str_input, categories]
+    if config is not None:
+        args.append(config)
+    return cast(str | snowpark.Column, call_sql_function(function, session, *args))
 
 
 ClassifyText = deprecated(

snowflake/cortex/_complete.py

@@ -14,6 +14,7 @@ from snowflake.cortex._util import (
     CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
     SnowflakeAuthenticationException,
     SnowflakeConfigurationException,
+    warn_cortex_deprecated,
 )
 from snowflake.ml._internal import telemetry
 from snowflake.snowpark import context, functions
@@ -331,12 +332,46 @@ def _return_stream_response(
                 pass
 
 
-def _complete_call_sql_function_snowpark(
-    function: str, *args: str | snowpark.Column | CompleteOptions
-) -> snowpark.Column:
+def _complete_call_sql_function_snowpark(function: str, *args: Any) -> snowpark.Column:
     return cast(snowpark.Column, functions.builtin(function)(*args))
 
 
+_SQL_MODEL_PARAMETER_KEYS = ("max_tokens", "temperature", "top_p", "guardrails")
+
+
+def _sql_ai_complete_args(
+    model: str | snowpark.Column,
+    prompt: str | list[ConversationMessage] | snowpark.Column,
+    options: CompleteOptions | snowpark.Column | None,
+) -> tuple[Any, ...]:
+    """Build positional SQL arguments for AI_COMPLETE.
+
+    Args:
+        model: Model name or Column of model names.
+        prompt: Prompt string, conversation history, or Column of prompts.
+        options: Complete options, a Column of options, or None.
+
+    Returns:
+        Positional AI_COMPLETE arguments with response format separate from
+        model parameters.
+    """
+    if options is None:
+        return (model, prompt)
+    if isinstance(options, snowpark.Column):
+        # Can't introspect a Column's keys at build time, so split response_format out of the options
+        # object in SQL: model_parameters is the object minus response_format, response_format is read
+        # out separately (NULL when absent, which AI_COMPLETE treats as "no response format").
+        model_parameters_col = functions.builtin("object_delete")(options, functions.lit("response_format"))
+        response_format_col = functions.get(options, functions.lit("response_format"))
+        return (model, prompt, model_parameters_col, response_format_col)
+    option_values = cast(dict[str, Any], options)
+    model_parameters = {key: option_values[key] for key in _SQL_MODEL_PARAMETER_KEYS if key in option_values}
+    response_format = options.get("response_format")
+    if response_format is None:
+        return (model, prompt, model_parameters) if model_parameters else (model, prompt)
+    return (model, prompt, model_parameters, response_format)
+
+
 def _complete_non_streaming_immediate(
     snow_api_xp_request_handler: Callable[..., dict[str, Any]] | None,
     model: str,
@@ -366,10 +401,7 @@ def _complete_non_streaming_impl(
     deadline: float | None = None,
 ) -> str | snowpark.Column:
     if isinstance(prompt, snowpark.Column):
-        if options is not None:
-            return _complete_call_sql_function_snowpark(function, model, prompt, options)
-        else:
-            return _complete_call_sql_function_snowpark(function, model, prompt)
+        return _complete_call_sql_function_snowpark(function, *_sql_ai_complete_args(model, prompt, options))
     if isinstance(model, snowpark.Column):
         raise ValueError("'model' cannot be a snowpark.Column when 'prompt' is a string.")
     if isinstance(options, snowpark.Column):
@@ -414,7 +446,7 @@ def _complete_impl(
     model: str | snowpark.Column,
     prompt: str | list[ConversationMessage] | snowpark.Column,
     snow_api_xp_request_handler: Callable[..., dict[str, Any]] | None = None,
-    function: str = "snowflake.cortex.complete",
+    function: str = "AI_COMPLETE",
     options: CompleteOptions | None = None,
     session: snowpark.Session | None = None,
     stream: bool = False,
@@ -480,7 +512,7 @@ def complete(
     Returns:
         A column of string responses.
     """
-
+    warn_cortex_deprecated("complete", "ai_complete")
     # Set the XP snow api function, if available.
     snow_api_xp_request_handler = None
     if is_in_stored_procedure():  # type: ignore[no-untyped-call]

snowflake/cortex/_embed_text_1024.py

@@ -1,11 +1,17 @@
-from typing import Union, cast
+from typing import cast
 
 from typing_extensions import deprecated
 
 from snowflake import snowpark
-from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    call_sql_function,
+    warn_cortex_deprecated,
+)
 from snowflake.ml._internal import telemetry
 
+_AI_EMBED_FUNCTION_NAME = "AI_EMBED"
+
 
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
@@ -26,8 +32,8 @@ def embed_text_1024(
     Returns:
         A column of vectors containing embeddings.
     """
-
-    return _embed_text_1024_impl("snowflake.cortex.embed_text_1024", model, text, session=session)
+    warn_cortex_deprecated("embed_text_1024", "ai_embed")
+    return _embed_text_1024_impl(_AI_EMBED_FUNCTION_NAME, model, text, session=session)
 
 
 def _embed_text_1024_impl(
@@ -36,7 +42,7 @@ def _embed_text_1024_impl(
     text: str | snowpark.Column,
     session: snowpark.Session | None = None,
 ) -> list[float] | snowpark.Column:
-    return cast(Union[list[float], snowpark.Column], call_sql_function(function, session, model, text))
+    return cast(list[float] | snowpark.Column, call_sql_function(function, session, model, text))
 
 
 EmbedText1024 = deprecated(

snowflake/cortex/_embed_text_768.py

@@ -1,11 +1,17 @@
-from typing import Union, cast
+from typing import cast
 
 from typing_extensions import deprecated
 
 from snowflake import snowpark
-from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    call_sql_function,
+    warn_cortex_deprecated,
+)
 from snowflake.ml._internal import telemetry
 
+_AI_EMBED_FUNCTION_NAME = "AI_EMBED"
+
 
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
@@ -26,8 +32,8 @@ def embed_text_768(
     Returns:
         A column of vectors containing embeddings.
     """
-
-    return _embed_text_768_impl("snowflake.cortex.embed_text_768", model, text, session=session)
+    warn_cortex_deprecated("embed_text_768", "ai_embed")
+    return _embed_text_768_impl(_AI_EMBED_FUNCTION_NAME, model, text, session=session)
 
 
 def _embed_text_768_impl(
@@ -36,7 +42,7 @@ def _embed_text_768_impl(
     text: str | snowpark.Column,
     session: snowpark.Session | None = None,
 ) -> list[float] | snowpark.Column:
-    return cast(Union[list[float], snowpark.Column], call_sql_function(function, session, model, text))
+    return cast(list[float] | snowpark.Column, call_sql_function(function, session, model, text))
 
 
 EmbedText768 = deprecated(

snowflake/cortex/_extract_answer.py

@@ -1,10 +1,39 @@
-from typing import Union, cast
+from typing import Any, cast
 
 from typing_extensions import deprecated
 
 from snowflake import snowpark
-from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    SnowflakeAuthenticationException,
+    warn_cortex_deprecated,
+)
 from snowflake.ml._internal import telemetry
+from snowflake.snowpark import context, functions
+
+_AI_EXTRACT_FUNCTION_NAME = "AI_EXTRACT"
+_ANSWER_KEY = "answer"
+
+
+def _legacy_extract_answer_column(
+    from_text: str | snowpark.Column,
+    response_format: dict[str, Any] | snowpark.Column,
+) -> snowpark.Column:
+    """Call AI_EXTRACT and reshape it to the legacy EXTRACT_ANSWER [{"answer": ..., "score": ...}] ARRAY."""
+    # AI_EXTRACT returns the score only when the `scores` argument is TRUE, and `scores` is a named-only
+    # argument that the Snowpark builtin wrapper can't pass positionally -- inject it as a raw SQL fragment.
+    # AI_EXTRACT(<text>, <response_format>, scores => TRUE) returns:
+    #   {"response": {"answer": ...}, "scoring": {"scores": {"answer": {"score": ...}}}}
+    extracted = functions.builtin(_AI_EXTRACT_FUNCTION_NAME)(
+        from_text, response_format, functions.sql_expr("scores => TRUE")
+    )
+    # Subscript access (not functions.get with a bare string, which Snowpark treats as a column
+    # identifier) extracts fields by literal key.
+    answer = extracted["response"][_ANSWER_KEY]
+    score = extracted["scoring"]["scores"][_ANSWER_KEY]["score"]
+    return functions.array_construct(
+        functions.object_construct(functions.lit(_ANSWER_KEY), answer, functions.lit("score"), score)
+    )
 
 
 @telemetry.send_api_usage_telemetry(
@@ -23,19 +52,32 @@ def extract_answer(
         session: The snowpark session to use. Will be inferred by context if not specified.
 
     Returns:
-        A column of strings containing answers.
-    """
+        A column of arrays, each [{"answer": ..., "score": ...}], matching legacy EXTRACT_ANSWER.
 
-    return _extract_answer_impl("snowflake.cortex.extract_answer", from_text, question, session=session)
+    Raises:
+        SnowflakeAuthenticationException: if no session is provided and none is active.
+    """
+    warn_cortex_deprecated("extract_answer", "ai_extract")
+    if isinstance(question, snowpark.Column):
+        response_format: dict[str, Any] | snowpark.Column = functions.object_construct(
+            functions.lit(_ANSWER_KEY), question
+        )
+    else:
+        response_format = {_ANSWER_KEY: question}
 
+    column = _legacy_extract_answer_column(from_text, response_format)
+    if isinstance(from_text, snowpark.Column) or isinstance(question, snowpark.Column):
+        return column
 
-def _extract_answer_impl(
-    function: str,
-    from_text: str | snowpark.Column,
-    question: str | snowpark.Column,
-    session: snowpark.Session | None = None,
-) -> str | snowpark.Column:
-    return cast(Union[str, snowpark.Column], call_sql_function(function, session, from_text, question))
+    # Immediate (literal) path: evaluate the expression and return the collected ARRAY as a JSON string,
+    # exactly as the legacy EXTRACT_ANSWER immediate path did.
+    session = session or context.get_active_session()
+    if session is None:
+        raise SnowflakeAuthenticationException(
+            """Session required. Provide the session through a session=... argument or ensure an active session is
+            available in your environment."""
+        )
+    return cast(str, session.create_dataframe([snowpark.Row()]).select(column).collect()[0][0])
 
 
 ExtractAnswer = deprecated(

snowflake/cortex/_sentiment.py

@@ -3,8 +3,43 @@ from typing import cast
 from typing_extensions import deprecated
 
 from snowflake import snowpark
-from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    SnowflakeAuthenticationException,
+    warn_cortex_deprecated,
+)
 from snowflake.ml._internal import telemetry
+from snowflake.snowpark import context, functions, types
+
+_AI_COMPLETE_FUNCTION_NAME = "AI_COMPLETE"
+
+# AI_SENTIMENT returns a categorical label (positive/negative/neutral/mixed), not the legacy FLOAT in
+# [-1, 1] that SNOWFLAKE.CORTEX.SENTIMENT produced. To keep the numeric return contract we score the
+# text with AI_COMPLETE and a constrained prompt instead. The score is model-derived, so exact values
+# differ from the legacy sentiment model.
+_SENTIMENT_MODEL = "llama3.1-8b"
+_SENTIMENT_PROMPT_PREFIX = (
+    "Rate the sentiment of the following text from -1.0 (very negative) to 1.0 (very positive). "
+    "Respond with only the score. Text: "
+)
+_SENTIMENT_MODEL_PARAMETERS = {"temperature": 0}
+_SENTIMENT_RESPONSE_FORMAT = {
+    "type": "json",
+    "schema": {"type": "object", "properties": {"score": {"type": "number"}}, "required": ["score"]},
+}
+
+
+def _sentiment_score_column(text: snowpark.Column) -> snowpark.Column:
+    """Build the AI_COMPLETE sentiment-score expression, clamped to the legacy [-1, 1] range."""
+    prompt = functions.concat(functions.lit(_SENTIMENT_PROMPT_PREFIX), text)
+    completion = functions.builtin(_AI_COMPLETE_FUNCTION_NAME)(
+        functions.lit(_SENTIMENT_MODEL),
+        prompt,
+        _SENTIMENT_MODEL_PARAMETERS,
+        _SENTIMENT_RESPONSE_FORMAT,
+    )
+    score = functions.get(functions.parse_json(completion), functions.lit("score")).cast(types.FloatType())
+    return functions.least(functions.lit(1.0), functions.greatest(functions.lit(-1.0), score))
 
 
 @telemetry.send_api_usage_telemetry(
@@ -19,17 +54,22 @@ def sentiment(text: str | snowpark.Column, session: snowpark.Session | None = No
 
     Returns:
         A column of floats. 1 represents positive sentiment, -1 represents negative sentiment.
-    """
-    return _sentiment_impl("snowflake.cortex.sentiment", text, session=session)
 
+    Raises:
+        SnowflakeAuthenticationException: if no session is provided and none is active.
+    """
+    warn_cortex_deprecated("sentiment", "ai_complete")
+    if isinstance(text, snowpark.Column):
+        return _sentiment_score_column(text)
 
-def _sentiment_impl(
-    function: str, text: str | snowpark.Column, session: snowpark.Session | None = None
-) -> float | snowpark.Column:
-
-    output = call_sql_function(function, session, text)
-    if isinstance(output, snowpark.Column):
-        return output
+    session = session or context.get_active_session()
+    if session is None:
+        raise SnowflakeAuthenticationException(
+            """Session required. Provide the session through a session=... argument or ensure an active session is
+            available in your environment."""
+        )
+    score_column = _sentiment_score_column(functions.lit(text))
+    output = session.create_dataframe([snowpark.Row()]).select(score_column).collect()[0][0]
     return float(cast(str, output))
 
 

snowflake/cortex/_summarize.py

@@ -1,11 +1,17 @@
-from typing import Union, cast
+from typing import cast
 
 from typing_extensions import deprecated
 
 from snowflake import snowpark
-from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    call_sql_function,
+    warn_cortex_deprecated,
+)
 from snowflake.ml._internal import telemetry
 
+_AI_SUMMARIZE_FUNCTION_NAME = "AI_SUMMARIZE"
+
 
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
@@ -23,8 +29,8 @@ def summarize(
     Returns:
         A column of string summaries.
     """
-
-    return _summarize_impl("snowflake.cortex.summarize", text, session=session)
+    warn_cortex_deprecated("summarize", "ai_complete")
+    return _summarize_impl(_AI_SUMMARIZE_FUNCTION_NAME, text, session=session)
 
 
 def _summarize_impl(
@@ -32,7 +38,7 @@ def _summarize_impl(
     text: str | snowpark.Column,
     session: snowpark.Session | None = None,
 ) -> str | snowpark.Column:
-    return cast(Union[str, snowpark.Column], call_sql_function(function, session, text))
+    return cast(str | snowpark.Column, call_sql_function(function, session, text))
 
 
 Summarize = deprecated("Summarize() is deprecated and will be removed in a future release. Use summarize() instead")(

snowflake/cortex/_translate.py

@@ -1,11 +1,17 @@
-from typing import Union, cast
+from typing import cast
 
 from typing_extensions import deprecated
 
 from snowflake import snowpark
-from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    call_sql_function,
+    warn_cortex_deprecated,
+)
 from snowflake.ml._internal import telemetry
 
+_AI_TRANSLATE_FUNCTION_NAME = "AI_TRANSLATE"
+
 
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
@@ -27,8 +33,8 @@ def translate(
     Returns:
         A column of string translations.
     """
-
-    return _translate_impl("snowflake.cortex.translate", text, from_language, to_language, session=session)
+    warn_cortex_deprecated("translate", "ai_translate")
+    return _translate_impl(_AI_TRANSLATE_FUNCTION_NAME, text, from_language, to_language, session=session)
 
 
 def _translate_impl(
@@ -38,7 +44,7 @@ def _translate_impl(
     to_language: str | snowpark.Column,
     session: snowpark.Session | None = None,
 ) -> str | snowpark.Column:
-    return cast(Union[str, snowpark.Column], call_sql_function(function, session, text, from_language, to_language))
+    return cast(str | snowpark.Column, call_sql_function(function, session, text, from_language, to_language))
 
 
 Translate = deprecated("Translate() is deprecated and will be removed in a future release. Use translate() instead")(

snowflake/cortex/_util.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, Union, cast
 
 from snowflake import snowpark
@@ -8,6 +9,40 @@ from snowflake.snowpark import context, functions
 CORTEX_FUNCTIONS_TELEMETRY_PROJECT = "CortexFunctions"
 
 
+_SNOWPARK_FUNCTIONS_DOC_URL = (
+    "https://docs.snowflake.com/en/developer-guide/snowpark/reference/python/latest/"
+    "snowpark/api/snowflake.snowpark.functions.{replacement}"
+)
+_MIGRATION_GUIDE_URL = "https://docs.snowflake.com/en/user-guide/snowflake-cortex/aisql-migrate-legacy-functions"
+
+# Functions that have already emitted their deprecation warning this process. Guarding on this keeps a
+# repeated call (e.g. inside a loop) from spamming the console -- each deprecated function warns once.
+_emitted_deprecations: set[str] = set()
+
+
+def warn_cortex_deprecated(name: str, replacement: str) -> None:
+    """Emit the deprecation warning for the named deprecated snowflake.cortex function, once per process.
+
+    Args:
+        name: The deprecated snowflake.cortex function name (e.g. "classify_text").
+        replacement: The snowflake.snowpark.functions replacement to point callers at (e.g. "ai_classify").
+    """
+    # Warn on use (not on package import) so callers of functions that are NOT deprecated -- e.g.
+    # Finetune -- are never warned.
+    if name in _emitted_deprecations:
+        return
+    _emitted_deprecations.add(name)
+    warnings.warn(
+        f"snowflake.cortex.{name} is deprecated and will be removed in a future release. "
+        f"Its backend is now redirected to Snowflake AI Functions. "
+        f"Use snowflake.snowpark.functions.{replacement} instead. "
+        f"API: {_SNOWPARK_FUNCTIONS_DOC_URL.format(replacement=replacement)} "
+        f"Migration guide: {_MIGRATION_GUIDE_URL}",
+        DeprecationWarning,
+        stacklevel=3,
+    )
+
+
 class SnowflakeAuthenticationException(Exception):
     """This exception is raised when there is an issue with Snowflake's configuration."""
 

snowflake/ml/_internal/env_utils.py

@@ -90,6 +90,9 @@ def _validate_pip_requirement_string(req_str: str) -> requirements.Requirement:
     Returns:
         A requirements.Requirement object containing the requirement information.
     """
+    # A requirement string must not carry surrounding whitespace: packaging>=26 accepts only a
+    # requirement that ends at the end of the string, so a trailing newline is a parse error.
+    req_str = req_str.strip()
     try:
         r = requirements.Requirement(req_str)
 
@@ -780,8 +783,9 @@ def load_requirements_file(path: pathlib.Path) -> list[requirements.Requirement]
     with open(path, encoding="utf-8") as f:
         reqs = f.readlines()
 
-    # Filter out pip options
-    reqs = [line for line in reqs if line.strip() and not line.strip().startswith(("-", "#"))]
+    # "-" lines are pip options (e.g. --extra-index-url) and "#" lines are comments; neither is a
+    # requirement. Strip as well, since a requirement string must not keep readlines()' newline.
+    reqs = [stripped for line in reqs if (stripped := line.strip()) and not stripped.startswith(("-", "#"))]
 
     return validate_pip_requirement_string_list(reqs)
 

snowflake/ml/_internal/file_utils.py

@@ -120,7 +120,21 @@ def make_archive(
     )
 
 
-def zip_python_package(zipfile_path: str, package_name: str, ignore_generated_py_file: bool = True) -> None:
+def zip_python_package(
+    zipfile_path: str,
+    package_name: str,
+    ignore_generated_py_file: bool = True,
+    *,
+    exclude_dir_names: tuple[str, ...] = (),
+) -> None:
+    """Zip an imported Python package.
+
+    Args:
+        zipfile_path: Destination zip path.
+        package_name: Dotted package name to zip.
+        ignore_generated_py_file: When True, omit generated bytecode and stub files.
+        exclude_dir_names: Directory names to omit anywhere in the package. The directory entry itself is not written.
+    """
     import importlib_resources
     from importlib_resources import abc as importlib_resources_abc
 
@@ -131,6 +145,7 @@ def zip_python_package(zipfile_path: str, package_name: str, ignore_generated_py
 
     base_dirs = package_name.split(".")
     assert len(base_dirs) >= 1, "Invalid package name."
+    excluded_dirs = set(exclude_dir_names)
 
     with zipfile.ZipFile(zipfile_path, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
         for i in range(len(base_dirs)):
@@ -141,6 +156,8 @@ def zip_python_package(zipfile_path: str, package_name: str, ignore_generated_py
         def _add_to_zip(
             zf: zipfile.ZipFile, path_info: importlib_resources_abc.Traversable, base_arcname: pathlib.PurePosixPath
         ) -> None:
+            if path_info.name in excluded_dirs:
+                return
             if path_info.is_file():
                 if ignore_generated_py_file and path_info.name.endswith(GENERATED_PY_FILE_EXT):
                     return