phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(vertex_ai.py): handle nested content dictionary for assistant message

Browse source 
openai allows the assistant content message to also be a list of dictionaries, handle that
This commit is contained in:
Krrish Dholakia 2024-05-24 22:47:34 -07:00
parent 799c63bf8e
commit 281906ff33
2 changed files with 44 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

26
litellm/llms/vertex_ai.py
View file

@ -376,17 +376,31 @@ def _gemini_convert_messages_with_history(messages: list) -> List[ContentType]:
assistant_content = []
## MERGE CONSECUTIVE ASSISTANT CONTENT ##
while msg_i < len(messages) and messages[msg_i]["role"] == "assistant":
assistant_text = (
messages[msg_i].get("content") or ""
) # either string or none
if assistant_text:
assistant_content.append(PartType(text=assistant_text))
if messages[msg_i].get(
if isinstance(messages[msg_i]["content"], list):
_parts = []
for element in messages[msg_i]["content"]:
if isinstance(element, dict):
if element["type"] == "text":
_part = PartType(text=element["text"])
_parts.append(_part)
elif element["type"] == "image_url":
image_url = element["image_url"]["url"]
_part = _process_gemini_image(image_url=image_url)
_parts.append(_part) # type: ignore
assistant_content.extend(_parts)
elif messages[msg_i].get(
"tool_calls", []
): # support assistant tool invoke convertion
assistant_content.extend(
convert_to_gemini_tool_call_invoke(messages[msg_i]["tool_calls"])
)
else:
assistant_text = (
messages[msg_i].get("content") or ""
) # either string or none
if assistant_text:
assistant_content.append(PartType(text=assistant_text))
msg_i += 1
if assistant_content:

24
litellm/tests/test_amazing_vertex_completion.py
View file

@ -975,3 +975,27 @@ def test_prompt_factory():
translated_messages = _gemini_convert_messages_with_history(messages=messages)
print(f"\n\ntranslated_messages: {translated_messages}\ntranslated_messages")
def test_prompt_factory_nested():
messages = [
{"role": "user", "content": [{"type": "text", "text": "hi"}]},
{
"role": "assistant",
"content": [
{"type": "text", "text": "Hi! 👋 \n\nHow can I help you today? 😊 \n"}
],
},
{"role": "user", "content": [{"type": "text", "text": "hi 2nd time"}]},
]
translated_messages = _gemini_convert_messages_with_history(messages=messages)
print(f"\n\ntranslated_messages: {translated_messages}\ntranslated_messages")
for message in translated_messages:
assert len(message["parts"]) == 1
assert "text" in message["parts"][0], "Missing 'text' from 'parts'"
assert isinstance(
message["parts"][0]["text"], str
), "'text' value not a string."