Merge pull request #126 from sudoleg/develop

introduce library

Author: sudoleg

.gitignore

@@ -17,6 +17,7 @@ transcript_audios/
 data/audio/*.mp3
 data/transcripts/*.txt
 scripts/tokens.ipynb
+scripts/persistance.ipynb
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

modules/helpers.py

@@ -17,6 +17,7 @@ def is_api_key_set() -> bool:
     return False
 
 
+@st.cache_data
 def is_api_key_valid(api_key: str):
     """
     Checks the validity of an OpenAI API key.
@@ -28,10 +29,6 @@ def is_api_key_valid(api_key: str):
         bool: True if the API key is valid, False if the API key is invalid.
     """
 
-    api_key_valid = os.getenv("OPENAI_API_KEY_VALID")
-    if api_key_valid:
-        return True
-
     openai.api_key = api_key
     try:
         openai.models.list()
@@ -49,7 +46,6 @@ def is_api_key_valid(api_key: str):
         return False
     else:
         logging.info("API key validation successful")
-        os.environ["OPENAI_API_KEY_VALID"] = "yes"
         return True
 
 

modules/persistance.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Literal
 
 from peewee import (
     BooleanField,
@@ -8,6 +9,7 @@
     IntegerField,
     Model,
     SqliteDatabase,
+    TextField,
     UUIDField,
 )
 
@@ -77,3 +79,43 @@ def delete_video(
             video.yt_video_id,
             str(e),
         )
+
+
+class LibraryEntry(BaseModel):
+    """Model for saved responses and summaries. Represents a table in a relational SQL database."""
+
+    ENTRY_TYPE_CHOICES = (
+        ("S", "Summary"),
+        ("A", "Answer"),
+    )
+
+    entry_type = CharField(max_length=1, choices=ENTRY_TYPE_CHOICES)
+    video = ForeignKeyField(Video, backref="lib_entries")
+    question = TextField(null=True)
+    text = TextField(null=False)
+
+
+def save_library_entry(
+    entry_type: Literal["S", "A"], question_text: str, response_text: str, video: Video
+):
+    """Saves a summary or answer entry to the library.
+
+    Args:
+        entry_type (str): Type of entry to save, "S" for summary or "A" for answer.
+        question_text (str): Text of the question (used only if entry_type is "A").
+        response_text (str): Text of the response or summary.
+        video (Video): The video object associated with the entry.
+    """
+    if entry_type == "S":
+        LibraryEntry.create(entry_type="S", video=video, text=response_text)
+    else:
+        LibraryEntry.create(
+            entry_type="A", video=video, question=question_text, text=response_text
+        )
+    logging.info("Saved library entry for video '%s'", video.title)
+
+
+def delete_library_entry(lib_entry: LibraryEntry):
+    """Deletes a library entry."""
+    LibraryEntry.delete_by_id(lib_entry)
+    logging.info("Deleted library entry for video '%s'", lib_entry.video.title)

modules/ui.py

@@ -127,3 +127,4 @@ def display_nav_menu():
     st.sidebar.page_link(page="main.py", label="Home")
     st.sidebar.page_link(page="pages/summary.py", label="Summary")
     st.sidebar.page_link(page="pages/chat.py", label="Chat")
+    st.sidebar.page_link(page="pages/library.py", label="Library")

pages/chat.py

@@ -19,7 +19,14 @@
     read_file,
     save_response_as_file,
 )
-from modules.persistance import SQL_DB, Transcript, Video, delete_video
+from modules.persistance import (
+    SQL_DB,
+    LibraryEntry,
+    Transcript,
+    Video,
+    delete_video,
+    save_library_entry,
+)
 from modules.rag import (
     CHUNK_SIZE_TO_K_MAPPING,
     embed_excerpts,
@@ -60,7 +67,7 @@
 # --- SQLite stuff ---
 SQL_DB.connect(reuse_if_open=True)
 # create tables if they don't already exist
-SQL_DB.create_tables([Video, Transcript], safe=True)
+SQL_DB.create_tables([Video, Transcript, LibraryEntry], safe=True)
 # --- end ---
 
 # --- Chroma ---
@@ -95,6 +102,26 @@ def refresh_page(message: str):
         st.rerun()
 
 
+# variable for holding the Video object
+saved_video: None | Video = None
+
+
+def save_response_to_lib():
+    """Wrapper func for saving responses to the library."""
+    try:
+        save_library_entry(
+            entry_type="A",
+            question_text=st.session_state.user_prompt,
+            response_text=st.session_state.response,
+            video=saved_video,
+        )
+    except Exception as e:
+        st.error("Saving failed! If you are a developer, see logs for details!")
+        logging.error("Error when saving library entry: %s", e)
+    else:
+        st.success("Saved answer to library successfully!")
+
+
 if (
     is_api_key_set()
     and is_api_key_valid(st.session_state.openai_api_key)
@@ -129,7 +156,7 @@ def refresh_page(message: str):
     )
     # --- end ---
 
-    # fetch saved videos from SQLite
+    # fetch all saved videos from SQLite
     saved_videos = Video.select()
 
     # create columns
@@ -139,7 +166,10 @@ def refresh_page(message: str):
         selected_video_title = st.selectbox(
             label="Select from already processed videos",
             placeholder="Choose a video",
-            options=[video.title for video in saved_videos],
+            # only videos with an associated transcript can be selected
+            options=[
+                video.title for video in saved_videos if video.transcripts.count() != 0
+            ],
             index=None,
             key="selected_video",
             help=get_default_config_value("help_texts.selected_video"),
@@ -148,7 +178,6 @@ def refresh_page(message: str):
             label="Or enter the URL of a new video:", disabled=is_video_selected()
         )
 
-        saved_video = None
         if is_video_selected():
             saved_video = Video.get(Video.title == selected_video_title)
 
@@ -218,16 +247,18 @@ def refresh_page(message: str):
                         url=url_input,
                     )
 
-                    # 1. fetch transcript from youtube
+                    # 1. save video in the database
                     saved_video = Video.create(
                         yt_video_id=extract_youtube_video_id(url_input),
                         link=url_input,
                         title=video_metadata["name"],
                         channel=video_metadata["channel"],
                         saved_on=dt.now(),
                     )
+                    # 2. fetch transcript from youtube
                     original_transcript = fetch_youtube_transcript(url_input)
 
+                    # 3. save transcript, ormore precisely, information about it, in the database
                     saved_transcript = Transcript.create(
                         video=saved_video,
                         original_token_num=num_tokens_from_string(
@@ -236,6 +267,7 @@ def refresh_page(message: str):
                         ),
                     )
 
+                    # 4. get an already existing or create a new collection in ChromaDB
                     collection = chroma_client.get_or_create_collection(
                         name=randomname.get_name(),
                         metadata={
@@ -245,7 +277,7 @@ def refresh_page(message: str):
                         },
                     )
 
-                    # 2. create excerpts. Either
+                    # 5. create excerpts. Either
                     #   - from original transcript
                     #   - or from whisper transcription if transcription checkbox is checked
                     if transcription_checkbox:
@@ -271,7 +303,7 @@ def refresh_page(message: str):
                             len_func="tokens",
                         )
 
-                    # 3. embed/index transcript excerpts
+                    # 6. embed/index transcript excerpts
                     Transcript.update(
                         {
                             Transcript.preprocessed: transcription_checkbox,
@@ -322,10 +354,10 @@ def refresh_page(message: str):
 
             prompt = st.chat_input(
                 placeholder="Ask a question or provide a topic covered in the video",
-                key="user_prompt",
             )
 
             if prompt:
+                st.session_state.user_prompt = prompt
                 with st.spinner("Generating answer..."):
                     try:
                         relevant_docs = find_relevant_documents(
@@ -340,13 +372,19 @@ def refresh_page(message: str):
                             llm=openai_chat_model,
                             relevant_docs=relevant_docs,
                         )
+                        st.session_state.response = response
                     except Exception as e:
                         logging.error(
                             "An unexpected error occurred: %s", str(e), exc_info=True
                         )
                         st.error(GENERAL_ERROR_MESSAGE)
                     else:
-                        st.write(response)
+                        st.write(st.session_state.response)
+                        st.button(
+                            label="Save this response to your library",
+                            on_click=save_response_to_lib,
+                            help="Unfortunately, the response disappears in this view after saving it to the library. However, it will be visible on the 'Library' page!",
+                        )
                         with st.expander(
                             label="Show chunks retrieved from index and provided to the model as context"
                         ):

pages/library.py

@@ -0,0 +1,58 @@
+import streamlit as st
+
+from modules.persistance import SQL_DB, LibraryEntry, delete_library_entry
+from modules.ui import display_nav_menu
+
+
+st.set_page_config("Library", layout="wide", initial_sidebar_state="auto")
+display_nav_menu()
+
+# --- SQLite stuff ---
+SQL_DB.connect(reuse_if_open=True)
+# create tables if they don't already exist
+SQL_DB.create_tables([LibraryEntry], safe=True)
+# --- end ---
+
+saved_lib_entries_summaries = (
+    LibraryEntry.select().where(LibraryEntry.entry_type == "S").execute()
+)
+saved_lib_entries_answers = (
+    LibraryEntry.select().where(LibraryEntry.entry_type == "A").execute()
+)
+
+
+def execute_entry_deletion(entry: LibraryEntry):
+    """Wrapper func for deleting a library entry."""
+    delete_library_entry(entry)
+    st.rerun()
+
+
+if saved_lib_entries_summaries:
+    st.header("Saved summaries")
+    for i, entry in enumerate(saved_lib_entries_summaries, 0):
+        st.caption(f"{entry.video.title} - {entry.video.channel}")
+        with st.expander("Show"):
+            st.write(entry.text)
+        if st.button(
+            label="Delete entry",
+            key=f"delete_summary_{i}",
+        ):
+            execute_entry_deletion(entry)
+        st.divider()
+else:
+    st.info("You don't have any saved summaries yet!")
+
+if saved_lib_entries_answers:
+    st.header("Saved answers")
+    for j, entry in enumerate(saved_lib_entries_answers, 0):
+        st.subheader(entry.question)
+        with st.expander("Show"):
+            st.write(entry.text)
+        if st.button(
+            label="Delete entry",
+            key=f"delete_answer_{j}",
+        ):
+            execute_entry_deletion(entry)
+        st.divider()
+else:
+    st.info("You don't have any saved answers yet!")