utils/database.py

# utils/database.py

# Update the imports first
from langchain_community.chat_models import ChatOpenAI
from langchain_core.messages import (
    HumanMessage,
    AIMessage,
    SystemMessage,
    BaseMessage
)
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.agents import AgentExecutor, Tool, create_openai_tools_agent
from langchain.agents.format_scratchpad.tools import format_to_tool_messages
from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain.vectorstores import FAISS
import os
import streamlit as st
import sqlite3
import traceback
import time
import io
import tempfile
from sqlite3 import Error
from threading import Lock
from typing import Dict, List, Optional, Any
from datetime import datetime
from threading import Lock

# Create a lock for database connection
conn_lock = Lock()

def create_connection(db_file):
    """
    Create a database connection to the SQLite database.

    Args:
        db_file (str): Path to the SQLite database file.
    Returns:
        sqlite3.Connection: Database connection object or None if an error occurs.
    """
    conn = None
    try:
        conn = sqlite3.connect(db_file, check_same_thread=False)
        return conn
    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

# utils/database.py
# Add this version of create_tables (replacing the existing one)

# utils/database.py

def create_tables(conn: sqlite3.Connection) -> None:
    """Create all necessary tables in the database."""
    try:
        with conn_lock:
            cursor = conn.cursor()
            
            # Force create collections tables first
            collections_tables = [
                '''
                CREATE TABLE IF NOT EXISTS collections (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL UNIQUE,
                    description TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                ''',
                '''
                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
                )
                '''
            ]
            
            # Execute collections tables creation separately
            for table_sql in collections_tables:
                try:
                    cursor.execute(table_sql)
                    conn.commit()
                except sqlite3.Error as e:
                    st.error(f"Error creating collections table: {e}")
                    st.error(f"SQL that failed: {table_sql}")
                    raise
            
            # Create other tables
            other_tables = [
                '''
                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
                )
                ''',
                '''
                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
                )
                ''',
                '''
                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
                )
                '''
            ]
            
            # Execute other tables creation
            for table_sql in other_tables:
                try:
                    cursor.execute(table_sql)
                    conn.commit()
                except sqlite3.Error as e:
                    st.error(f"Error creating table: {e}")
                    st.error(f"SQL that failed: {table_sql}")
                    raise
            
            # Create indices
            indices = [
                'CREATE INDEX IF NOT EXISTS idx_doc_name ON documents(name)',
                'CREATE INDEX IF NOT EXISTS idx_collection_name ON collections(name)',
                'CREATE INDEX IF NOT EXISTS idx_doc_collections ON document_collections(collection_id)'
            ]
            
            # Execute indices creation
            for index_sql in indices:
                try:
                    cursor.execute(index_sql)
                    conn.commit()
                except sqlite3.Error as e:
                    st.error(f"Error creating index: {e}")
                    st.error(f"SQL that failed: {index_sql}")
            
            # Verify collections table was created
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='collections'")
            if not cursor.fetchone():
                st.error("Failed to create collections table despite no errors")
                raise Exception("Collections table creation failed silently")
            
            conn.commit()
            
    except sqlite3.Error as e:
        st.error(f"Error in create_tables: {e}")
        raise
    except Exception as e:
        st.error(f"Unexpected error in create_tables: {e}")
        raise

def force_recreate_collections_tables(conn: sqlite3.Connection) -> bool:
    """Force recreate collections tables if they're missing."""
    try:
        with conn_lock:
            cursor = conn.cursor()
            
            # Drop existing tables if they exist
            cursor.execute("DROP TABLE IF EXISTS document_collections")
            cursor.execute("DROP TABLE IF EXISTS collections")
            
            # Create collections table
            cursor.execute('''
                CREATE TABLE collections (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL UNIQUE,
                    description TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Create document_collections table
            cursor.execute('''
                CREATE TABLE 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
                )
            ''')
            
            # Create indices
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_collection_name ON collections(name)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_doc_collections ON document_collections(collection_id)')
            
            conn.commit()
            return True
            
    except sqlite3.Error as e:
        st.error(f"Error recreating collections tables: {e}")
        return False


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):
    """
    Create necessary tables in the database.

    Args:
        conn (sqlite3.Connection): SQLite database connection.
    """
    try:
        sql_create_documents_table = '''
        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
        );
        '''

        sql_create_queries_table = '''
        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)
        );
        '''

        sql_create_annotations_table = '''
        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)
        );
        '''

        conn.execute(sql_create_documents_table)
        conn.execute(sql_create_queries_table)
        conn.execute(sql_create_annotations_table)
    except Error as e:
        st.error(f"Error: {e}")


def get_documents(conn):
    """
    Retrieve all documents from the database.

    Args:
        conn (sqlite3.Connection): SQLite database connection.
    Returns:
        tuple: (list of document contents, list of document names).
    """
    try:
        with conn_lock:
            cursor = conn.cursor()
            cursor.execute("SELECT content, name FROM documents")
            results = cursor.fetchall()

        if not results:
            return [], []

        # Separate contents and names
        document_contents = [row[0] for row in results]
        document_names = [row[1] for row in results]

        return document_contents, document_names

    except Error as e:
        st.error(f"Error retrieving documents: {e}")
        return [], []


def insert_document(conn, name, content):
    """
    Insert a new document into the database.

    Args:
        conn (sqlite3.Connection): SQLite database connection.
        name (str): Name of the document.
        content (str): Content of the document.
    Returns:
        int: ID of the inserted document, or None if insertion failed.
    """
    try:
        with conn_lock:
            cursor = conn.cursor()
            sql = '''INSERT INTO documents (name, content)
                    VALUES (?, ?)'''
            cursor.execute(sql, (name, content))
            conn.commit()
            return cursor.lastrowid

    except Error as e:
        st.error(f"Error inserting document: {e}")
        return None


def verify_vector_store(vector_store):
    """
    Verify that the vector store has documents loaded.

    Args:
        vector_store (FAISS): FAISS vector store instance.
    Returns:
        bool: True if vector store is properly initialized with documents.
    """
    try:
        # Try to perform a simple similarity search
        test_results = vector_store.similarity_search("test", k=1)
        return len(test_results) > 0
    except Exception as e:
        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.
    
    Args:
        uploaded_files (list): List of uploaded files
        **kwargs: Additional arguments including:
            - collection_id (int, optional): ID of the collection to add documents to
    """
    try:
        # Initialize session state variables if they don't exist
        if 'qa_system' not in st.session_state:
            st.session_state.qa_system = None
        if 'vector_store' not in st.session_state:
            st.session_state.vector_store = None

        # Create progress containers
        progress_container = st.empty()
        status_container = st.empty()
        details_container = st.empty()

        # Initialize progress bar
        progress_bar = progress_container.progress(0)
        status_container.info("🔄 Initializing document processing...")

        # Reset existing states
        st.session_state.vector_store = None
        st.session_state.qa_system = None

        # Initialize embeddings (10% progress)
        status_container.info("🔄 Initializing embeddings model...")
        embeddings = get_embeddings_model()
        if not embeddings:
            status_container.error("❌ Failed to initialize embeddings model")
            return False
        progress_bar.progress(10)

        # Process documents
        all_chunks = []
        documents = []
        document_names = []

        progress_per_file = 70 / len(uploaded_files)
        current_progress = 10

        collection_id = kwargs.get('collection_id')

        for idx, uploaded_file in enumerate(uploaded_files):
            file_name = uploaded_file.name
            status_container.info(f"🔄 Processing document {idx + 1}/{len(uploaded_files)}: {file_name}")

            # Create temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
                tmp_file.write(uploaded_file.getvalue())
                tmp_file.flush()

                # Process document with chunking
                chunks, content = process_document(tmp_file.name)

                # Store in database
                doc_id = insert_document(st.session_state.db_conn, file_name, content)
                if not doc_id:
                    status_container.error(f"❌ Failed to store document: {file_name}")
                    continue

                # Add to collection if specified
                if collection_id:
                    if not add_document_to_collection(st.session_state.db_conn, doc_id, collection_id):
                        status_container.warning(f"⚠️ Failed to add document to collection: {file_name}")

                # Add chunks with metadata
                for chunk in chunks:
                    chunk.metadata.update({
                        "source": file_name,
                        "document_id": doc_id,
                        "collection_id": collection_id if collection_id else None
                    })
                all_chunks.extend(chunks)

                documents.append(content)
                document_names.append(file_name)

            current_progress += progress_per_file
            progress_bar.progress(int(current_progress))

        # Initialize vector store with chunks
        status_container.info("🔄 Initializing vector store...")
        vector_store = FAISS.from_documents(
            all_chunks,
            embeddings
        )

        # Verify vector store
        status_container.info("🔄 Verifying document indexing...")
        details_container.text("✨ Performing final checks...")
        if not verify_vector_store(vector_store):
            status_container.error("❌ Vector store verification failed")
            return False

        # Initialize QA system (90-100% progress)
        status_container.info("🔄 Setting up QA system...")
        qa_system = initialize_qa_system(vector_store)
        if not qa_system:
            status_container.error("❌ Failed to initialize QA system")
            return False

        # Store in session state
        if collection_id:
            if 'vector_stores' not in st.session_state:
                st.session_state.vector_stores = {}
            st.session_state.vector_stores[collection_id] = vector_store
        else:
            st.session_state.vector_store = vector_store
        
        st.session_state.qa_system = qa_system

        # Complete!
        progress_bar.progress(100)
        status_container.success("✅ Documents processed successfully!")
        details_container.markdown(
            f"""
            🎉 **Ready to chat!**
            - Documents processed: {len(documents)}
            - Total content size: {sum(len(doc) for doc in documents) / 1024:.2f} KB
            - {"Added to collection" if collection_id else "Processed as standalone documents"}
            
            You can now start asking questions about your documents!
            """
        )

        # Add notification
        st.balloons()

        # Clean up progress display after 3 seconds
        time.sleep(3)
        progress_container.empty()
        status_container.empty()
        details_container.empty()

        return True

    except Exception as e:
        st.error(f"❌ Error processing documents: {str(e)}")
        if status_container:
            status_container.error(traceback.format_exc())
        # Reset states on error
        st.session_state.vector_store = None
        st.session_state.qa_system = None
        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]:
    """
    Get all documents with their metadata and collection info.
    
    Args:
        conn (sqlite3.Connection): Database connection
        
    Returns:
        List[Dict]: List of documents with their metadata
    """
    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_by_id(conn: sqlite3.Connection, document_id: int) -> Optional[Dict]:
    """
    Get a single document by its ID.
    
    Args:
        conn (sqlite3.Connection): Database connection
        document_id (int): ID of the document to retrieve
        
    Returns:
        Optional[Dict]: Document data if found, None otherwise
    """
    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
                WHERE d.id = ?
                GROUP BY d.id
            ''', (document_id,))
            
            row = cursor.fetchone()
            if row:
                return {
                    'id': row[0],
                    'name': row[1],
                    'content': row[2],
                    'upload_date': row[3],
                    'collections': row[4].split(',') if row[4] else []
                }
            return None
            
    except sqlite3.Error as e:
        st.error(f"Error retrieving document: {e}")
        return None

