Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from chat_langraph import system, workflow, HumanMessage, AIMessage, get_all_chat_ids, ToolMessage , SystemMessage
+from chat_langraph import system, workflow, HumanMessage, AIMessage, get_all_chat_ids, ToolMessage
 import uuid
 import os
 import base64
@@ -10,103 +10,114 @@ st.title("College Chatbot")
 TEMP_DIR = "/tmp"
 os.makedirs(TEMP_DIR, exist_ok=True)
 
+
 # --- Helpers ---
 
+def set_title(messages):
+    if messages:
+        title = "New Chat"
+        st.session_state.chat_dict[st.session_state.current_chat_id] = title
+
+
 def set_config():
-    """Returns the config for the current chat thread."""
     return {"configurable": {"thread_id": st.session_state.current_chat_id}}
 
+
 def load_session_state():
-    """Initializes the session state for chats."""
     if "chats" not in st.session_state:
         st.session_state.chats = get_all_chat_ids()
     if "current_chat_id" not in st.session_state:
-        if st.session_state.chats:
+        if len(st.session_state.chats) > 0:
             st.session_state.current_chat_id = st.session_state.chats[-1]
         else:
-            # Create a new chat if none exist
             new_id = str(uuid.uuid4())
             st.session_state.chats.append(new_id)
             st.session_state.current_chat_id = new_id
-            # Initialize the new chat with the system message
-            workflow.update_state(
-                {"configurable": {"thread_id": new_id}},
-                {"messages": [system]}
-            )
     if "chat_dict" not in st.session_state:
-        # Simple naming for chats for now
-        st.session_state.chat_dict = {chat_id: f"Chat {i+1}" for i, chat_id in enumerate(st.session_state.chats)}
+        st.session_state.chat_dict = {}
+
 
 def render_sidebar():
-    """Displays the chat history and a new chat button in the sidebar."""
     with st.sidebar:
         st.title("Chats")
         if st.button("➕ New Chat"):
             new_id = str(uuid.uuid4())
             st.session_state.chats.append(new_id)
             st.session_state.current_chat_id = new_id
-            st.session_state.chat_dict[new_id] = f"Chat {len(st.session_state.chats)}"
-            # Initialize the new chat with the system message
-            workflow.update_state(
-                {"configurable": {"thread_id": new_id}},
-                {"messages": [system]}
-            )
-            st.rerun() # Rerun to switch to the new chat
+            config = {"configurable": {"thread_id": new_id}}
+            workflow.update_state(config, {"messages": [system]})
+            st.session_state.chat_dict[new_id] = "New Chat"
 
-        st.divider()
         for chat_id in st.session_state.chats:
-            if st.button(st.session_state.chat_dict.get(chat_id, "New Chat"), key=chat_id, use_container_width=True):
+            if st.button(st.session_state.chat_dict.get(chat_id, "New Chat"), key=chat_id):
                 st.session_state.current_chat_id = chat_id
-                st.rerun()
 
-def create_download_link(file_path: str) -> str:
-    """Generates an HTML download link for a file."""
-    if not os.path.exists(file_path): return ""
-    with open(file_path, "rb") as f:
-        data = f.read()
-    b64 = base64.b64encode(data).decode()
-    return f'<a href="data:file/octet-stream;base64,{b64}" download="{os.path.basename(file_path)}">📥 Download {os.path.basename(file_path)}</a>'
+
+def create_download_link(file_path: str, label: str = None) -> str:
+    """Generate HTML download link for a file."""
+    if not os.path.exists(file_path):
+        return ""
+    try:
+        with open(file_path, "rb") as f:
+            data = f.read()
+        b64 = base64.b64encode(data).decode()
+        label = label or f"📥 Download {os.path.basename(file_path)}"
+        href = f'<a href="data:file/octet-stream;base64,{b64}" download="{os.path.basename(file_path)}">{label}</a>'
+        return href
+    except Exception as e:
+        return f"Error creating download link: {e}"
+
 
 def show_file(file_path: str):
-    """Displays a file inline with a download link."""
-    if not os.path.exists(file_path): return
+    """Show file inline + download link."""
+    if not os.path.exists(file_path):
+        return
     ext = os.path.splitext(file_path)[1].lower()
     if ext in [".png", ".jpg", ".jpeg"]:
-        st.image(file_path)
-    else:
-        try:
-            with open(file_path, "r", encoding="utf-8") as f:
-                st.code(f.read(), language=ext.lstrip("."))
-        except:
-            st.warning(f"Could not display file: {os.path.basename(file_path)}")
+        st.image(file_path, caption=os.path.basename(file_path))
+    elif ext in [".txt", ".py", ".java", ".cpp", ".md"]:
+        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+            st.code(f.read(), language=ext.lstrip("."))
     st.markdown(create_download_link(file_path), unsafe_allow_html=True)
 
 
 def render_tool_message(tool_message: ToolMessage):
