Update app.py

app.py

@@ -18,9 +18,9 @@ logging.basicConfig(
 )
 logger = logging.getLogger("CompanyChatbot")
 
-# Environment variables (for production use)
+# Environment variables
 HF_MODEL = os.environ.get("HF_MODEL", "HuggingFaceH4/zephyr-7b-beta")
-HF_API_TOKEN = os.environ.get("HF_API_TOKEN", None)  # Set your API token as env variable
+HF_API_TOKEN = os.environ.get("HF_API_TOKEN", None)
 COMPANY_NAME = os.environ.get("COMPANY_NAME", "Your Company")
 DEFAULT_SYSTEM_PROMPT = os.environ.get("DEFAULT_SYSTEM_PROMPT", 
     f"You are {COMPANY_NAME}'s professional AI assistant. Be helpful, accurate, and concise.")
@@ -41,92 +41,81 @@ def save_conversation(user_id, conversation):
         json.dump(conversation, f)
     logger.info(f"Saved conversation for user {user_id}")
 
-# Main chat function - fixed to properly handle message format
-def respond(
-    message, 
-    history,
-    system_message, 
-    max_tokens, 
-    temperature, 
-    top_p,
-    user_id
-):
+def generate_response(message, history, system_message, max_tokens, temperature, top_p, user_id):
+    messages = [{"role": "system", "content": system_message}]
+    
+    # Add history to messages
+    for h in history:
+        messages.append({"role": "user", "content": h[0]})
+        messages.append({"role": "assistant", "content": h[1]})
+    
+    # Add current message
+    messages.append({"role": "user", "content": message})
+    
+    # Generate response
+    try:
+        response = client.text_generation(
+            prompt=message,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+        )
+        return response
+    except Exception as e:
+        logger.error(f"Error generating response: {str(e)}")
+        return f"An error occurred: {str(e)}"
+
+def chat(message, history, system_message, max_tokens, temperature, top_p, user_id):
     if not message.strip():
-        return history + [[message, "I'm sorry, I didn't receive any input. How can I help you today?"]]
+        return history
     
     logger.info(f"User {user_id} sent message - Length: {len(message)}")
     
     try:
-        messages = [{"role": "system", "content": system_message}]
-        
-        for user_msg, assistant_msg in history:
-            if user_msg:
-                messages.append({"role": "user", "content": user_msg})
-            if assistant_msg:
-                messages.append({"role": "assistant", "content": assistant_msg})
-        
-        messages.append({"role": "user", "content": message})
-        
-        full_response = ""
-        start_time = datetime.now()
-        
-        # Ensure proper response format for streaming
-        current_history = history.copy()
-        
-        for message_chunk in client.chat_completion(
-            messages, 
-            max_tokens=max_tokens, 
-            stream=True, 
-            temperature=temperature, 
-            top_p=top_p,
-        ):
-            token = message_chunk.choices[0].delta.content
-            if token:
-                full_response += token
-                # Always ensure a list of lists with pairs
-                new_history = current_history + [[message, full_response]]
-                yield new_history
+        response = generate_response(
+            message, 
+            history, 
+            system_message, 
+            max_tokens, 
+            temperature, 
+            top_p, 
+            user_id
+        )
         
-        time_taken = (datetime.now() - start_time).total_seconds()
-        logger.info(f"Response generated for user {user_id} in {time_taken:.2f}s - Length: {len(full_response)}")
+        # Format response
+        new_history = list(history or [])
+        new_history.append([message, str(response)])
         
+        # Save conversation
         conversation_data = {
             "timestamp": datetime.now().isoformat(),
             "user_id": user_id,
-            "messages": messages,
-            "response": full_response,
-            "parameters": {
-                "max_tokens": max_tokens,
-                "temperature": temperature,
-                "top_p": top_p
-            },
-            "time_taken": time_taken
+            "message": message,
+            "response": str(response),
+            "history": new_history
         }
         save_conversation(user_id, conversation_data)
         
-        # Ensure final response is properly formatted
-        if not full_response:
-            return history
-        return current_history + [[message, full_response]]
+        return new_history
         
     except Exception as e:
