diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 2b92911b055..4c0258d9f4b 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -1,14 +1,15 @@ """ -Support for Snowflake REST API +Support for Snowflake REST API """ -from typing import TYPE_CHECKING, Any, List, Optional, Tuple +import json +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 ModelResponse +from litellm.types.utils import ChatCompletionMessageToolCall, Function, ModelResponse from ...openai_like.chat.transformation import OpenAIGPTConfig @@ -22,15 +23,25 @@ else: class SnowflakeConfig(OpenAIGPTConfig): """ - source: https://docs.snowflake.com/en/sql-reference/functions/complete-snowflake-cortex + 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 get_supported_openai_params(self, model: str) -> List: - return ["temperature", "max_tokens", "top_p", "response_format"] + 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, @@ -56,6 +67,57 @@ class SnowflakeConfig(OpenAIGPTConfig): 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]]]: + """ + Transform Snowflake tool calls to OpenAI format. + + Args: + content_list: Snowflake's content_list array containing text and tool_use items + + Returns: + Tuple of (text_content, tool_calls) + + Snowflake format in content_list: + { + "type": "tool_use", + "tool_use": { + "tool_use_id": "tooluse_...", + "name": "get_weather", + "input": {"location": "Paris"} + } + } + + OpenAI format (returned tool_calls): + ChatCompletionMessageToolCall( + id="tooluse_...", + type="function", + function=Function(name="get_weather", arguments='{"location": "Paris"}') + ) + """ + text_content = "" + tool_calls: List[ChatCompletionMessageToolCall] = [] + + for idx, content_item in enumerate(content_list): + if content_item.get("type") == "text": + text_content += content_item.get("text", "") + + ## TOOL CALLING + elif content_item.get("type") == "tool_use": + tool_use_data = content_item.get("tool_use", {}) + tool_call = ChatCompletionMessageToolCall( + id=tool_use_data.get("tool_use_id", ""), + type="function", + function=Function( + name=tool_use_data.get("name", ""), + arguments=json.dumps(tool_use_data.get("input", {})), + ), + ) + tool_calls.append(tool_call) + + return text_content, tool_calls if tool_calls else None + def transform_response( self, model: str, @@ -71,6 +133,7 @@ class SnowflakeConfig(OpenAIGPTConfig): json_mode: Optional[bool] = None, ) -> ModelResponse: response_json = raw_response.json() + logging_obj.post_call( input=messages, api_key="", @@ -78,6 +141,26 @@ class SnowflakeConfig(OpenAIGPTConfig): additional_args={"complete_input_dict": request_data}, ) + ## RESPONSE TRANSFORMATION + # Snowflake returns content_list (not content) with tool_use objects + # We need to transform this to OpenAI's format with content + tool_calls + if "choices" in response_json and len(response_json["choices"]) > 0: + choice = response_json["choices"][0] + if "message" in choice and "content_list" in choice["message"]: + content_list = choice["message"]["content_list"] + ( + text_content, + tool_calls, + ) = self._transform_tool_calls_from_snowflake_to_openai(content_list) + + # Update the choice message with OpenAI format + choice["message"]["content"] = text_content + if tool_calls: + choice["message"]["tool_calls"] = tool_calls + + # Remove Snowflake-specific content_list + del choice["message"]["content_list"] + returned_response = ModelResponse(**response_json) returned_response.model = "snowflake/" + (returned_response.model or "") @@ -150,6 +233,95 @@ class SnowflakeConfig(OpenAIGPTConfig): return api_base + def _transform_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Transform OpenAI tool format to Snowflake tool format. + + Args: + tools: List of tools in OpenAI format + + Returns: + List of tools in Snowflake format + + OpenAI format: + { + "type": "function", + "function": { + "name": "get_weather", + "description": "...", + "parameters": {...} + } + } + + Snowflake format: + { + "tool_spec": { + "type": "generic", + "name": "get_weather", + "description": "...", + "input_schema": {...} + } + } + """ + snowflake_tools: List[Dict[str, Any]] = [] + for tool in tools: + if tool.get("type") == "function": + function = tool.get("function", {}) + snowflake_tool: Dict[str, Any] = { + "tool_spec": { + "type": "generic", + "name": function.get("name"), + "input_schema": function.get( + "parameters", + {"type": "object", "properties": {}}, + ), + } + } + # Add description if present + if "description" in function: + snowflake_tool["tool_spec"]["description"] = function[ + "description" + ] + + snowflake_tools.append(snowflake_tool) + + return snowflake_tools + + def _transform_tool_choice( + self, tool_choice: Union[str, Dict[str, Any]] + ) -> Union[str, Dict[str, Any]]: + """ + Transform OpenAI tool_choice format to Snowflake format. + + Args: + tool_choice: Tool choice in OpenAI format (str or dict) + + Returns: + Tool choice in Snowflake format + + OpenAI format: + {"type": "function", "function": {"name": "get_weather"}} + + Snowflake format: + {"type": "tool", "name": ["get_weather"]} + + Note: String values ("auto", "required", "none") pass through unchanged. + """ + if isinstance(tool_choice, str): + # "auto", "required", "none" pass through as-is + return tool_choice + + if isinstance(tool_choice, dict): + if tool_choice.get("type") == "function": + function_name = tool_choice.get("function", {}).get("name") + if function_name: + return { + "type": "tool", + "name": [function_name], # Snowflake expects array + } + + return tool_choice + def transform_request( self, model: str, @@ -160,6 +332,18 @@ class SnowflakeConfig(OpenAIGPTConfig): ) -> dict: stream: bool = optional_params.pop("stream", None) or False extra_body = optional_params.pop("extra_body", {}) + + ## TOOL CALLING + # Transform tools from OpenAI format to Snowflake's tool_spec format + tools = optional_params.pop("tools", None) + if tools: + optional_params["tools"] = self._transform_tools(tools) + + # Transform tool_choice from OpenAI format to Snowflake's tool name array format + tool_choice = optional_params.pop("tool_choice", None) + if tool_choice: + optional_params["tool_choice"] = self._transform_tool_choice(tool_choice) + return { "model": model, "messages": messages, diff --git a/tests/llm_translation/test_snowflake.py b/tests/llm_translation/test_snowflake.py index 083081eaee3..12d738458c3 100644 --- a/tests/llm_translation/test_snowflake.py +++ b/tests/llm_translation/test_snowflake.py @@ -6,7 +6,7 @@ from dotenv import load_dotenv load_dotenv() import pytest -from litellm import completion, acompletion +from litellm import completion, acompletion, responses from litellm.exceptions import APIConnectionError @pytest.mark.parametrize("sync_mode", [True, False]) @@ -87,3 +87,70 @@ async def test_chat_completion_snowflake_stream(sync_mode): raise # Re-raise if it's a different APIConnectionError except Exception as e: pytest.fail(f"Error occurred: {e}") + + +@pytest.mark.skip(reason="Requires Snowflake credentials - run manually when needed") +def test_snowflake_tool_calling_responses_api(): + """ + Test Snowflake tool calling with Responses API. + Requires SNOWFLAKE_JWT and SNOWFLAKE_ACCOUNT_ID environment variables. + """ + import litellm + + # Skip if credentials not available + if not os.getenv("SNOWFLAKE_JWT") or not os.getenv("SNOWFLAKE_ACCOUNT_ID"): + pytest.skip("Snowflake credentials not available") + + litellm.drop_params = False # We now support tools! + + tools = [ + { + "type": "function", + "name": "get_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + } + }, + "required": ["location"], + }, + } + ] + + try: + # Test with tool_choice to force tool use + response = responses( + model="snowflake/claude-3-5-sonnet", + input="What's the weather in Paris?", + tools=tools, + tool_choice={"type": "function", "function": {"name": "get_weather"}}, + max_output_tokens=200, + ) + + assert response is not None + assert hasattr(response, "output") + assert len(response.output) > 0 + + # Verify tool call was made + tool_call_found = False + for item in response.output: + if hasattr(item, "type") and item.type == "function_call": + tool_call_found = True + assert item.name == "get_weather" + assert hasattr(item, "arguments") + print(f"✅ Tool call detected: {item.name}({item.arguments})") + break + + assert tool_call_found, "Expected tool call but none was found" + + except APIConnectionError as e: + if "JWT token is invalid" in str(e): + pytest.skip("Invalid Snowflake JWT token") + elif "Application failed to respond" in str(e) or "502" in str(e): + pytest.skip(f"Snowflake API unavailable: {e}") + else: + raise 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 new file mode 100644 index 00000000000..7422d03074e --- /dev/null +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -0,0 +1,315 @@ +""" +Unit tests for Snowflake chat transformation +Tests tool calling request/response transformations +""" + +import json +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.llms.snowflake.chat.transformation import SnowflakeConfig +from litellm.types.utils import ModelResponse + + +class TestSnowflakeToolTransformation: + """Test suite for Snowflake tool calling transformations""" + + def test_transform_request_with_tools(self): + """ + Test that OpenAI tool format is correctly transformed to Snowflake's tool_spec format. + """ + config = SnowflakeConfig() + + # OpenAI format tools + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + }, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + }, + }, + "required": ["location"], + }, + }, + } + ] + + optional_params = {"tools": tools} + + transformed_request = config.transform_request( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "What's the weather?"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + # Verify tools were transformed to Snowflake format + assert "tools" in transformed_request + assert len(transformed_request["tools"]) == 1 + + snowflake_tool = transformed_request["tools"][0] + 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 "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"] + + def test_transform_request_with_tool_choice(self): + """ + Test that OpenAI tool_choice format is correctly transformed to Snowflake format. + """ + config = SnowflakeConfig() + + # OpenAI format tool_choice + tool_choice = {"type": "function", "function": {"name": "get_weather"}} + + optional_params = {"tool_choice": tool_choice} + + transformed_request = config.transform_request( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "What's the weather?"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + # 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 + + def test_transform_request_with_string_tool_choice(self): + """ + Test that string tool_choice values pass through unchanged. + """ + config = SnowflakeConfig() + + for value in ["auto", "required", "none"]: + optional_params = {"tool_choice": value} + + transformed_request = config.transform_request( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Test"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert transformed_request["tool_choice"] == value + + def test_transform_response_with_tool_calls(self): + """ + Test that Snowflake's content_list with tool_use is transformed to OpenAI format. + """ + config = SnowflakeConfig() + + # Mock Snowflake response with tool call + mock_snowflake_response = { + "choices": [ + { + "message": { + "content_list": [ + {"type": "text", "text": ""}, + { + "type": "tool_use", + "tool_use": { + "tool_use_id": "tooluse_abc123", + "name": "get_weather", + "input": {"location": "Paris, France", "unit": "celsius"}, + }, + }, + ] + } + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + } + + response = httpx.Response( + status_code=200, + json=mock_snowflake_response, + headers={"Content-Type": "application/json"}, + ) + + model_response = ModelResponse( + choices=[litellm.Choices(index=0, message=litellm.Message())] + ) + + logging_obj = MagicMock() + + result = config.transform_response( + model="claude-3-5-sonnet", + raw_response=response, + model_response=model_response, + logging_obj=logging_obj, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding={}, + ) + + # General assertions + assert isinstance(result, ModelResponse) + assert len(result.choices) == 1 + + choice = result.choices[0] + assert isinstance(choice, litellm.Choices) + + # Message and tool_calls assertions + message = choice.message + assert isinstance(message, litellm.Message) + assert hasattr(message, "tool_calls") + assert isinstance(message.tool_calls, list) + assert len(message.tool_calls) == 1 + + # Specific tool_call assertions + tool_call = message.tool_calls[0] + assert isinstance(tool_call, litellm.utils.ChatCompletionMessageToolCall) + assert tool_call.id == "tooluse_abc123" + assert tool_call.type == "function" + assert tool_call.function.name == "get_weather" + + # Verify arguments are properly JSON serialized + arguments = json.loads(tool_call.function.arguments) + assert arguments["location"] == "Paris, France" + assert arguments["unit"] == "celsius" + + # Verify content_list was removed and content was set + assert message.content == "" + + def test_transform_response_with_mixed_content(self): + """ + Test that responses with both text and tool calls are handled correctly. + """ + config = SnowflakeConfig() + + # Mock Snowflake response with text and tool call + mock_snowflake_response = { + "choices": [ + { + "message": { + "content_list": [ + {"type": "text", "text": "Let me check the weather for you. "}, + { + "type": "tool_use", + "tool_use": { + "tool_use_id": "tooluse_xyz789", + "name": "get_weather", + "input": {"location": "Tokyo, Japan"}, + }, + }, + ] + } + } + ], + "usage": {"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40}, + } + + response = httpx.Response( + status_code=200, + json=mock_snowflake_response, + headers={"Content-Type": "application/json"}, + ) + + model_response = ModelResponse( + choices=[litellm.Choices(index=0, message=litellm.Message())] + ) + + logging_obj = MagicMock() + + result = config.transform_response( + model="claude-3-5-sonnet", + raw_response=response, + model_response=model_response, + logging_obj=logging_obj, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding={}, + ) + + # Verify text content was extracted + message = result.choices[0].message + assert message.content == "Let me check the weather for you. " + + # Verify tool call was also extracted + assert len(message.tool_calls) == 1 + assert message.tool_calls[0].function.name == "get_weather" + + def test_transform_response_without_tool_calls(self): + """ + Test that regular text responses (without tools) work correctly. + """ + config = SnowflakeConfig() + + # Mock Snowflake response without tool calls (standard response) + mock_snowflake_response = { + "choices": [ + { + "message": { + "content": "Hello! I'm doing well, thank you for asking.", + "role": "assistant", + } + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 15, "total_tokens": 25}, + } + + response = httpx.Response( + status_code=200, + json=mock_snowflake_response, + headers={"Content-Type": "application/json"}, + ) + + model_response = ModelResponse( + choices=[litellm.Choices(index=0, message=litellm.Message())] + ) + + logging_obj = MagicMock() + + result = config.transform_response( + model="mistral-7b", + raw_response=response, + model_response=model_response, + logging_obj=logging_obj, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding={}, + ) + + # Verify standard response works + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Hello! I'm doing well, thank you for asking." + + def test_get_supported_openai_params_includes_tools(self): + """ + Test that tools and tool_choice are in supported params. + """ + config = SnowflakeConfig() + supported_params = config.get_supported_openai_params("claude-3-5-sonnet") + + assert "tools" in supported_params + assert "tool_choice" in supported_params + assert "temperature" in supported_params + assert "max_tokens" in supported_params