Snowflake provider support: added embeddings, PAT, account_id (#15372)

* snowflake support PAT, account_id and embeddings

* format

* test embeddings

* format

* complete test

---------

Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>

litellm/__init__.py

@@ -1292,6 +1292,7 @@ from .llms.ovhcloud.chat.transformation import OVHCloudChatConfig
 from .llms.ovhcloud.embedding.transformation import OVHCloudEmbeddingConfig
 from .llms.cometapi.embed.transformation import CometAPIEmbeddingConfig
 from .llms.lemonade.chat.transformation import LemonadeChatConfig
+from .llms.snowflake.embedding.transformation import SnowflakeEmbeddingConfig
 from .main import *  # type: ignore
 from .integrations import *
 from .llms.custom_httpx.async_client_cleanup import close_litellm_async_clients

litellm/litellm_core_utils/get_llm_provider_logic.py

@@ -692,12 +692,12 @@ def _get_openai_compatible_provider_info(  # noqa: PLR0915
         )  # type: ignore
         dynamic_api_key = api_key or get_secret_str("NOVITA_API_KEY")
     elif custom_llm_provider == "snowflake":
-        api_base = (
-            api_base
-            or get_secret_str("SNOWFLAKE_API_BASE")
-            or f"https://{get_secret('SNOWFLAKE_ACCOUNT_ID')}.snowflakecomputing.com/api/v2/cortex/inference:complete"
-        )  # type: ignore
-        dynamic_api_key = api_key or get_secret_str("SNOWFLAKE_JWT")
+        (
+            api_base,
+            dynamic_api_key,
+        ) = litellm.SnowflakeConfig()._get_openai_compatible_provider_info(
+            api_base, api_key
+        )
     elif custom_llm_provider == "gradient_ai":
         (
             api_base,

litellm/llms/snowflake/chat/transformation.py

@@ -21,18 +21,7 @@ else:
     LiteLLMLoggingObj = Any
 
 
-class SnowflakeConfig(OpenAIGPTConfig):
-    """
-    Reference: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api
-
-    Snowflake Cortex LLM REST API supports function calling with specific models (e.g., Claude 3.5 Sonnet).
-    This config handles transformation between OpenAI format and Snowflake's tool_spec format.
-    """
-
-    @classmethod
-    def get_config(cls):
-        return super().get_config()
-
+class SnowflakeBaseConfig:
     def get_supported_openai_params(self, model: str) -> List[str]:
         return [
             "temperature",
@@ -67,6 +56,82 @@ class SnowflakeConfig(OpenAIGPTConfig):
                 optional_params[param] = value
         return optional_params
 
+    def _get_api_base(self, api_base, optional_params):
+        if not api_base:
+            if "account_id" in optional_params:
+                account_id = optional_params.pop("account_id")
+            else:
+                account_id = get_secret_str("SNOWFLAKE_ACCOUNT_ID")
+            if account_id is None:
+                raise ValueError("Missing snowflake account_id")
+            api_base = f"https://{account_id}.snowflakecomputing.com/api/v2"
+
+        api_base = api_base.rstrip("/")
+        if not api_base.endswith("/api/v2"):
+            api_base += "/api/v2"
+        return api_base
+
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+    ) -> dict:
+        """
+        Return headers to use for Snowflake completion request
+
+        Snowflake REST API Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api#api-reference
+        Expected headers:
+        {
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+            "Authorization": "Bearer " + <JWT>,
+            "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT"
+        }
+        """
+
+        auth_type = "KEYPAIR_JWT"
+
+        if "pat_key" in optional_params:
+            api_key = optional_params.pop("pat_key")
+            auth_type = "PROGRAMMATIC_ACCESS_TOKEN"
+
+        if api_key is None:
+            raise ValueError("Missing Snowflake JWT key")
+
+        headers.update(
+            {
+                "Content-Type": "application/json",
+                "Accept": "application/json",
+                "Authorization": "Bearer " + api_key,
+                "X-Snowflake-Authorization-Token-Type": auth_type,
+            }
+        )
+        return headers
+
+    def _get_openai_compatible_provider_info(
+        self, api_base: Optional[str], api_key: Optional[str]
+    ) -> Tuple[Optional[str], Optional[str]]:
+        dynamic_api_key = api_key or get_secret_str("SNOWFLAKE_JWT")
+        return api_base, dynamic_api_key
+
+
+class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
+    """
+    Reference: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api
+
+    Snowflake Cortex LLM REST API supports function calling with specific models (e.g., Claude 3.5 Sonnet).
+    This config handles transformation between OpenAI format and Snowflake's tool_spec format.
+    """
+
+    @classmethod
+    def get_config(cls):
+        return super().get_config()
+
     def _transform_tool_calls_from_snowflake_to_openai(
         self, content_list: List[Dict[str, Any]]
     ) -> Tuple[str, Optional[List[ChatCompletionMessageToolCall]]]:
@@ -169,53 +234,6 @@ class SnowflakeConfig(OpenAIGPTConfig):
             returned_response._hidden_params["model"] = model
         return returned_response
 
-    def validate_environment(
-        self,
-        headers: dict,
-        model: str,
-        messages: List[AllMessageValues],
-        optional_params: dict,
-        litellm_params: dict,
-        api_key: Optional[str] = None,
-        api_base: Optional[str] = None,
-    ) -> dict:
-        """
-        Return headers to use for Snowflake completion request
-
-        Snowflake REST API Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api#api-reference
-        Expected headers:
-        {
-            "Content-Type": "application/json",
-            "Accept": "application/json",
-            "Authorization": "Bearer " + <JWT>,
-            "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT"
-        }
-        """
-
-        if api_key is None:
-            raise ValueError("Missing Snowflake JWT key")
-
-        headers.update(
-            {
-                "Content-Type": "application/json",
-                "Accept": "application/json",
-                "Authorization": "Bearer " + api_key,
-                "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT",
-            }
-        )
-        return headers
-
-    def _get_openai_compatible_provider_info(
-        self, api_base: Optional[str], api_key: Optional[str]
-    ) -> Tuple[Optional[str], Optional[str]]:
-        api_base = (
-            api_base
-            or f"""https://{get_secret_str("SNOWFLAKE_ACCOUNT_ID")}.snowflakecomputing.com/api/v2/cortex/inference:complete"""
-            or get_secret_str("SNOWFLAKE_API_BASE")
-        )
-        dynamic_api_key = api_key or get_secret_str("SNOWFLAKE_JWT")
-        return api_base, dynamic_api_key
-
     def get_complete_url(
         self,
         api_base: Optional[str],
@@ -228,10 +246,10 @@ class SnowflakeConfig(OpenAIGPTConfig):
         """
         If api_base is not provided, use the default DeepSeek /chat/completions endpoint.
         """
-        if not api_base:
-            api_base = f"""https://{get_secret_str("SNOWFLAKE_ACCOUNT_ID")}.snowflakecomputing.com/api/v2/cortex/inference:complete"""
 
-        return api_base
+        api_base = self._get_api_base(api_base, optional_params)
+
+        return f"{api_base}/cortex/inference:complete"
 
     def _transform_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
         """
@@ -279,9 +297,7 @@ class SnowflakeConfig(OpenAIGPTConfig):
                 }
                 # Add description if present
                 if "description" in function:
-                    snowflake_tool["tool_spec"]["description"] = function[
-                        "description"
-                    ]
+                    snowflake_tool["tool_spec"]["description"] = function["description"]
 
                 snowflake_tools.append(snowflake_tool)
 

litellm/llms/snowflake/embedding/transformation.py

@@ -0,0 +1,72 @@
+from typing import Optional, Union
+
+import httpx
+
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
+from litellm.types.llms.openai import AllEmbeddingInputValues
+from litellm.types.utils import EmbeddingResponse
+
+
+from ..chat.transformation import SnowflakeBaseConfig
+
+from ..utils import SnowflakeException
+
+
+class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig):
+    """
+    source: https://docs.snowflake.com/developer-guide/snowflake-rest-api/reference/cortex-embed
+    """
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        model: str,
+        optional_params: dict,
+        litellm_params: dict,
+        stream: Optional[bool] = None,
+    ) -> str:
+        api_base = self._get_api_base(api_base, optional_params)
+
+        return f"{api_base}/cortex/inference:embed"
+
+    def transform_embedding_request(
+        self,
+        model: str,
+        input: AllEmbeddingInputValues,
+        optional_params: dict,
+        headers: dict,
+    ) -> dict:
+        return {"text": input, "model": model, **optional_params}
+
+    def transform_embedding_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: EmbeddingResponse,
+        logging_obj: LiteLLMLoggingObj,
+        api_key: Optional[str],
+        request_data: dict,
+        optional_params: dict,
+        litellm_params: dict,
+    ) -> EmbeddingResponse:
+        response_json = raw_response.json()
+        # convert embeddings to 1d array
+        for item in response_json["data"]:
+            item["embedding"] = item["embedding"][0]
+        returned_response = EmbeddingResponse(**response_json)
+
+        returned_response.model = "snowflake/" + (returned_response.model or "")
+
+        if model is not None:
+            returned_response._hidden_params["model"] = model
+        return returned_response
+
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
+    ) -> BaseLLMException:
+        return SnowflakeException(
+            message=error_message, status_code=status_code, headers=headers
+        )

litellm/llms/snowflake/utils.py

@@ -0,0 +1,7 @@
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+
+
+class SnowflakeException(BaseLLMException):
+    """Snowflake AI Endpoints exception handling class"""
+
+    pass

litellm/main.py

@@ -4779,6 +4779,22 @@ def embedding(  # noqa: PLR0915
                 print_verbose=print_verbose,
                 litellm_params=litellm_params_dict,
             )
+        elif custom_llm_provider == "snowflake":
+            api_key = api_key or get_secret_str("SNOWFLAKE_JWT")
+            response = base_llm_http_handler.embedding(
+                model=model,
+                input=input,
+                custom_llm_provider=custom_llm_provider,
+                api_base=api_base,
+                api_key=api_key,
+                logging_obj=logging,
+                timeout=timeout,
+                model_response=EmbeddingResponse(),
+                optional_params=optional_params,
+                client=client,
+                aembedding=aembedding,
+                litellm_params={},
+            )
         else:
             raise LiteLLMUnknownProvider(
                 model=model, custom_llm_provider=custom_llm_provider

litellm/utils.py

@@ -7213,6 +7213,8 @@ class ProviderConfigManager:
             return VolcEngineEmbeddingConfig()
         elif litellm.LlmProviders.OVHCLOUD == provider:
             return litellm.OVHCloudEmbeddingConfig()
+        elif litellm.LlmProviders.SNOWFLAKE == provider:
+            return litellm.SnowflakeEmbeddingConfig()
         elif litellm.LlmProviders.COMETAPI == provider:
             return litellm.CometAPIEmbeddingConfig()
         return None

tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py

@@ -3,11 +3,15 @@ Unit tests for Snowflake chat transformation
 Tests tool calling request/response transformations
 """
 
+import os
+import copy
 import json
+
+import pytest
+from unittest.mock import patch
 from unittest.mock import MagicMock
 
 import httpx
-import pytest
 
 import litellm
 from litellm.llms.snowflake.chat.transformation import SnowflakeConfig
@@ -66,7 +70,10 @@ class TestSnowflakeToolTransformation:
         assert "tool_spec" in snowflake_tool
         assert snowflake_tool["tool_spec"]["type"] == "generic"
         assert snowflake_tool["tool_spec"]["name"] == "get_weather"
-        assert snowflake_tool["tool_spec"]["description"] == "Get the current weather in a given location"
+        assert (
+            snowflake_tool["tool_spec"]["description"]
+            == "Get the current weather in a given location"
+        )
         assert "input_schema" in snowflake_tool["tool_spec"]
         assert snowflake_tool["tool_spec"]["input_schema"]["type"] == "object"
         assert "location" in snowflake_tool["tool_spec"]["input_schema"]["properties"]
@@ -93,7 +100,9 @@ class TestSnowflakeToolTransformation:
         # Verify tool_choice was transformed to Snowflake format
         assert "tool_choice" in transformed_request
         assert transformed_request["tool_choice"]["type"] == "tool"
-        assert transformed_request["tool_choice"]["name"] == ["get_weather"]  # Array format
+        assert transformed_request["tool_choice"]["name"] == [
+            "get_weather"
+        ]  # Array format
 
     def test_transform_request_with_string_tool_choice(self):
         """
@@ -132,7 +141,10 @@ class TestSnowflakeToolTransformation:
                                 "tool_use": {
                                     "tool_use_id": "tooluse_abc123",
                                     "name": "get_weather",
-                                    "input": {"location": "Paris, France", "unit": "celsius"},
+                                    "input": {
+                                        "location": "Paris, France",
+                                        "unit": "celsius",
+                                    },
                                 },
                             },
                         ]
@@ -207,7 +219,10 @@ class TestSnowflakeToolTransformation:
                 {
                     "message": {
                         "content_list": [
-                            {"type": "text", "text": "Let me check the weather for you. "},
+                            {
+                                "type": "text",
+                                "text": "Let me check the weather for you. ",
+                            },
                             {
                                 "type": "tool_use",
                                 "tool_use": {
@@ -300,7 +315,10 @@ class TestSnowflakeToolTransformation:
 
         # Verify standard response works
         assert isinstance(result, ModelResponse)
-        assert result.choices[0].message.content == "Hello! I'm doing well, thank you for asking."
+        assert (
+            result.choices[0].message.content
+            == "Hello! I'm doing well, thank you for asking."
+        )
 
     def test_get_supported_openai_params_includes_tools(self):
         """
@@ -313,3 +331,120 @@ class TestSnowflakeToolTransformation:
         assert "tool_choice" in supported_params
         assert "temperature" in supported_params
         assert "max_tokens" in supported_params
+
+
+class TestSnowFlakeCompletion:
+    model_name = "mistral"
+
+    messages = [
+        {"role": "system", "content": "hi"},
+        {"role": "user", "content": "the capital of France"},
+    ]
+
+    response = {
+        "choices": [
+            {
+                "message": {
+                    "content": "Paris",
+                    "content_list": [{"type": "text", "text": "Paris"}],
+                }
+            }
+        ],
+        "usage": {"prompt_tokens": 16, "completion_tokens": 18, "total_tokens": 34},
+    }
+
+    @patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+    def test_snowflake_no_account_id(self, mock_post):
+        mock_post().json.return_value = copy.deepcopy(self.response)
+
+        with pytest.raises(Exception):
+            litellm.completion(
+                f"snowflake/{self.model_name}",
+                messages=self.messages,
+                api_key="0000",
+            )
+
+    @patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+    def test_snowflake_no_api_key(self, mock_post):
+        mock_post().json.return_value = copy.deepcopy(self.response)
+
+        with pytest.raises(Exception):
+            litellm.completion(
+                f"snowflake/{self.model_name}",
+                messages=self.messages,
+                account_id="AAAA-BBBB",
+            )
+
+    @patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+    def test_snowflake_jwt_account_id(self, mock_post):
+        mock_post().json.return_value = copy.deepcopy(self.response)
+
+        response = litellm.completion(
+            f"snowflake/{self.model_name}",
+            messages=self.messages,
+            api_key="00000",
+            account_id="AAAA-BBBB",
+        )
+        assert len(response.choices) == 1
+        assert response.choices[0]["message"].content == "Paris"
+
+        # check request
+        post_kwargs = mock_post.call_args_list[-1][1]
+        body = json.loads(post_kwargs["data"])
+        assert body["model"] == self.model_name
+        assert "the capital of France" in str(body["messages"])
+
+        # JWT key was used
+        assert "00000" in post_kwargs["headers"]["Authorization"]
+        # account id was used
+        assert "AAAA-BBBB" in post_kwargs["url"]
+        # is completion
+        post_kwargs["url"].endswith("cortex/inference:complete")
+
+    @patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+    def test_snowflake_pat_key_account_id(self, mock_post):
+        mock_post().json.return_value = copy.deepcopy(self.response)
+
+        response = litellm.completion(
+            f"snowflake/{self.model_name}",
+            messages=self.messages,
+            pat_key="xxxxx",
+            account_id="AAAA-BBBB",
+        )
+        assert len(response.choices) == 1
+        assert response.choices[0]["message"].content == "Paris"
+
+        # PAT key was used
+        post_kwargs = mock_post.call_args_list[-1][1]
+        assert "xxxxx" in post_kwargs["headers"]["Authorization"]
+        assert (
+            post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"]
+            == "PROGRAMMATIC_ACCESS_TOKEN"
+        )
+
+        # account id was used
+        assert "AAAA-BBBB" in post_kwargs["url"]
+
+    @patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+    def test_snowflake_env(self, mock_post):
+        mock_post().json.return_value = copy.deepcopy(self.response)
+
+        os.environ["SNOWFLAKE_ACCOUNT_ID"] = "AAAA-BBBB"
+        os.environ["SNOWFLAKE_JWT"] = "00000"
+
+        response = litellm.completion(
+            f"snowflake/{self.model_name}",
+            messages=self.messages,
+        )
+
+        assert len(response.choices) == 1
+        assert response.choices[0]["message"].content == "Paris"
+
+        # JWT key was used
+        post_kwargs = mock_post.call_args_list[-1][1]
+        assert "00000" in post_kwargs["headers"]["Authorization"]
+        # account id was used
+        assert "AAAA-BBBB" in post_kwargs["url"]
+
+        os.environ.pop("SNOWFLAKE_ACCOUNT_ID", None)
+        os.environ.pop("SNOWFLAKE_JWT", None)

tests/test_litellm/llms/snowflake/embedding/test_snowflake_embedding.py

@@ -0,0 +1,121 @@
+import os
+import json
+import copy
+
+import pytest
+from unittest.mock import patch
+
+import litellm
+
+model_name = "snowflake-arctic-embed"
+
+embed_response = {
+    "object": "list",
+    "data": [
+        {
+            "object": "embedding",
+            "embedding": [[0.1, 0.2, 0.3]],
+            "index": 0,
+        }
+    ],
+    "model": model_name,
+    "usage": {"total_tokens": 4},
+}
+
+
+@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+def test_snowflake_no_account_id(mock_post):
+    mock_post().json.return_value = copy.deepcopy(embed_response)
+
+    with pytest.raises(Exception):
+        litellm.embedding(
+            f"snowflake/{model_name}",
+            input=["test"],
+            api_key="0000",
+        )
+
+
+@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+def test_snowflake_no_api_key(mock_post):
+    mock_post().json.return_value = copy.deepcopy(embed_response)
+
+    with pytest.raises(Exception):
+        litellm.embedding(
+            f"snowflake/{model_name}",
+            input=["test"],
+            account_id="AAAA-BBBB",
+        )
+
+
+@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+def test_snowflake_jwt_account_id(mock_post):
+    mock_post().json.return_value = copy.deepcopy(embed_response)
+
+    response = litellm.embedding(
+        f"snowflake/{model_name}",
+        input=["document"],
+        api_key="00000",
+        account_id="AAAA-BBBB",
+    )
+    assert len(response.data) == 1
+    assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
+
+    # check request
+    post_kwargs = mock_post.call_args_list[-1][1]
+    body = json.loads(post_kwargs["data"])
+    assert body["model"] == model_name
+    assert body["text"][0] == "document"
+
+    # JWT key was used
+    assert "00000" in post_kwargs["headers"]["Authorization"]
+    # account id was used
+    assert "AAAA-BBBB" in post_kwargs["url"]
+    # is embedding
+    post_kwargs["url"].endswith("cortex/inference:embed")
+
+
+@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+def test_snowflake_pat_key_account_id(mock_post):
+    mock_post().json.return_value = copy.deepcopy(embed_response)
+
+    response = litellm.embedding(
+        f"snowflake/{model_name}",
+        input=["document"],
+        pat_key="xxxxx",
+        account_id="AAAA-BBBB",
+    )
+    assert len(response.data) == 1
+    assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
+
+    # PAT key was used
+    post_kwargs = mock_post.call_args_list[-1][1]
+    assert "xxxxx" in post_kwargs["headers"]["Authorization"]
+    assert (
+        post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"]
+        == "PROGRAMMATIC_ACCESS_TOKEN"
+    )
+
+    # account id was used
+    assert "AAAA-BBBB" in post_kwargs["url"]
+
+
+@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
+def test_snowflake_env(mock_post):
+    mock_post().json.return_value = copy.deepcopy(embed_response)
+
+    os.environ["SNOWFLAKE_ACCOUNT_ID"] = "AAAA-BBBB"
+    os.environ["SNOWFLAKE_JWT"] = "00000"
+
+    response = litellm.embedding(f"snowflake/{model_name}", input=["document"])
+
+    assert len(response.data) == 1
+    assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
+
+    # JWT key was used
+    post_kwargs = mock_post.call_args_list[-1][1]
+    assert "00000" in post_kwargs["headers"]["Authorization"]
+    # account id was used
+    assert "AAAA-BBBB" in post_kwargs["url"]
+
+    os.environ.pop("SNOWFLAKE_ACCOUNT_ID", None)
+    os.environ.pop("SNOWFLAKE_JWT", None)