Update utils/database.py

utils/database.py

@@ -267,6 +267,83 @@ def generate_document_tags(content: str) -> List[str]:
         st.error(f"Error generating tags: {e}")
         return []
 
+
+def add_document_to_collection(conn: sqlite3.Connection, document_id: int, collection_id: int) -> bool:
+    """Link a document to a collection."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                INSERT OR IGNORE INTO document_collections (document_id, collection_id)
+                VALUES (?, ?)
+            ''', (document_id, collection_id))
+            conn.commit()
+            return True
+    except sqlite3.Error as e:
+        st.error(f"Error linking document to collection: {e}")
+        return False
+
+def get_collection_documents(conn: sqlite3.Connection, collection_id: int) -> List[Dict]:
+    """Get all documents in a collection."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT 
+                    d.id,
+                    d.name,
+                    d.content,
+                    d.upload_date
+                FROM documents d
+                JOIN document_collections dc ON d.id = dc.document_id
+                WHERE dc.collection_id = ?
+                ORDER BY d.upload_date DESC
+            ''', (collection_id,))
+            
+            documents = []
+            for row in cursor.fetchall():
+                documents.append({
+                    'id': row[0],
+                    'name': row[1],
+                    'content': row[2],
+                    'upload_date': row[3]
+                })
+            return documents
+    except sqlite3.Error as e:
+        st.error(f"Error retrieving collection documents: {e}")
+        return []
+
+def get_collections(conn: sqlite3.Connection) -> List[Dict]:
+    """Get all collections with document counts."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT 
+                    c.id,
+                    c.name,
+                    c.description,
+                    c.created_at,
+                    COUNT(DISTINCT dc.document_id) as doc_count
+                FROM collections c
+                LEFT JOIN document_collections dc ON c.id = dc.collection_id
+                GROUP BY c.id
+                ORDER BY c.name
+            ''')
+            
+            collections = []
+            for row in cursor.fetchall():
+                collections.append({
+                    'id': row[0],
+                    'name': row[1],
+                    'description': row[2],
+                    'created_at': row[3],
+                    'doc_count': row[4]
+                })
+            return collections
+    except sqlite3.Error as e:
+        st.error(f"Error retrieving collections: {e}")
+        return []
 def add_document_tags(conn: sqlite3.Connection, document_id: int, tags: List[str]) -> bool:
     """Add tags to a document."""
     try: