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

snowflake/cortex/_complete.py

@@ -1,7 +1,12 @@
-from typing import Optional, Union
+from typing import Iterator, Optional, Union
 
 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_rest_function,
+    call_sql_function,
+    process_rest_response,
+)
 from snowflake.ml._internal import telemetry
 
 
@@ -10,19 +15,35 @@ from snowflake.ml._internal import 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, snowpark.Column],
+    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.
         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.
     """
-
+    if stream is True and use_rest_api_experimental is False:
+        raise ValueError("If stream is set to True use_rest_api_experimental must also be set to True")
+    if use_rest_api_experimental:
+        response = call_rest_function("complete", model, prompt, session=session, stream=stream)
+        return process_rest_response(response)
     return _complete_impl("snowflake.cortex.complete", model, prompt, session=session)
 
 

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/_util.py

@@ -1,15 +1,34 @@
-from typing import Optional, Union, cast
+import json
+from typing import Iterator, Optional, Union, cast
+from urllib.parse import urljoin, urlparse
+
+import requests
 
 from snowflake import snowpark
+from snowflake.cortex._sse_client import SSEClient
 from snowflake.snowpark import context, functions
 
 CORTEX_FUNCTIONS_TELEMETRY_PROJECT = "CortexFunctions"
 
 
+class SSEParseException(Exception):
+    """This exception is raised when an invalid server sent event is received from the server."""
+
+    pass
+
+
+class SnowflakeAuthenticationException(Exception):
+    """This exception is raised when the session object does not have session.connection.rest.token attribute."""
+
+    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],
 ) -> Union[str, snowpark.Column]:
     handle_as_column = False
     for arg in args:
@@ -17,21 +36,29 @@ def call_sql_function(
             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]) -> 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],
 ) -> str:
     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:
@@ -40,3 +67,73 @@ def call_sql_function_immediate(
     empty_df = session.create_dataframe([snowpark.Row()])
     df = empty_df.select(functions.builtin(function)(*lit_args))
     return cast(str, df.collect()[0][0])
+
+
+def call_rest_function(
+    function: str,
+    model: Union[str, snowpark.Column],
+    prompt: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+    stream: bool = False,
+) -> requests.Response:
+    if session is None:
+        session = 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 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 == "":  # type: ignore[union-attr]
+        raise SnowflakeAuthenticationException("Snowflake session error: REST token is empty.")
+
+    url = urljoin(session.connection.host, f"api/v2/cortex/inference/{function}")
+    if urlparse(url).scheme == "":
+        url = "https://" + url
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f'Snowflake Token="{session.connection.rest.token}"',  # type: ignore[union-attr]
+        "Accept": "application/json, text/event-stream",
+    }
+
+    data = {
+        "model": model,
+        "messages": [{"content": prompt}],
+        "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 not stream:
+        try:
+            message = response.json()["choices"][0]["message"]
+            output = str(message.get("content", ""))
+            return output
+        except (KeyError, IndexError) as e:
+            raise SSEParseException("Failed to parse streamed response.") from e
+    else:
+        return _return_gen(response)
+
+
+def _return_gen(response: requests.Response) -> Iterator[str]:
+    client = SSEClient(response)
+    for event in client.events():
+        response_loaded = json.loads(event.data)
+        try:
+            delta = response_loaded["choices"][0]["delta"]
+            output = str(delta.get("content", ""))
+            yield output
+        except (KeyError, IndexError) as e:
+            raise SSEParseException("Failed to parse streamed response.") from e

snowflake/ml/_internal/lineage/lineage_utils.py

@@ -1,21 +1,11 @@
 import copy
 import functools
-from typing import Any, Callable, List
+from typing import Any, Callable, List, Optional
 
 from snowflake import snowpark
 from snowflake.ml._internal.lineage import data_source
 
-DATA_SOURCES_ATTR = "_data_sources"
-
-
-def _get_datasources(*args: Any) -> List[data_source.DataSource]:
-    """Helper method for extracting data sources attribute from DataFrames in an argument list"""
-    result = []
-    for arg in args:
-        srcs = getattr(arg, DATA_SOURCES_ATTR, None)
-        if isinstance(srcs, list) and all(isinstance(s, data_source.DataSource) for s in srcs):
-            result += srcs
-    return result
+_DATA_SOURCES_ATTR = "_data_sources"
 
 
 def _wrap_func(
@@ -32,6 +22,37 @@ def _wrap_func(
     return wrapped
 
 
+def _wrap_class_func(fn: Callable[..., snowpark.DataFrame]) -> Callable[..., snowpark.DataFrame]:
+    @functools.wraps(fn)
+    def wrapped(*args: Any, **kwargs: Any) -> snowpark.DataFrame:
+        df = fn(*args, **kwargs)
+        data_sources = get_data_sources(*args, *kwargs.values())
+        if data_sources:
+            patch_dataframe(df, data_sources, inplace=True)
+        return df
+
+    return wrapped
+
+
+def get_data_sources(*args: Any) -> Optional[List[data_source.DataSource]]:
+    """Helper method for extracting data sources attribute from DataFrames in an argument list"""
+    result: Optional[List[data_source.DataSource]] = None
+    for arg in args:
+        srcs = getattr(arg, _DATA_SOURCES_ATTR, None)
+        if isinstance(srcs, list) and all(isinstance(s, data_source.DataSource) for s in srcs):
+            if result is None:
+                result = []
+            result += srcs
+    return result
+
+
+def set_data_sources(obj: Any, data_sources: Optional[List[data_source.DataSource]]) -> None:
+    """Helper method for attaching data sources to an object"""
+    if data_sources:
+        assert all(isinstance(ds, data_source.DataSource) for ds in data_sources)
+    setattr(obj, _DATA_SOURCES_ATTR, data_sources)
+
+
 def patch_dataframe(
     df: snowpark.DataFrame, data_sources: List[data_source.DataSource], inplace: bool = False
 ) -> snowpark.DataFrame:
@@ -62,7 +83,7 @@ def patch_dataframe(
     ]
     if not inplace:
         df = copy.copy(df)
-    setattr(df, DATA_SOURCES_ATTR, data_sources)
+    set_data_sources(df, data_sources)
     for func in funcs:
         fn = getattr(df, func, None)
         if fn is not None:
@@ -70,18 +91,6 @@ def patch_dataframe(
     return df
 
 
-def _wrap_class_func(fn: Callable[..., snowpark.DataFrame]) -> Callable[..., snowpark.DataFrame]:
-    @functools.wraps(fn)
-    def wrapped(*args: Any, **kwargs: Any) -> snowpark.DataFrame:
-        df = fn(*args, **kwargs)
-        data_sources = _get_datasources(*args) + _get_datasources(*kwargs.values())
-        if data_sources:
-            patch_dataframe(df, data_sources, inplace=True)
-        return df
-
-    return wrapped
-
-
 # Class-level monkey-patches
 for klass, func_list in {
     snowpark.DataFrame: [

snowflake/ml/dataset/dataset.py

@@ -65,6 +65,20 @@ class DatasetVersion:
         comment: Optional[str] = self._get_property("comment")
         return comment
 
+    @property
+    def label_cols(self) -> List[str]:
+        metadata = self._get_metadata()
+        if metadata is None or metadata.label_cols is None:
+            return []
+        return metadata.label_cols
+
+    @property
+    def exclude_cols(self) -> List[str]:
+        metadata = self._get_metadata()
+        if metadata is None or metadata.exclude_cols is None:
+            return []
+        return metadata.exclude_cols
+
     def _get_property(self, property_name: str, default: Any = None) -> Any:
         if self._properties is None:
             sql_result = (
@@ -91,17 +105,6 @@ class DatasetVersion:
                 warnings.warn(f"Metadata parsing failed with error: {e}", UserWarning, stacklevel=2)
         return self._metadata
 
-    def _get_exclude_cols(self) -> List[str]:
-        metadata = self._get_metadata()
-        if metadata is None:
-            return []
-        cols = []
-        if metadata.exclude_cols:
-            cols.extend(metadata.exclude_cols)
-        if metadata.label_cols:
-            cols.extend(metadata.label_cols)
-        return cols
-
     def url(self) -> str:
         """Returns the URL of the DatasetVersion contents in Snowflake.
 
