fix: Snowflake PAT auth and Claude streaming support

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 " + <JWT>,
-            "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT"
+            "Authorization": "Bearer " + <JWT or PAT>,
+            "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

litellm/llms/snowflake/common_utils.py

@@ -15,20 +15,28 @@ def validate_environment(
         {
             "Content-Type": "application/json",
             "Accept": "application/json",
-            "Authorization": "Bearer " + <JWT>,
-            "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT"
+            "Authorization": "Bearer " + <JWT or PAT>,
+            "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

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)