sci/chat_page.py

import streamlit as st
from datetime import datetime, timedelta
from colorprinter.print_color import *


from _base_class import StreamlitBaseClass
from _rss import RSSReader
from projects_page import Project
from streamlit_chatbot import StreamlitChat, StreamlitBot


class BotChatPage(StreamlitBaseClass):
    """
    BotChatPage - A Streamlit interface for chatting with various AI assistants.
    This class provides a user interface for interacting with different types of AI bots
    (Research Assistant, Editor, Podcast) that can access and work with user's collections
    and projects.
    Attributes:
        username (str): The username of the current user.
        collection_name (str): Name of the selected collection.
        project_name (str): Name of the selected project.
        project (Project): Project instance the chat is associated with.
        chat (StreamlitChat): Chat instance for maintaining conversation history.
        role (str): The selected bot persona, default is "Research Assistant".
        page_name (str): Name of the current page ("Bot Chat").
        chat_key (str): Unique identifier for the current chat session.
        bot (StreamlitBot): Instance of the selected bot type.
    Methods:
        run(): Main method to render the chat interface and handle interactions.
        get_chat(role, new_chat): Retrieves existing chat or creates a new one.
        sidebar_actions(): Renders sidebar elements for selecting collections, projects, and chat options.
        remove_old_unsaved_chats(): Cleans up unsaved chats older than two weeks.
    """
    def __init__(self, username):
        super().__init__(username=username)
        self.collection_name = None
        self.project_name = None
        self.project: Project = None
        self.chat = None
        self.role = "Research Assistant"  # Default persona
        self.page_name = "Bot Chat"
        self.chat_key = None
        self.bot: StreamlitBot = None

        # Initialize attributes from session state if available
        if self.page_name in st.session_state:
            for k, v in st.session_state[self.page_name].items():
                setattr(self, k, v)

    def run(self):
        from streamlit_chatbot import EditorBot, ResearchAssistantBot, PodBot, StreamlitBot
        self.bot: StreamlitBot = None
        self.update_current_page("Bot Chat")
        self.remove_old_unsaved_chats()
        self.sidebar_actions()

        if self.collection_name or self.project:
            print_purple("Collection:", self.collection_name, "Project:", self.project_name)
            # If no chat exists, create a new Chat instance
            self.chat = self.get_chat(role=self.role)

            # Create a Bot instance with the Chat object
            if self.role == "Research Assistant":
                print_blue("Creating Research Assistant Bot")
                self.bot: ResearchAssistantBot = ResearchAssistantBot(
                    username=self.username,
                    chat=self.chat,
                    collection=self.collection_name,
                    project=self.project,
                    tools=[
                    "fetch_other_documents_tool",
                    "fetch_science_articles_tool",
                    "fetch_science_articles_and_other_documents_tool",
                    "conversational_response_tool"]
                )

            elif self.role == "Editor":
                self.bot: StreamlitBot = EditorBot(
                    username=self.username,
                    chat=self.chat,
                    collection=self.collection,
                    project=self.project,
                    tools=[
                    "fetch_other_documents_tool",
                    "fetch_notes_tool",
                    "conversational_response_tool"]
                )

            elif self.role == "Podcast":
                st.session_state["make_podcast"] = True
                # with st.sidebar:
                with st.sidebar:
                    with st.form("make_podcast_form"):
                        instructions = st.text_area(
                            "What should the podcast be about? Give a brief description, as if you were the producer."
                        )
                        start = st.form_submit_button("Make Podcast!")
                        if start:
                            bot = PodBot(
                                subject=self.project.name,
                                username=self.username,
                                chat=self.chat,
                                collection=self.collection,
                                project=self.project,
                                instructions=instructions                                
                            )

            # Save updated chat state to session state
            st.session_state[self.page_name] = {
                "collection": self.collection,
                "project": self.project,
                "chat": self.chat,
                "role": self.role,
            }

            # Run the bot (this will display chat history and process user input)
            if self.bot:
                self.bot.run()

        else: # If no collection or project is selected, use the conversational response bot           
            print_yellow("No collection or project selected. Using conversational response bot.")
            self.bot: StreamlitBot = StreamlitBot(
                username=self.username,
                chat=self.get_chat(),
                tools=["conversational_response_tool"],
            )
            self.bot.run()


    def get_chat(self, role="Research Assistant", new_chat=False):
        """
        Retrieves or creates a chat session.
        
        This method handles chat session management by either creating a new chat,
        retrieving an existing one from the database, or initializing a chat when
        none exists in the session state.
        
        Parameters:
        -----------
        role : str, optional
            The role assigned to the chat (default is "Research Assistant").
        new_chat : bool, optional
            If True, creates a new chat regardless of existing sessions (default is False).
            
        Returns:
        --------
        StreamlitChat
            A chat instance either newly created or retrieved from the database.
            
        Notes:
        ------
        - If new_chat is True, a new chat is always created
        - If no chat exists in session state, a new one is created
        - Otherwise, retrieves the existing chat from the database using the chat_key in session state
        """
        print_blue('CHAT TYPE:', role)
        if new_chat:
            chat = StreamlitChat(username=self.username, role=role)
            st.session_state['chat_key'] = chat._key
            print_blue("Creating new chat:", st.session_state['chat_key'])
        elif 'chat_key' not in st.session_state:
            chat = StreamlitChat(username=self.username, role=role)
            st.session_state['chat_key'] = chat._key
            print_blue("Creating new chat:", st.session_state['chat_key'])
        else:
            print_blue("Old chat:", st.session_state['chat_key'])
            chat_data = self.user_arango.db.collection("chats").get(st.session_state['chat_key'])
            chat = StreamlitChat.from_dict(chat_data)
        return chat

    def sidebar_actions(self):
        with st.sidebar:
            with st.form("select_chat"):
                self.collection = self.choose_collection("Article collection to use for chat:")
                self.project = self.choose_project("Project to use for chat:")
                submitted = st.form_submit_button("Select Collection/Project")

            with st.form("chat_settings"):
                if submitted or any([self.collection, self.project]):
                    if self.project:
                        self.role = st.selectbox(
                            "Choose Bot Role",
                            options=["Research Assistant", "Editor", "Podcast"],
                            index=0,
                        )
                    elif self.collection:
                        self.role = "Research Assistant"

                    # Load existing chats from the database
                    if self.project:
                        chat_history = list(
                            self.user_arango.db.aql.execute(
                                f'FOR doc IN chats FILTER doc["project"] == "{self.project}" RETURN {{"_key": doc["_key"], "name": doc["name"]}}'
                            )
                        )
                        # self.project = Project(username=self.username, project_name=self.project_name, user_arango=self.user_arango)
                    elif self.collection:
                        chat_history = list(
                            self.user_arango.db.aql.execute(
                                f'FOR doc IN chats FILTER doc["collection"] == "{self.collection}" RETURN {{"_key": doc["_key"], "name": doc["name"]}}'
                            )
                        )

                    chats = {i["name"]: i["_key"] for i in chat_history}
                    selected_chat = st.selectbox(
                        "Continue another chat", options=[""] + list(chats.keys()), index=None
                    )
                
                    if not self.role:
                        self.role == "Research Assistant"

                    start_chat = st.form_submit_button("Start Chat")
                    if start_chat:
                        if selected_chat:
                            st.session_state["chat_key"] = chats[selected_chat]
                            self.chat = self.get_chat()
                        else:
                            self.chat = self.get_chat(role=self.role, new_chat=True)
                        st.rerun()

    def remove_old_unsaved_chats(self):
        two_weeks_ago = datetime.now() - timedelta(weeks=2)
        q = f'FOR doc IN chats FILTER doc.saved == false AND doc.last_updated < "{two_weeks_ago.isoformat()}" RETURN doc'

        old_chats = self.user_arango.db.aql.execute(
            f'FOR doc IN chats RETURN doc'
        )
        old_chats = self.user_arango.db.aql.execute(
            f'FOR doc IN chats FILTER doc.saved == false AND doc.last_updated < "{two_weeks_ago.isoformat()}" RETURN doc'
        )
        for chat in old_chats:
            print_red(chat["_id"])
            self.user_arango.db.collection("chats").delete(chat["_key"])