diff --git a/tests/llm_translation/test_snowflake.py b/tests/llm_translation/test_snowflake.py index 484be1e054b..6861c2c7eca 100644 --- a/tests/llm_translation/test_snowflake.py +++ b/tests/llm_translation/test_snowflake.py @@ -1,171 +1,79 @@ import asyncio import json +import os import httpx from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm import completion, acompletion +from litellm import completion, acompletion, responses +from litellm.exceptions import APIConnectionError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -FAKE_API_BASE = "https://fake-snowflake.example.com/api/v2/cortex/inference:chat" +@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 -def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock: - mock = MagicMock(spec=httpx.Response) - mock.status_code = 200 - mock.headers = {"content-type": "application/json"} - mock.json.return_value = json_data - mock.text = json.dumps(json_data) - return mock + # 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! -def _chat_response() -> Dict[str, Any]: - return { - "id": "chatcmpl-snowflake-123", - "object": "chat.completion", - "created": 1700000000, - "model": "mistral-7b", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "The sky above is painted blue,\nWith clouds of white and morning dew.", + 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", + } }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 30, - "total_tokens": 40, - }, - } - - -def _streaming_chunks() -> List[str]: - base = { - "id": "chatcmpl-snowflake-stream-123", - "object": "chat.completion.chunk", - "created": 1700000000, - "model": "mistral-7b", - } - deltas = [ - {"role": "assistant", "content": "The"}, - {"content": " sky"}, - {"content": " is blue"}, + "required": ["location"], + }, + } ] - chunks = [] - for i, delta in enumerate(deltas): - finish = "stop" if i == len(deltas) - 1 else None - chunks.append( - json.dumps( - { - **base, - "choices": [ - {"index": 0, "delta": delta, "finish_reason": finish} - ], - } - ) + + 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, ) - return chunks + assert response is not None + assert hasattr(response, "output") + assert len(response.output) > 0 -@pytest.mark.parametrize("sync_mode", [True, False]) -def test_chat_completion_snowflake(sync_mode): - messages = [{"role": "user", "content": "Write me a poem about the blue sky"}] - mock_resp = _make_mock_response(_chat_response()) + # 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 - if sync_mode: - with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: - response = completion( - model="snowflake/mistral-7b", - messages=messages, - api_base=FAKE_API_BASE, - ) - mock_post.assert_called_once() - else: - with patch.object( - AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp - ) as mock_post: - response = asyncio.run( - acompletion( - model="snowflake/mistral-7b", - messages=messages, - api_base=FAKE_API_BASE, - ) - ) - mock_post.assert_called_once() + assert tool_call_found, "Expected tool call but none was found" - assert response is not None - assert response.choices[0].message.content is not None - assert "sky" in response.choices[0].message.content.lower() - assert response.usage.prompt_tokens == 10 - assert response.usage.completion_tokens == 30 - - -@pytest.mark.parametrize("sync_mode", [True, False]) -def test_chat_completion_snowflake_stream(sync_mode): - messages = [{"role": "user", "content": "Write me a poem about the blue sky"}] - raw_chunks = _streaming_chunks() - - if sync_mode: - - def _iter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.iter_lines.return_value = _iter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - - with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: - response = completion( - model="snowflake/mistral-7b", - messages=messages, - max_tokens=100, - stream=True, - api_base=FAKE_API_BASE, - ) - chunks_received = list(response) - mock_post.assert_called_once() - else: - - async def _aiter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.aiter_lines.return_value = _aiter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - - async def _run(): - with patch.object( - AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp - ) as mock_post: - resp = await acompletion( - model="snowflake/mistral-7b", - messages=messages, - max_tokens=100, - stream=True, - api_base=FAKE_API_BASE, - ) - received = [] - async for chunk in resp: - received.append(chunk) - mock_post.assert_called_once() - return received - - chunks_received = asyncio.run(_run()) - - assert len(chunks_received) > 0 - content = "".join( - c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content - ) - assert "sky" in content.lower() + 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 index 5bb4942dde6..5b18618fdf5 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 @@ -1,18 +1,22 @@ """ Unit tests for Snowflake chat transformation -Tests tool calling request/response transformations +Tests tool calling request/response transformations and chat completions """ +import asyncio import os import copy import json +from typing import Any, Dict, List -from unittest.mock import patch -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, patch, Mock, MagicMock import httpx +import pytest import litellm +from litellm import completion, acompletion +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.snowflake.chat.transformation import SnowflakeConfig from litellm.types.utils import ModelResponse @@ -438,3 +442,174 @@ class TestSnowFlakeCompletion: os.environ.pop("SNOWFLAKE_ACCOUNT_ID", None) os.environ.pop("SNOWFLAKE_JWT", None) + + +FAKE_API_BASE = "https://fake-snowflake.example.com/api/v2/cortex/inference:chat" + + +def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock: + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + mock.headers = {"content-type": "application/json"} + mock.json.return_value = json_data + mock.text = json.dumps(json_data) + return mock + + +def _chat_response() -> Dict[str, Any]: + return { + "id": "chatcmpl-snowflake-123", + "object": "chat.completion", + "created": 1700000000, + "model": "mistral-7b", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "The sky above is painted blue,\nWith clouds of white and morning dew.", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 30, + "total_tokens": 40, + }, + } + + +def _streaming_chunks() -> List[str]: + base = { + "id": "chatcmpl-snowflake-stream-123", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-7b", + } + deltas = [ + {"role": "assistant", "content": "The"}, + {"content": " sky"}, + {"content": " is blue"}, + ] + chunks = [] + for i, delta in enumerate(deltas): + finish = "stop" if i == len(deltas) - 1 else None + chunks.append( + json.dumps( + { + **base, + "choices": [ + {"index": 0, "delta": delta, "finish_reason": finish} + ], + } + ) + ) + return chunks + + +class TestSnowflakeChatCompletion: + """End-to-end chat completion tests (mocked HTTP).""" + + messages = [{"role": "user", "content": "Write me a poem about the blue sky"}] + + @pytest.mark.parametrize("sync_mode", [True, False]) + def test_chat_completion_snowflake(self, sync_mode): + mock_resp = _make_mock_response(_chat_response()) + + if sync_mode: + with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: + response = completion( + model="snowflake/mistral-7b", + messages=self.messages, + api_key="fake-jwt", + account_id="FAKE-ACCOUNT", + api_base=FAKE_API_BASE, + ) + mock_post.assert_called_once() + else: + with patch.object( + AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp + ) as mock_post: + response = asyncio.run( + acompletion( + model="snowflake/mistral-7b", + messages=self.messages, + api_key="fake-jwt", + account_id="FAKE-ACCOUNT", + api_base=FAKE_API_BASE, + ) + ) + mock_post.assert_called_once() + + assert response is not None + assert response.choices[0].message.content is not None + assert "sky" in response.choices[0].message.content.lower() + assert response.usage.prompt_tokens == 10 + assert response.usage.completion_tokens == 30 + + @pytest.mark.parametrize("sync_mode", [True, False]) + def test_chat_completion_snowflake_stream(self, sync_mode): + raw_chunks = _streaming_chunks() + + if sync_mode: + + def _iter_lines(): + for chunk in raw_chunks: + yield f"data: {chunk}" + yield "data: [DONE]" + + mock_resp = MagicMock() + mock_resp.iter_lines.return_value = _iter_lines() + mock_resp.status_code = 200 + mock_resp.headers = {"content-type": "text/event-stream"} + + with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: + response = completion( + model="snowflake/mistral-7b", + messages=self.messages, + max_tokens=100, + stream=True, + api_key="fake-jwt", + account_id="FAKE-ACCOUNT", + api_base=FAKE_API_BASE, + ) + chunks_received = list(response) + mock_post.assert_called_once() + else: + + async def _aiter_lines(): + for chunk in raw_chunks: + yield f"data: {chunk}" + yield "data: [DONE]" + + mock_resp = MagicMock() + mock_resp.aiter_lines.return_value = _aiter_lines() + mock_resp.status_code = 200 + mock_resp.headers = {"content-type": "text/event-stream"} + + async def _run(): + with patch.object( + AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp + ) as mock_post: + resp = await acompletion( + model="snowflake/mistral-7b", + messages=self.messages, + max_tokens=100, + stream=True, + api_key="fake-jwt", + account_id="FAKE-ACCOUNT", + api_base=FAKE_API_BASE, + ) + received = [] + async for chunk in resp: + received.append(chunk) + mock_post.assert_called_once() + return received + + chunks_received = asyncio.run(_run()) + + assert len(chunks_received) > 0 + content = "".join( + c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content + )