Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,7 +1,6 @@
 import streamlit as st
 from chat_langraph import system, workflow, HumanMessage, AIMessage, get_all_chat_ids, ToolMessage
 import uuid
-import re
 import os
 import base64
 
@@ -9,8 +8,8 @@ import base64
 TEMP_DIR = "/tmp"
 os.makedirs(TEMP_DIR, exist_ok=True)
 
+st.set_page_config(layout="wide")
 st.title("MY CHATBOT")
-st.set_page_config(layout='wide')
 
 
 def set_title(messages):
@@ -20,7 +19,7 @@ def set_title(messages):
 
 
 def set_config():
-    return {'configurable': {'thread_id': st.session_state.current_chat_id}}
+    return {"configurable": {"thread_id": st.session_state.current_chat_id}}
 
 
 def load_session_state():
@@ -44,7 +43,7 @@ def render_sidebar():
             new_id = str(uuid.uuid4())
             st.session_state.chats.append(new_id)
             st.session_state.current_chat_id = new_id
-            config = {'configurable': {'thread_id': new_id}}
+            config = {"configurable": {"thread_id": new_id}}
             workflow.update_state(config, {"messages": [system]})
             st.session_state.chat_dict[new_id] = "New Chat"
         for chat_id in st.session_state.chats:
@@ -55,7 +54,7 @@ def render_sidebar():
 def loadchats():
     if "current_chat_id" not in st.session_state:
         return []
-    config = {'configurable': {'thread_id': st.session_state.current_chat_id}}
+    config = {"configurable": {"thread_id": st.session_state.current_chat_id}}
     state = workflow.get_state(config)
     messages = state.values.get("messages", [])
     for message in messages:
@@ -68,8 +67,7 @@ def loadchats():
                 with st.chat_message("assistant"):
                     st.write(message.content)
         elif isinstance(message, ToolMessage):
-            with st.chat_message("assistant"):
-                st.info(f"Used a tool")
+            render_tool_message(message)
     return messages
 
 
@@ -84,13 +82,25 @@ 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)
+def render_tool_message(tool_message: ToolMessage):
+    """Render tool output including download links for any created files."""
+    with st.chat_message("assistant"):
+        st.info("Tool executed")
 
+        # If tool_message has an explicit "file" field in content/metadata
+        if hasattr(tool_message, "content") and isinstance(tool_message.content, dict):
+            file_paths = tool_message.content.get("files", [])
+            for file_path in file_paths:
+                if os.path.exists(file_path):
+                    if file_path.lower().endswith((".png", ".jpg", ".jpeg")):
+                        st.image(file_path, caption=os.path.basename(file_path))
+                    st.markdown(
+                        create_download_link(file_path, f"📥 Download {os.path.basename(file_path)}"),
+                        unsafe_allow_html=True,
+                    )
 
+
+# --- Main App Logic ---
 load_session_state()
 render_sidebar()
 
@@ -104,26 +114,15 @@ if "current_chat_id" in st.session_state:
         with st.chat_message("assistant"):
             response_placeholder = st.empty()
             full_response = ""
-            for messages, metadata in workflow.stream(
+            for message, metadata in workflow.stream(
                 {"messages": [system, HumanMessage(user_input)]},
-                config={'configurable': {'thread_id': st.session_state.current_chat_id}},
-                stream_mode="messages"
+                config={"configurable": {"thread_id": st.session_state.current_chat_id}},
+                stream_mode="messages",
             ):
-                if isinstance(messages, AIMessage) and messages.content.strip():
-                    full_response += messages.content
-
-                if isinstance(messages, ToolMessage) and messages.content.strip():
-                    st.info(f"Using appropriate tool")
-
-                    # 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()
-                        if os.path.exists(file_name):
-                            show_image_with_download(file_name)
+                if isinstance(message, AIMessage) and message.content.strip():
+                    full_response += message.content
+
+                if isinstance(message, ToolMessage):
+                    render_tool_message(message)
 
             response_placeholder.markdown(full_response)