fixed new methods in client.py

Author: anngoroshi

test2text/pages/controls/controls_page.py

@@ -13,10 +13,10 @@ def controls_page():
         def refresh_counts():
             with get_db_client() as db:
                 st.session_state["all_annotations_count"] = (
-                    db.count_all_entries_in_table("Annotations")
+                    db.count_all_entries("Annotations")
                 )
                 st.session_state["embedded_annotations_count"] = (
-                    db.count_embedded_entries_in_table("Annotations")
+                    db.count_notnull_entries("embedding",from_table="Annotations")
                 )
 
         refresh_counts()

test2text/services/db/client.py

@@ -88,7 +88,7 @@ def get_table_names(self):
         cursor.close()
         return tables
 
-    def get_column_values(self, columns: list[str], from_table: str):
+    def get_column_values(self, *columns: str, from_table: str):
         cursor = self.conn.execute(f"SELECT {', '.join(columns)} FROM {from_table}")
         return cursor.fetchall()
 
@@ -117,10 +117,10 @@ def count_all_entries(self, from_table: str) -> int:
         return count
 
     def count_notnull_entries(
-        self, columns: list[str], from_table: str
-    ) -> Union[int, None]:
+        self, *columns: str, from_table: str
+    ) -> int:
         count = self.conn.execute(
-            f"SELECT COUNT(*) FROM {from_table} WHERE {', '.join(columns)} IS NOT NULL"
+            f"SELECT COUNT(*) FROM {from_table} WHERE {' AND  '.join([column + ' IS NOT NULL' for column in columns])}"
         ).fetchone()[0]
         return count
 

test2text/services/embeddings/annotation_embeddings_controls.py

@@ -12,8 +12,8 @@ def embed_annotations(*_, embed_all=False, on_progress: OnProgress = None):
     with get_db_client() as db:
         from .embed import embed_annotations_batch
 
-        annotations_count = db.count_all_entries_in_table("Annotations")
-        embedded_annotations_count = db.count_embedded_entries_in_table("Annotations")
+        annotations_count = db.count_all_entries("Annotations")
+        embedded_annotations_count = db.count_notnull_entries("embedding",from_table="Annotations")
         if embed_all:
             annotations_to_embed = annotations_count
         else: