Isolate chats with different keys. Show bookkiping

pyproject.toml

@@ -4,7 +4,7 @@
   license = "MIT"
   name = "pyrobbot"
   readme = "README.md"
-  version = "0.1.1"
+  version = "0.1.2"
 
 [build-system]
   build-backend = "poetry.core.masonry.api"

pyrobbot/__init__.py

@@ -1,16 +1,22 @@
 #!/usr/bin/env python3
 """Unnoficial OpenAI API UI and CLI tool."""
+import hashlib
 import os
 import tempfile
 import uuid
+from dataclasses import dataclass
 from importlib.metadata import version
 from pathlib import Path
 
 import openai
 
 
-class GeneralConstants:
-    """General constants for the package."""
+@dataclass
+class GeneralDefinitions:
+    """General definitions for the package."""
+
+    # Keep track of the OpenAI API key available to the package at initialization
+    SYSTEM_ENV_OPENAI_API_KEY: str = None
 
     # Main package info
     RUN_ID = uuid.uuid4().hex
@@ -19,26 +25,43 @@ class GeneralConstants:
 
     # Main package directories
     PACKAGE_DIRECTORY = Path(__file__).parent
-    PACKAGE_CACHE_DIRECTORY = Path.home() / ".cache" / PACKAGE_NAME
     _PACKAGE_TMPDIR = tempfile.TemporaryDirectory()
     PACKAGE_TMPDIR = Path(_PACKAGE_TMPDIR.name)
-    CHAT_CACHE_DIR = PACKAGE_CACHE_DIRECTORY / "chats"
 
     # Constants related to the app
     APP_NAME = "pyRobBot"
     APP_DIR = PACKAGE_DIRECTORY / "app"
     APP_PATH = APP_DIR / "app.py"
     PARSED_ARGS_FILE = PACKAGE_TMPDIR / f"parsed_args_{RUN_ID}.pkl"
 
-    # Constants related to using the OpenAI API
-    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-    TOKEN_USAGE_DATABASE = PACKAGE_CACHE_DIRECTORY / "token_usage.db"
+    @staticmethod
+    def openai_key_hash():
+        """Return a hash of the OpenAI API key."""
+        if openai.api_key is None:
+            return "demo"
+        return hashlib.sha256(openai.api_key.encode("utf-8")).hexdigest()
+
+    @property
+    def package_cache_directory(self):
+        """Return the path to the package's cache directory."""
+        return (
+            Path.home() / ".cache" / self.PACKAGE_NAME / f"user_{self.openai_key_hash()}"
+        )
+
+    @property
+    def chat_cache_dir(self):
+        """Return the path to the package's cache directory."""
+        return self.package_cache_directory / "chats"
+
+    @property
+    def general_token_usage_db_path(self):
+        """Return the path to the package's token usage database."""
+        return self.package_cache_directory / "token_usage.db"
 
-    # Initialise the package's directories
-    PACKAGE_TMPDIR.mkdir(parents=True, exist_ok=True)
-    PACKAGE_CACHE_DIRECTORY.mkdir(parents=True, exist_ok=True)
-    CHAT_CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
+GeneralConstants = GeneralDefinitions(
+    SYSTEM_ENV_OPENAI_API_KEY=os.environ.get("OPENAI_API_KEY")
+)
 
 # Initialize the OpenAI API client
-openai.api_key = GeneralConstants.OPENAI_API_KEY
+openai.api_key = GeneralConstants.SYSTEM_ENV_OPENAI_API_KEY

pyrobbot/app/app_page_templates.py

@@ -3,6 +3,7 @@
 import uuid
 from abc import ABC, abstractmethod
 from json.decoder import JSONDecodeError
+from typing import TYPE_CHECKING
 
 import streamlit as st
 from loguru import logger
@@ -13,6 +14,9 @@
 from pyrobbot.chat_configs import ChatOptions
 from pyrobbot.openai_utils import CannotConnectToApiError
 
