app.py

import os
import streamlit as st
from openai import OpenAI
from PyPDF2 import PdfReader
import requests
from youtube_transcript_api import YouTubeTranscriptApi, NoTranscriptFound, TranscriptsDisabled
from urllib.parse import urlparse, parse_qs
from pinecone import Pinecone
import uuid
from dotenv import load_dotenv
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from itertools import islice
import unicodedata
import re
from tiktoken import encoding_for_model
import json
import datetime
import tiktoken
import pandas as pd
import io
from fpdf import FPDF
import tempfile
from PyPDF2 import PdfReader
import base64
from pathlib import Path
import numpy as np
from pymongo import MongoClient
import traceback
from docx import Document
import pandas as pd
import io
import time
import traceback

load_dotenv()

# Set up OpenAI client
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Set up Pinecone
pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))

# Get index name and URL from .env
index_name = os.getenv("PINECONE_INDEX_NAME")
index_url = os.getenv("PINECONE_INDEX_URL")

index = pc.Index(index_name, url=index_url)

# Set up MongoDB connection
mongo_client = MongoClient(os.getenv("MONGODB_URI"))
db = mongo_client.get_database("finance")
users_collection = db["users"]

def get_embedding(text):
    response = client.embeddings.create(input=text, model="text-embedding-3-large")
    return response.data[0].embedding

def chunk_text(text, content_type):
    sanitized_text = sanitize_text(text)
    if content_type == "YouTube":
        chunk_size = 2000  # Adjust this value as needed
        content_length = len(sanitized_text)
        return [sanitized_text[i:i+chunk_size] for i in range(0, content_length, chunk_size)]
    else:  # Default for PDF and Web Link
        chunk_size = 2000
        content_length = len(sanitized_text)
        return [sanitized_text[i:i+chunk_size] for i in range(0, content_length, chunk_size)]

def process_pdf(file):
    reader = PdfReader(file)
    text = []
    for page in reader.pages:
        text.append(page.extract_text())
    return " ".join(text)  # Join all pages into a single string

def process_web_link(url):
    response = requests.get(url)
    return chunk_text(response.text, "Web")

def process_youtube_link(url):
    try:
        video_id = extract_video_id(url)
        transcript = YouTubeTranscriptApi.get_transcript(video_id)
        full_text = " ".join([entry['text'] for entry in transcript])
        print("Transcript obtained from YouTube API")
        return chunk_text(full_text, "YouTube")  # Use chunk_text function to split large transcripts
    except NoTranscriptFound:
        print("No transcript found for this YouTube video.")
        return []
    except TranscriptsDisabled:
        print("Transcripts are disabled for this YouTube video.")
        return []
    except Exception as e:
        print(f"An error occurred while processing the YouTube link: {str(e)}")
        return []

def extract_video_id(url):
    parsed_url = urlparse(url)
    if parsed_url.hostname == 'youtu.be':
        return parsed_url.path[1:]
    if parsed_url.hostname in ('www.youtube.com', 'youtube.com'):
        if parsed_url.path == '/watch':
            return parse_qs(parsed_url.query)['v'][0]
        if parsed_url.path[:7] == '/embed/':
            return parsed_url.path.split('/')[2]
        if parsed_url.path[:3] == '/v/':
            return parsed_url.path.split('/')[2]
    return None

