chore: resolve suggestions + tests

Author: AstraBert

src/notebookllama/Home.py

@@ -150,28 +150,36 @@ def sync_create_podcast(file_content: str):
 st.markdown("---")
 st.markdown("## NotebookLlaMa - Home🦙")
 
+# Initialize session state BEFORE creating the text input
+if "workflow_results" not in st.session_state:
+    st.session_state.workflow_results = None
+if "document_title" not in st.session_state:
+    st.session_state.document_title = randomname.get_name(
+        adj=("music_theory", "geometry", "emotions"), noun=("cats", "food")
+    )
+
+# Use session_state as the value and update it when changed
 document_title = st.text_input(
     label="Document Title",
-    value=randomname.get_name(
-        adj=("music_theory", "geometry", "emotions"), noun=("cats", "food")
-    ),
+    value=st.session_state.document_title,
+    key="document_title_input",
 )
+
+# Update session state when the input changes
+if document_title != st.session_state.document_title:
+    st.session_state.document_title = document_title
+
 file_input = st.file_uploader(
     label="Upload your source PDF file!", accept_multiple_files=False
 )
 
-
-# Initialize session state
-if "workflow_results" not in st.session_state:
-    st.session_state.workflow_results = None
-
 if file_input is not None:
     # First button: Process Document
     if st.button("Process Document", type="primary"):
         with st.spinner("Processing document... This may take a few minutes."):
             try:
                 md_content, summary, q_and_a, bullet_points, mind_map = (
-                    sync_run_workflow(file_input, document_title)
+                    sync_run_workflow(file_input, st.session_state.document_title)
                 )
                 st.session_state.workflow_results = {
                     "md_content": md_content,

src/notebookllama/documents.py

@@ -37,11 +37,13 @@ def __init__(
         table_metadata: Optional[MetaData] = None,
     ):
         self.table_name: str = table_name or "documents"
-        self.table: Optional[Table] = None
+        self._table: Optional[Table] = None
         self._connection: Optional[Connection] = None
-        self.metadata: Optional[MetaData] = table_metadata or MetaData()
+        self.metadata: MetaData = cast(MetaData, table_metadata or MetaData())
         if engine or engine_url:
-            self._engine: Union[Engine, str] = engine or engine_url
+            self._engine: Union[Engine, str] = cast(
+                Union[Engine, str], engine or engine_url
+            )
         else:
             raise ValueError("One of engine or engine_setup_kwargs must be set")
 
@@ -51,14 +53,20 @@ def connection(self) -> Connection:
             self._connect()
         return cast(Connection, self._connection)
 
+    @property
+    def table(self) -> Table:
+        if not self._table:
+            self._create_table()
+        return cast(Table, self._table)
+
     def _connect(self) -> None:
         # move network calls outside of constructor
         if isinstance(self._engine, str):
             self._engine = create_engine(self._engine)
         self._connection = self._engine.connect()
 
     def _create_table(self) -> None:
-        self.table = Table(
+        self._table = Table(
             self.table_name,
             self.metadata,
             Column("id", Integer, primary_key=True, autoincrement=True),
@@ -69,11 +77,9 @@ def _create_table(self) -> None:
             Column("mindmap", Text),
             Column("bullet_points", Text),
         )
-        self.table.create(self.connection, checkfirst=True)
+        self._table.create(self.connection, checkfirst=True)
 
     def put_documents(self, documents: List[ManagedDocument]) -> None:
-        if not self.table:
-            self._create_table()
         for document in documents:
             stmt = insert(self.table).values(
                 document_name=document.document_name,
@@ -87,7 +93,7 @@ def put_documents(self, documents: List[ManagedDocument]) -> None:
         self.connection.commit()
 
     def get_documents(self, names: Optional[List[str]] = None) -> List[ManagedDocument]:
-        if not self.table_exists:
+        if self.table is None:
             self._create_table()
         if not names:
             stmt = select(self.table).order_by(self.table.c.id)
@@ -114,7 +120,7 @@ def get_documents(self, names: Optional[List[str]] = None) -> List[ManagedDocume
         return documents
 
     def get_names(self) -> List[str]:
-        if not self.table_exists:
+        if self.table is None:
             self._create_table()
         stmt = select(self.table)
         result = self.connection.execute(stmt)
@@ -124,4 +130,7 @@ def get_names(self) -> List[str]:
     def disconnect(self) -> None:
         if not self._connection:
             raise ValueError("Engine was never connected!")
-        self._engine.dispose(close=True)
+        if isinstance(self._engine, str):
+            pass
+        else:
+            self._engine.dispose(close=True)

tests/test_document_management.py

@@ -5,7 +5,7 @@
 from typing import List
 
 from src.notebookllama.documents import DocumentManager, ManagedDocument
-from sqlalchemy import text
+from sqlalchemy import text, Table
 
 ENV = load_dotenv()
 
@@ -63,10 +63,11 @@ def documents() -> List[ManagedDocument]:
 def test_document_manager(documents: List[ManagedDocument]) -> None:
     engine_url = f"postgresql+psycopg2://{os.getenv('pgql_user')}:{os.getenv('pgql_psw')}@localhost:5432/{os.getenv('pgql_db')}"
     manager = DocumentManager(engine_url=engine_url, table_name="test_documents")
-    assert not manager.table_exists
-    manager._execute(text("DROP TABLE IF EXISTS test_documents;"))
+    assert not manager.table
+    manager.connection.execute(text("DROP TABLE IF EXISTS test_documents;"))
+    manager.connection.commit()
     manager._create_table()
-    assert manager.table_exists
+    assert isinstance(manager.table, Table)
     manager.put_documents(documents=documents)
     names = manager.get_names()
     assert names == [doc.document_name for doc in documents]

tests/test_models.py

@@ -197,13 +197,4 @@ def test_managed_documents() -> None:
         mindmap="Hello -> World",
         bullet_points=". Hello, . World",
     )
-    assert d2.summary == "Test''s child"
-    with pytest.raises(ValidationError):
-        ManagedDocument(
-            document_name=1,
-            content="This is a test",
-            summary="Test's child",
-            q_and_a="Hello? World.",
-            mindmap="Hello -> World",
-            bullet_points=". Hello, . World",
-        )
+    assert d2.summary == "Test's child"