app.py

import gradio as gr
import gc
import os
import time
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Force CPU only
import uuid
import threading
import pandas as pd
import torch
from langchain.document_loaders import CSVLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import HuggingFacePipeline
from langchain.chains import LLMChain
from transformers import AutoTokenizer, AutoModelForCausalLM, T5Tokenizer, T5ForConditionalGeneration, BitsAndBytesConfig, pipeline
from langchain.prompts import PromptTemplate
from llama_cpp import Llama
import re
import datetime
import warnings
warnings.filterwarnings('ignore')

# Global model cache
MODEL_CACHE = {
    "model": None,
    "tokenizer": None,
    "init_lock": threading.Lock(),
    "model_name": None
}

# Create directories for user data
os.makedirs("user_data", exist_ok=True)
os.makedirs("performance_metrics", exist_ok=True)

# Model configuration dictionary
MODEL_CONFIG = {
    "Llama 2 Chat GGUF": {
        "name": "TheBloke/Llama-2-7B-Chat-GGUF",
        "description": "Llama 2 7B Chat model with good general performance",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "TinyLlama Chat GGUF": {
        "name": "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF",
        "description": "Lightweight model with 1.1B parameters, fast and efficient",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "Mistral Instruct GGUF": {
        "name": "TheBloke/Mistral-7B-Instruct-v0.2-GGUF",
        "description": "7B instruction-tuned model with excellent reasoning",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "DeepSeek Coder Instruct": {
        "name": "deepseek-ai/deepseek-coder-1.3b-instruct",
        "description": "1.3B model for code and data analysis",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "Qwen2.5 Coder Instruct GGUF": {
        "name": "Qwen/Qwen2.5-Coder-3B-Instruct-GGUF",
        "description": "3B model specialized for code and technical applications",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "DeepSeek Distill Qwen": {
        "name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
        "description": "1.5B distilled model with good balance of speed and quality",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "Flan T5 Small": {
        "name": "google/flan-t5-small",
        "description": "Lightweight T5 model optimized for instruction following",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
        "is_t5": True
    }
}

# Performance metrics tracking
class PerformanceTracker:
    def __init__(self):
        self.metrics_file = "performance_metrics/model_performance.csv"
        
        # Create metrics file if it doesn't exist
        if not os.path.exists(self.metrics_file):
            with open(self.metrics_file, "w") as f:
                f.write("timestamp,model,question,processing_time,response_length\n")
    
    def log_performance(self, model_name, question, processing_time, response):
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        response_length = len(response)
        
        with open(self.metrics_file, "a") as f:
            f.write(f'"{timestamp}","{model_name}","{question}",{processing_time},{response_length}\n')
        
        print(f"Logged performance for {model_name}: {processing_time:.2f}s")

# Initialize performance tracker
performance_tracker = PerformanceTracker()

def initialize_model_once(model_key):
    with MODEL_CACHE["init_lock"]:
        try:
            current_model = MODEL_CACHE["model_name"]
            if MODEL_CACHE["model"] is None or current_model != model_key:
                # Clear previous model
                if MODEL_CACHE["model"] is not None:
                    del MODEL_CACHE["model"]
                    if MODEL_CACHE["tokenizer"] is not None:
                        del MODEL_CACHE["tokenizer"]
                    
                # Force garbage collection
                gc.collect()
                torch.cuda.empty_cache() if torch.cuda.is_available() else None
                time.sleep(1)  # Give system time to release memory
                
                model_info = MODEL_CONFIG[model_key]
                model_name = model_info["name"]
                MODEL_CACHE["model_name"] = model_key

                print(f"Loading model: {model_name}")
                
                # Check if this is a GGUF model
                if "GGUF" in model_name:
                    # Download the model file first if it doesn't exist
                    from huggingface_hub import hf_hub_download
                    try:
                        # Try to find the GGUF file in the repo
                        repo_id = model_name
                        model_path = hf_hub_download(
                            repo_id=repo_id, 
                            filename="model.gguf"  # File name may differ
                        )
                    except Exception as e:
                        print(f"Couldn't find model.gguf, trying other filenames: {str(e)}")
                        # Try to find GGUF file with other names
                        import requests
                        from huggingface_hub import list_repo_files
                        
                        files = list_repo_files(repo_id)
                        gguf_files = [f for f in files if f.endswith('.gguf')]
                        
                        if not gguf_files:
                            raise ValueError(f"No GGUF files found in {repo_id}")
                        
                        # Use first GGUF file found
                        model_path = hf_hub_download(repo_id=repo_id, filename=gguf_files[0])
                    
                    # Load GGUF model with llama-cpp-python
                    MODEL_CACHE["model"] = Llama(
                        model_path=model_path,
                        n_ctx=2048,  # Smaller context for memory savings
                        n_batch=512,
                        n_threads=2  # Adjust for 2 vCPU
                    )
                    MODEL_CACHE["tokenizer"] = None  # GGUF doesn't need separate tokenizer
                    MODEL_CACHE["is_gguf"] = True
                
                # Handle T5 models
                elif model_info.get("is_t5", False):
                    MODEL_CACHE["tokenizer"] = T5Tokenizer.from_pretrained(model_name)
                    MODEL_CACHE["model"] = T5ForConditionalGeneration.from_pretrained(
                        model_name,
                        torch_dtype=model_info["dtype"],
                        device_map="auto" if torch.cuda.is_available() else None,
                        low_cpu_mem_usage=True
                    )
                    MODEL_CACHE["is_gguf"] = False
                
                # Handle standard HF models
                else:
                    # Only use quantization if CUDA is available
                    if torch.cuda.is_available():
                        quantization_config = BitsAndBytesConfig(
                            load_in_4bit=True,
                            bnb_4bit_compute_dtype=torch.float16,
                            bnb_4bit_quant_type="nf4",
                            bnb_4bit_use_double_quant=True
                        )
                        
                        MODEL_CACHE["tokenizer"] = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
                        MODEL_CACHE["model"] = AutoModelForCausalLM.from_pretrained(
                            model_name,
                            quantization_config=quantization_config,
                            torch_dtype=model_info["dtype"],
                            device_map="auto",
                            low_cpu_mem_usage=True,
                            trust_remote_code=True
                        )
                    else:
                        # For CPU-only environments, load without quantization
                        MODEL_CACHE["tokenizer"] = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
                        MODEL_CACHE["model"] = AutoModelForCausalLM.from_pretrained(
                            model_name,
                            torch_dtype=torch.float32,  # Use float32 for CPU
                            device_map=None,
                            low_cpu_mem_usage=True,
                            trust_remote_code=True
                        )
                    MODEL_CACHE["is_gguf"] = False
                    
                print(f"Model {model_name} loaded successfully")
                
                # Final verification that model loaded correctly
                if MODEL_CACHE["model"] is None:
                    print(f"WARNING: Model {model_name} appears to be None after loading")
                    # Try to free memory before returning
                    torch.cuda.empty_cache() if torch.cuda.is_available() else None
                    gc.collect()
            
        except Exception as e:
            # Reset model cache on error
            MODEL_CACHE["model"] = None
            MODEL_CACHE["tokenizer"] = None
            # Force garbage collection
            gc.collect()
            torch.cuda.empty_cache() if torch.cuda.is_available() else None
            import traceback
            print(f"Error loading model {model_key}: {str(e)}")
            print(traceback.format_exc())
            raise RuntimeError(f"Failed to load model {model_key}: {str(e)}")
                
    return MODEL_CACHE["tokenizer"], MODEL_CACHE["model"], MODEL_CACHE.get("is_gguf", False)

def get_fallback_model(current_model):
    """Get appropriate fallback model for problematic models"""
    fallback_map = {
        "Flan T5 Small": "Llama 2 Chat GGUF"
    }
    return fallback_map.get(current_model, "Llama 2 Chat GGUF")

# Optimized pipeline for models
def create_optimized_pipeline(model, tokenizer, model_key):
    """Optimized pipeline for models"""
    # Default pipeline for other models
    pipe = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer,
            max_new_tokens=256,
            temperature=0.3, 
            top_p=0.9,
            top_k=30,
            repetition_penalty=1.2,
            return_full_text=False,
    )
    return HuggingFacePipeline(pipeline=pipe)

def create_llm_pipeline(model_key):
    """Create a new pipeline using the specified model with better error handling"""
    try:
        print(f"Creating pipeline for model: {model_key}")
        fallback_model = get_fallback_model(model_key)  # Define fallback_model at the beginning
        tokenizer, model, is_gguf = initialize_model_once(model_key)

        # Additional check to ensure model was properly loaded
        if model is None:
            print(f"Model is None for {model_key}, falling back to alternate model")
            if fallback_model != model_key:
                print(f"Attempting to use fallback model: {fallback_model}")
                tokenizer, model, is_gguf = initialize_model_once(fallback_model)
                if model is None:
                    raise ValueError(f"Both original and fallback models failed to load")
            else:
                raise ValueError(f"Model is None and no fallback available")
        
        # Get the model info for reference
        model_info = MODEL_CONFIG.get(model_key, MODEL_CONFIG.get(fallback_model, {}))
        
        # For GGUF models from llama-cpp-python
        if is_gguf:
            # Create adapter to use GGUF model like HF pipeline
            from langchain.llms import LlamaCpp
            llm = LlamaCpp(
                model_path=model.model_path,
                temperature=0.3,
                max_tokens=256,  # Increased for more comprehensive answers
                top_p=0.9,
                n_ctx=2048,
                streaming=False
            )
            return llm
            
        # Create appropriate pipeline for HF models
        elif model_info.get("is_t5", False):
            print("Creating T5 pipeline")
            pipe = pipeline(
                "text2text-generation",
                model=model,
                tokenizer=tokenizer,
                max_new_tokens=256,  # Increased for more comprehensive answers
                temperature=0.3,
                top_p=0.9,
                # Remove return_full_text parameter for T5 models
            )
        else:
            # Use optimized pipeline for problematic model
            return create_optimized_pipeline(model, tokenizer, model_key)
        
        print("Pipeline created successfully")
        return HuggingFacePipeline(pipeline=pipe)
    except Exception as e:
        import traceback
        print(f"Error creating pipeline: {str(e)}")
        print(traceback.format_exc())
        raise RuntimeError(f"Failed to create pipeline: {str(e)}")

# add a reset function to clear models between sessions
def reset_model_cache():
    """Force clear all model cache"""
    with MODEL_CACHE["init_lock"]:
        if MODEL_CACHE["model"] is not None:
            del MODEL_CACHE["model"]
        if MODEL_CACHE["tokenizer"] is not None:
            del MODEL_CACHE["tokenizer"]
        MODEL_CACHE["model"] = None
        MODEL_CACHE["tokenizer"] = None
        MODEL_CACHE["model_name"] = None
        MODEL_CACHE["is_gguf"] = False
        gc.collect()
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        time.sleep(1)

# Modified handle_model_loading_error function
def handle_model_loading_error(model_key, session_id):
    """Handle model loading errors by providing alternative model suggestions or fallbacks"""
    # Get the appropriate fallback model
    fallback_model = get_fallback_model(model_key)
    
    # Try to load the fallback model automatically
    if fallback_model != model_key:
        print(f"Automatically trying fallback model: {fallback_model} for {model_key}")
        
        try:
            # Try to initialize the fallback model
            tokenizer, model, is_gguf = initialize_model_once(fallback_model)
            return tokenizer, model, is_gguf, f"Model {model_key} couldn't be loaded. Automatically switched to {fallback_model}."
        except Exception as e:
            print(f"Fallback model {fallback_model} also failed: {str(e)}")
            # If fallback fails, continue with regular suggestion logic
    
    # Regular suggestion logic for when fallbacks don't work or aren't applicable
    suggested_models = [
        "DeepSeek Coder Instruct",  # 1.3B model
        "TinyLlama Chat GGUF",           # 1.1B model 
        "Qwen2.5 Coder Instruct"    # Another option
    ]
    
    # Remove problematic models and current model from suggestions
    problem_models = ["Flan T5 Small"]
    suggested_models = [m for m in suggested_models if m not in problem_models and m != model_key]
    
    suggestions = ", ".join(suggested_models[:3])  # Only show top 3 suggestions
    return None, None, None, f"Unable to load model {model_key}. Please try another model such as: {suggestions}"

def create_conversational_chain(db, file_path, model_key):
    llm = create_llm_pipeline(model_key)
    
    # Load the file into pandas to get metadata about the CSV
    df = pd.read_csv(file_path)
    
    # Create improved prompt template that focuses on pure LLM analysis
    template = """
    You are an expert data analyst tasked with answering questions about a CSV file.

    CSV File Structure:
    - Total rows: {row_count}
    - Total columns: {column_count}
    - Columns names: {columns_list}

    Sample data (first few rows):
    {sample_data}

    Additional context from the document:
    {context}

    User Question: {question}

    Give only the direct answer to the question. Be extremely brief. Do not explain your reasoning, calculations, or codes.

    Your answer:
    """
    
    PROMPT = PromptTemplate(
        template=template,
        input_variables=["row_count", "column_count", "columns_list", "sample_data", "context", "question"]
    )
    
    # Create retriever
    retriever = db.as_retriever(search_kwargs={"k": 5})  # Increase k for better context
    
    # Process query with better error handling
    def process_query(query, chat_history):
        try:
            start_time = time.time()
            
            # Get information from dataframe for context
            columns_list = ", ".join(df.columns.tolist())
            sample_data = df.head(5).to_string()  # Show 5 rows for better context
            row_count = len(df)
            column_count = len(df.columns)
            
            # Get context from vector database
            docs = retriever.get_relevant_documents(query)
            context = "\n\n".join([doc.page_content for doc in docs])
            
            # Run the chain
            chain = LLMChain(llm=llm, prompt=PROMPT)
            raw_result = chain.run(
                row_count=row_count,
                column_count=column_count,
                columns_list=columns_list,
                sample_data=sample_data,
                context=context,
                question=query
            )
            
            # Clean the result
            cleaned_result = raw_result.strip()

            # Add special handling for T5 models
            if MODEL_CONFIG.get(model_key, {}).get("is_t5", False):
                # T5 models sometimes return lists instead of strings
                if isinstance(raw_result, list) and len(raw_result) > 0:
                    if isinstance(raw_result[0], dict) and "generated_text" in raw_result[0]:
                        raw_result = raw_result[0]["generated_text"]
                    else:
                        raw_result = str(raw_result[0])
            
            # If result is empty after cleaning, use a fallback
            if not cleaned_result:
                cleaned_result = "I couldn't process a complete answer to your question. Please try asking in a different way or provide more specific details about what you'd like to know about the data."
            
            processing_time = time.time() - start_time
            
            # Log performance metrics
            performance_tracker.log_performance(
                model_key, 
                query, 
                processing_time, 
                cleaned_result
            )
            
            # Add processing time to the response for comparison purposes
            result_with_metrics = f"{cleaned_result}\n\n[Processing time: {processing_time:.2f} seconds]"
                
            return {"answer": result_with_metrics}
            
        except Exception as e:
            import traceback
            print(f"Error in process_query: {str(e)}")
            print(traceback.format_exc())
            return {"answer": f"An error occurred while processing your question: {str(e)}"}
    
    return process_query

class ChatBot:
    def __init__(self, session_id, model_key="DeepSeek Coder Instruct"):
        self.session_id = session_id
        self.chat_history = []
        self.chain = None
        self.user_dir = f"user_data/{session_id}"
        self.csv_file_path = None
        self.model_key = model_key
        os.makedirs(self.user_dir, exist_ok=True)

    def process_file(self, file, model_key=None):
        if model_key:
            self.model_key = model_key
        
        if file is None:
            return "Please upload a CSV file first."

        try:
            start_time = time.time()
            print(f"Processing file using model: {self.model_key}")
            # Handle file from Gradio
            file_path = file.name if hasattr(file, 'name') else str(file)
            self.csv_file_path = file_path
            print(f"CSV file path: {file_path}")

            # Copy to user directory
            user_file_path = f"{self.user_dir}/uploaded.csv"

            # Verify the CSV can be loaded
            try:
                df = pd.read_csv(file_path)
                print(f"CSV verified: {df.shape[0]} rows, {len(df.columns)} columns")

                # Save a copy in user directory
                df.to_csv(user_file_path, index=False)
                self.csv_file_path = user_file_path
                print(f"CSV saved to {user_file_path}")
            except Exception as e:
                print(f"Error reading CSV: {str(e)}")
                return f"Error reading CSV: {str(e)}"

            # Load document with reduced chunk size for better memory usage
            try:
                loader = CSVLoader(file_path=user_file_path, encoding="utf-8", csv_args={
                    'delimiter': ','})
                data = loader.load()
                print(f"Documents loaded: {len(data)}")
            except Exception as e:
                print(f"Error loading documents: {str(e)}")
                return f"Error loading documents: {str(e)}"

            # Create vector database with optimized settings
            try:
                db_path = f"{self.user_dir}/db_faiss"
            
                # Use CPU-friendly embeddings with smaller dimensions
                embeddings = HuggingFaceEmbeddings(
                    model_name='sentence-transformers/all-MiniLM-L6-v2',
                    model_kwargs={'device': 'cpu'}
                )

                db = FAISS.from_documents(data, embeddings)
                db.save_local(db_path)
                print(f"Vector database created at {db_path}")
            except Exception as e:
                print(f"Error creating vector database: {str(e)}")
                return f"Error creating vector database: {str(e)}"

            # Create custom chain
            try:
                print(f"Creating conversation chain with model: {self.model_key}")
                self.chain = create_conversational_chain(db, self.csv_file_path, self.model_key)
                print("Chain created successfully")
            except Exception as e:
                import traceback
                print(f"Error creating chain: {str(e)}")
                print(traceback.format_exc())
                return f"Error creating chain: {str(e)}"

            # Add basic file info to chat history for context
            file_processing_time = time.time() - start_time
            file_info = f"CSV successfully loaded with {df.shape[0]} rows and {len(df.columns)} columns using model {self.model_key}. Columns: {', '.join(df.columns.tolist())}"
            self.chat_history.append(("System", file_info))

            return f"CSV file successfully processed with model {self.model_key}! You can now chat with the model to analyze the data.\n\n[Processing time: {file_processing_time:.2f} seconds]"
        except Exception as e:
            import traceback
            print(traceback.format_exc())
            return f"File processing error: {str(e)}"

    def change_model(self, model_key):
        """Change the model being used and recreate the chain if necessary"""
        try:
            if model_key == self.model_key:
                return f"Model {model_key} is already in use."
            
            print(f"Changing model from {self.model_key} to {model_key}")
            self.model_key = model_key
        
            # If we have an active session with a file already loaded, recreate the chain
            if self.csv_file_path and os.path.exists(self.csv_file_path):
                try:
                    # Load existing database
                    db_path = f"{self.user_dir}/db_faiss"
                    if not os.path.exists(db_path):
                        return f"Error: Database not found. Please upload the CSV file again."
                
                    print(f"Loading embeddings from {db_path}")
                    embeddings = HuggingFaceEmbeddings(
                        model_name='sentence-transformers/all-MiniLM-L6-v2',
                        model_kwargs={'device': 'cpu'}
                    )
                
                    # Add allow_dangerous_deserialization=True flag
                    db = FAISS.load_local(db_path, embeddings, allow_dangerous_deserialization=True)
                    print(f"FAISS database loaded successfully")
                
                    # Create new chain with the selected model
                    print(f"Creating new conversation chain with {model_key}")
                    self.chain = create_conversational_chain(db, self.csv_file_path, self.model_key)
                    print(f"Chain created successfully")
                
                    # Add notification to chat history
                    self.chat_history.append(("System", f"Model successfully changed to {model_key}."))
                
                    return f"Model successfully changed to {model_key}."
                except Exception as e:
                    import traceback
                    error_trace = traceback.format_exc()
                    print(f"Detailed error in change_model: {error_trace}")
                    return f"Error changing model: {str(e)}"
            else:
                # Just update the model key if no file is loaded yet
                print(f"No CSV file loaded yet, just updating model preference to {model_key}")
                return f"Model changed to {model_key}. Please upload a CSV file to begin."
        except Exception as e:
            import traceback
            error_trace = traceback.format_exc()
            print(f"Unexpected error in change_model: {error_trace}")
            return f"Unexpected error while changing model: {str(e)}"

    def chat(self, message, history):
        if self.chain is None:
            return "Please upload a CSV file first."

        try:
            # Process the question with the chain
            result = self.chain(message, self.chat_history)

            # Get the answer with fallback
            answer = result.get("answer", "Sorry, I couldn't generate an answer. Please try asking a different question.")

            # Ensure we never return empty
            if not answer or answer.strip() == "":
                answer = "Sorry, I couldn't generate an appropriate answer. Please try asking the question differently."
            
            # Update internal chat history
            self.chat_history.append((message, answer))

            # Return just the answer for Gradio
            return answer
        except Exception as e:
            import traceback
            print(traceback.format_exc())
            return f"Error: {str(e)}"

# UI Code
def create_gradio_interface():
    with gr.Blocks(title="CSVBot - Chat with CSV using AI Models") as interface:
        session_id = gr.State(lambda: str(uuid.uuid4()))
        chatbot_state = gr.State(lambda: None)
        model_selected = gr.State(lambda: False)  # Track if model is already in use
        
        # Get model choices
        model_choices = list(MODEL_CONFIG.keys())
        default_model = "DeepSeek Distill Qwen"  # Default model

        gr.HTML("<h1 style='text-align: center;'>CSVBot - Chat with CSV using AI Models</h1>")
        gr.HTML("<h3 style='text-align: center;'>AI Assistant to Analyze your CSV file for all your needs</h3>")

        with gr.Row():
            with gr.Column(scale=1):
                with gr.Group():
                    gr.Markdown("### Step 1: Choose AI Model")
                    model_dropdown = gr.Dropdown(
                        label="Model",
                        choices=model_choices,
                        value=default_model,
                        interactive=True
                    )
                    model_info = gr.Markdown(
                        value=f"**{default_model}**: {MODEL_CONFIG[default_model]['description']}"
                    )
                
                with gr.Group():
                    gr.Markdown("### Step 2: Upload and Process CSV")
                    file_input = gr.File(
                        label="Upload CSV Anda",
                        file_types=[".csv"]
                    )
                    process_button = gr.Button("Process CSV")
                
                reset_button = gr.Button("Reset Session (To Change Model)")

            with gr.Column(scale=2):
                chatbot_interface = gr.Chatbot(
                    label="Chat History",
                    # type="messages",
                    height=400
                )
                message_input = gr.Textbox(
                    label="Type your message",
                    placeholder="Ask questions about your CSV data...",
                    lines=2
                )
                submit_button = gr.Button("Send")
                clear_button = gr.Button("Clear Chat")

        # Update model info when selection changes
        def update_model_info(model_key):
            return f"**{model_key}**: {MODEL_CONFIG[model_key]['description']}"
            
        model_dropdown.change(
            fn=update_model_info,
            inputs=[model_dropdown],
            outputs=[model_info]
        )
        
        # Modified handle_process_file function
        def handle_process_file(file, model_key, sess_id):
            """Process uploaded file with fallback model handling"""
            if file is None:
                return None, None, False, "Please upload a CSV file first."
        
            try:
                chatbot = ChatBot(sess_id, model_key)
                result = chatbot.process_file(file)
                return chatbot, True, [(None, result)]
            except Exception as e:
                import traceback
                print(f"Error processing file with {model_key}: {str(e)}")
                print(traceback.format_exc())

                # Try with fallback model if original fails
                fallback = get_fallback_model(model_key)
                if fallback != model_key:
                    try:
                        print(f"Trying fallback model: {fallback}")
                        chatbot = ChatBot(sess_id, fallback)
                        result = chatbot.process_file(file)
                        message = f"Original model {model_key} failed. Using {fallback} instead.\n\n{result}"
                        return chatbot, True, [(None, message)]
                    except Exception as fallback_error:
                        print(f"Fallback model also failed: {str(fallback_error)}")
                        
                error_msg = f"Error with model {model_key}: {str(e)}\n\nPlease try another model."
                return None, False, [(None, error_msg)]

        process_button.click(
            fn=handle_process_file,
            inputs=[file_input, model_dropdown, session_id],
            outputs=[chatbot_state, model_selected, chatbot_interface]
        ).then(
            # Disable model dropdown after processing file
            fn=lambda selected: gr.update(interactive=not selected),
            inputs=[model_selected],
            outputs=[model_dropdown]
        )

        # Reset handler - enables model selection again
        def reset_session():
            reset_model_cache() # call reset model cache
            return None, False, [], gr.update(interactive=True)
            
        reset_button.click(
            fn=reset_session,
            inputs=[],
            outputs=[chatbot_state, model_selected, chatbot_interface, model_dropdown]
        )

        # Chat handlers
        def user_message_submitted(message, history, chatbot, sess_id):
            history = history + [(message, None)]
            return history, "", chatbot, sess_id

        def bot_response(history, chatbot, sess_id):
            if chatbot is None:
                chatbot = ChatBot(sess_id)
                history[-1] = (history[-1][0], "Please upload a CSV file first.")
                return chatbot, history

            user_message = history[-1][0]
            response = chatbot.chat(user_message, history[:-1])
            history[-1] = (user_message, response)
            return chatbot, history

        submit_button.click(
            fn=user_message_submitted,
            inputs=[message_input, chatbot_interface, chatbot_state, session_id],
            outputs=[chatbot_interface, message_input, chatbot_state, session_id]
        ).then(
            fn=bot_response,
            inputs=[chatbot_interface, chatbot_state, session_id],
            outputs=[chatbot_state, chatbot_interface]
        )

        message_input.submit(
            fn=user_message_submitted,
            inputs=[message_input, chatbot_interface, chatbot_state, session_id],
            outputs=[chatbot_interface, message_input, chatbot_state, session_id]
        ).then(
            fn=bot_response,
            inputs=[chatbot_interface, chatbot_state, session_id],
            outputs=[chatbot_state, chatbot_interface]
        )

        # Clear chat handler
        def handle_clear_chat(chatbot):
            if chatbot is not None:
                chatbot.chat_history = []
            return chatbot, []

        clear_button.click(
            fn=handle_clear_chat,
            inputs=[chatbot_state],
            outputs=[chatbot_state, chatbot_interface]
        )

    return interface

# Launch the interface
demo = create_gradio_interface()
demo.launch(share=True)