def process_upload(upload_type, file_or_link, file_name=None):
    print(f"Starting process_upload for {upload_type}")
    doc_id = str(uuid.uuid4())
    print(f"Generated doc_id: {doc_id}")

    if upload_type == "PDF":
        chunks = process_pdf(file_or_link)
        doc_name = file_name or "Uploaded PDF"
    elif upload_type == "Web Link":
        chunks = process_web_link(file_or_link)
        doc_name = file_or_link
    elif upload_type == "YouTube Link":
        chunks = process_youtube_link(file_or_link)
        doc_name = f"YouTube: {file_or_link}"
    else:
        print("Invalid upload type")
        return "Invalid upload type"

    vectors = []
    with ThreadPoolExecutor() as executor:
        futures = [executor.submit(process_chunk, chunk, doc_id, i, upload_type, doc_name) for i, chunk in enumerate(chunks)]
        
        for future in as_completed(futures):
            vectors.append(future.result())
            # Update progress
            progress = len(vectors) / len(chunks)
            st.session_state.upload_progress.progress(progress)
    
    print(f"Generated {len(vectors)} vectors")
    
    # Upsert vectors in batches
    batch_size = 100  # Adjust this value as needed
    for i in range(0, len(vectors), batch_size):
        batch = list(islice(vectors, i, i + batch_size))
        index.upsert(vectors=batch)
        print(f"Upserted batch {i//batch_size + 1} of {len(vectors)//batch_size + 1}")
    
    print("All vectors upserted to Pinecone")
    
    return f"Processing complete for {upload_type}. Document Name: {doc_name}"

def process_chunk(chunk, doc_id, i, upload_type, doc_name):
    # Sanitize the chunk text
    sanitized_chunk = sanitize_text(chunk)
    embedding = get_embedding(sanitized_chunk)
    return (f"{doc_id}_{i}", embedding, {
        "text": sanitized_chunk, 
        "type": upload_type,
        "doc_id": doc_id,
        "doc_name": doc_name,
        "chunk_index": i
    })

def sanitize_text(text):
    # Remove control characters and normalize Unicode
    return ''.join(ch for ch in unicodedata.normalize('NFKD', text) if unicodedata.category(ch)[0] != 'C')

def get_relevant_context(query, top_k=5):
    print(f"Getting relevant context for query: {query}")
    query_embedding = get_embedding(query)
    
    search_results = index.query(vector=query_embedding, top_k=top_k, include_metadata=True)
    print(f"Found {len(search_results['matches'])} relevant results")
    
    # Sort results by similarity score (higher is better)
    sorted_results = sorted(search_results['matches'], key=lambda x: x['score'], reverse=True)
    
    context = "\n".join([result['metadata']['text'] for result in sorted_results])
    return context, sorted_results

def truncate_context(context, max_tokens):
    enc = encoding_for_model("gpt-4o-mini")
    encoded = enc.encode(context)
    if len(encoded) > max_tokens:
        return enc.decode(encoded[:max_tokens])
    return context

def chat_with_ai(message):
    print(f"Chatting with AI, message: {message}")
    context, results = get_relevant_context(message)
    print(f"Retrieved context, length: {len(context)}")
    
    # Truncate context if it's too long
    max_tokens = 7000  # Leave some room for the system message and user query
    context = truncate_context(context, max_tokens)
    
    messages = [
        {"role": "system", "content": "You are a helpful assistant. Use the following information to answer the user's question, but don't mention the context directly in your response. If the information isn't in the context, say you don't know."},
        {"role": "system", "content": f"Context: {context}"},
        {"role": "user", "content": message}
    ]
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages
    )
    print("Received response from OpenAI")
    
    ai_response = response.choices[0].message.content
    
    # Prepare source information
    sources = [
        {
            "doc_id": result['metadata']['doc_id'],
            "doc_name": result['metadata']['doc_name'],
            "chunk_index": result['metadata']['chunk_index'],
            "text": result['metadata']['text'],
            "type": result['metadata']['type'],
            "score": result['score']
        }
        for result in results
    ]
    
    return ai_response, sources

def process_youtube_links(links):
    results = []
    for link in links:
        result = process_upload("YouTube Link", link.strip())
        results.append(result)
    return results

def process_excel(file):
    dfs = pd.read_excel(file, sheet_name=None)  # Read all sheets
    for sheet_name, df in dfs.items():
        df.columns = df.columns.astype(str)
        # Remove any unnamed columns
        df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
    return dfs