def get_recent_documents(conn: sqlite3.Connection, limit: int = 5) -> List[Dict]:
    """
    Get most recently uploaded documents.
    
    Args:
        conn (sqlite3.Connection): Database connection
        limit (int): Maximum number of documents to return
        
    Returns:
        List[Dict]: List of recent documents
    """
    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
                LIMIT ?
            ''', (limit,))
            
            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 recent documents: {e}")
        return []

def get_collections(conn: sqlite3.Connection) -> List[Dict]:
    """
    Get all collections with their document counts.
    
    Args:
        conn (sqlite3.Connection): Database connection
        
    Returns:
        List[Dict]: List of collections with metadata
    """
    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 get_collection_documents(conn: sqlite3.Connection, collection_id: int) -> List[Dict]:
    """
    Get all documents in a specific collection.
    
    Args:
        conn (sqlite3.Connection): Database connection
        collection_id (int): ID of the collection
        
    Returns:
        List[Dict]: List of documents in the 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 create_collection(conn: sqlite3.Connection, name: str, description: str = "") -> Optional[int]:
    """
    Create a new collection.
    
    Args:
        conn (sqlite3.Connection): Database connection
        name (str): Name of the collection
        description (str): Optional description
        
    Returns:
        Optional[int]: ID of the created collection if successful
    """
    try:
        with conn_lock:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO collections (name, description)
                VALUES (?, ?)
            ''', (name, description))
            conn.commit()
            return cursor.lastrowid
            
    except sqlite3.Error as e:
        st.error(f"Error creating collection: {e}")
        return None

def add_document_to_collection(conn: sqlite3.Connection, document_id: int, collection_id: int) -> bool:
    """
    Add a document to a collection.
    
    Args:
        conn (sqlite3.Connection): Database connection
        document_id (int): ID of the document
        collection_id (int): ID of the collection
        
    Returns:
        bool: True if successful
    """
    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 adding document to collection: {e}")
        return False

def process_document(file_path):
    """
    Process a PDF document with proper chunking.
    
    Args:
        file_path (str): Path to the PDF file
        
    Returns:
        tuple: (list of document chunks, full content of the document)
    """
    # Load PDF
    loader = PyPDFLoader(file_path)
    documents = loader.load()

    # Create text splitter
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
        separators=["\n\n", "\n", " ", ""]
    )

    # Split documents into chunks
    chunks = text_splitter.split_documents(documents)

    # Extract full content for database storage
    full_content = "\n".join(doc.page_content for doc in documents)

    return chunks, full_content


