Update utils/database.py

utils/database.py

@@ -50,7 +50,183 @@ def create_connection(db_file):
     except Error as e:
         st.error("Failed to connect to database. Please try again or contact support.")
     return None
+def create_connection(db_file: str) -> Optional[sqlite3.Connection]:
+    """Create a database connection."""
+    try:
+        conn = sqlite3.connect(db_file, check_same_thread=False)
+        return conn
+    except sqlite3.Error as e:
+        st.error(f"Error connecting to database: {e}")
+        return None
 
+def create_tables(conn: sqlite3.Connection) -> None:
+    """
+    Create all necessary tables in the database.
+    
+    Args:
+        conn (sqlite3.Connection): SQLite database connection
+    """
+    try:
+        with conn_lock:
+            # Documents table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS documents (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    name TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    upload_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            ''')
+            
+            # Queries table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS queries (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    query TEXT NOT NULL,
+                    response TEXT NOT NULL,
+                    document_id INTEGER,
+                    query_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE
+                )
+            ''')
+            
+            # Annotations table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS annotations (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    document_id INTEGER NOT NULL,
+                    annotation TEXT NOT NULL,
+                    page_number INTEGER,
+                    annotation_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE
+                )
+            ''')
+            
+            # Collections table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS collections (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    name TEXT NOT NULL UNIQUE,
+                    description TEXT,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            ''')
+            
+            # Document-Collections relationship table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS document_collections (
+                    document_id INTEGER,
+                    collection_id INTEGER,
+                    added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    PRIMARY KEY (document_id, collection_id),
+                    FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE,
+                    FOREIGN KEY (collection_id) REFERENCES collections (id) ON DELETE CASCADE
+                )
+            ''')
+            
+            # Add indices for better performance
+            conn.execute('CREATE INDEX IF NOT EXISTS idx_doc_name ON documents(name)')
+            conn.execute('CREATE INDEX IF NOT EXISTS idx_collection_name ON collections(name)')
+            conn.execute('CREATE INDEX IF NOT EXISTS idx_query_date ON queries(query_date)')
+            conn.execute('CREATE INDEX IF NOT EXISTS idx_doc_collections ON document_collections(collection_id)')
+            
+            conn.commit()
+            
+    except sqlite3.Error as e:
+        st.error(f"Error creating tables: {e}")
+        raise
+
+def get_all_documents(conn: sqlite3.Connection) -> List[Dict]:
+    """Get all documents with their metadata and collections."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT 
+                    d.id,
+                    d.name,
+                    d.content,
+                    d.upload_date,
+                    GROUP_CONCAT(c.name) as collections
+                FROM documents d
+                LEFT JOIN document_collections dc ON d.id = dc.document_id
+                LEFT JOIN collections c ON dc.collection_id = c.id
+                GROUP BY d.id
+                ORDER BY d.upload_date DESC
+            ''')
+            
+            documents = []
+            for row in cursor.fetchall():
+                documents.append({
+                    'id': row[0],
+                    'name': row[1],
+                    'content': row[2],
+                    'upload_date': row[3],
+                    'collections': row[4].split(',') if row[4] else []
+                })
+            return documents
+            
+    except sqlite3.Error as e:
+        st.error(f"Error retrieving documents: {e}")
+        return []
+
+def get_document_queries(conn: sqlite3.Connection, document_id: int) -> List[Dict]:
+    """Get all queries associated with a document."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT id, query, response, query_date
+                FROM queries
+                WHERE document_id = ?
+                ORDER BY query_date DESC
+            ''', (document_id,))
+            
+            queries = []
+            for row in cursor.fetchall():
+                queries.append({
+                    'id': row[0],
+                    'query': row[1],
+                    'response': row[2],
+                    'query_date': row[3]
+                })
+            return queries
+            
+    except sqlite3.Error as e:
+        st.error(f"Error retrieving document queries: {e}")
+        return []
+
+def add_query(conn: sqlite3.Connection, query: str, response: str, document_id: Optional[int] = None) -> bool:
+    """Add a new query and its response."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                INSERT INTO queries (query, response, document_id)
+                VALUES (?, ?, ?)
+            ''', (query, response, document_id))
+            conn.commit()
+            return True
+            
+    except sqlite3.Error as e:
+        st.error(f"Error adding query: {e}")
+        return False
+
+def add_annotation(conn: sqlite3.Connection, document_id: int, annotation: str, page_number: Optional[int] = None) -> bool:
+    """Add an annotation to a document."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                INSERT INTO annotations (document_id, annotation, page_number)
+                VALUES (?, ?, ?)
+            ''', (document_id, annotation, page_number))
+            conn.commit()
+            return True
+            
+    except sqlite3.Error as e:
+        st.error(f"Error adding annotation: {e}")
+        return False
 
 def create_tables(conn):
     """
@@ -169,9 +345,6 @@ def verify_vector_store(vector_store):
         st.error(f"Vector store verification failed: {e}")
         return False
 
-
-
-
 def handle_document_upload(uploaded_files, **kwargs):
     """
     Handle document upload with progress tracking and collection support.