app.py

import os
import tempfile
import gradio as gr
import openai
from typing import Optional, List
import hashlib
import base64
import json
import time
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize OpenAI client with error handling
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
    raise ValueError("OPENAI_API_KEY environment variable is not set")

# Initialize OpenAI client with older API syntax
openai.api_key = api_key

# Custom CSS for a beautiful, modern look
custom_css = """
html, body, .gradio-container {
    height: 100vh !important;
    min-height: 100vh !important;
    max-width: 100vw !important;
    margin: 0 !important;
    padding: 0 !important;
    font-family: 'Inter', 'Segoe UI', Arial, sans-serif;
    background: #f4f7fb;
    color: #222;
}

.centered-main {
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: flex-start;
    min-height: 100vh;
    width: 100vw;
    padding-top: 32px;
}

.compact-box {
    background: #fff;
    border-radius: 18px;
    box-shadow: 0 4px 24px rgba(0, 60, 180, 0.07), 0 1.5px 4px rgba(0,0,0,0.04);
    padding: 32px 32px 20px 32px;
    margin-bottom: 32px;
    width: 100%;
    max-width: 600px;
    margin-left: auto;
    margin-right: auto;
    border: 1.5px solid #e3e8f0;
}

.section-title {
    font-size: 1.25rem;
    font-weight: 700;
    margin-bottom: 18px;
    color: #1a237e;
    letter-spacing: 0.01em;
}

.upload-btn, .send-btn, .audio-btn, .reset-btn {
    background: linear-gradient(135deg, #1976D2 0%, #00bcd4 100%);
    color: white;
    border: none;
    padding: 12px 28px;
    border-radius: 24px;
    cursor: pointer;
    font-weight: 600;
    font-size: 16px;
    margin-top: 10px;
    margin-bottom: 10px;
    transition: all 0.2s;
    box-shadow: 0 2px 8px rgba(25, 118, 210, 0.08);
}
.upload-btn:hover, .send-btn:hover, .audio-btn:hover, .reset-btn:hover {
    background: linear-gradient(135deg, #00bcd4 0%, #1976D2 100%);
    box-shadow: 0 4px 16px rgba(0, 188, 212, 0.13);
}

.gradio-chatbot {
    border-radius: 14px !important;
    border: 1.5px solid #e3e8f0 !important;
    background: #f8fafc !important;
    padding: 12px !important;
    min-height: 350px !important;
    max-height: 400px !important;
    overflow-y: auto !important;
    margin-bottom: 10px;
}

.gradio-audio {
    margin-top: 12px;
    margin-bottom: 12px;
}

.textbox {
    border-radius: 12px !important;
    border: 1.5px solid #e3e8f0 !important;
    padding: 12px !important;
    font-size: 16px !important;
    margin-bottom: 10px;
    background: #f8fafc !important;
    color: #222 !important;
}
.textbox:focus {
    border-color: #1976D2 !important;
    box-shadow: 0 0 0 2px rgba(25, 118, 210, 0.13) !important;
}

.status-text {
    color: #1976D2;
    font-size: 15px;
    margin-top: 10px;
    font-weight: 500;
    background: #e3f2fd;
    border-radius: 8px;
    padding: 8px 12px;
}

/* File upload area */
input[type="file"]::-webkit-file-upload-button {
    background: #1976D2;
    color: #fff;
    border: none;
    border-radius: 8px;
    padding: 8px 18px;
    font-weight: 600;
    cursor: pointer;
}
input[type="file"]::-webkit-file-upload-button:hover {
    background: #00bcd4;
}

/* Only one main scroll */
body, .gradio-container, #root, #app {
    overflow: auto !important;
    height: 100vh !important;
}
#component-0, #component-1, #component-2, .chatbot, .chat-container {
    overflow: visible !important;
    height: auto !important;
    max-height: none !important;
}
"""

# Custom audio recorder component with improved styling
def create_audio_recorder():
    return gr.HTML("""
        <div class="audio-recorder">
            <button id="recordButton" class="record-button">
                <span class="record-icon">🎤</span>
                <span class="record-text">Start Recording</span>
            </button>
            <div id="recordingStatus" class="status-text"></div>
            <audio id="audioPlayback" controls style="display: none; margin-top: 10px;"></audio>
        </div>
        <script>
            let mediaRecorder;
            let audioChunks = [];
            let isRecording = false;
            const recordButton = document.getElementById('recordButton');
            const recordingStatus = document.getElementById('recordingStatus');
            const audioPlayback = document.getElementById('audioPlayback');

            recordButton.addEventListener('click', async () => {
                if (!isRecording) {
                    try {
                        const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
                        mediaRecorder = new MediaRecorder(stream);
                        audioChunks = [];
                        
                        mediaRecorder.ondataavailable = (event) => {
                            audioChunks.push(event.data);
                        };
                        
                        mediaRecorder.onstop = () => {
                            const audioBlob = new Blob(audioChunks, { type: 'audio/wav' });
                            const audioUrl = URL.createObjectURL(audioBlob);
                            audioPlayback.src = audioUrl;
                            audioPlayback.style.display = 'block';
                            
                            const reader = new FileReader();
                            reader.readAsDataURL(audioBlob);
                            reader.onloadend = () => {
                                const base64Audio = reader.result;
                                window.parent.postMessage({
                                    type: 'audio_data',
                                    data: base64Audio
                                }, '*');
                            };
                        };
                        
                        mediaRecorder.start();
                        isRecording = true;
                        recordButton.classList.add('recording');
                        recordButton.querySelector('.record-text').textContent = 'Stop Recording';
                        recordingStatus.textContent = 'Recording...';
                    } catch (err) {
                        console.error('Error accessing microphone:', err);
                        recordingStatus.textContent = 'Error accessing microphone';
                    }
                } else {
                    mediaRecorder.stop();
                    isRecording = false;
                    recordButton.classList.remove('recording');
                    recordButton.querySelector('.record-text').textContent = 'Start Recording';
                    recordingStatus.textContent = 'Recording saved';
                }
            });
        </script>
    """)

