snowflake-ml-python 2.3.0__py3-none-any.whl → 2.4.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/experiment/_entities/run.py

@@ -57,11 +57,19 @@ class Run:
             run_name=self.name.identifier(),
         )
 
-    def _warn_about_run_metadata_size(self, sql_error_msg: str) -> None:
+    def _handle_run_metadata_size_limit(self, sql_error_msg: str) -> None:
+        """Report that run metadata exceeded the size limit.
+
+        Args:
+            sql_error_msg: The error message returned by Snowflake.
+
+        Raises:
+            RuntimeWarning: If the experiment tracking instance was created with warnings_as_errors=True.
+        """
+        message = f"{sql_error_msg}. {METADATA_SIZE_WARNING_MESSAGE}"
+        if self._experiment_tracking._warnings_as_errors:
+            raise RuntimeWarning(message)
+        # Once the limit is hit every subsequent log attempt fails the same way, so warn only once.
         if not self._warned_about_metadata_size:
-            warnings.warn(
-                f"{sql_error_msg}. {METADATA_SIZE_WARNING_MESSAGE}",
-                RuntimeWarning,
-                stacklevel=2,
-            )
+            warnings.warn(message, RuntimeWarning, stacklevel=2)
             self._warned_about_metadata_size = True

snowflake/ml/experiment/callback/catboost.py

@@ -0,0 +1,75 @@
+from typing import TYPE_CHECKING, Any
+
+from snowflake.ml.experiment import utils
+
+if TYPE_CHECKING:
+    from snowflake.ml.experiment.experiment_tracking import ExperimentTracking
+
+
+class SnowflakeCatboostCallback:
+    """CatBoost callback that logs training metrics and parameters to a Snowflake ML Experiment.
+
+    CatBoost's callback protocol only exposes iteration number and metrics, not the model itself.
+    This callback logs metrics during training and, optionally, model parameters. To log the trained
+    model, call ``ExperimentTracking.log_model`` after ``fit`` returns::
+
+        callback = SnowflakeCatboostCallback(
+            exp,
+            params=model.get_params(),
+        )
+        model.fit(X, y, callbacks=[callback])
+        exp.log_model(model, model_name="my_model", signatures={"predict": sig})
+    """
+
+    def __init__(
+        self,
+        experiment_tracking: "ExperimentTracking",
+        *,
+        params: dict[str, Any] | None = None,
+        log_every_n_epochs: int = 1,
+    ) -> None:
+        """
+        Initialize the callback.
+
+        Args:
+            experiment_tracking: The Experiment Tracking instance to use for logging.
+            params: Model parameters to log at the start of training. Pass
+                ``model.get_params()`` before calling ``fit``. If None, no parameters are logged.
+            log_every_n_epochs: Frequency with which to log metrics. Must be a positive integer.
+                Default is 1, logging after every iteration. A metric is only logged for iterations
+                where CatBoost recomputed it, so a model trained with ``metric_period`` greater than 1
+                logs less often than this setting alone implies.
+
+        Raises:
+            ValueError: When ``log_every_n_epochs`` is not a positive integer.
+        """
+        self._experiment_tracking = experiment_tracking
+        self._params = params
+        if not (utils.is_integer(log_every_n_epochs) and log_every_n_epochs > 0):
+            raise ValueError("`log_every_n_epochs` must be a positive integer.")
+        self.log_every_n_epochs = log_every_n_epochs
+        self._params_logged = False
+        self._metric_value_counts: dict[str, int] = {}
+
+    def after_iteration(self, info: Any) -> bool:
+        # CatBoost numbers iterations from 1; normalize to 0-based to match the other callbacks.
+        epoch = info.iteration - 1
+
+        if self._params is not None and not self._params_logged:
+            self._experiment_tracking.log_params(utils.flatten_nested_params(self._params))
+            self._params_logged = True
+
+        should_log = epoch % self.log_every_n_epochs == 0
+        for dataset_name, metrics in info.metrics.items():
+            for metric_name, values in metrics.items():
+                metric_key = dataset_name + ":" + metric_name
+                # CatBoost runs callbacks on every iteration but only appends a metric value on the
+                # iterations where it recomputes one, which `metric_period` can make less frequent. A
+                # list that did not grow means the trailing value belongs to an earlier iteration, so
+                # logging it again would attribute a stale value to this step.
+                recomputed = len(values) > self._metric_value_counts.get(metric_key, 0)
+                self._metric_value_counts[metric_key] = len(values)
+                if should_log and recomputed:
+                    self._experiment_tracking.log_metric(key=metric_key, value=values[-1], step=epoch)
+
+        return True