feat(vertex_httpx.py): support the 'response_schema' param for older vertex ai models - pass as prompt (user-controlled)

if 'response_schema' is not supported for vertex model (e.g. gemini-1.5-flash) pass in prompt

litellm/__init__.py

@@ -749,6 +749,7 @@ from .utils import (
     create_pretrained_tokenizer,
     create_tokenizer,
     supports_function_calling,
+    supports_response_schema,
     supports_parallel_function_calling,
     supports_vision,
     supports_system_messages,

litellm/llms/prompt_templates/factory.py

@@ -2033,6 +2033,50 @@ def function_call_prompt(messages: list, functions: list):
     return messages
 
 
+def response_schema_prompt(model: str, response_schema: dict) -> str:
+    """
+    Decides if a user-defined custom prompt or default needs to be used
+
+    Returns the prompt str that's passed to the model as a user message
+    """
+    custom_prompt_details: Optional[dict] = None
+    response_schema_as_message = [
+        {"role": "user", "content": "{}".format(response_schema)}
+    ]
+    if f"{model}/response_schema_prompt" in litellm.custom_prompt_dict:
+
+        custom_prompt_details = litellm.custom_prompt_dict[
+            f"{model}/response_schema_prompt"
+        ]  # allow user to define custom response schema prompt by model
+    elif "response_schema_prompt" in litellm.custom_prompt_dict:
+        custom_prompt_details = litellm.custom_prompt_dict["response_schema_prompt"]
+
+    if custom_prompt_details is not None:
+        return custom_prompt(
+            role_dict=custom_prompt_details["roles"],
+            initial_prompt_value=custom_prompt_details["initial_prompt_value"],
+            final_prompt_value=custom_prompt_details["final_prompt_value"],
+            messages=response_schema_as_message,
+        )
+    else:
+        return default_response_schema_prompt(response_schema=response_schema)
+
+
+def default_response_schema_prompt(response_schema: dict) -> str:
+    """
+    Used if provider/model doesn't support 'response_schema' param.
+
+    This is the default prompt. Allow user to override this with a custom_prompt.
+    """
+    prompt_str = """Use this JSON schema: 
+    ```json 
+    {}
+    ```""".format(
+        response_schema
+    )
+    return prompt_str
+
+
 # Custom prompt template
 def custom_prompt(
     role_dict: dict,

litellm/llms/vertex_ai.py

@@ -12,6 +12,7 @@ import requests  # type: ignore
 from pydantic import BaseModel
 
 import litellm
+from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.core_helpers import map_finish_reason
 from litellm.llms.prompt_templates.factory import (
     convert_to_anthropic_image_obj,
@@ -328,11 +329,14 @@ def _gemini_convert_messages_with_history(messages: list) -> List[ContentType]:
     contents: List[ContentType] = []
 
     msg_i = 0
+    try:
         while msg_i < len(messages):
             user_content: List[PartType] = []
             init_msg_i = msg_i
             ## MERGE CONSECUTIVE USER CONTENT ##
-        while msg_i < len(messages) and messages[msg_i]["role"] in user_message_types:
+            while (
+                msg_i < len(messages) and messages[msg_i]["role"] in user_message_types
+            ):
                 if isinstance(messages[msg_i]["content"], list):
                     _parts: List[PartType] = []
                     for element in messages[msg_i]["content"]:
@@ -375,7 +379,9 @@ def _gemini_convert_messages_with_history(messages: list) -> List[ContentType]:
                     "tool_calls", []
                 ):  # support assistant tool invoke conversion
                     assistant_content.extend(
-                    convert_to_gemini_tool_call_invoke(messages[msg_i]["tool_calls"])
+                        convert_to_gemini_tool_call_invoke(
+                            messages[msg_i]["tool_calls"]
+                        )
                     )
                 else:
                     assistant_text = (
@@ -400,8 +406,9 @@ def _gemini_convert_messages_with_history(messages: list) -> List[ContentType]:
                         messages[msg_i]
                     )
                 )
-
         return contents
+    except Exception as e:
+        raise e
 
 
 def _get_client_cache_key(model: str, vertex_project: str, vertex_location: str):

litellm/llms/vertex_httpx.py

@@ -21,7 +21,10 @@ import litellm.litellm_core_utils.litellm_logging
 from litellm import verbose_logger
 from litellm.litellm_core_utils.core_helpers import map_finish_reason
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
-from litellm.llms.prompt_templates.factory import convert_url_to_base64
+from litellm.llms.prompt_templates.factory import (
+    convert_url_to_base64,
+    response_schema_prompt,
+)
 from litellm.llms.vertex_ai import _gemini_convert_messages_with_history
 from litellm.types.llms.openai import (
     ChatCompletionResponseMessage,
@@ -1011,6 +1014,22 @@ class VertexLLM(BaseLLM):
             if len(system_prompt_indices) > 0:
                 for idx in reversed(system_prompt_indices):
                     messages.pop(idx)
+
+        # Checks for 'response_schema' support - if passed in
+        if "response_schema" in optional_params:
+            supports_response_schema = litellm.supports_response_schema(
+                model=model, custom_llm_provider="vertex_ai"
+            )
+            if supports_response_schema is False:
+                user_response_schema_message = response_schema_prompt(
+                    model=model, response_schema=optional_params.get("response_schema")  # type: ignore
+                )
+                messages.append(
+                    {"role": "user", "content": user_response_schema_message}
+                )
+                optional_params.pop("response_schema")
+
+        try:
             content = _gemini_convert_messages_with_history(messages=messages)
             tools: Optional[Tools] = optional_params.pop("tools", None)
             tool_choice: Optional[ToolConfig] = optional_params.pop("tool_choice", None)
@@ -1040,6 +1059,8 @@ class VertexLLM(BaseLLM):
                 headers["Authorization"] = f"Bearer {auth_header}"
             if extra_headers is not None:
                 headers.update(extra_headers)
+        except Exception as e:
+            raise e
 
         ## LOGGING
         logging_obj.pre_call(

litellm/model_prices_and_context_window_backup.json

@@ -1538,6 +1538,7 @@
         "supports_system_messages": true,
         "supports_function_calling": true,
         "supports_tool_choice": true, 
+        "supports_response_schema": true, 
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
     },
     "gemini-1.5-pro-preview-0215": { 
@@ -1563,6 +1564,7 @@
         "supports_system_messages": true,
         "supports_function_calling": true,
         "supports_tool_choice": true, 
+        "supports_response_schema": true, 
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
     },
     "gemini-1.5-pro-preview-0409": {
@@ -1587,6 +1589,7 @@
         "mode": "chat",
         "supports_function_calling": true,
         "supports_tool_choice": true,
+        "supports_response_schema": true, 
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
     },
     "gemini-1.5-flash": {

litellm/tests/test_amazing_vertex_completion.py

@@ -880,10 +880,19 @@ Using this JSON schema:
         mock_call.assert_called_once()
 
 
-@pytest.mark.parametrize("provider", ["vertex_ai_beta"])  # "vertex_ai",
+@pytest.mark.parametrize(
+    "model, supports_response_schema",
+    [
+        ("vertex_ai_beta/gemini-1.5-pro-001", True),
+        ("vertex_ai_beta/gemini-1.5-flash", False),
+    ],
+)  # "vertex_ai",
 @pytest.mark.asyncio
-async def test_gemini_pro_json_schema_httpx(provider):
+async def test_gemini_pro_json_schema_httpx(model, supports_response_schema):
     load_vertex_ai_credentials()
+    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+    litellm.model_cost = litellm.get_model_cost_map(url="")
+
     litellm.set_verbose = True
     messages = [{"role": "user", "content": "List 5 cookie recipes"}]
     from litellm.llms.custom_httpx.http_handler import HTTPHandler
@@ -905,8 +914,8 @@ async def test_gemini_pro_json_schema_httpx(provider):
 
     with patch.object(client, "post", new=MagicMock()) as mock_call:
         try:
-            response = completion(
-                model="vertex_ai_beta/gemini-1.5-pro-001",
+            _ = completion(
+                model=model,
                 messages=messages,
                 response_format={
                     "type": "json_object",
@@ -914,14 +923,26 @@ async def test_gemini_pro_json_schema_httpx(provider):
                 },
                 client=client,
             )
-        except Exception as e:
+        except Exception:
             pass
 
         mock_call.assert_called_once()
         print(mock_call.call_args.kwargs)
         print(mock_call.call_args.kwargs["json"]["generationConfig"])
+
+        if supports_response_schema:
             assert (
-            "response_schema" in mock_call.call_args.kwargs["json"]["generationConfig"]
+                "response_schema"
+                in mock_call.call_args.kwargs["json"]["generationConfig"]
+            )
+        else:
+            assert (
+                "response_schema"
+                not in mock_call.call_args.kwargs["json"]["generationConfig"]
+            )
+            assert (
+                "Use this JSON schema:"
+                in mock_call.call_args.kwargs["json"]["contents"][0]["parts"][1]["text"]
             )
 
 

litellm/utils.py

@@ -1879,8 +1879,7 @@ def supports_response_schema(model: str, custom_llm_provider: Optional[str]) ->
     Returns:
     bool: True if the model supports response_schema, False otherwise.
 
-    Raises:
-    Exception: If the given model is not found in model_prices_and_context_window.json.
+    Does not raise error. Defaults to 'False'. Outputs logging.error.
     """
     try:
         ## GET LLM PROVIDER ##
@@ -1900,9 +1899,10 @@ def supports_response_schema(model: str, custom_llm_provider: Optional[str]) ->
             return True
         return False
     except Exception:
-        raise Exception(
+        verbose_logger.error(
             f"Model not in model_prices_and_context_window.json. You passed model={model}, custom_llm_provider={custom_llm_provider}."
         )
+        return False
 
 
 def supports_function_calling(model: str) -> bool:

model_prices_and_context_window.json

@@ -1538,6 +1538,7 @@
         "supports_system_messages": true,
         "supports_function_calling": true,
         "supports_tool_choice": true, 
+        "supports_response_schema": true, 
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
     },
     "gemini-1.5-pro-preview-0215": { 
@@ -1563,6 +1564,7 @@
         "supports_system_messages": true,
         "supports_function_calling": true,
         "supports_tool_choice": true, 
+        "supports_response_schema": true, 
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
     },
     "gemini-1.5-pro-preview-0409": {
@@ -1587,6 +1589,7 @@
         "mode": "chat",
         "supports_function_calling": true,
         "supports_tool_choice": true,
+        "supports_response_schema": true, 
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
     },
     "gemini-1.5-flash": {