Keep rich objects in chat message #7851
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
There was a problem hiding this comment.
Pull request overview
This PR modifies the chat UI to preserve rich Python objects (e.g., altair.Chart, UI elements) in the chat.value property, rather than only storing their HTML string representations. The key change is that while the frontend receives HTML-serialized content for display, the backend maintains the original rich objects in _chat_history and restores them when the frontend sends messages back.
Changes:
- Modified
_send_promptto store rich response objects in chat history instead of just their string representations - Updated
_convert_valueto prefer Python object content from_chat_historyover serialized content from the frontend when reconstructing message history - Changed
_send_promptreturn type fromstr | NonetoNonefor consistency since it now streams all responses
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| marimo/_plugins/ui/_impl/chat/chat.py | Core logic changes to preserve rich objects in _chat_history and restore them in _convert_value |
| marimo/_ai/_types.py | Updated comments to clarify that content can be rich Python objects while parts must be JSON-serializable |
| tests/_plugins/ui/_impl/chat/test_chat.py | Added comprehensive tests for rich object preservation and updated existing tests for new return type |
| marimo/_smoke_tests/ai/chat-variations.py | Added cell to manually test accessing rich content via chat.value[1].content |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if isinstance(msg, ChatMessage): | ||
| if prev_content is not None: | ||
| msg.content = prev_content | ||
| continue |
There was a problem hiding this comment.
The code on lines 410-413 appears to be unreachable dead code. The condition isinstance(msg, ChatMessage) will never be True in practice because the frontend always sends dict messages (as seen in test line 1219-1233). If this case were to occur, the code would also fail to append the message to the result list due to the continue statement on line 413, causing messages to be silently dropped. Consider removing this dead code block entirely.
| if isinstance(msg, ChatMessage): | |
| if prev_content is not None: | |
| msg.content = prev_content | |
| continue |
| def get_prev_content(idx: int) -> Any: | ||
| # Only get the prev content if messages are the same size | ||
| if len(messages) == len(self._chat_history): | ||
| return self._chat_history[idx].content | ||
| return None | ||
|
|
||
| result: list[ChatMessage] = [] | ||
| for i, msg in enumerate(messages): | ||
| prev_content = get_prev_content(i) | ||
| if isinstance(msg, ChatMessage): | ||
| if prev_content is not None: | ||
| msg.content = prev_content | ||
| continue | ||
|
|
||
| msg_id = msg.get("id") | ||
| role = msg.get("role", "user") | ||
| # Prefer the content in Python object format over the serialized content from the frontend, | ||
| # since this is the most accurate representation of the message and more valuable to the user in Python-land. | ||
| content = ( | ||
| prev_content | ||
| if prev_content is not None | ||
| else msg.get("content") | ||
| ) | ||
| result.append( |
There was a problem hiding this comment.
The _convert_value method matches messages between the frontend and backend purely by index position without verifying that message IDs match. This could lead to incorrect content being associated with messages if the frontend ever sends messages in a different order or with a different structure than what's stored in _chat_history. Consider adding ID-based matching or at least validating that the IDs match when both lists have the same length.
|
🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.19.3-dev47 |
This makes sure if the model returns rich content (e.g.
altair.Chart), then we keep those rich messages in thechat.valueto be later accessed.