snowflake-ml-python 1.7.0__tar.gz → 1.7.2__tar.gz

{snowflake_ml_python-1.7.0 → snowflake_ml_python-1.7.2}/CHANGELOG.md

@@ -1,6 +1,45 @@
 # Release History
 
-## 1.7.0
+## 1.7.2
+
+### Bug Fixes
+
+- Model Explainability: Fix issue that explain is enabled for scikit-learn pipeline
+whose task is UNKNOWN and fails later when invoked.
+
+### Behavior Changes
+
+### New Features
+
+- Registry: Support asynchronous model inference service creation with the `block` option
+  in `ModelVersion.create_service()` set to True by default.
+
+## 1.7.1 (2024-11-05)
+
+### Bug Fixes
+
+- Registry: Null value is now allowed in the dataframe used in model signature inference. Null values will be ignored
+ and others will be used to infer the signature.
+- Registry: Pandas Extension DTypes (`pandas.StringDType()`, `pandas.BooleanDType()`, etc.) are now supported in model
+signature inference.
+- Registry: Null value is now allowed in the dataframe used to predict.
+- Data: Fix missing `snowflake.ml.data.*` module exports in wheel
+- Dataset: Fix missing `snowflake.ml.dataset.*` module exports in wheel.
+- Registry: Fix the issue that `tf_keras.Model` is not recognized as keras model when logging.
+
+### Behavior Changes
+
+### New Features
+
+- Registry: Option to `enable_monitoring` set to False by default.  This will gate access to preview features of Model Monitoring.
+- Model Monitoring: `show_model_monitors` Registry method.  This feature is still in Private Preview.
+- Registry: Support `pd.Series` in input and output data.
+- Model Monitoring: `add_monitor` Registry method.  This feature is still in Private Preview.
+- Model Monitoring: `resume` and `suspend` ModelMonitor.  This feature is still in Private Preview.
+- Model Monitoring: `get_monitor` Registry method.  This feature is still in Private Preview.
+- Model Monitoring: `delete_monitor` Registry method.  This feature is still in Private Preview.
+
+## 1.7.0 (10-22-2024)
 
 ### Behavior Change
 

