snowflake-ml-python 1.5.2__py3-none-any.whl → 1.5.4__py3-none-any.whl

snowflake/cortex/__init__.py

@@ -1,4 +1,4 @@
-from snowflake.cortex._complete import Complete
+from snowflake.cortex._complete import Complete, CompleteOptions
 from snowflake.cortex._extract_answer import ExtractAnswer
 from snowflake.cortex._sentiment import Sentiment
 from snowflake.cortex._summarize import Summarize
@@ -6,6 +6,7 @@ from snowflake.cortex._translate import Translate
 
 __all__ = [
     "Complete",
+    "CompleteOptions",
     "ExtractAnswer",
     "Sentiment",
     "Summarize",

snowflake/cortex/_complete.py

@@ -1,35 +1,259 @@
-from typing import Optional, Union
+import json
+import logging
+from typing import Iterator, List, Optional, TypedDict, Union, cast
+from urllib.parse import urlunparse
+
+import requests
+from typing_extensions import NotRequired
 
 from snowflake import snowpark
-from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.cortex._sse_client import SSEClient
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    SnowflakeAuthenticationException,
+    SnowflakeConfigurationException,
+)
 from snowflake.ml._internal import telemetry
+from snowflake.snowpark import context, functions
+
+logger = logging.getLogger(__name__)
+
+
+class ConversationMessage(TypedDict):
+    """Represents an conversation interaction."""
+
+    role: str
+    """The role of the participant. For example, "user" or "assistant"."""
+
+    content: str
+    """The content of the message."""
+
+
+class CompleteOptions(TypedDict):
+    """Options configuring a snowflake.cortex.Complete call."""
+
+    max_tokens: NotRequired[int]
+    """ Sets the maximum number of output tokens in the response. Small values can result in
+    truncated responses. """
+    temperature: NotRequired[float]
+    """ A value from 0 to 1 (inclusive) that controls the randomness of the output of the language
+    model. A higher temperature (for example, 0.7) results in more diverse and random output, while a lower
+    temperature (such as 0.2) makes the output more deterministic and focused. """
+
+    top_p: NotRequired[float]
+    """ A value from 0 to 1 (inclusive) that controls the randomness and diversity of the language model,
+    generally used as an alternative to temperature. The difference is that top_p restricts the set of possible tokens
+    that the model outputs, while temperature influences which tokens are chosen at each step. """
+
+
+class ResponseParseException(Exception):
+    """This exception is raised when the server response cannot be parsed."""
+
+    pass
+
+
+def _call_complete_rest(
+    model: str,
+    prompt: Union[str, List[ConversationMessage]],
+    options: Optional[CompleteOptions] = None,
+    session: Optional[snowpark.Session] = None,
+    stream: bool = False,
+) -> requests.Response:
+    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."""
+        )
+
+    if session.connection.host is None or session.connection.host == "":
+        raise SnowflakeConfigurationException("Snowflake connection configuration does not specify 'host'")
+
+    if session.connection.rest is None or not hasattr(session.connection.rest, "token"):
+        raise SnowflakeAuthenticationException("Snowflake session error: REST token missing.")
+
+    if session.connection.rest.token is None or session.connection.rest.token == "":
+        raise SnowflakeAuthenticationException("Snowflake session error: REST token is empty.")
+
+    scheme = "https"
+    if hasattr(session.connection, "scheme"):
+        scheme = session.connection.scheme
+    url = urlunparse((scheme, session.connection.host, "api/v2/cortex/inference/complete", "", "", ""))
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f'Snowflake Token="{session.connection.rest.token}"',
+        "Accept": "application/json, text/event-stream",
+    }
+
+    data = {
+        "model": model,
+        "stream": stream,
+    }
+    if isinstance(prompt, List):
+        data["messages"] = prompt
+    else:
+        data["messages"] = [{"content": prompt}]
+
+    if options:
+        if "max_tokens" in options:
+            data["max_tokens"] = options["max_tokens"]
+            data["max_output_tokens"] = options["max_tokens"]
+        if "temperature" in options:
+            data["temperature"] = options["temperature"]
+        if "top_p" in options:
+            data["top_p"] = options["top_p"]
+
+    logger.debug(f"making POST request to {url} (model={model}, stream={stream})")
+    response = requests.post(
+        url,
+        json=data,
+        headers=headers,
+        stream=stream,
+    )
+    response.raise_for_status()
+    return response
+
+
+def _process_rest_response(response: requests.Response, stream: bool = False) -> Union[str, Iterator[str]]:
+    if stream:
+        return _return_stream_response(response)
+
+    try:
+        content = response.json()["choices"][0]["message"]["content"]
+        assert isinstance(content, str)
+        return content
+    except (KeyError, IndexError, AssertionError) as e:
+        # Unlike the streaming case, errors are not ignored because a message must be returned.
+        raise ResponseParseException("Failed to parse message from response.") from e
+
+
+def _return_stream_response(response: requests.Response) -> Iterator[str]:
+    client = SSEClient(response)
+    for event in client.events():
+        try:
+            yield json.loads(event.data)["choices"][0]["delta"]["content"]
+        except (json.JSONDecodeError, KeyError, IndexError):
+            # For the sake of evolution of the output format,
+            # ignore stream messages that don't match the expected format.
+            pass
+
+
+def _complete_call_sql_function_snowpark(
+    function: str, *args: Union[str, snowpark.Column, CompleteOptions]
+) -> snowpark.Column:
+    return cast(snowpark.Column, functions.builtin(function)(*args))
+
+
+def _complete_call_sql_function_immediate(
+    function: str,
+    model: str,
+    prompt: Union[str, List[ConversationMessage]],
+    options: Optional[CompleteOptions],
+    session: Optional[snowpark.Session],
+) -> str:
+    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."""
+        )
+
+    # https://docs.snowflake.com/en/sql-reference/functions/complete-snowflake-cortex
+    if options is not None or not isinstance(prompt, str):
+        if isinstance(prompt, List):
+            prompt_arg = prompt
+        else:
+            prompt_arg = [{"role": "user", "content": prompt}]
+        options = options or {}
+        lit_args = [
+            functions.lit(model),
+            functions.lit(prompt_arg),
+            functions.lit(options),
+        ]
+    else:
+        lit_args = [
+            functions.lit(model),
+            functions.lit(prompt),
+        ]
+
+    empty_df = session.create_dataframe([snowpark.Row()])
+    df = empty_df.select(functions.builtin(function)(*lit_args))
+    return cast(str, df.collect()[0][0])
+
+
+def _complete_sql_impl(
+    function: str,
+    model: Union[str, snowpark.Column],
+    prompt: Union[str, List[ConversationMessage], snowpark.Column],
+    options: Optional[Union[CompleteOptions, snowpark.Column]],
+    session: Optional[snowpark.Session],
+) -> Union[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)
+    if isinstance(model, snowpark.Column):
+        raise ValueError("'model' cannot be a snowpark.Column when 'prompt' is a string.")
+    if isinstance(options, snowpark.Column):
+        raise ValueError("'options' cannot be a snowpark.Column when 'prompt' is a string.")
+    return _complete_call_sql_function_immediate(function, model, prompt, options, session)
+
+
+def _complete_impl(
+    model: Union[str, snowpark.Column],
+    prompt: Union[str, List[ConversationMessage], snowpark.Column],
+    options: Optional[CompleteOptions] = None,
+    session: Optional[snowpark.Session] = None,
+    use_rest_api_experimental: bool = False,
+    stream: bool = False,
+    function: str = "snowflake.cortex.complete",
+) -> Union[str, Iterator[str], snowpark.Column]:
+    if use_rest_api_experimental:
+        if not isinstance(model, str):
+            raise ValueError("in REST mode, 'model' must be a string")
+        if not isinstance(prompt, str) and not isinstance(prompt, List):
+            raise ValueError("in REST mode, 'prompt' must be a string or a list of ConversationMessage")
+        response = _call_complete_rest(model, prompt, options, session=session, stream=stream)
+        return _process_rest_response(response, stream=stream)
+    if stream is True:
+        raise ValueError("streaming can only be enabled in REST mode, set use_rest_api_experimental=True")
+    return _complete_sql_impl(function, model, prompt, options, session)
 
 
-@snowpark._internal.utils.experimental(version="1.0.12")
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
 def Complete(
-    model: Union[str, snowpark.Column], prompt: Union[str, snowpark.Column], session: Optional[snowpark.Session] = None
-) -> Union[str, snowpark.Column]:
+    model: Union[str, snowpark.Column],
+    prompt: Union[str, List[ConversationMessage], snowpark.Column],
+    *,
+    options: Optional[CompleteOptions] = None,
+    session: Optional[snowpark.Session] = None,
+    use_rest_api_experimental: bool = False,
+    stream: bool = False,
+) -> Union[str, Iterator[str], snowpark.Column]:
     """Complete calls into the LLM inference service to perform completion.
 
     Args:
         model: A Column of strings representing model types.
         prompt: A Column of prompts to send to the LLM.
+        options: A instance of snowflake.cortex.CompleteOptions
         session: The snowpark session to use. Will be inferred by context if not specified.
+        use_rest_api_experimental (bool): Toggles between the use of SQL and REST implementation. This feature is
+            experimental and can be removed at any time.
+        stream (bool): Enables streaming. When enabled, a generator function is returned that provides the streaming
+            output as it is received. Each update is a string containing the new text content since the previous update.
+            The use of streaming requires the experimental use_rest_api_experimental flag to be enabled.
+
+    Raises:
+        ValueError: If `stream` is set to True and `use_rest_api_experimental` is set to False.
 
     Returns:
         A column of string responses.
     """
