From 88e93f9f64492cc6e9872c51feb3dbd37b0f46c2 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 6 Feb 2025 09:11:19 -0800
Subject: [PATCH] fix(utils.py): handle key error in msg validation

---
 litellm/utils.py                        |  4 ++++
 tests/litellm_utils_tests/test_utils.py | 12 ++++++++++++
 2 files changed, 16 insertions(+)

diff --git a/litellm/utils.py b/litellm/utils.py
index 7e66ad2b2203..d18179151575 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -5884,6 +5884,10 @@ def validate_chat_completion_user_messages(messages: List[AllMessageValues]):
                                 if item.get("type") not in ValidUserMessageContentTypes:
                                     raise Exception("invalid content type")
         except Exception as e:
+            if isinstance(e, KeyError):
+                raise Exception(
+                    f"Invalid message={m} at index {idx}. Please ensure all messages are valid OpenAI chat completion messages."
+                )
             if "invalid content type" in str(e):
                 raise Exception(
                     f"Invalid user message={m} at index {idx}. Please ensure all user messages are valid OpenAI chat completion messages."
diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py
index 4f5e1e2737ac..4a2f63b51dbf 100644
--- a/tests/litellm_utils_tests/test_utils.py
+++ b/tests/litellm_utils_tests/test_utils.py
@@ -1850,3 +1850,15 @@ def test_dict_to_response_format_helper():
         "ref_template": "/$defs/{model}",
     }
     _dict_to_response_format_helper(**args)
+
+
+def test_validate_user_messages_invalid_content_type():
+    from litellm.utils import validate_chat_completion_user_messages
+
+    messages = [{"content": [{"type": "invalid_type", "text": "Hello"}]}]
+
+    with pytest.raises(Exception) as e:
+        validate_chat_completion_user_messages(messages)
+
+    assert "Invalid message" in str(e)
+    print(e)