+if TYPE_CHECKING:
+    from pyrobbot.app.multipage import MultipageChatbotApp
+
 _AVATAR_FILES_DIR = GeneralConstants.APP_DIR / "data"
 _ASSISTANT_AVATAR_FILE_PATH = _AVATAR_FILES_DIR / "assistant_avatar.png"
 _USER_AVATAR_FILE_PATH = _AVATAR_FILES_DIR / "user_avatar.png"
@@ -27,7 +31,9 @@
 class AppPage(ABC):
     """Abstract base class for a page within a streamlit application."""
 
-    def __init__(self, sidebar_title: str = "", page_title: str = ""):
+    def __init__(
+        self, parent: "MultipageChatbotApp", sidebar_title: str = "", page_title: str = ""
+    ):
         """Initializes a new instance of the AppPage class.
 
         Args:
@@ -37,7 +43,8 @@ def __init__(self, sidebar_title: str = "", page_title: str = ""):
                 Defaults to an empty string.
         """
         self.page_id = str(uuid.uuid4())
-        self.page_number = st.session_state.get("n_created_pages", 0) + 1
+        self.parent = parent
+        self.page_number = self.parent.state.get("n_created_pages", 0) + 1
 
         chat_number_for_title = f"Chat #{self.page_number}"
         if page_title is _RecoveredChat:
@@ -55,9 +62,9 @@ def __init__(self, sidebar_title: str = "", page_title: str = ""):
     @property
     def state(self):
         """Return the state of the page, for persistence of data."""
-        if self.page_id not in st.session_state:
-            st.session_state[self.page_id] = {}
-        return st.session_state[self.page_id]
+        if self.page_id not in self.parent.state:
+            self.parent.state[self.page_id] = {}
+        return self.parent.state[self.page_id]
 
     @property
     def sidebar_title(self):
@@ -88,7 +95,7 @@ class ChatBotPage(AppPage):
     """Implement a chatbot page in a streamlit application, inheriting from AppPage."""
 
     def __init__(
-        self, chat_obj: Chat = None, sidebar_title: str = "", page_title: str = ""
+        self, parent, chat_obj: Chat = None, sidebar_title: str = "", page_title: str = ""
     ):
         """Initialize new instance of the ChatBotPage class with an optional Chat object.
 
@@ -99,7 +106,9 @@ def __init__(
             page_title (str): The title for the chatbot page.
                 Defaults to an empty string.
         """
-        super().__init__(sidebar_title=sidebar_title, page_title=page_title)
+        super().__init__(
+            parent=parent, sidebar_title=sidebar_title, page_title=page_title
+        )
 
         if chat_obj:
             self.chat_obj = chat_obj
@@ -152,15 +161,28 @@ def render_chat_history(self):
             with st.chat_message(role, avatar=self.avatars.get(role)):
                 st.markdown(message["content"])
 
-    def render(self):
+    def render_cost_estimate_page(self):
+        """Render the estimated costs information in the chat."""
+        general_df = self.chat_obj.general_token_usage_db.get_usage_balance_dataframe()
+        chat_df = self.chat_obj.token_usage_db.get_usage_balance_dataframe()
+        dfs = {"All Recorded Chats": general_df, "Current Chat": chat_df}
+
+        st.header(dfs["Current Chat"].attrs["description"], divider="rainbow")
+        with st.container():
+            for category, df in dfs.items():
+                st.subheader(f"**{category}**")
+                st.dataframe(df)
+                st.write()
+            st.caption(df.attrs["disclaimer"])
+
+    def _render_chatbot_page(self):
         """Render a chatbot page.
 
         Adapted from:
         <https://docs.streamlit.io/knowledge-base/tutorials/build-conversational-apps>
 
         """
-        st.title(self.title)
-        st.divider()
+        st.header(self.title, divider="rainbow")
 
         if self.chat_history:
             self.render_chat_history()
