From 4dbbec233cebd126c6286d574839e4922bff4ff4 Mon Sep 17 00:00:00 2001 From: Andrey Date: Tue, 18 Nov 2025 07:27:46 +0300 Subject: [PATCH] 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 950b7cef44f14b2db1429f6fbd32548a7c95d325. * fix: arize ui integration * need to remove a file This reverts commit d6c877b73ac763464f204b77135f3786342373b7. * 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 Co-authored-by: Mubashir Osmani Co-authored-by: Krish Dholakia Co-authored-by: mogith-pn <143642606+mogith-pn@users.noreply.github.com> --- docs/my-website/docs/providers/snowflake.md | 49 ++++--- litellm/__init__.py | 1 + .../get_llm_provider_logic.py | 12 +- litellm/llms/snowflake/chat/transformation.py | 97 ++------------ .../snowflake/embedding/transformation.py | 69 ++++++++++ litellm/llms/snowflake/utils.py | 118 +++++++++++++++++ litellm/main.py | 16 +++ litellm/utils.py | 2 + .../test_snowflake_chat_transformation.py | 124 +++++++++++++++++- .../embedding/test_snowflake_embedding.py | 96 ++++++++++++++ 10 files changed, 468 insertions(+), 116 deletions(-) create mode 100644 litellm/llms/snowflake/embedding/transformation.py create mode 100644 litellm/llms/snowflake/utils.py create mode 100644 tests/test_litellm/llms/snowflake/embedding/test_snowflake_embedding.py diff --git a/docs/my-website/docs/providers/snowflake.md b/docs/my-website/docs/providers/snowflake.md index 40deef87805..483bf939fe6 100644 --- a/docs/my-website/docs/providers/snowflake.md +++ b/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 diff --git a/litellm/__init__.py b/litellm/__init__.py index 6cf2aca3136..c86768490f3 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index fb25c5ed840..ef0ebe074d7 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/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, diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 4c0258d9f4b..62ede0aeaf8 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/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 " + , - "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) diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py new file mode 100644 index 00000000000..83716f3ef26 --- /dev/null +++ b/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 + ) diff --git a/litellm/llms/snowflake/utils.py b/litellm/llms/snowflake/utils.py new file mode 100644 index 00000000000..9d458f6ece3 --- /dev/null +++ b/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 " + , + "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 diff --git a/litellm/main.py b/litellm/main.py index 14d0b04b7b5..88c3f7bc55b 100644 --- a/litellm/main.py +++ b/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 diff --git a/litellm/utils.py b/litellm/utils.py index 783d462a7af..2249599dbc5 100644 --- a/litellm/utils.py +++ b/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: diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index 7422d03074e..c2527d8fbdc 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/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) diff --git a/tests/test_litellm/llms/snowflake/embedding/test_snowflake_embedding.py b/tests/test_litellm/llms/snowflake/embedding/test_snowflake_embedding.py new file mode 100644 index 00000000000..8f4551bc79b --- /dev/null +++ b/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)