improvement: only use 'parts' api in llms if necessary (#2797)

marimo-team · Nov 5, 2024 · 643a1e9 · 643a1e9
1 parent 4c20bbc
commit 643a1e9
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 65 deletions.
diff --git 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
@@ -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