-        error_msg = f"An error occurred: {str(e)}"
-        logger.error(f"Error generating response for user {user_id}: {str(e)}")
-        return history + [[message, error_msg]]
+        logger.error(f"Error in chat function: {str(e)}")
+        new_history = list(history or [])
+        new_history.append([message, f"An error occurred: {str(e)}"])
+        return new_history
 
-# Authentication function (replace with your actual auth system)
+# Authentication
 def authenticate(username, password):
-    # In production, this should check against your company's auth system
     valid_credentials = {
         "admin": {"password": "admin123", "role": "admin"},
         "user": {"password": "user123", "role": "user"}
-    }  # Example only
+    }
     
     if username in valid_credentials and valid_credentials[username]["password"] == password:
-        return True, str(uuid.uuid4()), valid_credentials[username]["role"]  # Return success, user_id, and role
+        return True, str(uuid.uuid4()), valid_credentials[username]["role"]
     return False, None, None
 
-# Login interface with error handling
+# Login interface
 def login(username, password):
     if not username or not password:
         return (
@@ -137,10 +126,9 @@ def login(username, password):
             gr.update(visible=True, value="Please enter both username and password")
         )
     
-    # Simulate processing delay for security (prevents timing attacks)
     time.sleep(0.5)
-    
     success, user_id, role = authenticate(username, password)
+    
     if success:
         return (
             gr.update(visible=False),
@@ -158,16 +146,9 @@ def login(username, password):
             gr.update(visible=True, value="Invalid username or password")
         )
 
-# CSS for better styling
+# CSS
 css = """
-body {
-    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-    background-color: #f9f9f9;
-}
-.container {
-    max-width: 1400px !important;
-    margin: auto;
-}
+.container { max-width: 1400px !important; margin: auto; }
 .setting-panel {
     background-color: #f0f4f8;
     border-radius: 10px;
@@ -179,64 +160,12 @@ body {
     box-shadow: 0 2px 6px rgba(0,0,0,0.1);
     background-color: white;
 }
-.company-header {
-    background-color: #2c3e50;
-    color: white;
-    padding: 15px;
-    border-radius: 10px 10px 0 0;
-    margin-bottom: 15px;
-}
-.footer {
-    text-align: center;
-    margin-top: 20px;
-    color: #666;
-    font-size: 0.8em;
-}
-.message-user {
-    background-color: #e6f7ff !important;
-    border-radius: 15px 15px 0 15px !important;
-}
-.message-bot {
-    background-color: #f0f0f0 !important;  
-    border-radius: 15px 15px 15px 0 !important;
-}
-.login-container {
-    max-width: 500px;
-    margin: 50px auto;
-    padding: 30px;
-    background-color: white;
-    border-radius: 10px;
-    box-shadow: 0 4px 10px rgba(0,0,0,0.1);
-}
-.login-header {
-    text-align: center;
-    margin-bottom: 30px;
-}
-.error-message {
-    color: #e74c3c;
-    background-color: #fdedeb;
-    padding: 10px;
-    border-radius: 5px;
-    margin-bottom: 15px;
-    font-size: 14px;
-}
-.role-badge {
-    font-size: 12px;
-    padding: 3px 8px;
-    border-radius: 10px;
-    margin-left: 10px;
-}
-.admin-badge {
-    background-color: #e74c3c;
-    color: white;
-}
-.user-badge {
-    background-color: #3498db;
-    color: white;
-}
-.setting-disabled {
-    opacity: 0.5;
+.setting-disabled input,
+.setting-disabled select,
+.setting-disabled textarea,
+.setting-disabled button {
     pointer-events: none;
+    opacity: 0.6;
 }
 """
 
@@ -245,188 +174,116 @@ with gr.Blocks(css=css, title=f"{COMPANY_NAME} AI Assistant") as demo:
     user_id = gr.State(None)
     user_role = gr.State(None)
     
-    with gr.Row():
-        gr.Markdown(f"<div class='company-header'><h1>{COMPANY_NAME} AI Assistant</h1></div>", elem_classes=["company-header"])
-    
-    # Login Group
+    # Login Interface
     with gr.Group(visible=True) as login_group:
-        with gr.Column(elem_classes=["login-container"]):
-            gr.Markdown(f"<div class='login-header'><h2>Welcome to {COMPANY_NAME}</h2><p>Please log in to continue</p></div>")
-            
-            # Error message container (initially hidden)
-            error_message = gr.Markdown(visible=False, value="", elem_classes=["error-message"])
-            
-            username = gr.Textbox(label="Username", placeholder="Enter your username")
-            password = gr.Textbox(label="Password", type="password", placeholder="Enter your password")
-            
-            with gr.Row():
-                login_button = gr.Button("Login", variant="primary", size="lg")
+        username = gr.Textbox(label="Username")
+        password = gr.Textbox(label="Password", type="password")
+        login_button = gr.Button("Login")
+        error_message = gr.Markdown(visible=False)
     
-    # Chat Group (initially hidden)
+    # Chat Interface
     with gr.Group(visible=False) as chat_group:
         with gr.Row():
-            # Left side: Settings Panel
-            with gr.Column(scale=1, elem_classes=["setting-panel"]):
-                # Role indicator
-                role_indicator = gr.Markdown("", elem_id="role-indicator")
-                
-                gr.Markdown("### Configuration")
+            # Settings Panel
+            with gr.Column(scale=1):
+                role_display = gr.Markdown("")
                 system_message = gr.Textbox(
+                    label="System Message",
                     value=DEFAULT_SYSTEM_PROMPT,
-                    label="System Instructions",
-                    lines=4
+                    lines=3
                 )
                 
-                with gr.Group(elem_id="admin-settings") as admin_settings:
+                with gr.Group() as settings_group:
                     max_tokens = gr.Slider(
                         minimum=1,
                         maximum=2048,
                         value=512,
                         step=1,
-                        label="Max Response Length"
+                        label="Max Tokens"
                     )
                     temperature = gr.Slider(
                         minimum=0.1,
                         maximum=1.0,
                         value=0.7,
                         step=0.1,
-                        label="Temperature (Creativity)"
+                        label="Temperature"
                     )
                     top_p = gr.Slider(
                         minimum=0.1,
                         maximum=1.0,
                         value=0.95,
                         step=0.05,
-                        label="Top-p (Variation)"
+                        label="Top P"
                     )
-                    
-                    gr.Markdown("---")
-                
-                gr.Markdown("### Chat Actions")
-                with gr.Row():
-                    clear_btn = gr.Button("Clear Chat", variant="secondary")
-                    export_btn = gr.Button("Export Conversation", variant="secondary")
                 
-                # Logout button
-                logout_btn = gr.Button("Logout", variant="stop")
-                
-                gr.Markdown(f"<div class='footer'>© {datetime.now().year} {COMPANY_NAME}. All rights reserved.</div>")
+                clear_button = gr.Button("Clear Chat")
+                logout_button = gr.Button("Logout")
             
-            # Right side: Chat Interface  
-            with gr.Column(scale=2, elem_classes=["chat-container"]):
-                chatbot = gr.Chatbot(elem_classes=["chatbox"])
-                with gr.Row():
-                    msg = gr.Textbox(
-                        show_label=False,
-                        placeholder="Type your message here...",
-                        container=False
-                    )
-                    submit_btn = gr.Button("Send", variant="primary")
+            # Chat Panel
+            with gr.Column(scale=2):
+                chatbot = gr.Chatbot()
+                msg = gr.Textbox(
+                    show_label=False,
+                    placeholder="Type your message here..."
+                )
+                submit_button = gr.Button("Send")
     
-    # Event handlers
-    def update_role_display(role):
-        if role == "admin":
-            return f"<h3>Role: <span class='role-badge admin-badge'>Administrator</span></h3>"
-        else:
-            return f"<h3>Role: <span class='role-badge user-badge'>Standard User</span></h3>"
+    # Event Handlers
+    def update_interface(role):
+        settings_class = [] if role == "admin" else ["setting-disabled"]
+        return (
+            f"Current Role: {role.title()}",
+            gr.update(elem_classes=settings_class)
+        )
     
