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

snowflake/cortex/__init__.py

@@ -1,11 +1,14 @@
-from snowflake.cortex._complete import Complete
+from snowflake.cortex._classify_text import ClassifyText
+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
 from snowflake.cortex._translate import Translate
 
 __all__ = [
+    "ClassifyText",
     "Complete",
+    "CompleteOptions",
     "ExtractAnswer",
     "Sentiment",
     "Summarize",

snowflake/cortex/_classify_text.py

@@ -0,0 +1,36 @@
+from typing import List, Optional, Union
+
+from snowflake import snowpark
+from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.ml._internal import telemetry
+
+
+@telemetry.send_api_usage_telemetry(
+    project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+)
+def ClassifyText(
+    str_input: Union[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.
+
+    Args:
+        str_input: A Column of strings to classify.
+        categories: A list of candidate categories to classify the INPUT text into.
+        session: The snowpark session to use. Will be inferred by context if not specified.
+
+    Returns:
+        A column of classification responses.
+    """
+
+    return _classify_text_impl("snowflake.cortex.classify_text", str_input, categories, session=session)
+
+
+def _classify_text_impl(
+    function: str,
+    str_input: Union[str, snowpark.Column],
+    categories: Union[List[str], snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[str, snowpark.Column]:
+    return call_sql_function(function, session, str_input, categories)

snowflake/cortex/_complete.py

@@ -1,37 +1,299 @@
-from typing import Iterator, Optional, Union
+import json
+import logging
+import time
+from typing import Any, Callable, 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._sse_client import SSEClient
 from snowflake.cortex._util import (
     CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
-    call_rest_function,
-    call_sql_function,
-    process_rest_response,
+    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
+
+
+_MAX_RETRY_SECONDS = 30
+
+
+def retry(func: Callable[..., requests.Response]) -> Callable[..., requests.Response]:
+    def inner(*args: Any, **kwargs: Any) -> requests.Response:
+        deadline = cast(Optional[float], kwargs["deadline"])
+        kwargs = {key: value for key, value in kwargs.items() if key != "deadline"}
+        expRetrySeconds = 0.5
+        while True:
+            if deadline is not None and time.time() > deadline:
+                raise TimeoutError()
+            response = func(*args, **kwargs)
+            if response.status_code >= 200 and response.status_code < 300:
+                return response
+            retry_status_codes = [429, 503, 504]
+            if response.status_code not in retry_status_codes:
+                response.raise_for_status()
+            logger.debug(f"request failed with status code {response.status_code}, retrying")
+
+            # Formula: delay(i) = max(RetryAfterHeader, min(2^i, _MAX_RETRY_SECONDS)).
+            expRetrySeconds = min(2 * expRetrySeconds, _MAX_RETRY_SECONDS)
+            retrySeconds = expRetrySeconds
+            retryAfterHeader = response.headers.get("retry-after")
+            if retryAfterHeader is not None:
+                retrySeconds = max(retrySeconds, int(retryAfterHeader))
+            logger.debug(f"sleeping for {retrySeconds}s before retrying")
+            time.sleep(retrySeconds)
+
+    return inner
+
+
+@retry
+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})")
+    return requests.post(
+        url,
+        json=data,
+        headers=headers,
+        stream=stream,
+    )
+
+
+def _process_rest_response(
+    response: requests.Response,
+    stream: bool = False,
+    deadline: Optional[float] = None,
+) -> Union[str, Iterator[str]]:
+    if stream:
+        return _return_stream_response(response, deadline)
+
+    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, deadline: Optional[float]) -> Iterator[str]:
+    client = SSEClient(response)
+    for event in client.events():
+        if deadline is not None and time.time() > deadline:
+            raise TimeoutError()
+        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",
+    timeout: Optional[float] = None,
+    deadline: Optional[float] = None,
+) -> Union[str, Iterator[str], snowpark.Column]:
+    if timeout is not None and deadline is not None:
+        raise ValueError('only one of "timeout" and "deadline" must be set')
+    if timeout is not None:
+        deadline = time.time() + timeout
+    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, deadline=deadline)
+        assert response.status_code >= 200 and response.status_code < 300
+        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],
+    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,
+    timeout: Optional[float] = None,
+    deadline: Optional[float] = None,
 ) -> 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.
+        timeout (float): Timeout in seconds to retry failed REST requests.
+        deadline (float): Time in seconds since the epoch (as returned by time.time()) to retry failed REST requests.
 
     Raises:
         ValueError: If `stream` is set to True and `use_rest_api_experimental` is set to False.
@@ -39,18 +301,16 @@ def Complete(
     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)
-
-
-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=options,
+            session=session,
+            use_rest_api_experimental=use_rest_api_experimental,
+            stream=stream,
+            timeout=timeout,
+            deadline=deadline,
+        )
+    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/_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,24 +1,19 @@
-import json
-from typing import Iterator, Optional, Union, cast
-from urllib.parse import urljoin, urlparse
-
-import requests
+from typing import Dict, List, Optional, Union, cast
 
 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."""
+class SnowflakeAuthenticationException(Exception):
+    """This exception is raised when there is an issue with Snowflake's configuration."""
 
     pass
 
 
-class SnowflakeAuthenticationException(Exception):
-    """This exception is raised when the session object does not have session.connection.rest.token attribute."""
+class SnowflakeConfigurationException(Exception):
+    """This exception is raised when there is an issue with Snowflake's configuration."""
 
     pass
 
@@ -28,9 +23,10 @@ class SnowflakeAuthenticationException(Exception):
 def call_sql_function(
     function: str,
     session: Optional[snowpark.Session],
-    *args: Union[str, snowpark.Column],
+    *args: Union[str, List[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
@@ -43,17 +39,18 @@ def call_sql_function(
     )
 
 
-def _call_sql_function_column(function: str, *args: Union[str, snowpark.Column]) -> snowpark.Column:
+def _call_sql_function_column(
+    function: str, *args: Union[str, List[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],
+    *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
 ) -> str:
-    if session is None:
-        session = context.get_active_session()
+    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
@@ -67,73 +64,3 @@ 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/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