-
-    return _complete_impl("snowflake.cortex.complete", model, prompt, session=session)
-
-
-def _complete_impl(
-    function: str,
-    model: Union[str, snowpark.Column],
-    prompt: Union[str, snowpark.Column],
-    session: Optional[snowpark.Session] = None,
-) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, model, prompt)
+    try:
+        return _complete_impl(model, prompt, options, session, use_rest_api_experimental, stream)
+    except ValueError as err:
+        raise err

snowflake/cortex/_extract_answer.py

@@ -5,7 +5,6 @@ from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_
 from snowflake.ml._internal import telemetry
 
 
-@snowpark._internal.utils.experimental(version="1.0.12")
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )

snowflake/cortex/_sentiment.py

@@ -5,7 +5,6 @@ from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_
 from snowflake.ml._internal import telemetry
 
 
-@snowpark._internal.utils.experimental(version="1.0.12")
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )

snowflake/cortex/_sse_client.py

@@ -0,0 +1,81 @@
+from typing import Iterator, cast
+
+import requests
+
+
+class Event:
+    def __init__(self, event: str = "message", data: str = "") -> None:
+        self.event = event
+        self.data = data
+
+    def __str__(self) -> str:
+        s = f"{self.event} event"
+        if self.data:
+            s += f", {len(self.data)} bytes"
+        else:
+            s += ", no data"
+        return s
+
+
+class SSEClient:
+    def __init__(self, response: requests.Response) -> None:
+
+        self.response = response
+
+    def _read(self) -> Iterator[str]:
+
+        lines = b""
+        for chunk in self.response:
+            for line in chunk.splitlines(True):
+                lines += line
+                if lines.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")):
+                    yield cast(str, lines)
+                    lines = b""
+        if lines:
+            yield cast(str, lines)
+
+    def events(self) -> Iterator[Event]:
+        for raw_event in self._read():
+            event = Event()
+            # splitlines() only uses \r and \n
+            for line in raw_event.splitlines():
+
+                line = cast(bytes, line).decode("utf-8")
+
+                data = line.split(":", 1)
+                field = data[0]
+
+                if len(data) > 1:
+                    # "If value starts with a single U+0020 SPACE character,
+                    # remove it from value. .strip() would remove all white spaces"
+                    if data[1].startswith(" "):
+                        value = data[1][1:]
+                    else:
+                        value = data[1]
+                else:
+                    value = ""
+
+                # The data field may come over multiple lines and their values
+                # are concatenated with each other.
+                if field == "data":
+                    event.data += value + "\n"
+                elif field == "event":
+                    event.event = value
+
+            if not event.data:
+                continue
+
+            # If the data field ends with a newline, remove it.
+            if event.data.endswith("\n"):
+                event.data = event.data[0:-1]  # Replace trailing newline - rstrip would remove multiple.
+
+            # Empty event names default to 'message'
+            event.event = event.event or "message"
+
+            if event.event != "message":  # ignore anything but “message” or default event
+                continue
+
+            yield event
+
+    def close(self) -> None:
+        self.response.close()