@@ -230,4 +252,11 @@ def render(self):
 
                     self.title = title
                     self.sidebar_title = title
-                    st.title(title)
+                    st.header(title, divider="rainbow")
+
+    def render(self):
+        """Render the app's chatbot or costs page, depending on user choice."""
+        if st.session_state["toggle_show_costs"]:
+            self.render_cost_estimate_page()
+        else:
+            self._render_chatbot_page()

pyrobbot/app/multipage.py

@@ -1,6 +1,6 @@
 """Code for the creation streamlit apps with dynamically created pages."""
 import contextlib
-from abc import ABC, abstractmethod
+from abc import ABC, abstractmethod, abstractproperty
 
 import openai
 import streamlit as st
@@ -28,21 +28,24 @@ def __init__(self, **kwargs) -> None:
     @property
     def n_created_pages(self):
         """Return the number of pages created by the app, including deleted ones."""
-        return st.session_state.get("n_created_pages", 0)
+        return self.state.get("n_created_pages", 0)
 
     @n_created_pages.setter
     def n_created_pages(self, value):
-        st.session_state["n_created_pages"] = value
+        self.state["n_created_pages"] = value
 
     @property
     def pages(self) -> dict[AppPage]:
         """Return the pages of the app."""
-        if "available_pages" not in st.session_state:
-            st.session_state["available_pages"] = {}
-        return st.session_state["available_pages"]
+        if "available_pages" not in self.state:
+            self.state["available_pages"] = {}
+        return self.state["available_pages"]
 
-    def add_page(self, page: AppPage, selected: bool = True):
+    def add_page(self, page: AppPage, selected: bool = True, **page_obj_kwargs):
         """Add a page to the app."""
+        if page is None:
+            page = AppPage(parent=self, **page_obj_kwargs)
+
         self.pages[page.page_id] = page
         self.n_created_pages += 1
         if selected:
@@ -60,24 +63,28 @@ def remove_page(self, page: AppPage):
 
     def register_selected_page(self, page: AppPage):
         """Register a page as selected."""
-        st.session_state["selected_page"] = page
+        self.state["selected_page"] = page
 
     @property
     def selected_page(self) -> ChatBotPage:
         """Return the selected page."""
-        if "selected_page" not in st.session_state:
+        if "selected_page" not in self.state:
             return next(iter(self.pages.values()))
-        return st.session_state["selected_page"]
-
-    @abstractmethod
-    def handle_ui_page_selection(self, **kwargs):
-        """Control page selection in the UI sidebar."""
+        return self.state["selected_page"]
 
     def render(self, **kwargs):
         """Render the multipage app with focus on the selected page."""
         self.handle_ui_page_selection(**kwargs)
         self.selected_page.render()
-        st.session_state["last_rendered_page"] = self.selected_page.page_id
+        self.state["last_rendered_page"] = self.selected_page.page_id
+
+    @abstractproperty
+    def state(self):
+        """Return the state of the app, for persistence of data."""
+
+    @abstractmethod
+    def handle_ui_page_selection(self, **kwargs):
+        """Control page selection in the UI sidebar."""
 
 
 class MultipageChatbotApp(AbstractMultipageApp):
