pygeai 0.1.6__py3-none-any.whl → 0.6.0b15__py3-none-any.whl

pygeai/chat/ui.py

@@ -0,0 +1,678 @@
+import logging
+
+import streamlit as st
+import argparse
+import json
+import os
+from datetime import datetime
+from pygeai.chat.session import AgentChatSession
+from pygeai.core.utils.console import Console
+from pygeai.core.common.config import get_settings
+from pygeai.lab.managers import AILabManager
+from pygeai.lab.models import FilterSettings, AgentList
+import sys
+
+logger = logging.getLogger(__name__)
+
+
+global SESSION_FILE_PATH
+
+
+def parse_args():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description="Streamlit chat interface for pygeai agent")
+    parser.add_argument("--agent-name", "-n", required=True, help="Name of the agent to interact with")
+    args, unknown = parser.parse_known_args()  # Ignore Streamlit's args
+    return args
+
+
+def save_session_to_file(messages, file_path):
+    """Helper function to save session to a server-side file."""
+    try:
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        with open(file_path, 'w') as f:
+            json.dump(messages, f, indent=2)
+        logger.info(f"Session automatically saved to {file_path}")
+        return True
+    except Exception as e:
+        logger.error(f"Error saving session to {file_path}: {e}")
+        st.error(f"Failed to auto-save session: {e}")
+        return False
+
+
+def get_unique_file_path(base_path):
+    """Generate a unique file path by appending a numeric suffix if the file exists."""
+    if not os.path.exists(base_path):
+        return base_path
+
+    directory, filename = os.path.split(base_path)
+    name, ext = os.path.splitext(filename)
+    counter = 1
+
+    new_path = base_path
+    while os.path.exists(new_path):
+        new_filename = f"{name}_{counter}{ext}"
+        new_path = os.path.join(directory, new_filename)
+        counter += 1
+
+    return new_path
+
+
+def get_session_file_path(agent_name, custom_filename=None):
+    """Generate the session file path with date and agent name, or use a custom filename."""
+    if custom_filename:
+        if not custom_filename.endswith('.json'):
+            custom_filename += '.json'
+        return os.path.join("chats", custom_filename)
+    current_date = datetime.now().strftime("%Y-%m-%d")
+    return os.path.join("chats", f"chat_session_{agent_name}_{current_date}.json")
+
+
+def list_session_files():
+    """List all JSON files in the 'chats' directory."""
+    chats_dir = "chats"
+    if not os.path.exists(chats_dir):
+        return []
+    try:
+        return [f for f in os.listdir(chats_dir) if f.endswith('.json')]
+    except Exception as e:
+        logger.error(f"Error listing session files in {chats_dir}: {e}")
+        return []
+
+
+def load_session_from_file(file_path):
+    """Load session data from a specified file."""
+    try:
+        with open(file_path, 'r') as f:
+            data = json.load(f)
+            if isinstance(data, list):
+                return data, True
+            else:
+                return None, False
+    except Exception as e:
+        logger.error(f"Error loading session from {file_path}: {e}")
+        return None, False
+
+
+def delete_session_file(file_path):
+    """Helper function to delete a session file from the file system."""
+    try:
+        if os.path.exists(file_path):
+            os.remove(file_path)
+            logger.info(f"Session file deleted: {file_path}")
+            return True, f"Session file {os.path.basename(file_path)} deleted successfully."
+        else:
+            logger.warning(f"Session file not found: {file_path}")
+            return False, f"Session file {os.path.basename(file_path)} not found."
+    except Exception as e:
+        logger.error(f"Error deleting session file {file_path}: {e}")
+        return False, f"Error deleting session file {os.path.basename(file_path)}: {str(e)}"
+
+
+def get_alias_list():
+    """Get list of available aliases from settings."""
+    try:
+        settings = get_settings()
+        aliases = list(settings.list_aliases().keys())
+        return ["-"] + aliases  # Add "-" as default no-selection option
+    except Exception as e:
+        logger.error(f"Error fetching alias list: {e}")
+        st.error(f"Failed to fetch alias list: {e}")
+        return ["-"]
+
+
+def get_agent_list(alias, project_id):
+    """Get list of agents for a given alias and project ID."""
+    try:
+        if alias == "-" or not project_id:
+            return ["-"]
+        ai_lab_manager = AILabManager(alias=alias)
+        filter_settings = FilterSettings(
+            allow_external=False,
+            allow_drafts=True,
+            access_scope="private"
+        )
+        result = ai_lab_manager.get_agent_list(
+            project_id=project_id,
+            filter_settings=filter_settings
+        )
+        if isinstance(result, AgentList) and result.agents:
+            # Store full agent data for preview
+            st.session_state.agent_data = {f"{agent.name} (ID: {agent.id})": agent for agent in result.agents}
+            return ["-"] + [f"{agent.name} (ID: {agent.id})" for agent in result.agents]
+        else:
+            st.error(f"No agents found for project ID {project_id} or errors occurred: {result.errors if hasattr(result, 'errors') else 'Unknown error'}")
+            return ["-"]
+    except Exception as e:
+        logger.error(f"Error fetching agents for project ID {project_id} with alias {alias}: {e}")
+        st.error(f"Failed to fetch agents for project: {e}")
+        return ["-"]
+
+
+def save_recent_agents(agent_name):
+    """Save a recently used agent to session state or file."""
+    if "recent_agents" not in st.session_state:
+        st.session_state.recent_agents = []
+    if agent_name and agent_name not in st.session_state.recent_agents:
+        st.session_state.recent_agents = [agent_name] + st.session_state.recent_agents[:4]  # Keep top 5 recent
+        # Optionally save to a file for persistence
+        try:
+            with open("recent_agents.json", "w") as f:
+                json.dump(st.session_state.recent_agents, f)
+        except Exception as e:
+            logger.error(f"Error saving recent agents: {e}")
+
+
+def load_recent_agents():
+    """Load recent agents from session state or file."""
+    if "recent_agents" not in st.session_state:
+        try:
+            with open("recent_agents.json", "r") as f:
+                st.session_state.recent_agents = json.load(f)
+        except Exception:
+            st.session_state.recent_agents = []
+    return st.session_state.recent_agents
+
+
+def run_streamlit_chat():
+    """Run a Streamlit chat interface for the specified agent."""
+    args = parse_args()
+    initial_agent_name = args.agent_name
+    try:
+        # Initialize session state for multiple tabs and current agent
+        if "sessions" not in st.session_state:
+            st.session_state.sessions = {
+                initial_agent_name: {
+                    "messages": [],
+                    "chat_session": AgentChatSession(initial_agent_name),
+                    "custom_filename": "",
+                    "active": True
+                }
+            }
+        if "active_tab" not in st.session_state:
+            st.session_state.active_tab = initial_agent_name
+        if "agent_data" not in st.session_state:
+            st.session_state.agent_data = {}
+        if "search_term" not in st.session_state:
+            st.session_state.search_term = ""
+        # if "editing_message" not in st.session_state:
+        #    st.session_state.editing_message = None
+        # if "edit_content" not in st.session_state:
+        #    st.session_state.edit_content = ""
+        # if "regenerate_response" not in st.session_state:
+        #    st.session_state.regenerate_response = False
+
+        # Load recent agents
+        load_recent_agents()
+
+        # Load global vars
+        global current_session
+        global SESSION_FILE_PATH
+
+        # Get current session for active tab
+        current_session = st.session_state.sessions[st.session_state.active_tab]
+        if "editing_message" not in current_session:
+            current_session["editing_message"] = None
+        if "edit_content" not in current_session:
+            current_session["edit_content"] = ""
+        if "regenerate_response" not in current_session:
+            current_session["regenerate_response"] = False
+        SESSION_FILE_PATH = get_session_file_path(st.session_state.active_tab, current_session.get("custom_filename", ""))
+
+        # Initialize messages for new session if not already set
+        if not current_session["messages"]:
+            if os.path.exists(SESSION_FILE_PATH):
+                try:
+                    with open(SESSION_FILE_PATH, 'r') as f:
+                        restored_data = json.load(f)
+                        if isinstance(restored_data, list):
+                            current_session["messages"] = restored_data
+                            st.success(f"Session automatically restored from {SESSION_FILE_PATH}")
+                        else:
+                            logger.warning(f"Invalid session data in {SESSION_FILE_PATH}. Starting fresh.")
+                            intro = current_session["chat_session"].get_answer(
+                                ["You're about to speak to a user. Introduce yourself in a clear and concise manner, "
+                                 "stating who you are and what you do. Nothing else."]
+                            )
+                            if "Agent not found" in str(intro):
+                                st.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                                logger.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                                return
+                            current_session["messages"] = [{"role": "assistant", "content": intro}]
+                except Exception as e:
+                    logger.error(f"Error auto-restoring session from {SESSION_FILE_PATH}: {e}")
+                    intro = current_session["chat_session"].get_answer(
+                        ["You're about to speak to a user. Introduce yourself in a clear and concise manner, "
+                         "stating who you are and what you do. Nothing else."]
+                    )
+                    if "Agent not found" in str(intro):
+                        st.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                        logger.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                        return
+                    current_session["messages"] = [{"role": "assistant", "content": intro}]
+            else:
+                intro = current_session["chat_session"].get_answer(
+                    ["You're about to speak to a user. Introduce yourself in a clear and concise manner, "
+                     "stating who you are and what you do. Nothing else."]
+                )
+                if "Agent not found" in str(intro):
+                    st.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                    logger.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                    return
+                current_session["messages"] = [{"role": "assistant", "content": intro}]
+                save_session_to_file(current_session["messages"], SESSION_FILE_PATH)
+                st.rerun()
+
+        # Set page title
+        st.title(f"Chat with {st.session_state.active_tab}")
+
+        # Display the session save path in the main area for better readability
+        st.info(f"Session will be saved as: {SESSION_FILE_PATH}", icon="ℹ️")
+
+        # Chat History Search
+        search_term = st.text_input("Search Chat History", value=st.session_state.search_term, placeholder="Enter keyword to filter messages")
+        if search_term != st.session_state.search_term:
+            st.session_state.search_term = search_term
+            st.rerun()
+
+        # Single Session Chat
+        current_session = st.session_state.sessions[st.session_state.active_tab]
+        if "editing_message" not in current_session:
+            current_session["editing_message"] = None
+        if "edit_content" not in current_session:
+            current_session["edit_content"] = ""
+        render_chat_history(
+            current_session["messages"],
+            search_term,
+            tab_key_prefix=st.session_state.active_tab,
+            session=current_session
+        )
+
+        # Sidebar for session management and agent selection
+        with st.sidebar:
+            st.header("Session Management")
+
+            with st.expander("Help: Session Saving & Restoring", expanded=False):
+                st.markdown("""
+                **Session Saving & Restoring Explained:**
+                - **Auto-Save Session**: When toggled on, your chat history is automatically saved to a server-side file in the 'chats' directory after each message or action.
+                - **Custom Session Filename**: Enter a custom name for the session file to save it with a specific identifier.
+                - **Save Session (JSON)**: Download a local copy of your chat history.
+                - **Restore Session (JSON)**: Upload a previously saved JSON file to load a specific chat history.
+                - **Available Sessions**: List of saved session files with previews. Click to load into the current tab.
+                - **Reset Chat**: Clears the current session and starts fresh with the agent's introduction.
+                """)
+
+            # Custom filename input for the current session
+            custom_filename = st.text_input("Custom Session Filename (optional)", value=current_session.get("custom_filename", ""), placeholder="e.g., my_custom_session")
+            if custom_filename != current_session.get("custom_filename", ""):
+                current_session["custom_filename"] = custom_filename
+                SESSION_FILE_PATH = get_session_file_path(st.session_state.active_tab, custom_filename)
+                if current_session["messages"]:
+                    save_session_to_file(current_session["messages"], SESSION_FILE_PATH)
+                    st.rerun()
+
+            uploaded_file = st.file_uploader("Restore Session (JSON)", type=["json"])
+            if uploaded_file is not None and "session_restored" not in st.session_state:
+                try:
+                    restored_data = json.load(uploaded_file)
+                    if isinstance(restored_data, list):
+                        current_session["messages"] = restored_data
+                        st.session_state.session_restored = True
+                        save_session_to_file(current_session["messages"], SESSION_FILE_PATH)
+                        st.success(f"Session restored from {uploaded_file.name}")
+                        st.rerun()
+                    else:
+                        st.error("Invalid session file: Must contain a list of messages in JSON format.")
+                except json.JSONDecodeError:
+                    st.error("Invalid JSON format in uploaded file.")
+                except Exception as e:
+                    st.error(f"Error restoring session: {e}")
+                    logger.error(f"Error restoring session: {e}")
+
+            if current_session["messages"]:
+                session_json = json.dumps(current_session["messages"], indent=2)
+                current_date = datetime.now().strftime("%Y-%m-%d")
+                st.download_button(
+                    label="Save Session (JSON)",
+                    data=session_json,
+                    file_name=f"chat_session_{st.session_state.active_tab}_{current_date}.json",
+                    mime="application/json"
+                )
+
+            auto_save = st.toggle("Auto-Save Session", value=st.session_state.get("auto_save", True))
+            st.session_state.auto_save = auto_save
+
+            # Available Sessions with Preview
+            st.subheader("Available Sessions")
+            session_files = list_session_files()
+            session_feedback_placeholder = st.empty()
+            if session_files:
+                st.markdown("Click a file to load the session, or use the trash icon to delete:")
+                for file in session_files:
+                    col1, col2 = st.columns([3, 1])
+                    with col1:
+                        if st.button(file, key=f"load_{file}", use_container_width=True):
+                            file_path = os.path.join("chats", file)
+                            loaded_data, success = load_session_from_file(file_path)
+                            if success:
+                                current_session["messages"] = loaded_data
+                                st.success(f"Session loaded from {file}")
+                                if st.session_state.auto_save:
+                                    save_session_to_file(current_session["messages"], SESSION_FILE_PATH)
+                                st.rerun()
+                            else:
+                                st.error(f"Failed to load session from {file}")
+                    with col2:
+                        if st.button("🗑️", key=f"delete_{file}", help=f"Delete {file}"):
+                            file_path = os.path.join("chats", file)
+                            if os.path.abspath(file_path) == os.path.abspath(SESSION_FILE_PATH):
+                                with session_feedback_placeholder.container():
+                                    st.markdown("---")
+                                    st.error(
+                                        "Cannot delete the currently active session file. Please switch to another session or reset the chat first.",
+                                        icon="🚫")
+                                    st.markdown("---")
+                            else:
+                                try:
+                                    os.remove(file_path)
+                                    logger.info(f"Deleted session file: {file_path}")
+                                    with session_feedback_placeholder.container():
+                                        st.markdown("---")
+                                        st.success(f"Deleted {file}", icon="✅")
+                                        st.markdown("---")
+                                    st.rerun()
+                                except Exception as e:
+                                    logger.error(f"Error deleting session file {file_path}: {e}")
+                                    with session_feedback_placeholder.container():
+                                        st.markdown("---")
+                                        st.error(f"Failed to delete {file}: {e}", icon="❌")
+                                        st.markdown("---")
+            else:
+                st.markdown("No saved sessions found in 'chats' directory.")
+
+            # New Feature: Agent Selection
+            st.header("Switch Agent")
+            with st.expander("Help: Switching Agents", expanded=False):
+                st.markdown("""
+                **Switching Agents Explained:**
+                - **Select Alias**: Choose a profile (alias) to access specific API configurations.
+                - **Enter Project ID**: Provide the ID of the project to list available agents.
+                - **Select Agent**: Pick an agent to chat with. You'll be prompted to confirm before switching.
+                - Switching to a new agent starts a fresh session in a new tab.
+                """)
+
+            # Initialize session state for alias, project ID, and agent selection
+            if "selected_alias" not in st.session_state:
+                st.session_state.selected_alias = "-"
+            if "project_id_input" not in st.session_state:
+                st.session_state.project_id_input = ""
+            if "selected_agent" not in st.session_state:
+                st.session_state.selected_agent = "-"
+            if "confirm_switch" not in st.session_state:
+                st.session_state.confirm_switch = False
+
+            # Alias Selection
+            alias_list = get_alias_list()
+            selected_alias = st.selectbox("Select Alias (Profile)", alias_list, index=alias_list.index(st.session_state.selected_alias))
+            if selected_alias != st.session_state.selected_alias:
+                st.session_state.selected_alias = selected_alias
+                st.session_state.selected_agent = "-"
+                st.rerun()
+
+            # Project ID Input
+            project_id_input = st.text_input("Enter Project ID", value=st.session_state.project_id_input, placeholder="e.g., 2ca6883f-6778-40bb-bcc1-85451fb11107")
+            if project_id_input != st.session_state.project_id_input:
+                st.session_state.project_id_input = project_id_input
+                st.session_state.selected_agent = "-"
+                st.rerun()
+
+            # Agent Selection
+            agent_list = get_agent_list(st.session_state.selected_alias, st.session_state.project_id_input) if st.session_state.selected_alias != "-" and st.session_state.project_id_input else ["-"]
+            selected_agent = st.selectbox("Select Agent", agent_list, index=agent_list.index(st.session_state.selected_agent) if st.session_state.selected_agent in agent_list else 0)
+            if selected_agent != st.session_state.selected_agent:
+                st.session_state.selected_agent = selected_agent
+                st.session_state.confirm_switch = False
+
+            # Agent Info Preview
+            if st.session_state.selected_agent != "-" and st.session_state.selected_agent in st.session_state.agent_data:
+                agent_obj = st.session_state.agent_data.get(st.session_state.selected_agent)
+                if agent_obj and getattr(agent_obj, 'description', None):
+                    st.markdown("**Agent Description:**")
+                    st.markdown(f"{agent_obj.description[:200]}{'...' if len(agent_obj.description) > 200 else ''}")
+
+            # Favorite/Recent Agents List
+            st.subheader("Recent Agents")
+            recent_agents = load_recent_agents()
+            if recent_agents:
+                for agent in recent_agents:
+                    if st.button(f"Switch to {agent}", key=f"recent_{agent}"):
+                        if agent != st.session_state.active_tab:
+                            if agent not in st.session_state.sessions:
+                                st.session_state.sessions[agent] = {
+                                    "messages": [],
+                                    "chat_session": AgentChatSession(agent),
+                                    "custom_filename": "",
+                                    "active": True
+                                }
+                            st.session_state.active_tab = agent
+                            save_recent_agents(agent)
+                            st.rerun()
+            else:
+                st.markdown("No recent agents found.")
+
+            # Confirmation for switching agent
+            if st.session_state.selected_agent != "-" and st.session_state.selected_agent.split(" (ID: ")[0] != st.session_state.active_tab:
+                if not st.session_state.confirm_switch:
+                    st.warning(f"Do you wish to chat with {st.session_state.selected_agent.split(' (ID: ')[0]}? This will start a new session in a new tab.", icon="⚠️")
+                    col1, col2 = st.columns(2)
+                    with col1:
+                        if st.button("Confirm Switch"):
+                            st.session_state.confirm_switch = True
+                            new_agent_name = st.session_state.selected_agent.split(" (ID: ")[0]
+                            if new_agent_name not in st.session_state.sessions:
+                                st.session_state.sessions[new_agent_name] = {
+                                    "messages": [],
+                                    "chat_session": AgentChatSession(new_agent_name),
+                                    "custom_filename": "",
+                                    "active": True
+                                }
+                            st.session_state.active_tab = new_agent_name
+                            save_recent_agents(new_agent_name)
+                            current_session = st.session_state.sessions[new_agent_name]
+                            if not current_session["messages"]:
+                                intro = current_session["chat_session"].get_answer(
+                                    ["You're about to speak to a user. Introduce yourself in a clear and concise manner, "
+                                     "stating who you are and what you do. Nothing else."]
+                                )
+                                if "Agent not found" in str(intro):
+                                    st.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                                    logger.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                                    return
+                                current_session["messages"] = [{"role": "assistant", "content": intro}]
+                            SESSION_FILE_PATH = get_session_file_path(new_agent_name, current_session.get("custom_filename", ""))
+                            if st.session_state.auto_save:
+                                save_session_to_file(current_session["messages"], SESSION_FILE_PATH)
+                            st.rerun()
+                    with col2:
+                        if st.button("Cancel"):
+                            st.session_state.selected_agent = "-"
+                            st.session_state.confirm_switch = False
+                            st.rerun()
+                else:
+                    st.success(f"Switched to agent {st.session_state.selected_agent.split(' (ID: ')[0]}!", icon="✅")
+
+        # Reset chat button for current tab
+        if st.button("Reset Chat"):
+            current_session["messages"] = []
+            intro = current_session["chat_session"].get_answer(
+                ["You're about to speak to a user. Introduce yourself in a clear and concise manner, "
+                 "stating who you are and what you do. Nothing else."]
+            )
+            if "Agent not found" in str(intro):
+                st.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                logger.error("The specified agent doesn't seem to exist. Please review the name and try again.")
+                return
+            current_session["messages"] = [{"role": "assistant", "content": intro}]
+            if st.session_state.auto_save:
+                save_session_to_file(current_session["messages"], SESSION_FILE_PATH)
+            st.rerun()
+
+        error_container = st.empty()
+
+        # "Complete Answer" button logic
+        if (current_session["messages"] and
+            current_session["messages"][-1]["role"] == "assistant" and
+            "complete_answer_triggered" not in st.session_state):
+            if st.button("Complete Answer"):
+                st.session_state.complete_answer_triggered = True
+                last_assistant_message = current_session["messages"][-1]["content"]
+                continuation_prompt = (
+                    f"The previous answer was: '{last_assistant_message}'. "
+                    "It seems incomplete. Please continue and complete the answer."
+                )
+                current_session["messages"].append({"role": "user", "content": continuation_prompt})
+                with st.chat_message("user"):
+                    st.markdown(continuation_prompt)
+                with st.chat_message("assistant"):
+                    with st.spinner("Continuing answer..."):
+                        response_placeholder = st.empty()
+                        continued_answer = ""
+                        result = current_session["chat_session"].stream_answer(current_session["messages"])
+                        for chunk in result:
+                            continued_answer += chunk
+                            sanitized_answer = continued_answer
+                            response_placeholder.markdown(f'{sanitized_answer}')
+                        current_session["messages"].append({"role": "assistant", "content": continued_answer})
+                if st.session_state.auto_save:
+                    save_session_to_file(current_session["messages"], SESSION_FILE_PATH)
+                del st.session_state.complete_answer_triggered
+                st.rerun()
+
+        # Chat input for current tab
+        if user_input := st.chat_input(f"Ask {st.session_state.active_tab}"):
+            if not user_input.strip():
+                logger.warning(f"Empty input submitted for agent {st.session_state.active_tab}")
+                with error_container.container():
+                    st.error(f"Unable to communicate with the agent {st.session_state.active_tab}")
+                return
+            error_container.empty()
+
+            # Append user message to session
+            with st.chat_message("user"):
+                st.markdown(f"{user_input}")
+            current_session["messages"].append({"role": "user", "content": user_input})
+
+            # Save session if auto_save is enabled
+            if st.session_state.auto_save:
+                save_session_to_file(current_session["messages"], SESSION_FILE_PATH)
+
+            # Stream assistant response
+            with st.chat_message("assistant"):
+                with st.spinner("Processing..."):
+                    response_placeholder = st.empty()
+                    answer = ""
+                    result = current_session["chat_session"].stream_answer(current_session["messages"])
+                    for chunk in result:
+                        answer += chunk
+                        sanitized_answer = answer
+                        response_placeholder.markdown(f'{sanitized_answer}')
+
+                    # Append the complete response to session
+                    current_session["messages"].append({"role": "assistant", "content": answer})
+
+            # Save session again after assistant response
+            if st.session_state.auto_save:
+                save_session_to_file(current_session["messages"], SESSION_FILE_PATH)
+                st.rerun()
+    except Exception as e:
+        st.error(f"An unexpected error occurred: {e}")
+        logger.error(f"An unexpected error occurred: {e}")
+        Console.write_stderr("An unexpected error has occurred. Please contact the developers.")
+        sys.exit(1)
+
+
+def render_chat_history(messages, search_term="", tab_key_prefix="", session=None):
+    """Render chat history with search filter, copy, and edit buttons."""
+
+    global auto_save
+
+    filtered_messages = []
+    if search_term:
+        search_term = search_term.lower()
+        filtered_messages = [msg for msg in messages if search_term in msg["content"].lower()]
+    else:
+        filtered_messages = messages
+
+    if search_term and not filtered_messages:
+        st.info("No messages match your search term.", icon="ℹ️")
+    elif search_term:
+        st.info(f"Showing {len(filtered_messages)} messages matching '{search_term}'", icon="🔍")
+
+    # Find the index of the last user message in the filtered list for edit button display
+    last_user_msg_index = None
+    for i, msg in enumerate(filtered_messages):
+        if msg["role"] == "user":
+            last_user_msg_index = i
+
+    for i, message in enumerate(filtered_messages):
+        with st.chat_message(message["role"]):
+            col1, col2 = st.columns([5, 1])
+            with col1:
+                st.markdown(message["content"])
+            with col2:
+                # Edit Button - Only for the last user message
+                if message["role"] == "user" and i == last_user_msg_index:
+                    if st.button("✏️", key=f"{tab_key_prefix}_edit_{i}", help="Edit this message"):
+                        session["editing_message"] = i
+                        session["edit_content"] = message["content"]
+                        st.rerun()
+
+    # Handle Message Editing
+    if session and session.get("editing_message") is not None:
+        with st.container():
+            st.markdown("**Editing Message:**")
+            new_content = st.text_area(
+                "Edit your message:",
+                value=session.get("edit_content", ""),
+                key=f"{tab_key_prefix}_edit_area"
+            )
+            col1, col2 = st.columns(2)
+            with col1:
+                if st.button("Save Edit", key=f"{tab_key_prefix}_save_edit"):
+                    # Update the message content
+                    messages[session["editing_message"]]["content"] = new_content
+                    # Check if there's a subsequent agent response to remove
+                    if session["editing_message"] + 1 < len(messages) and messages[session["editing_message"] + 1]["role"] == "assistant":
+                        messages.pop(session["editing_message"] + 1)  # Remove the old response
+                        session["regenerate_response"] = True  # Flag to regenerate response
+                    session["editing_message"] = None
+                    session["edit_content"] = ""
+                    st.rerun()
+            with col2:
+                if st.button("Cancel Edit", key=f"{tab_key_prefix}_cancel_edit"):
+                    session["editing_message"] = None
+                    session["edit_content"] = ""
+                    session["regenerate_response"] = False
+                    st.rerun()
+
+    # Regenerate agent response if needed after edit
+    if session and session.get("regenerate_response"):
+        with st.chat_message("assistant"):
+            with st.spinner("Regenerating response..."):
+                response_placeholder = st.empty()
+                answer = ""
+                result = current_session["chat_session"].stream_answer(messages)
+                for chunk in result:
+                    answer += chunk
+                    sanitized_answer = answer
+                    response_placeholder.markdown(f'{sanitized_answer}')
+                messages.append({"role": "assistant", "content": answer})
+        if st.session_state.auto_save:
+            save_session_to_file(messages, SESSION_FILE_PATH)
+        session["regenerate_response"] = False
+        st.rerun()
+
+
+if __name__ == "__main__":
+    run_streamlit_chat()