feat: add library feature, closes #145

sudoleg · sudoleg · commit e6ec6396545c · 2024-10-26T13:36:20.000+02:00
diff --git a/modules/persistance.py b/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)
diff --git a/modules/ui.py b/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")
diff --git a/pages/chat.py b/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"
                         ):
diff --git a/pages/library.py b/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!")
diff --git a/pages/summary.py b/pages/summary.py
@@ -1,15 +1,18 @@
 import logging
+from datetime import datetime as dt
 
 import streamlit as st
 from langchain_community.callbacks.openai_info import OpenAICallbackHandler
 from langchain_openai import ChatOpenAI
 
 from modules.helpers import (
+    extract_youtube_video_id,
     get_default_config_value,
     is_api_key_set,
     is_api_key_valid,
     save_response_as_file,
 )
+from modules.persistance import SQL_DB, LibraryEntry, Video, save_library_entry
 from modules.summary import TranscriptTooLongForModelException, get_transcript_summary
 from modules.ui import (
     GENERAL_ERROR_MESSAGE,
@@ -27,6 +30,12 @@
     get_video_metadata,
 )
 
+# --- SQLite stuff ---
+SQL_DB.connect(reuse_if_open=True)
+# create tables if they don't already exist
+SQL_DB.create_tables([Video, LibraryEntry], safe=True)
+# --- end ---
+
 st.set_page_config("Summaries", layout="wide", initial_sidebar_state="auto")
 display_api_key_warning()
 
@@ -36,6 +45,25 @@
 display_link_to_repo("summary")
 # --- end ---
 
+# variable for holding the Video object
+saved_video: None | Video = None
+
+
+def save_summary_to_lib():
+    """Wrapper func for saving summaries to the library."""
+    try:
+        save_library_entry(
+            entry_type="S",
+            question_text=None,
+            response_text=st.session_state.summary,
+            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 summary to library successfully!")
+
 
 @st.dialog(title="Transcript too long", width="large")
 def display_dialog(message: str):
@@ -85,6 +113,13 @@ def display_dialog(message: str):
     with col2:
         if summarize_button:
             try:
+                saved_video = Video.create(
+                    yt_video_id=extract_youtube_video_id(url_input),
+                    link=url_input,
+                    title=vid_metadata["name"],
+                    channel=vid_metadata["channel"],
+                    saved_on=dt.now(),
+                )
                 transcript = fetch_youtube_transcript(url_input)
                 cb = OpenAICallbackHandler()
                 llm = ChatOpenAI(
@@ -102,7 +137,12 @@ def display_dialog(message: str):
                         )
                     else:
                         resp = get_transcript_summary(transcript, llm)
-                st.markdown(resp)
+                    st.session_state.summary = resp
+                st.markdown(st.session_state.summary)
+
+                # button for saving summary to library
+                st.button(label="Save summary to library", on_click=save_summary_to_lib)
+
                 st.caption(
                     f"The estimated cost for the request is: {cb.total_cost:.4f}$"
                 )
