From 643a1e902d8341c30967c3790ed03d69ce8ee8f1 Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Tue, 5 Nov 2024 14:15:46 -0500 Subject: [PATCH] improvement: only use 'parts' api in llms if necessary (#2797) --- marimo/_plugins/ui/_impl/chat/convert.py | 94 ++++++++++--------- .../ui/_impl/chat/test_chat_convert.py | 30 ++---- 2 files changed, 59 insertions(+), 65 deletions(-) diff --git a/marimo/_plugins/ui/_impl/chat/convert.py b/marimo/_plugins/ui/_impl/chat/convert.py index 6a00208c829..91fe7c72e5f 100644 --- a/marimo/_plugins/ui/_impl/chat/convert.py +++ b/marimo/_plugins/ui/_impl/chat/convert.py @@ -13,29 +13,31 @@ def convert_to_openai_messages( openai_messages: List[Dict[Any, Any]] = [] for message in messages: - parts: List[Dict[Any, Any]] = [] + if not message.attachments: + openai_messages.append( + {"role": message.role, "content": message.content} + ) + continue + # Handle attachments + parts: List[Dict[Any, Any]] = [] parts.append({"type": "text", "text": message.content}) + for attachment in message.attachments: + content_type = attachment.content_type or "text/plain" - if message.attachments: - for attachment in message.attachments: - content_type = attachment.content_type or "text/plain" - - if content_type.startswith("image"): - parts.append( - { - "type": "image_url", - "image_url": {"url": attachment.url}, - } - ) - elif content_type.startswith("text"): - parts.append( - {"type": "text", "text": _extract_text(attachment.url)} - ) - else: - raise ValueError( - f"Unsupported content type {content_type}" - ) + if content_type.startswith("image"): + parts.append( + { + "type": "image_url", + "image_url": {"url": attachment.url}, + } + ) + elif content_type.startswith("text"): + parts.append( + {"type": "text", "text": _extract_text(attachment.url)} + ) + else: + raise ValueError(f"Unsupported content type {content_type}") openai_messages.append({"role": message.role, "content": parts}) @@ -48,34 +50,36 @@ def convert_to_anthropic_messages( anthropic_messages: List[Dict[Any, Any]] = [] for message in messages: - parts: List[Dict[Any, Any]] = [] + if not message.attachments: + anthropic_messages.append( + {"role": message.role, "content": message.content} + ) + continue + # Handle attachments + parts: List[Dict[Any, Any]] = [] parts.append({"type": "text", "text": message.content}) + for attachment in message.attachments: + content_type = attachment.content_type or "text/plain" + if content_type.startswith("image"): + parts.append( + { + "type": "image", + "source": { + "type": "base64", + "media_type": content_type, + "data": _extract_data(attachment.url), + }, + } + ) - if message.attachments: - for attachment in message.attachments: - content_type = attachment.content_type or "text/plain" - if content_type.startswith("image"): - parts.append( - { - "type": "image", - "source": { - "type": "base64", - "media_type": content_type, - "data": _extract_data(attachment.url), - }, - } - ) - - elif content_type.startswith("text"): - parts.append( - {"type": "text", "text": _extract_text(attachment.url)} - ) - - else: - raise ValueError( - f"Unsupported content type {content_type}" - ) + elif content_type.startswith("text"): + parts.append( + {"type": "text", "text": _extract_text(attachment.url)} + ) + + else: + raise ValueError(f"Unsupported content type {content_type}") anthropic_messages.append({"role": message.role, "content": parts}) diff --git a/tests/_plugins/ui/_impl/chat/test_chat_convert.py b/tests/_plugins/ui/_impl/chat/test_chat_convert.py index 58d584b7693..8182b967bd4 100644 --- a/tests/_plugins/ui/_impl/chat/test_chat_convert.py +++ b/tests/_plugins/ui/_impl/chat/test_chat_convert.py @@ -68,11 +68,10 @@ def test_convert_to_openai_messages(sample_messages: List[ChatMessage]): # Check assistant message assert result[1]["role"] == "assistant" - assert len(result[1]["content"]) == 1 - assert result[1]["content"][0] == { - "type": "text", - "text": "Sure, I'd be happy to help. What's your question?", - } + assert ( + result[1]["content"] + == "Sure, I'd be happy to help. What's your question?" + ) def test_convert_to_anthropic_messages(sample_messages: List[ChatMessage]): @@ -102,11 +101,10 @@ def test_convert_to_anthropic_messages(sample_messages: List[ChatMessage]): # Check assistant message assert result[1]["role"] == "assistant" - assert len(result[1]["content"]) == 1 - assert result[1]["content"][0] == { - "type": "text", - "text": "Sure, I'd be happy to help. What's your question?", - } + assert ( + result[1]["content"] + == "Sure, I'd be happy to help. What's your question?" + ) def test_convert_to_google_messages(sample_messages: List[ChatMessage]): @@ -170,20 +168,12 @@ def test_message_without_attachments(): openai_result = convert_to_openai_messages(messages) assert len(openai_result) == 1 assert openai_result[0]["role"] == "user" - assert len(openai_result[0]["content"]) == 1 - assert openai_result[0]["content"][0] == { - "type": "text", - "text": "Just a simple message", - } + assert openai_result[0]["content"] == "Just a simple message" anthropic_result = convert_to_anthropic_messages(messages) assert len(anthropic_result) == 1 assert anthropic_result[0]["role"] == "user" - assert len(anthropic_result[0]["content"]) == 1 - assert anthropic_result[0]["content"][0] == { - "type": "text", - "text": "Just a simple message", - } + assert anthropic_result[0]["content"] == "Just a simple message" google_result = convert_to_google_messages(messages) assert len(google_result) == 1