test: mock fireworks ai test - unstable api

This commit is contained in:
Krrish Dholakia 2025-01-22 18:52:11 -08:00
parent 55546f403b
commit 049915c14b

View file

@ -93,6 +93,57 @@ class TestFireworksAIChatCompletion(BaseLLMChatTest):
"""
pass
@pytest.mark.parametrize(
"response_format",
[
{"type": "json_object"},
{"type": "text"},
],
)
@pytest.mark.flaky(retries=6, delay=1)
def test_json_response_format(self, response_format):
"""
Test that the JSON response format is supported by the LLM API
"""
from litellm.utils import supports_response_schema
from openai import OpenAI
from unittest.mock import patch
client = OpenAI()
base_completion_call_args = self.get_base_completion_call_args()
litellm.set_verbose = True
messages = [
{
"role": "system",
"content": "Your output should be a JSON object with no additional properties. ",
},
{
"role": "user",
"content": "Respond with this in json. city=San Francisco, state=CA, weather=sunny, temp=60",
},
]
with patch.object(
client.chat.completions.with_raw_response, "create"
) as mock_post:
response = self.completion_function(
**base_completion_call_args,
messages=messages,
response_format=response_format,
client=client,
)
mock_post.assert_called_once()
if response_format["type"] == "json_object":
assert (
mock_post.call_args.kwargs["response_format"]["type"]
== "json_object"
)
else:
assert mock_post.call_args.kwargs["response_format"]["type"] == "text"
class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest):
def get_base_audio_transcription_call_args(self) -> dict: