snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

snowflake/cortex/_classify_text.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union, cast
+from typing import Optional, Union, cast
 
 from typing_extensions import deprecated
 
@@ -12,7 +12,7 @@ from snowflake.ml._internal import telemetry
 )
 def classify_text(
     str_input: Union[str, snowpark.Column],
-    categories: Union[List[str], snowpark.Column],
+    categories: Union[list[str], snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
     """Use the LLM inference service to classify the INPUT text into one of the target CATEGORIES.
@@ -32,7 +32,7 @@ def classify_text(
 def _classify_text_impl(
     function: str,
     str_input: Union[str, snowpark.Column],
-    categories: Union[List[str], snowpark.Column],
+    categories: Union[list[str], snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
     return cast(Union[str, snowpark.Column], call_sql_function(function, session, str_input, categories))

snowflake/cortex/_complete.py

@@ -1,11 +1,13 @@
 import json
 import logging
 import time
+import typing
 from io import BytesIO
-from typing import Any, Callable, Dict, Iterator, List, Optional, TypedDict, Union, cast
+from typing import Any, Callable, Iterator, Optional, TypedDict, Union, cast
 from urllib.parse import urlunparse
 
 import requests
+from snowflake.core.rest import RESTResponse
 from typing_extensions import NotRequired, deprecated
 
 from snowflake import snowpark
@@ -28,7 +30,7 @@ class ResponseFormat(TypedDict):
 
     type: str
     """The response format type (e.g. "json")"""
-    schema: Dict[str, Any]
+    schema: dict[str, Any]
     """The schema defining the structure of the response. For json it should be a valid json schema object"""
 
 
@@ -69,7 +71,27 @@ class CompleteOptions(TypedDict):
 class ResponseParseException(Exception):
     """This exception is raised when the server response cannot be parsed."""
 
-    pass
+
+class MidStreamException(Exception):
+    """The SSE (Server-sent Event) stream can contain error messages in the middle of the stream,
+    using the “error” event type. This exception is raised when there is such a mid-stream error.
+    """
+
+    def __init__(
+        self,
+        reason: typing.Optional[str] = None,
+        http_resp: typing.Optional["RESTResponse"] = None,
+        request_id: typing.Optional[str] = None,
+    ) -> None:
+        message = ""
+        if reason is not None:
+            message = reason
+        if http_resp:
+            message = f"Error in stream (HTTP Response: {http_resp.status}) - {http_resp.reason}"
+        if request_id != "":
+            # add request_id to error message
+            message += f" (Request ID: {request_id})"
+        super().__init__(message)
 
 
 class GuardrailsOptions(TypedDict):
@@ -112,7 +134,7 @@ def retry(func: Callable[..., requests.Response]) -> Callable[..., requests.Resp
     return inner
 
 
-def _make_common_request_headers() -> Dict[str, str]:
+def _make_common_request_headers() -> dict[str, str]:
     headers = {
         "Content-Type": "application/json",
         "Accept": "application/json, text/event-stream",
@@ -120,6 +142,18 @@ def _make_common_request_headers() -> Dict[str, str]:
     return headers
 
 
+def _get_request_id(resp: dict[str, Any]) -> Optional[Any]:
+    request_id = None
+    if "headers" in resp:
+        for key, value in resp["headers"].items():
+            # Note: There is some whitespace in the headers making it not possible
+            # to directly index the header reliably.
+            if key.strip().lower() == "x-snowflake-request-id":
+                request_id = value
+                break
+    return request_id
+
+
 def _validate_response_format_object(options: CompleteOptions) -> None:
     """Validate the response format object for structured-output mode.
 
@@ -148,14 +182,14 @@ def _validate_response_format_object(options: CompleteOptions) -> None:
 
 def _make_request_body(
     model: str,
-    prompt: Union[str, List[ConversationMessage]],
+    prompt: Union[str, list[ConversationMessage]],
     options: Optional[CompleteOptions] = None,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     data = {
         "model": model,
         "stream": True,
     }
-    if isinstance(prompt, List):
+    if isinstance(prompt, list):
         data["messages"] = prompt
     else:
         data["messages"] = [{"content": prompt}]
@@ -182,19 +216,13 @@ def _make_request_body(
 
 # XP endpoint returns a dict response which needs to be converted to a format which can
 # be consumed by the SSEClient. This method does that.
-def _xp_dict_to_response(raw_resp: Dict[str, Any]) -> requests.Response:
+def _xp_dict_to_response(raw_resp: dict[str, Any]) -> requests.Response:
 
     response = requests.Response()
     response.status_code = int(raw_resp["status"])
     response.headers = raw_resp["headers"]
 
-    request_id = None
-    for key, value in raw_resp["headers"].items():
-        # Note: there is some whitespace in the headers making it not possible
-        # to directly index the header reliably.
-        if key.strip().lower() == "x-snowflake-request-id":
-            request_id = value
-            break
+    request_id = _get_request_id(raw_resp)
 
     data = raw_resp["content"]
     try:
@@ -222,9 +250,9 @@ def _xp_dict_to_response(raw_resp: Dict[str, Any]) -> requests.Response:
 
 @retry
 def _call_complete_xp(
-    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]],
+    snow_api_xp_request_handler: Optional[Callable[..., dict[str, Any]]],
     model: str,
-    prompt: Union[str, List[ConversationMessage]],
+    prompt: Union[str, list[ConversationMessage]],
     options: Optional[CompleteOptions] = None,
     deadline: Optional[float] = None,
 ) -> requests.Response:
@@ -238,7 +266,7 @@ def _call_complete_xp(
 @retry
 def _call_complete_rest(
     model: str,
-    prompt: Union[str, List[ConversationMessage]],
+    prompt: Union[str, list[ConversationMessage]],
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
 ) -> requests.Response:
@@ -276,7 +304,12 @@ def _call_complete_rest(
     )
 
 
-def _return_stream_response(response: requests.Response, deadline: Optional[float]) -> Iterator[str]:
+def _return_stream_response(
+    response: requests.Response,
+    deadline: Optional[float],
+    session: Optional[snowpark.Session] = None,
+) -> Iterator[str]:
+    request_id = _get_request_id(dict(response.headers))
     client = SSEClient(response)
     for event in client.events():
         if deadline is not None and time.time() > deadline:
@@ -294,7 +327,7 @@ def _return_stream_response(response: requests.Response, deadline: Optional[floa
             # This is the case of midstream errors which were introduced specifically for structured output.
             # TODO: discuss during code review
             if parsed_resp.get("error"):
-                yield json.dumps(parsed_resp)
+                raise MidStreamException(reason=response.text, request_id=request_id)
             else:
                 pass
 
@@ -306,9 +339,9 @@ def _complete_call_sql_function_snowpark(
 
 
 def _complete_non_streaming_immediate(
-    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]],
+    snow_api_xp_request_handler: Optional[Callable[..., dict[str, Any]]],
     model: str,
-    prompt: Union[str, List[ConversationMessage]],
+    prompt: Union[str, list[ConversationMessage]],
     options: Optional[CompleteOptions],
     session: Optional[snowpark.Session] = None,
     deadline: Optional[float] = None,
@@ -325,10 +358,10 @@ def _complete_non_streaming_immediate(
 
 
 def _complete_non_streaming_impl(
-    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]],
+    snow_api_xp_request_handler: Optional[Callable[..., dict[str, Any]]],
     function: str,
     model: Union[str, snowpark.Column],
-    prompt: Union[str, List[ConversationMessage], snowpark.Column],
+    prompt: Union[str, list[ConversationMessage], snowpark.Column],
     options: Optional[Union[CompleteOptions, snowpark.Column]],
     session: Optional[snowpark.Session] = None,
     deadline: Optional[float] = None,
@@ -355,9 +388,9 @@ def _complete_non_streaming_impl(
 
 
 def _complete_rest(
-    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]],
+    snow_api_xp_request_handler: Optional[Callable[..., dict[str, Any]]],
     model: str,
-    prompt: Union[str, List[ConversationMessage]],
+    prompt: Union[str, list[ConversationMessage]],
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
     deadline: Optional[float] = None,
@@ -375,13 +408,13 @@ def _complete_rest(
     else:
         response = _call_complete_rest(model=model, prompt=prompt, options=options, session=session, deadline=deadline)
     assert response.status_code >= 200 and response.status_code < 300
-    return _return_stream_response(response, deadline)
+    return _return_stream_response(response, deadline, session)
 
 
 def _complete_impl(
     model: Union[str, snowpark.Column],
-    prompt: Union[str, List[ConversationMessage], snowpark.Column],
-    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]] = None,
+    prompt: Union[str, list[ConversationMessage], snowpark.Column],
+    snow_api_xp_request_handler: Optional[Callable[..., dict[str, Any]]] = None,
     function: str = "snowflake.cortex.complete",
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
@@ -396,7 +429,7 @@ def _complete_impl(
     if stream:
         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):
+        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")
         return _complete_rest(
             snow_api_xp_request_handler=snow_api_xp_request_handler,
@@ -422,7 +455,7 @@ def _complete_impl(
 )
 def complete(
     model: Union[str, snowpark.Column],
-    prompt: Union[str, List[ConversationMessage], snowpark.Column],
+    prompt: Union[str, list[ConversationMessage], snowpark.Column],
     *,
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,

snowflake/cortex/_embed_text_1024.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union, cast
+from typing import Optional, Union, cast
 
 from typing_extensions import deprecated
 
@@ -14,7 +14,7 @@ def embed_text_1024(
     model: Union[str, snowpark.Column],
     text: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
-) -> Union[List[float], snowpark.Column]:
+) -> Union[list[float], snowpark.Column]:
     """Calls into the LLM inference service to embed the text.
 
     Args:
@@ -35,8 +35,8 @@ def _embed_text_1024_impl(
     model: Union[str, snowpark.Column],
     text: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
-) -> Union[List[float], snowpark.Column]:
-    return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))
+) -> Union[list[float], snowpark.Column]:
+    return cast(Union[list[float], snowpark.Column], call_sql_function(function, session, model, text))
 
 
 EmbedText1024 = deprecated(

snowflake/cortex/_embed_text_768.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union, cast
+from typing import Optional, Union, cast
 
 from typing_extensions import deprecated
 
@@ -14,7 +14,7 @@ def embed_text_768(
     model: Union[str, snowpark.Column],
     text: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
-) -> Union[List[float], snowpark.Column]:
+) -> Union[list[float], snowpark.Column]:
     """Calls into the LLM inference service to embed the text.
 
     Args:
@@ -35,8 +35,8 @@ def _embed_text_768_impl(
     model: Union[str, snowpark.Column],
     text: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
-) -> Union[List[float], snowpark.Column]:
-    return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))
+) -> Union[list[float], snowpark.Column]:
+    return cast(Union[list[float], snowpark.Column], call_sql_function(function, session, model, text))
 
 
 EmbedText768 = deprecated(

snowflake/cortex/_finetune.py

@@ -1,6 +1,6 @@
 import json
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import (
@@ -53,7 +53,7 @@ class FinetuneStatus:
     created_on: Optional[int] = None
     """Creation timestamp of the Fine-tuning job in milliseconds."""
 
-    error: Optional[Dict[str, Any]] = None
+    error: Optional[dict[str, Any]] = None
     """Error message propagated from the job."""
 
     finished_on: Optional[int] = None
@@ -62,7 +62,7 @@ class FinetuneStatus:
     progress: Optional[float] = None
     """Progress made as a fraction of total [0.0,1.0]."""
 
-    training_result: Optional[List[Dict[str, Any]]] = None
+    training_result: Optional[list[dict[str, Any]]] = None
     """Detailed metrics report for a completed training."""
 
     trained_tokens: Optional[int] = None
@@ -135,7 +135,7 @@ class FinetuneJob:
         """
         result_string = _finetune_impl(operation="DESCRIBE", session=self._session, function_args=[self.status.id])
 
-        result = FinetuneStatus(**cast(Dict[str, Any], _try_load_json(result_string)))
+        result = FinetuneStatus(**cast(dict[str, Any], _try_load_json(result_string)))
         return result
 
 
@@ -167,7 +167,7 @@ class Finetune:
         base_model: str,
         training_data: Union[str, snowpark.DataFrame],
         validation_data: Optional[Union[str, snowpark.DataFrame]] = None,
-        options: Optional[Dict[str, Any]] = None,
+        options: Optional[dict[str, Any]] = None,
     ) -> FinetuneJob:
         """Create a new fine-tuning runs.
 
@@ -240,7 +240,7 @@ class Finetune:
         project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
         subproject=CORTEX_FINETUNE_TELEMETRY_SUBPROJECT,
     )
-    def list_jobs(self) -> List["FinetuneJob"]:
+    def list_jobs(self) -> list["FinetuneJob"]:
         """Show current and past fine-tuning runs.
 
         Returns:
@@ -253,7 +253,7 @@ class Finetune:
         return [FinetuneJob(session=self._session, status=FinetuneStatus(**run_status)) for run_status in result]
 
 
-def _try_load_json(json_string: str) -> Union[Dict[Any, Any], List[Any]]:
+def _try_load_json(json_string: str) -> Union[dict[Any, Any], list[Any]]:
     try:
         result = json.loads(str(json_string))
     except json.JSONDecodeError as e:
@@ -269,5 +269,5 @@ def _try_load_json(json_string: str) -> Union[Dict[Any, Any], List[Any]]:
     return result
 
 
-def _finetune_impl(operation: str, session: Optional[snowpark.Session], function_args: List[Any]) -> str:
+def _finetune_impl(operation: str, session: Optional[snowpark.Session], function_args: list[Any]) -> str:
     return call_sql_function_literals(_CORTEX_FINETUNE_SYSTEM_FUNCTION_NAME, session, operation, *function_args)

snowflake/cortex/_util.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.ml._internal.exceptions import error_codes, exceptions
@@ -11,22 +11,18 @@ 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, List[str], snowpark.Column, Dict[str, Union[int, float]]],
-) -> Union[str, List[float], snowpark.Column]:
+    *args: Union[str, list[str], snowpark.Column, dict[str, Union[int, float]]],
+) -> Union[str, list[float], snowpark.Column]:
     handle_as_column = False
 
     for arg in args:
@@ -34,15 +30,15 @@ def call_sql_function(
             handle_as_column = True
 
     if handle_as_column:
-        return cast(Union[str, List[float], snowpark.Column], _call_sql_function_column(function, *args))
+        return cast(Union[str, list[float], snowpark.Column], _call_sql_function_column(function, *args))
     return cast(
-        Union[str, List[float], snowpark.Column],
+        Union[str, list[float], snowpark.Column],
         _call_sql_function_immediate(function, session, *args),
     )
 
 
 def _call_sql_function_column(
-    function: str, *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]]
+    function: str, *args: Union[str, list[str], snowpark.Column, dict[str, Union[int, float]]]
 ) -> snowpark.Column:
     return cast(snowpark.Column, functions.builtin(function)(*args))
 
@@ -50,8 +46,8 @@ def _call_sql_function_column(
 def _call_sql_function_immediate(
     function: str,
     session: Optional[snowpark.Session],
-    *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
-) -> Union[str, List[float]]:
+    *args: Union[str, list[str], snowpark.Column, dict[str, Union[int, float]]],
+) -> Union[str, list[float]]:
     session = session or context.get_active_session()
     if session is None:
         raise SnowflakeAuthenticationException(

snowflake/ml/_internal/env.py

@@ -1,8 +1,9 @@
+import os
 import platform
 
-from snowflake.ml import version
-
 SOURCE = "SnowML"
-VERSION = version.VERSION
 PYTHON_VERSION = platform.python_version()
 OS = platform.system()
+IN_ML_RUNTIME_ENV_VAR = "IN_SPCS_ML_RUNTIME"
+IN_ML_RUNTIME = os.getenv(IN_ML_RUNTIME_ENV_VAR)
+USE_OPTIMIZED_DATA_INGESTOR = "USE_OPTIMIZED_DATA_INGESTOR"