class AdvancedRAG:
    def __init__(self):
        self.thread_id: Optional[str] = None
        self.assistant_id: Optional[str] = os.getenv("ASSISTANT_ID")
        if not self.assistant_id:
            raise ValueError("ASSISTANT_ID environment variable is not set")

    def create_thread(self) -> str:
        thread = openai.beta.threads.create()
        self.thread_id = thread.id
        return self.thread_id

    def upload_document(self, file) -> str:
        try:
            if file is None:
                raise ValueError("No file provided")

            # The file is now a bytes object directly
            file_bytes = file
            if not file_bytes:
                raise ValueError("Empty file")

            # Create a temporary file with a generic extension
            with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as tmp:
                tmp.write(file_bytes)
                tmp.flush()
                
                # Upload the file to OpenAI
                with open(tmp.name, "rb") as file_obj:
                    file_obj = openai.files.create(
                        file=file_obj,
                        purpose="assistants"
                    )

            # Create a new thread if one doesn't exist
            if not self.thread_id:
                self.create_thread()

            # Add the file to the thread with the required tools object
            openai.beta.threads.messages.create(
                thread_id=self.thread_id,
                role="user",
                content="I have uploaded a document. Please analyze it.",
                attachments=[{"file_id": file_obj.id, "tools": [{"type": "file_search"}]}]
            )

            return file_obj.id
        except Exception as e:
            raise Exception(f"Error uploading document: {str(e)}")
        finally:
            # Clean up temporary file
            if 'tmp' in locals():
                try:
                    os.unlink(tmp.name)
                except:
                    pass

    def ask_question(self, question: str) -> str:
        try:
            if not self.thread_id:
                self.create_thread()

            # Add the question to the thread
            openai.beta.threads.messages.create(
                thread_id=self.thread_id,
                role="user",
                content=question
            )
            
            # Create a run
            run = openai.beta.threads.runs.create(
                thread_id=self.thread_id,
                assistant_id=self.assistant_id
            )
            
            # Wait for the run to complete
            waited = 0
            while True:
                run_status = openai.beta.threads.runs.retrieve(
                    thread_id=self.thread_id,
                    run_id=run.id
                )
                if run_status.status == 'completed':
                    break
                elif run_status.status == 'failed':
                    raise Exception("Run failed")
                time.sleep(0.2)
                waited += 0.2
                if waited > 60:
                    raise Exception("Run timed out after 60 seconds.")

            # Get the latest message
            messages = openai.beta.threads.messages.list(
                thread_id=self.thread_id,
                order='desc',
                limit=1
            )
            if not messages.data:
                return "No response received from the assistant."
            return messages.data[0].content[0].text.value
        except Exception as e:
            return f"[Error: {str(e)}]"

    def transcribe_audio(self, audio_file):
        try:
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
                tmp.write(audio_file.read())
                tmp.flush()
                tmp_path = tmp.name
            with open(tmp_path, "rb") as audio:
                transcript = openai.audio.transcriptions.create(
                    model="whisper-1",
                    file=audio,
                    language="en"
                )
            os.remove(tmp_path)
            return transcript.text
        except Exception as e:
            return f"[Error transcribing audio: {str(e)}]"

# Initialize RAG system
rag = AdvancedRAG()

def process_file(file):
    if file is None:
        return "Please upload a file first."
    try:
        rag.upload_document(file)
        return "File uploaded successfully! You can now ask questions about the document."
    except Exception as e:
        return f"Error uploading file: {str(e)}"

def process_question(question, history):
    if not rag.thread_id:
        return "Please upload a document first.", history, "", None
    try:
        response = rag.ask_question(question)
        history.append({"role": "user", "content": question})
        history.append({"role": "assistant", "content": response})
        return "", history, "", None
    except Exception as e:
        return "", history + [{"role": "assistant", "content": f"Error: {str(e)}"}], "", None

