app.py

import os
import gradio as gr
from huggingface_hub import InferenceClient
import firebase_admin
from firebase_admin import credentials, firestore, auth
import uuid
import tempfile
import json

class XylariaChat:
    def __init__(self):
        # Securely load HuggingFace and Firebase tokens from environment variables
        self.hf_token = os.environ.get("HF_TOKEN")
        firebase_cred_json = os.environ.get("FIREBASE_SERVICE_ACCOUNT")
        
        # Optional: Add more robust error handling for missing environment variables
        if not self.hf_token:
            print("Warning: HuggingFace token not found. Some functionality may be limited.")
        
        # Firebase initialization
        self.db = None
        if firebase_cred_json:
            try:
                # Convert the JSON string to a temporary credentials file
                with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as temp_cred_file:
                    # Add error handling for JSON parsing
                    try:
                        firebase_creds = json.loads(firebase_cred_json)
                        json.dump(firebase_creds, temp_cred_file)
                        temp_cred_file.close()
                        
                        firebase_cred = credentials.Certificate(temp_cred_file.name)
                        firebase_admin.initialize_app(firebase_cred)
                        self.db = firestore.client()
                    except json.JSONDecodeError:
                        print("Error: Invalid Firebase credentials JSON")
                    
                # Remove the temporary credentials file
                os.unlink(temp_cred_file.name)
            except Exception as e:
                print(f"Firebase initialization error: {e}")
        
        # Initialize the inference client (only if token is available)
        self.client = None
        if self.hf_token:
            try:
                self.client = InferenceClient(
                    model="Qwen/QwQ-32B-Preview", 
                    api_key=self.hf_token
                )
            except Exception as e:
                print(f"Inference client initialization error: {e}")
        
        # Initialize conversation history
        self.conversation_history = []

    def signup(self, username, email, password):
        """User signup method"""
        if not self.db:
            return "Firebase is not initialized. Signup unavailable."
        
        try:
            # Create user in Firebase Authentication
            user = auth.create_user(
                email=email,
                password=password,
                display_name=username
            )
            
            # Store additional user info in Firestore
            user_ref = self.db.collection('users').document(user.uid)
            user_ref.set({
                'username': username,
                'email': email,
                'created_at': firestore.SERVER_TIMESTAMP
            })
            
            return f"Successfully created account for {username}"
        except Exception as e:
            return f"Signup error: {str(e)}"

    def login(self, email, password):
        """User login method"""
        if not self.db:
            return "Firebase is not initialized. Login unavailable."
        
        try:
            # Authenticate user with Firebase
            user = auth.get_user_by_email(email)
            
            # In a real-world scenario, you'd use Firebase Auth's sign-in method
            # This is a simplified version for demonstration
            return f"Login successful for {user.display_name}"
        except Exception as e:
            return f"Login error: {str(e)}"

    def save_message(self, user_id, message, sender):
        """Save message to Firestore"""
        if not self.db:
            return
        
        try:
            messages_ref = self.db.collection('conversations').document(user_id)
            messages_ref.collection('messages').add({
                'text': message,
                'sender': sender,
                'timestamp': firestore.SERVER_TIMESTAMP
            })
        except Exception as e:
            print(f"Error saving message: {e}")

    def get_response(self, user_input, chat_history):
        # Check if client is initialized
        if not self.client:
            return "AI response is currently unavailable."
        
        # Prepare messages with conversation context
        messages = [
            {"role": "system", "content": """You are Xylaria 1.4 Senoa, an AI assistant developed by SK MD Saad Amin. 
            - You're NOT MADE BY OPENAI OR ANY OTHER INSTITUTION
            - Be friendly and chatty, use emoji sometimes."""},
            *[{"role": "user" if msg[0] else "assistant", "content": msg[1]} for msg in chat_history]
        ]
        
        # Add current user input
        messages.append({"role": "user", "content": user_input})
        
        # Generate response with streaming
        try:
            stream = self.client.chat.completions.create(
                messages=messages,
                temperature=0.5,
                max_tokens=10240,
                top_p=0.7,
                stream=True
            )
            
            return stream
        
        except Exception as e:
            return f"Error generating response: {str(e)}"

    def create_interface(self):
        # Dynamic CSS for OS theme detection
        custom_css = """
        /* OS Theme Detection */
        @media (prefers-color-scheme: dark) {
            .gradio-container {
                background-color: #121212;
                color: #ffffff;
            }
        }
        
        @media (prefers-color-scheme: light) {
            .gradio-container {
                font-family: 'Inter', sans-serif;
                background-color: #f3f4f6;
                color: #1f2937;
            }
        }
        """

        def streaming_response(message, chat_history):
            # Unique user ID (in a real app, you'd use actual user authentication)
            user_id = str(uuid.uuid4())
            
            # Save user message to Firestore
            self.save_message(user_id, message, 'user')
            
            # Get AI response stream
            response_stream = self.get_response(message, chat_history)
            
            # If it's an error or string, return immediately
            if isinstance(response_stream, str):
                return "", chat_history + [[message, response_stream]]
            
            # Prepare for streaming response
            full_response = ""
            updated_history = chat_history + [[message, ""]]
            
            # Streaming output
            try:
                for chunk in response_stream:
                    if chunk.choices[0].delta.content:
                        chunk_content = chunk.choices[0].delta.content
                        full_response += chunk_content
                        
                        # Update the last message in chat history with partial response
                        updated_history[-1][1] = full_response
                        yield "", updated_history
            except Exception as e:
                return "", updated_history + [["", f"Error in response: {str(e)}"]]
            
            # Save bot message to Firestore
            self.save_message(user_id, full_response, 'bot')
            
            return "", updated_history

        with gr.Blocks(theme='soft', css=custom_css) as demo:
            # Authentication Tab
            with gr.Tabs():
                with gr.TabItem("Login"):
                    with gr.Column():
                        login_email = gr.Textbox(label="Email")
                        login_password = gr.Textbox(label="Password", type="password")
                        login_btn = gr.Button("Login")
                        login_output = gr.Textbox(label="Login Status")
                
                with gr.TabItem("Signup"):
                    with gr.Column():
                        signup_username = gr.Textbox(label="Username")
                        signup_email = gr.Textbox(label="Email")
                        signup_password = gr.Textbox(label="Password", type="password")
                        signup_btn = gr.Button("Sign Up")
                        signup_output = gr.Textbox(label="Signup Status")
            
            # Chat Interface
            with gr.Column(scale=1):
                chatbot = gr.Chatbot(
                    label="Xylaria 1.4 Senoa",
                    height=500,
                    show_copy_button=True,
                    type="messages"  # Added to resolve deprecation warning
                )
                
                # Input row with minimalist design
                with gr.Row():
                    txt = gr.Textbox(
                        show_label=False, 
                        placeholder="Type your message...", 
                        scale=4
                    )
                    btn = gr.Button("Send", scale=1)
                
                # Clear conversation button
                clear = gr.Button("Clear Conversation")
            
            # Authentication Event Handlers
            login_btn.click(
                fn=self.login, 
                inputs=[login_email, login_password], 
                outputs=[login_output]
            )
            
            signup_btn.click(
                fn=self.signup, 
                inputs=[signup_username, signup_email, signup_password], 
                outputs=[signup_output]
            )
            
            # Submit functionality with streaming
            btn.click(
                fn=streaming_response, 
                inputs=[txt, chatbot], 
                outputs=[txt, chatbot]
            )
            txt.submit(
                fn=streaming_response, 
                inputs=[txt, chatbot], 
                outputs=[txt, chatbot]
            )
            clear.click(lambda: None, None, chatbot)

        return demo

# Note: Actual launch should be done with proper environment variable setup
xylaria_chat = XylariaChat()
xylaria_chat.create_interface().launch(share=True)