snowflake/cortex/_summarize.py

@@ -5,7 +5,6 @@ from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_
 from snowflake.ml._internal import telemetry
 
 
-@snowpark._internal.utils.experimental(version="1.0.12")
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )

snowflake/cortex/_translate.py

@@ -5,7 +5,6 @@ from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_
 from snowflake.ml._internal import telemetry
 
 
-@snowpark._internal.utils.experimental(version="1.0.12")
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )

snowflake/cortex/_util.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union, cast
+from typing import Dict, Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.snowpark import context, functions
@@ -6,32 +6,56 @@ from snowflake.snowpark import context, functions
 CORTEX_FUNCTIONS_TELEMETRY_PROJECT = "CortexFunctions"
 
 
+class SnowflakeAuthenticationException(Exception):
+    """This exception is raised when there is an issue with Snowflake's configuration."""
+
+    pass
+
+
+class SnowflakeConfigurationException(Exception):
+    """This exception is raised when there is an issue with Snowflake's configuration."""
+
+    pass
+
+
 # Calls a sql function, handling both immediate (e.g. python types) and batch
 # (e.g. snowpark column and literal type modes).
 def call_sql_function(
-    function: str, session: Optional[snowpark.Session], *args: Union[str, snowpark.Column]
+    function: str,
+    session: Optional[snowpark.Session],
+    *args: Union[str, snowpark.Column, Dict[str, Union[int, float]]],
 ) -> Union[str, snowpark.Column]:
     handle_as_column = False
