diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 4c0258d9f4b..3b122793cdb 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -3,13 +3,15 @@ """ import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +import time +from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple, Union import httpx +from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator 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 litellm.types.utils import ChatCompletionMessageToolCall, Function, ModelResponse, ModelResponseStream from ...openai_like.chat.transformation import OpenAIGPTConfig @@ -187,20 +189,28 @@ def validate_environment( { "Content-Type": "application/json", "Accept": "application/json", - "Authorization": "Bearer " + , - "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT" + "Authorization": "Bearer " + , + "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT" or "PROGRAMMATIC_ACCESS_TOKEN" } """ if api_key is None: - raise ValueError("Missing Snowflake JWT key") + raise ValueError("Missing Snowflake JWT or PAT key") + + # Detect if using PAT token (prefixed with "pat/") + token_type = "KEYPAIR_JWT" + token = api_key + + if api_key.startswith("pat/"): + token_type = "PROGRAMMATIC_ACCESS_TOKEN" + token = api_key[4:] # Strip "pat/" prefix headers.update( { "Content-Type": "application/json", "Accept": "application/json", - "Authorization": "Bearer " + api_key, - "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT", + "Authorization": "Bearer " + token, + "X-Snowflake-Authorization-Token-Type": token_type, } ) return headers @@ -351,3 +361,54 @@ def transform_request( **optional_params, **extra_body, } + + def get_model_response_iterator( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], + sync_stream: bool, + json_mode: Optional[bool] = False, + ) -> Any: + """ + Return custom streaming handler for Snowflake that handles missing 'created' field. + + Some Snowflake models (like claude-sonnet-4-5) may not include the 'created' field + in their streaming responses, so we need a more defensive chunk parser. + """ + return SnowflakeChatCompletionStreamingHandler( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + + +class SnowflakeChatCompletionStreamingHandler(BaseModelResponseIterator): + """ + Custom streaming handler for Snowflake Cortex API. + + Handles cases where the 'created' field might be missing from streaming chunks, + which can occur with certain Snowflake models (e.g., claude-sonnet-4-5). + """ + + def chunk_parser(self, chunk: dict) -> ModelResponseStream: + """ + Parse a streaming chunk from Snowflake, providing defaults for missing fields. + + Args: + chunk: The JSON chunk from Snowflake's streaming response + + Returns: + ModelResponseStream with all required fields populated + """ + try: + # Use current time as default if 'created' field is missing + created = chunk.get("created", int(time.time())) + + return ModelResponseStream( + id=chunk["id"], + object="chat.completion.chunk", + created=created, + model=chunk["model"], + choices=chunk["choices"], + ) + except Exception as e: + raise e diff --git a/litellm/llms/snowflake/common_utils.py b/litellm/llms/snowflake/common_utils.py index 40c8270f95f..883b9fa3706 100644 --- a/litellm/llms/snowflake/common_utils.py +++ b/litellm/llms/snowflake/common_utils.py @@ -15,20 +15,28 @@ def validate_environment( { "Content-Type": "application/json", "Accept": "application/json", - "Authorization": "Bearer " + , - "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT" + "Authorization": "Bearer " + , + "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT" or "PROGRAMMATIC_ACCESS_TOKEN" } """ if JWT is None: - raise ValueError("Missing Snowflake JWT key") + raise ValueError("Missing Snowflake JWT or PAT key") + + # Detect if using PAT token (prefixed with "pat/") + token_type = "KEYPAIR_JWT" + token = JWT + + if JWT.startswith("pat/"): + token_type = "PROGRAMMATIC_ACCESS_TOKEN" + token = JWT[4:] # Strip "pat/" prefix headers.update( { "Content-Type": "application/json", "Accept": "application/json", - "Authorization": "Bearer " + JWT, - "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT", + "Authorization": "Bearer " + token, + "X-Snowflake-Authorization-Token-Type": token_type, } ) return headers 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..0677e4cf8ab 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 @@ -313,3 +313,172 @@ def test_get_supported_openai_params_includes_tools(self): assert "tool_choice" in supported_params assert "temperature" in supported_params assert "max_tokens" in supported_params + + +class TestSnowflakeAuthenticationHeaders: + """Test suite for Snowflake authentication header handling""" + + def test_validate_environment_with_jwt(self): + """ + Test that JWT tokens are handled correctly with KEYPAIR_JWT header. + """ + config = SnowflakeConfig() + headers = {} + + jwt_token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.test" + + result_headers = config.validate_environment( + headers=headers, + model="mistral-7b", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=jwt_token, + api_base=None, + ) + + assert result_headers["Authorization"] == f"Bearer {jwt_token}" + assert result_headers["X-Snowflake-Authorization-Token-Type"] == "KEYPAIR_JWT" + assert result_headers["Content-Type"] == "application/json" + assert result_headers["Accept"] == "application/json" + + def test_validate_environment_with_pat_token(self): + """ + Test that PAT tokens with pat/ prefix are handled correctly. + The pat/ prefix should be stripped and PROGRAMMATIC_ACCESS_TOKEN should be used. + """ + config = SnowflakeConfig() + headers = {} + + pat_token = "pat/abc123xyz789" + expected_token = "abc123xyz789" + + result_headers = config.validate_environment( + headers=headers, + model="mistral-7b", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=pat_token, + api_base=None, + ) + + assert result_headers["Authorization"] == f"Bearer {expected_token}" + assert result_headers["X-Snowflake-Authorization-Token-Type"] == "PROGRAMMATIC_ACCESS_TOKEN" + assert result_headers["Content-Type"] == "application/json" + assert result_headers["Accept"] == "application/json" + + def test_validate_environment_missing_api_key(self): + """ + Test that missing API key raises ValueError. + """ + config = SnowflakeConfig() + headers = {} + + with pytest.raises(ValueError, match="Missing Snowflake JWT or PAT key"): + config.validate_environment( + headers=headers, + model="mistral-7b", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + +class TestSnowflakeStreamingHandler: + """Test suite for Snowflake streaming response handling""" + + def test_chunk_parser_with_created_field(self): + """ + Test that streaming chunks with 'created' field are parsed correctly. + This is the standard case for models like mistral-7b and llama3.3. + """ + from litellm.llms.snowflake.chat.transformation import ( + SnowflakeChatCompletionStreamingHandler, + ) + + handler = SnowflakeChatCompletionStreamingHandler( + streaming_response=iter([]), + sync_stream=True, + json_mode=False, + ) + + chunk = { + "id": "chatcmpl-123", + "created": 1234567890, + "model": "mistral-7b", + "choices": [ + { + "index": 0, + "delta": {"content": "Hello"}, + "finish_reason": None, + } + ], + } + + result = handler.chunk_parser(chunk) + + assert result.id == "chatcmpl-123" + assert result.created == 1234567890 + assert result.model == "mistral-7b" + assert result.object == "chat.completion.chunk" + assert len(result.choices) == 1 + + def test_chunk_parser_without_created_field(self): + """ + Test that streaming chunks WITHOUT 'created' field are parsed correctly. + This handles the case for Claude models (sonnet-3.5, sonnet-4-5) which + don't include the 'created' field in their streaming responses. + """ + from litellm.llms.snowflake.chat.transformation import ( + SnowflakeChatCompletionStreamingHandler, + ) + + handler = SnowflakeChatCompletionStreamingHandler( + streaming_response=iter([]), + sync_stream=True, + json_mode=False, + ) + + # Chunk without 'created' field (like claude-sonnet-4-5) + chunk = { + "id": "chatcmpl-456", + "model": "claude-sonnet-4-5", + "choices": [ + { + "index": 0, + "delta": {"content": "Hi there"}, + "finish_reason": None, + } + ], + } + + result = handler.chunk_parser(chunk) + + assert result.id == "chatcmpl-456" + assert result.created is not None # Should have a default timestamp + assert isinstance(result.created, int) # Should be an integer timestamp + assert result.model == "claude-sonnet-4-5" + assert result.object == "chat.completion.chunk" + assert len(result.choices) == 1 + + def test_get_model_response_iterator(self): + """ + Test that SnowflakeConfig returns the custom streaming handler. + """ + from litellm.llms.snowflake.chat.transformation import ( + SnowflakeChatCompletionStreamingHandler, + SnowflakeConfig, + ) + + config = SnowflakeConfig() + + handler = config.get_model_response_iterator( + streaming_response=iter([]), + sync_stream=True, + json_mode=False, + ) + + assert isinstance(handler, SnowflakeChatCompletionStreamingHandler)