@@ -168,7 +171,7 @@ class Dataset:
                         fully_qualified_name=self._fully_qualified_name,
                         version=v.name,
                         url=v.url(),
-                        exclude_cols=v._get_exclude_cols(),
+                        exclude_cols=(v.label_cols + v.exclude_cols),
                     )
                 ],
             )

snowflake/ml/dataset/dataset_factory.py

@@ -16,8 +16,7 @@ def create_from_dataframe(
     **version_kwargs: Any,
 ) -> dataset.Dataset:
     """
-    Create a new versioned Dataset from a DataFrame and returns
-    a DatasetReader for the newly created Dataset version.
+    Create a new versioned Dataset from a DataFrame.
 
     Args:
         session: The Snowpark Session instance to use.
@@ -39,7 +38,7 @@ def create_from_dataframe(
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def load_dataset(session: snowpark.Session, name: str, version: str) -> dataset.Dataset:
     """
-    Load a versioned Dataset into a DatasetReader.
+    Load a versioned Dataset.
 
     Args:
         session: The Snowpark Session instance to use.
@@ -47,7 +46,7 @@ def load_dataset(session: snowpark.Session, name: str, version: str) -> dataset.
         version: The dataset version name.
 
     Returns:
-        A DatasetReader object.
+        A Dataset object.
     """
     ds: dataset.Dataset = dataset.Dataset.load(session, name).select_version(version)
     return ds

snowflake/ml/feature_store/feature_store.py

@@ -920,7 +920,7 @@ class FeatureStore:
         try:
             if output_type == "table":
                 table_name = f"{name}_{version}"
-                result_df.write.mode("errorifexists").save_as_table(table_name)  # type: ignore[call-overload]
+                result_df.write.mode("errorifexists").save_as_table(table_name)
                 ds_df = self._session.table(table_name)
                 return ds_df
             else:
@@ -1762,7 +1762,7 @@ class FeatureStore:
                 f"""
                     SELECT * FROM TABLE(
                         {self._config.database}.INFORMATION_SCHEMA.TAG_REFERENCES_INTERNAL(
-                            TAG_NAME => '{_FEATURE_STORE_OBJECT_TAG}'
+                            TAG_NAME => '{self._get_fully_qualified_name(_FEATURE_STORE_OBJECT_TAG)}'
                         )
                     ) LIMIT 1;
                 """

snowflake/ml/model/_client/sql/model_version.py

@@ -272,7 +272,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
                 actual_schema_name.identifier(),
                 tmp_table_name,
             )
-            input_df.write.save_as_table(  # type: ignore[call-overload]
+            input_df.write.save_as_table(
                 table_name=INTERMEDIATE_TABLE_NAME,
                 mode="errorifexists",
                 table_type="temporary",
@@ -348,7 +348,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
                 actual_schema_name.identifier(),
                 tmp_table_name,
             )
-            input_df.write.save_as_table(  # type: ignore[call-overload]
+            input_df.write.save_as_table(
                 table_name=INTERMEDIATE_TABLE_NAME,
                 mode="errorifexists",
                 table_type="temporary",

snowflake/ml/model/_model_composer/model_composer.py

@@ -182,9 +182,9 @@ class ModelComposer:
     def _get_data_sources(
         self, model: model_types.SupportedModelType, sample_input_data: Optional[model_types.SupportedDataType] = None
     ) -> Optional[List[data_source.DataSource]]:
-        data_sources = getattr(model, lineage_utils.DATA_SOURCES_ATTR, None)
+        data_sources = lineage_utils.get_data_sources(model)
         if not data_sources and sample_input_data is not None:
-            data_sources = getattr(sample_input_data, lineage_utils.DATA_SOURCES_ATTR, None)
+            data_sources = lineage_utils.get_data_sources(sample_input_data)
         if isinstance(data_sources, list) and all(isinstance(item, data_source.DataSource) for item in data_sources):
             return data_sources
         return None

snowflake/ml/model/_model_composer/model_method/infer_table_function.py_template

@@ -74,4 +74,6 @@ dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
 class {function_name}:
     @vectorized(input=pd.DataFrame)
     def end_partition(self, df: pd.DataFrame) -> pd.DataFrame:
-        return runner(df)
+        df.columns = input_cols
+        input_df = df.astype(dtype=dtype_map)
+        return runner(input_df[input_cols])

snowflake/ml/model/_packager/model_meta/_core_requirements.py

@@ -6,6 +6,6 @@ REQUIREMENTS = [
     "packaging>=20.9,<24",
     "pandas>=1.0.0,<3",
     "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.11.1,<2,!=1.12.0",
+    "snowflake-snowpark-python>=1.15.0,<2",
     "typing-extensions>=4.1.0,<5"
 ]

snowflake/ml/model/_packager/model_runtime/_snowml_inference_alternative_requirements.py

@@ -5,6 +5,6 @@ REQUIREMENTS = [
     "packaging>=20.9,<24",
     "pandas>=1.0.0,<3",
     "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.11.1,<2,!=1.12.0",
+    "snowflake-snowpark-python>=1.15.0,<2",
     "typing-extensions>=4.1.0,<5"
 ]

snowflake/ml/model/_signatures/builtins_handler.py

@@ -1,3 +1,4 @@
+import datetime
 from collections import abc
 from typing import Literal, Sequence
 
@@ -24,7 +25,7 @@ class ListOfBuiltinHandler(base_handler.BaseDataHandler[model_types._SupportedBu
             # String is a Sequence but we take them as an whole
             if isinstance(element, abc.Sequence) and not isinstance(element, str):
                 can_handle = ListOfBuiltinHandler.can_handle(element)
-            elif not isinstance(element, (int, float, bool, str)):
+            elif not isinstance(element, (int, float, bool, str, datetime.datetime)):
                 can_handle = False
                 break
         return can_handle

snowflake/ml/model/_signatures/core.py

@@ -53,6 +53,8 @@ class DataType(Enum):
     STRING = ("string", spt.StringType, np.str_)
     BYTES = ("bytes", spt.BinaryType, np.bytes_)
 
+    TIMESTAMP_NTZ = ("datetime64[ns]", spt.TimestampType, "datetime64[ns]")
+
     def as_snowpark_type(self) -> spt.DataType:
         """Convert to corresponding Snowpark Type.
 