def analyze_and_generate_formulas(main_df, other_dfs):
    # Focus on the 'DETAILS' column and the month columns
    details_column = main_df.columns[0]
    month_columns = main_df.columns[1:-1]  # Exclude the 'Total' column
    
    main_summary = f"DETAILS column data: {main_df[details_column].tolist()}\n"
    for month in month_columns:
        main_summary += f"{month} column data: {main_df[month].tolist()}\n"
    
    other_sheets_summary = ""
    for name, df in other_dfs.items():
        if len(df.columns) > 0:
            other_sheets_summary += f"\nSheet '{name}' structure:\n"
            other_sheets_summary += f"Columns: {df.columns.tolist()}\n"
            other_sheets_summary += f"First few rows:\n{df.head().to_string()}\n"
    
    prompt = f"""Analyze the following Excel data and generate Python formulas to fill missing values:

Main Sheet Structure:
{main_summary}

Other Sheets:
{other_sheets_summary}

Provide Python formulas using pandas to fill missing values in the month columns of the main sheet.
Use pandas and numpy functions where appropriate. If a value cannot be determined, use None.
Return a dictionary with column names as keys and formulas as values.
Example of the expected format: {{'Apr'24': "df['Apr'24'].fillna(method='ffill')"}}
"""

    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "You are a data analysis expert skilled in creating concise formulas for data filling."},
                {"role": "user", "content": prompt}
            ]
        )

        formulas = eval(response.choices[0].message.content.strip())
        print(f"Generated formulas: {formulas}")
        
        return formulas

    except Exception as e:
        print(f"Error in analyze_and_generate_formulas: {str(e)}")
        print(traceback.format_exc())
        return {}

def apply_formulas(main_df, other_dfs, formulas):
    filled_df = main_df.copy()
    
    for column, formula in formulas.items():
        if column in filled_df.columns:
            try:
                print(f"Applying formula for column {column}: {formula}")
                
                # Create a local namespace with all dataframes
                namespace = {'df': filled_df, 'np': np, 'pd': pd}
                
                # Execute the formula in the namespace
                exec(f"result = {formula}", namespace)
                
                # Apply the result to the column
                filled_df[column] = namespace['result']
                print(f"Successfully applied formula for column {column}")
            except Exception as e:
                print(f"Error applying formula for column {column}: {str(e)}")
                print(traceback.format_exc())
    
    return filled_df

def excel_to_pdf(df):
    pdf = FPDF(orientation='L', unit='mm', format='A4')
    pdf.add_page()
    pdf.set_font("Arial", size=8)
    
    # Calculate column widths
    col_widths = [pdf.get_string_width(str(col)) + 6 for col in df.columns]
    
    # Write header
    for i, col in enumerate(df.columns):
        pdf.cell(col_widths[i], 10, str(col), border=1)
    pdf.ln()
    
    # Write data
    for _, row in df.iterrows():
        for i, value in enumerate(row):
            pdf.cell(col_widths[i], 10, str(value), border=1)
        pdf.ln()
    
    # Save to a temporary file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
        pdf.output(temp_file.name)
        return temp_file.name

def pdf_to_text(pdf_path):
    with open(pdf_path, 'rb') as file:
        pdf = PdfReader(file)
        text = []
        for page in pdf.pages:
            text.append(page.extract_text())
    return text

def get_user_feedback(user_id):
    user = users_collection.find_one({"_id": user_id})
    return user.get("feedback", "") if user else ""

def get_category_reports():
    return {
        "Default": [],  # Changed from "None" to "Default"
        "Sales KPI": [
            "Monthwise Sales Table", "Customer-wise Sales Table (top 10)", "Qty Sales",
            "Customer-wise Churn", "Avg Sales per Customer", "Product-wise Sales Rate",
            "Geography-wise", "Trend Analysis", "Graphs", "Month-wise Comparison"
        ],
        "Expenses KPI": [
            "Vendor-wise Comparison", "Year on Year Monthwise", "Division-wise"
        ],
        "Purchase Register": [
            "Vendor-wise Monthwise", "Monthwise", "Material-wise Purchase Rate Analysis"
        ],
        "Balance Sheet": [
            "Year on Year Comparison", "Ratios"
        ]
    }

