Render data visualizations in the chat

Author: chriscarrollsmith

routers/chat.py

@@ -338,27 +338,42 @@ async def iterate_stream(s, response_id: str = "") -> AsyncGenerator[str, None]:
                                     yield sse_format("textDelta", wrap_for_oob_swap(current_item_id, event.delta))
 
                             case ResponseOutputTextAnnotationAddedEvent():
-                                if event.annotation and current_item_id:
+                                # Use event.item_id (not current_item_id) because annotations
+                                # may fire after other output items (e.g. code interpreter)
+                                # have changed current_item_id to a different element.
+                                annotation_target_id = event.item_id or current_item_id
+                                if event.annotation and annotation_target_id:
                                     if event.annotation["type"] == "file_citation":
                                         filename = event.annotation["filename"]
                                         # Emit a literal HTML anchor to avoid markdown parsing edge cases
                                         encoded_filename = url_quote(filename, safe="")
                                         file_url_path = files_router.url_path_for("download_stored_file", file_name=encoded_filename)
                                         citation = f"(<a href=\"{file_url_path}\">†</a>)"
-                                        yield sse_format("textDelta", wrap_for_oob_swap(current_item_id, citation))
+                                        yield sse_format("textDelta", wrap_for_oob_swap(annotation_target_id, citation))
                                     elif event.annotation["type"] == "container_file_citation":
                                         container_id = event.annotation["container_id"]
                                         file_id = event.annotation["file_id"]
-                                        file = await client.containers.files.retrieve(file_id, container_id=container_id)
-                                        container_file_path = file.path
+                                        filename = event.annotation.get("filename", "")
                                         file_url_path = files_router.url_path_for("download_container_file", container_id=container_id, file_id=file_id)
-                                        replacement_payload = f"sandbox:{container_file_path}|{file_url_path}"
-                                        yield sse_format("textReplacement", wrap_for_oob_swap(current_item_id, replacement_payload))
+                                        # Check if the file is an image by extension
+                                        image_extensions = (".png", ".jpg", ".jpeg", ".gif", ".svg", ".webp")
+                                        if filename.lower().endswith(image_extensions):
+                                            img_html = (
+                                                f'<div class="imageOutput">'
+                                                f'<img src="{file_url_path}" alt="Code interpreter output" />'
+                                                f'</div>'
+                                            )
+                                            yield sse_format("imageOutput", img_html)
+                                        else:
+                                            file = await client.containers.files.retrieve(file_id, container_id=container_id)
+                                            container_file_path = file.path
+                                            replacement_payload = f"sandbox:{container_file_path}|{file_url_path}"
+                                            yield sse_format("textReplacement", wrap_for_oob_swap(annotation_target_id, replacement_payload))
                                     elif event.annotation["type"] == "url_citation":
                                         url = event.annotation["url"]
                                         title = event.annotation.get("title", url)
                                         citation = f'(<a href="{escape(url)}" target="_blank" rel="noopener noreferrer">{escape(title)}</a>)'
-                                        yield sse_format("textDelta", wrap_for_oob_swap(current_item_id, citation))
+                                        yield sse_format("textDelta", wrap_for_oob_swap(annotation_target_id, citation))
                                     else:
                                         logger.error(f"Unhandled annotation type: {event.annotation['type']}")
 

routers/files.py

@@ -1,4 +1,5 @@
 import logging
+import mimetypes
 from typing import Literal
 from fastapi import (
     APIRouter, Request, UploadFile, File, HTTPException, Depends, Path, Form
@@ -248,14 +249,23 @@ async def download_container_file(
         client.base_url = f"https://api.openai.com/v1/containers/{container_id}"
         file_content = await client.files.content(file_id)
         client.base_url = "https://api.openai.com/v1"
-        
+
         if not hasattr(file_content, 'content'):
             raise HTTPException(status_code=500, detail="File content not available")
-            
-        # Use stream_file_content helper
+
+        filename = file.path.split("/")[-1] or file_id
+        # Serve images inline with correct Content-Type so <img> tags work
+        mime_type, _ = mimetypes.guess_type(filename)
+        if mime_type and mime_type.startswith("image/"):
+            return StreamingResponse(
+                stream_file_content(file_content.content),
+                media_type=mime_type,
+                headers={"Content-Disposition": f'inline; filename="{filename}"'}
+            )
+
         return StreamingResponse(
-            stream_file_content(file_content.content), # Assuming stream_file_content handles bytes
-            headers={"Content-Disposition": f'attachment; filename="{file.path.split("/")[-1] or file_id}"'}
+            stream_file_content(file_content.content),
+            headers={"Content-Disposition": f'attachment; filename="{filename}"'}
         )
     except Exception as e:
         logger.error(f"Error downloading file {file_id} from OpenAI: {e}")

static/styles.css

@@ -542,7 +542,8 @@ pre {
   overflow-wrap: break-word;
 }
 
-.assistantMessage img {
+.assistantMessage img,
+.imageOutput img {
   max-width: 100%;
   margin: 8px 0px 8px 0px;
   border-radius: 8px;