Update src/streamlit_app.py

src/streamlit_app.py

@@ -85,6 +85,13 @@ def create_download_link(file_path: str, label: str = "Download file") -> str:
         return f"Error creating download link: {e}"
 
 
+def show_image_with_download(file_path: str):
+    """Display image in chat and show a download button."""
+    if os.path.exists(file_path) and file_path.endswith(".png"):
+        st.image(file_path, caption=os.path.basename(file_path))
+        st.markdown(create_download_link(file_path, "📥 Download Image"), unsafe_allow_html=True)
+
+
 load_session_state()
 render_sidebar()
 
@@ -109,18 +116,15 @@ if "current_chat_id" in st.session_state:
                 if isinstance(messages, ToolMessage) and messages.content.strip():
                     st.info(f"Using appropriate tool")
 
-                    # If tool output contains file creation info → show download link
+                    # Detect generated PNG files and show + download them
+                    file_matches = re.findall(r"/tmp/[\w\.\-_]+\.png", messages.content)
+                    for file_path in file_matches:
+                        show_image_with_download(file_path)
+
+                    # Existing logic for "Content saved to"
                     if "Content saved to" in messages.content:
                         file_name = messages.content.replace("Content saved to ", "").strip()
-                        file_path = os.path.join(TEMP_DIR, file_name)
-                        if os.path.exists(file_path):
-                            st.markdown(
-                                create_download_link(file_path, "📥 Download File"),
-                                unsafe_allow_html=True
-                            )
+                        if os.path.exists(file_name):
+                            show_image_with_download(file_name)
 
             response_placeholder.markdown(full_response)
-
-
-
-