def analyze_excel_with_gpt(df, sheet_name, user_feedback, category, reports_needed, use_assistants_api=False):
    if use_assistants_api:
        return process_excel_with_assistant(df, category, reports_needed, user_feedback)
    else:
        # Existing OCR-based analysis code
        prompt = f"""Analyze the following Excel data from sheet '{sheet_name}':

        {df.to_string()}

        User's previous feedback and insights:
        {user_feedback}

        """

        if category != "general":
            prompt += f"""Please provide analysis and insights based on the following required reports for the category '{category}':
            {', '.join(reports_needed)}

            Please provide:
            1. A comprehensive overview of the data focusing on the {category} category
            2. Key observations and trends related to the required reports
            3. Any anomalies, interesting patterns, or correlations relevant to the {category}
            4. Suggestions for further analysis or visualization based on the required reports
            5. Address any previous feedback or insights mentioned above, if applicable

            Focus on providing a thorough analysis of all aspects of the data relevant to the {category} and the specified reports."""
        else:
            prompt += """Please provide a general analysis of the data, including:
            1. A comprehensive overview of the data
            2. Key observations and trends
            3. Any anomalies, interesting patterns, or correlations
            4. Suggestions for further analysis or visualization
            5. Address any previous feedback or insights mentioned above, if applicable

            Focus on providing a thorough analysis of all aspects of the data."""

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": f"You are a data analyst expert in interpreting Excel data for {'general' if category == 'general' else category} analysis."},
                {"role": "user", "content": prompt}
            ]
        )
        
        return response.choices[0].message.content

def analyze_document_with_gpt(document_content, user_feedback, category, reports_needed, use_assistants_api=False, file_id=None):
    if use_assistants_api:
        assistant = client.beta.assistants.create(
            name="Document Analyzer",
            instructions=f"You are a document analysis expert. Analyze the uploaded document and provide insights based on the category: {category}.",
            model="gpt-4-1106-preview"
        )
        
        thread = client.beta.threads.create()
        
        message = client.beta.threads.messages.create(
            thread_id=thread.id,
            role="user",
            content=f"Analyze the document with file ID: {file_id}. Category: {category}. Required reports: {', '.join(reports_needed)}. User feedback: {user_feedback}",
            file_ids=[file_id]
        )
        
        run = client.beta.threads.runs.create(
            thread_id=thread.id,
            assistant_id=assistant.id
        )
        
        while run.status != "completed":
            run = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
            time.sleep(1)
        
        messages = client.beta.threads.messages.list(thread_id=thread.id)
        return messages.data[0].content[0].text.value
    else:
        # Existing OCR-based analysis code
        prompt = f"""Analyze the following document content:

        {document_content}

        User's previous feedback and insights:
        {user_feedback}

        """

        if category != "general":
            prompt += f"""Please provide analysis and insights based on the following required reports for the category '{category}':
            {', '.join(reports_needed)}

            Please provide:
            1. A comprehensive overview of the content focusing on the {category} category
            2. Key points and main ideas related to the required reports
            3. Any interesting patterns or unique aspects relevant to the {category}
            4. Suggestions for further analysis or insights based on the required reports
            5. Any limitations of the analysis due to the document format or OCR process
            6. Address any previous feedback or insights mentioned above, if applicable

            Focus on providing a thorough analysis of all aspects of the content relevant to the {category} and the specified reports."""
        else:
            prompt += """Please provide a general analysis of the document content, including:
            1. A comprehensive overview of the content
            2. Key points and main ideas
            3. Any interesting patterns or unique aspects
            4. Suggestions for further analysis or insights
            5. Any limitations of the analysis due to the document format or OCR process
            6. Address any previous feedback or insights mentioned above, if applicable

            Focus on providing a thorough analysis of all aspects of the content."""

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": f"You are a data analyst expert in interpreting complex document content for {'general' if category == 'general' else category} analysis."},
                {"role": "user", "content": prompt}
            ]
        )
        
        return response.choices[0].message.content

def process_uploaded_file(uploaded_file):
    file_type = uploaded_file.type
    if file_type in ["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "application/vnd.ms-excel"]:
        # Process Excel file
        dfs = process_excel(uploaded_file)
        return "excel", dfs
    elif file_type == "application/pdf":
        # Process PDF file using PyPDF2
        try:
            pdf_reader = PdfReader(uploaded_file)
            text = ""
            for page in pdf_reader.pages:
                text += page.extract_text()
            return "text", text
        except Exception as e:
            st.error(f"Error processing PDF: {str(e)}")
            print(f"Error processing PDF: {str(e)}")
            print(traceback.format_exc())
            return None, None
    elif file_type in ["application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/msword"]:
        # Process Word document
        try:
            doc = Document(io.BytesIO(uploaded_file.read()))
            text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
            return "text", text
        except Exception as e:
            st.error(f"Error processing Word document: {str(e)}")
            print(f"Error processing Word document: {str(e)}")
            print(traceback.format_exc())
            return None, None
    else:
        st.error(f"Unsupported file type: {file_type}. Please upload an Excel file, PDF, or Word document.")
        return None, None

def chat_with_data(data, user_question, data_type):
    if data_type == "excel":
        df_string = data.to_string()
        data_description = "Excel sheet data"
    else:  # PDF or Word document
        df_string = data
        data_description = "document content"
    
    prompt = f"""You are an AI assistant specialized in analyzing {data_description}. You have access to the following data:

{df_string}

Based on this data, please answer the following question:
{user_question}

Provide a detailed and accurate answer based on the given data. If the answer cannot be directly inferred from the data, provide the best possible response based on the available information and your general knowledge about data analysis."""

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": f"You are a data analysis expert skilled in interpreting {data_description}."},
            {"role": "user", "content": prompt}
        ]
    )
    
    return response.choices[0].message.content

def extract_challan_data(pdf_text):
    data = {}
    patterns = {
        'ITNS No.': r'ITNS No\.\s*:\s*(\d+)',
        'TAN': r'TAN\s*:\s*(\w+)',
        'Name': r'Name\s*:\s*(.+)',
        'Assessment Year': r'Assessment Year\s*:\s*(\d{4}-\d{2})',
        'Financial Year': r'Financial Year\s*:\s*(\d{4}-\d{2})',
        'Amount': r'Amount \(in Rs\.\)\s*:\s*₹\s*([\d,]+)',
        'CIN': r'CIN\s*:\s*(\w+)',
        'Date of Deposit': r'Date of Deposit\s*:\s*(\d{2}-\w{3}-\d{4})',
        'Challan No': r'Challan No\s*:\s*(\d+)',
    }
    
    for key, pattern in patterns.items():
        match = re.search(pattern, pdf_text)
        if match:
            data[key] = match.group(1)
        else:
            data[key] = 'N/A'
    
    return data

def process_challan_pdfs(pdf_files):
    all_data = []
    for pdf_file in pdf_files:
        pdf_text = process_pdf(pdf_file)
        challan_data = extract_challan_data(pdf_text)
        all_data.append(challan_data)
    
    df = pd.DataFrame(all_data)
    return df

def process_file_with_assistant(file, file_type, category, reports_needed, user_feedback):
    print(f"Starting {file_type} processing with Assistant")
    try:
        # Upload the file to OpenAI
        uploaded_file = client.files.create(
            file=file,
            purpose='assistants'
        )
        print(f"File uploaded successfully. File ID: {uploaded_file.id}")

        # Create an assistant
        assistant = client.beta.assistants.create(
            name=f"{file_type} Analyzer",
            instructions=f"You are an expert in analyzing {file_type} files, focusing on {category}. Provide insights and summaries of the content based on the following reports: {', '.join(reports_needed)}. Consider the user's previous feedback: {user_feedback}",
            model="gpt-4o",
            tools=[{"type": "file_search"}]
        )
        print(f"Assistant created. Assistant ID: {assistant.id}")

        # Create a thread
        thread = client.beta.threads.create()
        print(f"Thread created. Thread ID: {thread.id}")

        # Add a message to the thread with the file attachment
        message = client.beta.threads.messages.create(
            thread_id=thread.id,
            role="user",
            content=f"Please analyze this  file and provide insights for the {category} category, focusing on the following reports: {', '.join(reports_needed)}.",
            attachments=[
                {"file_id": uploaded_file.id, "tools": [{"type": "file_search"}]}
            ]
        )
        print(f"Message added to thread. Message ID: {message.id}")

        # Run the assistant
        run = client.beta.threads.runs.create(
            thread_id=thread.id,
            assistant_id=assistant.id
        )
        print(f"Run created. Run ID: {run.id}")

        # Wait for the run to complete
        while run.status != 'completed':
            run = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
            print(f"Run status: {run.status}")
            time.sleep(1)

        # Retrieve the messages
        messages = client.beta.threads.messages.list(thread_id=thread.id)
        
        # Extract the assistant's response
        analysis_result = next((msg.content[0].text.value for msg in messages if msg.role == 'assistant'), None)
        
        print(f"{file_type} analysis completed successfully")
        return analysis_result

    except Exception as e:
        print(f"Error in process_file_with_assistant: {str(e)}")
        print(traceback.format_exc())
        return None

