diff --git a/litellm/__init__.py b/litellm/__init__.py index dddcc6c2574..1fce5aa1903 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 275f2a63a1b..b68795a55d6 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/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, diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 4c0258d9f4b..68a0d3ce141 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/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 " + , + "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 " + , - "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) diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py new file mode 100644 index 00000000000..6173ae453ed --- /dev/null +++ b/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 + ) diff --git a/litellm/llms/snowflake/utils.py b/litellm/llms/snowflake/utils.py new file mode 100644 index 00000000000..b26ac1d8d2d --- /dev/null +++ b/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 diff --git a/litellm/main.py b/litellm/main.py index fa505f3fac5..e9ff9af328a 100644 --- a/litellm/main.py +++ b/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 diff --git a/litellm/utils.py b/litellm/utils.py index dd42f4aebf1..0b12603093a 100644 --- a/litellm/utils.py +++ b/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 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..83483c5260d 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,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) 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..3929618e6fa --- /dev/null +++ b/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)