Update utils/database.py

utils/database.py

@@ -496,7 +496,155 @@ def handle_document_upload(uploaded_files, **kwargs):
         st.session_state.chat_ready = False
         return False
 
+# Add these to your database.py file
 
+def remove_from_collection(conn: sqlite3.Connection, document_id: int, collection_id: int) -> bool:
+    """
+    Remove a document from a collection.
+    
+    Args:
+        conn (sqlite3.Connection): Database connection
+        document_id (int): ID of the document to remove
+        collection_id (int): ID of the collection
+        
+    Returns:
+        bool: True if successful
+    """
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                DELETE FROM document_collections 
+                WHERE document_id = ? AND collection_id = ?
+            ''', (document_id, collection_id))
+            conn.commit()
+            return True
+            
+    except sqlite3.Error as e:
+        st.error(f"Error removing document from collection: {e}")
+        return False
+
+def update_collection(conn: sqlite3.Connection, collection_id: int, name: Optional[str] = None, 
+                     description: Optional[str] = None) -> bool:
+    """
+    Update collection details.
+    
+    Args:
+        conn (sqlite3.Connection): Database connection
+        collection_id (int): ID of the collection to update
+        name (Optional[str]): New name for the collection
+        description (Optional[str]): New description for the collection
+        
+    Returns:
+        bool: True if successful
+    """
+    try:
+        with conn_lock:
+            updates = []
+            params = []
+            
+            if name is not None:
+                updates.append("name = ?")
+                params.append(name)
+            if description is not None:
+                updates.append("description = ?")
+                params.append(description)
+                
+            if not updates:
+                return True  # Nothing to update
+                
+            params.append(collection_id)
+            
+            cursor = conn.cursor()
+            cursor.execute(f'''
+                UPDATE collections 
+                SET {", ".join(updates)}
+                WHERE id = ?
+            ''', params)
+            
+            conn.commit()
+            return True
+            
+    except sqlite3.Error as e:
+        st.error(f"Error updating collection: {e}")
+        return False
+
+def search_documents(conn: sqlite3.Connection, query: str, 
+                    collection_id: Optional[int] = None, 
+                    filters: Optional[Dict] = None) -> List[Dict]:
+    """
+    Search documents using fuzzy matching and filters.
+    
+    Args:
+        conn (sqlite3.Connection): Database connection
+        query (str): Search query
+        collection_id (Optional[int]): Filter by collection
+        filters (Optional[Dict]): Additional filters
+        
+    Returns:
+        List[Dict]: List of matching documents
+    """
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            
+            # Base query
+            sql = """
+                SELECT DISTINCT
+                    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
+            """
+            
+            params = []
+            where_clauses = []
+            
+            # Add collection filter if specified
+            if collection_id:
+                where_clauses.append("dc.collection_id = ?")
+                params.append(collection_id)
+            
+            # Add date filters if specified
+            if filters and 'date_range' in filters:
+                start_date, end_date = filters['date_range']
+                where_clauses.append("d.upload_date BETWEEN ? AND ?")
+                params.extend([start_date, end_date])
+            
+            # Add text search
+            if query:
+                where_clauses.append("(d.name LIKE ? OR d.content LIKE ?)")
+                search_term = f"%{query}%"
+                params.extend([search_term, search_term])
+            
+            # Combine WHERE clauses
+            if where_clauses:
+                sql += " WHERE " + " AND ".join(where_clauses)
+            
+            sql += " GROUP BY d.id ORDER BY d.upload_date DESC"
+            
+            # Execute query
+            cursor.execute(sql, params)
+            
+            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 searching documents: {e}")
+        return []
 
 def get_all_documents(conn: sqlite3.Connection) -> List[Dict]:
     """