# Streamlit UI
st.set_page_config(layout="wide")
st.title("Document Processing, Chat, Excel Filling, and Analysis")

# Add login/signup system
if "user" not in st.session_state:
    st.session_state.user = None

def login(username, password):
    user = users_collection.find_one({"username": username})
    if user and user.get("password") == password:
        return user
    return None

def signup(username, password):
    existing_user = users_collection.find_one({"username": username})
    if existing_user:
        return False
    users_collection.insert_one({
        "username": username, 
        "password": password,
        "feedback": []  # Initialize an empty list for feedback
    })
    return True

def store_feedback(username, feedback):
    users_collection.update_one(
        {"username": username},
        {"$push": {"feedback": feedback}},
        upsert=True
    )

# Login/Signup form
if not st.session_state.user:
    tab1, tab2 = st.tabs(["Login", "Sign Up"])
    
    with tab1:
        st.subheader("Login")
        login_username = st.text_input("Username", key="login_username")
        login_password = st.text_input("Password", type="password", key="login_password")
        if st.button("Login"):
            user = login(login_username, login_password)
            if user:
                st.session_state.user = user
                st.success("Logged in successfully!")
                st.rerun()
            else:
                st.error("Invalid username or password")
    
    with tab2:
        st.subheader("Sign Up")
        signup_username = st.text_input("Username", key="signup_username")
        signup_password = st.text_input("Password", type="password", key="signup_password")
        if st.button("Sign Up"):
            if signup(signup_username, signup_password):
                st.success("Account created successfully! Please log in.")
            else:
                st.error("Username already exists")
 
