From bcbc44d731c00788d31c32644d794266e9192386 Mon Sep 17 00:00:00 2001
From: Cole Murray <colemurray.cs@gmail.com>
Date: Thu, 26 Sep 2024 10:26:42 -0700
Subject: [PATCH] Add Support for Custom Providers in Vision and Function Call
 Utils (#5688)

* Add Support for Custom Providers in Vision and Function Call Utils Lookup

* Remove parallel function call due to missing model info param

* Add Unit Tests for Vision and Function Call Changes
---
 litellm/tests/test_utils.py | 34 +++++++++++++++++++++++++++++
 litellm/utils.py            | 43 ++++++++++++++++++++++++-------------
 2 files changed, 62 insertions(+), 15 deletions(-)

diff --git a/litellm/tests/test_utils.py b/litellm/tests/test_utils.py
index 75c31c30230..060e69e7690 100644
--- a/litellm/tests/test_utils.py
+++ b/litellm/tests/test_utils.py
@@ -762,6 +762,40 @@ def test_supports_response_schema(model, expected_bool):
 
     assert expected_bool == response
 
+@pytest.mark.parametrize(
+    "model, expected_bool",
+    [
+        ("gpt-3.5-turbo", True),
+        ("gpt-4", True),
+        ("command-nightly", False),
+        ("gemini-pro", True),
+    ],
+)
+def test_supports_function_calling(model, expected_bool):
+    """
+    Unit test for 'supports_function_calling' helper function.
+    """
+    from litellm.utils import supports_function_calling
+    response = supports_function_calling(model=model, custom_llm_provider=None)
+    assert expected_bool == response
+
+@pytest.mark.parametrize(
+    "model, expected_bool",
+    [
+        ("gpt-4-vision-preview", True),
+        ("gpt-3.5-turbo", False),
+        ("claude-3-opus-20240229", True),
+        ("gemini-pro-vision", True),
+        ("command-nightly", False),
+    ],
+)
+def test_supports_vision(model, expected_bool):
+    """
+    Unit test for 'supports_vision' helper function.
+    """
+    from litellm.utils import supports_vision
+    response = supports_vision(model=model, custom_llm_provider=None)
+    assert expected_bool == response
 
 def test_usage_object_null_tokens():
     """
diff --git a/litellm/utils.py b/litellm/utils.py
index a63c25393ac..1a6d799fb67 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -2148,50 +2148,64 @@ def supports_response_schema(model: str, custom_llm_provider: Optional[str]) ->
         return False
 
 
-def supports_function_calling(model: str) -> bool:
+def supports_function_calling(model: str, custom_llm_provider: Optional[str] = None) -> bool:
     """
     Check if the given model supports function calling and return a boolean value.
 
     Parameters:
     model (str): The model name to be checked.
+    custom_llm_provider (Optional[str]): The provider to be checked.
 
     Returns:
     bool: True if the model supports function calling, False otherwise.
 
     Raises:
-    Exception: If the given model is not found in model_prices_and_context_window.json.
+    Exception: If the given model is not found or there's an error in retrieval.
     """
+    try:
+        model, custom_llm_provider, _, _ = litellm.get_llm_provider(
+            model=model, custom_llm_provider=custom_llm_provider
+        )
+
+        model_info = litellm.get_model_info(
+            model=model, custom_llm_provider=custom_llm_provider
+        )
 
-    if model in litellm.model_cost:
-        model_info = litellm.model_cost[model]
         if model_info.get("supports_function_calling", False) is True:
             return True
         return False
-    else:
+    except Exception as e:
         raise Exception(
-            f"Model not supports function calling. You passed model={model}."
+            f"Model not found or error in checking function calling support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {str(e)}"
         )
 
-
-def supports_vision(model: str):
+def supports_vision(model: str, custom_llm_provider: Optional[str] = None) -> bool:
     """
     Check if the given model supports vision and return a boolean value.
 
     Parameters:
     model (str): The model name to be checked.
+    custom_llm_provider (Optional[str]): The provider to be checked.
 
     Returns:
     bool: True if the model supports vision, False otherwise.
-
-    Raises:
-    Exception: If the given model is not found in model_prices_and_context_window.json.
     """
-    if model in litellm.model_cost:
-        model_info = litellm.model_cost[model]
+    try:
+        model, custom_llm_provider, _, _ = litellm.get_llm_provider(
+            model=model, custom_llm_provider=custom_llm_provider
+        )
+
+        model_info = litellm.get_model_info(
+            model=model, custom_llm_provider=custom_llm_provider
+        )
+
         if model_info.get("supports_vision", False) is True:
             return True
         return False
-    else:
+    except Exception as e:
+        verbose_logger.error(
+            f"Model not found or error in checking vision support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {str(e)}"
+        )
         return False
 
 
@@ -2218,7 +2232,6 @@ def supports_parallel_function_calling(model: str):
             f"Model not supports parallel function calling. You passed model={model}."
         )
 
-
 ####### HELPER FUNCTIONS ################
 def _update_dictionary(existing_dict: Dict, new_dict: dict) -> dict:
     for k, v in new_dict.items():