{snowflake_ml_python-1.7.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.7.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: snowflake-ml-python
-Version: 1.7.0
+Version: 1.7.2
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -236,12 +236,14 @@ Requires-Dist: absl-py<2,>=0.15
 Requires-Dist: anyio<4,>=3.5.0
 Requires-Dist: cachetools<6,>=3.1.1
 Requires-Dist: cloudpickle>=2.0.0
+Requires-Dist: cryptography
 Requires-Dist: fsspec[http]<2024,>=2022.11
 Requires-Dist: importlib_resources<7,>=6.1.1
 Requires-Dist: numpy<2,>=1.23
-Requires-Dist: packaging<24,>=20.9
+Requires-Dist: packaging<25,>=20.9
 Requires-Dist: pandas<3,>=1.0.0
 Requires-Dist: pyarrow
+Requires-Dist: pyjwt<3,>=2.0.0
 Requires-Dist: pytimeparse<2,>=1.1.8
 Requires-Dist: pyyaml<7,>=6.0
 Requires-Dist: retrying<2,>=1.3.3
@@ -255,7 +257,7 @@ Requires-Dist: typing-extensions<5,>=4.1.0
 Requires-Dist: xgboost<3,>=1.7.3
 Provides-Extra: all
 Requires-Dist: catboost<2,>=1.2.0; extra == "all"
-Requires-Dist: lightgbm<5,>=3.3.5; extra == "all"
+Requires-Dist: lightgbm<5,>=4.1.0; extra == "all"
 Requires-Dist: mlflow<2.4,>=2.1.0; extra == "all"
 Requires-Dist: peft<1,>=0.5.0; extra == "all"
 Requires-Dist: sentence-transformers<3,>=2.2.2; extra == "all"
@@ -269,7 +271,7 @@ Requires-Dist: transformers<5,>=4.32.1; extra == "all"
 Provides-Extra: catboost
 Requires-Dist: catboost<2,>=1.2.0; extra == "catboost"
 Provides-Extra: lightgbm
-Requires-Dist: lightgbm<5,>=3.3.5; extra == "lightgbm"
+Requires-Dist: lightgbm<5,>=4.1.0; extra == "lightgbm"
 Provides-Extra: llm
 Requires-Dist: peft<1,>=0.5.0; extra == "llm"
 Provides-Extra: mlflow
@@ -301,7 +303,7 @@ and deployment process, and includes two key components.
 
 ### Snowpark ML Development
 
-[Snowpark ML Development](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#snowpark-ml-development)
+[Snowpark ML Development](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#ml-modeling)
 provides a collection of python APIs enabling efficient ML model development directly in Snowflake:
 
 1. Modeling API (`snowflake.ml.modeling`) for data preprocessing, feature engineering and model training in Snowflake.
@@ -315,14 +317,21 @@ their native data loader formats.
 1. FileSet API: FileSet provides a Python fsspec-compliant API for materializing data into a Snowflake internal stage
 from a query or Snowpark Dataframe along with a number of convenience APIs.
 
-### Snowpark Model Management [Public Preview]
+### Snowflake MLOps
 
-[Snowpark Model Management](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#snowpark-ml-ops) complements
-the Snowpark ML Development API, and provides model management capabilities along with integrated deployment into Snowflake.
+Snowflake MLOps contains suit of tools and objects to make ML development cycle. It complements
+the Snowpark ML Development API, and provides end to end development to deployment within Snowflake.
 Currently, the API consists of:
 
-1. Registry: A python API for managing models within Snowflake which also supports deployment of ML models into Snowflake
-as native MODEL object running with Snowflake Warehouse.
+1. [Registry](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#snowflake-model-registry): A python API
+  allows secure deployment and management of models in Snowflake, supporting models trained both inside and outside of
+  Snowflake.
+2. [Feature Store](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#snowflake-feature-store): A fully
+  integrated solution for defining, managing, storing and discovering ML features derived from your data. The
+  Snowflake Feature Store supports automated, incremental refresh from batch and streaming data sources, so that
+  feature pipelines need be defined only once to be continuously updated with new data.
+3. [Datasets](https://docs.snowflake.com/developer-guide/snowflake-ml/overview#snowflake-datasets): Dataset provide an
+  immutable, versioned snapshot of your data suitable for ingestion by your machine learning models.
 
 ## Getting started
 
@@ -370,9 +379,64 @@ conda install \
 Note that until a `snowflake-ml-python` package version is available in the official Snowflake conda channel, there may
 be compatibility issues. Server-side functionality that `snowflake-ml-python` depends on may not yet be released.
 
+### Verifying the package
+
+1. Install cosign.
+  This example is using golang installation: [installing-cosign-with-go](https://edu.chainguard.dev/open-source/sigstore/cosign/how-to-install-cosign/#installing-cosign-with-go).
+1. Download the file from the repository like [pypi](https://pypi.org/project/snowflake-ml-python/#files).
+1. Download the signature files from the [release tag](https://github.com/snowflakedb/snowflake-ml-python/releases/tag/1.7.0).
+1. Verify signature on projects signed using Jenkins job:
+
+   ```sh
+   cosign verify-blob snowflake_ml_python-1.7.0.tar.gz --key snowflake-ml-python-1.7.0.pub --signature resources.linux.snowflake_ml_python-1.7.0.tar.gz.sig
+
+   cosign verify-blob snowflake_ml_python-1.7.0.tar.gz --key snowflake-ml-python-1.7.0.pub --signature resources.linux.snowflake_ml_python-1.7.0
+   ```
+
+NOTE: Version 1.7.0 is used as example here. Please choose the the latest version.
+
 # Release History
 
-## 1.7.0
+## 1.7.2
+
+### Bug Fixes
+
+- Model Explainability: Fix issue that explain is enabled for scikit-learn pipeline
+whose task is UNKNOWN and fails later when invoked.
+
+### Behavior Changes
+
+### New Features
+
+- Registry: Support asynchronous model inference service creation with the `block` option
+  in `ModelVersion.create_service()` set to True by default.
+
+## 1.7.1 (2024-11-05)
+
+### Bug Fixes
+
+- Registry: Null value is now allowed in the dataframe used in model signature inference. Null values will be ignored
+ and others will be used to infer the signature.
+- Registry: Pandas Extension DTypes (`pandas.StringDType()`, `pandas.BooleanDType()`, etc.) are now supported in model
+signature inference.
+- Registry: Null value is now allowed in the dataframe used to predict.
+- Data: Fix missing `snowflake.ml.data.*` module exports in wheel
+- Dataset: Fix missing `snowflake.ml.dataset.*` module exports in wheel.
+- Registry: Fix the issue that `tf_keras.Model` is not recognized as keras model when logging.
+
+### Behavior Changes
+
+### New Features
+
+- Registry: Option to `enable_monitoring` set to False by default.  This will gate access to preview features of Model Monitoring.
+- Model Monitoring: `show_model_monitors` Registry method.  This feature is still in Private Preview.
+- Registry: Support `pd.Series` in input and output data.
+- Model Monitoring: `add_monitor` Registry method.  This feature is still in Private Preview.
+- Model Monitoring: `resume` and `suspend` ModelMonitor.  This feature is still in Private Preview.
+- Model Monitoring: `get_monitor` Registry method.  This feature is still in Private Preview.
+- Model Monitoring: `delete_monitor` Registry method.  This feature is still in Private Preview.
+
+## 1.7.0 (10-22-2024)
 
 ### Behavior Change
 

{snowflake_ml_python-1.7.0 → snowflake_ml_python-1.7.2}/README.md

@@ -12,7 +12,7 @@ and deployment process, and includes two key components.
 
 ### Snowpark ML Development
 
-[Snowpark ML Development](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#snowpark-ml-development)
+[Snowpark ML Development](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#ml-modeling)
 provides a collection of python APIs enabling efficient ML model development directly in Snowflake:
 
 1. Modeling API (`snowflake.ml.modeling`) for data preprocessing, feature engineering and model training in Snowflake.
@@ -26,14 +26,21 @@ their native data loader formats.
 1. FileSet API: FileSet provides a Python fsspec-compliant API for materializing data into a Snowflake internal stage
 from a query or Snowpark Dataframe along with a number of convenience APIs.
 
-### Snowpark Model Management [Public Preview]
+### Snowflake MLOps
 
-[Snowpark Model Management](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#snowpark-ml-ops) complements
-the Snowpark ML Development API, and provides model management capabilities along with integrated deployment into Snowflake.
+Snowflake MLOps contains suit of tools and objects to make ML development cycle. It complements
+the Snowpark ML Development API, and provides end to end development to deployment within Snowflake.
 Currently, the API consists of:
 
-1. Registry: A python API for managing models within Snowflake which also supports deployment of ML models into Snowflake
-as native MODEL object running with Snowflake Warehouse.
+1. [Registry](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#snowflake-model-registry): A python API
+  allows secure deployment and management of models in Snowflake, supporting models trained both inside and outside of
+  Snowflake.
+2. [Feature Store](https://docs.snowflake.com/en/developer-guide/snowpark-ml/index#snowflake-feature-store): A fully
+  integrated solution for defining, managing, storing and discovering ML features derived from your data. The
+  Snowflake Feature Store supports automated, incremental refresh from batch and streaming data sources, so that
+  feature pipelines need be defined only once to be continuously updated with new data.
+3. [Datasets](https://docs.snowflake.com/developer-guide/snowflake-ml/overview#snowflake-datasets): Dataset provide an
+  immutable, versioned snapshot of your data suitable for ingestion by your machine learning models.
 
 ## Getting started
 
@@ -80,3 +87,19 @@ conda install \
 
 Note that until a `snowflake-ml-python` package version is available in the official Snowflake conda channel, there may
 be compatibility issues. Server-side functionality that `snowflake-ml-python` depends on may not yet be released.
+
+### Verifying the package
+
+1. Install cosign.
+  This example is using golang installation: [installing-cosign-with-go](https://edu.chainguard.dev/open-source/sigstore/cosign/how-to-install-cosign/#installing-cosign-with-go).
+1. Download the file from the repository like [pypi](https://pypi.org/project/snowflake-ml-python/#files).
+1. Download the signature files from the [release tag](https://github.com/snowflakedb/snowflake-ml-python/releases/tag/1.7.0).
+1. Verify signature on projects signed using Jenkins job:
+
+   ```sh
+   cosign verify-blob snowflake_ml_python-1.7.0.tar.gz --key snowflake-ml-python-1.7.0.pub --signature resources.linux.snowflake_ml_python-1.7.0.tar.gz.sig
+
+   cosign verify-blob snowflake_ml_python-1.7.0.tar.gz --key snowflake-ml-python-1.7.0.pub --signature resources.linux.snowflake_ml_python-1.7.0
+   ```
+
+NOTE: Version 1.7.0 is used as example here. Please choose the the latest version.

{snowflake_ml_python-1.7.0 → snowflake_ml_python-1.7.2}/pyproject.toml

@@ -8,7 +8,7 @@ description = "The machine learning client library that is used for interacting
 classifiers = [ "Development Status :: 5 - Production/Stable", "Environment :: Console", "Environment :: Other Environment", "Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Information Technology", "Intended Audience :: System Administrators", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Topic :: Database", "Topic :: Software Development", "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Application Frameworks", "Topic :: Software Development :: Libraries :: Python Modules", "Topic :: Scientific/Engineering :: Information Analysis",]
 requires-python = ">=3.9, <3.12"
 dynamic = [ "version", "readme",]
-dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<4", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "fsspec[http]>=2022.11,<2024", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<2", "packaging>=20.9,<24", "pandas>=1.0.0,<3", "pyarrow", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2022.11,<2024", "scikit-learn>=1.4,<1.6", "scipy>=1.9,<2", "snowflake-connector-python[pandas]>=3.5.0,<4", "snowflake-snowpark-python>=1.17.0,<2", "sqlparse>=0.4,<1", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<3",]
+dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<4", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "cryptography", "fsspec[http]>=2022.11,<2024", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<2", "packaging>=20.9,<25", "pandas>=1.0.0,<3", "pyarrow", "pyjwt>=2.0.0, <3", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2022.11,<2024", "scikit-learn>=1.4,<1.6", "scipy>=1.9,<2", "snowflake-connector-python[pandas]>=3.5.0,<4", "snowflake-snowpark-python>=1.17.0,<2", "sqlparse>=0.4,<1", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<3",]
 [[project.authors]]
 name = "Snowflake, Inc"
 email = "support@snowflake.com"
@@ -24,9 +24,9 @@ Issues = "https://github.com/snowflakedb/snowflake-ml-python/issues"
 Changelog = "https://github.com/snowflakedb/snowflake-ml-python/blob/master/CHANGELOG.md"
 
 [project.optional-dependencies]
-all = [ "catboost>=1.2.0, <2", "lightgbm>=3.3.5,<5", "mlflow>=2.1.0,<2.4", "peft>=0.5.0,<1", "sentence-transformers>=2.2.2,<3", "sentencepiece>=0.1.95,<1", "shap>=0.46.0,<1", "tensorflow>=2.10,<3", "tokenizers>=0.10,<1", "torch>=2.0.1,<2.3.0", "torchdata>=0.4,<1", "transformers>=4.32.1,<5",]
+all = [ "catboost>=1.2.0, <2", "lightgbm>=4.1.0, <5", "mlflow>=2.1.0,<2.4", "peft>=0.5.0,<1", "sentence-transformers>=2.2.2,<3", "sentencepiece>=0.1.95,<1", "shap>=0.46.0,<1", "tensorflow>=2.10,<3", "tokenizers>=0.10,<1", "torch>=2.0.1,<2.3.0", "torchdata>=0.4,<1", "transformers>=4.32.1,<5",]
 catboost = [ "catboost>=1.2.0, <2",]
-lightgbm = [ "lightgbm>=3.3.5,<5",]
+lightgbm = [ "lightgbm>=4.1.0, <5",]
 llm = [ "peft>=0.5.0,<1",]
 mlflow = [ "mlflow>=2.1.0,<2.4",]
 shap = [ "shap>=0.46.0,<1",]

{snowflake_ml_python-1.7.0 → snowflake_ml_python-1.7.2}/snowflake/cortex/__init__.py

@@ -3,6 +3,7 @@ from snowflake.cortex._complete import Complete, CompleteOptions
 from snowflake.cortex._embed_text_768 import EmbedText768
 from snowflake.cortex._embed_text_1024 import EmbedText1024
 from snowflake.cortex._extract_answer import ExtractAnswer
+from snowflake.cortex._finetune import Finetune, FinetuneJob, FinetuneStatus
 from snowflake.cortex._sentiment import Sentiment
 from snowflake.cortex._summarize import Summarize
 from snowflake.cortex._translate import Translate
@@ -14,6 +15,9 @@ __all__ = [
     "EmbedText768",
     "EmbedText1024",
     "ExtractAnswer",
+    "Finetune",
+    "FinetuneJob",
+    "FinetuneStatus",
     "Sentiment",
     "Summarize",
     "Translate",

{snowflake_ml_python-1.7.0 → snowflake_ml_python-1.7.2}/snowflake/cortex/_complete.py

@@ -1,7 +1,8 @@
 import json
 import logging
 import time
-from typing import Any, Callable, Iterator, List, Optional, TypedDict, Union, cast
+from io import BytesIO
+from typing import Any, Callable, Dict, Iterator, List, Optional, TypedDict, Union, cast
 from urllib.parse import urlunparse
 
 import requests
@@ -16,8 +17,10 @@ from snowflake.cortex._util import (
 )
 from snowflake.ml._internal import telemetry
 from snowflake.snowpark import context, functions
+from snowflake.snowpark._internal.utils import is_in_stored_procedure
 
 logger = logging.getLogger(__name__)
+_REST_COMPLETE_URL = "/api/v2/cortex/inference:complete"
 
 
 class ConversationMessage(TypedDict):
@@ -84,6 +87,76 @@ def retry(func: Callable[..., requests.Response]) -> Callable[..., requests.Resp
     return inner
 
 
+def _make_common_request_headers() -> Dict[str, str]:
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json, text/event-stream",
+    }
+    return headers
+
+
+def _make_request_body(
+    model: str,
+    prompt: Union[str, List[ConversationMessage]],
+    options: Optional[CompleteOptions] = None,
+) -> Dict[str, Any]:
+    data = {
+        "model": model,
+        "stream": True,
+    }
+    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"]
+    return data
+
+
+# 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:
+    response = requests.Response()
+    response.status_code = int(raw_resp["status"])
+    response.headers = raw_resp["headers"]
+
+    data = raw_resp["content"]
+    data = json.loads(data)
+    # Convert the dictionary to a string format that resembles the SSE event format
+    # For example, if the dict is {'event': 'message', 'data': 'your data'}, it should be formatted like this:
+    sse_format_data = ""
+    for event in data:
+        event_type = event.get("event", "message")
+        event_data = event.get("data", "")
+        event_data = json.dumps(event_data)
+        sse_format_data += f"event: {event_type}\ndata: {event_data}\n\n"  # Add each event with new lines
+
+    response.raw = BytesIO(sse_format_data.encode("utf-8"))
+    return response
+
+
+@retry
+def _call_complete_xp(
+    model: str,
+    prompt: Union[str, List[ConversationMessage]],
+    options: Optional[CompleteOptions] = None,
+    deadline: Optional[float] = None,
+) -> requests.Response:
+    headers = _make_common_request_headers()
+    body = _make_request_body(model, prompt, options)
+    import _snowflake
+
+    raw_resp = _snowflake.send_snow_api_request("POST", _REST_COMPLETE_URL, {}, headers, body, {}, deadline)
+    return _xp_dict_to_response(raw_resp)
+
+
 @retry
 def _call_complete_rest(
     model: str,
@@ -110,36 +183,16 @@ def _call_complete_rest(
     scheme = "https"
     if hasattr(session.connection, "scheme"):
         scheme = session.connection.scheme
-    url = urlunparse((scheme, session.connection.host, "api/v2/cortex/inference:complete", "", "", ""))
+    url = urlunparse((scheme, session.connection.host, _REST_COMPLETE_URL, "", "", ""))
 
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f'Snowflake Token="{session.connection.rest.token}"',
-        "Accept": "application/json, text/event-stream",
-    }
-
-    data = {
-        "model": model,
-        "stream": True,
-    }
-    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"]
+    headers = _make_common_request_headers()
+    headers["Authorization"] = f'Snowflake Token="{session.connection.rest.token}"'
 
+    body = _make_request_body(model, prompt, options)
     logger.debug(f"making POST request to {url} (model={model})")
     return requests.post(
         url,
-        json=data,
+        json=body,
         headers=headers,
         stream=True,
     )
@@ -164,49 +217,24 @@ def _complete_call_sql_function_snowpark(
     return cast(snowpark.Column, functions.builtin(function)(*args))
 
 
-def _complete_call_sql_function_immediate(
-    function: str,
+def _complete_non_streaming_immediate(
     model: str,
     prompt: Union[str, List[ConversationMessage]],
     options: Optional[CompleteOptions],
-    session: Optional[snowpark.Session],
+    session: Optional[snowpark.Session] = None,
+    deadline: Optional[float] = None,
 ) -> 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."""
-        )
+    response = _complete_rest(model=model, prompt=prompt, options=options, session=session, deadline=deadline)
+    return "".join(response)
 
-    # 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(
+def _complete_non_streaming_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],
+    session: Optional[snowpark.Session] = None,
+    deadline: Optional[float] = None,
 ) -> Union[str, snowpark.Column]:
     if isinstance(prompt, snowpark.Column):
         if options is not None:
@@ -217,7 +245,24 @@ def _complete_sql_impl(
         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)
+    return _complete_non_streaming_immediate(
+        model=model, prompt=prompt, options=options, session=session, deadline=deadline
+    )
+
+
+def _complete_rest(
+    model: str,
+    prompt: Union[str, List[ConversationMessage]],
+    options: Optional[CompleteOptions] = None,
+    session: Optional[snowpark.Session] = None,
+    deadline: Optional[float] = None,
+) -> Iterator[str]:
+    if is_in_stored_procedure():  # type: ignore[no-untyped-call]
+        response = _call_complete_xp(model=model, prompt=prompt, options=options, deadline=deadline)
+    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)
 
 
 def _complete_impl(
@@ -239,10 +284,8 @@ def _complete_impl(
             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, deadline=deadline)
-        assert response.status_code >= 200 and response.status_code < 300
-        return _return_stream_response(response, deadline)
-    return _complete_sql_impl(function, model, prompt, options, session)
+        return _complete_rest(model=model, prompt=prompt, options=options, session=session, deadline=deadline)
+    return _complete_non_streaming_impl(function, model, prompt, options, session, deadline)
 
 
 @telemetry.send_api_usage_telemetry(