Update src/streamlit_app.py

src/streamlit_app.py

@@ -10,113 +10,103 @@ 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 len(st.session_state.chats) > 0:
+        if st.session_state.chats:
             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:
-        st.session_state.chat_dict = {}
-
+        # 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)}
 
 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
-            config = {"configurable": {"thread_id": new_id}}
-            workflow.update_state(config, {"messages": [system]})
-            st.session_state.chat_dict[new_id] = "New Chat"
+            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
 
+        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):
+            if st.button(st.session_state.chat_dict.get(chat_id, "New Chat"), key=chat_id, use_container_width=True):
                 st.session_state.current_chat_id = chat_id
+                st.rerun()
 
-
-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 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 show_file(file_path: str):
-    """Show file inline + download link."""
-    if not os.path.exists(file_path):
-        return
+    """Displays a file inline with a 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, 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.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.markdown(create_download_link(file_path), unsafe_allow_html=True)
 
 
 def render_tool_message(tool_message: ToolMessage):
-    """Render tool execution based on tool name instead of message content."""
-    file_related_keywords = ["read", "write", "file", "save", "export", "create"]
+    """Renders the output of a tool call."""
     with st.chat_message("assistant"):
-        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)
+        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```")
 
 
-def loadchats():
+def load_and_display_chats():
+    """Loads and displays the messages for the current chat."""
     if "current_chat_id" not in st.session_state:
-        return []
-    config = {"configurable": {"thread_id": st.session_state.current_chat_id}}
-    state = workflow.get_state(config)
+        return
+    
+    state = workflow.get_state(set_config())
     messages = state.values.get("messages", [])
+    
     for message in messages:
-        if isinstance(message, HumanMessage):
+        # Skip the system message in the UI
+        if isinstance(message, SystemMessage):
+            continue
+        elif isinstance(message, HumanMessage):
             with st.chat_message("human"):
                 st.write(message.content)
         elif isinstance(message, AIMessage):
@@ -124,42 +114,41 @@ def loadchats():
                 st.write(message.content)
         elif isinstance(message, ToolMessage):
             render_tool_message(message)
-    return messages
-
 
 # --- Main Chat Flow ---
 load_session_state()
 render_sidebar()
 
-if "current_chat_id" in st.session_state:
-    loadchats()
-    user_input = st.chat_input("Your message:")
-    if user_input:
-        with st.chat_message("human"):
-            st.write(user_input)
-
-        with st.chat_message("assistant"):
-            # ... (inside with st.chat_message("assistant"))
-            response_placeholder = st.empty()
-            full_response = ""
-
-# Use stream_mode='values' for a cleaner loop
-            for chunk in workflow.stream(
-                {"messages": [HumanMessage(user_input)]}, # System message should already be in state
-                config=set_config()):
-    # The output of a streaming node is the key of the node itself
-                     if "llmresponse" in chunk:
-                         ai_message = chunk["llmresponse"]["messages"][-1]
-                         full_response += ai_message.content
-                         response_placeholder.markdown(full_response + "▌") # Add a cursor effect
-    
-    # Handle tool calls if they appear in the stream
-                     elif "tool_node" in chunk:
-        # Since we are streaming, render the tool message right away
-        # And clear the response placeholder for the next text chunk
-                         response_placeholder.empty() 
-                         tool_messages = chunk["tool_node"]["messages"]
-                         for tool_message in tool_messages:
-                             render_tool_message(tool_message)
-
-            response_placeholder.markdown(full_response) # Final update without cursor
+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()):
+            # 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)