app_streamlit.py

import os
import streamlit as st
import asyncio
import tempfile
from pathlib import Path
import sys
import requests
import logging
import traceback
import time
from threading import Thread
from queue import Queue
import sqlite3
import hashlib
import json
import shutil
from datetime import datetime
import chromadb

# Configure logging
logging.basicConfig(level=logging.INFO, 
                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Add the RAG-Agent source code to the Python path
rag_agent_path = os.path.join(os.path.dirname(__file__), "rag-agent/src")
sys.path.append(rag_agent_path)
try:
    # Import Ollama with error handling
    import ollama
    logger.info("Successfully imported ollama package")
except ImportError as e:
    logger.error(f"Failed to import ollama package: {str(e)}")
    st.error(f"Failed to import ollama package: {str(e)}")
    sys.exit(1)

try:
    # Import local modules from the rag-agent source
    from ollama_pure_chat import OllamaChat
    from fileloader import Loader_Local
    logger.info("Successfully imported modules from rag-agent")
except ImportError as e:
    logger.error(f"Failed to import modules from rag-agent: {str(e)}")
    st.error(f"Failed to import modules from rag-agent: {str(e)}")
    sys.exit(1)

    # Mark imports as completed
    st.session_state.imports_done = True
else:
    # Silent imports when already done
    import ollama
    from ollama_pure_chat import OllamaChat
    from fileloader import Loader_Local

# Check if Ollama API is available
def is_ollama_available():
    """Check if the Ollama server is running and accessible"""
    # Only check once per session unless explicitly refreshed
    if "ollama_available" not in st.session_state or "ollama_checked_time" not in st.session_state:
        try:
            response = requests.get("http://localhost:11434/api/tags", timeout=2)
            status = response.status_code == 200
            st.session_state.ollama_available = status
            st.session_state.ollama_checked_time = time.time()
            return status
        except (requests.ConnectionError, requests.Timeout):
            logger.warning("Ollama server is not running or not accessible")
            st.session_state.ollama_available = False
            st.session_state.ollama_checked_time = time.time()
            return False
        except Exception as e:
            logger.error(f"Unexpected error checking Ollama availability: {str(e)}")
            st.session_state.ollama_available = False
            st.session_state.ollama_checked_time = time.time()
            return False
    else:
        # Check if we need to refresh (every 60 seconds)
        if time.time() - st.session_state.ollama_checked_time > 60:
            try:
                response = requests.get("http://localhost:11434/api/tags", timeout=2)
                status = response.status_code == 200
                st.session_state.ollama_available = status
                st.session_state.ollama_checked_time = time.time()
                return status
            except Exception:
                # Just keep existing status on error
                return st.session_state.ollama_available
        return st.session_state.ollama_available

# Default model to use if none is specified
default_model = "qwen2.5:latest"

# Persistent storage paths
PERSISTENT_STORAGE_BASE = os.getenv("PERSISTENT_STORAGE_BASE",'localdatabase')
USERS_DB_PATH = os.path.join(PERSISTENT_STORAGE_BASE, "users.db")


# Ensure persistent storage directories exist
os.makedirs(PERSISTENT_STORAGE_BASE, exist_ok=True)

# Function to get user-specific persistent storage path
def get_user_storage_path(username):
    user_path = os.path.join(PERSISTENT_STORAGE_BASE, username)
    os.makedirs(user_path, exist_ok=True)
    return user_path

# Modify init_user_database to use persistent path
def init_user_database():
    """Initialize SQLite database for user management"""
    conn = sqlite3.connect(USERS_DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            username TEXT PRIMARY KEY,
            password TEXT NOT NULL,
            databases TEXT DEFAULT '["default"]',
            last_login DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

# Modify get_available_databases to use persistent storage
def get_available_databases():
    """Get available databases for the current user from ChromaDB"""
    if not hasattr(st.session_state, 'username'):
        return []
    
    try:
        # Get user-specific storage path
        user_storage_path = get_user_storage_path(st.session_state.username)
        
        # List all directories in the user's storage path
        # Each directory represents a database
        if os.path.exists(user_storage_path):
            databases = [d for d in os.listdir(user_storage_path) 
                        if os.path.isdir(os.path.join(user_storage_path, d))]
            
            # If no databases found, return empty list
            if not databases:
                return []
                
            return databases
        else:
            return []
    except Exception as e:
        logger.error(f"Error retrieving databases from ChromaDB: {str(e)}")
        return []

# Modify get_context to improve retrieval
def get_context(question):
    """Improved context retrieval with better error handling and persistent storage support"""
    if not init_loader():
        st.error("Document loader could not be initialized.")
        return "", []
    
    context = ""
    context_chunks = []
    
    try:
        # Get the current database
        current_db = st.session_state.current_database
        
        # Get user-specific storage path
        user_storage_path = get_user_storage_path(st.session_state.username)
        
        # Set persistent directory for the loader
        st.session_state.loader.persist_directory = os.path.join(user_storage_path, current_db)
        
        # Get the top_k value from session state (default to 5 if not set)
        top_k = st.session_state.get("retrieval_top_k", 5)
        
        # Show debug info if enabled
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            logger.info(f"Getting context for question: {question}")
            logger.info(f"Using database: {current_db}")
            logger.info(f"Persist directory: {os.path.join(user_storage_path, current_db)}")
            logger.info(f"Using top_k: {top_k}")
        
        # First, check if there are temporary collections
        temp_collections_exist = hasattr(st.session_state, 'temp_collections') and st.session_state.temp_collections
        
        # Check if we have temporary collections
        if temp_collections_exist:
            st.session_state.topk=1000
            if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                logger.info(f"Found {len(st.session_state.temp_collections)} temporary collections")
            
            # Get documents from all temporary collections
            for temp_coll_name, file_info in st.session_state.temp_collections.items():
                try:
                    display_name = file_info.get('file_name', temp_coll_name)
                    
                    if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                        logger.info(f"Retrieving from temporary collection: {temp_coll_name}")
                    
                    # Get ALL documents from temporary collections regardless of relevance
                    # This ensures we prioritize temporary files even if relevance is low
                    try:
                        # First try to get relevant documents using the query
                        results = st.session_state.loader(
                            question,
                            temp_coll_name,
                            top_k=top_k
                        )
                        
                        # If no documents found with query, get all documents from the collection
                        if not results or not results.get('documents') or not results['documents'][0]:
                            if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                                logger.info(f"No relevant documents found in {temp_coll_name}, retrieving all documents")
                            
                            # Get the ChromaDB client directly
                            chroma_client = chromadb.PersistentClient(path=st.session_state.loader.persist_directory)
                            try:
                                # Get the collection
                                collection = chroma_client.get_collection(name=temp_coll_name)
                                
                                # Get all documents from the collection (up to 10)
                                all_docs = collection.get(limit=10)
                                
                                if all_docs and 'documents' in all_docs and all_docs['documents']:
                                    # Create a results structure similar to the loader's output
                                    results = {
                                        'ids': [all_docs['ids']],
                                        'documents': [all_docs['documents']],
                                        'metadatas': [all_docs['metadatas'] if 'metadatas' in all_docs else [{}] * len(all_docs['documents'])],
                                        'distances': [[1.0] * len(all_docs['documents'])]  # Placeholder distances
                                    }
                                    
                                    if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                                        logger.info(f"Retrieved {len(all_docs['documents'])} documents directly from collection")
                            except Exception as coll_error:
                                logger.error(f"Error accessing collection directly: {str(coll_error)}")
                    
                    except Exception as query_error:
                        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                            logger.error(f"Error querying temporary collection {temp_coll_name}: {str(query_error)}")
                        continue
                    
                    # Process retrieved documents if any
                    if results and 'documents' in results and results['documents'] and results['documents'][0]:
                        # Add documents to the context for the model
                        for i, doc in enumerate(results['documents'][0]):
                            # Truncate very long documents to prevent context overflow
                            #truncated_doc = doc[:1000] + '...' if len(doc) > 1000 else doc
                            
                            context += f"doc_{i} from {display_name} (Temporary File): {doc}\n"
                            context_chunks.append(f"**Document Chunk {i+1} from {display_name} (Temporary File):**\n{doc}")
                    
                except Exception as collection_error:
                    logger.error(f"Error retrieving from temporary collection {temp_coll_name}: {str(collection_error)}")
                    if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                        logger.error(f"Traceback: {traceback.format_exc()}")
            
            # If we found context from temporary collections, don't search database collections
            if context_chunks:
                if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                    logger.info(f"Found {len(context_chunks)} chunks from temporary collections. Skipping database search.")
                
                # Log context retrieval details
                logger.info(f"Retrieved {len(context_chunks)} context chunks from temporary files for query: {question}")
                return context, context_chunks
            
            # If we reached here, no context was found in temporary collections
            if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                logger.info("No context found in temporary collections despite forced retrieval. Searching database collections.")
        
        # If we have no temporary collections or found no context in them, search database collections
        database_collections = chromadb.PersistentClient(path=st.session_state.loader.persist_directory).list_collections()
        collections_to_search = []
        
        # Check if a specific collection is selected for querying

        if hasattr(st.session_state, 'active_collection') and st.session_state.active_collection:
            # Only search the selected collection
            selected_collection = st.session_state.active_collection
            if selected_collection in database_collections:
                collections_to_search.append({
                    "id": selected_collection,  # Use the collection name directly as the ID
                    "name": selected_collection,
                    "is_temporary": False
                })
                if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                    logger.info(f"Searching only the selected collection: {selected_collection}")
        else:
            # Add all database collections to search
            for collection_name in database_collections:
                collections_to_search.append({
                    "id": collection_name,  # Use the collection name directly as the ID
                    "name": collection_name,
                    "is_temporary": False
                })
        
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            logger.info(f"Searching {len(collections_to_search)} database collections")
        
        # If there are no collections at all, return empty context
        if not collections_to_search:
            st.info("No collections available to search.")
            return "", []
        
        # Search in database collections
        for collection in collections_to_search:
            try:
                collection_id = collection["id"]
                display_name = collection["name"]
                
                # Try to retrieve context
                if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                    logger.info(f"Querying database collection: {collection_id}")
                
                # Get collection results, handling errors gracefully
                try:
                    results = st.session_state.loader(
                        question, 
                        collection_id, 
                        top_k=top_k  # Use the configurable top_k value
                    )
                except Exception as query_error:
                    if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                        logger.error(f"Error querying collection {collection_id}: {str(query_error)}")
                    continue  # Skip this collection and try the next one
                
                # Check if results are valid and contain documents
                if results and 'documents' in results and results['documents'] and results['documents'][0]:
                    # Add documents to the context for the model
                    for i, doc in enumerate(results['documents']):
                        # Truncate very long documents to prevent context overflow
                        #truncated_doc = doc[:1000] + '...' if len(doc) > 1000 else doc
                        
                        context += f"doc_{i} from {display_name} (Database): {doc}\n"
                        context_chunks.append(f"**Document Chunk {i+1} from {display_name} (Database):**\n{doc}")
                
            except Exception as collection_error:
                logger.error(f"Error retrieving from database collection {collection['name']}: {str(collection_error)}")
                if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                    logger.error(f"Traceback: {traceback.format_exc()}")
        
        # Log context retrieval detailsQA
        logger.info(f"Retrieved {len(context_chunks)} context chunks for query: {question}")
        
        # If no context found, return empty
        if not context_chunks:
            st.info("No relevant context found for the query.")
    
    except Exception as e:
        logger.error(f"Comprehensive error getting context: {str(e)}")
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            logger.error(f"Traceback: {traceback.format_exc()}")
        st.error(f"Error retrieving context: {str(e)}")
    
    return context, context_chunks

# Modify show_create_database_modal to use persistent storage
def show_create_database_modal():
    """Create a new database for the current user"""
    if not hasattr(st.session_state, 'username'):
        st.error("Please log in first")
        return
    
    with st.form(key="create_database_form"):
        st.subheader("Create New Database")
        new_db_name = st.text_input("Database Name")
        submit_button = st.form_submit_button(label="Create")
        
        if submit_button and new_db_name:
            try:
                # Check if database already exists in session state
                current_databases = st.session_state.available_databases if hasattr(st.session_state, 'available_databases') else []
                
                # Check if database already exists
                if new_db_name in current_databases:
                    st.error(f"Database '{new_db_name}' already exists")
                else:
                    # Add new database to available databases list
                    current_databases.append(new_db_name)
                    st.session_state.available_databases = current_databases

                    # Create a directory for the new ChromaDB database in user's persistent storage
                    user_storage_path = get_user_storage_path(st.session_state.username)
                    database_path = os.path.join(user_storage_path, new_db_name)
                    os.makedirs(database_path, exist_ok=True)
                    
                    # Update session state
                    st.session_state.current_database = new_db_name
                    st.success(f"ChromaDB database '{new_db_name}' created for file loading!")
                    
                    # Initialize collections for this ChromaDB database
                    if "database_collections" not in st.session_state:
                        st.session_state.database_collections = {}
                    st.session_state.database_collections[new_db_name] = {}
                    
                    # Refresh the page to update database list
                    st.rerun()
            
            except Exception as e:
                        st.error(f"Error creating ChromaDB database: {str(e)}")

# Modify login_user to update last login and load persistent collections
def login_user(username, password):
    """Authenticate user and load persistent collections"""
    if len(username) < 3:
        return False, "Username must be at least 3 characters long"
    
    # Hash the password
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    
    try:
        conn = sqlite3.connect(USERS_DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute("SELECT * FROM users WHERE username = ? AND password = ?", 
                       (username, hashed_password))
        user = cursor.fetchone()
        
        if user:
            # Update last login timestamp
            cursor.execute(
                "UPDATE users SET last_login = CURRENT_TIMESTAMP WHERE username = ?", 
                (username,)
            )
            conn.commit()
            
            # Load persistent collections for the user
            user_storage_path = get_user_storage_path(username)
            
            # Initialize database collections from persistent storage
            database_collections = {}
            
            # Get user's databases
            cursor.execute("SELECT databases FROM users WHERE username = ?", (username,))
            result = cursor.fetchone()
            
            if result and result[0]:
                databases = json.loads(result[0])
                
                # Load collections for each database
                for db_name in databases:
                    db_path = os.path.join(user_storage_path, db_name)
                    if os.path.exists(db_path):
                        # Here you might want to add logic to load collection metadata
                        database_collections[db_name] = {}
            
            # Load temporary collections
            temp_collections = {}
            file_collections = {}
            
            # Load any persisted temporary collections from ChromaDB
            if init_loader():
                
                
                # Check each database directory for temporary collections
                for db_name in ['default']:
                    db_path = os.path.join(user_storage_path, db_name)
                    if os.path.exists(db_path):
                        # Set up the loader to check this database
                        
                        try:
                            # Get all collection names from ChromaDB
                            chroma_client = chromadb.PersistentClient(path=db_path)
                            all_collections = chroma_client.list_collections()
                            
                            # Find temporary collections for this user
                            for collection in all_collections:
                                collection_name = collection
                                
                                # Only load temp collections for this user
                                if collection_name.startswith(f"temp_{username}_"):
                                    # Extract file ID from collection name
                                    file_id = collection_name.replace(f"temp_{username}_", "")
                                    
                                    # Try to get collection metadata
                                    try:
                                        collection_obj = chroma_client.get_collection(name=collection_name)
                                        # Get any file in this collection to extract metadata
                                        results = collection_obj.get(limit=1)
                                        
                                        if results and results['metadatas'] and results['metadatas'][0]:
                                            source = results['metadatas'][0].get('source', '')
                                            file_name = os.path.basename(source)
                                            file_extension = Path(file_name).suffix.lstrip(".").lower()
                                            
                                            # Store in temp_collections
                                            temp_collections[collection_name] = {
                                                "file_name": file_name,
                                                "file_type": file_extension,
                                                "added_date": "Previously uploaded",
                                                "is_temporary": True
                                            }
                                            
                                            # Add to file_collections for retrieval
                                            file_collections[file_id] = collection_name
                                    except Exception as e:
                                        logger.error(f"Error loading temporary collection metadata: {str(e)}")
                                else:
                                    # Process regular collections (not temporary)
                                    try:
                                        collection_obj = chroma_client.get_collection(name=collection_name)
                                        # Get any file in this collection to extract metadata
                                        results = collection_obj.get(limit=1)
                                        
                                        if results and results['metadatas'] and results['metadatas'][0]:
                                            source = results['metadatas'][0].get('source', '')
                                            file_name = os.path.basename(source)
                                            file_extension = Path(file_name).suffix.lstrip(".").lower()
                                            
                                            # Get current database name
                                            current_db = db_name
                                            
                                            # Initialize database collection if not exists
                                            if current_db not in database_collections:
                                                database_collections[current_db] = {}
                                            
                                            # Store in database_collections
                                            if collection_name not in database_collections[current_db]:
                                                database_collections[current_db][collection_name] = {
                                                    "file_name": file_name,
                                                    "file_type": file_extension,
                                                    "added_date": "Previously uploaded",
                                                    "is_temporary": False,
                                                    "files": [file_name]
                                                }
                                            else:
                                                # Add file to existing collection if not already there
                                                if file_name not in database_collections[current_db][collection_name]["files"]:
                                                    database_collections[current_db][collection_name]["files"].append(file_name)
                                    except Exception as e:
                                        logger.error(f"Error loading regular collection metadata: {str(e)}")
                        except Exception as e:
                                logger.error(f"Error listing collections in database {db_name}: {str(e)}")
                                
                                conn.close()
            
            return True, "Login successful", database_collections, temp_collections, file_collections
        else:
            conn.close()
            return False, "Invalid username or password", {}, {}, {}
    except Exception as e:
        return False, f"Login error: {str(e)}", {}, {}, {}

# Initialize the loader only once
def init_loader():
    if not hasattr(st.session_state, 'loader') or st.session_state.loader is None:
        try:
            st.session_state.loader = Loader_Local()
            logger.info("Successfully initialized Loader_Local")
            return True
        except Exception as e:
            logger.error(f"Failed to initialize Loader_Local: {str(e)}")
            return False
    return True

# Modify main to handle persistent collections
def main():
    # Set page config
    st.set_page_config(
        page_title="RAG-Agent Chat",
        page_icon="🤖",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Initialize session state variables
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    if "initialized" not in st.session_state:
        st.session_state.initialized = False
    
    if "selected_model" not in st.session_state:
        st.session_state.selected_model = default_model
    
    if "enable_retrieval" not in st.session_state:
        st.session_state.enable_retrieval = True
    
    if "file_collections" not in st.session_state:
        st.session_state.file_collections = {}
        
    if "retrieval_top_k" not in st.session_state:
        st.session_state.retrieval_top_k = 5
    
    # Check Ollama status only once per session load
    # Use session variable instead of calling API multiple times
    if "ollama_available" not in st.session_state:
        # Only make the API call once
        ollama_status = is_ollama_available()
    else:
        ollama_status = st.session_state.ollama_available
        
    # Initialize the document loader
    loader_success = init_loader()
    if not loader_success:
        logger.warning("Failed to initialize document loader. Some features may not be available.")
    
    # Initialize the database
    init_user_database()
    
    # Check if the user is authenticated
    if "authenticated" not in st.session_state or not st.session_state.authenticated:
        show_login_page()
        return
    
    # Ensure chat_instance exists and is initialized at the very beginning
    if "chat_instance" not in st.session_state or st.session_state.chat_instance is None:
        if ollama_status:
            try:
                # Create a new chat instance
                model_name = st.session_state.selected_model
                logger.info(f"Initializing chat instance with model: {model_name}")
                st.session_state.chat_instance = OllamaChat(
                    model=model_name,
                    system_prompt=st.session_state.system_prompt
                )
                st.session_state.initialized = True
                logger.info(f"Successfully initialized chat instance with model: {model_name}")
            except Exception as e:
                logger.error(f"Error initializing chat instance: {str(e)}\n{traceback.format_exc()}")
                st.session_state.chat_instance = None
                st.session_state.initialized = False
                # We'll show an error to the user in the UI later
        else:
            logger.warning("Ollama server is not available. Chat functionality will be limited.")
            st.session_state.chat_instance = None
            st.session_state.initialized = False
    
    # Set current database to first available database or default
    available_dbs = get_available_databases()

    if available_dbs:
        if "current_database" not in st.session_state:
            st.session_state.current_database = available_dbs[0]
    
    # Load persistent collections
    if "database_collections" not in st.session_state:
        user_storage_path = get_user_storage_path(st.session_state.username)
        database_collections = {}
        
        for db_name in available_dbs:
            db_path = os.path.join(user_storage_path, db_name)
            if os.path.exists(db_path):
                # Load collections from files.db
                files_db_path = os.path.join(db_path, "files.db")
                if os.path.exists(files_db_path):
                    try:
                        files_conn = sqlite3.connect(files_db_path)
                        files_cursor = files_conn.cursor()
                        files_cursor.execute("SELECT * FROM files")
                        files = files_cursor.fetchall()
                        
                        database_collections[db_name] = {}
                        for file in files:
                            file_id = file[0]
                            file_name = file[1]
                            file_type = file[2]
                            collection_name = file[3]
                            added_date = file[4]
                            description = file[5]
                            
                            database_collections[db_name][collection_name] = {
                                "collection_name": collection_name,
                                "file_name": file_name,
                                "file_type": file_type,
                                "added_date": added_date,
                                "description": description or ""
                            }
                        
                        files_conn.close()
                    except Exception as e:
                        logger.error(f"Error loading files from database {db_name}: {str(e)}")
                        database_collections[db_name] = {}
                else:
                    database_collections[db_name] = {}
        
        st.session_state.database_collections = database_collections
    
    # Sidebar Configuration
    with st.sidebar:
        st.title("RAG-Agent Chat")
        
        # User info
        st.markdown(f"**Logged in as:** {st.session_state.username}")
        if st.button("Logout"):
            st.session_state.authenticated = False
            st.rerun()
            
        st.divider()
        
        # Add debug mode toggle
        if "debug_mode" not in st.session_state:
            st.session_state.debug_mode = False
            
        debug_mode = st.toggle("Debug Mode", value=st.session_state.debug_mode, key="debug_toggle")
        if debug_mode != st.session_state.debug_mode:
            st.session_state.debug_mode = debug_mode
            if debug_mode:
                st.info("Debug mode enabled. Additional logging information will be shown.")
                # Set logging level to DEBUG when debug mode is on
                logger.setLevel(logging.DEBUG)
            else:
                # Reset to INFO level when debug mode is off
                logger.setLevel(logging.INFO)
        
        # Check if Ollama is running - show clear status
        if not ollama_status:
            st.error("⚠️ Ollama server is not running")
            st.info("Please start Ollama by running 'ollama serve' in a terminal")
            
            # Add manual check button
            if st.button("Check Ollama Status"):
                # Force refresh the Ollama status
                if "ollama_available" in st.session_state:
                    del st.session_state.ollama_available
                if "ollama_checked_time" in st.session_state:
                    del st.session_state.ollama_checked_time
                if "cached_models" in st.session_state:
                    del st.session_state.cached_models
                st.rerun()
        else:
            st.success("✅ LLM  is running")
        # System Prompt Configuration
        st.markdown("### System Prompt")
        
        # Default system prompt if not set
        if "system_prompt" not in st.session_state:
            st.session_state.system_prompt = "You are a helpful chatbot assistant designed to answer questions about the given context. Always keep it short and precise. ."
        
        # Text area for system prompt
        system_prompt = st.text_area(
            "Customize system prompt:",
            value=st.session_state.system_prompt,
            height=100,
            help="Customize the system prompt that guides the model's behavior. The <context> tag will be replaced with retrieved documents."
        )
        
        # Update system prompt if changed
        if system_prompt != st.session_state.system_prompt:
            st.session_state.system_prompt = system_prompt
            
            # Update chat instance with new system prompt if initialized
            if hasattr(st.session_state, 'chat_instance') and st.session_state.chat_instance:
                try:
                    st.session_state.chat_instance.system_prompt = system_prompt
                    st.success("System prompt updated successfully")
                    logger.info("System prompt updated")
                except Exception as e:
                    logger.error(f"Error updating system prompt: {str(e)}")
                    st.error(f"Error updating system prompt: {str(e)}")

            
        st.divider()
        
        # Retrieval Toggle
        st.markdown("### Document Retrieval")
        retrieval_status = "🔍 ON" if st.session_state.enable_retrieval else "🚫 OFF"
        toggle_label = f"Database Retrieval: {retrieval_status}"
        
        if st.toggle(toggle_label, value=st.session_state.enable_retrieval, key="retrieval_toggle"):
            st.session_state.enable_retrieval = True
        else:
            st.session_state.enable_retrieval = False
        
        st.divider()
        
        # Model Selection Dropdown
        st.markdown("### Model Selection")
        
        # Only show model dropdown if Ollama is running
        if ollama_status:
            # Get available models
            models = get_available_models()
            
            # Create dropdown for model selection
            if models:
                selected_model = st.selectbox(
                    "Select a model:",
                    models,
                    index=models.index(st.session_state.selected_model) if st.session_state.selected_model in models else 0
                )
            
            # Apply model change
            if selected_model != st.session_state.selected_model:
                try:
                    # Update the session state
                    st.session_state.selected_model = selected_model
                    
                    # Create new chat instance with selected model
                    st.session_state.chat_instance = OllamaChat(
                        model=selected_model,
                            system_prompt=st.session_state.system_prompt
                    )
                    st.session_state.initialized = True
                    st.success(f"Model changed to: {selected_model}")
                    logger.info(f"Model changed to: {selected_model}")
                except Exception as e:
                    logger.error(f"Error changing model: {str(e)}\n{traceback.format_exc()}")
                    st.error(f"Error changing model: {str(e)}")
            else:
                st.warning("No models available. Please pull models using Ollama CLI.")
        else:
            st.warning("Model selection unavailable - Ollama is not running")
        
        st.divider()
        
        # ChromaDB Collections Section
        with st.expander("ChromaDB Collections", expanded=True):
            st.markdown("### Add Documents to ChromaDB")
            
            # File Upload for ChromaDB
            uploaded_files = st.file_uploader(
                "Upload documents to add to ChromaDB:",
                accept_multiple_files=True,
                type=["pdf", "txt", "csv", "md", "doc", "docx"],
                key="chromadb_uploader"
            )
            
            # Get current database and existing collections
            st.session_state.current_database = 'default'
            current_db = st.session_state.current_database
            db_path = get_user_storage_path(st.session_state.username)
            current_db_path = os.path.join(db_path, current_db)
            client = chromadb.PersistentClient(path=current_db_path)
            existing_collections = [col for col in client.list_collections() if not col.startswith("temp")]
            
            # Toggle options for extraction features
            st.markdown("### Extraction Options")
            col1, col2 , col3= st.columns(3)
            st.session_state.extract_figures = False
            st.session_state.extract_tables = True
            with col1:
                extract_figures = st.toggle(
                    "Extract Figures",
                    value=st.session_state.get('extract_figures', False),
                    help="Enable to extract figures from documents (may increase processing time)"
                )
                st.session_state.extract_figures = extract_figures
                
            with col2:
                extract_tables = st.toggle(
                    "Extract Tables",
                    value=st.session_state.get('extract_tables', True),
                    help="Enable to extract tables from documents"
                )
                st.session_state.extract_tables = extract_tables
            with col3:
                topk= st.slider(
                    "Top K",
                    min_value=1,
                    max_value=100,
                    value=st.session_state.get('topk', 5),
                    help="Number of documents to retrieve"
                )
                st.session_state.topk = topk
            # Show available collections even before upload
            if existing_collections:
                st.markdown("### Available Collections")
                selected_collection = st.selectbox(
                    "Select Existing Collection:", 
                    existing_collections,
                    key="upload_existing_collection"
                )
                st.session_state.active_collection = selected_collection
            # Process files for ChromaDB
            if uploaded_files:
                st.markdown("### Configure Collection for Uploaded Files")

                # Collection selection
                collection_choice = st.radio(
                    "Choose Collection Option", 
                    ["Create New Collection", "Use Existing Collection"],
                    key="upload_collection_choice"
                )
                
                if collection_choice == "Use Existing Collection":
                    # Dropdown for existing collections
                    pass
                if collection_choice == "Create New Collection":
                    # Input for custom collection name
                    st.markdown("#### Create New Collection")
                    st.info("Enter a descriptive name for your collection. This name should reflect the content or topic of your documents.")
                    
                    new_collection_name = st.text_input(
                        "Collection Name", 
                        key="upload_new_collection_name",
                        help="Choose a descriptive name related to the document content. For example: 'machine_learning', 'research_papers', 'company_policies', etc."
                    )
                    
                    # Validate collection name
                    if new_collection_name and new_collection_name in existing_collections:
                        st.warning(f"Collection '{new_collection_name}' already exists. Files will be added to this existing collection.")
                    
                    # Optional: Add a description or tags
                    collection_description = st.text_area(
                        "Collection Description (Optional)", 
                        key="upload_collection_description",
                        help="Provide additional context about what this collection contains or represents."
                    )
                    
                    selected_collection = new_collection_name
                
                # Process the files
                if st.button("Process Files", key="upload_process_files_btn", disabled=not (selected_collection and selected_collection.strip())):
                    with st.spinner("Processing files..."):
                        success_count = 0
                        for file in uploaded_files:
                            success = process_file(file, selected_collection, 
                                                  collection_description if collection_choice == "Create New Collection" else None)
                            if success:
                                success_count += 1
                        
                        if success_count == len(uploaded_files):
                            st.success(f"Successfully processed all {success_count} files to collection '{selected_collection}'")
                        elif success_count > 0:
                            st.warning(f"Processed {success_count} out of {len(uploaded_files)} files to collection '{selected_collection}'")
                        else:
                            st.error("Failed to process any files. Please check the logs for details.")
            
            # Show help text 
            st.markdown("#### How to use file upload")
            st.markdown("""
            1. Upload one or more documents using the file uploader
            2. Choose to create a new collection or use an existing one
            3. If creating a new collection, enter a descriptive name for your documents
            4. Click "Process Files" to add the documents to the collection
            5. Use the chat interface to ask questions about your documents
            """)
            
            # Show available collections in current database
            collections = get_collections_for_database(st.session_state.current_database)
            if collections:
                st.markdown("### Available Collections")
                collection_names = list(collections.keys())
                
                # Allow selecting a collection
                selected_collection = st.selectbox(
                    "Select Collection to Use:",
                    collection_names,
                    key="selected_chromadb_collection"
                )
                
                if selected_collection:
                    st.session_state.active_collection = selected_collection
                    st.success(f"Using collection: {selected_collection}")
                    
                    # Show collection details
                    if selected_collection in collections:
                        coll_info = collections[selected_collection]
                        st.markdown(f"**Collection:** {selected_collection}")
                        
                        # Show all files in the collection if available
                        if "files" in coll_info and len(coll_info["files"]) > 0:
                            st.markdown(f"**Files:** {len(coll_info['files'])} document(s)")
                            
                            # Display first 5 files directly
                            files_to_show = min(5, len(coll_info["files"]))
                            for idx, file in enumerate(coll_info["files"][:files_to_show]):
                                st.markdown(f"- {file}")
                                
                            # If there are more files, provide option to view more
                            remaining_files = len(coll_info["files"]) - files_to_show
                            if remaining_files > 0:
                                display_option = st.radio(
                                    f"Show all files ({remaining_files} more)",
                                    ["Hide Additional Files", "Show All Files"],
                                    key=f"show_all_{selected_collection}",
                                    horizontal=True
                                )
                                if display_option == "Show All Files":
                                    for idx, file in enumerate(coll_info["files"][files_to_show:]):
                                        st.markdown(f"- {file}")
                        else:
                            st.markdown(f"**File:** {coll_info.get('file_name', 'Unknown')}")
                            
                        st.markdown(f"**Type:** {coll_info.get('file_type', 'Unknown')}")
                        st.markdown(f"**Added:** {coll_info.get('added_date', 'Unknown')}")
                        

            else:
                st.info("No collections available in the current database. Upload documents to create collections.")
        
        st.divider()
        
        # File Upload Section for temporary QA
        st.markdown("### Upload Documents")
        
        # Track upload state to avoid reprocessing
        if "upload_state" not in st.session_state:
            st.session_state.upload_state = {}
            
        if "temp_collections" not in st.session_state:
            st.session_state.temp_collections = {}
        
        # Create an expander to show currently loaded temporary files
        with st.expander("Temporary Files in Memory", expanded=True):
            if st.session_state.temp_collections:
                for temp_coll_name, file_info in st.session_state.temp_collections.items():
                    st.markdown(f"📄 **{file_info['file_name']}** - *Collection: {temp_coll_name}*")
                
                # Add button to clear all temporary collections
                if st.button("Clear All Temporary Files"):
                    st.session_state.temp_collections = {}
                    st.success("All temporary files cleared from memory")
                    try:
                        collections = st.session_state.loader.vectorstore.list_collections()
                        for collection in collections:
                                if collection.startswith("temp_"):
                                        st.session_state.loader.vectorstore.delete_collection(collection)
                    except Exception as e:
                        st.error(f"Error clearing temporary files: {str(e)}")
                    st.rerun()
            else:
                st.info("No temporary files loaded. Upload files below for quick Q&A without saving to your database.")
            
        uploaded_files = st.file_uploader(
            "Upload documents for temporary QA (not saved to database):",
            accept_multiple_files=True,
            type=["pdf", "txt", "csv", "md", "doc", "docx"]
        )
        
        # Process newly uploaded files
        if uploaded_files:
            with st.spinner("Processing files for quick Q&A..."):
                # Add a progress bar
                progress_bar = st.progress(0)
                status_text = st.empty()
                
                # Track processing status
                all_processed = True
                
                total_files = len(uploaded_files)
                for idx, file in enumerate(uploaded_files):
                    # Update progress
                    progress = float(idx) / float(total_files)
                    progress_bar.progress(progress)
                    status_text.text(f"Processing file {idx+1}/{total_files}: {file.name}")
                    
                    # Create a unique identifier for this file
                file_id = file.name.replace(" ", "_").replace(".", "_")
                
                # Only process if it's a new file we haven't seen before
                if file_id not in st.session_state.upload_state:
                        # Create a temporary collection name specific to this user and file
                        temp_collection_name = f"temp_{st.session_state.username}_{file_id}"
                            
                        # Mark as being processed
                        st.session_state.upload_state[file_id] = "processing"
                        
                        # Create a temporary file
                        temp_dir = tempfile.mkdtemp()
                        temp_file_path = os.path.join(temp_dir, file.name)
                        
                        try:
                            # Write file to temporary location
                            with open(temp_file_path, "wb") as f:
                                f.write(file.getbuffer())
                            
                            # Get file extension
                            file_extension = Path(file.name).suffix.lstrip(".").lower()
                            
                            # Get user-specific storage path
                            user_storage_path = get_user_storage_path(st.session_state.username)
                            
                            # Set persistent directory for the loader
                            current_db = st.session_state.current_database
                            current_db_path = os.path.join(user_storage_path, current_db)
                            st.session_state.loader.persist_directory = current_db_path
                            
                            # Process the file with timeout handling
                            logger.info(f"Processing temporary file: {file.name} into collection: {temp_collection_name}")
                            status_text.text(f"Processing file {idx+1}/{total_files}: {file.name} - Generating chunks and embeddings...")
                            
                            # Import the needed modules for timeout handling
                            import concurrent.futures
                            import time
                            
                            # Define a function to process with timeout
                            def process_with_timeout(timeout=120):  # 2-minute timeout
                                try:
                                    # Store local references to avoid session state access in thread
                                    local_loader = st.session_state.loader
                                    
                                    # Define a completely isolated thread function - no Streamlit dependencies
                                    def process_file_thread():
                                        try:
                                            # No session state access in this function
                                            result = local_loader.create_or_insert_collection(
                                                temp_file_path,
                                                temp_collection_name,
                                                extract_figures=st.session_state.get('extract_figures', False),
                                                extract_tables=st.session_state.get('extract_tables', True),
                                                file_extension=file_extension,
                                                parser_type='local_parser'

                                            )
                                            return result
                                        except Exception as e:
                                            # Don't use logger inside thread
                                            print(f"Error in processing thread: {str(e)}")
                                            return False

                                    # Try to create the temporary collection in a separate thread
                                    with concurrent.futures.ThreadPoolExecutor() as executor:
                                        future = executor.submit(process_file_thread)
                                        
                                        # Poll and update status while waiting
                                        start_time = time.time()
                                        processing = True
                                        while processing and time.time() - start_time < timeout:
                                            if future.done():
                                                processing = False
                                                return future.result()
                                            else:
                                                # Only update the status message from the main thread
                                                elapsed = time.time() - start_time
                                                # This will work because status_text is in the main thread context
                                                status_text.text(f"Processing file {idx+1}/{total_files}: {file.name} - {elapsed:.1f}s elapsed...")
                                                # Add a small sleep to prevent UI from freezing
                                                time.sleep(0.5)
                                        
                                        # If we get here, we've timed out
                                        if not future.done():
                                            future.cancel()
                                            raise TimeoutError("Processing timed out")

                                except Exception as e:
                                    logger.error(f"Error in processing with timeout: {str(e)}")
                                    return False
                            
                            # Process with timeout
                            success = process_with_timeout(6000)
                            
                            if success:
                                # Store in file_collections for retrieval
                                if "file_collections" not in st.session_state:
                                    st.session_state.file_collections = {}
                                
                                st.session_state.file_collections[file_id] = temp_collection_name
                                
                                # Store metadata for display
                                st.session_state.temp_collections[temp_collection_name] = {
                                    "file_name": file.name,
                                    "file_type": file_extension,
                                    "added_date": time.strftime("%Y-%m-%d %H:%M:%S"),
                                    "is_temporary": True
                                }
                                
                                # Mark as processed
                                st.session_state.upload_state[file_id] = "processed"
                                status_text.text(f"✅ Processed file {idx+1}/{total_files}: {file.name}")
                            else:
                                all_processed = False
                                status_text.text(f"❌ Failed to process file {idx+1}/{total_files}: {file.name}")
                                st.session_state.upload_state[file_id] = "failed"
                        
                        except TimeoutError:
                            all_processed = False
                            logger.error(f"Processing timed out for file: {file.name}")
                            st.error(f"Processing timed out for file: {file.name} - File may be too large or complex")
                            st.session_state.upload_state[file_id] = "failed"
                        except Exception as e:
                            all_processed = False
                            logger.error(f"Error processing temporary file: {str(e)}")
                            st.error(f"Error processing file: {str(e)}")
                            st.session_state.upload_state[file_id] = "failed"
                        finally:
                            # Clean up temporary files
                            try:
                                if os.path.exists(temp_file_path):
                                    os.remove(temp_file_path)
                                if os.path.exists(temp_dir):
                                    os.rmdir(temp_dir)
                            except Exception as cleanup_error:
                                logger.warning(f"Failed to clean up temporary files: {str(cleanup_error)}")
                        
                        # Complete the progress bar
                        progress_bar.progress(1.0)
                        status_text.text("Processing completed")
                        
                        if all_processed:
                            st.success(f"Successfully processed all {total_files} files")
                        else:
                            st.warning("Some files failed to process. See error messages above.")
                            
                        # Force UI to refresh after a short delay
                        time.sleep(1)  # Give user time to see completion message
                        st.rerun()
        
        # Help text for temporary uploads
        st.markdown("""
        **Temporary uploads:**
        - Files uploaded here are processed immediately for Q&A
        - Context is retrieved automatically when you ask questions
        - These files remain available between sessions
        - They don't appear in your regular database collections
        - Use the database upload section above for permanent storage
        """)
                
        st.divider()
        
        # Display collections dropdown
        show_collections_dropdown()
        
        st.divider()
        
        # Clear Chat Button
        if st.button("🗑️ Clear Chat History"):
            st.session_state.chat_history = []
            st.session_state.messages = []
            #st.session_state.file_collections = {}
            #st.session_state.upload_state = {}
            
            # Don't reinitialize these - just keep them
            # st.session_state.loader - keep this
            # st.session_state.chat_instance - keep this
            
            st.success("Chat history and documents have been cleared.")
        
        # Show Models Information
        with st.expander("🤖 Available Models", expanded=True):
            if ollama_status:
                models = get_available_models()
                if models:
                    st.markdown("### Ollama Models")
                    for model in models:
                        if model == st.session_state.selected_model:
                            st.markdown(f"✅ **{model}** (selected)")
                        else:
                            st.markdown(f"• {model}")
                    
                    # Add a button to refresh models list
                    if st.button("Refresh Models"):
                        if "cached_models" in st.session_state:
                            del st.session_state.cached_models
                        st.rerun()
                else:
                    st.warning("No models found. Is Ollama running?")
            else:
                st.warning("Cannot retrieve models - Ollama is not running")
        
        # Help Information
        with st.expander("ℹ️ Help & Information"):
            st.markdown("""
            ### RAG-Agent Chat Help
            
            **Features:**
            - Chat with selected Ollama model
            - Upload documents to ask questions about them
            - Toggle database retrieval on/off
            - Choose from different models
            - Select or create databases
            
            **Tips:**
            - Upload multiple files to build a larger knowledge base
            - Be specific with your questions for better answers
            - Toggle database retrieval off if you want pure model responses
            - Clear chat history to start fresh
            - Use different databases for different projects
            """)
            
    # Main Chat Area
    chat_container = st.container()
    
    with chat_container:
        # Display header
        st.title("RAG-Agent Chat")

        # Show database info in the header
        if hasattr(st.session_state, 'current_database'):
            st.markdown(f"**Current Database:** {st.session_state.current_database}")
    
        # Display Ollama status
        if not st.session_state.ollama_available:
            st.error("⚠️ Ollama server is not available. Please start it by running 'ollama serve' in a terminal and refresh this page.")
        elif "chat_instance" not in st.session_state or st.session_state.chat_instance is None:
            st.warning("⚠️ Chat instance not initialized. Please check that Ollama is running correctly and the selected model is available.")
            
            # Add a retry button
            if st.button("Retry Connection"):
                try:
                    # Attempt to reconnect to Ollama
                    if is_ollama_available():
                        model_name = st.session_state.selected_model
                        st.session_state.chat_instance = OllamaChat(
                            model=model_name,
                            system_prompt=st.session_state.system_prompt
                        )
                        st.session_state.initialized = True
                        st.success(f"Successfully connected to Ollama with model: {model_name}")
                        st.rerun()
                    else:
                        st.error("Ollama is still not available. Please start it and try again.")
                except Exception as e:
                            st.error(f"Failed to initialize chat: {str(e)}")
    
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
                st.write(message["content"])
    
    # Display retrieved context
    context_container = st.container()
    
    # Chat input - enable only if chat is initialized and Ollama is running
    chat_disabled = not st.session_state.ollama_available or not st.session_state.initialized
        
    if chat_disabled:
        st.warning("Chat input disabled. Please ensure Ollama is running and a model is loaded.")
        prompt = None
    else:
        prompt = st.chat_input("Ask a question about your documents...")
    
    if prompt:
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Display assistant response
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            context = ""
            
            # Get context if retrieval is enabled
            if hasattr(st.session_state, 'enable_retrieval') and st.session_state.enable_retrieval and "file_collections" in st.session_state:
                with context_container:
                    context_status = st.empty()
                    context_status.info("🔍 Retrieving relevant context...")
                    
                    try:
                        context, context_chunks = get_context(prompt)
                        
                        if context_chunks:
                            st.success(f"🔍 Found {min(st.session_state.topk,len(context_chunks))} and enhanced to {len(context_chunks)} relevant document chunks.")
                            with st.expander("Show Retrieved Context"):
                                for i, chunk in enumerate(context_chunks):
                                    st.markdown(chunk)
                        else:
                            context_status.info("🔍 No relevant document chunks found for this query.")
                    except Exception as e:
                        context_status.error(f"Error retrieving context: {str(e)}")
                        logger.error(f"Error retrieving context: {str(e)}\n{traceback.format_exc()}")
            
            # Generate response
            full_response = ""
                
                # Double check initialization before proceeding
            if "chat_instance" not in st.session_state or st.session_state.chat_instance is None:
                    message_placeholder.error("Chat instance is not initialized. Please try refreshing the page.")
                    return
            
            try:
                message_placeholder.markdown("Thinking...")
                
                # Get chat history in the format expected by OllamaChat
                chat_history = []
                if hasattr(st.session_state, 'chat_history'):
                    chat_history = st.session_state.chat_history
                else:
                    st.session_state.chat_history = chat_history
                        
                # Generate response in a streaming fashion
                result_queue = generate_llm_response(prompt, chat_history, context)
                
                # Stream the response
                timeout_count = 0
                while True:
                    try:
                        result_type, content = result_queue.get(timeout=0.1)
                        
                        if result_type == "chunk":
                            full_response += content
                            message_placeholder.markdown(full_response + "▌")
                            timeout_count = 0  # Reset timeout counter on successful retrieval
                        elif result_type == "done":
                            message_placeholder.markdown(full_response)
                            
                            # Update chat history
                            chat_history.append(f"User: {prompt}\n{full_response}")
                            
                            # Keep only the last 10 interactions
                            if len(chat_history) > 10:
                                chat_history = chat_history[-10:]
                            
                            st.session_state.chat_history = chat_history
                            break
                        elif result_type == "error":
                            message_placeholder.error(f"Error: {content}")
                            break
                    except Exception as queue_error:
                        # No result yet, continue waiting but with a timeout
                        timeout_count += 1
                        # If we've waited too long (5 seconds) with no response, show status
                        if timeout_count > 50 and timeout_count % 50 == 0:  # Every ~5 seconds
                            message_placeholder.markdown(f"{full_response}... _still thinking_▌")
                        continue
                
            except Exception as e:
                logger.error(f"Error generating response: {str(e)}\n{traceback.format_exc()}")
                message_placeholder.error(f"Error: {str(e)}")
            
            # Add assistant response to chat history
            if full_response:
                st.session_state.messages.append({"role": "assistant", "content": full_response})

# Function to show login page
def show_login_page():
    # Initialize user database
    init_user_database()
    
    st.title("RAG-Agent Chat Login")
    
    # Center the login form
    col1, col2, col3 = st.columns([1, 2, 1])
    
    with col2:
        login_mode = st.radio("Choose Action", ["Login", "Register"])
        
        if login_mode == "Login":
            st.markdown("### Login")
            username = st.text_input("Username")
            password = st.text_input("Password", type="password")
            
            if st.button("Login"):
                success, message, database_collections, temp_collections, file_collections = login_user(username, password)
                if success:
                    st.session_state.authenticated = True
                    st.session_state.username = username
                    
                    # Set persistent collections
                    st.session_state.database_collections = database_collections
                    st.session_state.temp_collections = {}
                    st.session_state.file_collections = file_collections
                    
                    # Set current database
                    available_dbs = list(database_collections.keys()) if database_collections else []
                    if available_dbs:
                        st.session_state.current_database = available_dbs[0]
                    
                    
                else:
                    st.error(message)
        
        else:
            st.markdown("### Register")
            new_username = st.text_input("Choose Username")
            new_password = st.text_input("Choose Password", type="password")
            confirm_password = st.text_input("Confirm Password", type="password")
            
            if st.button("Register"):
                if new_password != confirm_password:
                    st.error("Passwords do not match")
                else:
                    success, message = register_user(new_username, new_password)
                    if success:
                        st.success(message)
                        # Optionally, automatically log in after registration
                        st.session_state.authenticated = True
                        st.session_state.username = new_username
                        
                        # Initialize default database for new user
                        user_storage_path = get_user_storage_path(new_username)
                        # Create default database directory if it doesn't exist
                        default_db_path = os.path.join(user_storage_path, "default")
                        os.makedirs(default_db_path, exist_ok=True)
                        # Get collections from user path
                        database_collections = {}
                        available_dbs = get_available_databases()
                        for db_name in available_dbs:
                            database_collections[db_name] = {}
                        st.session_state.database_collections = database_collections
                        st.session_state.current_database = 'default'
                        
                        # Initialize empty temp collections
                        st.session_state.temp_collections = {}
                        st.session_state.file_collections = {}
                        
                        st.rerun()
                    else:
                        st.error(message)

# Function to get available Ollama models
def get_available_models():
    """Retrieve list of available Ollama models"""
    # Only fetch models once per session, rather than using a TTL-based cache
    if "cached_models" not in st.session_state:
        try:
            # Check if Ollama is available first
            if not is_ollama_available():
                logger.warning("Ollama server is not running. Cannot retrieve models.")
                st.session_state.cached_models = [default_model]  # Store in session state
                return st.session_state.cached_models
            
            # Use ollama library to list models
            response = ollama.list()
            
            # Extract model names
            if response and 'models' in response:
                models = [model['model'] for model in response['models']]
                
                # Ensure default model is in the list
                if default_model not in models:
                    models.insert(0, default_model)
                
                st.session_state.cached_models = models  # Store in session state
                logger.info(f"Retrieved {len(models)} models from Ollama")
                return models
            else:
                logger.warning("No models found. Using default model.")
                st.session_state.cached_models = [default_model]  # Store in session state
                return st.session_state.cached_models
        
        except Exception as e:
            logger.error(f"Error retrieving Ollama models: {str(e)}")
            st.session_state.cached_models = [default_model]  # Store in session state
            return st.session_state.cached_models
    else:
        # Return cached models
        return st.session_state.cached_models

# Function to show database editor
def show_database_editor():
    st.markdown(f"**Current Database:** {st.session_state.current_database}")
    
    # Get collections for current database
    collections = get_collections_for_database(st.session_state.current_database)
    
    if collections:
        st.markdown("**Collections:**")
        
        # Add search filter for collections
        filter_query = st.text_input("Filter collections", key="collection_filter")
        filtered_collections = {k: v for k, v in collections.items() 
                               if not filter_query or filter_query.lower() in k.lower()}
        
        # Display collections in a table format if there are many
        if len(filtered_collections) > 5:
            collection_data = []
            for coll_name, coll_info in filtered_collections.items():
                file_name = coll_info.get("file_name", "Unknown")
                file_type = coll_info.get("file_type", "Unknown")
                added_date = coll_info.get("added_date", "Unknown")
                collection_data.append({
                    "Collection": coll_name,  # Show collection name first
                    "File": file_name,
                    "Type": file_type,
                    "Added": added_date
                })
            
            # Create a DataFrame for display
            import pandas as pd
            df = pd.DataFrame(collection_data)
            st.dataframe(df)
            
            # Handle collection actions
            st.markdown("**Collection Actions:**")
            col1, col2 = st.columns(2)
            with col1:
                collection_to_edit = st.selectbox("Select collection", list(filtered_collections.keys()))
            with col2:
                action = st.selectbox("Action", ["View Details", "Delete"])
            
            if st.button("Execute Action"):
                if action == "Delete":
                    if delete_collection(st.session_state.current_database, collection_to_edit):
                        st.success(f"Deleted collection: {collection_to_edit}")
                        st.rerun()
                elif action == "View Details":
                    show_collection_details(st.session_state.current_database, collection_to_edit)
        else:
            # For fewer collections, show more detailed individual cards
            for coll_name, coll_info in filtered_collections.items():
                with st.container():
                    # Headline using collection name instead of file name
                    st.markdown(f"📄 **Collection: {coll_name}**")
                    col1, col2 = st.columns([3, 1])
                    
                    with col1:
                        # Show all files in the collection if available
                        if "files" in coll_info and len(coll_info["files"]) > 0:
                            st.markdown(f"**Files:** {len(coll_info['files'])} document(s)")
                            
                            # Use a radio button instead of checkbox to avoid nesting issues
                            display_option = st.radio(
                                f"Files in '{coll_name}'",
                                ["Hide Files", "Show Files"],
                                key=f"display_files_{coll_name}",
                                horizontal=True
                            )
                            
                            if display_option == "Show Files":
                                # Display files directly in the main container
                                max_files_displayed = 10
                                for idx, file in enumerate(coll_info["files"][:max_files_displayed]):
                                    st.markdown(f"{idx+1}. {file}")
                                
                                # If there are more files than can comfortably be shown
                                if len(coll_info["files"]) > max_files_displayed:
                                    st.markdown(f"... and {len(coll_info['files']) - max_files_displayed} more files")
                        else:
                            # Fallback to single file
                            st.markdown(f"**File:** {coll_info.get('file_name', 'Unknown')}")
                            
                        # Show type information
                        if "entries" in coll_info:
                            # This is a text input collection
                            st.markdown(f"**Type:** Text Collection")
                            st.markdown(f"**Entries:** {coll_info.get('entries', 1)}")
                        else:
                            # This is a file-based collection
                            st.markdown(f"**Type:** {coll_info.get('file_type', 'Unknown')}")
                                
                            st.markdown(f"**Added:** {coll_info.get('added_date', 'Unknown')}")
                        if "modified_date" in coll_info:
                            st.markdown(f"**Last Modified:** {coll_info.get('modified_date', 'Unknown')}")
                        
                        # Add notes or description field that can be edited
                        if "description" not in coll_info:
                            coll_info["description"] = ""
                        
                        new_description = st.text_area(f"Notes for '{coll_name}'", 
                                                      value=coll_info.get("description", ""), 
                                                      key=f"desc_{coll_name}")
                        
                        if new_description != coll_info.get("description", ""):
                            st.session_state.database_collections[st.session_state.current_database][coll_name]["description"] = new_description
                            st.success("Notes saved!")
                    
                    with col2:
                        if st.button("Delete", key=f"del_{coll_name}"):
                            if delete_collection(st.session_state.current_database, coll_name):
                                st.success(f"Deleted collection: {coll_name}")
                                st.rerun()
                        
                        # Button to view document chunks for debugging
                        if st.button("Inspect", key=f"inspect_{coll_name}"):
                            inspect_collection(coll_name)
    else:
        st.info("No collections in this database yet. Upload documents to create collections.")
    
    # Rename and Import/Export sections
    st.markdown("### Database Management")
    
    # Rename Database
    rename_container = st.container()
    with rename_container:
        new_name = st.text_input("Rename Current Database", key="rename_db_input")
        if st.button("Rename Database"):
            if new_name and new_name != st.session_state.current_database:
                if new_name not in get_available_databases():
                    rename_database(st.session_state.current_database, new_name)
                    st.success(f"Database renamed to '{new_name}'")
                    st.rerun()
                else:
                    st.error(f"Database '{new_name}' already exists")
    
    # Import/Export
    import_export_container = st.container()
    with import_export_container:
        col1, col2 = st.columns(2)
        with col1:
            st.download_button(
                "Export Database",
                data=get_database_export_data(st.session_state.current_database),
                file_name=f"{st.session_state.current_database}_export.json",
                mime="application/json"
            )
        with col2:
            st.file_uploader("Import Database", type=["json"], key="db_import_file", 
                            on_change=import_database_callback)

# Function to inspect a collection
def inspect_collection(collection_name):
    if not init_loader():
        st.error("Loader is not initialized")
        return
    
    collection_id = f"{collection_name}"
    
    try:
        # Get collection details - handle this more gracefully
        try:
            # Try to get the collection directly from ChromaDB
            user_storage_path = get_user_storage_path(st.session_state.username)
            current_db = st.session_state.current_database
            current_db_path = os.path.join(user_storage_path, current_db)
            
            # Create a local client to avoid thread issues
            chroma_client = chromadb.PersistentClient(path=current_db_path)
            
            # Check if collection exists
            try:
                collection = chroma_client.get_collection(name=collection_id)
                
                # Get collection data safely
                collection_data = {}
                try:
                    # Get count
                    count = collection.count()
                    collection_data["count"] = count
                    
                    # Get a sample of documents if available
                    if count > 0:
                        sample = collection.get(limit=min(5, count))
                        
                        # Only include important fields to avoid overwhelming the display
                        collection_data["sample"] = {
                            "ids": sample.get("ids", [])[:5],
                            "metadatas_count": len(sample.get("metadatas", [])),
                            "documents_preview": [doc[:100] + "..." if len(doc) > 100 else doc 
                                               for doc in sample.get("documents", [])[:3]]
                        }
                except Exception as e:
                    collection_data["error"] = f"Error retrieving data: {str(e)}"
                
                st.json(collection_data)
                
            except Exception as coll_error:
                st.error(f"Collection '{collection_id}' not found: {str(coll_error)}")
                
        except Exception as e:
            st.error(f"Error accessing ChromaDB: {str(e)}")
            
    except Exception as e:
        st.error(f"Error inspecting collection: {str(e)}")
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            st.error(traceback.format_exc())

# Function to show collection details
def show_collection_details(database_name, collection_name):
    if database_name in st.session_state.database_collections:
        if collection_name in st.session_state.database_collections[database_name]:
            coll_info = st.session_state.database_collections[database_name][collection_name]
            
            st.markdown(f"### Collection: {collection_name}")
            st.markdown(f"**File:** {coll_info.get('file_name', 'Unknown')}")
            st.markdown(f"**Type:** {coll_info.get('file_type', 'Unknown')}")
            st.markdown(f"**Added:** {coll_info.get('added_date', 'Unknown')}")
            st.markdown(f"**Description:** {coll_info.get('description', '')}")
            
            # Add option to inspect collection chunks
            if st.button("Inspect Collection Data"):
                inspect_collection(collection_name)
        else:
            st.error(f"Collection {collection_name} not found")
    else:
        st.error(f"Database {database_name} not found")

# Helper function to export database data
def get_database_export_data(database_name):
    if database_name in st.session_state.database_collections:
        export_data = {
            "database_name": database_name,
            "collections": st.session_state.database_collections[database_name],
            "export_date": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        return json.dumps(export_data, indent=2)
    else:
        return json.dumps({"error": "Database not found"})

# Callback for database import
def import_database_callback():
    uploaded_file = st.session_state.db_import_file
    if uploaded_file is not None:
        try:
            content = uploaded_file.getvalue().decode('utf-8')
            import_data = json.loads(content)
            
            # Check if it's a valid export
            if "database_name" in import_data and "collections" in import_data:
                db_name = import_data["database_name"]
                
                # Add database if it doesn't exist
                if db_name not in st.session_state.available_databases:
                    st.session_state.available_databases.append(db_name)
                
                # Add collections
                if "database_collections" not in st.session_state:
                    st.session_state.database_collections = {}
                
                st.session_state.database_collections[db_name] = import_data["collections"]
                
                # Switch to the imported database
                st.session_state.current_database = db_name
                st.success(f"Successfully imported database: {db_name}")
                
                # Clear the file uploader
                st.session_state.db_import_file = None
            else:
                st.error("Invalid database export format")
        except Exception as e:
            st.error(f"Error importing database: {str(e)}")
            
# Function to get collections for a database
def get_collections_for_database(database_name):
    # In a real app, this would query actual collections
    if "database_collections" not in st.session_state:
        st.session_state.database_collections = {}
    
    if database_name not in st.session_state.database_collections:
        st.session_state.database_collections[database_name] = {}
    
    return st.session_state.database_collections[database_name]

# Function to delete a collection
def delete_collection(database_name, collection_name):
    # Delete from session state
    if database_name in st.session_state.database_collections:
        if collection_name in st.session_state.database_collections[database_name]:
            del st.session_state.database_collections[database_name][collection_name]
            
    # Also delete from file_collections if present
    collection_value = f"collection_{collection_name}"
    keys_to_delete = []
    for key, value in st.session_state.file_collections.items():
        if value == collection_value:
            keys_to_delete.append(key)
            
    for key in keys_to_delete:
        del st.session_state.file_collections[key]
    
    # In a real app, this would also delete from the actual database
    return True

# Function to rename a database
def rename_database(old_name, new_name):
    """Rename a database for the current user"""
    if not hasattr(st.session_state, 'username'):
        st.error("Please log in first")
        return False
    
    try:
        conn = sqlite3.connect(USERS_DB_PATH)
        cursor = conn.cursor()
        
        # Retrieve current databases
        cursor.execute("SELECT databases FROM users WHERE username = ?", 
                       (st.session_state.username,))
        result = cursor.fetchone()
        
        import json
        current_databases = json.loads(result[0]) if result and result[0] else []
        
        # Check if old database exists and new name is unique
        if old_name not in current_databases:
            st.error(f"Database '{old_name}' does not exist")
            return False
        
        if new_name in current_databases:
            st.error(f"Database '{new_name}' already exists")
            return False
        
        # Rename the database
        current_databases = [new_name if db == old_name else db for db in current_databases]
        
        # Update user's databases in SQLite
        cursor.execute(
            "UPDATE users SET databases = ? WHERE username = ?", 
            (json.dumps(current_databases), st.session_state.username)
        )
        conn.commit()
        
        # Update session state
        st.session_state.current_database = new_name
        
        # Rename collections if they exist in session state

    
    except Exception as e:
        st.error(f"Error renaming database: {str(e)}")
        return False
    finally:
        conn.close()
    if "database_collections" in st.session_state and old_name in st.session_state.database_collections:
        st.session_state.database_collections[new_name] = st.session_state.database_collections[old_name]
        del st.session_state.database_collections[old_name]
    
        return True
# User registration function
def register_user(username, password):
    """Register a new user in the SQLite database"""
    if len(username) < 3:
        return False, "Username must be at least 3 characters long"
    
    # Hash the password
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    
    try:
        conn = sqlite3.connect(USERS_DB_PATH)
        cursor = conn.cursor()
        
        try:
            cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", 
                           (username, hashed_password))
            conn.commit()
            return True, "Registration successful"
        except sqlite3.IntegrityError:
            return False, "Username already exists"
        finally:
            conn.close()
    except Exception as e:
        return False, f"Registration error: {str(e)}"

# Process file uploads
def process_file(file, collection_name=None, collection_description=None):
    """
    Process an uploaded file and insert it into a ChromaDB collection.
    
    Args:
        file (UploadedFile): The file uploaded by the user
        collection_name (str, optional): Name of the collection to insert the file into
        collection_description (str, optional): Description of the collection
    
    Returns:
        bool: True if successful, False otherwise
    """
    if not init_loader():
        st.error("Document loader could not be initialized. Cannot process files.")
        return False
    
    logger.info(f"Processing uploaded file: {file.name}")
    
    # Create a temporary file
    temp_dir = tempfile.mkdtemp()
    temp_file_path = os.path.join(temp_dir, file.name)
    
    try:
        # Write file to temporary location
        with open(temp_file_path, "wb") as f:
            f.write(file.getbuffer())
        
        # Get file extension
        file_extension = Path(file.name).suffix.lstrip(".").lower()
        
        # Get user-specific storage path
        user_storage_path = get_user_storage_path(st.session_state.username)
        
        # If collection name is not provided, we need to prompt the user
        if collection_name is None:
            # Get existing collections for the current database
            current_db = st.session_state.current_database
            db_path = get_user_storage_path(st.session_state.username)
            current_db = os.path.join(db_path, current_db)
            existing_collections = list(st.session_state.database_collections.get(current_db, {}).keys())
            
            # Create a container for collection selection
            collection_container = st.container()
            
            with collection_container:
                # Radio button to choose between existing and new collection
                collection_choice = st.radio(
                    "Choose Collection", 
                    ["Use Existing Collection", "Create New Collection"],
                    key="collection_selection_radio"
                )
                
                if collection_choice == "Use Existing Collection":
                    # Dropdown for existing collections
                    if existing_collections:
                        selected_collection = st.selectbox(
                            "Select Existing Collection:", 
                            existing_collections,
                            key="existing_collection_select"
                        )
                    else:
                        st.warning("No existing collections. Please create a new one.")
                        collection_choice = "Create New Collection"
                
                if collection_choice == "Create New Collection":
                    # Input for new collection name with more guidance
                    st.info(f"Uploading file: {file.name}")
                    st.info(f"File type: {file_extension}")
                    
                    # Emphasize entering a meaningful collection name
                    st.markdown("### Enter a Meaningful Collection Name")
                    st.markdown("This name will be used to identify your collection. Choose a name that describes the content or purpose of the document.")
                    
                    # Provide an example rather than a suggestion based on file type
                    example_collection_name = "my_topic_collection"
                    
                    # Add an input for custom collection name with more context
                    new_collection_name = st.text_input(
                        "Collection Name", 
                        value=example_collection_name,
                        key="new_collection_name_input",
                        help="Choose a descriptive name related to the document content. For example: 'machine_learning', 'research_papers', 'company_policies', etc."
                    )
                    
                    # Optional: Add a description or tags
                    collection_description = st.text_area(
                        "Collection Description (Optional)", 
                        key="collection_description_input",
                        help="Provide additional context about what this collection contains or represents."
                    )
                    
                    # Validate collection name
                    if not new_collection_name:
                        st.error("Collection name cannot be empty")
                        return False
                    
                    # Check if collection name already exists
                    if new_collection_name in existing_collections:
                        st.error(f"Collection '{new_collection_name}' already exists")
                        return False
                    
                    selected_collection = new_collection_name
                
                # Confirm button
                if st.button("Confirm Collection", key="confirm_collection_btn"):
                    return process_file_with_collection(file, selected_collection, collection_description if collection_choice == "Create New Collection" else None)
                
                # If not confirmed, stop execution
                st.stop()
        else:
            # If collection name is provided, process directly
            return process_file_with_collection(file, collection_name, collection_description)
            
    except Exception as e:
        logger.error(f"Unexpected error processing file: {str(e)}\n{traceback.format_exc()}")
        st.error(f"Unexpected error processing file '{file.name}': {str(e)}")
        return False
    finally:
        # Clean up
        try:
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
            if os.path.exists(temp_dir):
                os.rmdir(temp_dir)
        except Exception as cleanup_error:
            logger.warning(f"Failed to clean up temporary files: {str(cleanup_error)}")
            
def process_file_with_collection(file, collection_name, collection_description=None):
    """
    Internal function to process a file with a specified collection
    """
    if not collection_name:
        st.error("Collection name is required")
        return False
        
    logger.info(f"Processing file {file.name} for collection: {collection_name}")
    
    # Create a temporary file
    temp_dir = tempfile.mkdtemp()
    temp_file_path = os.path.join(temp_dir, file.name)
    
    try:
        # Write file to temporary location
        with open(temp_file_path, "wb") as f:
            f.write(file.getbuffer())
            
        # Get file extension
        file_extension = Path(file.name).suffix.lstrip(".").lower()
        
        # Get user-specific storage path and current database
        user_storage_path = get_user_storage_path(st.session_state.username)
        current_db = st.session_state.current_database
        
        # Ensure the current database directory exists
        current_db_path = os.path.join(user_storage_path, current_db)
        os.makedirs(current_db_path, exist_ok=True)
        
        # Set persistent directory for the loader
        st.session_state.loader.persist_directory = current_db_path
        
        # Process the file
        try:
            # Add detailed logging for debugging
            logger.info(f"Attempting to create/insert collection for file: {file.name}")
            logger.info(f"Persistent directory: {current_db_path}")
            logger.info(f"Collection name: {collection_name}")
            
            # Try to create or insert the collection
            try:
                st.session_state.loader.create_or_insert_collection(
                    temp_file_path,
                    collection_name,
                    extension=file_extension,
                    parser_type='local_parser',
                    extract_figures=st.session_state.get('extract_figures', False),
                    extract_tables=st.session_state.get('extract_tables', True)
                )
            except Exception as embedding_error:
                # More detailed error logging
                logger.error(f"Embedding generation failed: {str(embedding_error)}")
                logger.error(f"Traceback: {traceback.format_exc()}")
                
                # Provide more user-friendly error message
                st.error(f"Failed to generate embeddings for the file. Possible reasons:")
                st.error("1. Unsupported file type")
                st.error("2. File is empty or corrupted")
                st.error("3. Embedding model is not working correctly")
                
                # Attempt to diagnose the issue
                try:
                    # Check file size and content
                    file_size = os.path.getsize(temp_file_path)
                    logger.info(f"File size: {file_size} bytes")
                    
                    # Try to read file content
                    with open(temp_file_path, 'rb') as f:
                        sample_content = f.read(1024)  # Read first 1KB
                        logger.info(f"Sample file content: {sample_content}")
                except Exception as diag_error:
                    logger.error(f"Diagnostic check failed: {str(diag_error)}")
                
                return False
            
            # Store the collection name
            if "file_collections" not in st.session_state:
                st.session_state.file_collections = {}
            
            # Use the collection name as the key
            st.session_state.file_collections[collection_name] = collection_name
            
            # Store in database collections
            if "database_collections" not in st.session_state:
                st.session_state.database_collections = {}
            if current_db not in st.session_state.database_collections:
                st.session_state.database_collections[current_db] = {}
            
            # Get or create collection metadata
            if collection_name in st.session_state.database_collections[current_db]:
                # Update existing collection metadata
                collection_metadata = st.session_state.database_collections[current_db][collection_name]
                # Update the modified date
                collection_metadata["modified_date"] = time.strftime("%Y-%m-%d %H:%M:%S")
                # Add this file to the collection's file list
                files = collection_metadata.get("files", [])
                files.append(file.name)
                collection_metadata["files"] = files
            else:
                # Create new collection metadata
                collection_metadata = {
                    "collection_name": collection_name,
                    "file_name": file.name,  # Main filename for display
                    "files": [file.name],    # List of all files in collection
                    "file_type": file_extension,
                    "added_date": time.strftime("%Y-%m-%d %H:%M:%S")
                }
                
                # Add description if provided
                if collection_description:
                    collection_metadata["description"] = collection_description
            
            # Store the updated metadata
            st.session_state.database_collections[current_db][collection_name] = collection_metadata
            
            logger.info(f"File processed successfully: {file.name}")
            return True
            
        except Exception as process_error:
            logger.error(f"Error processing file: {str(process_error)}\n{traceback.format_exc()}")
            st.error(f"Error processing file '{file.name}': {str(process_error)}")
            return False
    
    except Exception as e:
        logger.error(f"Unexpected error in process_file_with_collection: {str(e)}")
        st.error(f"Error processing file: {str(e)}")
        return False
    finally:
        # Clean up temporary files
        try:
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
            if os.path.exists(temp_dir):
                os.rmdir(temp_dir)
        except Exception as cleanup_error:
            logger.warning(f"Failed to clean up temporary files: {str(cleanup_error)}")

# Function to generate LLM response
def generate_llm_response(prompt, chat_history, context=""):
    """Generate a response from the LLM with streaming support"""
    result_queue = Queue()
    
    # Make local copies of everything needed before starting the thread
    # This completely isolates the thread from needing st.session_state
    if "chat_instance" not in st.session_state or st.session_state.chat_instance is None:
        return Queue_with_error("Chat instance not initialized. Please refresh the page and try again.")
        
    # Store all needed data in local variables to avoid st.session_state in thread
    local_system_prompt = ""
    if context:
        local_system_prompt = f"{st.session_state.system_prompt} <<Context>>: {context}"
    else:
        local_system_prompt = f"{st.session_state.system_prompt}"
    
    # Get a reference to the model outside the thread
    # We must set the system prompt before the thread starts
    chat_instance = st.session_state.chat_instance
    chat_instance.system_prompt = local_system_prompt
    
    # Add additional timeout to make sure the thread terminates
    response_timeout = 120  # seconds
    
    # Model name for logging only
    model_name = getattr(chat_instance, "model", "unknown")
    
    # Define a helper function to create a queue with an error
    def Queue_with_error(error_message):
        error_queue = Queue()
        error_queue.put(("error", error_message))
        return error_queue
    
    # Simple thread function that has zero Streamlit dependencies
    def generate_in_thread():
        try:
            # Define a simple async function that only streams from chat_async
            async def stream_response():
                try:
                    # Add a timeout to the async operation
                    async with asyncio.timeout(response_timeout):
                        # Direct call to chat_async with minimal processing
                        async for chunk in chat_instance.chat_async(
                            user_message=prompt,
                            previous_messages=chat_history,
                            context=context
                        ):
                            if chunk and chunk.strip():
                                result_queue.put(("chunk", chunk))
                        result_queue.put(("done", ""))
                except asyncio.TimeoutError:
                    result_queue.put(("error", f"Response generation timed out after {response_timeout} seconds"))
                except Exception as e:
                    # Don't use logger here to avoid StreamlitAPIException
                    result_queue.put(("error", f"Error in async streaming: {str(e)}"))
                    
            # Get a new event loop - create one if needed
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            # Run the coroutine and close the loop when done
            try:
                loop.run_until_complete(stream_response())
            except Exception as loop_error:
                result_queue.put(("error", f"Loop error: {str(loop_error)}"))
            finally:
                loop.close()
                
        except Exception as e:
            # Don't use logger here to avoid StreamlitAPIException
            result_queue.put(("error", f"Thread error: {str(e)}"))
    
    # Log before starting thread (in the main thread where it's safe)
    logger.debug(f"Starting chat with model {model_name}")
    
    # Run in a separate thread to not block the UI
    thread = Thread(target=generate_in_thread, daemon=True)
    thread.start()
    
    return result_queue

# Function to show collections dropdown in sidebar
def show_collections_dropdown():
    """Display the current database collections in a dropdown menu"""
    if "database_collections" not in st.session_state:
        return
    
    current_db = st.session_state.current_database
    collections = st.session_state.database_collections.get(current_db, {})
    
    if not collections:
        st.info("No collections available in the current database.")
        return
    
    st.markdown("### Collections")
    
    # Extract collection names and details for display
    collection_options = {}
    for coll_name, coll_info in collections.items():
        # Prioritize collection name (which is the key) instead of file name
        # This shows users the collection names they entered
        display_name = coll_name
        collection_options[display_name] = coll_name
    
    # Create the dropdown
    selected_display_name = st.selectbox(
        "Select Collection:",
        options=list(collection_options.keys()),
        key="collection_dropdown"
    )
    
    # Get the actual collection name from the selected display name
    if selected_display_name:
        selected_collection = collection_options[selected_display_name]
        
        # Show collection details
        if st.button("View Collection Details", key="view_coll_details"):
            # Get the collection info
            coll_info = collections[selected_collection]
            
            st.markdown(f"**Collection:** {selected_collection}")
            
            # Show files count and list without using expanders
            if "files" in coll_info and len(coll_info["files"]) > 0:
                st.markdown(f"**Files:** {len(coll_info['files'])} document(s)")
                
                # Display first 5 files directly
                files_to_show = min(5, len(coll_info["files"]))
                for idx, file in enumerate(coll_info["files"][:files_to_show]):
                    st.markdown(f"- {file}")
                    
                # If there are more files, provide option to view more
                remaining_files = len(coll_info["files"]) - files_to_show
                if remaining_files > 0:
                    display_option = st.radio(
                        f"Show all files ({remaining_files} more)",
                        ["Hide Additional Files", "Show All Files"],
                        key=f"show_all_{selected_collection}",
                        horizontal=True
                    )
                    if display_option == "Show All Files":
                        for idx, file in enumerate(coll_info["files"][files_to_show:]):
                            st.markdown(f"- {file}")
            else:
                st.markdown(f"**File:** {coll_info.get('file_name', 'Unknown')}")
                
            st.markdown(f"**Type:** {coll_info.get('file_type', 'Unknown')}")
            st.markdown(f"**Added:** {coll_info.get('added_date', 'Unknown')}")
            
            # Show description if available
            if "description" in coll_info and coll_info["description"]:
                st.markdown(f"**Description:** {coll_info['description']}")
            
            # Inspect button
            if st.button("Inspect Collection Data", key="inspect_from_dropdown"):
                inspect_collection(selected_collection)
    
    return

def process_text_input(input_text, collection_name, collection_description=None):
    """
    Process text input from a textbox and insert it into a ChromaDB collection.
    
    Args:
        input_text (str): The text entered by the user
        collection_name (str): Name of the collection to insert the text into
        collection_description (str, optional): Description of the collection
    
    Returns:
        bool: True if successful, False otherwise
    """
    if not init_loader():
        st.error("Document loader could not be initialized. Cannot process text.")
        return False
    
    logger.info(f"Processing direct text input for collection: {collection_name}")
    
    # Get current database
    current_db = st.session_state.current_database
    
    # Get existing collections for the current database
    existing_collections = list(st.session_state.database_collections.get(current_db, {}).keys())
    
    # Add timestamp to make each text entry unique
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    temp_filename = f"text_input_{timestamp}.txt"
    
    try:
        # Get user-specific storage path
        user_storage_path = get_user_storage_path(st.session_state.username)
        current_db_path = os.path.join(user_storage_path, current_db)
        
        # Format the text with proper paragraph breaks for better processing
        formatted_text = input_text.replace('\n\n', '\n-----\n').replace('\n', ' ').replace('-----', '\n\n')
        
        # Create a temporary file to save the text
        with tempfile.NamedTemporaryFile(mode="w+", suffix=".txt", delete=False) as temp_file:
            temp_file.write(formatted_text)
            temp_file_path = temp_file.name
        
        # Log details for debugging
        logger.info(f"Created temporary file for text: {temp_file_path}")
        
        # Set persistent directory for the loader
        st.session_state.loader.persist_directory = current_db_path
        
        # Try to create or insert the collection
        try:
            # Notify user about processing
            st.info("Processing text... This may take a moment.")
            
            # Process the text file
            success = st.session_state.loader.create_or_insert_collection(
                temp_file_path,
                collection_name,
                extension="txt",
                parser_type='local_parser',
                extract_figures=st.session_state.get('extract_figures', False),
                extract_tables=st.session_state.get('extract_tables', True)
            )
            
            if not success:
                st.error("Failed to process text. No chunks could be generated.")
                return False
                
        except Exception as embedding_error:
            logger.error(f"Error generating embeddings: {str(embedding_error)}")
            st.error(f"Failed to generate embeddings for the text: {str(embedding_error)}")
            return False
        
        # Store the collection in session state
        if "file_collections" not in st.session_state:
            st.session_state.file_collections = {}
        
        # Use the collection name as the key
        st.session_state.file_collections[collection_name] = collection_name
        
        # Store in database collections
        if "database_collections" not in st.session_state:
            st.session_state.database_collections = {}
            
        if current_db not in st.session_state.database_collections:
            st.session_state.database_collections[current_db] = {}
        
        # Create or update collection metadata
        if collection_name in st.session_state.database_collections[current_db]:
            # Update existing collection metadata
            collection_metadata = st.session_state.database_collections[current_db][collection_name]
            # Update the modified date
            collection_metadata["modified_date"] = time.strftime("%Y-%m-%d %H:%M:%S")
            # Track number of entries
            collection_metadata["entries"] = collection_metadata.get("entries", 1) + 1
        else:
            # Create new collection metadata
            collection_metadata = {
                "collection_name": collection_name,
                "file_name": temp_filename,
                "file_type": "txt",
                "added_date": time.strftime("%Y-%m-%d %H:%M:%S"),
                "entries": 1
            }
            
            # Add description if provided
            if collection_description:
                collection_metadata["description"] = collection_description
        
        # Store collection metadata
        st.session_state.database_collections[current_db][collection_name] = collection_metadata
        
        logger.info(f"Text processed successfully and added to collection: {collection_name}")
        return True
        
    except Exception as e:
        logger.error(f"Error processing text input: {str(e)}\n{traceback.format_exc()}")
        st.error(f"Error processing text: {str(e)}")
        return False
    finally:
        # Clean up the temporary file
        try:
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
        except Exception as cleanup_error:
            logger.warning(f"Failed to clean up temporary file: {str(cleanup_error)}")

if __name__ == "__main__":
    main()