refactor: address review feedback on invalid role validation

- 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
This commit is contained in:
xbrxr03 2026-06-22 02:40:20 -04:00
parent e165c29d5c
commit 72c197d74f
3 changed files with 8 additions and 4 deletions

View file

@ -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"}

View file

@ -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:

View file

@ -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