+
     for arg in args:
         if isinstance(arg, snowpark.Column):
             handle_as_column = True
 
     if handle_as_column:
-        return cast(Union[str, snowpark.Column], call_sql_function_column(function, *args))
-    return cast(Union[str, snowpark.Column], call_sql_function_immediate(function, session, *args))
+        return cast(Union[str, snowpark.Column], _call_sql_function_column(function, *args))
+    return cast(
+        Union[str, snowpark.Column],
+        _call_sql_function_immediate(function, session, *args),
+    )
 
 
-def call_sql_function_column(function: str, *args: Union[str, snowpark.Column]) -> snowpark.Column:
+def _call_sql_function_column(
+    function: str, *args: Union[str, snowpark.Column, Dict[str, Union[int, float]]]
+) -> snowpark.Column:
     return cast(snowpark.Column, functions.builtin(function)(*args))
 
 
-def call_sql_function_immediate(
-    function: str, session: Optional[snowpark.Session], *args: Union[str, snowpark.Column]
+def _call_sql_function_immediate(
+    function: str,
+    session: Optional[snowpark.Session],
+    *args: Union[str, snowpark.Column, Dict[str, Union[int, float]]],
 ) -> str:
+    session = session or context.get_active_session()
     if session is None:
-        session = context.get_active_session()
-    if session is None:
-        raise Exception("No session available in the current context nor specified as an argument.")
+        raise SnowflakeAuthenticationException(
+            """Session required. Provide the session through a session=... argument or ensure an active session is
+            available in your environment."""
+        )
 
     lit_args = []
     for arg in args:

snowflake/ml/_internal/container_services/image_registry/http_client.py

@@ -64,13 +64,20 @@ class ImageRegistryHttpClient:
     operations. For general use of a retryable HTTP client, consider using the "retryable_http" module.
     """
 
-    def __init__(self, *, session: snowpark.Session, repo_url: str) -> None:
+    def __init__(self, *, repo_url: str, session: Optional[snowpark.Session] = None, no_cred: bool = False) -> None:
         self._repo_url = repo_url
-        self._session_token_manager = session_token_manager.SessionTokenManager(session)
         self._retryable_http = retryable_http.get_http_client()
-        self._bearer_token = ""
+        self._no_cred = no_cred
+
+        if not self._no_cred:
+            self._bearer_token = ""
+            assert session is not None
+            self._session_token_manager = session_token_manager.SessionTokenManager(session)
 
     def _with_bearer_token_header(self, headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
+        if self._no_cred:
+            return {} if not headers else headers.copy()
+
         if not self._bearer_token:
             self._fetch_bearer_token()
         assert self._bearer_token

snowflake/ml/_internal/container_services/image_registry/imagelib.py

@@ -13,6 +13,7 @@ This library only supports a limited set of features:
 It's recommended to use this library to copy previously tested images using sha256 to avoid surprises
 with respect to compatibility.
 """
+
 import dataclasses
 import hashlib
 import io
@@ -152,7 +153,8 @@ class BlobTransfer:
     src_image: ImageDescriptor
     dest_image: ImageDescriptor
     manifest: Manifest
-    image_registry_http_client: image_registry_http_client.ImageRegistryHttpClient
+    src_image_registry_http_client: image_registry_http_client.ImageRegistryHttpClient
+    dest_image_registry_http_client: image_registry_http_client.ImageRegistryHttpClient
 
     def upload_all_blobs(self) -> None:
         blob_digests = self.manifest.get_blob_digests()
@@ -169,7 +171,7 @@ class BlobTransfer:
         """
         Check if the blob already exists in the destination registry.
         """
-        resp = self.image_registry_http_client.head(self.dest_image.blob_link(blob_digest), headers={})
+        resp = self.dest_image_registry_http_client.head(self.dest_image.blob_link(blob_digest), headers={})
         return resp.status_code != 200
 
     def _fetch_blob(self, blob_digest: str) -> Tuple[io.BytesIO, int]:
@@ -178,7 +180,7 @@ class BlobTransfer:
         """
         src_blob_link = self.src_image.blob_link(blob_digest)
         headers = {_CONTENT_LENGTH_HEADER: "0"}
-        resp = self.image_registry_http_client.get(src_blob_link, headers=headers)
+        resp = self.src_image_registry_http_client.get(src_blob_link, headers=headers)
 
         assert resp.status_code == 200, f"Blob GET failed with code {resp.status_code}"
         assert _CONTENT_LENGTH_HEADER in resp.headers, f"Blob does not contain {_CONTENT_LENGTH_HEADER}"
@@ -189,7 +191,7 @@ class BlobTransfer:
         """
         Obtain the upload URL from the destination registry.
         """
-        response = self.image_registry_http_client.post(self.dest_image.blob_upload_link())
+        response = self.dest_image_registry_http_client.post(self.dest_image.blob_upload_link())
         assert (
             response.status_code == 202
         ), f"Failed to get the upload URL to destination. Status {response.status_code}. {str(response.content)}"
@@ -216,14 +218,14 @@ class BlobTransfer:
             headers[_CONTENT_RANGE_HEADER] = f"{start_byte}-{end_byte}"
             headers[_CONTENT_LENGTH_HEADER] = str(chunk_length)
 
-            resp = self.image_registry_http_client.patch(next_loc, headers=headers, data=chunk)
+            resp = self.dest_image_registry_http_client.patch(next_loc, headers=headers, data=chunk)
             assert resp.status_code == 202, f"Blob PATCH failed with code {resp.status_code}"
 
             next_loc = resp.headers[_LOCATION_HEADER]
             start_byte += chunk_length
 
         # Finalize the upload
-        resp = self.image_registry_http_client.put(f"{next_loc}&digest={blob_digest}")
+        resp = self.dest_image_registry_http_client.put(f"{next_loc}&digest={blob_digest}")
         assert resp.status_code == 201, f"Blob PUT failed with code {resp.status_code}"
 
     def _transfer(self, blob_digest: str) -> None:
@@ -340,21 +342,32 @@ def copy_image(
     src_image: ImageDescriptor,
     dest_image: ImageDescriptor,
     arch: _Arch,
-    retryable_http: image_registry_http_client.ImageRegistryHttpClient,
+    src_retryable_http: image_registry_http_client.ImageRegistryHttpClient,
+    dest_retryable_http: image_registry_http_client.ImageRegistryHttpClient,
 ) -> None:
     logger.debug(f"Pulling image manifest for {src_image}")
 
     # 1. Get the manifest
-    manifest = get_manifest(src_image, arch, retryable_http)
+    manifest = get_manifest(src_image, arch, src_retryable_http)
     logger.debug(f"Manifest pulled for {src_image} with digest {manifest.manifest_digest}")
 
     # 2: Retrieve all blob digests from manifest; fetch blob based on blob digest, then upload blob.
-    blob_transfer = BlobTransfer(src_image, dest_image, manifest, image_registry_http_client=retryable_http)
+    blob_transfer = BlobTransfer(
+        src_image,
+        dest_image,
+        manifest,
+        src_image_registry_http_client=src_retryable_http,
+        dest_image_registry_http_client=dest_retryable_http,
+    )
     blob_transfer.upload_all_blobs()
 
     # 3. Upload the manifest
     logger.debug(f"All blobs copied successfully. Copying manifest for {src_image} to {dest_image}")
-    put_manifest(dest_image, manifest, retryable_http)
+    put_manifest(
+        dest_image,
+        manifest,
+        dest_retryable_http,
+    )
 
     logger.debug(f"Image {src_image} copied to {dest_image}")
 

snowflake/ml/_internal/container_services/image_registry/registry_client.py

@@ -201,6 +201,12 @@ class ImageRegistryClient:
         )
         # TODO[shchen]: Remove the imagelib, instead rely on the copy image system function later.
         imagelib.copy_image(
-            src_image=src_image, dest_image=dest_image, arch=arch, retryable_http=self.image_registry_http_client
+            src_image=src_image,
+            dest_image=dest_image,
+            arch=arch,
+            src_retryable_http=image_registry_http_client.ImageRegistryHttpClient(
+                repo_url=src_image.registry_name, no_cred=True
+            ),
+            dest_retryable_http=self.image_registry_http_client,
         )
         logger.info("Image copy completed successfully")

snowflake/ml/_internal/exceptions/dataset_errors.py

@@ -1,11 +1,11 @@
 # Error code from Snowflake Python Connector.
-ERRNO_OBJECT_ALREADY_EXISTS = "002002"
-ERRNO_OBJECT_NOT_EXIST = "002043"
-ERRNO_FILES_ALREADY_EXISTING = "001030"
-ERRNO_VERSION_ALREADY_EXISTS = "092917"
-ERRNO_DATASET_NOT_EXIST = "399019"
-ERRNO_DATASET_VERSION_NOT_EXIST = "399012"
-ERRNO_DATASET_VERSION_ALREADY_EXISTS = "399020"
+ERRNO_OBJECT_ALREADY_EXISTS = 2002
+ERRNO_OBJECT_NOT_EXIST = 2043
+ERRNO_FILES_ALREADY_EXISTING = 1030
+ERRNO_VERSION_ALREADY_EXISTS = 92917
+ERRNO_DATASET_NOT_EXIST = 399019
+ERRNO_DATASET_VERSION_NOT_EXIST = 399012
+ERRNO_DATASET_VERSION_ALREADY_EXISTS = 399020
 
 
 class DatasetError(Exception):

snowflake/ml/_internal/exceptions/fileset_errors.py

@@ -1,7 +1,7 @@
 # Error code from Snowflake Python Connector.
-ERRNO_FILE_EXIST_IN_STAGE = "001030"
-ERRNO_DOMAIN_NOT_EXIST = "002003"
-ERRNO_STAGE_NOT_EXIST = "391707"
+ERRNO_FILE_EXIST_IN_STAGE = 1030
+ERRNO_DOMAIN_NOT_EXIST = 2003
+ERRNO_STAGE_NOT_EXIST = 391707
 
 
 class FileSetError(Exception):

snowflake/ml/_internal/exceptions/sql_error_codes.py

@@ -0,0 +1,6 @@
+"""SQL Error Codes"""
+
+# SQL compilation error: Object ''{0}'' does not exist or not authorized.
+OBJECT_NOT_EXIST = 2001
+# SQL compilation error: Object ''{0}'' already exists.
+OBJECT_ALREADY_EXISTS = 2002