enable multiple chat sessions

Author: lfnovo

migrations/3.surrealql

@@ -1,4 +1,11 @@
-REMOVE FUNCTION fn::vector_search;
+
+DEFINE TABLE IF NOT EXISTS chat_session SCHEMALESS;
+
+DEFINE TABLE IF NOT EXISTS refers_to
+TYPE RELATION 
+FROM chat_session TO notebook;
+
+REMOVE FUNCTION IF EXISTS fn::vector_search;
 
 DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_count: int, $sources: bool, $show_notes: bool, $min_similarity: float) {
     let $source_embedding_search = 
@@ -16,7 +23,6 @@ DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_cou
         )}
         ELSE { [] };
 
-    -- Busca em source_insight com threshold
     let $source_insight_search = 
         IF $sources {(
             SELECT 
@@ -67,10 +73,10 @@ DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_cou
 };
 
 
-REMOVE FUNCTION fn::text_search;
+REMOVE FUNCTION IF EXISTS fn::text_search;
 
 
- DEFINE FUNCTION IF NOT EXISTS fn::text_search($query_text: string, $match_count: int, $sources:bool, $show_notes:bool) {
+DEFINE FUNCTION IF NOT EXISTS fn::text_search($query_text: string, $match_count: int, $sources:bool, $show_notes:bool) {
 
     let $source_title_search = 
         IF $sources {(

migrations/3_down.surrealql

@@ -1,3 +1,8 @@
+REMOVE TABLE IF EXISTS chat_session;
+
+REMOVE TABLE IF EXISTS refers_to;
+
+
 REMOVE FUNCTION fn::vector_search;
 
 

open_notebook/domain/notebook.py

@@ -1,11 +1,9 @@
-import os
 from typing import Any, ClassVar, Dict, List, Literal, Optional
 
 from loguru import logger
 from pydantic import BaseModel, Field, field_validator
 
 from open_notebook.database.repository import (
-    repo_create,
     repo_query,
 )
 from open_notebook.domain.base import ObjectModel
@@ -68,6 +66,27 @@ def notes(self) -> List["Note"]:
             logger.exception(e)
             raise DatabaseOperationError(e)
 
+    @property
+    def chat_sessions(self) -> List["ChatSession"]:
+        try:
+            srcs = repo_query(f"""
+                select * from (
+                    select
+                    <- chat_session as chat_session
+                    from refers_to
+                    where out={self.id}
+                    fetch chat_session
+                )
+                order by chat_session.updated desc
+            """)
+            return (
+                [ChatSession(**src["chat_session"][0]) for src in srcs] if srcs else []
+            )
+        except Exception as e:
+            logger.error(f"Error fetching notes for notebook {self.id}: {str(e)}")
+            logger.exception(e)
+            raise DatabaseOperationError(e)
+
 
 class Asset(BaseModel):
     file_path: Optional[str] = None
@@ -99,6 +118,22 @@ def get_context(
         else:
             return dict(id=self.id, title=self.title, insights=self.insights)
 
+    @property
+    def embedded_chunks(self) -> int:
+        try:
+            result = repo_query(
+                f"""
+                select count() as chunks from source_embedding where source={self.id} GROUP ALL
+                """
+            )
+            if len(result) == 0:
+                return 0
+            return result[0]["chunks"]
+        except Exception as e:
+            logger.error(f"Error fetching insights for source {self.id}: {str(e)}")
+            logger.exception(e)
+            raise DatabaseOperationError(f"Failed to count chunks for source: {str(e)}")
+
     @property
     def insights(self) -> List[SourceInsight]:
         try:
@@ -118,24 +153,6 @@ def add_to_notebook(self, notebook_id: str) -> Any:
             raise InvalidInputError("Notebook ID must be provided")
         return self.relate("reference", notebook_id)
 
-    def save_chunks(self, text: str) -> None:
-        if not text:
-            raise InvalidInputError("Text cannot be empty")
-        try:
-            chunks = split_text(text, chunk=500000, overlap=1000)
-            logger.debug(f"Split into {len(chunks)} chunks")
-            for i, chunk in enumerate(chunks):
-                logger.debug(f"Saving chunk {i}")
-                data = {"source": self.id, "order": i, "content": surreal_clean(chunk)}
-                repo_create(
-                    "source_chunk",
-                    data,
-                )
-        except Exception as e:
-            logger.exception(e)
-            logger.error(f"Error saving chunks for source {self.id}: {str(e)}")
-            raise DatabaseOperationError(e)
-
     def vectorize(self) -> None:
         EMBEDDING_MODEL = model_manager.embedding_model
 
@@ -144,8 +161,6 @@ def vectorize(self) -> None:
                 return
             chunks = split_text(
                 self.full_text,
-                chunk=int(os.environ.get("EMBEDDING_CHUNK_SIZE", 1000)),
-                overlap=int(os.environ.get("EMBEDDING_CHUNK_OVERLAP", 1000)),
             )
             logger.debug(f"Split into {len(chunks)} chunks")
 
@@ -166,26 +181,26 @@ def vectorize(self) -> None:
             logger.exception(e)
             raise DatabaseOperationError(e)
 
-    @classmethod
-    def search(cls, query: str) -> List[Dict[str, Any]]:
-        if not query:
-            raise InvalidInputError("Search query cannot be empty")
-        try:
-            result = repo_query(
-                """
-                SELECT * omit full_text
-                FROM source
-                WHERE string::lowercase(title) CONTAINS $query or title @@ $query 
-                OR string::lowercase(summary) CONTAINS $query or summary @@ $query 
-                OR string::lowercase(full_text) CONTAINS $query or full_text @@ $query 
-            """,
-                {"query": query},
-            )
-            return result
-        except Exception as e:
-            logger.error(f"Error searching sources: {str(e)}")
-            logger.exception(e)
-            raise DatabaseOperationError("Failed to search sources")
+    # @classmethod
+    # def search(cls, query: str) -> List[Dict[str, Any]]:
+    #     if not query:
+    #         raise InvalidInputError("Search query cannot be empty")
+    #     try:
+    #         result = repo_query(
+    #             """
+    #             SELECT * omit full_text
+    #             FROM source
+    #             WHERE string::lowercase(title) CONTAINS $query or title @@ $query
+    #             OR string::lowercase(summary) CONTAINS $query or summary @@ $query
+    #             OR string::lowercase(full_text) CONTAINS $query or full_text @@ $query
+    #         """,
+    #             {"query": query},
+    #         )
+    #         return result
+    #     except Exception as e:
+    #         logger.error(f"Error searching sources: {str(e)}")
+    #         logger.exception(e)
+    #         raise DatabaseOperationError("Failed to search sources")
 
     def add_insight(self, insight_type: str, content: str) -> Any:
         EMBEDDING_MODEL = model_manager.embedding_model
@@ -246,6 +261,16 @@ def get_embedding_content(self) -> Optional[str]:
         return self.content
 
 
+class ChatSession(ObjectModel):
+    table_name: ClassVar[str] = "chat_session"
+    title: Optional[str] = None
+
+    def relate_to_notebook(self, notebook_id: str) -> Any:
+        if not notebook_id:
+            raise InvalidInputError("Notebook ID must be provided")
+        return self.relate("refers_to", notebook_id)
+
+
 def text_search(keyword: str, results: int, source: bool = True, note: bool = True):
     if not keyword:
         raise InvalidInputError("Search keyword cannot be empty")
@@ -263,18 +288,6 @@ def text_search(keyword: str, results: int, source: bool = True, note: bool = Tr
         raise DatabaseOperationError(e)
 
 
-# def hybrid_search(
-#     keyword_search: List[str],
-#     embed_search: List[str],
-#     results: int = 50,
-#     source: bool = True,
-#     note: bool = True,
-# ):
-# EMBEDDING_MODEL = model_manager.embedding_model
-# embed1_vector = EMBEDDING_MODEL.embed(embed1) if embed1 else None
-
-
-# todo: mover o embedding pra ca
 def vector_search(keyword: str, results: int, source: bool = True, note: bool = True):
     if not keyword:
         raise InvalidInputError("Search keyword cannot be empty")

pages/2_📒_Notebooks.py

@@ -10,11 +10,14 @@
 setup_page("📒 Open Notebook")
 
 
-def notebook_header(current_notebook):
+def notebook_header(current_notebook: Notebook):
+    """
+    Defines the header of the notebook page, including the ability to edit the notebook's name and description.
+    """
     c1, c2, c3 = st.columns([8, 2, 2])
     c1.header(current_notebook.name)
     if c2.button("Back to the list", icon="🔙"):
-        st.session_state["current_notebook"] = None
+        st.session_state["current_notebook_id"] = None
         st.rerun()
 
     if c3.button("Refresh", icon="🔄"):
@@ -49,20 +52,20 @@ def notebook_header(current_notebook):
                 st.toast("Notebook unarchived", icon="🗃️")
         if c3.button("Delete forever", type="primary", icon="☠️"):
             current_notebook.delete()
-            st.session_state["current_notebook"] = None
+            st.session_state["current_notebook_id"] = None
             st.rerun()
 
 
-def notebook_page(current_notebook_id):
-    current_notebook: Notebook = Notebook.get(current_notebook_id)
-    if not current_notebook:
-        st.error("Notebook not found")
-        return
-    if current_notebook_id not in st.session_state.keys():
-        st.session_state[current_notebook_id] = current_notebook
+def notebook_page(current_notebook: Notebook):
+    # Guarantees that we have an entry for this notebook in the session state
+    if current_notebook.id not in st.session_state:
+        st.session_state[current_notebook.id] = {"notebook": current_notebook}
+
+    # sets up the active session
+    current_session = setup_stream_state(
+        current_notebook=current_notebook,
+    )
 
-    session_id = st.session_state["active_session"]
-    st.session_state[session_id]["notebook"] = current_notebook
     sources = current_notebook.sources
     notes = current_notebook.notes
 
@@ -74,18 +77,18 @@ def notebook_page(current_notebook_id):
         with sources_tab:
             with st.container(border=True):
                 if st.button("Add Source", icon="➕"):
-                    add_source(session_id)
+                    add_source(current_notebook.id)
                 for source in sources:
-                    source_card(session_id=session_id, source=source)
+                    source_card(source=source, notebook_id=current_notebook.id)
 
         with notes_tab:
             with st.container(border=True):
                 if st.button("Write a Note", icon="📝"):
-                    add_note(session_id)
+                    add_note(current_notebook.id)
                 for note in notes:
-                    note_card(session_id=session_id, note=note)
+                    note_card(note=note, notebook_id=current_notebook.id)
     with chat_tab:
-        chat_sidebar(session_id=session_id)
+        chat_sidebar(current_notebook=current_notebook, current_session=current_session)
 
 
 def notebook_list_item(notebook):
@@ -96,40 +99,50 @@ def notebook_list_item(notebook):
         )
         st.write(notebook.description)
         if st.button("Open", key=f"open_notebook_{notebook.id}"):
-            setup_stream_state(notebook.id)
-            st.session_state["current_notebook"] = notebook.id
+            st.session_state["current_notebook_id"] = notebook.id
             st.rerun()
 
 
-if "current_notebook" not in st.session_state:
-    st.session_state["current_notebook"] = None
+if "current_notebook_id" not in st.session_state:
+    st.session_state["current_notebook_id"] = None
 
-if st.session_state["current_notebook"]:
-    notebook_page(st.session_state["current_notebook"])
+# todo: get the notebook, check if it exists and if it's archived
+if st.session_state["current_notebook_id"]:
+    current_notebook: Notebook = Notebook.get(st.session_state["current_notebook_id"])
+    if not current_notebook:
+        st.error("Notebook not found")
+        st.stop()
+    notebook_page(current_notebook)
     st.stop()
 
 st.title("📒 My Notebooks")
-st.caption("Here are all your notebooks")
-
-notebooks = Notebook.get_all(order_by="updated desc")
-
-for notebook in notebooks:
-    if notebook.archived:
-        continue
-    notebook_list_item(notebook)
+st.caption(
+    "Notebooks are a great way to organize your thoughts, ideas, and sources. You can create notebooks for different research topics and projects, to create new articles, etc. "
+)
 
 with st.expander("➕ **New Notebook**"):
     new_notebook_title = st.text_input("New Notebook Name")
-    new_notebook_description = st.text_area("Description")
+    new_notebook_description = st.text_area(
+        "Description",
+        placeholder="Explain the purpose of this notebook. The more details the better.",
+    )
     if st.button("Create a new Notebook", icon="➕"):
         notebook = Notebook(
             name=new_notebook_title, description=new_notebook_description
         )
         notebook.save()
-        st.rerun()
+        st.toast("Notebook created successfully", icon="📒")
 
+notebooks = Notebook.get_all(order_by="updated desc")
 archived_notebooks = [nb for nb in notebooks if nb.archived]
+
+for notebook in notebooks:
+    if notebook.archived:
+        continue
+    notebook_list_item(notebook)
+
 if len(archived_notebooks) > 0:
     with st.expander(f"**🗃️ {len(archived_notebooks)} archived Notebooks**"):
+        st.write("ℹ Archived Notebooks can still be accessed and used in search.")
         for notebook in archived_notebooks:
             notebook_list_item(notebook)