mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-06 10:21:32 +00:00
snowflake test fix
This commit is contained in:
Ishaan Jaffer
@@ -1,6 +1,9 @@
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
import json
|
||||
import httpx
|
||||
from typing import Any, Dict, List
|
||||
from unittest.mock import Mock, MagicMock, patch
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
@@ -8,85 +11,201 @@ import pytest
|
||||
|
||||
from litellm import completion, acompletion, responses
|
||||
from litellm.exceptions import APIConnectionError
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
||||
|
||||
def mock_snowflake_chat_response() -> Dict[str, Any]:
|
||||
"""
|
||||
Mock response for Snowflake chat completion.
|
||||
"""
|
||||
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.\nA canvas vast, serene and bright,\nThat fills my heart with pure delight.",
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 30,
|
||||
"total_tokens": 40,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def mock_snowflake_streaming_response_chunks() -> List[str]:
|
||||
"""
|
||||
Mock streaming response chunks for Snowflake.
|
||||
"""
|
||||
return [
|
||||
json.dumps({
|
||||
"id": "chatcmpl-snowflake-stream-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "mistral-7b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "content": "The"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}),
|
||||
json.dumps({
|
||||
"id": "chatcmpl-snowflake-stream-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "mistral-7b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": " sky"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}),
|
||||
json.dumps({
|
||||
"id": "chatcmpl-snowflake-stream-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "mistral-7b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": " is blue"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
}),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_completion_snowflake(sync_mode):
|
||||
try:
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Write me a poem about the blue sky",
|
||||
},
|
||||
]
|
||||
def test_chat_completion_snowflake(sync_mode):
|
||||
"""
|
||||
Test Snowflake chat completion with mocked HTTP responses.
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Write me a poem about the blue sky",
|
||||
},
|
||||
]
|
||||
|
||||
if sync_mode:
|
||||
mock_response = Mock(spec=httpx.Response)
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = mock_snowflake_chat_response()
|
||||
|
||||
if sync_mode:
|
||||
sync_handler = HTTPHandler()
|
||||
with patch.object(HTTPHandler, "post", return_value=mock_response):
|
||||
response = completion(
|
||||
model="snowflake/mistral-7b",
|
||||
messages=messages,
|
||||
api_base = "https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions"
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions",
|
||||
client=sync_handler,
|
||||
)
|
||||
print(response)
|
||||
assert response is not None
|
||||
else:
|
||||
response = await acompletion(
|
||||
model="snowflake/mistral-7b",
|
||||
messages=messages,
|
||||
api_base = "https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions"
|
||||
assert response.choices[0].message.content is not None
|
||||
assert "sky" in response.choices[0].message.content.lower()
|
||||
else:
|
||||
async_handler = AsyncHTTPHandler()
|
||||
with patch.object(AsyncHTTPHandler, "post", return_value=mock_response):
|
||||
import asyncio
|
||||
response = asyncio.run(
|
||||
acompletion(
|
||||
model="snowflake/mistral-7b",
|
||||
messages=messages,
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions",
|
||||
client=async_handler,
|
||||
)
|
||||
)
|
||||
print(response)
|
||||
assert response is not None
|
||||
except APIConnectionError as e:
|
||||
# Skip test if Snowflake API is unavailable (502 error)
|
||||
if "Application failed to respond" in str(e) or "502" in str(e):
|
||||
pytest.skip(f"Snowflake API unavailable: {e}")
|
||||
else:
|
||||
raise # Re-raise if it's a different APIConnectionError
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
assert response.choices[0].message.content is not None
|
||||
assert "sky" in response.choices[0].message.content.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
async def test_chat_completion_snowflake_stream(sync_mode):
|
||||
try:
|
||||
set_verbose = True
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Write me a poem about the blue sky",
|
||||
},
|
||||
]
|
||||
def test_chat_completion_snowflake_stream(sync_mode):
|
||||
"""
|
||||
Test Snowflake streaming chat completion with mocked HTTP responses.
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Write me a poem about the blue sky",
|
||||
},
|
||||
]
|
||||
|
||||
if sync_mode:
|
||||
sync_handler = HTTPHandler()
|
||||
mock_chunks = mock_snowflake_streaming_response_chunks()
|
||||
|
||||
if sync_mode is False:
|
||||
response = await acompletion(
|
||||
def mock_iter_lines():
|
||||
for chunk in mock_chunks:
|
||||
for line in [f"data: {chunk}", "data: [DONE]"]:
|
||||
yield line
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.iter_lines.side_effect = mock_iter_lines
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch.object(HTTPHandler, "post", return_value=mock_response):
|
||||
response = completion(
|
||||
model="snowflake/mistral-7b",
|
||||
messages=messages,
|
||||
max_tokens=100,
|
||||
stream=True,
|
||||
api_base = "https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions"
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions",
|
||||
client=sync_handler,
|
||||
)
|
||||
|
||||
async for chunk in response:
|
||||
print(chunk)
|
||||
else:
|
||||
response = completion(
|
||||
model="snowflake/mistral-7b",
|
||||
messages=messages,
|
||||
max_tokens=100,
|
||||
stream=True,
|
||||
api_base = "https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions"
|
||||
)
|
||||
|
||||
chunks_received = []
|
||||
for chunk in response:
|
||||
print(chunk)
|
||||
except APIConnectionError as e:
|
||||
# Skip test if Snowflake API is unavailable (502 error)
|
||||
if "Application failed to respond" in str(e) or "502" in str(e):
|
||||
pytest.skip(f"Snowflake API unavailable: {e}")
|
||||
else:
|
||||
raise # Re-raise if it's a different APIConnectionError
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
chunks_received.append(chunk)
|
||||
|
||||
assert len(chunks_received) > 0
|
||||
else:
|
||||
async_handler = AsyncHTTPHandler()
|
||||
mock_chunks = mock_snowflake_streaming_response_chunks()
|
||||
|
||||
async def mock_iter_lines():
|
||||
for chunk in mock_chunks:
|
||||
for line in [f"data: {chunk}", "data: [DONE]"]:
|
||||
yield line
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.iter_lines.side_effect = mock_iter_lines
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch.object(AsyncHTTPHandler, "post", return_value=mock_response):
|
||||
import asyncio
|
||||
|
||||
async def test_async_stream():
|
||||
response = await acompletion(
|
||||
model="snowflake/mistral-7b",
|
||||
messages=messages,
|
||||
max_tokens=100,
|
||||
stream=True,
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions",
|
||||
client=async_handler,
|
||||
)
|
||||
|
||||
chunks_received = []
|
||||
async for chunk in response:
|
||||
chunks_received.append(chunk)
|
||||
|
||||
assert len(chunks_received) > 0
|
||||
|
||||
asyncio.run(test_async_stream())
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Requires Snowflake credentials - run manually when needed")
|
||||
|
||||
Reference in New Issue
Block a user