Update utils/database.py

utils/database.py

@@ -364,6 +364,357 @@ def add_query(conn: sqlite3.Connection, query: str, response: str, document_id:
         st.error(f"Error adding query: {e}")
         return False
 
+# Add to utils/database.py
+
+import sqlite3
+from typing import List, Dict, Optional
+from datetime import datetime
+from langchain_core.messages import HumanMessage, AIMessage
+import streamlit as st
+
+def create_chat_tables(conn: sqlite3.Connection) -> None:
+    """Create necessary tables for chat management."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            
+            # Create chats table
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS chats (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    title TEXT NOT NULL,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    collection_id INTEGER,
+                    FOREIGN KEY (collection_id) REFERENCES collections (id) ON DELETE SET NULL
+                )
+            ''')
+            
+            # Create chat messages table
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS chat_messages (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    chat_id INTEGER NOT NULL,
+                    role TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    metadata TEXT,  -- Store metadata as JSON string
+                    FOREIGN KEY (chat_id) REFERENCES chats (id) ON DELETE CASCADE
+                )
+            ''')
+            
+            conn.commit()
+            
+    except sqlite3.Error as e:
+        st.error(f"Error creating chat tables: {e}")
+        raise
+
+def create_new_chat(conn: sqlite3.Connection, title: str, collection_id: Optional[int] = None) -> Optional[int]:
+    """Create a new chat session."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                INSERT INTO chats (title, collection_id, created_at, last_updated)
+                VALUES (?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+            ''', (title, collection_id))
+            
+            conn.commit()
+            return cursor.lastrowid
+            
+    except sqlite3.Error as e:
+        st.error(f"Error creating new chat: {e}")
+        return None
+
+def save_chat_message(conn: sqlite3.Connection, 
+                     chat_id: int, 
+                     role: str, 
+                     content: str, 
+                     metadata: Optional[Dict] = None) -> Optional[int]:
+    """Save a chat message to the database."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            
+            # Convert metadata to JSON string if present
+            metadata_str = json.dumps(metadata) if metadata else None
+            
+            # Insert message
+            cursor.execute('''
+                INSERT INTO chat_messages (chat_id, role, content, metadata, timestamp)
+                VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
+            ''', (chat_id, role, content, metadata_str))
+            
+            # Update chat last_updated timestamp
+            cursor.execute('''
+                UPDATE chats 
+                SET last_updated = CURRENT_TIMESTAMP 
+                WHERE id = ?
+            ''', (chat_id,))
+            
+            conn.commit()
+            return cursor.lastrowid
+            
+    except sqlite3.Error as e:
+        st.error(f"Error saving chat message: {e}")
+        return None
+
+def get_all_chats(conn: sqlite3.Connection) -> List[Dict]:
+    """Retrieve all chat sessions."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT 
+                    c.id,
+                    c.title,
+                    c.created_at,
+                    c.last_updated,
+                    c.collection_id,
+                    COUNT(m.id) as message_count,
+                    MAX(m.timestamp) as last_message
+                FROM chats c
+                LEFT JOIN chat_messages m ON c.id = m.chat_id
+                GROUP BY c.id
+                ORDER BY c.last_updated DESC
+            ''')
+            
+            chats = []
+            for row in cursor.fetchall():
+                chats.append({
+                    'id': row[0],
+                    'title': row[1],
+                    'created_at': row[2],
+                    'last_updated': row[3],
+                    'collection_id': row[4],
+                    'message_count': row[5],
+                    'last_message': row[6]
+                })
+            return chats
+            
+    except sqlite3.Error as e:
+        st.error(f"Error retrieving chats: {e}")
+        return []
+
+def get_chat_messages(conn: sqlite3.Connection, chat_id: int) -> List[Dict]:
+    """Retrieve all messages for a specific chat."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT id, role, content, metadata, timestamp
+                FROM chat_messages
+                WHERE chat_id = ?
+                ORDER BY timestamp
+            ''', (chat_id,))
+            
+            messages = []
+            for row in cursor.fetchall():
+                # Parse metadata JSON if present
+                metadata = json.loads(row[3]) if row[3] else None
+                
+                # Convert to appropriate message type
+                if row[1] == 'human':
+                    message = HumanMessage(content=row[2])
+                else:
+                    message = AIMessage(content=row[2], additional_kwargs={'metadata': metadata})
+                
+                messages.append(message)
+            
+            return messages
+            
+    except sqlite3.Error as e:
+        st.error(f"Error retrieving chat messages: {e}")
+        return []
+
+def delete_chat(conn: sqlite3.Connection, chat_id: int) -> bool:
+    """Delete a chat session and all its messages."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            # Messages will be automatically deleted due to CASCADE
+            cursor.execute('DELETE FROM chats WHERE id = ?', (chat_id,))
+            conn.commit()
+            return True
+            
+    except sqlite3.Error as e:
+        st.error(f"Error deleting chat: {e}")
+        return False
+
+def update_chat_title(conn: sqlite3.Connection, chat_id: int, new_title: str) -> bool:
+    """Update the title of a chat session."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                UPDATE chats 
+                SET title = ?, last_updated = CURRENT_TIMESTAMP 
+                WHERE id = ?
+            ''', (new_title, chat_id))
+            conn.commit()
+            return True
+            
+    except sqlite3.Error as e:
+        st.error(f"Error updating chat title: {e}")
+        return False
+
+def get_chat_by_id(conn: sqlite3.Connection, chat_id: int) -> Optional[Dict]:
+    """Retrieve a specific chat session by ID."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT 
+                    c.id,
+                    c.title,
+                    c.created_at,
+                    c.last_updated,
+                    c.collection_id,
+                    COUNT(m.id) as message_count
+                FROM chats c
+                LEFT JOIN chat_messages m ON c.id = m.chat_id
+                WHERE c.id = ?
+                GROUP BY c.id
+            ''', (chat_id,))
+            
+            row = cursor.fetchone()
+            if row:
+                return {
+                    'id': row[0],
+                    'title': row[1],
+                    'created_at': row[2],
+                    'last_updated': row[3],
+                    'collection_id': row[4],
+                    'message_count': row[5]
+                }
+            return None
+            
+    except sqlite3.Error as e:
+        st.error(f"Error retrieving chat: {e}")
+        return None
+
+def export_chat_history(conn: sqlite3.Connection, chat_id: int) -> Optional[Dict]:
+    """Export a chat session with all its messages."""
+    try:
+        chat = get_chat_by_id(conn, chat_id)
+        if not chat:
+            return None
+            
+        messages = get_chat_messages(conn, chat_id)
+        
+        return {
+            'chat_info': chat,
+            'messages': [
+                {
+                    'role': 'human' if isinstance(msg, HumanMessage) else 'assistant',
+                    'content': msg.content,
+                    'metadata': msg.additional_kwargs.get('metadata') if isinstance(msg, AIMessage) else None
+                }
+                for msg in messages
+            ]
+        }
+            
+    except Exception as e:
+        st.error(f"Error exporting chat history: {e}")
+        return None
+
+def import_chat_history(conn: sqlite3.Connection, chat_data: Dict) -> Optional[int]:
+    """Import a chat session from exported data."""
+    try:
+        with conn_lock:
+            # Create new chat
+            chat_id = create_new_chat(
+                conn,
+                chat_data['chat_info']['title'],
+                chat_data['chat_info'].get('collection_id')
+            )
+            
+            if not chat_id:
+                return None
+                
+            # Import messages
+            for msg in chat_data['messages']:
+                save_chat_message(
+                    conn,
+                    chat_id,
+                    msg['role'],
+                    msg['content'],
+                    msg.get('metadata')
+                )
+                
+            return chat_id
+            
+    except Exception as e:
+        st.error(f"Error importing chat history: {e}")
+        return None
+
+
+# utils/database.py
+def create_chat_tables(conn):
+    """Create tables for chat management."""
+    cursor = conn.cursor()
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS chats (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            title TEXT NOT NULL,
+            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+        )
+    ''')
+    
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS chat_messages (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            chat_id INTEGER,
+            role TEXT NOT NULL,
+            content TEXT NOT NULL,
+            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+            FOREIGN KEY (chat_id) REFERENCES chats (id) ON DELETE CASCADE
+        )
+    ''')
+    conn.commit()
+
+def save_chat(conn, chat_title: str, messages: List[Dict]):
+    """Save chat history."""
+    cursor = conn.cursor()
+    cursor.execute('INSERT INTO chats (title) VALUES (?)', (chat_title,))
+    chat_id = cursor.lastrowid
+    
+    for msg in messages:
+        cursor.execute('''
+            INSERT INTO chat_messages (chat_id, role, content)
+            VALUES (?, ?, ?)
+        ''', (chat_id, msg['role'], msg['content']))
+    
+    conn.commit()
+    return chat_id
+
+# components/chat.py
+def display_chat_manager():
+    """Display chat management interface."""
+    st.sidebar.markdown("### Chat Management")
+    
+    # New chat button
+    if st.sidebar.button("New Chat"):
+        st.session_state.messages = []
+        st.session_state.current_chat_id = None
+    
+    # Save current chat
+    if st.session_state.messages and st.sidebar.button("Save Chat"):
+        chat_title = st.sidebar.text_input("Chat Title", 
+                                         value=f"Chat {datetime.now().strftime('%Y-%m-%d %H:%M')}")
+        if chat_title:
+            save_chat(st.session_state.db_conn, chat_title, st.session_state.messages)
+            st.sidebar.success("Chat saved!")
+    
+    # Load previous chats
+    chats = get_all_chats(st.session_state.db_conn)
+    if chats:
+        st.sidebar.markdown("### Previous Chats")
+        for chat in chats:
+            if st.sidebar.button(f"📜 {chat['title']}", key=f"chat_{chat['id']}"):
+                st.session_state.messages = get_chat_messages(st.session_state.db_conn, chat['id'])
+                st.session_state.current_chat_id = chat['id']
+                st.rerun()
 def add_annotation(conn: sqlite3.Connection, document_id: int, annotation: str, page_number: Optional[int] = None) -> bool:
     """Add an annotation to a document."""
     try: