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

* added oauth mcp to docs

* added azure ai/grok-4 model family

* Revert "added oauth mcp to docs"

This reverts commit 950b7cef44.

* fix: arize ui integration

* need to remove a file

This reverts commit d6c877b73a.

* snowflake support PAT, account_id and embeddings

* format

* test embeddings

* format

* complete test

* fix: add arize from ui

* updated clarifai functions to openai compatible (#15615)

* fix: npm build errors

* update tests

* SnowflakeBaseConfig moved to utils

* rename pat_key => api_ke

* key_type=PAT => 'pat/key'

* fix if api_key is None

* doc update

* doc update

---------

Co-authored-by: mubashir1osmani <mubashir.osmani777@gmail.com>
Co-authored-by: Mubashir Osmani <ilikewafflesomcuh@gmail.com>
Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
Co-authored-by: mogith-pn <143642606+mogith-pn@users.noreply.github.com>

docs/my-website/docs/providers/snowflake.md

@@ -3,20 +3,15 @@ import TabItem from '@theme/TabItem';
 
 
 # Snowflake
-| Property | Details |
-|-------|-------|
-| Description | The Snowflake Cortex LLM REST API lets you access the COMPLETE function via HTTP POST requests|
-| Provider Route on LiteLLM | `snowflake/` |
-| Link to Provider Doc | [Snowflake ↗](https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api) |
-| Base URL | `https://{account-id}.snowflakecomputing.com/api/v2/cortex/inference:complete` |
-| Supported OpenAI Endpoints | `/chat/completions`, `/completions` |
+| Property                   | Details                                                                                                   |
+|----------------------------|-----------------------------------------------------------------------------------------------------------|
+| Description                | The Snowflake Cortex LLM REST API lets you access the COMPLETE and EMBED functions via HTTP POST requests |
+| Provider Route on LiteLLM  | `snowflake/`                                                                                              |
+| Link to Provider Doc       | [Snowflake ↗](https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api)              |
+| Base URLs                  | `https://{account-id}.snowflakecomputing.com/api/v2/cortex/inference:complete`,`https://{account-id}.snowflakecomputing.com/api/v2/cortex/inference:embed`|
+| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/embeddings`                                                        |
 
 
-
-Currently, Snowflake's REST API does not have an endpoint for `snowflake-arctic-embed` embedding models. If you want to use these embedding models with Litellm, you can call them through our Hugging Face provider. 
-
-Find the Arctic Embed models [here](https://huggingface.co/collections/Snowflake/arctic-embed-661fd57d50fab5fc314e4c18) on Hugging Face.
-
 ## Supported OpenAI Parameters
 ```
     "temperature",
@@ -29,6 +24,9 @@ Find the Arctic Embed models [here](https://huggingface.co/collections/Snowflake
 
 Snowflake does have API keys. Instead, you access the Snowflake API with your JWT token and account identifier.
 
+It is also possible to use [programmatic access tokens](https://docs.snowflake.com/en/user-guide/programmatic-access-tokens) (PAT). It can be defined by using 'pat/' prefix
+
+
 ```python
 import os 
 os.environ["SNOWFLAKE_JWT"] = "YOUR JWT"
@@ -37,17 +35,38 @@ os.environ["SNOWFLAKE_ACCOUNT_ID"] = "YOUR ACCOUNT IDENTIFIER"
 ## Usage
 
 ```python
-from litellm import completion
+from litellm import completion, embedding
 
 ## set ENV variables
-os.environ["SNOWFLAKE_JWT"] = "YOUR JWT"
+os.environ["SNOWFLAKE_JWT"] = "JWT_TOKEN"
 os.environ["SNOWFLAKE_ACCOUNT_ID"] = "YOUR ACCOUNT IDENTIFIER"
 
-# Snowflake call
+# Snowflake completion call
 response = completion(
     model="snowflake/mistral-7b", 
     messages = [{ "content": "Hello, how are you?","role": "user"}]
 )
+
+# Snowflake embedding call
+response = embedding(
+    model="snowflake/mistral-7b", 
+    input = ["My text"]
+)
+
+# Pass`api_key` and `account_id` as parameters
+response = completion(
+    model="snowflake/mistral-7b", 
+    messages = [{ "content": "Hello, how are you?","role": "user"}],
+    account_id="AAAA-BBBB",
+    api_key="JWT_TOKEN"
+)
+
+# using PAT
+response = completion(
+    model="snowflake/mistral-7b", 
+    messages = [{ "content": "Hello, how are you?","role": "user"}],
+    api_key="pat/PAT_TOKEN"
+)
 ```
 
 ## Usage with LiteLLM Proxy 

litellm/__init__.py

@@ -1356,6 +1356,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

@@ -693,12 +693,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

@@ -7,12 +7,14 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import httpx
 
-from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import ChatCompletionMessageToolCall, Function, ModelResponse
 
 from ...openai_like.chat.transformation import OpenAIGPTConfig
 
+from ..utils import SnowflakeBaseConfig
+
+
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
 
@@ -21,7 +23,7 @@ else:
     LiteLLMLoggingObj = Any
 
 
-class SnowflakeConfig(OpenAIGPTConfig):
+class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
     """
     Reference: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api
 
@@ -33,40 +35,6 @@ class SnowflakeConfig(OpenAIGPTConfig):
     def get_config(cls):
         return super().get_config()
 
-    def get_supported_openai_params(self, model: str) -> List[str]:
-        return [
-            "temperature",
-            "max_tokens",
-            "top_p",
-            "response_format",
-            "tools",
-            "tool_choice",
-        ]
-
-    def map_openai_params(
-        self,
-        non_default_params: dict,
-        optional_params: dict,
-        model: str,
-        drop_params: bool,
-    ) -> dict:
-        """
-        If any supported_openai_params are in non_default_params, add them to optional_params, so they are used in API call
-
-        Args:
-            non_default_params (dict): Non-default parameters to filter.
-            optional_params (dict): Optional parameters to update.
-            model (str): Model name for parameter support check.
-
-        Returns:
-            dict: Updated optional_params with supported non-default parameters.
-        """
-        supported_openai_params = self.get_supported_openai_params(model)
-        for param, value in non_default_params.items():
-            if param in supported_openai_params:
-                optional_params[param] = value
-        return optional_params
-
     def _transform_tool_calls_from_snowflake_to_openai(
         self, content_list: List[Dict[str, Any]]
     ) -> Tuple[str, Optional[List[ChatCompletionMessageToolCall]]]:
@@ -169,53 +137,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 +149,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 +200,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,69 @@
+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 ..utils import SnowflakeException, SnowflakeBaseConfig
+
+
+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,118 @@
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple
+
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.llms.openai import AllMessageValues
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
+
+    LiteLLMLoggingObj = _LiteLLMLoggingObj
+else:
+    LiteLLMLoggingObj = Any
+
+
+class SnowflakeException(BaseLLMException):
+    """Snowflake AI Endpoints exception handling class"""
+
+    pass
+
+
+class SnowflakeBaseConfig:
+    def get_supported_openai_params(self, model: str) -> List[str]:
+        return [
+            "temperature",
+            "max_tokens",
+            "top_p",
+            "response_format",
+            "tools",
+            "tool_choice",
+        ]
+
+    def map_openai_params(
+        self,
+        non_default_params: dict,
+        optional_params: dict,
+        model: str,
+        drop_params: bool,
+    ) -> dict:
+        """
+        If any supported_openai_params are in non_default_params, add them to optional_params, so they are used in API call
+
+        Args:
+            non_default_params (dict): Non-default parameters to filter.
+            optional_params (dict): Optional parameters to update.
+            model (str): Model name for parameter support check.
+
+        Returns:
+            dict: Updated optional_params with supported non-default parameters.
+        """
+        supported_openai_params = self.get_supported_openai_params(model)
+        for param, value in non_default_params.items():
+            if param in supported_openai_params:
+                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 api_key is None:
+            raise ValueError("Missing Snowflake JWT key")
+        else:
+            pat_key_prefix = "pat/"
+            if api_key.startswith(pat_key_prefix):
+                api_key = api_key[len(pat_key_prefix) :]
+                auth_type = "PROGRAMMATIC_ACCESS_TOKEN"
+
+        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

litellm/main.py

@@ -4822,6 +4822,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

@@ -7268,6 +7268,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()
         elif litellm.LlmProviders.SAGEMAKER == provider:

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

@@ -3,11 +3,14 @@ Unit tests for Snowflake chat transformation
 Tests tool calling request/response transformations
 """
 
+import os
+import copy
 import json
+
+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 +69,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 +99,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 +140,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 +218,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 +314,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 +330,98 @@ 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_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
+        assert 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,
+            api_key="pat/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,96 @@
+import os
+import json
+import copy
+
+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_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
+    assert 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"],
+        api_key="pat/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)