@@ -87,43 +94,55 @@ class MultipageChatbotApp(AbstractMultipageApp):
 
     """
 
-    def init_openai_client(self):
+    @property
+    def state(self):
+        """Return the state of the app, for persistence of data."""
+        app_state_id = f"app_state_{self.openai_api_key}"
+        if app_state_id not in st.session_state:
+            st.session_state[app_state_id] = {}
+        return st.session_state[app_state_id]
+
+    def init_chat_credentials(self):
         """Initializes the OpenAI client with the API key provided in the Streamlit UI."""
-        # Initialize the OpenAI API client
         placeholher = (
             "OPENAI_API_KEY detected"
-            if GeneralConstants.OPENAI_API_KEY
+            if GeneralConstants.SYSTEM_ENV_OPENAI_API_KEY
             else "You need this to use the chat"
         )
-        openai_api_key = st.text_input(
+        self.openai_api_key = st.text_input(
             label="OpenAI API Key (required)",
             placeholder=placeholher,
             key="openai_api_key",
             type="password",
-            help="[OpenAI API auth key](https://platform.openai.com/account/api-keys)",
+            help="[OpenAI API auth key](https://platform.openai.com/account/api-keys). "
+            + "Chats created with this key won't be visible to people using other keys.",
         )
         openai.api_key = (
-            openai_api_key if openai_api_key else GeneralConstants.OPENAI_API_KEY
+            self.openai_api_key
+            if self.openai_api_key
+            else GeneralConstants.SYSTEM_ENV_OPENAI_API_KEY
         )
         if not openai.api_key:
             st.write(":red[You need to provide a key to use the chat]")
 
-    def add_page(self, page: ChatBotPage = None, selected: bool = True, **kwargs):
+    def add_page(
+        self, page: ChatBotPage = None, selected: bool = True, **page_obj_kwargs
+    ):
         """Adds a new ChatBotPage to the app.
 
         If no page is specified, a new instance of ChatBotPage is created and added.
 
         Args:
             page: The ChatBotPage to be added. If None, a new page is created.
             selected: Whether the added page should be selected immediately.
-            **kwargs: Additional keyword arguments for ChatBotPage creation.
+            **page_obj_kwargs: Additional keyword arguments for ChatBotPage creation.
 
         Returns:
             The result of the superclass's add_page method.
 
         """
         if page is None:
-            page = ChatBotPage(**kwargs)
+            page = ChatBotPage(parent=self, **page_obj_kwargs)
         return super().add_page(page=page, selected=selected)
 
     def get_widget_previous_value(self, widget_key, default=None):
@@ -145,7 +164,7 @@ def get_saved_chat_cache_dir_paths(self):
         return sorted(
             (
                 directory
-                for directory in GeneralConstants.CHAT_CACHE_DIR.glob("chat_*/")
+                for directory in GeneralConstants.chat_cache_dir.glob("chat_*/")
                 if next(directory.iterdir(), False)
             ),
             key=lambda fpath: fpath.stat().st_mtime,
@@ -185,17 +204,23 @@ def render(self, **kwargs):
         """Renders the multipage chatbot app in the  UI according to the selected page."""
         with st.sidebar:
             st.title(GeneralConstants.APP_NAME)
-            self.init_openai_client()
+            self.init_chat_credentials()
             # Create a sidebar with tabs for chats and settings
             tab1, tab2 = st.tabs(["Chats", "Settings for Current Chat"])
             self.sidebar_tabs = {"chats": tab1, "settings": tab2}
             with tab1:
+                # Add button to show the costs table
+                st.toggle(
+                    key="toggle_show_costs",
+                    label=":moneybag:",
+                    help="Show estimated token usage and associated costs",
+                )
                 # Add button to create a new chat
                 new_chat_button = st.button(label=":heavy_plus_sign:  New Chat")
 
                 # Reopen chats from cache (if any)
-                if not st.session_state.get("saved_chats_reloaded", False):
-                    st.session_state["saved_chats_reloaded"] = True
+                if not self.state.get("saved_chats_reloaded", False):
+                    self.state["saved_chats_reloaded"] = True
                     for cache_dir_path in self.get_saved_chat_cache_dir_paths():
                         try:
                             chat = Chat.from_cache(cache_dir=cache_dir_path)
@@ -208,6 +233,7 @@ def render(self, **kwargs):
                             continue
 
                         new_page = ChatBotPage(
+                            parent=self,
                             chat_obj=chat,
                             page_title=chat.metadata.get("page_title", _RecoveredChat),
                             sidebar_title=chat.metadata.get("sidebar_title"),