From f20a8ff3647c051333f8809bc06f1e66e892ca9c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 11 Jul 2024 13:04:18 -0700 Subject: [PATCH] test get mode info for gemini/gemini-1.5-flash --- litellm/tests/test_get_model_info.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_get_model_info.py b/litellm/tests/test_get_model_info.py index 3fd6a6d22f..687aa062f1 100644 --- a/litellm/tests/test_get_model_info.py +++ b/litellm/tests/test_get_model_info.py @@ -1,13 +1,16 @@ # What is this? ## Unit testing for the 'get_model_info()' function -import os, sys, traceback +import os +import sys +import traceback sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path +import pytest + import litellm from litellm import get_model_info -import pytest def test_get_model_info_simple_model_name(): @@ -37,3 +40,9 @@ def test_get_model_info_custom_llm_with_same_name_vllm(): pytest.fail("Expected get model info to fail for an unmapped model/provider") except Exception: pass + + +def test_get_model_info_shows_correct_supports_vision(): + info = litellm.get_model_info("gemini/gemini-1.5-flash") + print("info", info) + assert info["supports_vision"] is True