def display_vector_store_info():
    """
    Display information about the current vector store state.
    """
    if 'vector_store' not in st.session_state:
        st.info("ℹ️ No documents loaded yet.")
        return

    try:
        # Get the vector store from session state
        vector_store = st.session_state.vector_store

        # Get basic stats
        test_query = vector_store.similarity_search("test", k=1)
        doc_count = len(test_query)

        # Create an expander for detailed info
        with st.expander("📊 Knowledge Base Status"):
            col1, col2 = st.columns(2)

            with col1:
                st.metric(
                    label="Documents Loaded",
                    value=doc_count
                )

            with col2:
                st.metric(
                    label="System Status",
                    value="Ready" if verify_vector_store(vector_store) else "Not Ready"
                )

            # Display sample queries
            if verify_vector_store(vector_store):
                st.markdown("### 🔍 Sample Document Snippets")
                sample_docs = vector_store.similarity_search("", k=3)
                for i, doc in enumerate(sample_docs, 1):
                    with st.container():
                        st.markdown(f"**Snippet {i}:**")
                        st.text(doc.page_content[:200] + "...")

    except Exception as e:
        st.error(f"Error displaying vector store info: {e}")
        st.error(traceback.format_exc())


def initialize_qa_system(vector_store):
    """
    Initialize QA system with optimized retrieval.

    Args:
        vector_store (FAISS): FAISS vector store instance.
    Returns:
        dict: QA system chain or None if initialization fails.
    """
    try:
        llm = ChatOpenAI(
            temperature=0.5,
            model_name="gpt-4",
            max_tokens=4000,  # Explicitly set max tokens
            api_key=os.environ.get("OPENAI_API_KEY")
        )

        # Optimize retriever settings
        retriever = vector_store.as_retriever(
            search_kwargs={
                "k": 3,  # Retrieve fewer, more relevant chunks
                "fetch_k": 5  # Consider more candidates before selecting top k
            }
        )

        # Create a template that enforces clean formatting
        prompt = ChatPromptTemplate.from_messages([
            ("system", """
You are an expert consultant specializing in analyzing Request for Proposal (RFP) documents. Your goal is to assist users by providing clear, concise, and professional insights based on the content provided. Please adhere to the following guidelines when crafting your responses:

Begin with a summary that highlights the key findings or answers the main query.

Structured Format: Use clear and descriptive section headers to organize the information logically.

Bullet Points: Utilize bullet points for lists or complex information to enhance readability.

Source Attribution: Cite specific sections or page numbers from the RFP document when referencing information.

Professional Formatting: Maintain a clean and professional layout using Markdown formatting, such as headings, bullet points, bold, italics, and tables where appropriate.

Use Markdown Syntax: Ensure the response is fully formatted using Markdown for optimal readability in the chat.

Focused Content: Keep your responses concise and directly related to the user's query, avoiding unnecessary information.

Scope Awareness: If a query falls outside the provided information or context, politely acknowledge this and suggest consulting the relevant sections or additional sources.

Confidentiality: Respect the confidentiality of the information provided and avoid sharing any sensitive data beyond the scope of the query.

Tone and Language: Use formal and professional language, ensuring clarity and precision in your responses.

Accuracy: Double-check all information for accuracy and completeness before providing it to the user.
"""),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}\n\nContext: {context}")
        ])

        def get_chat_history(inputs):
            chat_history = inputs.get("chat_history", [])
            if not isinstance(chat_history, list):
                return []
            return [msg for msg in chat_history if isinstance(msg, BaseMessage)]

        def get_context(inputs):
            docs = retriever.get_relevant_documents(inputs["input"])
            context_parts = []
            for doc in docs:
                source = doc.metadata.get('source', 'Unknown source')
                context_parts.append(f"\nFrom {source}:\n{doc.page_content}")
            return "\n".join(context_parts)

        chain = (
            {
                "context": get_context,
                "chat_history": get_chat_history,
                "input": lambda x: x["input"]
            }
            | prompt
            | llm
        )

        return chain

    except Exception as e:
        st.error(f"Error initializing QA system: {e}")
        return None


# FAISS vector store initialization
def initialize_faiss(embeddings, documents, document_names):
    """
    Initialize FAISS vector store.

    Args:
        embeddings (Embeddings): Embeddings model to use.
        documents (list): List of document contents.
        document_names (list): List of document names.
    Returns:
        FAISS: FAISS vector store instance or None if initialization fails.
    """
    try:
        from langchain.vectorstores import FAISS

        vector_store = FAISS.from_texts(
            documents,
            embeddings,
            metadatas=[{"source": name} for name in document_names],
        )
        return vector_store
    except Exception as e:
        st.error(f"Error initializing FAISS: {e}")
        return None


# Embeddings model retrieval
@st.cache_resource
def get_embeddings_model():
    """
    Get the embeddings model.

    Returns:
        Embeddings: Embeddings model instance or None if loading fails.
    """
    try:
        from langchain.embeddings import HuggingFaceEmbeddings

        model_name = "sentence-transformers/all-MiniLM-L6-v2"
        embeddings = HuggingFaceEmbeddings(model_name=model_name)
        return embeddings
    except Exception as e:
        st.error(f"Error loading embeddings model: {e}")
        return None