-    """Renders the output of a tool call."""
+    """Render tool execution based on tool name instead of message content."""
+    file_related_keywords = ["read", "write", "file", "save", "export", "create"]
     with st.chat_message("assistant"):
-        st.info(f"🧰 Tool Used: `{tool_message.name}`")
-        # For file-creating tools, we can find and show the file
-        if tool_message.name in ["write_file", "plot_graph"] and "saved at" in str(tool_message.content):
-            file_path = str(tool_message.content).split("saved at")[-1].strip()
-            show_file(file_path)
-        else: # For other tools, just show the content
-            st.markdown(f"```\n{tool_message.content}\n```")
+        tool_name = getattr(tool_message, "name", "").lower()
+        st.info(f"🧰 Tool used: {tool_name or 'Unknown Tool'}")
+
+        # Check if this is a file-related tool
+        if any(k in tool_name for k in file_related_keywords):
+            # Find all files in TEMP_DIR (freshly modified ones)
+            created_files = sorted(
+                [os.path.join(TEMP_DIR, f) for f in os.listdir(TEMP_DIR)],
+                key=lambda x: os.path.getmtime(x),
+                reverse=True,
+            )
+            if created_files:
+                st.success("📄 File(s) created by tool:")
+                for file_path in created_files[:3]:  # show up to 3 recent
+                    show_file(file_path)
+            else:
+                st.warning("No new file detected in /tmp.")
+        else:
+            # Non-file tools
+            if isinstance(tool_message.content, str):
+                st.write(tool_message.content)
+            elif isinstance(tool_message.content, dict):
+                st.json(tool_message.content)
 
 
-def load_and_display_chats():
-    """Loads and displays the messages for the current chat."""
+def loadchats():
     if "current_chat_id" not in st.session_state:
-        return
-    
-    state = workflow.get_state(set_config())
+        return []
+    config = {"configurable": {"thread_id": st.session_state.current_chat_id}}
+    state = workflow.get_state(config)
     messages = state.values.get("messages", [])
-    
     for message in messages:
-        # Skip the system message in the UI
-        if isinstance(message, SystemMessage):
-            continue
-        elif isinstance(message, HumanMessage):
+        if isinstance(message, HumanMessage):
             with st.chat_message("human"):
                 st.write(message.content)
         elif isinstance(message, AIMessage):
@@ -114,41 +125,32 @@ def load_and_display_chats():
                 st.write(message.content)
         elif isinstance(message, ToolMessage):
             render_tool_message(message)
+    return messages
+
 
 # --- Main Chat Flow ---
+
 load_session_state()
 render_sidebar()
 
-load_and_display_chats()
-
-if user_input := st.chat_input("Your message:"):
-    with st.chat_message("human"):
-        st.write(user_input)
-
-    with st.chat_message("assistant"):
-        response_placeholder = st.empty()
-        full_response = ""
-
-        # The input to stream is just the new message.
-        # LangGraph loads the history from the checkpointer.
-        input_for_graph = {"messages": [HumanMessage(content=user_input)]}
-
-        # Use stream_mode='values' to get outputs from each node as they run
-        for chunk in workflow.stream(input_for_graph, config=set_config() , stream_mode = "messages"):
-            # Check for output from the LLM node
-            if "llmresponse" in chunk:
-                ai_message_chunk = chunk["llmresponse"]["messages"][-1]
-                full_response += ai_message_chunk.content
-                response_placeholder.markdown(full_response + "▌") # Add a cursor effect
-
-            # Check for output from the tool node
-            elif "tool_node" in chunk:
-                # A tool was called. We can render its output.
-                # First, clear the placeholder for the new text that will come
-                response_placeholder.empty()
-                tool_messages = chunk["tool_node"]["messages"]
-                for tool_message in tool_messages:
-                    render_tool_message(tool_message)
-
-        # Final update to remove the cursor
-        response_placeholder.markdown(full_response)
+if "current_chat_id" in st.session_state:
+    loadchats()
+    if user_input := st.chat_input("Your message:"):
+        with st.chat_message("human"):
+            st.write(user_input)
+
+        with st.chat_message("assistant"):
+            response_placeholder = st.empty()
+            full_response = ""
+
+            for message, metadata in workflow.stream(
+                {"messages": [system, HumanMessage(user_input)]},
+                config={"configurable": {"thread_id": st.session_state.current_chat_id}},
+                stream_mode="messages",
+            ):
+                if isinstance(message, AIMessage):
+                    full_response += message.content or ""
+                elif isinstance(message, ToolMessage):
+                    render_tool_message(message)
+
+            response_placeholder.markdown(full_response)