diff --git a/litellm/utils.py b/litellm/utils.py index 8ef55758ef1..344d28254ec 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7260,6 +7260,16 @@ def shorten_message_to_fit_limit(message, tokens_needed, model: str | None, rais return message +def _extract_system_message_content(content: Union[str, list, Any]) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + return "".join( + part.get("text", "") for part in content if isinstance(part, dict) and part.get("type") == "text" + ) + return "" + + # LiteLLM token trimmer # this code is borrowed from https://github.com/KillianLucas/tokentrim/blob/main/tokentrim/tokentrim.py # Credits for this code go to Killian Lucas @@ -7305,7 +7315,7 @@ def trim_messages( for message in messages: if message["role"] == "system": system_message += "\n" if system_message else "" - system_message += message["content"] + system_message += _extract_system_message_content(message.get("content", "")) ## Handle Tool Call ## - check if last message is a tool response, return as is - https://github.com/BerriAI/litellm/issues/4931 tool_messages: Final = [] diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b0951b2b211..69279aad175 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -58,6 +58,7 @@ from litellm.utils import ( get_prompt_cache_min_tokens, is_cached_message, is_prompt_caching_valid_prompt, + trim_messages, ) # Adds the parent directory to the system path @@ -122,6 +123,26 @@ def test_usage_openai_cache_write_tokens_populates_both_names(): assert usage.prompt_tokens_details.cache_creation_tokens == 800 +def test_trim_messages_system_message_list_content(): + messages = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Always answer in French."}, + ], + }, + {"role": "user", "content": "Hello there, how are you?"}, + ] + + trimmed_messages = trim_messages(messages, model="gpt-4", max_tokens=5) + + assert ( + trimmed_messages[0]["content"] + == "You are a helpful assistant.Always answer in French." + ) + + def test_usage_anthropic_cache_creation_maps_to_cache_write_tokens(): """Anthropic/Bedrock report the top-level cache_creation_input_tokens field. It must be normalized onto the OpenAI cache_write_tokens name as well as the