if st.session_state.user:
    st.write(f"Welcome, {st.session_state.user['username']}!")
    
    # Create four tabs
    tab1, tab2, tab3, tab4 = st.tabs(["Upload, Chat, and Source", "Excel Processing", "Excel Analysis and Chat", "Challan Processing"])

    with tab1:
        st.subheader("Upload")
        
        # PDF upload
        uploaded_files = st.file_uploader("Choose one or more PDF files", type="pdf", accept_multiple_files=True)
        
        # Web Link input
        web_link = st.text_input("Enter a Web Link")
        
        # YouTube Links input
        youtube_links = st.text_area("Enter YouTube Links (one per line)")
        
        if st.button("Process All"):
            st.session_state.upload_progress = st.progress(0)
            with st.spinner("Processing uploads..."):
                results = []
                if uploaded_files:
                    for file in uploaded_files:
                        pdf_result = process_upload("PDF", file, file.name)
                        results.append(pdf_result)
                if web_link:
                    web_result = process_upload("Web Link", web_link)
                    results.append(web_result)
                if youtube_links:
                    youtube_links_list = re.split(r'[\n\r]+', youtube_links.strip())
                    youtube_results = process_youtube_links(youtube_links_list)
                    results.extend(youtube_results)
                
                if results:
                    for result in results:
                        st.success(result)
                else:
                    st.warning("No content uploaded. Please provide at least one input.")
            st.session_state.upload_progress.empty()

        st.subheader("Chat")
        user_input = st.text_input("Ask a question about the uploaded content:")
        if st.button("Send"):
            if user_input:
                print(f"Sending user input: {user_input}")
                st.session_state.chat_progress = st.progress(0)
                response, sources = chat_with_ai(user_input)
                st.session_state.chat_progress.progress(1.0)
                st.markdown("**You:** " + user_input)
                st.markdown("**AI:** " + response)
                
                # Store sources in session state for display in the Source Chunks section
                st.session_state.sources = sources
                st.session_state.chat_progress.empty()
            else:
                print("Empty user input")
                st.warning("Please enter a question.")

        st.subheader("Source Chunks")
        if 'sources' in st.session_state and st.session_state.sources:
            for i, source in enumerate(st.session_state.sources, 1):
                with st.expander(f"Source {i} - {source['type']} ({source['doc_name']}) - Score: {source['score']}"):
                    st.markdown(f"**Chunk Index:** {source['chunk_index']}")
                    st.text(source['text'])
        else:
            st.info("Ask a question to see source chunks here.")

    with tab2:
        st.subheader("Excel Processing")
        uploaded_excel = st.file_uploader("Choose an Excel file", type=["xlsx", "xls"])
        
        if uploaded_excel is not None:
            dfs = process_excel(uploaded_excel)
            
            # Display preview of each sheet
            for sheet_name, df in dfs.items():
                if not df.empty:
                    st.write(f"Preview of {sheet_name}:")
                    st.dataframe(df.head())
            
            # Select the main sheet for processing
            main_sheet = st.selectbox("Select the main sheet to fill", list(dfs.keys()))
            main_df = dfs[main_sheet]
            
            if st.button("Fill Missing Data"):
                with st.spinner("Analyzing data and generating formulas..."):
                    other_dfs = {name: df for name, df in dfs.items() if name != main_sheet}
                    formulas = analyze_and_generate_formulas(main_df, other_dfs)
                    
                    if formulas:
                        st.write("Generated Formulas:")
                        for column, formula in formulas.items():
                            st.code(f"{column}: {formula}")
                        
                        filled_df = apply_formulas(main_df, other_dfs, formulas)
                        
                        st.write("Filled Excel Data:")
                        st.dataframe(filled_df)
                        
                        # Provide download link for the filled Excel file
                        buffer = io.BytesIO()
                        with pd.ExcelWriter(buffer, engine='xlsxwriter') as writer:
                            filled_df.to_excel(writer, index=False, sheet_name=main_sheet)
                            # Also save other sheets
                            for sheet_name, df in dfs.items():
                                if sheet_name != main_sheet:
                                    df.to_excel(writer, index=False, sheet_name=sheet_name)
                        buffer.seek(0)
                        st.download_button(
                            label="Download Filled Excel",
                            data=buffer,
                            file_name="filled_excel.xlsx",
                            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                        )
                    else:
                        st.warning("No formulas were generated. The data might not have clear patterns for filling missing values.")

    with tab3:
        st.subheader("Excel and Document Analysis")
        uploaded_file = st.file_uploader(
            "Choose an Excel file, PDF, or Word document for analysis",
            type=["xlsx", "xls", "pdf", "docx", "doc"],
            key="excel_analysis_uploader"
        )
        
        if uploaded_file is not None:
            file_type, content = process_uploaded_file(uploaded_file)
            
            if file_type is not None and content is not None:
                if file_type == "excel":
                    dfs = content
                    sheet_names = list(dfs.keys())
                    selected_sheet = st.selectbox("Select a sheet for analysis", sheet_names)
                    
                    df_to_analyze = dfs[selected_sheet]
                    st.write(f"Preview of {selected_sheet}:")
                    st.dataframe(df_to_analyze.head())
                    
                    st.session_state.analyzed_data = df_to_analyze
                    analysis_method = "OCR"  # Default to OCR for Excel files
                elif file_type == "text":
                    st.write("Document content preview:")
                    preview_text = content[:500] + "..."
                    st.text(preview_text)  # Show first 500 characters
                    
                    # Add option to choose between OCR and Assistants API for PDF/Word
                    analysis_method = st.radio("Choose analysis method:", ("OCR", "OpenAI Assistants API"))
                    
                    st.session_state.analyzed_data = content
                
                # Add category selection with "Default" option
                categories = list(get_category_reports().keys())
                if "Default" in categories:
                    categories.remove("Default")
                categories = ["Default"] + categories
                selected_category = st.selectbox("Select analysis category", categories)
                
                if st.button("Analyze with GPT"):
                    with st.spinner("Analyzing data... This may take a while for large datasets."):
                        user_feedback = get_user_feedback(st.session_state.user["_id"])
                        reports_needed = get_category_reports().get(selected_category, [])
                        
                        if file_type == "excel":
                            analysis_result = analyze_excel_with_gpt(st.session_state.analyzed_data, selected_sheet, user_feedback, selected_category, reports_needed)
                        else:  # PDF or Word document
                            if analysis_method == "OpenAI Assistants API":
                                analysis_result = process_file_with_assistant(uploaded_file, "PDF", selected_category, reports_needed, user_feedback)
                            else:
                                analysis_result = analyze_document_with_gpt(st.session_state.analyzed_data, user_feedback, selected_category, reports_needed)
                        
                        st.markdown("## Analysis Results")
                        st.markdown(analysis_result)
                        
                        st.session_state.analysis_result = analysis_result
                        
                        if file_type == "excel":
                            pdf_path = excel_to_pdf(st.session_state.analyzed_data)
                            with open(pdf_path, "rb") as pdf_file:
                                pdf_bytes = pdf_file.read()
                            st.download_button(
                                label="Download Excel PDF version",
                                data=pdf_bytes,
                                file_name="excel_data.pdf",
                                mime="application/pdf"
                            )
                
                # Feedback section
                st.markdown("## Feedback")
                new_feedback = st.text_area("Provide feedback or additional insights about the analysis:")
                if st.button("Submit Feedback"):
                    if new_feedback:
                        user = users_collection.find_one({"_id": st.session_state.user["_id"]})
                        existing_feedback = user.get("feedback", "")
                        
                        updated_feedback = f"{existing_feedback}\n{new_feedback}" if existing_feedback else new_feedback
                        
                        users_collection.update_one(
                            {"_id": st.session_state.user["_id"]},
                            {"$set": {"feedback": updated_feedback}}
                        )
                        st.success("Feedback submitted successfully!")
                    else:
                        st.warning("Please enter some feedback before submitting.")
                
                # Chat with Data section
                st.markdown("## Chat with Data")
                with st.form(key='chat_form'):
                    user_question = st.text_input("Ask a question about the data:")
                    chat_submit_button = st.form_submit_button(label='Get Answer')
                    
                    if chat_submit_button:
                        if user_question:
                            with st.spinner("Analyzing your question..."):
                                answer = chat_with_data(st.session_state.analyzed_data, user_question, file_type)
                                st.markdown("### Answer")
                                st.markdown(answer)
                        else:
                            st.warning("Please enter a question about the data.")
            else:
                st.error("Unable to process the uploaded file. Please check the file format and try again.")

    with tab4:
        st.subheader("Challan Processing")
        
        challan_pdfs = st.file_uploader(
            "Choose Challan PDF files",
            type="pdf",
            accept_multiple_files=True,
            key="challan_processing_uploader"
        )
        
        if st.button("Process Challan PDFs"):
            if challan_pdfs:
                with st.spinner("Processing Challan PDFs..."):
                    challan_df = process_challan_pdfs(challan_pdfs)
                    st.write("Challan Data:")
                    st.dataframe(challan_df)
                    
                    buffer = io.BytesIO()
                    with pd.ExcelWriter(buffer, engine='xlsxwriter') as writer:
                        challan_df.to_excel(writer, index=False, sheet_name='Challan Data')
                    buffer.seek(0)
                    st.download_button(
                        label="Download Challan Excel",
                        data=buffer,
                        file_name="challan_data.xlsx",
                        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                    )
                    
                    st.success("Challan PDFs processed successfully")
            else:
                st.warning("No Challan PDFs uploaded. Please choose at least one PDF file.")