-    def handle_role_permissions(role):
-        # Only disable settings for non-admin users
-        if role == "admin":
-            return gr.update(visible=True, elem_classes=[])  # Remove any disabled classes
-        else:
-            return gr.update(visible=True, elem_classes=["setting-disabled"])
+    def chat_wrapper(message, history, uid, role):
+        if not message.strip():
+            return history
+            
+        system_msg = system_message.value
+        max_tok = max_tokens.value if role == "admin" else 512
+        temp = temperature.value if role == "admin" else 0.7
+        tp = top_p.value if role == "admin" else 0.95
+        
+        return chat(message, history, system_msg, max_tok, temp, tp, uid)
     
-    # Login handler with role-based permission setting
+    # Login flow
     login_button.click(
         login,
         inputs=[username, password],
         outputs=[login_group, chat_group, user_id, user_role, error_message]
     ).then(
-        update_role_display,
-        inputs=[user_role],
-        outputs=[role_indicator]
-    ).then(
-        handle_role_permissions,
+        update_interface,
         inputs=[user_role],
-        outputs=[admin_settings]
+        outputs=[role_display, settings_group]
     )
     
-    # Chat functionality
-    def chat_with_saved_params(message, history, uid, role):
-        # Get current values of parameters
-        system_msg = system_message.value
-        tokens = max_tokens.value if role == "admin" else 512  # Default value for non-admin
-        temp = temperature.value if role == "admin" else 0.7   # Default value for non-admin
-        topp = top_p.value if role == "admin" else 0.95       # Default value for non-admin
-        
-        return respond(message, history, system_msg, tokens, temp, topp, uid)
-    
-    msg_and_submit = msg.submit(
-        chat_with_saved_params,
+    # Chat flow
+    submit_button.click(
+        chat_wrapper,
         inputs=[msg, chatbot, user_id, user_role],
-        outputs=[chatbot],
-        show_progress=True
-    )
+        outputs=[chatbot]
+    ).then(lambda: "", None, msg)
     
-    submit_click = submit_btn.click(
-        chat_with_saved_params,
+    msg.submit(
+        chat_wrapper,
         inputs=[msg, chatbot, user_id, user_role],
-        outputs=[chatbot],
-        show_progress=True
-    )
-    
-    # Clear the textbox when message is sent
-    msg_and_submit.then(lambda: gr.update(value=""), None, [msg])
-    submit_click.then(lambda: gr.update(value=""), None, [msg])
-    
-    # Clear chat button - fixed to properly reset the history list
-    clear_btn.click(lambda: [], None, chatbot, queue=False)
-    
-    # Export conversation
-    def export_conversation(chat_history, uid):
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        filename = f"conversations/export_{uid}_{timestamp}.json"
-        os.makedirs(os.path.dirname(filename), exist_ok=True)
-        with open(filename, 'w') as f:
-            json.dump(chat_history, f)
-        logger.info(f"Exported conversation for user {uid}")
-        return gr.update(value=f"Conversation exported to {filename}", visible=True)
-    
-    export_btn.click(
-        export_conversation,
-        inputs=[chatbot, user_id],
-        outputs=[error_message]
-    )
+        outputs=[chatbot]
+    ).then(lambda: "", None, msg)
     
-    # Logout functionality
-    def logout():
-        return gr.update(visible=True), gr.update(visible=False), None, None
+    # Clear chat
+    clear_button.click(lambda: None, None, chatbot)
     
-    logout_btn.click(
-        logout,
+    # Logout
+    logout_button.click(
+        lambda: (
+            gr.update(visible=True),
+            gr.update(visible=False),
+            None,
+            None
+        ),
         outputs=[login_group, chat_group, user_id, user_role]
     )
 
-# Generate CSS file
-with open("styles.css", "w") as f:
-    f.write(css)
-
-css = css.replace(
-    ".setting-disabled {",
-    ".setting-disabled input, .setting-disabled button, .setting-disabled select, .setting-disabled textarea {"
-)
-
 if __name__ == "__main__":
-    # Check if we're running in production
-    if os.environ.get("PRODUCTION", "false").lower() == "true":
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=int(os.environ.get("PORT", 7860)),
-            share=False,
-            show_error=False,
-            auth=None,  # We handle auth in the app
-        )
-    else:
-        # Development mode
-        demo.launch(share=True)
+    demo.launch(share=True)