@@ -78,6 +80,13 @@ class DataType(Enum):
             Corresponding DataType.
         """
         np_to_snowml_type_mapping = {i._numpy_type: i for i in DataType}
+
+        # Add datetime types:
+        datetime_res = ["Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns"]
+
+        for res in datetime_res:
+            np_to_snowml_type_mapping[f"datetime64[{res}]"] = DataType.TIMESTAMP_NTZ
+
         for potential_type in np_to_snowml_type_mapping.keys():
             if np.can_cast(np_type, potential_type, casting="no"):
                 # This is used since the same dtype might represented in different ways.
@@ -247,9 +256,12 @@ class FeatureSpec(BaseFeatureSpec):
             result_type = spt.ArrayType(result_type)
         return result_type
 
-    def as_dtype(self) -> npt.DTypeLike:
+    def as_dtype(self) -> Union[npt.DTypeLike, str]:
         """Convert to corresponding local Type."""
         if not self._shape:
+            # scalar dtype: use keys from `np.sctypeDict` to prevent unit-less dtype 'datetime64'
+            if "datetime64" in self._dtype._value:
+                return self._dtype._value
             return self._dtype._numpy_type
         return np.object_
 

snowflake/ml/model/_signatures/pandas_handler.py

@@ -147,6 +147,8 @@ class PandasDataFrameHandler(base_handler.BaseDataHandler[pd.DataFrame]):
                     specs.append(core.FeatureSpec(dtype=core.DataType.STRING, name=ft_name))
                 elif isinstance(data[df_col].iloc[0], bytes):
                     specs.append(core.FeatureSpec(dtype=core.DataType.BYTES, name=ft_name))
+            elif isinstance(data[df_col].iloc[0], np.datetime64):
+                specs.append(core.FeatureSpec(dtype=core.DataType.TIMESTAMP_NTZ, name=ft_name))
             else:
                 specs.append(core.FeatureSpec(dtype=core.DataType.from_numpy_type(df_col_dtype), name=ft_name))
         return specs

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -107,6 +107,9 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         if not features:
             features = pandas_handler.PandasDataFrameHandler.infer_signature(df, role="input")
         # Role will be no effect on the column index. That is to say, the feature name is the actual column name.
+        if keep_order:
+            df = df.reset_index(drop=True)
+            df[infer_template._KEEP_ORDER_COL_NAME] = df.index
         sp_df = session.create_dataframe(df)
         column_names = []
         columns = []
@@ -122,7 +125,4 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
 
         sp_df = sp_df.with_columns(column_names, columns)
 
-        if keep_order:
-            sp_df = sp_df.with_column(infer_template._KEEP_ORDER_COL_NAME, F.monotonically_increasing_id())
-
         return sp_df

snowflake/ml/model/model_signature.py

@@ -168,6 +168,8 @@ def _validate_numpy_array(
             max_v <= np.finfo(feature_type._numpy_type).max  # type: ignore[arg-type]
             and min_v >= np.finfo(feature_type._numpy_type).min  # type: ignore[arg-type]
         )
+    elif feature_type in [core.DataType.TIMESTAMP_NTZ]:
+        return np.issubdtype(arr.dtype, np.datetime64)
     else:
         return np.can_cast(arr.dtype, feature_type._numpy_type, casting="no")
 

snowflake/ml/model/type_hints.py

@@ -54,6 +54,7 @@ _SupportedNumpyDtype = Union[
     "np.bool_",
     "np.str_",
     "np.bytes_",
+    "np.datetime64",
 ]
 _SupportedNumpyArray = npt.NDArray[_SupportedNumpyDtype]
 _SupportedBuiltinsList = Sequence[_SupportedBuiltins]