diff --git a/tests/llm_translation/test_together_ai.py b/tests/llm_translation/test_together_ai.py index c371caefa5e..174c8187374 100644 --- a/tests/llm_translation/test_together_ai.py +++ b/tests/llm_translation/test_together_ai.py @@ -23,26 +23,20 @@ class TestTogetherAI(BaseLLMChatTest): pass @pytest.mark.parametrize( - "model, expected_bool", + "model, supports_response_format", [ ("meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo", True), ("nvidia/Llama-3.1-Nemotron-70B-Instruct-HF", False), ], ) def test_get_supported_response_format_together_ai( - self, model: str, expected_bool: bool + self, model: str, supports_response_format: bool ) -> None: os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") optional_params = litellm.get_supported_openai_params( model, custom_llm_provider="together_ai" ) - # Mapped provider assert isinstance(optional_params, list) - - if expected_bool: - assert "response_format" in optional_params - assert "tools" in optional_params - else: - assert "response_format" not in optional_params - assert "tools" not in optional_params + assert ("response_format" in optional_params) is supports_response_format + assert "tools" in optional_params