def process_audio(audio_file, history):
    if audio_file is None:
        return "Please record or upload an audio file.", history, "", None
    try:
        transcript = rag.transcribe_audio(audio_file)
        if not rag.thread_id:
            return "Please upload a document first.", history, "", None
        response = rag.ask_question(transcript)
        history.append({"role": "user", "content": f"🎤 {transcript}"})
        history.append({"role": "assistant", "content": response})
        return "", history, "", None
    except Exception as e:
        history.append({"role": "user", "content": f"🎤 [Error transcribing audio: {str(e)}]"})
        history.append({"role": "assistant", "content": "It seems there was an error while transcribing audio due to a technical issue. If there's anything specific from the document or any other questions you have regarding the content, please let me know, and I can assist you with that information."})
        return "", history, "", None

def process_audio_base64(audio_base64, history):
    if not audio_base64:
        return "Please record an audio message first.", history
    try:
        # Convert base64 to audio file
        audio_data = base64.b64decode(audio_base64.split(',')[1])
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            tmp.write(audio_data)
            tmp.flush()
            tmp_path = tmp.name

        # Transcribe audio
        with open(tmp_path, "rb") as audio_file:
            transcript = openai.audio.transcriptions.create(
                model="whisper-1",
                file=audio_file,
                language="en"
            )
        os.remove(tmp_path)

        if not rag.thread_id:
            return "Please upload a document first.", history

        # Get response from assistant
        response = rag.ask_question(transcript.text)
        history.append((f"🎤 {transcript.text}", response))
        return "", history
    except Exception as e:
        return "", history + [("Audio input", f"Error: {str(e)}")]

# Create Gradio interface with improved layout
with gr.Blocks(css=custom_css, title="Document Q&A System") as demo:
    gr.Markdown("""
    # <span style='color:#1976D2;'>Document Q&A System</span>
    <div style='text-align:center; color:#1976D2; margin-bottom:18px;'>Upload a document, ask questions, or use your voice!</div>
    """)

    # Define shared components at the top
    chatbot = gr.Chatbot(height=350, elem_classes="gradio-chatbot", label=None, type="messages")
    file_output = gr.Textbox(label="Upload Status", interactive=False, elem_classes="textbox")
    question = gr.Textbox(label="Type your question and press Enter", placeholder="Ask a question about your document...", elem_classes="textbox")
    audio_input = gr.Audio(type="filepath", label="Record or Upload Audio", elem_classes="gradio-audio", visible=False)
    audio_status = gr.Textbox(label="Audio Status", interactive=False, elem_classes="status-text", visible=False)

    # Unified card for upload, chat, and voice
    with gr.Row(elem_classes="centered-main"):
        with gr.Column():
            with gr.Group(elem_classes="compact-box"):
                gr.Markdown("<div class='section-title'>Document Q&A</div>")
                file_input = gr.File(label="Upload Document", file_types=[".pdf", ".txt", ".doc", ".docx"], file_count="single", type="binary", elem_classes="upload-btn")
                reset_btn = gr.Button("Reset Chat & Upload New Document", elem_classes="reset-btn")
                file_input.change(process_file, file_input, file_output)
                def reset_all():
                    rag.thread_id = None
                    return "", [], "", None
                reset_btn.click(reset_all, None, [file_output, chatbot, question, audio_input])
                file_output.render()
                chatbot.render()
                with gr.Row():
                    question.render()
                    send_btn = gr.Button("Send", elem_classes="send-btn")
                    mic_btn = gr.Button("🎤", elem_classes="audio-btn")
                send_btn.click(process_question, [question, chatbot], [question, chatbot, question, audio_input])
                question.submit(process_question, [question, chatbot], [question, chatbot, question, audio_input])
                audio_input.change(process_audio, [audio_input, chatbot], [audio_input, chatbot, question, audio_input])
                def show_audio():
                    return {audio_input: gr.update(visible=True), audio_status: gr.update(visible=True)}
                mic_btn.click(show_audio, None, [audio_input, audio_status])

    # Add JavaScript for audio handling
    demo.load(
        fn=None,
        inputs=None,
        outputs=None,
        js="""
        function() {
            window.addEventListener('message', function(event) {
                if (event.data.type === 'audio_data') {
                    const audioData = event.data.data;
                    const byteString = atob(audioData.split(',')[1]);
                    const mimeString = audioData.split(',')[0].split(':')[1].split(';')[0];
                    const ab = new ArrayBuffer(byteString.length);
                    const ia = new Uint8Array(ab);
                    for (let i = 0; i < byteString.length; i++) {
                        ia[i] = byteString.charCodeAt(i);
                    }
                    const blob = new Blob([ab], {type: mimeString});
                    const file = new File([blob], "recording.wav", {type: mimeString});
                    
                    const audioInput = document.querySelector('input[type="file"]');
                    const dataTransfer = new DataTransfer();
                    dataTransfer.items.add(file);
                    audioInput.files = dataTransfer.files;
                    audioInput.dispatchEvent(new Event('change', { bubbles: true }));
                }
            });
        }
        """
    )

if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )