From 72c197d74f644ba0751f9592df027e4c6bfcca61 Mon Sep 17 00:00:00 2001 From: xbrxr03 Date: Mon, 22 Jun 2026 02:40:20 -0400 Subject: [PATCH] refactor: address review feedback on invalid role validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Move VALID_MESSAGE_ROLES to constants.py per project convention - Remove unused idx variable from enumerate - Fix weak OR assertion — split into two independent assertions with descriptive error messages --- litellm/constants.py | 3 +++ litellm/utils.py | 5 ++--- tests/litellm_utils_tests/test_invalid_role_validation.py | 4 +++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c0e265c0e4a..5c1258d4e4d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1791,3 +1791,6 @@ ADVISOR_TOOL_DESCRIPTION: str = ( "want to verify your reasoning, or face a complex decision. " "Describe your question or challenge clearly in the 'question' field." ) + +# Valid roles for OpenAI chat completion messages +VALID_MESSAGE_ROLES = {"system", "user", "assistant", "tool", "function", "developer"} diff --git a/litellm/utils.py b/litellm/utils.py index 99db50caf95..ae56d93ff15 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -73,6 +73,7 @@ from litellm.constants import ( MINIMUM_PROMPT_CACHE_TOKEN_COUNT, OPENAI_EMBEDDING_PARAMS, TOOL_CHOICE_OBJECT_TOKEN_COUNT, + VALID_MESSAGE_ROLES, ) _CachingHandlerResponse = None @@ -8211,8 +8212,6 @@ def convert_list_message_to_dict(messages: List): return new_messages -VALID_MESSAGE_ROLES = {"system", "user", "assistant", "tool", "function", "developer"} - def validate_and_fix_openai_messages(messages: List): """ @@ -8221,7 +8220,7 @@ def validate_and_fix_openai_messages(messages: List): Handles missing role for assistant messages and rejects invalid roles. """ new_messages = [] - for idx, message in enumerate(messages): + for message in messages: if not message.get("role"): message["role"] = "assistant" elif message["role"] not in VALID_MESSAGE_ROLES: diff --git a/tests/litellm_utils_tests/test_invalid_role_validation.py b/tests/litellm_utils_tests/test_invalid_role_validation.py index ee2b53ceed5..1778ee86a0d 100644 --- a/tests/litellm_utils_tests/test_invalid_role_validation.py +++ b/tests/litellm_utils_tests/test_invalid_role_validation.py @@ -67,7 +67,9 @@ class TestInvalidRoleValidation: error_text = str(exc_info.value) # No traceback leakage assert "Traceback" not in error_text - assert ".py\"" not in error_text or "litellm" not in error_text + # No internal path disclosure + assert ".py\"" not in error_text, f"Internal file path leaked: {error_text}" + assert "litellm" not in error_text or "Invalid role" in error_text, f"litellm internal path leaked: {error_text}" # No internal path disclosure assert "/usr/lib/" not in error_text assert "site-packages" not in error_text